scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Merge pull request #37439 from BerriAI/litellm_decrease_anys_fable_round3

Browse source 
chore(typing): drop 1.3k basedpyright errors across 42 Any hotspot files
This commit is contained in:
Mateo Wang 2026-08-19 12:07:26 -07:00 • committed by GitHub
commit 4d100bdd89
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
45 changed files with 1660 additions and 651 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

24
basedpyright-code-budget.json
View file

@ -1,12 +1,12 @@
{
"reportAny": {
"limit": 22343
"limit": 19955
},
"reportArgumentType": {
"limit": 2578
"limit": 2566
},
"reportAssignmentType": {
"limit": 323
"limit": 320
},
"reportAttributeAccessIssue": {
"limit": 488
@ -24,7 +24,7 @@
"limit": 19
},
"reportExplicitAny": {
"limit": 6991
"limit": 6049
},
"reportFunctionMemberAccess": {
"limit": 7
@ -54,10 +54,10 @@
"limit": 0
},
"reportMissingParameterType": {
"limit": 5681
"limit": 5663
},
"reportMissingTypeArgument": {
"limit": 15605
"limit": 15557
},
"reportMissingTypeStubs": {
"limit": 40
@ -99,19 +99,19 @@
"limit": 0
},
"reportUnknownArgumentType": {
"limit": 44709
"limit": 44655
},
"reportUnknownLambdaType": {
"limit": 112
"limit": 109
},
"reportUnknownMemberType": {
"limit": 39154
"limit": 39043
},
"reportUnknownParameterType": {
"limit": 19944
"limit": 19887
},
"reportUnknownVariableType": {
"limit": 30772
"limit": 30574
},
"reportUnnecessaryCast": {
"limit": 117
@ -123,7 +123,7 @@
"limit": 5
},
"reportUnnecessaryIsInstance": {
"limit": 851
"limit": 836
},
"reportUntypedBaseClass": {
"limit": 0

45
enterprise/litellm_enterprise/proxy/vector_stores/endpoints.py
View file

@ -10,7 +10,8 @@ All /vector_store management endpoints
import copy
import json
from typing import List, Optional
from collections.abc import Mapping
from typing import TYPE_CHECKING, Final, List, Optional, Protocol
from fastapi import APIRouter, Depends, HTTPException
@ -32,9 +33,35 @@ from litellm.types.vector_stores import (
)
from litellm.vector_stores.vector_store_registry import VectorStoreRegistry
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
router = APIRouter()
class ManagedVectorStoreRow(Protocol):
"""A ``litellm_managedvectorstorestable`` row as returned by Prisma."""
def model_dump(self) -> LiteLLM_ManagedVectorStore: ...
class ManagedVectorStoreTable(Protocol):
"""The Prisma actions namespace for ``litellm_managedvectorstorestable``."""
async def find_unique(self, where: Mapping[str, str | None]) -> ManagedVectorStoreRow | None: ...
async def create(self, data: Mapping[str, object]) -> ManagedVectorStoreRow: ...
async def delete(self, where: Mapping[str, str | None]) -> ManagedVectorStoreRow | None: ...
async def update(self, where: Mapping[str, str | None], data: Mapping[str, object]) -> ManagedVectorStoreRow: ...
def managed_vector_store_table(prisma_client: "PrismaClient") -> ManagedVectorStoreTable:
"""The Prisma table actions for managed vector stores, behind a typed surface."""
return prisma_client.db.litellm_managedvectorstorestable
########################################################
# Management Endpoints
########################################################
@ -66,7 +93,7 @@ async def new_vector_store(
try:
# Check if vector store already exists
existing_vector_store = (
await prisma_client.db.litellm_managedvectorstorestable.find_unique(
await managed_vector_store_table(prisma_client).find_unique(
where={"vector_store_id": vector_store.get("vector_store_id")}
)
)
@ -92,7 +119,7 @@ async def new_vector_store(
del vector_store["litellm_params"]
_new_vector_store = (
await prisma_client.db.litellm_managedvectorstorestable.create(
await managed_vector_store_table(prisma_client).create(
data={
**vector_store,
"litellm_params": litellm_params_json,
@ -213,7 +240,7 @@ async def delete_vector_store(
try:
# Check if vector store exists
existing_vector_store = (
await prisma_client.db.litellm_managedvectorstorestable.find_unique(
await managed_vector_store_table(prisma_client).find_unique(
where={"vector_store_id": data.vector_store_id}
)
)
@ -224,7 +251,7 @@ async def delete_vector_store(
)
# Delete vector store
await prisma_client.db.litellm_managedvectorstorestable.delete(
await managed_vector_store_table(prisma_client).delete(
where={"vector_store_id": data.vector_store_id}
)
@ -288,7 +315,7 @@ async def get_vector_store_info(
return {"vector_store": vector_store_pydantic_obj}
vector_store = (
await prisma_client.db.litellm_managedvectorstorestable.find_unique(
await managed_vector_store_table(prisma_client).find_unique(
where={"vector_store_id": data.vector_store_id}
)
)
@ -298,7 +325,7 @@ async def get_vector_store_info(
detail=f"Vector store with ID {data.vector_store_id} not found",
)
vector_store_dict = vector_store.model_dump() # type: ignore[attr-defined]
vector_store_dict = vector_store.model_dump()
return {"vector_store": vector_store_dict}
except Exception as e:
verbose_proxy_logger.exception(f"Error getting vector store info: {str(e)}")
@ -322,13 +349,13 @@ async def update_vector_store(
try:
update_data = data.model_dump(exclude_unset=True)
vector_store_id = update_data.pop("vector_store_id")
vector_store_id: Final[str] = update_data.pop("vector_store_id")
if update_data.get("vector_store_metadata") is not None:
update_data["vector_store_metadata"] = safe_dumps(
update_data["vector_store_metadata"]
)
updated = await prisma_client.db.litellm_managedvectorstorestable.update(
updated = await managed_vector_store_table(prisma_client).update(
where={"vector_store_id": vector_store_id},
data=update_data,
)

112
litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py
View file

@ -7,9 +7,10 @@ import hashlib
import json
import time
from collections.abc import AsyncIterator
from typing import Any, Final, NamedTuple, cast
from typing import Any, Final, NamedTuple, Protocol
import httpx
from typing_extensions import NotRequired, ReadOnly, TypedDict
from litellm._logging import verbose_logger
from litellm.a2a_protocol.providers.watsonx_orchestrate.transformation import (
@ -38,11 +39,59 @@ class WXORequestParams(NamedTuple):
thread_id: str | None
class WXOLitellmParams(TypedDict, total=False):
"""litellm_params keys read when routing an A2A request to watsonx Orchestrate."""
cp4d_host: ReadOnly[str]
instance_id: ReadOnly[str]
wxo_agent_id: ReadOnly[str]
api_key: ReadOnly[str]
username: ReadOnly[str | None]
auth_mode: ReadOnly[str]
thread_id: ReadOnly[str | None]
class _IBMCloudTokenBody(TypedDict):
"""Fields read from the IBM Cloud IAM token response."""
access_token: ReadOnly[str]
expires_in: ReadOnly[NotRequired[int]]
class _CP4DTokenBody(TypedDict):
"""Fields read from the CP4D authorize response."""
token: ReadOnly[str]
expiration: ReadOnly[NotRequired[float]]
class _WXORun(TypedDict, total=False):
"""Fields the handler reads from a WXO run object or run event."""
status: ReadOnly[str]
run_id: ReadOnly[str]
id: ReadOnly[str]
class _SSELineSource(Protocol):
def aiter_lines(self) -> AsyncIterator[str]: ...
class _WXOView(TypedDict, total=False):
"""Typed reads of otherwise untyped watsonx Orchestrate and httpx values."""
ibm_cloud_token: ReadOnly[_IBMCloudTokenBody]
cp4d_token: ReadOnly[_CP4DTokenBody]
run: ReadOnly[_WXORun]
content_type: ReadOnly[str]
sse_source: ReadOnly[_SSELineSource]
class WatsonxOrchestrateHandler:
@staticmethod
def _http_client(timeout: float = 90.0) -> AsyncHTTPHandler:
return get_async_httpx_client(
llm_provider=cast(Any, httpxSpecialProvider.A2AProvider),
llm_provider=httpxSpecialProvider.A2AProvider,
params={"timeout": timeout},
)
@ -57,7 +106,7 @@ class WatsonxOrchestrateHandler:
return hashlib.sha256(material.encode()).hexdigest()
@staticmethod
def _cp4d_token_ttl_seconds(expiration: Any, now_wall: float | None = None) -> int:
def _cp4d_token_ttl_seconds(expiration: float, now_wall: float | None = None) -> int:
# CP4D returns expiration as absolute Unix epoch seconds, not a duration.
expires_at: Final = int(expiration)
wall: Final = now_wall if now_wall is not None else time.time()
@ -90,9 +139,9 @@ class WatsonxOrchestrateHandler:
headers={"Content-Type": "application/x-www-form-urlencoded"},
)
response.raise_for_status()
payload = response.json()
token = str(payload["access_token"])
ttl_s = int(payload.get("expires_in", 3600))
iam_payload: Final[_WXOView] = {"ibm_cloud_token": response.json()}
token = str(iam_payload["ibm_cloud_token"]["access_token"])
ttl_s = int(iam_payload["ibm_cloud_token"].get("expires_in", 3600))
else:
if not username:
raise ValueError("'username' is required in litellm_params when auth_mode='cp4d'")
@ -103,9 +152,9 @@ class WatsonxOrchestrateHandler:
headers={"Content-Type": "application/json"},
)
response.raise_for_status()
payload = response.json()
token = str(payload["token"])
expiration: Final = payload.get("expiration")
cp4d_payload: Final[_WXOView] = {"cp4d_token": response.json()}
token = str(cp4d_payload["cp4d_token"]["token"])
expiration: Final = cp4d_payload["cp4d_token"].get("expiration")
if expiration is None:
ttl_s = 3600
else:
@ -118,6 +167,16 @@ class WatsonxOrchestrateHandler:
del _token_cache[stale_key]
return token
@staticmethod
def _run_body(response: httpx.Response) -> _WXORun:
view: Final[_WXOView] = {"run": response.json()}
return view["run"]
@staticmethod
def _decode_run_event(payload: str | bytes) -> _WXORun:
view: Final[_WXOView] = {"run": json.loads(payload)}
return view["run"]
@staticmethod
async def _poll_run(
base_url: str,
@ -126,14 +185,14 @@ class WatsonxOrchestrateHandler:
client: AsyncHTTPHandler,
max_attempts: int = _MAX_POLL_ATTEMPTS,
interval_s: float = _POLL_INTERVAL_S,
) -> dict[str, Any]:
) -> _WXORun:
url: Final = f"{base_url}/v1/orchestrate/runs/{run_id}"
for attempt in range(max_attempts):
await asyncio.sleep(interval_s)
response = await client.get(url, headers=auth_headers)
response.raise_for_status()
result: dict[str, Any] = response.json()
result = WatsonxOrchestrateHandler._run_body(response)
status = result.get("status", "")
verbose_logger.debug("WXO: Poll %s/%s run='%s' status='%s'", attempt + 1, max_attempts, run_id, status)
if status in WatsonxOrchestrateTransformation.TERMINAL_STATES:
@ -145,11 +204,11 @@ class WatsonxOrchestrateHandler:
@staticmethod
async def _get_successful_run_data(
run_data: dict[str, Any],
run_data: _WXORun,
base_url: str,
auth_headers: dict[str, str],
client: AsyncHTTPHandler,
) -> dict[str, Any]:
) -> _WXORun:
status = run_data.get("status", "")
if status not in WatsonxOrchestrateTransformation.TERMINAL_STATES:
run_id: Final = run_data.get("run_id") or run_data.get("id") or ""
@ -170,15 +229,16 @@ class WatsonxOrchestrateHandler:
@staticmethod
async def _accumulate_wxo_sse_text(response: Any) -> str:
source: Final[_WXOView] = {"sse_source": response}
accumulated_text = ""
async for line in response.aiter_lines():
async for line in source["sse_source"].aiter_lines():
if not line.startswith("data:"):
continue
data_str = line[5:].strip()
if not data_str or data_str == "[DONE]":
continue
try:
event = json.loads(data_str)
event = WatsonxOrchestrateHandler._decode_run_event(data_str)
except json.JSONDecodeError:
continue
chunk_text = WatsonxOrchestrateTransformation.extract_text_from_wxo_result(event)
@ -187,7 +247,7 @@ class WatsonxOrchestrateHandler:
return accumulated_text
@staticmethod
def _extract_litellm_params(litellm_params: dict[str, Any]) -> WXORequestParams:
def _extract_litellm_params(litellm_params: WXOLitellmParams) -> WXORequestParams:
cp4d_host: Final = litellm_params.get("cp4d_host") or ""
instance_id: Final = litellm_params.get("instance_id") or ""
wxo_agent_id: Final = litellm_params.get("wxo_agent_id") or ""
@ -215,9 +275,9 @@ class WatsonxOrchestrateHandler:
@staticmethod
async def handle_non_streaming(
request_id: str,
params: dict[str, Any],
litellm_params: dict[str, Any],
) -> dict[str, Any]:
params: dict[str, object],
litellm_params: WXOLitellmParams,
) -> dict[str, object]:
wxo: Final = WatsonxOrchestrateHandler._extract_litellm_params(litellm_params)
client: Final = WatsonxOrchestrateHandler._http_client(timeout=90.0)
@ -246,7 +306,8 @@ class WatsonxOrchestrateHandler:
headers=auth_headers,
)
run_response.raise_for_status()
run_data: dict[str, Any] = run_response.json()
started: Final[_WXOView] = {"run": run_response.json()}
run_data: _WXORun = started["run"]
run_data = await WatsonxOrchestrateHandler._get_successful_run_data(
run_data=run_data,
@ -261,11 +322,11 @@ class WatsonxOrchestrateHandler:
@staticmethod
async def handle_streaming(
request_id: str,
params: dict[str, Any],
litellm_params: dict[str, Any],
params: dict[str, object],
litellm_params: WXOLitellmParams,
chunk_size: int = 50,
delay_ms: int = 10,
) -> AsyncIterator[dict[str, Any]]:
) -> AsyncIterator[dict[str, object]]:
wxo: Final = WatsonxOrchestrateHandler._extract_litellm_params(litellm_params)
client: Final = WatsonxOrchestrateHandler._http_client(timeout=120.0)
@ -316,10 +377,11 @@ class WatsonxOrchestrateHandler:
yield chunk
return
content_type: Final = response.headers.get("content-type", "").lower()
header_view: Final[_WXOView] = {"content_type": response.headers.get("content-type", "")}
content_type: Final = header_view["content_type"].lower()
if "text/event-stream" not in content_type:
response_body: Final = await response.aread()
result = json.loads(response_body)
result = WatsonxOrchestrateHandler._decode_run_event(response_body)
result = await WatsonxOrchestrateHandler._get_successful_run_data(
run_data=result,
base_url=base_url,

47
litellm/caching/caching_handler.py
View file

@ -18,7 +18,7 @@ import asyncio
import datetime
import inspect
import time
from collections.abc import AsyncGenerator, AsyncIterator, Callable, Generator
from collections.abc import AsyncGenerator, AsyncIterator, Callable, Generator, Mapping
from typing import TYPE_CHECKING, Any, Final, Optional, TypeVar
from pydantic import BaseModel
@ -106,7 +106,7 @@ def _is_chat_completion_cached_dict(cached_result: dict) -> bool:
return "choices" in cached_result
def _should_defer_streaming_cache_hit_callbacks(*, kwargs: dict[str, Any]) -> bool:
def _should_defer_streaming_cache_hit_callbacks(*, kwargs: dict[str, object]) -> bool:
"""
When stream=True, do not run success callbacks at cache-hit time.
@ -119,11 +119,21 @@ def _should_defer_streaming_cache_hit_callbacks(*, kwargs: dict[str, Any]) -> bo
return kwargs.get("stream", False) is True
def _prompt_tokens_details_as_mapping(details: "PromptTokensDetailsWrapper") -> Mapping[str, object]:
"""Dump prompt token details to an opaque field mapping, tolerating non-pydantic stand-ins."""
return details.model_dump(exclude_none=True) if hasattr(details, "model_dump") else {}
def _request_cache_key(request_kwargs: Mapping[str, Any]) -> str | None:
"""Read the caller-supplied ``cache_key`` off the request kwargs."""
return request_kwargs.get("cache_key", None)
class LLMCachingHandler:
def __init__(
self,
original_function: Callable,
request_kwargs: dict[str, Any],
request_kwargs: dict[str, object],
start_time: datetime.datetime,
):
from litellm.caching import DualCache, RedisCache
@ -150,7 +160,7 @@ class LLMCachingHandler:
start_time: datetime.datetime,
call_type: str,
kwargs: dict[str, Any],
args: tuple[Any, ...] | None = None,
args: tuple[object, ...] | None = None,
) -> CachingHandlerResponse | None:
"""
Internal method to get from the cache.
@ -289,7 +299,7 @@ class LLMCachingHandler:
start_time: datetime.datetime,
call_type: str,
kwargs: dict[str, Any],
args: tuple[Any, ...] | None = None,
args: tuple[object, ...] | None = None,
) -> CachingHandlerResponse:
cached_result: Any | None = None
@ -366,7 +376,7 @@ class LLMCachingHandler:
return CachingHandlerResponse(cached_result=cached_result)
return CachingHandlerResponse(cached_result=cached_result)
def handle_kwargs_input_list_or_str(self, kwargs: dict[str, Any]) -> list[str]:
def handle_kwargs_input_list_or_str(self, kwargs: dict[str, object]) -> list[str]:
"""
Handles the input of kwargs['input'] being a list or a string
"""
@ -548,8 +558,8 @@ class LLMCachingHandler:
if details2 is None:
return details1
dict1: Final = details1.model_dump(exclude_none=True) if hasattr(details1, "model_dump") else {}
dict2: Final = details2.model_dump(exclude_none=True) if hasattr(details2, "model_dump") else {}
dict1: Final = _prompt_tokens_details_as_mapping(details1)
dict2: Final = _prompt_tokens_details_as_mapping(details2)
merged: Final[dict] = {}
for key in set(dict1.keys()) | set(dict2.keys()):
@ -671,7 +681,9 @@ class LLMCachingHandler:
cache_hit=cache_hit,
)
async def _retrieve_from_cache(self, call_type: str, kwargs: dict[str, Any], args: tuple[Any, ...]) -> Any | None:
async def _retrieve_from_cache(
self, call_type: str, kwargs: dict[str, object], args: tuple[object, ...]
) -> Any | None:
"""
Internal method to
- get cache key
@ -727,7 +739,8 @@ class LLMCachingHandler:
cached_result = None
else:
request_kwargs: Final = new_kwargs.copy()
request_cache_key: Final = request_kwargs.pop("cache_key", None)
request_cache_key: Final = _request_cache_key(request_kwargs)
request_kwargs.pop("cache_key", None)
if litellm.cache._supports_async() is True:
## check if dual cache is supported ##
self.preset_cache_key = request_cache_key or litellm.cache.get_cache_key(**request_kwargs)
@ -749,10 +762,10 @@ class LLMCachingHandler:
self,
cached_result: Any,
call_type: str,
kwargs: dict[str, Any],
kwargs: dict[str, object],
logging_obj: LiteLLMLoggingObj,
model: str,
args: tuple[Any, ...],
args: tuple[object, ...],
custom_llm_provider: str | None = None,
) -> (
ModelResponse
@ -948,7 +961,7 @@ class LLMCachingHandler:
result: Any,
original_function: Callable,
kwargs: dict[str, Any],
args: tuple[Any, ...] | None = None,
args: tuple[object, ...] | None = None,
):
"""
Internal method to check the type of the result & cache used and adds the result to the cache accordingly
@ -1013,8 +1026,8 @@ class LLMCachingHandler:
def sync_set_cache(
self,
result: Any,
kwargs: dict[str, Any],
args: tuple[Any, ...] | None = None,
kwargs: dict[str, object],
args: tuple[object, ...] | None = None,
):
"""
Sync internal method to add the result to the cache
@ -1204,8 +1217,8 @@ class LLMCachingHandler:
def convert_args_to_kwargs(
original_function: Callable,
args: tuple[Any, ...] | None = None,
) -> dict[str, Any]:
args: tuple[object, ...] | None = None,
) -> dict[str, object]:
# Get the signature of the original function
signature: Final = inspect.signature(original_function)

43
litellm/cost_calculator.py
View file

@ -102,6 +102,7 @@ from litellm.types.utils import (
LlmProviders,
LlmProvidersSet,
ModelInfo,
PromptTokensDetailsWrapper,
ServiceTier,
StandardBuiltInToolsParams,
TranscriptionUsageDurationObject,
@ -286,7 +287,7 @@ def _transcription_usage_has_token_details(
prompt_tokens_val: Final = getattr(usage_block, "prompt_tokens", 0) or 0
completion_tokens_val: Final = getattr(usage_block, "completion_tokens", 0) or 0
prompt_details: Final = getattr(usage_block, "prompt_tokens_details", None)
prompt_details: Final[PromptTokensDetailsWrapper | None] = getattr(usage_block, "prompt_tokens_details", None)
if prompt_details is not None:
audio_token_count: Final = getattr(prompt_details, "audio_tokens", 0) or 0
@ -375,7 +376,7 @@ def cost_per_token(
_is_anthropic_style = False
if usage_object is not None:
_pt_details: Final = getattr(usage_object, "prompt_tokens_details", None)
_pt_details: Final[PromptTokensDetailsWrapper | None] = getattr(usage_object, "prompt_tokens_details", None)
if _pt_details is not None:
_cache_read_tokens = float(getattr(_pt_details, "cached_tokens", 0) or 0)
# OpenAI-compatible providers report cache-write tokens under
@ -385,8 +386,8 @@ def cost_per_token(
getattr(_pt_details, "cache_write_tokens", 0) or getattr(_pt_details, "cache_creation_tokens", 0) or 0
)
_anthropic_read: Final = getattr(usage_object, "cache_read_input_tokens", None)
_anthropic_create: Final = getattr(usage_object, "cache_creation_input_tokens", None)
_anthropic_read: Final[int | None] = getattr(usage_object, "cache_read_input_tokens", None)
_anthropic_create: Final[int | None] = getattr(usage_object, "cache_creation_input_tokens", None)
if _anthropic_read is not None or _anthropic_create is not None:
_is_anthropic_style = True
if _anthropic_read is not None:
@ -703,7 +704,7 @@ def get_replicate_completion_pricing(completion_response: dict, total_time=0.0):
return a100_80gb_price_per_second_public * total_time / 1000
def has_hidden_params(obj: Any) -> bool:
def has_hidden_params(obj: object) -> bool:
return hasattr(obj, "_hidden_params")
@ -728,7 +729,7 @@ def _get_provider_for_cost_calc(
def _select_model_name_for_cost_calc(
model: str | None,
completion_response: Any | None,
completion_response: object | None,
base_model: str | None = None,
custom_pricing: bool | None = None,
custom_llm_provider: str | None = None,
@ -804,7 +805,7 @@ def _model_contains_known_llm_provider(model: str) -> bool:
return _provider_prefix in LlmProvidersSet
def _get_response_model(completion_response: Any) -> str | None:
def _get_response_model(completion_response: object) -> str | None:
"""
Extract the model name from a completion response object.
@ -866,8 +867,18 @@ def _normalize_service_tier(service_tier: object) -> str | None:
return service_tier
def _extract_service_tier(source: object) -> str | None:
"""Read a raw ``service_tier`` off a response body or usage object, dict or pydantic model alike."""
if isinstance(source, BaseModel):
return getattr(source, "service_tier", None)
elif isinstance(source, dict):
return source.get("service_tier")
return None
def _get_usage_object(
completion_response: Any,
completion_response: object,
) -> Usage | None:
usage_obj: Final = cast(
Usage | ResponseAPIUsage | dict | BaseModel,
@ -1110,7 +1121,7 @@ def _store_cost_breakdown_in_logging_obj(
def completion_cost(
completion_response=None,
completion_response: object | None = None,
model: str | None = None,
prompt="",
messages: list = [],
@ -1197,19 +1208,13 @@ def completion_cost(
# Extract service_tier from completion_response if not provided
if service_tier is None and completion_response is not None:
if isinstance(completion_response, BaseModel):
service_tier = getattr(completion_response, "service_tier", None)
elif isinstance(completion_response, dict):
service_tier = completion_response.get("service_tier")
service_tier = _extract_service_tier(completion_response)
service_tier = _normalize_service_tier(service_tier)
# Extract service_tier from usage object if not provided
if service_tier is None and cost_per_token_usage_object is not None:
if isinstance(cost_per_token_usage_object, BaseModel):
service_tier = getattr(cost_per_token_usage_object, "service_tier", None)
elif isinstance(cost_per_token_usage_object, dict):
service_tier = cost_per_token_usage_object.get("service_tier")
service_tier = _extract_service_tier(cost_per_token_usage_object)
service_tier = _normalize_service_tier(service_tier)
@ -1412,7 +1417,7 @@ def completion_cost(
if completion_response is not None and isinstance(completion_response, RerankResponse):
meta_obj = completion_response.meta
if meta_obj is not None:
billed_units = meta_obj.get("billed_units", {}) or {}
billed_units: RerankBilledUnits = meta_obj.get("billed_units") or {}
else:
billed_units = {}
@ -1801,7 +1806,7 @@ def response_cost_calculator(
def ocr_cost(
model: str,
custom_llm_provider: str | None,
response: Any | None = None,
response: object | None = None,
) -> tuple[float, float]:
"""
Args:

48
litellm/google_genai/adapters/transformation.py
View file

@ -1,5 +1,5 @@
import json
from collections.abc import AsyncIterator, Iterator
from collections.abc import AsyncIterator, Iterator, Sequence
from typing import Any, Final, TypedDict, cast
from typing_extensions import ReadOnly
@ -27,6 +27,7 @@ from litellm.types.utils import (
ModelResponse,
ModelResponseStream,
StreamingChoices,
Usage,
)
@ -43,6 +44,29 @@ class _GenAIPart(TypedDict, total=False):
functionCall: ReadOnly[dict[str, object]]
class _GenAIFunctionDeclaration(TypedDict, total=False):
name: ReadOnly[str]
description: ReadOnly[str]
parametersJsonSchema: ReadOnly[dict[str, object]]
class _GenAITool(TypedDict, total=False):
functionDeclarations: ReadOnly[list[_GenAIFunctionDeclaration]]
class _GenAIFunctionCallingConfig(TypedDict, total=False):
mode: ReadOnly[str]
class _GenAIToolConfig(TypedDict, total=False):
functionCallingConfig: ReadOnly[_GenAIFunctionCallingConfig]
def _decode_tool_call_arguments(raw_arguments: str) -> object:
"""Decode a tool call's JSON-encoded arguments into the value Google GenAI expects."""
return json.loads(raw_arguments)
class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
"""
Wrapper for streaming Google GenAI generate_content responses.
@ -51,7 +75,7 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
sent_first_chunk: bool = False
# State tracking for accumulating partial tool calls
accumulated_tool_calls: dict[str, dict[str, str]]
accumulated_tool_calls: dict[int, dict[str, str]]
def __init__(self, completion_stream: object):
self.sent_first_chunk = False
@ -108,7 +132,7 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
try:
# For tool calls with no arguments, accumulated_args will be "", which is not valid JSON.
# We default to an empty JSON object in this case.
parsed_args = json.loads(tool_call_data["arguments"] or "{}")
parsed_args = _decode_tool_call_arguments(tool_call_data["arguments"] or "{}")
function_call_part: _GenAIPart = {
"functionCall": {
"name": tool_call_data["name"] or "undefined_tool_name",
@ -319,7 +343,7 @@ class GoogleGenAIAdapter:
def _transform_google_genai_tools_to_openai(
self,
tools: list[dict[str, Any]],
tools: Sequence[_GenAITool],
) -> list[ChatCompletionToolParam]:
"""Transform Google GenAI tools to OpenAI tools format"""
openai_tools: Final[list[dict[str, object]]] = []
@ -346,7 +370,7 @@ class GoogleGenAIAdapter:
def _transform_google_genai_tool_config_to_openai(
self,
tool_config: dict[str, Any],
tool_config: _GenAIToolConfig,
) -> ChatCompletionToolChoiceValues | None:
"""Transform Google GenAI tool_config to OpenAI tool_choice"""
function_calling_config: Final = tool_config.get("functionCallingConfig", {})
@ -563,7 +587,7 @@ class GoogleGenAIAdapter:
parts = self._transform_openai_delta_to_google_genai_parts_with_accumulation(choice.delta, wrapper)
else:
parts = []
finish_reason = getattr(choice, "finish_reason", None)
finish_reason: str | None = getattr(choice, "finish_reason", None)
else:
# Fallback for generic choice objects
message_content: Final = getattr(choice, "delta", {}).get("content", "")
@ -625,7 +649,11 @@ class GoogleGenAIAdapter:
for tool_call in message.tool_calls:
if hasattr(tool_call, "function") and tool_call.function:
try:
args = json.loads(tool_call.function.arguments) if tool_call.function.arguments else {}
args = (
_decode_tool_call_arguments(tool_call.function.arguments)
if tool_call.function.arguments
else {}
)
except json.JSONDecodeError:
args = {}
@ -661,7 +689,7 @@ class GoogleGenAIAdapter:
continue
# 3. Use `index` as the primary key for accumulation
tool_call_index = getattr(tool_call, "index", None)
tool_call_index: int | None = getattr(tool_call, "index", None)
if tool_call_index is None:
continue # Index is essential for tracking streaming tool calls
@ -695,7 +723,7 @@ class GoogleGenAIAdapter:
# 5. Attempt to parse arguments even if name hasn't arrived.
try:
# Attempt to parse the accumulated arguments string
parsed_args = json.loads(accumulated_args)
parsed_args = _decode_tool_call_arguments(accumulated_args)
# If parsing succeeds, but we don't have a name yet, wait.
# The part will be created by a later chunk that brings the name.
@ -729,7 +757,7 @@ class GoogleGenAIAdapter:
return mapping.get(finish_reason, "STOP")
def _map_usage(self, usage: Any) -> dict[str, int]:
def _map_usage(self, usage: Usage | None) -> dict[str, int]:
"""Map OpenAI usage to Google GenAI usage format"""
return {
"promptTokenCount": getattr(usage, "prompt_tokens", 0) or 0,

30
litellm/integrations/galileo.py
View file

@ -60,13 +60,13 @@ class LLMResponse(BaseModel):
default=None,
description="Total cost of the LLM call in USD as computed by LiteLLM.",
)
output_logprobs: dict[str, Any] | None = Field(
output_logprobs: dict[str, object] | None = Field(
default=None,
description="Optional. When available, logprobs are used to compute Uncertainty.",
)
created_at: str = Field(..., description='timestamp constructed in "%Y-%m-%dT%H:%M:%S" format')
tags: list[str] | None = None
user_metadata: dict[str, Any] | None = None
user_metadata: dict[str, object] | None = None
class GalileoObserve(CustomLogger):
@ -238,13 +238,13 @@ class GalileoObserve(CustomLogger):
return created_at
@staticmethod
def _token_metrics_from_record(record: Mapping[str, Any]) -> dict[str, Any]:
def _token_metrics_from_record(record: Mapping[str, Any]) -> dict[str, object]:
num_input_tokens: Final = int(record.get("num_input_tokens") or 0)
num_output_tokens: Final = int(record.get("num_output_tokens") or 0)
num_total_tokens = int(record.get("num_total_tokens") or 0)
if num_total_tokens == 0 and (num_input_tokens or num_output_tokens):
num_total_tokens = num_input_tokens + num_output_tokens
metrics: Final[dict[str, Any]] = {
metrics: Final[dict[str, object]] = {
"num_input_tokens": num_input_tokens,
"num_output_tokens": num_output_tokens,
"num_total_tokens": num_total_tokens,
@ -260,10 +260,10 @@ class GalileoObserve(CustomLogger):
*,
trace_id: str,
span_id: str,
) -> dict[str, Any]:
) -> dict[str, object]:
created_at: Final = GalileoObserve._normalize_created_at(record.get("created_at", ""))
span: Final[dict[str, Any]] = {
span: Final[dict[str, object]] = {
"type": "llm",
"id": span_id,
"trace_id": trace_id,
@ -287,7 +287,7 @@ class GalileoObserve(CustomLogger):
return span
@staticmethod
def _record_to_v2_trace(record: Mapping[str, Any]) -> dict[str, Any]:
def _record_to_v2_trace(record: Mapping[str, Any]) -> dict[str, object]:
trace_id: Final = str(uuid.uuid4())
span_id: Final = str(uuid.uuid4())
created_at: Final = GalileoObserve._normalize_created_at(record.get("created_at", ""))
@ -307,8 +307,8 @@ class GalileoObserve(CustomLogger):
"spans": [GalileoObserve._record_to_v2_span(record, trace_id=trace_id, span_id=span_id)],
}
def _build_traces_payload(self, records: Sequence[Mapping[str, Any]]) -> dict[str, Any]:
payload: Final[dict[str, Any]] = {
def _build_traces_payload(self, records: Sequence[Mapping[str, object]]) -> dict[str, object]:
payload: Final[dict[str, object]] = {
"traces": [self._record_to_v2_trace(record) for record in records],
"logging_method": "api_direct",
"reliable": False,
@ -318,7 +318,7 @@ class GalileoObserve(CustomLogger):
payload["log_stream_id"] = self.log_stream_id
return payload
def _get_ingest_request(self) -> tuple[str, dict[str, Any]] | None:
def _get_ingest_request(self) -> tuple[str, dict[str, object]] | None:
if not self.base_url or not self.project_id:
return None
@ -427,9 +427,9 @@ class GalileoObserve(CustomLogger):
pass
@staticmethod
def _build_prompt(kwargs: Mapping[str, Any]) -> dict[str, Any]:
def _build_prompt(kwargs: Mapping[str, Any]) -> dict[str, object]:
optional_params: Final[Mapping[str, object]] = kwargs.get("optional_params", {}) or {}
prompt: Final[dict[str, Any]] = {"messages": kwargs.get("messages")}
prompt: Final[dict[str, object]] = {"messages": kwargs.get("messages")}
if optional_params.get("functions") is not None:
prompt["functions"] = optional_params["functions"]
if optional_params.get("tools") is not None:
@ -451,7 +451,7 @@ class GalileoObserve(CustomLogger):
return json.dumps(value, default=_json_default)
@staticmethod
def _prompt_to_input_text(prompt: Mapping[str, Any]) -> str:
def _prompt_to_input_text(prompt: Mapping[str, object]) -> str:
messages: Final[object] = prompt.get("messages")
if messages is not None:
text: Final = GalileoObserve._input_text_from_messages(messages)
@ -464,7 +464,7 @@ class GalileoObserve(CustomLogger):
if response_obj.choices and len(response_obj.choices) > 0:
message: Final = response_obj["choices"][0]["message"]
if hasattr(message, "json"):
message_json: Final = message.json()
message_json: Final[object] = message.json()
if isinstance(message_json, str):
return json.loads(message_json)
return message_json
@ -488,7 +488,7 @@ class GalileoObserve(CustomLogger):
return None
@staticmethod
def _langfuse_style_rerank_prompt(kwargs: Mapping[str, object]) -> dict[str, Any]:
def _langfuse_style_rerank_prompt(kwargs: Mapping[str, object]) -> dict[str, object]:
"""Match Langfuse rerank input: prompt = {"messages": kwargs.get("messages")}."""
return {"messages": kwargs.get("messages")}

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

@ -5,7 +5,7 @@ import traceback
from collections.abc import Callable, Iterable, Mapping
from datetime import datetime
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, cast
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast
from packaging.version import Version
@ -49,10 +49,21 @@ else:
_DENIED_STEERING_KEYS: Final = frozenset({"headers", "endpoint", "caching_groups", "previous_models"})
_NO_METADATA: Final[Mapping[str, Any]] = MappingProxyType({})
_NO_METADATA: Final[Mapping[str, object]] = MappingProxyType({})
_REDACTED_PROXY_HEADERS: Final[frozenset[str]] = frozenset({"authorization", "cookie", "referer"})
def _object_mapping(value: object) -> Mapping[str, object] | None:
"""Return ``value`` as an opaque mapping when it is a dict."""
return value if isinstance(value, dict) else None
class _UsageObject(Protocol):
"""Token-count surface the Langfuse logger reads off a response usage payload."""
def get(self, key: Literal["cache_creation_input_tokens", "cache_read_input_tokens"], /) -> int | None: ...
def _extract_cache_read_input_tokens(usage_obj) -> int:
"""
Extract cache_read_input_tokens from usage object.
@ -82,6 +93,11 @@ def _extract_cache_read_input_tokens(usage_obj) -> int:
return cache_read_input_tokens
def _logging_id(start_time: datetime | None, response_obj: object) -> str | None:
"""Typed view of the timestamped response id Langfuse uses as the generation id."""
return litellm.utils.get_logging_id(start_time, response_obj)
def _as_steering_flag(value: object) -> bool:
"""A string ``str_to_bool`` does not recognise falls back to its truthiness."""
if isinstance(value, str):
@ -222,7 +238,7 @@ class LangFuseLogger:
return langfuse_client
@staticmethod
def add_metadata_from_header(litellm_params: dict, metadata: dict) -> dict:
def add_metadata_from_header(litellm_params: dict, metadata: dict) -> dict[str, object]:
"""
Adds metadata from proxy request headers to Langfuse logging if keys start with "langfuse_"
and overwrites litellm_params.metadata if already included.
@ -494,7 +510,7 @@ class LangFuseLogger:
def _log_langfuse_v2(
self,
user_id: str | None,
metadata: dict,
metadata: dict[str, object],
litellm_params: dict,
output: str | dict | list | None,
start_time: datetime | None,
@ -519,7 +535,7 @@ class LangFuseLogger:
else []
)
allowlisted_metadata: Final[StandardLoggingMetadata | dict[str, Any]] = (
allowlisted_metadata: Final[StandardLoggingMetadata | Mapping[str, object]] = (
standard_logging_object["metadata"] if standard_logging_object is not None else _NO_METADATA
)
end_user_id: Final = allowlisted_metadata.get("user_api_key_end_user_id", None)
@ -531,11 +547,12 @@ class LangFuseLogger:
# Clean Metadata before logging - never log raw metadata
# the raw metadata can contain circular references which leads to infinite recursion
# we clean out all extra litellm metadata params before logging
clean_metadata: dict[str, Any] = {}
clean_metadata: dict[str, object] = {}
if prompt_management_metadata is not None:
clean_metadata["prompt_management_metadata"] = prompt_management_metadata
if isinstance(metadata, dict):
for key, value in metadata.items():
metadata_entries: Final = _object_mapping(metadata)
if metadata_entries is not None:
for key, value in metadata_entries.items():
# generate langfuse tags - Default Tags sent to Langfuse from LiteLLM Proxy
if (
litellm.langfuse_default_tags is not None
@ -705,8 +722,8 @@ class LangFuseLogger:
usage_details = None
if response_obj is not None:
if hasattr(response_obj, "id") and response_obj.get("id", None) is not None:
generation_id = litellm.utils.get_logging_id(start_time, response_obj)
_usage_obj: Final = getattr(response_obj, "usage", None)
generation_id = _logging_id(start_time, response_obj)
_usage_obj: Final[_UsageObject | None] = getattr(response_obj, "usage", None)
if _usage_obj:
# Safely get usage values, defaulting None to 0 for Langfuse compatibility.
@ -811,7 +828,7 @@ class LangFuseLogger:
@staticmethod
def _get_chat_content_for_langfuse(
response_obj: ModelResponse,
):
) -> str | None:
"""
Get the chat content for Langfuse logging
"""
@ -1078,7 +1095,7 @@ def log_provider_specific_information_as_span(
None
"""
_hidden_params: Final = clean_metadata.get("hidden_params", None)
_hidden_params: Final[Mapping[str, object] | None] = clean_metadata.get("hidden_params", None)
if _hidden_params is None:
return

29
litellm/integrations/otel/logger.py
View file

@ -62,8 +62,13 @@ from litellm.integrations.otel.plumbing.providers import (
from litellm.integrations.otel.plumbing.routing import TenantTracerCache
if TYPE_CHECKING:
from opentelemetry.metrics import MeterProvider
from litellm.caching.dual_cache import DualCache
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.services import ServiceLoggerPayload
from litellm.types.utils import (
CallTypesLiteral,
StandardLoggingGuardrailInformation,
StandardLoggingPayload,
)
@ -140,7 +145,7 @@ class OpenTelemetryV2(CustomLogger):
callback_name: str | None = None,
tracer_provider: TracerProvider | None = None,
logger_provider: LoggerProvider | None = None,
meter_provider: Any | None = None,
meter_provider: "MeterProvider | None" = None,
**kwargs: Any,
) -> None:
super().__init__(**kwargs)
@ -162,7 +167,7 @@ class OpenTelemetryV2(CustomLogger):
self._open_llm_calls: OrderedDict[str, _LLMCallSpan] = OrderedDict()
self._init_otel_logger_on_litellm_proxy()
def _init_metrics(self, meter_provider: Any | None) -> "GenAIMetricRecorder | None":
def _init_metrics(self, meter_provider: "MeterProvider | None") -> "GenAIMetricRecorder | None":
"""Create the six GenAI histograms when metrics are enabled, else ``None``.
``meter_provider`` is an explicit override (tests inject one); otherwise the
@ -340,7 +345,7 @@ class OpenTelemetryV2(CustomLogger):
def _emit_mcp_tool_call(
self,
kwargs: Mapping[str, Any],
kwargs: Mapping[str, object],
start_time: datetime | float | None,
end_time: datetime | float | None,
) -> bool:
@ -417,7 +422,7 @@ class OpenTelemetryV2(CustomLogger):
def _close_llm_call(
self,
kwargs: Mapping[str, Any],
kwargs: Mapping[str, object],
start_time: datetime | float | None,
end_time: datetime | float | None,
) -> Span | None:
@ -474,7 +479,7 @@ class OpenTelemetryV2(CustomLogger):
async def async_service_success_hook(
self,
payload: Any,
payload: "ServiceLoggerPayload",
parent_otel_span: Span | None = None,
start_time: datetime | float | None = None,
end_time: datetime | float | None = None,
@ -491,7 +496,7 @@ class OpenTelemetryV2(CustomLogger):
async def async_service_failure_hook(
self,
payload: Any,
payload: "ServiceLoggerPayload",
error: str | None = "",
parent_otel_span: Span | None = None,
start_time: datetime | float | None = None,
@ -509,7 +514,7 @@ class OpenTelemetryV2(CustomLogger):
def _emit_service(
self,
payload: Any,
payload: "ServiceLoggerPayload",
*,
parent_otel_span: Span | None,
start_time: datetime | float | None,
@ -559,7 +564,7 @@ class OpenTelemetryV2(CustomLogger):
# / errors are the FastAPI instrumentor's job, so we don't touch it here.
# ====================================================================== #
def seed_request_identity(self, user_api_key_dict: Any, model: Any = None) -> None:
def seed_request_identity(self, user_api_key_dict: object, model: str | None = None) -> None:
"""Attach request-identity Baggage to the current context + server span.
Seeding identity into Baggage makes **every** span emitted afterwards for
@ -615,10 +620,10 @@ class OpenTelemetryV2(CustomLogger):
async def async_pre_call_hook(
self,
user_api_key_dict: Any,
cache: Any,
user_api_key_dict: "UserAPIKeyAuth",
cache: "DualCache",
data: dict,
call_type: Any,
call_type: "CallTypesLiteral",
) -> dict:
self.seed_request_identity(
user_api_key_dict,
@ -790,7 +795,7 @@ def emit_guardrail_span(entry: "StandardLoggingGuardrailInformation") -> None:
pass
def seed_request_identity(user_api_key_dict: Any, model: Any = None) -> None:
def seed_request_identity(user_api_key_dict: object, model: str | None = None) -> None:
logger: Final = _registered_v2_logger()
if logger is not None:
logger.seed_request_identity(user_api_key_dict, model=model)

85
litellm/integrations/prometheus.py
View file

@ -9,7 +9,9 @@ import os
import sys
from collections.abc import Awaitable, Callable, Mapping, Sequence
from datetime import datetime, timedelta
from typing import TYPE_CHECKING, Any, Final, Literal, cast
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, cast
from pydantic import BaseModel
import litellm
from litellm._logging import print_verbose, verbose_logger
@ -38,6 +40,7 @@ from litellm.proxy._types import (
LiteLLM_UserTable,
UserAPIKeyAuth,
)
from litellm.repositories.base_repository import BaseRepository
from litellm.repositories.organization_repository import OrganizationRepository
from litellm.repositories.team_repository import TeamRepository
from litellm.repositories.user_repository import UserRepository
@ -58,6 +61,9 @@ if TYPE_CHECKING:
else:
AsyncIOScheduler = Any
_BudgetRowT: Final = TypeVar("_BudgetRowT")
_TableRowT: Final = TypeVar("_TableRowT", bound=BaseModel)
_DEFAULT_BUDGET_METRICS_PER_REQUEST_TIMEOUT: Final = 5.0
_NON_ENUM_METRIC_LABELS: Final[frozenset[str]] = frozenset(
@ -73,6 +79,36 @@ _NON_ENUM_METRIC_LABELS: Final[frozenset[str]] = frozenset(
)
class _PaginatedPrismaTable(Protocol[_TableRowT]):
"""The slice of a prisma table action surface used for budget-metric pagination."""
async def find_many(
self,
*,
skip: int,
take: int,
order: Mapping[str, str],
include: Mapping[str, bool] | None = None,
) -> list[_TableRowT]: ...
async def count(self) -> int: ...
def _paginated_table(repository: BaseRepository[_TableRowT]) -> _PaginatedPrismaTable[_TableRowT]:
"""View a repository's prisma table through the pagination surface budget metrics need."""
return repository.table
class _OrgBudgetRow(Protocol):
"""The budget columns joined onto an organization row."""
@property
def max_budget(self) -> float | None: ...
@property
def budget_reset_at(self) -> datetime | None: ...
class _ExcludedLabelMetric:
"""Proxies a prometheus metric whose declared ``labelnames`` had globally
excluded labels removed, dropping those labels from every ``labels(...)``
@ -1531,7 +1567,7 @@ class PrometheusLogger(CustomLogger):
cache_creation_detail_tokens: Final = PrometheusLogger._resolve_cache_write_tokens(prompt_details)
detail_metrics: Final[list[tuple[Any, DEFINED_PROMETHEUS_METRICS, Any]]] = [
detail_metrics: Final[list[tuple[Any, DEFINED_PROMETHEUS_METRICS, object]]] = [
(
self.litellm_input_cached_tokens_metric,
"litellm_input_cached_tokens_metric",
@ -1584,7 +1620,7 @@ class PrometheusLogger(CustomLogger):
if not isinstance(usage_object, dict):
return
media_metrics: Final[list[tuple[Any, DEFINED_PROMETHEUS_METRICS, Any]]] = [
media_metrics: Final[list[tuple[Any, DEFINED_PROMETHEUS_METRICS, object]]] = [
(
self.litellm_video_duration_seconds_metric,
"litellm_video_duration_seconds_metric",
@ -1606,7 +1642,7 @@ class PrometheusLogger(CustomLogger):
def _inc_sparse_usage_counters(
self,
counters_with_values: list[tuple[Any, DEFINED_PROMETHEUS_METRICS, Any]],
counters_with_values: Sequence[tuple[Any, DEFINED_PROMETHEUS_METRICS, object]],
enum_values: UserAPIKeyLabelValues,
label_context: PrometheusLabelFactoryContext | None = None,
) -> None:
@ -2133,7 +2169,7 @@ class PrometheusLogger(CustomLogger):
def _extract_status_code(
self,
kwargs: dict | None = None,
enum_values: Any | None = None,
enum_values: UserAPIKeyLabelValues | None = None,
exception: Exception | None = None,
) -> int | None:
"""
@ -2151,7 +2187,7 @@ class PrometheusLogger(CustomLogger):
Returns:
Status code as integer if found, None otherwise
"""
status_code = None
status_code: int | None = None
# Try from enum_values first (most common in our callbacks)
if enum_values and hasattr(enum_values, "status_code") and enum_values.status_code:
@ -2225,8 +2261,8 @@ class PrometheusLogger(CustomLogger):
def _should_skip_metrics_for_invalid_key(
self,
kwargs: dict | None = None,
user_api_key_dict: Any | None = None,
enum_values: Any | None = None,
user_api_key_dict: UserAPIKeyAuth | None = None,
enum_values: UserAPIKeyLabelValues | None = None,
standard_logging_payload: dict | StandardLoggingPayload | None = None,
exception: Exception | None = None,
) -> bool:
@ -2391,7 +2427,7 @@ class PrometheusLogger(CustomLogger):
for all successful requests (both streaming and non-streaming).
"""
def _safe_get(self, obj: Any, key: str, default: Any = None) -> Any:
def _safe_get(self, obj: Any, key: str, default: object = None) -> Any:
"""Get value from dict or Pydantic model."""
if obj is None:
return default
@ -3273,8 +3309,8 @@ class PrometheusLogger(CustomLogger):
async def _initialize_budget_metrics(
self,
data_fetch_function: Callable[..., Awaitable[tuple[list[Any], int | None]]],
set_metrics_function: Callable[[list[Any]], Awaitable[None]],
data_fetch_function: Callable[..., Awaitable[tuple[list[_BudgetRowT], int | None]]],
set_metrics_function: Callable[[list[_BudgetRowT]], Awaitable[None]],
data_type: Literal["teams", "keys", "users", "orgs"],
):
"""
@ -3393,12 +3429,12 @@ class PrometheusLogger(CustomLogger):
async def fetch_users(page_size: int, page: int) -> tuple[list[LiteLLM_UserTable], int | None]:
skip: Final = (page - 1) * page_size
users: Final = await UserRepository(prisma_client).table.find_many(
users: Final = await _paginated_table(UserRepository(prisma_client)).find_many(
skip=skip,
take=page_size,
order={"created_at": "desc"},
)
total_count: Final = await UserRepository(prisma_client).table.count()
total_count: Final = await _paginated_table(UserRepository(prisma_client)).count()
return users, total_count
await self._initialize_budget_metrics(
@ -3419,13 +3455,13 @@ class PrometheusLogger(CustomLogger):
async def fetch_orgs(page_size: int, page: int) -> tuple[list, int | None]:
skip: Final = (page - 1) * page_size
orgs: Final = await OrganizationRepository(prisma_client).table.find_many(
orgs: Final = await _paginated_table(OrganizationRepository(prisma_client)).find_many(
skip=skip,
take=page_size,
order={"created_at": "desc"},
include={"litellm_budget_table": True},
)
total_count: Final = await OrganizationRepository(prisma_client).table.count()
total_count: Final = await _paginated_table(OrganizationRepository(prisma_client)).count()
return orgs, total_count
await self._initialize_budget_metrics(
@ -3488,7 +3524,7 @@ class PrometheusLogger(CustomLogger):
try:
# Get total user count
total_users: Final = await UserRepository(prisma_client).table.count()
total_users: Final = await _paginated_table(UserRepository(prisma_client)).count()
self.litellm_total_users_metric.set(total_users)
verbose_logger.debug("Prometheus: set litellm_total_users to %s", total_users)
@ -3497,13 +3533,13 @@ class PrometheusLogger(CustomLogger):
verbose_logger.debug("Prometheus: set litellm_active_users to %s", billable_users)
# Get total team count
total_teams: Final = await TeamRepository(prisma_client).table.count()
total_teams: Final = await _paginated_table(TeamRepository(prisma_client)).count()
self.litellm_teams_count_metric.set(total_teams)
verbose_logger.debug("Prometheus: set litellm_teams_count to %s", total_teams)
except Exception as e:
verbose_logger.exception("Error initializing user/team count metrics: %s", e)
async def _set_key_list_budget_metrics(self, keys: list[str | UserAPIKeyAuth]):
async def _set_key_list_budget_metrics(self, keys: list[str | UserAPIKeyAuth | LiteLLM_DeletedVerificationToken]):
"""Helper function to set budget metrics for a list of keys"""
for key in keys:
if isinstance(key, UserAPIKeyAuth):
@ -3522,7 +3558,7 @@ class PrometheusLogger(CustomLogger):
async def _set_org_list_budget_metrics(self, orgs: list):
"""Helper function to set budget metrics for a list of orgs"""
for org in orgs:
budget_table = getattr(org, "litellm_budget_table", None)
budget_table: _OrgBudgetRow | None = getattr(org, "litellm_budget_table", None)
self._set_org_budget_metrics(
org_id=org.organization_id or "",
org_alias=org.organization_alias or "",
@ -4051,6 +4087,11 @@ class PrometheusLogger(CustomLogger):
verbose_proxy_logger.debug("Starting Prometheus Metrics on /metrics (no authentication)")
def _label_source(enum_values: UserAPIKeyLabelValues) -> Mapping[str, object]:
"""Flatten the label values into the opaque name/value mapping the label filters read."""
return enum_values.model_dump()
def _prometheus_labels_from_context(
supported_enum_labels: list[str],
ctx: PrometheusLabelFactoryContext,
@ -4098,7 +4139,7 @@ def prometheus_label_factory(
return _prometheus_labels_from_context(supported_enum_labels, label_context)
# Extract dictionary from Pydantic object
enum_dict: Final = enum_values.model_dump()
enum_dict: Final = _label_source(enum_values)
# Filter supported labels and sanitize values to prevent breaking
# the Prometheus text format (e.g. U+2028 Line Separator in label values)
@ -4154,7 +4195,7 @@ def get_custom_labels_from_metadata(metadata: dict) -> dict[str, str]:
keys_parts = key.split(".")
# Traverse through the dictionary using the parts
value: Any = metadata
value: object = metadata
for part in keys_parts:
if isinstance(value, dict):
value = value.get(part, None) # Get the value, return None if not found
@ -4171,7 +4212,7 @@ def get_custom_labels_from_metadata(metadata: dict) -> dict[str, str]:
def _get_combined_custom_metadata_from_standard_logging_payload(
standard_logging_payload: dict | None,
) -> dict[str, Any]:
) -> dict[str, object]:
"""
Combine the metadata sources that can supply custom Prometheus labels.

32
litellm/integrations/websearch_interception/handler.py
View file

@ -12,6 +12,8 @@ import uuid
from collections.abc import AsyncIterator, Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, TypedDict, cast
from typing_extensions import ReadOnly
import litellm
from litellm._logging import verbose_logger
from litellm.anthropic_interface import messages as anthropic_messages
@ -90,6 +92,20 @@ class _SearchToolConfig(TypedDict, total=False):
litellm_params: Mapping[str, object] | None
class _DeploymentKwargsView(TypedDict):
"""Typed reads of the untyped request kwargs seen by the deployment hook."""
custom_llm_provider: ReadOnly[str]
litellm_params: ReadOnly[Mapping[str, object]]
model: ReadOnly[str]
class _UserAuthView(TypedDict):
"""Typed read of the optional team attached to the caller's auth object."""
team_id: ReadOnly[str | None]
class WebSearchInterceptionLogger(CustomLogger):
"""
CustomLogger that intercepts WebSearch tool calls for models that don't
@ -265,7 +281,9 @@ class WebSearchInterceptionLogger(CustomLogger):
)
return response
async def async_pre_call_deployment_hook(self, kwargs: dict[str, Any], call_type: CallTypes | None) -> dict | None:
async def async_pre_call_deployment_hook(
self, kwargs: dict[str, Any], call_type: CallTypes | None
) -> dict[str, object] | None:
"""
Pre-call hook to convert native Anthropic web_search tools to regular tools.
@ -275,12 +293,17 @@ class WebSearchInterceptionLogger(CustomLogger):
"""
# Check if this is for an enabled provider
# Try top-level kwargs first, then nested litellm_params, then derive from model name
custom_llm_provider = kwargs.get("custom_llm_provider", "") or kwargs.get("litellm_params", {}).get(
kwargs_view: Final[_DeploymentKwargsView] = {
"custom_llm_provider": kwargs.get("custom_llm_provider", ""),
"litellm_params": kwargs.get("litellm_params", {}),
"model": kwargs.get("model", ""),
}
custom_llm_provider = kwargs_view["custom_llm_provider"] or kwargs_view["litellm_params"].get(
"custom_llm_provider", ""
)
if not custom_llm_provider:
try:
_, custom_llm_provider, _, _ = litellm.get_llm_provider(model=kwargs.get("model", ""))
_, custom_llm_provider, _, _ = litellm.get_llm_provider(model=kwargs_view["model"])
except Exception:
custom_llm_provider = ""
if custom_llm_provider not in self.enabled_providers:
@ -1422,7 +1445,8 @@ class WebSearchInterceptionLogger(CustomLogger):
valid_token=user_api_key_auth,
)
team_id: Final = getattr(user_api_key_auth, "team_id", None)
auth_view: Final[_UserAuthView] = {"team_id": getattr(user_api_key_auth, "team_id", None)}
team_id: Final = auth_view["team_id"]
if team_id:
from litellm.proxy.proxy_server import (
prisma_client,

79
litellm/litellm_core_utils/realtime_streaming.py
View file

@ -1,7 +1,9 @@
import asyncio
import json
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, Protocol, cast
from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict, cast
from typing_extensions import ReadOnly
import litellm
from litellm._logging import verbose_logger
@ -32,13 +34,52 @@ class _ClientWebSocketExceptions(Protocol):
ConnectionClosed: type[Exception]
class _ClientWebSocket(Protocol):
class _ASGIScope(TypedDict, total=False):
"""The part of an ASGI connection scope this module reads."""
headers: ReadOnly[Sequence[tuple[bytes | str, bytes | str]]]
class _ClientEventItem(TypedDict, total=False):
"""The ``item`` payload of a client ``conversation.item.create`` frame."""
type: ReadOnly[str]
role: ReadOnly[str]
output: ReadOnly[object]
content: ReadOnly[Sequence[object]]
class _ClientEventFrame(TypedDict, total=False):
"""The fields the proxy reads from a client realtime frame."""
type: ReadOnly[str]
item: ReadOnly[_ClientEventItem]
session: ReadOnly[Mapping[str, object]]
class _ResponseDoneBody(TypedDict, total=False):
"""The ``response`` body of a ``response.done`` event, as read for spend logging."""
output: ReadOnly[Sequence[Mapping[str, object]]]
class _ScopedWebSocket(Protocol):
@property
def scope(self) -> _ASGIScope: ...
class _ClientWebSocket(_ScopedWebSocket, Protocol):
exceptions: _ClientWebSocketExceptions
async def send_text(self, data: str) -> None: ...
async def receive_text(self) -> str: ...
def _decode_json_object(payload: str) -> Mapping[str, object]:
"""Decode a realtime frame into its top-level field mapping."""
return json.loads(payload)
class RealtimeEventNormalizer(Protocol):
def should_drop(self, event: object) -> bool: ...
def normalize(self, event: dict) -> dict: ...
@ -294,7 +335,7 @@ class RealTimeStreaming:
try:
if event_obj.get("type") != "response.done":
return
response: Final = cast(dict[str, Any], event_obj.get("response", {}))
response: Final = cast(_ResponseDoneBody, event_obj.get("response", {}))
item: Mapping[str, object]
for item in response.get("output", []):
if item.get("type") == "function_call":
@ -353,7 +394,7 @@ class RealTimeStreaming:
sent = False
for msg in transformed:
try:
msg_obj = json.loads(msg)
msg_obj = _decode_json_object(msg)
except (json.JSONDecodeError, TypeError):
msg_obj = None
if isinstance(msg_obj, dict) and self.provider_config.is_setup_message(msg_obj):
@ -399,7 +440,7 @@ class RealTimeStreaming:
return message
try:
message_obj: Final[Mapping[str, object]] = json.loads(message)
message_obj: Final = _decode_json_object(message)
except (json.JSONDecodeError, TypeError):
return message
@ -468,7 +509,7 @@ class RealTimeStreaming:
for message in messages:
try:
msg_type = json.loads(message).get("type")
msg_type = _decode_json_object(message).get("type")
except (json.JSONDecodeError, TypeError):
collapsed.extend(pending_appends)
pending_appends = []
@ -502,14 +543,14 @@ class RealTimeStreaming:
if self._backend_setup_complete and not self._flushing_pending_messages_until_setup:
return False
try:
msg_obj: Final[Mapping[str, object]] = json.loads(message)
msg_obj: Final = _decode_json_object(message)
except (json.JSONDecodeError, TypeError):
return False
return msg_obj.get("type") in RealTimeStreaming._CLIENT_AUDIO_BUFFER_TYPES
def _buffer_pending_message_until_setup(self, message: str) -> None:
try:
msg_type = json.loads(message).get("type")
msg_type = _decode_json_object(message).get("type")
except (json.JSONDecodeError, TypeError):
msg_type = None
@ -602,7 +643,7 @@ class RealTimeStreaming:
``return_new_content_delta_events`` modality lookup, ...).
"""
try:
message_obj: Final = json.loads(transformed_message)
message_obj: Final = _decode_json_object(transformed_message)
if "setup" in message_obj:
self.session_configuration_request = transformed_message
except (json.JSONDecodeError, TypeError):
@ -930,7 +971,7 @@ class RealTimeStreaming:
def _parse_backend_event(raw_response: str) -> dict[str, object] | None:
"""Parse a backend frame once. Returns None for non-JSON or non-object frames."""
try:
event: Final = json.loads(raw_response)
event: Final = _decode_json_object(raw_response)
except (json.JSONDecodeError, TypeError):
return None
return event if isinstance(event, dict) else None
@ -1030,14 +1071,14 @@ class RealTimeStreaming:
await self.log_messages()
@staticmethod
def _detect_beta_header(websocket: Any) -> bool:
def _detect_beta_header(websocket: _ScopedWebSocket) -> bool:
"""Return True if the client sent 'OpenAI-Beta: realtime=v1'.
Checks the raw ASGI scope headers so it works for both FastAPI WebSocket
objects and any test doubles that expose a .scope dict.
"""
try:
headers: Final[Sequence[tuple[bytes | str, bytes | str]]] = websocket.scope.get("headers", [])
headers: Final = websocket.scope.get("headers", [])
for name, value in headers:
if isinstance(name, bytes):
name = name.decode("latin-1")
@ -1183,6 +1224,7 @@ class RealTimeStreaming:
return item
async def client_ack_messages(self):
client_event: _ClientEventFrame
try:
while True:
message = await self.websocket.receive_text()
@ -1194,11 +1236,12 @@ class RealTimeStreaming:
from litellm.types.guardrails import GuardrailEventHooks
msg_obj = json.loads(message)
msg_type = msg_obj.get("type")
client_event = msg_obj
msg_type = client_event.get("type")
if msg_type == "conversation.item.create":
# Check user text messages for prompt injection
item = msg_obj.get("item", {})
item = client_event.get("item", {})
# Check function_call_output first so a client cannot
# bypass the tool-result guardrail by also setting
# role="user" on a function_call_output item.
@ -1297,7 +1340,7 @@ class RealTimeStreaming:
and not self._guardrail_turn_detection_update_sent
and self._has_audio_transcription_guardrails()
):
session: object = msg_obj.setdefault("session", {})
session: Mapping[str, object] | None = msg_obj.setdefault("session", {})
if isinstance(session, dict):
existing_td = session.get("turn_detection")
if not isinstance(existing_td, dict):
@ -1324,7 +1367,7 @@ class RealTimeStreaming:
and not guardrail_turn_detection_injected
and self._has_audio_transcription_guardrails()
):
session = msg_obj.get("session")
session = client_event.get("session")
if isinstance(session, dict):
td_overridden = False
flat_td = session.get("turn_detection")
@ -1367,14 +1410,14 @@ class RealTimeStreaming:
# the upstream is in GA mode. Beta upstreams expect the flat
# session shape unchanged.
if msg_type == "session.update" and not self._backend_uses_beta_protocol:
session = msg_obj.get("session", {})
session = client_event.get("session", {})
if isinstance(session, dict):
session = self._remap_beta_session_to_ga(session)
msg_obj["session"] = session
message = json.dumps(msg_obj)
if msg_type == "session.update" and self._event_normalizer:
session = msg_obj.get("session")
session = client_event.get("session")
if isinstance(session, dict):
msg_obj["session"] = self._event_normalizer.patch_outgoing_session(session)
message = json.dumps(msg_obj)

9
litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py
View file

@ -4,7 +4,7 @@ Handler for the Anthropic v1/messages -> OpenAI Responses API path.
Used when the target model is an OpenAI or Azure model.
"""
from collections.abc import AsyncIterator, Coroutine
from collections.abc import AsyncIterator, Coroutine, Mapping
from typing import Any, Final
import litellm
@ -25,6 +25,11 @@ from .transformation import LiteLLMAnthropicToResponsesAPIAdapter
_ADAPTER: Final = LiteLLMAnthropicToResponsesAPIAdapter()
def _forwarded_kwargs(extra_kwargs: Mapping[str, object] | None) -> Mapping[str, object]:
"""The litellm-specific kwargs forwarded verbatim onto the Responses API request."""
return extra_kwargs or {}
def _build_responses_kwargs(
*,
max_tokens: int,
@ -100,7 +105,7 @@ def _build_responses_kwargs(
# Forward litellm-specific kwargs (api_key, api_base, logging obj, etc.)
excluded: Final = {"anthropic_messages"}
for key, value in (extra_kwargs or {}).items():
for key, value in _forwarded_kwargs(extra_kwargs).items():
if key == "litellm_logging_obj" and value is not None:
from litellm.litellm_core_utils.litellm_logging import (
Logging as LiteLLMLoggingObject,

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

@ -1,4 +1,7 @@
from typing import TYPE_CHECKING, Any, Final, Optional
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, Optional, Protocol, TypeAlias
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
@ -40,7 +43,7 @@ def _generic_passthrough_handler() -> BaseTranslation:
_StringHolder = tuple[Any, str | int]
def _collect_strings(node: Any, holders: list[_StringHolder]) -> None:
def _collect_strings(node: object, holders: list[_StringHolder]) -> None:
"""
Record a (container, key) holder for every non-empty string value nested
under an arbitrary JSON node, so prompt content a caller hides in fields
@ -48,7 +51,7 @@ def _collect_strings(node: Any, holders: list[_StringHolder]) -> None:
and can be written back in place. Iterative to avoid unbounded recursion
on deeply nested payloads.
"""
stack: Final[list[Any]] = [node]
stack: Final[list[object]] = [node]
while stack:
current = stack.pop()
if isinstance(current, dict):
@ -129,7 +132,7 @@ def _extract_converse_texts(
def _extract_converse_output_texts(
content_blocks: list[Any],
content_blocks: Sequence[object],
) -> tuple[list[str], list[_StringHolder]]:
"""
Collect user-visible text from Bedrock Converse output content blocks.
@ -178,10 +181,34 @@ def _write_back_texts(
container[key] = guardrailed_texts[idx]
_DeltaHolder = tuple[Any, Any, str | int]
_GroupKey: TypeAlias = str | tuple[str, int]
def _collect_stream_delta_text_holders(delta: Any) -> list[_DeltaHolder]:
class _TextContainer(Protocol):
"""JSON object whose ``key`` entry holds a guardrailable text string."""
def __getitem__(self, key: str, /) -> str: ...
def __setitem__(self, key: str, value: str, /) -> None: ...
_DeltaHolder = tuple[_GroupKey, _TextContainer, str]
class _StreamFrame(TypedDict):
"""One raw event-stream frame plus the guardrailable texts it carries."""
raw: ReadOnly[bytes]
texts: ReadOnly[Sequence[tuple[_GroupKey, str]]]
def _unpack_uint32(buffer: bytes) -> int:
import struct
return struct.unpack("!I", buffer)[0]
def _collect_stream_delta_text_holders(delta: object) -> list[_DeltaHolder]:
"""
Collect the user-visible text strings a Bedrock Converse ``contentBlockDelta``
can carry, matching the coverage of the non-streaming output handler.
@ -238,11 +265,11 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation):
from botocore.eventstream import EventStreamBuffer
frames: Final[list[dict]] = []
frames: Final[list[_StreamFrame]] = []
offset = 0
while offset + 16 <= len(body_bytes):
total_length = struct.unpack("!I", body_bytes[offset : offset + 4])[0]
total_length = _unpack_uint32(body_bytes[offset : offset + 4])
if total_length < 16 or offset + total_length > len(body_bytes):
break
frame_raw = body_bytes[offset : offset + total_length]
@ -263,10 +290,10 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation):
frames.append({"raw": frame_raw, "texts": []})
continue
texts: list[tuple[Any, str]] = []
texts: list[tuple[_GroupKey, str]] = []
if event_type == "contentBlockDelta":
try:
payload_dict = _json.loads(payload_bytes)
payload_dict: dict[str, object] = _json.loads(payload_bytes)
texts = [
(group_key, container[key])
for group_key, container, key in _collect_stream_delta_text_holders(payload_dict.get("delta"))
@ -282,9 +309,9 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation):
trailing_bytes: Final = body_bytes[offset:]
group_order: Final[list[Any]] = []
group_members: Final[dict[Any, list[tuple[int, int]]]] = {}
group_texts: Final[dict[Any, list[str]]] = {}
group_order: Final[list[_GroupKey]] = []
group_members: Final[dict[_GroupKey, list[tuple[int, int]]]] = {}
group_texts: Final[dict[_GroupKey, list[str]]] = {}
for frame_idx, frame in enumerate(frames):
for local_idx, (group_key, text) in enumerate(frame["texts"]):
if group_key not in group_members:
@ -351,8 +378,8 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation):
continue
frame_raw = frame["raw"]
orig_total = struct.unpack("!I", frame_raw[0:4])[0]
orig_hdrs_len = struct.unpack("!I", frame_raw[4:8])[0]
orig_total = _unpack_uint32(frame_raw[0:4])
orig_hdrs_len = _unpack_uint32(frame_raw[4:8])
headers_bytes = frame_raw[12 : 12 + orig_hdrs_len]
try:
@ -386,7 +413,7 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation):
data: dict,
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None,
) -> Any:
) -> Mapping[str, object]:
endpoint: Final = data.get("endpoint", "")
body: Final = data.get("data")
@ -428,12 +455,12 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation):
async def process_output_response(
self,
response: Any,
response: object,
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None,
user_api_key_dict: Any | None = None,
user_api_key_dict: Optional["UserAPIKeyAuth"] = None,
request_data: dict | None = None,
) -> Any:
) -> object:
endpoint: Final = (request_data or {}).get("endpoint", "")
if endpoint and not _is_converse_endpoint(endpoint):
return await _generic_passthrough_handler().process_output_response(

69
litellm/llms/gemini/vector_stores/transformation.py
View file

@ -5,9 +5,11 @@ Implements the transformation between LiteLLM's unified vector store API
and Google Gemini's File Search API.
"""
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final
import httpx
from typing_extensions import ReadOnly, TypedDict
from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig
from litellm.llms.gemini.common_utils import (
@ -35,6 +37,61 @@ else:
LiteLLMLoggingObj = Any
class GeminiRetrievedContext(TypedDict, total=False):
"""Passage Gemini retrieved from a File Search store."""
text: ReadOnly[str]
uri: ReadOnly[str]
title: ReadOnly[str]
class GeminiGroundingChunk(TypedDict, total=False):
"""One source Gemini grounded its answer on."""
retrievedContext: ReadOnly[GeminiRetrievedContext]
class GeminiGroundingSegment(TypedDict, total=False):
"""Span of the generated answer a grounding support refers to."""
text: ReadOnly[str]
class GeminiGroundingSupport(TypedDict, total=False):
"""Citation linking an answer span to the grounding chunks that back it."""
segment: ReadOnly[GeminiGroundingSegment]
groundingChunkIndices: ReadOnly[Sequence[int]]
confidenceScores: ReadOnly[Sequence[float]]
class GeminiFileSearchGroundingMetadata(TypedDict, total=False):
"""Grounding metadata Gemini returns for a File Search candidate."""
groundingChunks: ReadOnly[Sequence[GeminiGroundingChunk]]
groundingSupports: ReadOnly[Sequence[GeminiGroundingSupport]]
class GeminiFileSearchCandidate(TypedDict, total=False):
"""One candidate of a Gemini File Search ``generateContent`` response."""
groundingMetadata: ReadOnly[GeminiFileSearchGroundingMetadata]
class GeminiFileSearchResponse(TypedDict, total=False):
"""Body of a ``generateContent`` call made with the File Search tool."""
candidates: ReadOnly[Sequence[GeminiFileSearchCandidate]]
class GeminiFileSearchStore(TypedDict, total=False):
"""Body of a Gemini ``fileSearchStores`` create response."""
name: ReadOnly[str]
displayName: ReadOnly[str]
createTime: ReadOnly[str]
class GeminiVectorStoreConfig(BaseVectorStoreConfig):
"""
Vector store configuration for Google Gemini File Search.
@ -110,7 +167,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig):
api_base: str,
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
extra_body: dict[str, Any] | None = None,
extra_body: Mapping[str, object] | None = None,
) -> tuple[str, dict]:
"""
Transform search request to Gemini's generateContent format.
@ -133,7 +190,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig):
url: Final = f"{api_base}/models/{model}:generateContent"
# Build file_search tool configuration (using snake_case as per Gemini docs)
file_search_config: Final[dict[str, Any]] = {"file_search_store_names": [vector_store_id]}
file_search_config: Final[dict[str, object]] = {"file_search_store_names": [vector_store_id]}
# Add metadata filter if provided
metadata_filter: Final = vector_store_search_optional_params.get("filters")
@ -178,7 +235,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig):
Extracts grounding metadata and citations from the response.
"""
try:
response_data: Final = response.json()
response_data: Final[GeminiFileSearchResponse] = response.json()
results: Final[list[VectorStoreSearchResult]] = []
# Extract candidates and grounding metadata
@ -246,7 +303,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig):
)
)
query: Final = litellm_logging_obj.model_call_details.get("query", "")
query: Final[str] = litellm_logging_obj.model_call_details.get("query", "")
return VectorStoreSearchResponse(
object="vector_store.search_results.page",
@ -273,7 +330,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig):
# API key is passed via x-goog-api-key header (set in validate_environment)
request_body: Final[dict[str, Any]] = {}
request_body: Final[dict[str, object]] = {}
# Add display name if provided
name: Final = vector_store_create_optional_params.get("name")
@ -287,7 +344,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig):
Transform Gemini's fileSearchStore response to standard format.
"""
try:
response_data: Final = response.json()
response_data: Final[GeminiFileSearchStore] = response.json()
# Extract store name (format: fileSearchStores/xxxxxxx)
store_name: Final = response_data.get("name", "")

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

@ -26,7 +26,9 @@ without the optional STT extras installed.
import asyncio
import inspect
from typing import TYPE_CHECKING, Any, Final
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 (
get_audio_file_name,
@ -62,6 +64,45 @@ _DEFAULT_CHUNK_BYTES: Final = _DEFAULT_CHUNK_SAMPLES * 2 # int16 = 2 bytes/samp
_RIVA_INSTALL_HINT = "NVIDIA Riva client is not installed. Install with `pip install 'litellm[stt-nvidia-riva]'`."
class _RivaAuth(Protocol):
"""Opaque ``riva.client.Auth`` handle."""
class _AsrService(Protocol):
@property
def streaming_response_generator(self) -> Callable[..., Iterable[object]]: ...
class _EndpointingConfig(Protocol):
"""Opaque ``EndpointingConfig`` protobuf message."""
class _EndpointingConfigField(Protocol):
CopyFrom: Callable[[_EndpointingConfig], None]
class _RecognitionConfig(Protocol):
@property
def endpointing_config(self) -> _EndpointingConfigField: ...
class _StreamingRecognitionConfig(Protocol):
"""Opaque ``StreamingRecognitionConfig`` protobuf message."""
class _AudioEncoding(Protocol):
@property
def LINEAR_PCM(self) -> object: ...
def _auth_factory(riva_module: ModuleType) -> Callable[..., _RivaAuth]:
return riva_module.Auth
def _audio_encoding(riva_asr_module: ModuleType) -> _AudioEncoding:
return riva_asr_module.AudioEncoding
class NvidiaRivaAudioTranscription:
"""Sync + async entry point for Riva ASR."""
@ -206,7 +247,9 @@ class NvidiaRivaAudioTranscription:
riva_asr_module=riva_asr_module,
recognition_config_dict=recognition_config_dict,
)
streaming_config = riva_asr_module.StreamingRecognitionConfig(config=recognition_config, interim_results=False)
streaming_config: Final[_StreamingRecognitionConfig] = riva_asr_module.StreamingRecognitionConfig(
config=recognition_config, interim_results=False
)
logging_obj.pre_call(
input=None,
@ -223,9 +266,9 @@ class NvidiaRivaAudioTranscription:
)
try:
asr_service: Final = riva_module.ASRService(auth_obj)
asr_service: Final[_AsrService] = riva_module.ASRService(auth_obj)
audio_chunks: Final = self._iter_audio_chunks(resampled.pcm_bytes)
stream_kwargs: Final[dict[str, Any]] = {
stream_kwargs: Final[dict[str, object]] = {
"audio_chunks": audio_chunks,
"streaming_config": streaming_config,
}
@ -274,11 +317,11 @@ class NvidiaRivaAudioTranscription:
def _construct_auth(
self,
riva_module: Any,
riva_module: ModuleType,
api_base: str,
api_key: str | None,
optional_params: dict,
) -> Any:
) -> _RivaAuth:
"""
Build a ``riva.client.Auth`` object.
@ -300,20 +343,22 @@ class NvidiaRivaAudioTranscription:
metadata.append(("authorization", f"Bearer {api_key}"))
try:
return riva_module.Auth(uri=api_base, use_ssl=use_ssl, metadata_args=metadata)
return _auth_factory(riva_module)(uri=api_base, use_ssl=use_ssl, metadata_args=metadata)
except TypeError:
# Older riva-client signatures used positional-only args.
return riva_module.Auth(None, use_ssl, api_base, metadata)
return _auth_factory(riva_module)(None, use_ssl, api_base, metadata)
def _build_recognition_config_proto(self, riva_asr_module: Any, recognition_config_dict: dict[str, Any]):
def _build_recognition_config_proto(
self, riva_asr_module: ModuleType, recognition_config_dict: dict[str, Any]
) -> _RecognitionConfig:
encoding_name: Final = (recognition_config_dict.get("encoding") or "LINEAR_PCM").upper()
encoding_enum: Final = getattr(
riva_asr_module.AudioEncoding,
encoding_enum: Final[object] = getattr(
_audio_encoding(riva_asr_module),
encoding_name,
riva_asr_module.AudioEncoding.LINEAR_PCM,
_audio_encoding(riva_asr_module).LINEAR_PCM,
)
config: Final = riva_asr_module.RecognitionConfig(
config: Final[_RecognitionConfig] = riva_asr_module.RecognitionConfig(
encoding=encoding_enum,
sample_rate_hertz=int(recognition_config_dict["sample_rate_hertz"]),
language_code=recognition_config_dict["language_code"],
@ -329,7 +374,7 @@ class NvidiaRivaAudioTranscription:
endpointing: Final = recognition_config_dict.get("endpointing_config")
if isinstance(endpointing, dict) and endpointing:
try:
ep: Final = riva_asr_module.EndpointingConfig(**endpointing)
ep: Final[_EndpointingConfig] = riva_asr_module.EndpointingConfig(**endpointing)
config.endpointing_config.CopyFrom(ep)
except Exception:
# If the user supplied an unknown EndpointingConfig field
@ -340,7 +385,7 @@ class NvidiaRivaAudioTranscription:
return config
@staticmethod
def _supports_timeout_kwarg(callable_obj: Any) -> bool:
def _supports_timeout_kwarg(callable_obj: Callable[..., object]) -> bool:
try:
sig: Final = inspect.signature(callable_obj)
except (TypeError, ValueError):
@ -359,14 +404,14 @@ class NvidiaRivaAudioTranscription:
yield chunk
@staticmethod
def _collect_final_results(stream) -> list[dict[str, Any]]:
def _collect_final_results(stream) -> list[dict[str, object]]:
"""
Walk the gRPC stream, ignore empty / non-final chunks, and return a
list of normalized final-result dicts. Matching the user's note: the
``id`` blocks with no ``results`` are streaming heartbeats and must
be skipped.
"""
final_results: Final[list[dict[str, Any]]] = []
final_results: Final[list[dict[str, object]]] = []
for response in stream:
results = getattr(response, "results", None) or []
for result in results:
@ -391,7 +436,7 @@ class NvidiaRivaAudioTranscription:
return final_results
def _import_riva():
def _import_riva() -> tuple[ModuleType, ModuleType]:
"""
Lazy import of ``riva.client`` and ``riva.client.proto.riva_asr_pb2``.

14
litellm/llms/oci/chat/cohere.py
View file

@ -84,9 +84,9 @@ def adapt_messages_to_cohere_standard(
tool_calls_raw: Any = msg.get("tool_calls") or []
for tc in tool_calls_raw:
tc_id = tc.get("id", "")
raw_args: Any = tc.get("function", {}).get("arguments", "{}")
raw_args = tc.get("function", {}).get("arguments", "{}")
try:
params: dict[str, Any] = json.loads(raw_args) if isinstance(raw_args, str) else raw_args
params: dict[str, object] = json.loads(raw_args) if isinstance(raw_args, str) else raw_args
except json.JSONDecodeError:
params = {}
tool_call_lookup[tc_id] = CohereToolCall(
@ -111,10 +111,10 @@ def adapt_messages_to_cohere_standard(
if role == "assistant" and msg.get("tool_calls"):
tool_calls = []
for tc in msg["tool_calls"]: # pyright: ignore[reportOptionalIterable] # truthiness check above rules out None
raw_arguments: Any = tc.get("function", {}).get("arguments", {})
raw_arguments = tc.get("function", {}).get("arguments", {})
if isinstance(raw_arguments, str):
try:
arguments: dict[str, Any] = json.loads(raw_arguments)
arguments: dict[str, object] = json.loads(raw_arguments)
except json.JSONDecodeError:
arguments = {}
else:
@ -211,7 +211,7 @@ def handle_cohere_response(
response_text: Final = cohere_response.chatResponse.text
finish_reason: Final = _normalize_oci_finish_reason(cohere_response.chatResponse.finishReason)
tool_calls: list[dict[str, Any]] | None = None
tool_calls: list[dict[str, object]] | None = None
if cohere_response.chatResponse.toolCalls:
tool_calls = [
{
@ -232,7 +232,7 @@ def handle_cohere_response(
# ``"tool_calls" in message`` (rather than truthiness) incorrectly conclude
# that tool calls were attempted. Matches the generic handler's behaviour,
# which only sets ``message.tool_calls`` when tool calls are present.
message: Final[dict[str, Any]] = {"role": "assistant", "content": content}
message: Final[dict[str, object]] = {"role": "assistant", "content": content}
if tool_calls is not None:
message["tool_calls"] = tool_calls
@ -317,7 +317,7 @@ def handle_cohere_stream_chunk(
# passing them through is the only chance to surface them.
cohere_tool_calls = None if (is_terminal_consolidation and prior_tool_calls_emitted) else typed_chunk.toolCalls
tool_calls: list[dict[str, Any]] | None = None
tool_calls: list[dict[str, object]] | None = None
if cohere_tool_calls:
tool_calls = [
{

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

@ -28,10 +28,13 @@ Output: response.output is List[GenericResponseOutputItem] where each has:
- text: str
"""
from collections.abc import Sequence
from typing import TYPE_CHECKING, Any, Final, Union, cast
from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall
from openai.types.responses.tool_param import FunctionToolParam
from pydantic import BaseModel
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.completion_extras.litellm_responses_transformation.transformation import (
@ -45,6 +48,7 @@ from litellm.types.llms.openai import (
AllMessageValues,
ChatCompletionToolCallChunk,
ChatCompletionToolParam,
OpenAIMcpServerTool,
ResponsesAPIStreamEvents,
)
from litellm.types.responses.main import (
@ -56,10 +60,26 @@ 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.proxy._types import UserAPIKeyAuth
from litellm.types.llms.openai import ResponseInputParam
from litellm.types.utils import ResponsesAPIResponse
class ResponseOutputEnvelope(TypedDict, total=False):
"""Dict form of a Responses API response, as far as guardrail write-back reads it."""
output: ReadOnly[Sequence[object]]
model: ReadOnly[str | None]
class ResponsesStreamChunk(TypedDict, total=False):
"""Responses API streaming event, as far as the accumulated-stream helpers read it."""
type: ReadOnly[str]
text: ReadOnly[str]
class OpenAIResponsesHandler(BaseTranslation):
"""
Handler for processing OpenAI Responses API with guardrails.
@ -91,8 +111,8 @@ class OpenAIResponsesHandler(BaseTranslation):
self,
data: dict,
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: Any | None = None,
) -> Any:
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
) -> dict[str, object]:
"""
Process input by applying guardrails to text content.
@ -108,7 +128,7 @@ class OpenAIResponsesHandler(BaseTranslation):
# Handle simple string input
if isinstance(input_data, str):
inputs = GenericGuardrailAPIInputs(texts=[input_data])
original_tools: list[dict[str, Any]] = []
original_tools: list[dict[str, object]] = []
# Extract and transform tools if present
if "tools" in data and data["tools"]:
@ -142,7 +162,7 @@ class OpenAIResponsesHandler(BaseTranslation):
texts_to_check: Final[list[str]] = []
images_to_check: Final[list[str]] = []
task_mappings: Final[list[tuple[int, int | None]]] = []
original_tools_list: Final[list[dict[str, Any]]] = list(data.get("tools") or [])
original_tools_list: Final[list[dict[str, object]]] = list(data.get("tools") or [])
# Step 1: Extract all text content, images, and tools
for msg_idx, message in enumerate(input_data):
@ -211,7 +231,7 @@ class OpenAIResponsesHandler(BaseTranslation):
def _extract_and_transform_tools(
self,
tools: list[dict[str, Any]],
tools: list[FunctionToolParam | OpenAIMcpServerTool],
tools_to_check: list[ChatCompletionToolParam],
) -> None:
"""
@ -228,7 +248,7 @@ class OpenAIResponsesHandler(BaseTranslation):
) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools)
tools_to_check.extend(cast(list[ChatCompletionToolParam], transformed_tools))
def _remap_tools_to_responses_api_format(self, guardrailed_tools: list[Any]) -> list[dict[str, Any]]:
def _remap_tools_to_responses_api_format(self, guardrailed_tools: list[Any]) -> list[dict[str, object]]:
"""
Remap guardrail-returned tools (Chat Completion format) back to
Responses API request tool format.
@ -239,9 +259,9 @@ class OpenAIResponsesHandler(BaseTranslation):
def _merge_tools_after_guardrail(
self,
original_tools: list[dict[str, Any]],
remapped: list[dict[str, Any]],
) -> list[dict[str, Any]]:
original_tools: list[dict[str, object]],
remapped: list[dict[str, object]],
) -> list[dict[str, object]]:
"""
Merge remapped guardrailed tools with original tools that were not sent
to the guardrail (e.g. web_search, web_search_preview), preserving order.
@ -250,7 +270,7 @@ class OpenAIResponsesHandler(BaseTranslation):
"""
if not original_tools:
return remapped
result: Final[list[dict[str, Any]]] = []
result: Final[list[dict[str, object]]] = []
j = 0
for tool in original_tools:
if isinstance(tool, dict) and tool.get("type") in (
@ -269,8 +289,8 @@ class OpenAIResponsesHandler(BaseTranslation):
def _apply_guardrailed_tools_to_data(
self,
data: dict,
original_tools: list[dict[str, Any]],
guardrailed_tools: list[Any] | None,
original_tools: list[dict[str, object]],
guardrailed_tools: list[ChatCompletionToolParam] | None,
) -> None:
"""Remap guardrailed tools to Responses API format and merge with original, then set data['tools']."""
if guardrailed_tools is not None:
@ -279,7 +299,7 @@ class OpenAIResponsesHandler(BaseTranslation):
def _extract_input_text_and_images(
self,
message: Any, # Can be Dict[str, Any] or ResponseInputParam
message: Any,
msg_idx: int,
texts_to_check: list[str],
images_to_check: list[str],
@ -348,12 +368,12 @@ class OpenAIResponsesHandler(BaseTranslation):
async def process_output_response(
self,
response: "ResponsesAPIResponse",
response: Union["ResponsesAPIResponse", ResponseOutputEnvelope],
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: Any | None = None,
user_api_key_dict: Any | None = None,
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
user_api_key_dict: "UserAPIKeyAuth | None" = None,
request_data: dict | None = None,
) -> Any:
) -> Union["ResponsesAPIResponse", ResponseOutputEnvelope]:
"""
Process output response by applying guardrails to text content and tool calls.
@ -381,6 +401,7 @@ class OpenAIResponsesHandler(BaseTranslation):
# Track (output_item_index, content_index) for each text
# Handle both dict and Pydantic object responses
response_output: Sequence[object]
if isinstance(response, dict):
response_output = response.get("output", [])
elif hasattr(response, "output"):
@ -426,7 +447,7 @@ class OpenAIResponsesHandler(BaseTranslation):
if tool_calls_to_check:
inputs["tool_calls"] = tool_calls_to_check
# Include model information from the response if available
response_model = None
response_model: str | None = None
if isinstance(response, dict):
response_model = response.get("model")
elif hasattr(response, "model"):
@ -458,8 +479,8 @@ class OpenAIResponsesHandler(BaseTranslation):
self,
responses_so_far: list[Any],
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: Any | None = None,
user_api_key_dict: Any | None = None,
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
user_api_key_dict: "UserAPIKeyAuth | None" = None,
request_data: dict | None = None,
) -> list[Any]:
"""
@ -488,10 +509,10 @@ class OpenAIResponsesHandler(BaseTranslation):
# final chunk; iterate output items, apply guardrail, write back. #
# ------------------------------------------------------------------ #
if final_chunk.get("type") == "response.completed":
response_obj: Final = final_chunk.get("response") or {}
response_obj: Final[ResponseOutputEnvelope] = final_chunk.get("response") or {}
if not hasattr(response_obj, "get"):
return responses_so_far
outputs: Final[list[Any]] = response_obj.get("output") or []
outputs: Final[Sequence[object]] = response_obj.get("output") or []
texts_to_check: Final[list[str]] = []
tool_calls_to_check: Final[list[ChatCompletionToolCallChunk]] = []
@ -586,7 +607,7 @@ class OpenAIResponsesHandler(BaseTranslation):
)
return responses_so_far
def _check_streaming_has_ended(self, responses_so_far: list[Any]) -> bool:
def _check_streaming_has_ended(self, responses_so_far: Sequence[ResponsesStreamChunk]) -> bool:
"""
Check if the streaming has ended.
"""
@ -599,7 +620,7 @@ class OpenAIResponsesHandler(BaseTranslation):
}
return responses_so_far[-1].get("type") in terminal_types
def get_streaming_string_so_far(self, responses_so_far: list[Any]) -> str:
def get_streaming_string_so_far(self, responses_so_far: Sequence[ResponsesStreamChunk]) -> str:
"""
Get the string so far from the responses so far.
"""
@ -641,7 +662,7 @@ class OpenAIResponsesHandler(BaseTranslation):
def _extract_output_text_and_images(
self,
output_item: Any,
output_item: object,
output_idx: int,
texts_to_check: list[str],
images_to_check: list[str],
@ -724,7 +745,7 @@ class OpenAIResponsesHandler(BaseTranslation):
async def _apply_guardrail_responses_to_output(
self,
response: Union["ResponsesAPIResponse", dict[Any, Any]],
response: Union["ResponsesAPIResponse", ResponseOutputEnvelope],
responses: list[str],
task_mappings: list[tuple[int, int]],
) -> None:

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

@ -9,9 +9,11 @@ Ref: https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-rest-api
"""
import json
from typing import TYPE_CHECKING, Any, Final
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict
import httpx
from typing_extensions import ReadOnly
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk
from litellm.types.utils import (
@ -44,6 +46,47 @@ _CLAUDE_MODEL_PREFIXES: Final = (
)
class _AnthropicContentBlock(TypedDict, total=False):
type: ReadOnly[str]
text: ReadOnly[str]
id: ReadOnly[str]
name: ReadOnly[str]
input: ReadOnly[Mapping[str, object]]
class _AnthropicUsageBlock(TypedDict, total=False):
input_tokens: ReadOnly[int]
output_tokens: ReadOnly[int]
class _AnthropicMessagesResponse(TypedDict, total=False):
id: ReadOnly[str]
model: ReadOnly[str]
stop_reason: ReadOnly[str]
content: ReadOnly[Sequence[_AnthropicContentBlock]]
usage: ReadOnly[_AnthropicUsageBlock]
class _ChatCompletionsResponse(Protocol):
"""Response view that decodes the Cortex chat-completions body as a field mapping."""
def json(self) -> Mapping[str, object]: ...
class _MessagesResponse(Protocol):
"""Response view that decodes the Cortex messages body in Anthropic shape."""
def json(self) -> _AnthropicMessagesResponse: ...
def _decoded_chat_completions(response: _ChatCompletionsResponse) -> Mapping[str, object]:
return response.json()
def _decoded_messages(response: _MessagesResponse) -> _AnthropicMessagesResponse:
return response.json()
def _is_claude_model(model: str) -> bool:
"""Return True if model name (after stripping snowflake/ prefix) is a Claude model."""
name: Final = model.lower().removeprefix("snowflake/")
@ -129,7 +172,7 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
for tool in tools:
if tool.get("type") == "function" and "function" in tool:
func = tool["function"]
anthropic_tool: dict[str, Any] = {
anthropic_tool: dict[str, object] = {
"name": func.get("name", ""),
}
if "description" in func:
@ -173,7 +216,7 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
elif role == "assistant":
tool_calls = msg.get("tool_calls") if isinstance(msg, dict) else getattr(msg, "tool_calls", None)
if tool_calls:
content_blocks: list[dict[str, Any]] = []
content_blocks: list[dict[str, object]] = []
if content:
content_blocks.append({"type": "text", "text": content})
for tc in tool_calls:
@ -310,7 +353,7 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
model_name: Final = model.removeprefix("snowflake/")
body: Final[dict[str, Any]] = {
body: Final[dict[str, object]] = {
"model": model_name,
"messages": conversation,
"stream": stream,
@ -336,7 +379,7 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
messages: list[AllMessageValues],
optional_params: dict,
litellm_params: dict,
encoding: Any,
encoding: object,
api_key: str | None = None,
json_mode: bool | None = None,
) -> ModelResponse:
@ -356,7 +399,7 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
messages: list[AllMessageValues],
) -> ModelResponse:
"""Parse standard OpenAI chat completions response."""
response_json: Final = raw_response.json()
response_json: Final = _decoded_chat_completions(raw_response)
logging_obj.post_call(
input=messages,
@ -383,7 +426,7 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
messages: list[AllMessageValues],
) -> ModelResponse:
"""Parse Anthropic Messages response into OpenAI format."""
response_json: Final = raw_response.json()
response_json: Final = _decoded_messages(raw_response)
logging_obj.post_call(
input=messages,
@ -447,10 +490,10 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
def get_model_response_iterator(
self,
streaming_response: Any,
streaming_response: object,
sync_stream: bool,
json_mode: bool | None = False,
) -> Any:
) -> "SnowflakeStreamingHandler":
return SnowflakeStreamingHandler(
streaming_response=streaming_response,
sync_stream=sync_stream,
@ -468,7 +511,7 @@ class SnowflakeStreamingHandler(BaseModelResponseIterator):
def __init__(
self,
streaming_response: Any,
streaming_response: object,
sync_stream: bool,
json_mode: bool | None = False,
):

124
litellm/llms/soniox/audio_transcription/handler.py
View file

@ -18,10 +18,11 @@ handler (analogous to the OpenAI / Azure transcription handlers).
import asyncio
import math
import time
from collections.abc import Coroutine
from collections.abc import Coroutine, Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final
import httpx
from typing_extensions import ReadOnly, TypedDict
from litellm.litellm_core_utils.audio_utils.utils import (
get_audio_file_name,
@ -57,6 +58,49 @@ else:
LiteLLMLoggingObj = Any
class _TranscriptionMeta(TypedDict, total=False):
"""Fields the handler reads from a Soniox transcription object."""
status: ReadOnly[str]
error_message: ReadOnly[str]
error_type: ReadOnly[str]
audio_duration_ms: ReadOnly[float]
class _IdentifiedResource(TypedDict):
"""Soniox create/upload response, carrying the new resource id."""
id: ReadOnly[str]
class _SonioxErrorBody(TypedDict, total=False):
"""Fields the handler reads from a Soniox error response body."""
error_message: ReadOnly[object]
error: ReadOnly[object]
class _SonioxJsonView(TypedDict, total=False):
"""Typed reads of decoded Soniox JSON response bodies."""
resource: ReadOnly[_IdentifiedResource]
transcription: ReadOnly[_TranscriptionMeta]
transcript: ReadOnly[Mapping[str, object]]
error: ReadOnly[_SonioxErrorBody]
class _HandlerOptions(TypedDict):
"""Handler-only options pulled out of ``optional_params``."""
poll_interval: ReadOnly[float]
max_attempts: ReadOnly[int]
cleanup: ReadOnly[Sequence[str]]
filename_override: ReadOnly[str | None]
audio_url: ReadOnly[str | None]
file_id: ReadOnly[str | None]
response_format: ReadOnly[str | None]
class SonioxAudioTranscriptionHandler:
"""Orchestrates the Soniox async transcription flow."""
@ -78,9 +122,9 @@ class SonioxAudioTranscriptionHandler:
api_base: str | None,
client: HTTPHandler | AsyncHTTPHandler | None = None,
atranscription: bool = False,
headers: dict[str, Any] | None = None,
headers: dict[str, str] | None = None,
provider_config: SonioxAudioTranscriptionConfig | None = None,
) -> TranscriptionResponse | Coroutine[Any, Any, TranscriptionResponse]:
) -> TranscriptionResponse | Coroutine[object, object, TranscriptionResponse]:
"""Sync/async dispatch for Soniox transcription requests.
Note: ``max_retries`` is accepted for signature compatibility with
@ -134,12 +178,12 @@ class SonioxAudioTranscriptionHandler:
api_key: str | None,
api_base: str | None,
provider_config: SonioxAudioTranscriptionConfig,
headers: dict[str, Any],
headers: dict[str, str],
) -> tuple[
dict[str, str], # auth headers
str, # api_base (no trailing slash)
dict[str, Any], # body for POST /v1/transcriptions (without file_id/audio_url)
dict[str, Any], # handler-only options (poll interval, cleanup, ...)
dict[str, object], # body for POST /v1/transcriptions (without file_id/audio_url)
_HandlerOptions, # handler-only options (poll interval, cleanup, ...)
]:
# Validate env -> auth headers.
auth_headers: Final = provider_config.validate_environment(
@ -184,32 +228,31 @@ class SonioxAudioTranscriptionHandler:
clamped_poll_interval: Final = max(SONIOX_MIN_POLL_INTERVAL, min(poll_interval, SONIOX_MAX_POLL_INTERVAL))
clamped_max_attempts: Final = max(1, min(max_attempts, SONIOX_MAX_POLL_ATTEMPTS))
handler_opts: Final[dict[str, Any]] = {
# response_format is handled by LiteLLM post-processing, not Soniox.
handler_opts: Final[_HandlerOptions] = {
"poll_interval": clamped_poll_interval,
"max_attempts": clamped_max_attempts,
"cleanup": cleanup,
"filename_override": filename_override,
"audio_url": params.pop("audio_url", None),
"file_id": params.pop("file_id", None),
"response_format": params.pop("response_format", None),
}
# Soniox does not accept `language` directly; map_openai_params should
# already have translated it, but drop any leftover to be safe.
params.pop("language", None)
# response_format is handled by LiteLLM post-processing, not Soniox.
handler_opts["response_format"] = params.pop("response_format", None)
return auth_headers, base_url, params, handler_opts
def _build_create_body(
self,
model: str,
optional_params: dict,
handler_opts: dict[str, Any],
optional_params: Mapping[str, object],
handler_opts: _HandlerOptions,
file_id: str | None,
) -> dict[str, Any]:
body: Final[dict[str, Any]] = {"model": model}
) -> dict[str, object]:
body: Final[dict[str, object]] = {"model": model}
# Soniox-native passthrough fields
for key, value in optional_params.items():
if value is None:
@ -224,7 +267,7 @@ class SonioxAudioTranscriptionHandler:
return body
@staticmethod
def _redact_body_for_logging(body: dict[str, Any]) -> dict[str, Any]:
def _redact_body_for_logging(body: dict[str, object]) -> dict[str, object]:
"""Return a shallow copy of ``body`` with secret fields redacted.
Soniox's create-transcription body can include
@ -248,7 +291,7 @@ class SonioxAudioTranscriptionHandler:
logging_obj: LiteLLMLoggingObj,
api_key: str | None,
api_base: str,
body: dict[str, Any],
body: dict[str, object],
) -> None:
try:
logging_obj.pre_call(
@ -270,8 +313,8 @@ class SonioxAudioTranscriptionHandler:
logging_obj: LiteLLMLoggingObj,
audio_file: FileTypes | None,
api_key: str | None,
body: dict[str, Any],
original_response: Any,
body: dict[str, object],
original_response: Mapping[str, object],
) -> None:
try:
logging_obj.post_call(
@ -285,6 +328,11 @@ class SonioxAudioTranscriptionHandler:
# observability integration must never break a real Soniox call.
pass
@staticmethod
def _transcription_meta(response: httpx.Response) -> _TranscriptionMeta:
polled: Final[_SonioxJsonView] = {"transcription": response.json()}
return polled["transcription"]
@staticmethod
def _raise_for_response(
response: httpx.Response,
@ -293,8 +341,8 @@ class SonioxAudioTranscriptionHandler:
) -> None:
if response.status_code >= 400:
try:
payload: Final = response.json()
message = payload.get("error_message") or payload.get("error") or response.text
payload: Final[_SonioxJsonView] = {"error": response.json()}
message = payload["error"].get("error_message") or payload["error"].get("error") or response.text
except Exception:
message = response.text
raise provider_config.get_error_class(
@ -319,7 +367,7 @@ class SonioxAudioTranscriptionHandler:
api_key: str | None,
api_base: str | None,
client: HTTPHandler | None,
headers: dict[str, Any],
headers: dict[str, str],
provider_config: SonioxAudioTranscriptionConfig,
) -> TranscriptionResponse:
auth_headers, base_url, opt_params, handler_opts = self._prepare(
@ -378,7 +426,8 @@ class SonioxAudioTranscriptionHandler:
timeout=timeout,
)
self._raise_for_response(create_resp, provider_config, "create transcription")
transcription_id = create_resp.json()["id"]
created: Final[_SonioxJsonView] = {"resource": create_resp.json()}
transcription_id = created["resource"]["id"]
transcription_meta: Final = self._sync_poll_until_completed(
http_client=http_client,
@ -397,9 +446,9 @@ class SonioxAudioTranscriptionHandler:
timeout=timeout,
)
self._raise_for_response(transcript_resp, provider_config, "fetch transcript")
transcript: Final = transcript_resp.json()
fetched: Final[_SonioxJsonView] = {"transcript": transcript_resp.json()}
payload: Final = {"transcription": transcription_meta, "transcript": transcript}
payload: Final = {"transcription": transcription_meta, "transcript": fetched["transcript"]}
response: Final = provider_config._build_response_from_payload(
payload,
model_response=model_response,
@ -454,7 +503,8 @@ class SonioxAudioTranscriptionHandler:
timeout=timeout,
)
self._raise_for_response(resp, provider_config, "upload file")
return resp.json()["id"]
uploaded: Final[_SonioxJsonView] = {"resource": resp.json()}
return uploaded["resource"]["id"]
def _sync_poll_until_completed(
self,
@ -466,7 +516,7 @@ class SonioxAudioTranscriptionHandler:
max_attempts: int,
timeout: float,
provider_config: SonioxAudioTranscriptionConfig,
) -> dict[str, Any]:
) -> _TranscriptionMeta:
for _ in range(max_attempts):
resp = http_client.get(
url=f"{base_url}/v1/transcriptions/{transcription_id}",
@ -474,7 +524,7 @@ class SonioxAudioTranscriptionHandler:
timeout=timeout,
)
self._raise_for_response(resp, provider_config, "poll transcription")
data = resp.json()
data = self._transcription_meta(resp)
status = data.get("status")
if status == "completed":
return data
@ -502,7 +552,7 @@ class SonioxAudioTranscriptionHandler:
http_client: HTTPHandler,
base_url: str,
auth_headers: dict[str, str],
cleanup: list[str],
cleanup: Sequence[str],
file_id_to_cleanup: str | None,
transcription_id: str | None,
timeout: float,
@ -548,7 +598,7 @@ class SonioxAudioTranscriptionHandler:
api_key: str | None,
api_base: str | None,
client: AsyncHTTPHandler | None,
headers: dict[str, Any],
headers: dict[str, str],
provider_config: SonioxAudioTranscriptionConfig,
) -> TranscriptionResponse:
import litellm
@ -610,7 +660,8 @@ class SonioxAudioTranscriptionHandler:
timeout=timeout,
)
self._raise_for_response(create_resp, provider_config, "create transcription")
transcription_id = create_resp.json()["id"]
created: Final[_SonioxJsonView] = {"resource": create_resp.json()}
transcription_id = created["resource"]["id"]
transcription_meta: Final = await self._async_poll_until_completed(
http_client=http_client,
@ -629,9 +680,9 @@ class SonioxAudioTranscriptionHandler:
timeout=timeout,
)
self._raise_for_response(transcript_resp, provider_config, "fetch transcript")
transcript: Final = transcript_resp.json()
fetched: Final[_SonioxJsonView] = {"transcript": transcript_resp.json()}
payload: Final = {"transcription": transcription_meta, "transcript": transcript}
payload: Final = {"transcription": transcription_meta, "transcript": fetched["transcript"]}
response: Final = provider_config._build_response_from_payload(
payload,
model_response=model_response,
@ -685,7 +736,8 @@ class SonioxAudioTranscriptionHandler:
timeout=timeout,
)
self._raise_for_response(resp, provider_config, "upload file")
return resp.json()["id"]
uploaded: Final[_SonioxJsonView] = {"resource": resp.json()}
return uploaded["resource"]["id"]
async def _async_poll_until_completed(
self,
@ -697,7 +749,7 @@ class SonioxAudioTranscriptionHandler:
max_attempts: int,
timeout: float,
provider_config: SonioxAudioTranscriptionConfig,
) -> dict[str, Any]:
) -> _TranscriptionMeta:
for _ in range(max_attempts):
resp = await http_client.get(
url=f"{base_url}/v1/transcriptions/{transcription_id}",
@ -705,7 +757,7 @@ class SonioxAudioTranscriptionHandler:
timeout=timeout,
)
self._raise_for_response(resp, provider_config, "poll transcription")
data = resp.json()
data = self._transcription_meta(resp)
status = data.get("status")
if status == "completed":
return data
@ -733,7 +785,7 @@ class SonioxAudioTranscriptionHandler:
http_client: AsyncHTTPHandler,
base_url: str,
auth_headers: dict[str, str],
cleanup: list[str],
cleanup: Sequence[str],
file_id_to_cleanup: str | None,
transcription_id: str | None,
timeout: float,

30
litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py
View file

@ -9,10 +9,11 @@ import os
import re
from collections.abc import Mapping, Sequence
from pathlib import PurePosixPath
from typing import Any, Final, TypeAlias, TypedDict
from typing import Any, Final, TypedDict
from urllib.parse import quote
import httpx
from typing_extensions import ReadOnly, Required
# Tool names emitted from OpenAPI specs must work across all major LLM providers.
# OpenAI/Anthropic/Bedrock all enforce a character class roughly equivalent to
@ -47,11 +48,17 @@ from litellm.proxy._experimental.mcp_server.tool_registry import (
global_mcp_tool_registry,
)
_OpenAPIParameter: TypeAlias = Mapping[str, Any]
class _OpenAPIJSONSchema(TypedDict, total=False):
properties: Mapping[str, object]
type: ReadOnly[str]
class _OpenAPIParameter(TypedDict, total=False):
name: Required[ReadOnly[str]]
description: ReadOnly[str]
required: ReadOnly[bool]
schema: ReadOnly[_OpenAPIJSONSchema]
class _OpenAPIMediaType(TypedDict, total=False):
@ -241,7 +248,7 @@ def resolve_operation_params(
operation: _OpenAPIOperation,
path_item: _OpenAPIPathItem,
components: _OpenAPIComponents,
) -> dict[str, Any]:
) -> _OpenAPIOperation:
"""Return a copy of *operation* with fully-resolved, merged parameters.
Handles two common patterns in real-world OpenAPI specs:
@ -261,12 +268,11 @@ def resolve_operation_params(
op_level: Final = _resolve_param_list(operation.get("parameters", []), component_params)
op_keys: Final = {(p["name"], p.get("in")) for p in op_level}
merged: Final = [p for p in path_level if (p["name"], p.get("in")) not in op_keys] + op_level
result: Final = dict(operation)
result["parameters"] = merged
result: Final[_OpenAPIOperation] = {**operation, "parameters": merged}
return result
def extract_parameters(operation: Mapping[str, Any]) -> tuple[Sequence[str], Sequence[str], Sequence[str]]:
def extract_parameters(operation: _OpenAPIOperation) -> tuple[Sequence[str], Sequence[str], Sequence[str]]:
"""Extract parameter names from OpenAPI operation."""
path_params: Final = []
query_params: Final = []
@ -292,7 +298,7 @@ def extract_parameters(operation: Mapping[str, Any]) -> tuple[Sequence[str], Seq
return path_params, query_params, body_params
def build_input_schema(operation: Mapping[str, Any]) -> dict[str, Any]:
def build_input_schema(operation: _OpenAPIOperation) -> dict[str, object]:
"""Build MCP input schema from OpenAPI operation."""
properties: Final = {}
required: Final = []
@ -389,7 +395,7 @@ def _merge_openapi_tool_request_headers(
def create_tool_function(
path: str,
method: str,
operation: Mapping[str, Any],
operation: _OpenAPIOperation,
base_url: str,
headers: dict[str, str] | None = None,
):
@ -443,7 +449,7 @@ def create_tool_function(
url = url.replace("{{" + param_name + "}}", safe_value)
# Build query params using original parameter names
params: Final[dict[str, Any]] = {}
params: Final[dict[str, object]] = {}
for param_name in query_params:
param_value = kwargs.get(param_name, "")
if param_value:
@ -451,7 +457,7 @@ def create_tool_function(
params[param_name] = param_value
# Build request body
json_body: dict[str, Any] | None = None
json_body: dict[str, object] | None = None
if body_params:
# Try "body" first (most common), then check all body param names
body_value = kwargs.get("body", {})
@ -492,7 +498,7 @@ def create_tool_function(
def register_tools_from_openapi(spec: Mapping[str, Any], base_url: str) -> None:
"""Register MCP tools from OpenAPI specification."""
paths: Final[Mapping[str, Mapping[str, Any]]] = spec.get("paths", {})
paths: Final[Mapping[str, Mapping[str, _OpenAPIOperation]]] = spec.get("paths", {})
used_names: Final = set()
for path, path_item in paths.items():

23
litellm/proxy/_experimental/mcp_server/rest_endpoints.py
View file

@ -41,6 +41,7 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
if TYPE_CHECKING:
from mcp.types import CallToolResult
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._experimental.mcp_server.db import OAuthCredentialPayload
from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers
from litellm.types.mcp import MCPAuth
@ -108,7 +109,7 @@ if MCP_AVAILABLE:
########################################################
############ MCP Server REST API Routes #################
async def _safe_fire_mcp_tool_call_logging(
logging_obj: Any | None,
logging_obj: "LiteLLMLoggingObj | None",
result: "CallToolResult",
start_time: datetime,
end_time: datetime,
@ -158,7 +159,7 @@ if MCP_AVAILABLE:
data: dict[str, Any],
tool_name: str,
user_api_key_dict: UserAPIKeyAuth,
) -> Any:
) -> "CallToolResult":
"""Handle the virtual ``mcp_tool_search`` / ``mcp_tool_call`` REST tools (gated on
``mcp_tool_search_enabled``). Kept out of ``call_tool_rest_api`` so that endpoint stays a single
dispatch. An upstream 401 raised by the virtual ``mcp_tool_call`` propagates unhandled to the
@ -298,8 +299,8 @@ if MCP_AVAILABLE:
"""
if not _is_v1_resolved_oauth2_server(server):
return None
user_id: Final = getattr(user_api_key_dict, "user_id", None)
server_id: Final = getattr(server, "server_id", None)
user_id: Final[str | None] = getattr(user_api_key_dict, "user_id", None)
server_id: Final[str | None] = getattr(server, "server_id", None)
if not user_id or not server_id:
return None
try:
@ -343,7 +344,7 @@ if MCP_AVAILABLE:
Returns a dict keyed by server_id. Used to avoid N+1 DB queries when
iterating over multiple OAuth2 MCP servers.
"""
user_id: Final = getattr(user_api_key_dict, "user_id", None)
user_id: Final[str | None] = getattr(user_api_key_dict, "user_id", None)
if not user_id:
return {}
try:
@ -664,7 +665,7 @@ if MCP_AVAILABLE:
"message": "Successfully retrieved tools",
}
def _as_query_str(value: Any) -> str | None:
def _as_query_str(value: object) -> str | None:
"""Coerce an Optional[str] Query param to str|None, dropping unresolved FastAPI defaults."""
return value if isinstance(value, str) else None
@ -935,8 +936,8 @@ if MCP_AVAILABLE:
user_api_key_dict = await acting_user_auth(user_api_key_dict)
data = await request.json()
tool_name: Final = data.get("name")
tool_arguments: Final = data.get("arguments") or {}
tool_name: Final[str | None] = data.get("name")
tool_arguments: Final[dict[str, object]] = data.get("arguments") or {}
from litellm.proxy._experimental.mcp_server.tool_search import (
MCP_TOOL_CALL_TOOL_NAME,
@ -947,7 +948,7 @@ if MCP_AVAILABLE:
return await _handle_virtual_mcp_tool(request, data, tool_name, user_api_key_dict)
# Validate required parameters early
server_id: Final = data.get("server_id")
server_id: Final[str | None] = data.get("server_id")
if not server_id:
raise HTTPException(
status_code=400,
@ -1123,11 +1124,11 @@ if MCP_AVAILABLE:
async def _execute_with_mcp_client(
request: NewMCPServerRequest,
operation: Callable[..., Awaitable[Any]],
operation: Callable[..., Awaitable[Mapping[str, object]]],
mcp_auth_header: str | dict[str, str] | None = None,
oauth2_headers: dict[str, str] | None = None,
raw_headers: dict[str, str] | None = None,
) -> dict:
) -> Mapping[str, object]:
"""
Create a temporary MCP client from *request*, run *operation*, and return the result.

69
litellm/proxy/client/cli/commands/keys.py
View file

@ -1,5 +1,6 @@
import builtins
import json
from collections.abc import Mapping, Sequence
from datetime import datetime
from typing import Any, Final, Literal
@ -7,10 +8,30 @@ import click
import requests
import rich
from rich.table import Table
from typing_extensions import ReadOnly, TypedDict
from ...keys import KeysManagementClient
class _CliContext(TypedDict):
"""Values the top-level CLI group stores on the click context."""
base_url: ReadOnly[str]
api_key: ReadOnly[str | None]
class _CliContextView(TypedDict):
obj: ReadOnly[_CliContext]
class _KeyRowsView(TypedDict):
rows: ReadOnly[Sequence[Mapping[str, object]]]
class _JsonBodyView(TypedDict):
body: ReadOnly[object]
@click.group()
def keys():
"""Manage API keys for the LiteLLM proxy server"""
@ -53,7 +74,8 @@ def list(
return_full_object: bool,
):
"""List all API keys"""
client: Final = KeysManagementClient(ctx.obj["base_url"], ctx.obj["api_key"])
context: Final[_CliContextView] = {"obj": ctx.obj}
client: Final = KeysManagementClient(context["obj"]["base_url"], context["obj"]["api_key"])
response: Final = client.list(
page=page,
size=size,
@ -70,14 +92,16 @@ def list(
if output_format == "json":
rich.print_json(data=response)
else:
rich.print(f"Showing {len(response.get('keys', []))} keys out of {response.get('total_count', 0)}")
listed: Final[_KeyRowsView] = {"rows": response.get("keys", [])}
rich.print(f"Showing {len(listed['rows'])} keys out of {response.get('total_count', 0)}")
table: Final = Table(title="API Keys")
table.add_column("Key Hash", style="cyan")
table.add_column("Alias", style="green")
table.add_column("User ID", style="magenta")
table.add_column("Team ID", style="yellow")
table.add_column("Spend", style="red")
for key in response.get("keys", []):
key_rows: Final[_KeyRowsView] = {"rows": response.get("keys", [])}
for key in key_rows["rows"]:
table.add_row(
str(key.get("token", "")),
str(key.get("key_alias", "")),
@ -116,7 +140,8 @@ def generate(
config: str | None,
):
"""Generate a new API key"""
client: Final = KeysManagementClient(ctx.obj["base_url"], ctx.obj["api_key"])
context: Final[_CliContextView] = {"obj": ctx.obj}
client: Final = KeysManagementClient(context["obj"]["base_url"], context["obj"]["api_key"])
try:
models_list: Final = [m.strip() for m in models.split(",")] if models else None
aliases_dict: Final = json.loads(aliases) if aliases else None
@ -139,8 +164,8 @@ def generate(
except requests.exceptions.HTTPError as e:
click.echo(f"Error: HTTP {e.response.status_code}", err=True)
try:
error_body: Final = e.response.json()
rich.print_json(data=error_body)
error_body: Final[_JsonBodyView] = {"body": e.response.json()}
rich.print_json(data=error_body["body"])
except json.JSONDecodeError:
click.echo(e.response.text, err=True)
raise click.Abort()
@ -152,7 +177,8 @@ def generate(
@click.pass_context
def delete(ctx: click.Context, keys: str | None, key_aliases: str | None):
"""Delete API keys by key or alias"""
client: Final = KeysManagementClient(ctx.obj["base_url"], ctx.obj["api_key"])
context: Final[_CliContextView] = {"obj": ctx.obj}
client: Final = KeysManagementClient(context["obj"]["base_url"], context["obj"]["api_key"])
keys_list: Final = [k.strip() for k in keys.split(",")] if keys else None
aliases_list: Final = [a.strip() for a in key_aliases.split(",")] if key_aliases else None
try:
@ -161,8 +187,8 @@ def delete(ctx: click.Context, keys: str | None, key_aliases: str | None):
except requests.exceptions.HTTPError as e:
click.echo(f"Error: HTTP {e.response.status_code}", err=True)
try:
error_body: Final = e.response.json()
rich.print_json(data=error_body)
error_body: Final[_JsonBodyView] = {"body": e.response.json()}
rich.print_json(data=error_body["body"])
except json.JSONDecodeError:
click.echo(e.response.text, err=True)
raise click.Abort()
@ -189,10 +215,10 @@ def _parse_created_since_filter(created_since: str | None) -> datetime | None:
def _fetch_all_keys_with_pagination(
source_client: KeysManagementClient, source_base_url: str
) -> builtins.list[dict[str, Any]]:
) -> Sequence[Mapping[str, object]]:
"""Fetch all keys from source instance using pagination."""
click.echo(f"Fetching keys from source server: {source_base_url}")
source_keys: Final = []
source_keys: Final[builtins.list[Mapping[str, object]]] = []
page = 1
page_size: Final = 100 # Use a larger page size to minimize API calls
@ -200,7 +226,7 @@ def _fetch_all_keys_with_pagination(
source_response = source_client.list(return_full_object=True, page=page, size=page_size)
# source_client.list() returns Dict[str, Any] when return_request is False (default)
assert isinstance(source_response, dict), "Expected dict response from list API"
page_keys = source_response.get("keys", [])
page_keys: Sequence[Mapping[str, object]] = source_response.get("keys", [])
if not page_keys:
break
@ -218,15 +244,15 @@ def _fetch_all_keys_with_pagination(
def _filter_keys_by_created_since(
source_keys: builtins.list[dict[str, Any]],
source_keys: Sequence[Mapping[str, object]],
created_since_dt: datetime | None,
created_since: str,
) -> builtins.list[dict[str, Any]]:
) -> Sequence[Mapping[str, object]]:
"""Filter keys by created_since date if specified."""
if not created_since_dt:
return source_keys
filtered_keys: Final = []
filtered_keys: Final[builtins.list[Mapping[str, object]]] = []
for key in source_keys:
key_created_at = key.get("created_at")
if key_created_at:
@ -248,7 +274,7 @@ def _filter_keys_by_created_since(
return filtered_keys
def _display_dry_run_table(source_keys: builtins.list[dict[str, Any]]) -> None:
def _display_dry_run_table(source_keys: Sequence[Mapping[str, object]]) -> None:
"""Display a table of keys that would be imported in dry-run mode."""
click.echo("\n--- DRY RUN MODE ---")
table: Final = Table(title="Keys that would be imported")
@ -271,7 +297,7 @@ def _display_dry_run_table(source_keys: builtins.list[dict[str, Any]]) -> None:
rich.print(table)
def _prepare_key_import_data(key: dict[str, Any]) -> dict[str, Any]:
def _prepare_key_import_data(key: Mapping[str, object]) -> dict[str, Any]:
"""Prepare key data for import by extracting relevant fields."""
import_data: Final = {}
@ -293,7 +319,7 @@ def _prepare_key_import_data(key: dict[str, Any]) -> dict[str, Any]:
def _import_keys_to_destination(
source_keys: builtins.list[dict[str, Any]], dest_client: KeysManagementClient
source_keys: Sequence[Mapping[str, object]], dest_client: KeysManagementClient
) -> tuple[int, int]:
"""Import each key to the destination instance and return counts."""
imported_count = 0
@ -351,7 +377,8 @@ def import_keys(
# Create clients for both source and destination
source_client: Final = KeysManagementClient(source_base_url, source_api_key)
dest_client: Final = KeysManagementClient(ctx.obj["base_url"], ctx.obj["api_key"])
context: Final[_CliContextView] = {"obj": ctx.obj}
dest_client: Final = KeysManagementClient(context["obj"]["base_url"], context["obj"]["api_key"])
try:
# Get all keys from source instance with pagination
@ -383,8 +410,8 @@ def import_keys(
except requests.exceptions.HTTPError as e:
click.echo(f"Error: HTTP {e.response.status_code}", err=True)
try:
error_body: Final = e.response.json()
rich.print_json(data=error_body)
error_body: Final[_JsonBodyView] = {"body": e.response.json()}
rich.print_json(data=error_body["body"])
except json.JSONDecodeError:
click.echo(e.response.text, err=True)
raise click.Abort()

90
litellm/proxy/guardrails/guardrail_hooks/aim/aim.py
View file

@ -7,10 +7,11 @@
import asyncio
import json
import os
from collections.abc import AsyncGenerator
from typing import TYPE_CHECKING, Any, Final
from collections.abc import AsyncGenerator, AsyncIterator, Mapping, Sequence
from typing import TYPE_CHECKING, Final, TypeAlias
from pydantic import BaseModel
from typing_extensions import NotRequired, ReadOnly, TypedDict
from websockets.asyncio.client import ClientConnection, connect
from litellm import DualCache
@ -31,8 +32,7 @@ from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.utils import (
CallTypesLiteral,
Choices,
EmbeddingResponse,
ImageResponse,
LLMResponseTypes,
ModelResponse,
ModelResponseStream,
)
@ -45,6 +45,58 @@ class AimGuardrailMissingSecrets(Exception):
pass
class AimRequiredAction(TypedDict):
"""The ``required_action`` block of an Aim ``/fw/v1/analyze`` response."""
action_type: ReadOnly[NotRequired[str]]
detection_message: ReadOnly[str]
class AimAnalysisResult(TypedDict):
"""The ``analysis_result`` block of an Aim ``/fw/v1/analyze`` response."""
policy_drill_down: ReadOnly[Mapping[str, object]]
class AimRedactedMessage(TypedDict):
"""One entry of Aim's ``redacted_chat.all_redacted_messages``."""
role: ReadOnly[str]
content: ReadOnly[str]
class AimRedactedChat(TypedDict):
"""The ``redacted_chat`` block of an Aim ``/fw/v1/analyze`` response."""
all_redacted_messages: ReadOnly[Sequence[AimRedactedMessage]]
class AimAnalyzeResponse(TypedDict):
"""Body returned by Aim's ``POST /fw/v1/analyze``."""
required_action: ReadOnly[AimRequiredAction]
analysis_result: ReadOnly[AimAnalysisResult]
redacted_chat: ReadOnly[NotRequired[AimRedactedChat]]
class AimOutputGuardrailResult(TypedDict, total=False):
"""Outcome of inspecting one model completion with Aim."""
detection_message: ReadOnly[str]
redacted_output: ReadOnly[str]
class AimStreamMessage(TypedDict, total=False):
"""One frame of Aim's ``/fw/v1/analyze/stream`` websocket protocol."""
verified_chunk: ReadOnly[Mapping[str, object]]
done: ReadOnly[bool]
blocking_message: ReadOnly[str]
AimStreamChunk: TypeAlias = BaseModel | Mapping[str, object] | str | bytes
class AimGuardrail(CustomGuardrail):
@classmethod
def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]:
@ -110,7 +162,7 @@ class AimGuardrail(CustomGuardrail):
json={"messages": self._build_aim_inspection_messages(data)},
)
response.raise_for_status()
res: Final = response.json()
res: Final[AimAnalyzeResponse] = response.json()
required_action: Final = res.get("required_action")
action_type: Final = required_action and required_action.get("action_type", None)
if action_type is None:
@ -145,7 +197,7 @@ class AimGuardrail(CustomGuardrail):
openai_code=openai_code,
)
def _handle_block_action(self, analysis_result: Any, required_action: Any) -> None:
def _handle_block_action(self, analysis_result: AimAnalysisResult, required_action: AimRequiredAction) -> None:
detection_message: Final = required_action.get("detection_message", None)
verbose_proxy_logger.info(
"Aim: Violation detected enabled policies: {policies}".format(
@ -154,7 +206,7 @@ class AimGuardrail(CustomGuardrail):
)
raise self._rejection(detection_message, openai_code="content_policy_violation")
def _anonymize_request(self, res: Any, data: dict) -> dict:
def _anonymize_request(self, res: AimAnalyzeResponse, data: dict) -> dict:
verbose_proxy_logger.info("Aim: anonymize action")
redacted_chat: Final = res.get("redacted_chat")
if not redacted_chat:
@ -185,7 +237,7 @@ class AimGuardrail(CustomGuardrail):
async def call_aim_guardrail_on_output(
self, request_data: dict, output: str, hook: str, key_alias: str | None
) -> dict | None:
) -> AimOutputGuardrailResult | None:
user_email: Final = request_data.get("metadata", {}).get("headers", {}).get("x-aim-user-email")
call_id: Final = request_data.get("litellm_call_id")
response: Final = await self.async_handler.post(
@ -202,7 +254,7 @@ class AimGuardrail(CustomGuardrail):
},
)
response.raise_for_status()
res: Final = response.json()
res: Final[AimAnalyzeResponse] = response.json()
required_action: Final = res.get("required_action")
action_type: Final = required_action and required_action.get("action_type", None)
if action_type and action_type == "block_action":
@ -213,7 +265,9 @@ class AimGuardrail(CustomGuardrail):
return {"redacted_output": redacted_chat["all_redacted_messages"][-1]["content"]}
return {"redacted_output": output}
def _handle_block_action_on_output(self, analysis_result: Any, required_action: Any) -> dict | None:
def _handle_block_action_on_output(
self, analysis_result: AimAnalysisResult, required_action: AimRequiredAction
) -> AimOutputGuardrailResult | None:
detection_message: Final = required_action.get("detection_message", None)
verbose_proxy_logger.info(
"Aim: detected: {detected}, enabled policies: {policies}".format(
@ -260,8 +314,8 @@ class AimGuardrail(CustomGuardrail):
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
response: Any | ModelResponse | EmbeddingResponse | ImageResponse,
) -> Any:
response: LLMResponseTypes,
) -> LLMResponseTypes:
if not (isinstance(response, ModelResponse) and response.choices):
return response
# Inspect every choice — when ``n>1`` the additional completions
@ -289,9 +343,11 @@ class AimGuardrail(CustomGuardrail):
for choice, aim_output_guardrail_result in zip(choices_to_inspect, results):
if isinstance(aim_output_guardrail_result, BaseException):
raise aim_output_guardrail_result
if aim_output_guardrail_result and aim_output_guardrail_result.get("detection_message"):
if aim_output_guardrail_result and (
detection_message := aim_output_guardrail_result.get("detection_message")
):
raise self._rejection(
aim_output_guardrail_result.get("detection_message"),
detection_message,
openai_code="content_policy_violation",
)
if aim_output_guardrail_result and aim_output_guardrail_result.get("redacted_output"):
@ -301,7 +357,7 @@ class AimGuardrail(CustomGuardrail):
async def async_post_call_streaming_iterator_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
response,
response: AsyncIterator[AimStreamChunk],
request_data: dict,
) -> AsyncGenerator[ModelResponseStream, None]:
user_email: Final = request_data.get("metadata", {}).get("headers", {}).get("x-aim-user-email")
@ -317,7 +373,7 @@ class AimGuardrail(CustomGuardrail):
) as websocket:
sender: Final = asyncio.create_task(self.forward_the_stream_to_aim(websocket, response))
while True:
result = json.loads(await websocket.recv())
result: AimStreamMessage = json.loads(await websocket.recv())
if verified_chunk := result.get("verified_chunk"):
yield ModelResponseStream.model_validate(verified_chunk)
else:
@ -334,7 +390,7 @@ class AimGuardrail(CustomGuardrail):
async def forward_the_stream_to_aim(
self,
websocket: ClientConnection,
response_iter,
response_iter: AsyncIterator[AimStreamChunk],
) -> None:
async for chunk in response_iter:
if isinstance(chunk, BaseModel):

79
litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py
View file

@ -7,10 +7,13 @@ and provide safe, sandboxed functionality for common guardrail operations.
import json
import re
from collections.abc import Mapping, Sequence
from typing import Any, Final
from urllib.parse import urlparse
import httpx
from pydantic import JsonValue
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
@ -21,7 +24,7 @@ from litellm.types.llms.custom_http import httpxSpecialProvider
# =============================================================================
def allow() -> dict[str, Any]:
def allow() -> dict[str, object]:
"""
Allow the request/response to proceed unchanged.
@ -31,7 +34,7 @@ def allow() -> dict[str, Any]:
return {"action": "allow"}
def block(reason: str, detection_info: dict[str, Any] | None = None) -> dict[str, Any]:
def block(reason: str, detection_info: Mapping[str, object] | None = None) -> dict[str, object]:
"""
Block the request/response with a reason.
@ -42,17 +45,17 @@ def block(reason: str, detection_info: dict[str, Any] | None = None) -> dict[str
Returns:
Dict indicating the request should be blocked
"""
result: Final[dict[str, Any]] = {"action": "block", "reason": reason}
result: Final[dict[str, object]] = {"action": "block", "reason": reason}
if detection_info:
result["detection_info"] = detection_info
return result
def modify(
texts: list[str] | None = None,
images: list[Any] | None = None,
tool_calls: list[Any] | None = None,
) -> dict[str, Any]:
texts: Sequence[str] | None = None,
images: Sequence[object] | None = None,
tool_calls: Sequence[object] | None = None,
) -> dict[str, object]:
"""
Modify the request/response content.
@ -64,7 +67,7 @@ def modify(
Returns:
Dict indicating the content should be modified
"""
result: Final[dict[str, Any]] = {"action": "modify"}
result: Final[dict[str, object]] = {"action": "modify"}
if texts is not None:
result["texts"] = texts
if images is not None:
@ -161,7 +164,15 @@ def regex_find_all(text: str, pattern: str, flags: int = 0) -> list[str]:
# =============================================================================
def json_parse(text: str) -> Any | None:
class JsonSchemaNode(TypedDict, total=False):
"""Subset of JSON Schema keywords understood by the built-in validator."""
type: ReadOnly[str]
required: ReadOnly[Sequence[str]]
properties: ReadOnly[Mapping[str, "JsonSchemaNode"]]
def json_parse(text: str) -> JsonValue:
"""
Parse a JSON string into a Python object.
@ -178,7 +189,7 @@ def json_parse(text: str) -> Any | None:
return None
def json_stringify(obj: Any) -> str:
def json_stringify(obj: object) -> str:
"""
Convert a Python object to a JSON string.
@ -195,7 +206,7 @@ def json_stringify(obj: Any) -> str:
return ""
def json_schema_valid(obj: Any, schema: dict[str, Any]) -> bool:
def json_schema_valid(obj: JsonValue, schema: JsonSchemaNode) -> bool:
"""
Validate an object against a JSON schema.
@ -226,7 +237,7 @@ def json_schema_valid(obj: Any, schema: dict[str, Any]) -> bool:
return False
def _basic_json_schema_validate(obj: Any, schema: dict[str, Any], max_depth: int = 50) -> bool:
def _basic_json_schema_validate(obj: JsonValue, schema: JsonSchemaNode, max_depth: int = 50) -> bool:
"""
Basic JSON schema validation without external library.
Handles: type, required, properties
@ -234,7 +245,7 @@ def _basic_json_schema_validate(obj: Any, schema: dict[str, Any], max_depth: int
Uses an iterative approach with a stack to avoid recursion limits.
max_depth limits nesting to prevent infinite loops from circular schemas.
"""
type_map: Final[dict[str, type | tuple[type, ...]]] = {
type_map: Final[Mapping[str, type | tuple[type, ...]]] = {
"object": dict,
"array": list,
"string": str,
@ -245,7 +256,7 @@ def _basic_json_schema_validate(obj: Any, schema: dict[str, Any], max_depth: int
}
# Stack of (obj, schema, depth) tuples to process
stack: Final[list[tuple[Any, dict[str, Any], int]]] = [(obj, schema, 0)]
stack: Final[list[tuple[JsonValue, JsonSchemaNode, int]]] = [(obj, schema, 0)]
while stack:
current_obj, current_schema, depth = stack.pop()
@ -257,19 +268,19 @@ def _basic_json_schema_validate(obj: Any, schema: dict[str, Any], max_depth: int
# Check type
schema_type = current_schema.get("type")
if schema_type:
expected_type = type_map.get(schema_type)
expected_type: type | tuple[type, ...] | None = type_map.get(schema_type)
if expected_type is not None and not isinstance(current_obj, expected_type):
return False
# Check required fields and properties for dicts
if isinstance(current_obj, dict):
required = current_schema.get("required", [])
required: Sequence[str] = current_schema.get("required", [])
for field in required:
if field not in current_obj:
return False
# Queue property validations
properties = current_schema.get("properties", {})
properties: Mapping[str, JsonSchemaNode] = current_schema.get("properties", {})
for prop_name, prop_schema in properties.items():
if prop_name in current_obj:
stack.append((current_obj[prop_name], prop_schema, depth + 1))
@ -358,7 +369,17 @@ _HTTP_DEFAULT_TIMEOUT: Final = 30.0
_HTTP_MAX_TIMEOUT: Final = 60.0
def _http_error_response(error: str) -> dict[str, Any]:
class HttpResponseResult(TypedDict):
"""Outcome of an HTTP primitive call, as handed back to custom code."""
status_code: ReadOnly[int]
body: ReadOnly[JsonValue]
headers: ReadOnly[Mapping[str, str]]
success: ReadOnly[bool]
error: ReadOnly[str | None]
def _http_error_response(error: str) -> HttpResponseResult:
"""Create a standardized error response for HTTP requests."""
return {
"status_code": 0,
@ -369,9 +390,9 @@ def _http_error_response(error: str) -> dict[str, Any]:
}
def _http_success_response(response: httpx.Response) -> dict[str, Any]:
def _http_success_response(response: httpx.Response) -> HttpResponseResult:
"""Create a standardized success response from an httpx Response."""
parsed_body: Any
parsed_body: JsonValue
try:
parsed_body = response.json()
except (json.JSONDecodeError, ValueError):
@ -387,8 +408,8 @@ def _http_success_response(response: httpx.Response) -> dict[str, Any]:
def _prepare_http_body(
body: Any | None,
) -> tuple[dict[str, Any] | None, str | None]:
body: JsonValue,
) -> tuple[dict[str, JsonValue] | None, str | None]:
"""Prepare body arguments for HTTP request - returns (json_body, data_body)."""
if body is None:
return None, None
@ -405,9 +426,9 @@ async def http_request(
url: str,
method: str = "GET",
headers: dict[str, str] | None = None,
body: Any | None = None,
body: JsonValue = None,
timeout: float | None = None,
) -> dict[str, Any]:
) -> HttpResponseResult:
"""
Make an async HTTP request to an external service.
@ -491,7 +512,7 @@ async def _execute_http_request(
method: str,
url: str,
headers: dict[str, str] | None,
body: Any | None,
body: JsonValue,
timeout: float,
) -> httpx.Response:
"""Execute the HTTP request using the appropriate client method."""
@ -515,7 +536,7 @@ async def http_get(
url: str,
headers: dict[str, str] | None = None,
timeout: float | None = None,
) -> dict[str, Any]:
) -> HttpResponseResult:
"""
Make an async HTTP GET request.
@ -534,10 +555,10 @@ async def http_get(
async def http_post(
url: str,
body: Any | None = None,
body: JsonValue = None,
headers: dict[str, str] | None = None,
timeout: float | None = None,
) -> dict[str, Any]:
) -> HttpResponseResult:
"""
Make an async HTTP POST request.
@ -755,7 +776,7 @@ def trim(text: str) -> str:
# =============================================================================
def get_custom_code_primitives() -> dict[str, Any]:
def get_custom_code_primitives() -> dict[str, object]:
"""
Get all primitives to inject into the custom code environment.

46
litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py
View file

@ -2,7 +2,7 @@ from __future__ import annotations
import os
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypedDict
from urllib.parse import urlparse
from uuid import uuid4
@ -76,6 +76,36 @@ class _HiddenlayerChoice(TypedDict, total=False):
message: ReadOnly[_HiddenlayerChoiceMessage]
class _HiddenlayerV2Output(TypedDict, total=False):
messages: ReadOnly[Sequence[_HiddenlayerOutputMessage]]
choices: ReadOnly[Sequence[_HiddenlayerChoice]]
class _LoggedCallDetails(Protocol):
"""Logging object view that exposes its untyped call details with the shape this guardrail reads."""
@property
def model_call_details(self) -> Mapping[str, _LoggedCallLitellmParams]: ...
class _TokenPayloadSource(Protocol):
"""Response view that decodes the HiddenLayer OAuth token body as a string mapping."""
def json(self) -> Mapping[str, str]: ...
def _logged_request_headers(logging_obj: _LoggedCallDetails) -> Mapping[str, str]:
return logging_obj.model_call_details.get("litellm_params", {}).get("metadata", {}).get("headers", {})
def _token_payload(response: _TokenPayloadSource) -> Mapping[str, str]:
return response.json()
def _header_value(headers: Mapping[str, str], key: str, default: str) -> str:
return headers.get(key, default)
def is_saas(host: str) -> bool:
"""Checks whether the connection is to the SaaS platform"""
@ -102,7 +132,7 @@ def _get_jwt(auth_url, api_id, api_key) -> str:
f"Unable to get authentication credentials for the HiddenLayer API - invalid response: {resp.json()}"
)
return resp.json()["access_token"]
return _token_payload(resp)["access_token"]
class HiddenlayerGuardrail(CustomGuardrail):
@ -176,10 +206,7 @@ class HiddenlayerGuardrail(CustomGuardrail):
# from the logger object on the response from the model.
headers = request_data.get("proxy_server_request", {}).get("headers", {})
if not headers and logging_obj and logging_obj.model_call_details:
logged_litellm_params: Final[_LoggedCallLitellmParams] = logging_obj.model_call_details.get(
"litellm_params", {}
)
headers = logged_litellm_params.get("metadata", {}).get("headers", {})
headers = _logged_request_headers(logging_obj)
hl_request_metadata["requester_id"] = headers.get("hl-requester-id") or "LiteLLM"
project_id: Final = headers.get("hl-project-id")
@ -418,8 +445,9 @@ class HiddenlayerGuardrailV2(CustomGuardrail):
response: Final = await self._call_hiddenlayer(payload, input_type, hl_headers)
output: Final = response.json()
evaluated_output: Final[_HiddenlayerV2Output] = output
if response.headers.get("hl-runtime-action", "").lower() == "block":
if _header_value(response.headers, "hl-runtime-action", "").lower() == "block":
raise HTTPException(
status_code=400,
detail={
@ -432,7 +460,7 @@ class HiddenlayerGuardrailV2(CustomGuardrail):
if input_type == "request":
inputs["structured_messages"] = output
modified_messages: Final[Sequence[_HiddenlayerOutputMessage]] = output.get("messages", [])
modified_messages: Final[Sequence[_HiddenlayerOutputMessage]] = evaluated_output.get("messages", [])
for message in modified_messages:
content = message.get("content", "")
if isinstance(content, list):
@ -447,7 +475,7 @@ class HiddenlayerGuardrailV2(CustomGuardrail):
inputs["texts"] = new_texts
elif input_type == "response" and inputs.get("texts"):
redacted_choices: Final[Sequence[_HiddenlayerChoice]] = output.get("choices", [{}])
redacted_choices: Final[Sequence[_HiddenlayerChoice]] = evaluated_output.get("choices", [{}])
inputs["texts"] = [redacted_choices[-1].get("message", {}).get("content", "")]
elif input_type == "response" and inputs.get("tool_calls"):
inputs["tool_calls"] = output

49
litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/purview_dlp.py
View file

@ -10,9 +10,9 @@ Supports three modes:
import asyncio
import threading
import uuid
from collections.abc import AsyncGenerator
from collections.abc import AsyncGenerator, AsyncIterable, Mapping, Sequence
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final, Union, cast
from typing import TYPE_CHECKING, Any, Final, cast
import httpx
from fastapi import HTTPException
@ -36,14 +36,15 @@ from litellm.types.utils import (
from .base import PurviewGuardrailBase
if TYPE_CHECKING:
from litellm.caching.dual_cache import DualCache
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.llms.openai import AllMessageValues
from litellm.types.proxy.guardrails.guardrail_hooks.base import (
GuardrailConfigModel,
)
from litellm.types.utils import (
CallTypesLiteral,
EmbeddingResponse,
ImageResponse,
LLMResponseTypes,
)
@ -63,7 +64,7 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail):
client_secret: str,
purview_app_name: str = "LiteLLM",
user_id_field: str = "user_id",
**kwargs: Any,
**kwargs: object,
):
super().__init__(
tenant_id=tenant_id,
@ -104,7 +105,7 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail):
activity: str,
request_data: dict[str, Any],
block_on_violation: bool = True,
) -> dict[str, Any]:
) -> dict[str, object]:
"""Evaluate content against Purview DLP policies.
Args:
@ -119,7 +120,7 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail):
"""
start_time: Final = datetime.now()
status: GuardrailStatus = "success"
response: dict[str, Any] = {}
response: dict[str, object] = {}
try:
etag, _ = await self._compute_protection_scopes(user_id)
@ -149,7 +150,7 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail):
upstream_status: Final = exc.response.status_code
client_status: Final = 502 if upstream_status in (401, 403) else upstream_status
headers: dict[str, str] | None = None
retry_after: Final = exc.response.headers.get("retry-after")
retry_after: Final[str | None] = exc.response.headers.get("retry-after")
if retry_after:
headers = {"Retry-After": retry_after}
raise HTTPException(
@ -205,7 +206,7 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail):
return response
@staticmethod
def _extract_responses_api_function_call_args(result: Any) -> list[str]:
def _extract_responses_api_function_call_args(result: object) -> list[str]:
"""Return tool-call argument strings from a ``ResponsesAPIResponse.output``.
``ResponsesAPIResponse.output_text`` only aggregates ``output_text``
@ -215,7 +216,7 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail):
chat (``ModelResponse``) path.
"""
args: Final[list[str]] = []
output: Final = getattr(result, "output", None)
output: Final[Sequence[object] | None] = getattr(result, "output", None)
if not output:
return args
for item in output:
@ -230,7 +231,7 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail):
args.append(arguments)
return args
def _completion_response_text_parts(self, result: Any) -> list[str]:
def _completion_response_text_parts(self, result: object) -> list[str]:
"""Collect non-empty text segments from chat, text completions, or responses API.
Includes assistant message content *and* model-generated tool-call
@ -266,7 +267,7 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail):
parts.extend(self._extract_tool_call_args_from_message(msg))
return parts
def _assemble_responses_api_from_chunks(self, chunks: list[Any]) -> tuple[bool, ResponsesAPIResponse | None]:
def _assemble_responses_api_from_chunks(self, chunks: Sequence[object]) -> tuple[bool, ResponsesAPIResponse | None]:
"""Extract the final ``ResponsesAPIResponse`` from a buffered Responses API stream.
Returns a ``(is_responses_api_stream, assembled)`` tuple so the caller
@ -314,7 +315,7 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail):
input=input_data if input_data is not None else "",
responses_api_request=data,
)
return self.get_prompt_text_for_dlp(cast(list[Any], messages))
return self.get_prompt_text_for_dlp(cast(list["AllMessageValues"], messages))
except Exception:
verbose_proxy_logger.warning(
"Purview DLP: failed to transform responses API input",
@ -338,8 +339,8 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail):
def _resolve_user_id_for_blocking(
self,
data: dict[str, Any],
user_api_key_dict: Any,
data: Mapping[str, object],
user_api_key_dict: "UserAPIKeyAuth",
) -> str:
"""Resolve user ID for blocking (pre_call / post_call) DLP hooks.
@ -386,10 +387,10 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail):
async def async_pre_call_hook(
self,
user_api_key_dict: "UserAPIKeyAuth",
cache: Any,
cache: "DualCache",
data: dict[str, Any],
call_type: "CallTypesLiteral",
) -> dict[str, Any] | None:
) -> dict[str, object] | None:
"""Check user prompt against Purview DLP policies before LLM call."""
user_id: Final = self._resolve_user_id_for_blocking(data, user_api_key_dict)
@ -423,7 +424,7 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail):
else:
messages: Final[list | None] = data.get("messages")
if messages:
prompt_text = self.get_prompt_text_for_dlp(cast(list[Any], messages))
prompt_text = self.get_prompt_text_for_dlp(cast(list["AllMessageValues"], messages))
if not prompt_text:
return data
@ -446,8 +447,8 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail):
self,
data: dict,
user_api_key_dict: "UserAPIKeyAuth",
response: Union[Any, ModelResponse, "EmbeddingResponse", "ImageResponse"],
) -> Any:
response: "LLMResponseTypes",
) -> "LLMResponseTypes":
"""Check LLM response against Purview DLP policies (non-streaming only).
Streaming responses are handled by ``async_post_call_streaming_iterator_hook``
@ -472,7 +473,7 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail):
async def async_post_call_streaming_iterator_hook(
self,
user_api_key_dict: "UserAPIKeyAuth",
response: Any,
response: AsyncIterable[ModelResponseStream],
request_data: dict,
) -> AsyncGenerator[ModelResponseStream, None]:
"""Check streaming LLM responses against Purview DLP policies.
@ -592,7 +593,7 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail):
# Logging-only hook — audit without blocking
# ------------------------------------------------------------------
def logging_hook(self, kwargs: dict, result: Any, call_type: str) -> tuple[dict, Any]:
def logging_hook(self, kwargs: dict, result: object, call_type: str) -> tuple[dict, object]:
"""Fire-and-forget async audit logging; returns original (kwargs, result) immediately.
In the proxy's async success path, litellm independently calls both
@ -640,7 +641,7 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail):
return kwargs, result
async def async_logging_hook(self, kwargs: dict, result: Any, call_type: str) -> tuple[dict, Any]:
async def async_logging_hook(self, kwargs: dict, result: object, call_type: str) -> tuple[dict, object]:
"""Send both prompt and response to Purview for audit logging.
Errors are logged but never raised — this mode is non-blocking.
@ -670,7 +671,7 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail):
else:
messages: Final = kwargs.get("messages")
if messages:
prompt_text = self.get_prompt_text_for_dlp(cast(list[Any], messages))
prompt_text = self.get_prompt_text_for_dlp(cast(list["AllMessageValues"], messages))
if prompt_text:
await self._check_content(

53
litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py
View file

@ -1,9 +1,11 @@
import asyncio
import base64
import os
from typing import TYPE_CHECKING, Any, Final, Literal, Optional
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Final, Literal, Optional
from fastapi import HTTPException
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_guardrail import (
@ -26,6 +28,41 @@ class PromptSecurityGuardrailMissingSecrets(Exception):
pass
class _ProtectVerdict(TypedDict, total=False):
"""One side (``prompt`` or ``response``) of an ``/api/protect`` verdict."""
action: ReadOnly[str]
violations: ReadOnly[Sequence[str]]
modified_messages: ReadOnly[Sequence[Mapping[str, object]]]
modified_text: ReadOnly[str]
class _ProtectResult(TypedDict, total=False):
prompt: ReadOnly[_ProtectVerdict | None]
response: ReadOnly[_ProtectVerdict | None]
class _ProtectResponse(TypedDict, total=False):
result: ReadOnly[_ProtectResult]
class _SanitizeUploadResponse(TypedDict, total=False):
jobId: ReadOnly[str]
class _SanitizeMetadata(TypedDict, total=False):
action: ReadOnly[str]
violations: ReadOnly[Sequence[str]]
class _SanitizeStatusResponse(TypedDict, total=False):
"""One poll of ``/api/sanitizeFile``."""
status: ReadOnly[str]
content: ReadOnly[str]
metadata: ReadOnly[_SanitizeMetadata]
class PromptSecurityGuardrail(CustomGuardrail):
@classmethod
def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]:
@ -199,7 +236,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
json=payload,
)
response.raise_for_status()
res: Final = response.json()
res: Final[_ProtectResponse] = response.json()
self._log_api_response(
url=f"{self.api_base}/api/protect",
@ -261,7 +298,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
json=payload,
)
response.raise_for_status()
res: Final = response.json()
res: Final[_ProtectResponse] = response.json()
self._log_api_response(
url=f"{self.api_base}/api/protect",
@ -290,7 +327,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
return inputs
def _extract_texts_from_messages(self, messages: list) -> list[str]:
def _extract_texts_from_messages(self, messages: Sequence[Mapping[str, object]]) -> list[str]:
"""Extract text content from messages."""
texts: Final = []
for message in messages:
@ -379,7 +416,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
files=files,
)
upload_response.raise_for_status()
upload_result: Final = upload_response.json()
upload_result: Final[_SanitizeUploadResponse] = upload_response.json()
job_id: Final = upload_result.get("jobId")
self._log_api_response(
@ -409,7 +446,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
params={"jobId": job_id},
)
poll_response.raise_for_status()
result = poll_response.json()
result: _SanitizeStatusResponse = poll_response.json()
self._log_api_response(
url=f"{self.api_base}/api/sanitizeFile",
@ -656,7 +693,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
method: str,
url: str,
headers: dict,
payload: Any,
payload: object,
) -> None:
verbose_proxy_logger.debug(
"Prompt Security request %s %s headers=%s payload=%s",
@ -670,7 +707,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
self,
url: str,
status_code: int,
payload: Any,
payload: object,
) -> None:
verbose_proxy_logger.debug(
"Prompt Security response %s status=%s payload=%s",

40
litellm/proxy/guardrails/guardrail_hooks/tool_permission.py
View file

@ -1,6 +1,6 @@
import json
import re
from collections.abc import AsyncGenerator, Sequence
from collections.abc import AsyncGenerator, AsyncIterable, Mapping, Sequence
from typing import Any, Final, Literal
from fastapi import HTTPException
@ -41,6 +41,16 @@ from litellm.types.utils import (
GUARDRAIL_NAME: Final = "tool_permission"
def _object_mapping(value: object) -> Mapping[str, object] | None:
"""Return ``value`` as an opaque mapping when it is a dict."""
return value if isinstance(value, dict) else None
def _object_list(value: object) -> Sequence[object] | None:
"""Return ``value`` as an opaque sequence when it is a list."""
return value if isinstance(value, list) else None
class ToolPermissionGuardrail(CustomGuardrail):
def __init__(
self,
@ -274,12 +284,12 @@ class ToolPermissionGuardrail(CustomGuardrail):
def _parse_tool_call_arguments(
self, tool_call: ChatCompletionMessageToolCall
) -> tuple[dict[str, Any] | None, str | None]:
) -> tuple[Mapping[str, object] | None, str | None]:
arguments: Final = getattr(tool_call.function, "arguments", None)
if not arguments:
return None, "missing arguments"
parsed_arguments: Any = {}
parsed_arguments: object = {}
try:
if isinstance(arguments, str):
parsed_arguments = json.loads(arguments)
@ -306,9 +316,9 @@ class ToolPermissionGuardrail(CustomGuardrail):
def _collect_argument_paths(
self,
value: Any,
value: object,
current_path: str,
collected: dict[str, list[Any]],
collected: dict[str, list[object]],
depth: int = 0,
) -> None:
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
@ -316,13 +326,15 @@ class ToolPermissionGuardrail(CustomGuardrail):
if depth > DEFAULT_MAX_RECURSE_DEPTH:
return
if isinstance(value, dict):
for key, sub_value in value.items():
mapping_value: Final = _object_mapping(value)
list_value: Final = _object_list(value)
if mapping_value is not None:
for key, sub_value in mapping_value.items():
next_path = f"{current_path}.{key}" if current_path else key
self._collect_argument_paths(sub_value, next_path, collected, depth + 1)
elif isinstance(value, list):
elif list_value is not None:
list_path: Final = f"{current_path}[]" if current_path else "[]"
for item in value:
for item in list_value:
self._collect_argument_paths(item, list_path, collected, depth + 1)
else:
if not current_path:
@ -332,7 +344,7 @@ class ToolPermissionGuardrail(CustomGuardrail):
def _patterns_match_for_rule(
self,
*,
arguments: dict[str, Any],
arguments: Mapping[str, object],
rule: ToolPermissionRule,
tool_name: str | None,
) -> tuple[bool, str | None]:
@ -340,7 +352,7 @@ class ToolPermissionGuardrail(CustomGuardrail):
if not compiled_patterns:
return True, None
path_value_map: Final[dict[str, list[Any]]] = {}
path_value_map: Final[dict[str, list[object]]] = {}
self._collect_argument_paths(arguments, "", path_value_map)
for path, compiled_pattern in compiled_patterns.items():
@ -493,14 +505,14 @@ class ToolPermissionGuardrail(CustomGuardrail):
)
@staticmethod
def _get_anthropic_content_blocks(response: object) -> tuple[Any, ...] | None:
def _get_anthropic_content_blocks(response: object) -> tuple[object, ...] | None:
if not isinstance(response, dict):
return None
content: Final[object] = response.get("content")
return tuple(content) if isinstance(content, list) else None
def _extract_tool_calls_from_anthropic_content(
self, content: tuple[Any, ...]
self, content: tuple[object, ...]
) -> tuple[ChatCompletionMessageToolCall, ...]:
return tuple(
tool_call for block in content if (tool_call := self._anthropic_tool_use_to_tool_call(block)) is not None
@ -852,7 +864,7 @@ class ToolPermissionGuardrail(CustomGuardrail):
async def async_post_call_streaming_iterator_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
response: Any,
response: AsyncIterable[ModelResponseStream],
request_data: dict,
) -> AsyncGenerator[ModelResponseStream, None]:
"""

56
litellm/proxy/guardrails/guardrail_hooks/vigil_guard/vigil_guard.py
View file

@ -1,8 +1,9 @@
from collections.abc import Awaitable
from collections.abc import Awaitable, Mapping, Sequence
from json import JSONDecodeError
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, cast
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, TypeAlias, cast
import httpx
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.exceptions import GuardrailRaisedException
@ -50,7 +51,23 @@ _METADATA_ALLOWLIST: Final = (
"org_id",
)
_FallbackMode = Literal["fail_closed", "fail_open"]
_FallbackMode: TypeAlias = Literal["fail_closed", "fail_open"]
_MetadataValue: TypeAlias = str | int | float | Sequence[str | int | float]
class _AnalyzePayload(TypedDict):
"""Request body posted to the Vigil Guard analyze endpoint."""
text: ReadOnly[str]
source: ReadOnly[str]
mode: ReadOnly[str]
metadata: ReadOnly[Mapping[str, _MetadataValue]]
class _AnalysisView(TypedDict):
"""Typed read of the analyze endpoint's decoded JSON body."""
analysis: ReadOnly[Mapping[str, object]]
class _AsyncPostHandler(Protocol):
@ -59,7 +76,7 @@ class _AsyncPostHandler(Protocol):
*,
url: str,
headers: dict[str, str],
json: dict[str, Any],
json: _AnalyzePayload,
timeout: httpx.Timeout,
) -> Awaitable[httpx.Response]: ...
@ -244,7 +261,7 @@ class VigilGuardGuardrail(CustomGuardrail):
exc: Exception,
inputs: GenericGuardrailAPIInputs,
source: str,
final_texts: list[Any],
final_texts: list[str],
final_tool_calls: Any,
) -> GenericGuardrailAPIInputs:
if self.unreachable_fallback == "fail_open":
@ -271,7 +288,7 @@ class VigilGuardGuardrail(CustomGuardrail):
@staticmethod
def _build_output(
inputs: GenericGuardrailAPIInputs,
final_texts: list[Any],
final_texts: list[str],
final_tool_calls: Any,
) -> GenericGuardrailAPIInputs:
# When nothing was changed, return the input shape verbatim so the guardrail
@ -292,7 +309,7 @@ class VigilGuardGuardrail(CustomGuardrail):
return guardrailed
@staticmethod
def _tool_call_arguments(tool_calls: Any) -> list[tuple[int, str]]:
def _tool_call_arguments(tool_calls: Sequence[object] | None) -> list[tuple[int, str]]:
pairs: Final[list[tuple[int, str]]] = []
if isinstance(tool_calls, list):
for index, tool_call in enumerate(tool_calls):
@ -312,8 +329,8 @@ class VigilGuardGuardrail(CustomGuardrail):
updated[index] = tool_call
return updated
async def _analyze(self, text: str, source: str, metadata: dict[str, Any]) -> dict[str, Any]:
payload: Final = {
async def _analyze(self, text: str, source: str, metadata: Mapping[str, _MetadataValue]) -> Mapping[str, object]:
payload: Final[_AnalyzePayload] = {
"text": text,
"source": source,
"mode": "full",
@ -325,9 +342,12 @@ class VigilGuardGuardrail(CustomGuardrail):
"Content-Type": "application/json",
}
response: Final = await self._post_with_retry(endpoint, headers, payload)
return response.json()
decoded: Final[_AnalysisView] = {"analysis": response.json()}
return decoded["analysis"]
async def _post_with_retry(self, endpoint: str, headers: dict[str, str], payload: dict[str, Any]) -> httpx.Response:
async def _post_with_retry(
self, endpoint: str, headers: dict[str, str], payload: _AnalyzePayload
) -> httpx.Response:
for attempt in range(2):
try:
response = await self.async_handler.post(
@ -364,7 +384,7 @@ class VigilGuardGuardrail(CustomGuardrail):
)
@staticmethod
def _build_block_reason(analysis: dict[str, Any]) -> str:
def _build_block_reason(analysis: Mapping[str, object]) -> str:
for key in ("blockMessage", "decisionReason"):
value = analysis.get(key)
if isinstance(value, str) and value.strip():
@ -377,14 +397,16 @@ class VigilGuardGuardrail(CustomGuardrail):
return "Blocked by policy"
@staticmethod
def _resolve_sanitized_text(original: str, analysis: dict[str, Any]) -> str:
def _resolve_sanitized_text(original: str, analysis: Mapping[str, object]) -> str:
for key in ("sanitizedText", "outputText"):
value = analysis.get(key)
if isinstance(value, str):
return value
return original
def _collect_metadata(self, request_data: dict, logging_obj: Optional["LiteLLMLoggingObj"]) -> dict[str, Any]:
def _collect_metadata(
self, request_data: dict, logging_obj: Optional["LiteLLMLoggingObj"]
) -> Mapping[str, _MetadataValue]:
sources: Final[list[dict]] = []
if isinstance(request_data, dict):
sources.append(request_data)
@ -393,7 +415,7 @@ class VigilGuardGuardrail(CustomGuardrail):
if isinstance(nested, dict):
sources.append(nested)
collected: Final[dict[str, Any]] = {}
collected: Final[dict[str, _MetadataValue]] = {}
for field in _METADATA_ALLOWLIST:
for source in sources:
if field in source and source[field] is not None:
@ -409,7 +431,7 @@ class VigilGuardGuardrail(CustomGuardrail):
return collected
@staticmethod
def _clamp_metadata_value(value: Any) -> Any:
def _clamp_metadata_value(value: Any) -> _MetadataValue | None:
if isinstance(value, bool):
return None
if isinstance(value, str):
@ -417,7 +439,7 @@ class VigilGuardGuardrail(CustomGuardrail):
if isinstance(value, (int, float)):
return value
if isinstance(value, list):
clamped: Final[list[Any]] = []
clamped: Final[list[str | int | float]] = []
for item in value[:_METADATA_ARRAY_MAX_ITEMS]:
if isinstance(item, bool):
continue

37
litellm/proxy/guardrails/guardrail_registry.py
View file

@ -3,12 +3,12 @@
import asyncio
import importlib
import os
from collections.abc import Callable, Iterator, Mapping
from collections.abc import Callable, Iterator, Mapping, Sequence
from datetime import datetime, timezone
from itertools import chain, count
from typing import Any, Final, Literal, Optional, Protocol, cast
from typing import Final, Literal, Optional, Protocol, cast
from pydantic import ValidationError
from pydantic import BaseModel, ValidationError
import litellm
from litellm import Router
@ -68,6 +68,19 @@ class _GuardrailRowLike(Protocol):
def __iter__(self) -> Iterator[tuple[str, object]]: ...
class _GuardrailTableActions(Protocol):
async def create(self, *, data: Mapping[str, object]) -> _GuardrailRowLike: ...
async def delete(self, *, where: Mapping[str, str]) -> object: ...
async def update(self, *, where: Mapping[str, str], data: Mapping[str, object]) -> _GuardrailRowLike: ...
async def find_many(self, *, where: Mapping[str, str], order: Mapping[str, str]) -> Sequence[BaseModel]: ...
async def find_unique(self, *, where: Mapping[str, str]) -> BaseModel | None: ...
def _guardrail_table(prisma_client: PrismaClient) -> _GuardrailTableActions:
"""Typed view of the guardrails table actions exposed by the Prisma repository."""
return GuardrailsRepository(prisma_client).table
guardrail_initializer_registry: Final = {
SupportedGuardrailIntegrations.BEDROCK.value: initialize_bedrock,
SupportedGuardrailIntegrations.LAKERA.value: initialize_lakera,
@ -279,7 +292,7 @@ class GuardrailRegistry:
try:
guardrail_name: Final = guardrail.get("guardrail_name")
# Properly serialize LitellmParams Pydantic model to dict
litellm_params_obj: Final[Any] = guardrail.get("litellm_params", {})
litellm_params_obj: Final = guardrail.get("litellm_params", {})
if hasattr(litellm_params_obj, "model_dump"):
litellm_params_dict = litellm_params_obj.model_dump()
else:
@ -288,7 +301,7 @@ class GuardrailRegistry:
guardrail_info: Final[str] = safe_dumps(guardrail.get("guardrail_info", {}))
# Create guardrail in DB
created_guardrail: Final[_GuardrailRowLike] = await GuardrailsRepository(prisma_client).table.create(
created_guardrail: Final[_GuardrailRowLike] = await _guardrail_table(prisma_client).create(
data={
"guardrail_name": guardrail_name,
"litellm_params": litellm_params,
@ -312,7 +325,7 @@ class GuardrailRegistry:
"""
try:
# Delete from DB
await GuardrailsRepository(prisma_client).table.delete(where={"guardrail_id": guardrail_id})
await _guardrail_table(prisma_client).delete(where={"guardrail_id": guardrail_id})
return {"message": f"Guardrail {guardrail_id} deleted successfully"}
except Exception as e:
@ -325,7 +338,7 @@ class GuardrailRegistry:
try:
guardrail_name: Final = guardrail.get("guardrail_name")
# Properly serialize LitellmParams Pydantic model to dict
litellm_params_obj: Final[Any] = guardrail.get("litellm_params", {})
litellm_params_obj: Final = guardrail.get("litellm_params", {})
if hasattr(litellm_params_obj, "model_dump"):
litellm_params_dict = litellm_params_obj.model_dump()
else:
@ -334,7 +347,7 @@ class GuardrailRegistry:
guardrail_info: Final[str] = safe_dumps(guardrail.get("guardrail_info", {}))
# Update in DB
updated_guardrail: Final[_GuardrailRowLike] = await GuardrailsRepository(prisma_client).table.update(
updated_guardrail: Final[_GuardrailRowLike] = await _guardrail_table(prisma_client).update(
where={"guardrail_id": guardrail_id},
data={
"guardrail_name": guardrail_name,
@ -358,7 +371,7 @@ class GuardrailRegistry:
Only rows with status == "active" are returned (pending_review and rejected are excluded).
"""
try:
guardrails_from_db: Final = await GuardrailsRepository(prisma_client).table.find_many(
guardrails_from_db: Final = await _guardrail_table(prisma_client).find_many(
where={"status": "active"},
order={"created_at": "desc"},
)
@ -376,9 +389,7 @@ class GuardrailRegistry:
Get a guardrail by its ID from the database
"""
try:
guardrail: Final = await GuardrailsRepository(prisma_client).table.find_unique(
where={"guardrail_id": guardrail_id}
)
guardrail: Final = await _guardrail_table(prisma_client).find_unique(where={"guardrail_id": guardrail_id})
if not guardrail:
return None
@ -392,7 +403,7 @@ class GuardrailRegistry:
Get a guardrail by its name from the database
"""
try:
guardrail: Final = await GuardrailsRepository(prisma_client).table.find_unique(
guardrail: Final = await _guardrail_table(prisma_client).find_unique(
where={"guardrail_name": guardrail_name}
)

38
litellm/proxy/hooks/litellm_skills/main.py
View file

@ -27,7 +27,7 @@ Usage:
import base64
import json
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final
from typing import TYPE_CHECKING, Any, Final, Protocol
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
@ -43,6 +43,30 @@ if TYPE_CHECKING:
from litellm.llms.litellm_proxy.skills.sandbox_executor import SkillsSandboxExecutor
class _ToolCallFunction(Protocol):
@property
def name(self) -> str: ...
@property
def arguments(self) -> str: ...
class _ChatToolCall(Protocol):
@property
def id(self) -> str: ...
@property
def function(self) -> _ToolCallFunction: ...
class _ChatMessage(Protocol):
@property
def content(self) -> str | None: ...
@property
def tool_calls(self) -> Sequence[_ChatToolCall] | None: ...
class SkillsInjectionHook(CustomLogger):
"""
Pre/Post-call hook that processes skills from container.skills parameter.
@ -443,7 +467,7 @@ class SkillsInjectionHook(CustomLogger):
async def _execute_code_loop_messages_api(
self,
data: dict,
response: Any,
response: object,
skill_files: dict[str, bytes],
) -> LLMResponseTypes | None:
"""
@ -673,7 +697,7 @@ print('No executable skill module found')
async def _execute_code_loop(
self,
data: dict,
response: Any,
response: object,
skill_files: dict[str, bytes],
) -> LLMResponseTypes:
"""
@ -714,8 +738,8 @@ print('No executable skill module found')
for iteration in range(self.max_iterations):
# OpenAI format response has choices[0].message
assistant_message = current_response.choices[0].message
stop_reason = current_response.choices[0].finish_reason
assistant_message: _ChatMessage = current_response.choices[0].message
stop_reason: str | None = current_response.choices[0].finish_reason
# Build assistant message for conversation history
assistant_msg_dict: dict[str, object] = {
@ -784,14 +808,14 @@ print('No executable skill module found')
async def _execute_code_tool(
self,
tool_call: Any,
tool_call: _ChatToolCall,
skill_files: dict[str, bytes],
executor: "SkillsSandboxExecutor",
generated_files: list[dict[str, object]],
) -> str:
"""Execute a litellm_code_execution tool call and return result string."""
try:
args: Final = json.loads(tool_call.function.arguments)
args: Final[Mapping[str, str]] = json.loads(tool_call.function.arguments)
code: Final[str] = args.get("code", "")
verbose_proxy_logger.debug("SkillsInjectionHook: Executing code (%s chars)", len(code))

133
litellm/proxy/management_endpoints/auto_router_endpoints.py
View file

@ -7,7 +7,7 @@ POST /auto_router/test_routing - Route one prompt through an unsaved complexity-
from collections.abc import Mapping, Sequence
from datetime import datetime, timedelta, timezone
from types import MappingProxyType
from typing import TYPE_CHECKING, Annotated, Final
from typing import TYPE_CHECKING, Annotated, Final, Protocol
from pydantic import BaseModel, TypeAdapter
@ -29,6 +29,7 @@ from litellm.proxy.auth.auth_checks import (
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.db.autorouter_session_rollup import AUTOROUTER_BENCHMARKS_SQL
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
from litellm.repositories.base_repository import SupportsModelDump
from litellm.repositories.team_repository import TeamRepository
from litellm.router_strategy.complexity_router import ComplexityRouter
from litellm.types.management_endpoints.auto_router_endpoints import (
@ -61,6 +62,77 @@ else:
router: Final = APIRouter()
class _TeamTable(Protocol):
async def find_unique(self, *, where: Mapping[str, object]) -> SupportsModelDump | None: ...
class _VerificationTokenRow(Protocol):
@property
def token(self) -> str: ...
@property
def key_alias(self) -> str | None: ...
@property
def key_name(self) -> str | None: ...
class _VerificationTokenTable(Protocol):
async def find_unique(self, *, where: Mapping[str, object]) -> _VerificationTokenRow | None: ...
async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_VerificationTokenRow]: ...
class _ShadowEvalJobRow(Protocol):
@property
def id(self) -> str: ...
class _ShadowEvalJobTable(Protocol):
async def find_unique(self, *, where: Mapping[str, object]) -> _ShadowEvalJobRow | None: ...
async def find_first(self, *, where: Mapping[str, object]) -> _ShadowEvalJobRow | None: ...
async def find_many(
self, *, where: Mapping[str, object], order: Mapping[str, str], take: int
) -> Sequence[_ShadowEvalJobRow]: ...
async def create(self, data: Mapping[str, object]) -> _ShadowEvalJobRow: ...
async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> _ShadowEvalJobRow | None: ...
class _ShadowEvalAttemptRow(Protocol):
@property
def error(self) -> str | None: ...
class _ShadowEvalAttemptTable(Protocol):
async def find_first(
self, *, where: Mapping[str, object], order: Mapping[str, str]
) -> _ShadowEvalAttemptRow | None: ...
def _team_table(prisma_client: "PrismaClient") -> _TeamTable:
return TeamRepository(prisma_client).table
def _verification_tokens(prisma_client: "PrismaClient") -> _VerificationTokenTable:
return prisma_client.db.litellm_verificationtoken
def _shadow_eval_jobs(prisma_client: "PrismaClient") -> _ShadowEvalJobTable:
return prisma_client.db.litellm_shadowevaljob
def _shadow_eval_attempts(prisma_client: "PrismaClient") -> _ShadowEvalAttemptTable:
return prisma_client.db.litellm_shadowevalattempt
async def _query_raw(prisma_client: "PrismaClient", query: str, *args: object) -> Sequence[Mapping[str, object]]:
return await prisma_client.db.query_raw(query, *args)
async def _authorize_routing_test(user_api_key_dict: UserAPIKeyAuth, team_id: str | None) -> None:
"""Allow exactly the callers who could create this router.
@ -92,7 +164,7 @@ async def _authorize_routing_test(user_api_key_dict: UserAPIKeyAuth, team_id: st
},
)
team_row: Final = await TeamRepository(prisma_client).table.find_unique(
team_row: Final = await _team_table(prisma_client).find_unique(
where={"team_id": team_id}, # mutable-ok: Prisma query filters are dict-shaped
)
if team_row is None:
@ -342,6 +414,26 @@ def _benchmark_totals(row: _SessionAggRow) -> AutoRouterBenchmarkTotals:
)
def _benchmark_group(row: _SessionAggRow) -> AutoRouterBenchmarkGroup:
totals: Final = _benchmark_totals(row)
return AutoRouterBenchmarkGroup(
router_name=row.router_name,
router_type=row.router_type,
tier_turns=row.tier_turns,
sessions=totals.sessions,
turns=totals.turns,
avg_turns_per_session=totals.avg_turns_per_session,
avg_session_seconds=totals.avg_session_seconds,
avg_tokens_per_session=totals.avg_tokens_per_session,
spend=totals.spend,
saved_spend=totals.saved_spend,
baseline_spend=totals.baseline_spend,
saved_pct=totals.saved_pct,
saved_per_session=totals.saved_per_session,
cache=totals.cache,
)
def _summed_agg_row(rows: Sequence[_SessionAggRow]) -> _SessionAggRow:
return _SessionAggRow(
router_name="",
@ -407,21 +499,14 @@ async def get_auto_router_benchmarks(
if end_day < start_day:
raise HTTPException(status_code=400, detail="end_date must not be earlier than start_date")
raw_rows: Final = await prisma_client.db.query_raw(
raw_rows: Final = await _query_raw(
prisma_client,
AUTOROUTER_BENCHMARKS_SQL,
start_day.isoformat(),
(end_day + timedelta(days=1)).isoformat(),
)
rows: Final = _SESSION_AGG_ROWS.validate_python(raw_rows or ())
groups: Final = tuple(
AutoRouterBenchmarkGroup(
router_name=row.router_name,
router_type=row.router_type,
tier_turns=row.tier_turns,
**_benchmark_totals(row).model_dump(),
)
for row in rows
)
groups: Final = tuple(_benchmark_group(row) for row in rows)
return AutoRouterBenchmarksResponse(
start_date=start_day.strftime("%Y-%m-%d"),
end_date=end_day.strftime("%Y-%m-%d"),
@ -584,7 +669,7 @@ async def _with_key_labels(
so the UI can say whose traffic a job shadows. Deleted keys resolve to None."""
if not responses:
return ()
key_rows: Final = await prisma_client.db.litellm_verificationtoken.find_many(
key_rows: Final = await _verification_tokens(prisma_client).find_many(
where={"token": {"in": sorted({response.api_key_id for response in responses})}} # mutable-ok: Prisma filter
)
labels: Final[Mapping[str, tuple[str | None, str | None]]] = {
@ -608,12 +693,12 @@ async def _shadow_eval_results(prisma_client: "PrismaClient", job_id: str) -> Sh
"for the turns the router sent to X, did X beat the baseline" in reverse. Reads are
bounded by the job's own attempts (<= max_turns) via the job_id index."""
by_tier: Final = _ATTEMPT_AGG_ROWS.validate_python(
await prisma_client.db.query_raw(_ATTEMPT_AGG_BY_TIER_SQL, job_id) or ()
await _query_raw(prisma_client, _ATTEMPT_AGG_BY_TIER_SQL, job_id) or ()
)
if not by_tier:
return None
by_model: Final = _ATTEMPT_AGG_ROWS.validate_python(
await prisma_client.db.query_raw(_ATTEMPT_AGG_BY_MODEL_SQL, job_id) or ()
await _query_raw(prisma_client, _ATTEMPT_AGG_BY_MODEL_SQL, job_id) or ()
)
total_turns: Final = sum(r.turn_count for r in by_tier)
return ShadowEvalResult(
@ -661,7 +746,7 @@ async def start_shadow_eval(
_validate_plain_model(llm_router, data.judge_model, "judge_model")
if data.baseline_model is not None:
_validate_plain_model(llm_router, data.baseline_model, "baseline_model")
key_row: Final = await prisma_client.db.litellm_verificationtoken.find_unique(
key_row: Final = await _verification_tokens(prisma_client).find_unique(
where={"token": data.api_key_id} # mutable-ok: Prisma filter
)
if key_row is None:
@ -677,7 +762,7 @@ async def start_shadow_eval(
# still holds its slot in the per-key, per-direction partial unique index until
# stamped; free it so a new eval can start. Sweeping both directions is deliberate.
await prisma_client.db.execute_raw(_SWEEP_FINISHED_JOBS_SQL, data.api_key_id)
active: Final = await prisma_client.db.litellm_shadowevaljob.find_first(
active: Final = await _shadow_eval_jobs(prisma_client).find_first(
where={ # mutable-ok: Prisma filter
"api_key_id": data.api_key_id,
"direction": data.direction,
@ -691,7 +776,7 @@ async def start_shadow_eval(
)
now: Final = datetime.now(timezone.utc)
try:
job: Final = await prisma_client.db.litellm_shadowevaljob.create(
job: Final = await _shadow_eval_jobs(prisma_client).create(
data={ # mutable-ok: Prisma payload
"api_key_id": data.api_key_id,
"router_name": data.router_name,
@ -735,7 +820,7 @@ async def list_shadow_eval_jobs(
_require_admin_viewer(user_api_key_dict, "view shadow evals")
if prisma_client is None:
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
records: Final = await prisma_client.db.litellm_shadowevaljob.find_many(
records: Final = await _shadow_eval_jobs(prisma_client).find_many(
where={"api_key_id": api_key_id} if api_key_id else {}, # mutable-ok: Prisma filter
order={"created_at": "desc"}, # mutable-ok: Prisma order
take=limit,
@ -762,15 +847,15 @@ async def get_shadow_eval_job(
_require_admin_viewer(user_api_key_dict, "view shadow evals")
if prisma_client is None:
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
record: Final = await prisma_client.db.litellm_shadowevaljob.find_unique(
record: Final = await _shadow_eval_jobs(prisma_client).find_unique(
where={"id": job_id} # mutable-ok: Prisma filter
)
if record is None:
raise HTTPException(status_code=404, detail=f"No shadow eval job {job_id}")
totals: Final = _ATTEMPT_TOTALS_ROWS.validate_python(
await prisma_client.db.query_raw(_ATTEMPT_TOTALS_SQL, job_id) or ()
await _query_raw(prisma_client, _ATTEMPT_TOTALS_SQL, job_id) or ()
)
latest_error: Final = await prisma_client.db.litellm_shadowevalattempt.find_first(
latest_error: Final = await _shadow_eval_attempts(prisma_client).find_first(
where={"job_id": job_id, "outcome": "error"}, # mutable-ok: Prisma filter
order={"created_at": "desc"}, # mutable-ok: Prisma order
)
@ -804,7 +889,7 @@ async def stop_shadow_eval_job(
_require_admin_writer(user_api_key_dict, "stop a shadow eval")
if prisma_client is None:
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
record: Final = await prisma_client.db.litellm_shadowevaljob.find_unique(
record: Final = await _shadow_eval_jobs(prisma_client).find_unique(
where={"id": job_id} # mutable-ok: Prisma filter
)
if record is None:
@ -812,7 +897,7 @@ async def stop_shadow_eval_job(
current: Final = ShadowEvalJobResponse.model_validate(record, from_attributes=True)
if current.status != "running":
raise HTTPException(status_code=400, detail=f"Job {job_id} is already {current.status}")
updated: Final = await prisma_client.db.litellm_shadowevaljob.update(
updated: Final = await _shadow_eval_jobs(prisma_client).update(
where={"id": job_id}, # mutable-ok: Prisma filter
data={"stopped_at": datetime.now(timezone.utc)}, # mutable-ok: Prisma payload
)

40
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -194,6 +194,13 @@ class _PrismaTableActions(Protocol[_PrismaRowT]):
data: Mapping[str, object],
) -> _PrismaRowT | None: ...
async def upsert(
self,
*,
where: Mapping[str, object],
data: Mapping[str, object],
) -> _PrismaRowT: ...
class _UserRowLike(Protocol):
user_id: str | None
@ -209,24 +216,43 @@ class _TxTables(Protocol):
litellm_proxymodeltable: _PrismaTableActions[object]
class _TableSource(Protocol[_PrismaRowT]):
"""Repository view that exposes its untyped Prisma ``table`` with a concrete row type."""
@property
def table(self) -> _PrismaTableActions[_PrismaRowT]: ...
def _table_of(source: _TableSource[_PrismaRowT]) -> _PrismaTableActions[_PrismaRowT]:
return source.table
def _prisma_table(
repository: BaseRepository[_RepositoryModelT],
) -> _PrismaTableActions[_RepositoryModelT]:
return repository.table
return _table_of(repository)
def _deleted_verification_token_table(
prisma_client: PrismaClient,
) -> _PrismaTableActions[LiteLLM_DeletedVerificationToken]:
return DeletedVerificationTokenRepository(prisma_client).table
return _table_of(DeletedVerificationTokenRepository(prisma_client))
def _deprecated_verification_token_table(prisma_client: PrismaClient) -> _PrismaTableActions[object]:
return _table_of(DeprecatedVerificationTokenRepository(prisma_client))
def _user_table(prisma_client: PrismaClient) -> _PrismaTableActions[_UserRowLike]:
return _table_of(UserRepository(prisma_client))
def _credentials_table(prisma_client: PrismaClient) -> _PrismaTableActions[CredentialItem]:
return CredentialsRepository(prisma_client).table
return _table_of(CredentialsRepository(prisma_client))
def _config_table(prisma_client: PrismaClient) -> _PrismaTableActions[ConfigParam]:
return ConfigRepository(prisma_client).table
return _table_of(ConfigRepository(prisma_client))
async def _check_custom_key_allowed(custom_key_value: str | None) -> None:
@ -4656,7 +4682,7 @@ async def _insert_deprecated_key(
try:
revoke_at: Final = datetime.now(timezone.utc) + timedelta(seconds=grace_seconds)
await DeprecatedVerificationTokenRepository(prisma_client).table.upsert(
await _deprecated_verification_token_table(prisma_client).upsert(
where={"token": old_token_hash},
data={
"create": {
@ -6059,13 +6085,13 @@ async def _list_key_helper(
total_pages: Final = -(-total_count // size) # Ceiling division
# Fetch user information if expand includes "user"
user_map = {}
user_map = dict[str | None, _UserRowLike]()
if expand and "user" in expand:
user_ids: Final = [key.user_id for key in keys if key.user_id]
created_by_ids: Final = [key.created_by for key in keys if key.created_by]
all_ids: Final = list(set(user_ids + created_by_ids)) # Remove duplicates
if all_ids:
users: Final[Sequence[_UserRowLike]] = await UserRepository(prisma_client).table.find_many(
users: Final[Sequence[_UserRowLike]] = await _user_table(prisma_client).find_many(
where={"user_id": {"in": all_ids}}
)
user_map = {user.user_id: user for user in users}

12
litellm/proxy/management_endpoints/model_management_endpoints.py
View file

@ -114,13 +114,14 @@ class UpdatePublicModelGroupsRequest(BaseModel):
class _ProxyModelRow(Protocol):
model_id: str
model_name: str
litellm_params: Mapping[str, object]
model_info: Mapping[str, object] | None
def model_dump_json(self, *, exclude_none: bool = False) -> str: ...
class _ProxyModelTable(Protocol):
def find_unique(self, *, where: Mapping[str, object]) -> Awaitable[_ProxyModelRow | None]: ...
def find_unique(self, *, where: Mapping[str, object]) -> Awaitable[BaseModel | None]: ...
def find_many(self, *, where: Mapping[str, object]) -> Awaitable[Sequence[_ProxyModelRow]]: ...
@ -182,10 +183,7 @@ def _model_alias_table(prisma_client: PrismaClient) -> _ModelAliasTable:
async def get_db_model(model_id: str, prisma_client: PrismaClient) -> Deployment | None:
db_model: Final = cast(
BaseModel | None,
await _proxy_model_table(prisma_client).find_unique(where={"model_id": model_id}),
)
db_model: Final = await _proxy_model_table(prisma_client).find_unique(where={"model_id": model_id})
if not db_model:
return None
@ -1577,7 +1575,7 @@ async def delete_model(
},
)
model_in_db: Final = await ModelRepository(prisma_client).table.find_unique(where={"model_id": model_info.id})
model_in_db: Final = await _proxy_model_table(prisma_client).find_unique(where={"model_id": model_info.id})
if model_in_db is None:
raise HTTPException(
status_code=400,
@ -1914,7 +1912,7 @@ async def update_model(
)
_model_id: str | None = None
_model_info: Final = getattr(model_params, "model_info", None)
_model_info: Final[ModelInfo | None] = getattr(model_params, "model_info", None)
if _model_info is None:
raise Exception("model_info not provided")

126
litellm/proxy/management_helpers/utils.py
View file

@ -1,9 +1,9 @@
# What is this?
## Helper utils for the management endpoints (keys/users/teams)
from collections.abc import Callable
from collections.abc import Callable, Mapping, MutableMapping, Sequence
from datetime import datetime
from functools import wraps
from typing import Any, Final
from typing import Any, Final, Protocol
from fastapi import HTTPException, Request
from pydantic import BaseModel
@ -23,6 +23,7 @@ from litellm.proxy._types import ( # key request types; user request types; tea
LiteLLM_UserTable,
ManagementEndpointLoggingPayload,
Member,
Span,
SSOUserDefinedValues,
UpdateCustomerRequest,
UpdateKeyRequest,
@ -39,7 +40,53 @@ from litellm.repositories.table_repositories import TeamMembershipRepository
from litellm.repositories.user_repository import UserRepository
def get_new_internal_user_defaults(user_id: str, user_email: str | None = None) -> dict:
class _PrismaRecord(Protocol):
"""Row surface the management helpers read back from Prisma."""
def model_dump(self) -> Mapping[str, object]: ...
class _PrismaUserRecord(Protocol):
"""User row surface the management helpers read back from Prisma."""
user_id: str
def model_dump(self) -> Mapping[str, object]: ...
class _PrismaBudgetRecord(Protocol):
"""Budget row surface the management helpers read back from Prisma."""
budget_id: str
def model_dump(self) -> Mapping[str, object]: ...
class _PrismaBudgetTable(Protocol):
"""Budget table actions the management helpers issue."""
async def create(self, *, data: Mapping[str, object]) -> _PrismaBudgetRecord: ...
async def find_unique(self, *, where: Mapping[str, object]) -> _PrismaBudgetRecord | None: ...
class _PrismaUserTable(Protocol):
"""User table actions the management helpers issue."""
async def update_many(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> int: ...
async def upsert(
self, *, where: Mapping[str, object], data: Mapping[str, Mapping[str, object]]
) -> _PrismaUserRecord | None: ...
class _PrismaTeamMembershipTable(Protocol):
"""Team membership table actions the management helpers issue."""
async def create(self, *, data: Mapping[str, object], include: Mapping[str, bool]) -> _PrismaRecord: ...
def get_new_internal_user_defaults(user_id: str, user_email: str | None = None) -> dict[str, object]:
user_info: Final = litellm.default_internal_user_params or {}
returned_dict: Final[SSOUserDefinedValues] = {
@ -95,7 +142,7 @@ async def handle_budget_for_entity(
_budget_data: Final = {k: v for k, v in _json_data.items() if k in budget_params}
# Check if budget_id is explicitly provided in the data
data_budget_id: Final = getattr(data, "budget_id", None)
data_budget_id: Final[str | None] = getattr(data, "budget_id", None)
# Case 1: Creating new entity - no existing budget_id
if existing_budget_id is None:
@ -107,7 +154,7 @@ async def handle_budget_for_entity(
budget_row: Final = LiteLLM_BudgetTable(**_budget_data)
new_budget_data: Final = prisma_client.jsonify_object(budget_row.model_dump(exclude_none=True))
_budget: Final = await BudgetRepository(prisma_client).table.create(
_budget: Final[_PrismaBudgetRecord] = await BudgetRepository(prisma_client).table.create(
data={
**new_budget_data,
"created_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
@ -173,9 +220,8 @@ async def _clone_team_default_budget_for_member(
member while keeping the default's other limits, so an admin can set a
member's reset cadence without discarding the team default's max_budget.
"""
default_budget: Final = await BudgetRepository(prisma_client).table.find_unique(
where={"budget_id": default_team_budget_id}
)
budget_table: Final[_PrismaBudgetTable] = BudgetRepository(prisma_client).table
default_budget: Final = await budget_table.find_unique(where={"budget_id": default_team_budget_id})
if default_budget is None:
return None
@ -202,7 +248,7 @@ async def _clone_team_default_budget_for_member(
if cloned_data.get("budget_duration"):
cloned_data["budget_reset_at"] = get_budget_reset_time(cloned_data["budget_duration"])
new_budget: Final = await BudgetRepository(prisma_client).table.create(data=cloned_data)
new_budget: Final[_PrismaBudgetRecord] = await BudgetRepository(prisma_client).table.create(data=cloned_data)
return new_budget.budget_id
@ -238,7 +284,7 @@ async def _resolve_member_budget_id(
if not has_explicit_limit and budget_duration is None:
return None
budget_data: Final[dict] = {
budget_data: Final[dict[str, object]] = {
"created_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
"updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
}
@ -249,7 +295,8 @@ async def _resolve_member_budget_id(
if budget_duration is not None:
budget_data["budget_duration"] = budget_duration
budget_data["budget_reset_at"] = get_budget_reset_time(budget_duration=budget_duration)
response: Final = await BudgetRepository(prisma_client).table.create(data=budget_data)
budget_table: Final[_PrismaBudgetTable] = BudgetRepository(prisma_client).table
response: Final = await budget_table.create(data=budget_data)
return response.budget_id
@ -262,7 +309,8 @@ async def _append_team_id_if_absent(prisma_client: PrismaClient, user_id: str, t
number of teams a user belongs to). Teams added concurrently for a different
team id are unaffected, since each update filters on its own team id.
"""
await UserRepository(prisma_client).table.update_many(
user_table: Final[_PrismaUserTable] = UserRepository(prisma_client).table
await user_table.update_many(
where={"user_id": user_id, "NOT": {"teams": {"has": team_id}}},
data={"teams": {"push": [team_id]}},
)
@ -300,7 +348,8 @@ async def add_new_member(
# Prisma only compiles an upsert down to INSERT ... ON CONFLICT when it
# is non-empty, and falls back to a racy SELECT-then-INSERT when it is
# not, so this re-states user_id as a no-op rather than being empty.
_returned_user = await UserRepository(prisma_client).table.upsert(
user_table: Final[_PrismaUserTable] = UserRepository(prisma_client).table
_returned_user: _PrismaUserRecord | None = await user_table.upsert(
where={"user_id": new_member.user_id},
data={
"create": {"teams": [team_id], **new_user_defaults},
@ -314,7 +363,7 @@ async def add_new_member(
new_user_defaults = get_new_internal_user_defaults(user_id=str(uuid.uuid4()), user_email=new_member.user_email)
## user email is not unique acc. to prisma schema -> future improvement
### for now: check if it exists in db, if not - insert it
existing_user_row: Final[list | None] = await prisma_client.get_data(
existing_user_row: Final[list[_PrismaUserRecord] | None] = await prisma_client.get_data(
key_val={"user_email": new_member.user_email},
table_name="user",
query_type="find_all",
@ -346,7 +395,8 @@ async def add_new_member(
)
if _budget_id and returned_user is not None and returned_user.user_id is not None:
_returned_team_membership: Final = await TeamMembershipRepository(prisma_client).table.create(
membership_table: Final[_PrismaTeamMembershipTable] = TeamMembershipRepository(prisma_client).table
_returned_team_membership: Final = await membership_table.create(
data={
"team_id": team_id,
"user_id": returned_user.user_id,
@ -469,8 +519,18 @@ async def send_management_endpoint_alert(
)
def _redacted_env_var(entry: Any) -> dict:
get: Final = entry.get if isinstance(entry, dict) else lambda k: getattr(entry, k, None)
def _object_mapping(value: object) -> Mapping[str, object] | None:
"""Return ``value`` as an opaque mapping when it is a dict."""
return value if isinstance(value, dict) else None
def _object_list(value: object) -> Sequence[object] | None:
"""Return ``value`` as an opaque sequence when it is a list."""
return value if isinstance(value, list) else None
def _redacted_env_var(entry: object) -> dict[str, object]:
get: Final[Callable[[str], object]] = entry.get if isinstance(entry, dict) else lambda k: getattr(entry, k, None)
return {
"name": get("name"),
"scope": get("scope"),
@ -479,25 +539,28 @@ def _redacted_env_var(entry: Any) -> dict:
}
def _redact_record_env_vars(record: Any) -> Any:
def _redact_record_env_vars(record: object) -> object:
"""Return ``record`` with its ``env_vars[].value`` blanked.
Copies rather than mutating, because the record aliases the live response
object that is also returned to the caller. Records without an ``env_vars``
list are returned unchanged.
"""
env_vars: Final = record.get("env_vars") if isinstance(record, dict) else getattr(record, "env_vars", None)
if not isinstance(env_vars, list):
record_map: Final = _object_mapping(record)
env_vars: Final = _object_list(
record_map.get("env_vars") if record_map is not None else getattr(record, "env_vars", None)
)
if env_vars is None:
return record
redacted: Final = [_redacted_env_var(entry) for entry in env_vars]
if isinstance(record, dict):
return {**record, "env_vars": redacted}
if record_map is not None:
return {**record_map, "env_vars": redacted}
if isinstance(record, BaseModel):
return record.model_copy(update={"env_vars": redacted})
return record
def _redact_env_var_values(response: dict) -> None:
def _redact_env_var_values(response: MutableMapping[str, object]) -> None:
"""Blank ``env_vars[].value`` in a management response before telemetry.
MCP endpoints return decrypted ``scope="global"`` env var values so the admin
@ -507,18 +570,19 @@ def _redact_env_var_values(response: dict) -> None:
create/update) and nested under ``items`` (the submissions queue), so both are
scrubbed. Names, scopes, and descriptions are kept so traces stay useful.
"""
if isinstance(response.get("env_vars"), list):
response["env_vars"] = [_redacted_env_var(entry) for entry in response["env_vars"]]
env_vars: Final = _object_list(response.get("env_vars"))
if env_vars is not None:
response["env_vars"] = [_redacted_env_var(entry) for entry in env_vars]
items: Final = response.get("items")
if isinstance(items, list):
items: Final = _object_list(response.get("items"))
if items is not None:
response["items"] = [_redact_record_env_vars(item) for item in items]
async def _emit_management_endpoint_otel_span(
func: Callable,
kwargs: dict,
parent_otel_span: Any,
parent_otel_span: Span | None,
start_time: datetime,
end_time: datetime,
result: Any = None,
@ -571,10 +635,10 @@ async def _emit_management_endpoint_otel_span(
}
)
_response: dict | None = None
_response: dict[str, object] | None = None
if exception is None and result is not None:
try:
raw: Final = dict(result)
raw: Final[Mapping[str, object]] = dict(result)
_response = {k: v for k, v in raw.items() if k not in _CREDENTIAL_FIELDS}
_redact_env_var_values(_response)
except Exception:
@ -623,7 +687,7 @@ def management_endpoint_wrapper(func):
user_api_key_dict=user_api_key_dict,
function_name=func.__name__,
)
parent_otel_span = getattr(user_api_key_dict, "parent_otel_span", None)
parent_otel_span: Span | None = getattr(user_api_key_dict, "parent_otel_span", None)
if parent_otel_span is not None:
await _emit_management_endpoint_otel_span(
func=func,

6
litellm/proxy/proxy_server.py
View file

@ -12992,9 +12992,9 @@ async def _filter_models_by_team_id(
async def _find_model_by_id(
model_id: str,
search: str | None,
llm_router,
prisma_client,
proxy_config,
llm_router: Router | None,
prisma_client: PrismaClient | None,
proxy_config: "ProxyConfig",
) -> tuple[list, int | None]:
"""Find a model by its ID and optionally filter by search term."""
found_model = None

32
litellm/proxy/spend_tracking/budget_reservation.py
View file

@ -105,8 +105,8 @@ def _key_reservation_should_release_for_throttle(counter_key: str, valid_token:
async def _apply_over_budget_reservation_policy(
counter: _BudgetCounter,
valid_token: UserAPIKeyAuth | None,
entry: dict[str, Any],
applied_entries: list[dict[str, Any]],
entry: dict[str, float | str],
applied_entries: list[dict[str, float | str]],
reservation_cost: float,
current_spend: float,
) -> float:
@ -156,7 +156,7 @@ async def reserve_budget_for_request(
user_api_key_cache: DualCache,
proxy_logging_obj: ProxyLogging,
end_user_id: str | None = None,
end_user_object: Any | None = None,
end_user_object: object = None,
apply_user_budget_to_team_keys: bool = False,
fail_closed_budget_enforcement: bool = False,
) -> dict | None:
@ -194,7 +194,7 @@ async def reserve_budget_for_request(
if reservation_cost is None or reservation_cost <= 0:
return None
applied_entries: Final[list[dict[str, Any]]] = []
applied_entries: Final[list[dict[str, float | str]]] = []
try:
for counter in counters:
entry = _counter_to_reservation_entry(
@ -334,7 +334,7 @@ async def _get_budget_counters(
user_api_key_cache: DualCache,
proxy_logging_obj: ProxyLogging,
end_user_id: str | None = None,
end_user_object: Any | None = None,
end_user_object: object = None,
apply_user_budget_to_team_keys: bool = False,
) -> list[_BudgetCounter]:
counters: Final[list[_BudgetCounter]] = []
@ -443,7 +443,7 @@ async def _get_budget_counters(
async def _get_end_user_budget_counter(
valid_token: UserAPIKeyAuth,
end_user_id: str | None,
end_user_object: Any | None,
end_user_object: object,
) -> _BudgetCounter | None:
end_user_id = end_user_id or valid_token.end_user_id
if end_user_id is None:
@ -608,7 +608,7 @@ def _get_budget_limit_counters(
entity_prefix: str,
entity_type: str,
entity_id: str,
budget_limits: Sequence[Any] | None,
budget_limits: Sequence[object] | None,
fallback_spend: float,
) -> list[_BudgetCounter]:
counters: Final[list[_BudgetCounter]] = []
@ -855,7 +855,7 @@ async def _resize_applied_reservation(
def _counter_to_reservation_entry(
counter: _BudgetCounter,
reserved_cost: float,
) -> dict[str, Any]:
) -> dict[str, float | str]:
return {
"counter_key": counter.counter_key,
"entity_type": counter.entity_type,
@ -983,7 +983,7 @@ def _input_cost_for_cost_info(
request_body: dict,
route: str,
model: str,
model_info: dict[str, Any],
model_info: Mapping[str, object],
) -> float | None:
input_tokens: Final = _estimate_input_tokens(
request_body=request_body,
@ -1027,7 +1027,7 @@ def _max_cost_for_cost_info(
request_body: dict,
route: str,
model: str,
model_info: dict[str, Any],
model_info: Mapping[str, object],
) -> float | None:
image_cost: Final = _estimate_image_generation_cost(
request_body=request_body,
@ -1086,7 +1086,7 @@ def _max_cost_for_cost_info(
def _estimate_image_generation_cost(
request_body: dict,
model_info: dict[str, Any],
model_info: Mapping[str, object],
) -> float | None:
"""
Reserve `n × per-image cost` for image-generation requests so concurrent
@ -1125,7 +1125,7 @@ def _estimate_image_generation_cost(
def _get_model_cost_info(
model: str,
llm_router: Router | None,
) -> dict[str, Any] | None:
) -> Mapping[str, object] | None:
if llm_router is not None:
model_group_info: Final = llm_router.get_model_group_info(model_group=model)
if model_group_info is not None:
@ -1136,7 +1136,7 @@ def _get_model_cost_info(
def _get_model_cost_infos(
model: str,
llm_router: Router | None,
) -> list[dict[str, Any]]:
) -> Sequence[Mapping[str, object]]:
"""Cost-info candidates to estimate a request against for one model group.
Reservation runs before routing, so the deployment that will serve the request
@ -1181,7 +1181,7 @@ def _deployment_tiered_pricing_table(
def _get_deployment_tiered_pricing_tables(
model: str,
llm_router: Router | None,
) -> list[list[dict]]:
) -> Sequence[Sequence[Mapping[str, object]]]:
if llm_router is None:
return []
deployments: Final = llm_router.get_model_list(model_name=model) or []
@ -1196,7 +1196,7 @@ def _estimate_input_tokens(
request_body: dict,
route: str,
model: str,
model_info: dict[str, Any],
model_info: Mapping[str, object],
) -> int | None:
try:
if "messages" in request_body:
@ -1233,7 +1233,7 @@ DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK: Final = 16384
def _estimate_output_tokens(
request_body: dict,
route: str,
model_info: dict[str, Any],
model_info: Mapping[str, object],
) -> int | None:
if _is_input_only_route(route=route):
return 0

55
litellm/repositories/config_repository.py
View file

@ -10,12 +10,41 @@ import asyncio
import copy
import json
import os
from typing import Any, Final, Literal, cast
from collections.abc import Mapping, Sequence
from typing import Any, Final, Literal, Protocol, cast
from litellm._logging import verbose_proxy_logger
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper
class _ConfigRow(Protocol):
@property
def param_name(self) -> str: ...
@property
def param_value(self) -> object: ...
class _ConfigTable(Protocol):
async def find_unique(self, *, where: Mapping[str, str]) -> _ConfigRow | None: ...
async def find_many(self) -> Sequence[_ConfigRow]: ...
async def upsert(self, *, where: Mapping[str, str], data: Mapping[str, Mapping[str, str]]) -> _ConfigRow: ...
async def delete(self, *, where: Mapping[str, str]) -> _ConfigRow | None: ...
class _ConfigDb(Protocol):
@property
def litellm_config(self) -> _ConfigTable: ...
class _PrismaHandle(Protocol):
@property
def db(self) -> _ConfigDb: ...
class ConfigParam:
"""Simple wrapper for config parameter from DB."""
@ -38,18 +67,22 @@ class ConfigRepository:
self._prisma_client = prisma_client
@property
def prisma_client(self) -> Any:
def prisma_client(self) -> _PrismaHandle:
if self._prisma_client is None:
raise RuntimeError("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys")
return self._prisma_client
@property
def table(self) -> Any:
def _config_table(self) -> _ConfigTable:
return self.prisma_client.db.litellm_config
@property
def table(self) -> Any:
return self._config_table
async def get_param(self, param_name: str) -> ConfigParam | None:
"""Get a config parameter from the database."""
record: Final = await self.table.find_unique(where={"param_name": param_name})
record: Final = await self._config_table.find_unique(where={"param_name": param_name})
if record is None:
return None
param_value = record.param_value
@ -60,7 +93,7 @@ class ConfigRepository:
async def set_param(self, param_name: str, param_value: Any) -> ConfigParam:
"""Set a config parameter in the database."""
value_json: Final = json.dumps(param_value) if not isinstance(param_value, str) else param_value
await self.table.upsert(
await self._config_table.upsert(
where={"param_name": param_name},
data={
"create": {"param_name": param_name, "param_value": value_json},
@ -72,15 +105,15 @@ class ConfigRepository:
async def delete_param(self, param_name: str) -> bool:
"""Delete a config parameter from the database."""
try:
await self.table.delete(where={"param_name": param_name})
await self._config_table.delete(where={"param_name": param_name})
return True
except Exception:
return False
async def get_all_params(self) -> dict[str, Any]:
async def get_all_params(self) -> dict[str, object]:
"""Get all config parameters from the database."""
records: Final = await self.table.find_many()
result: Final = {}
records: Final = await self._config_table.find_many()
result: Final[dict[str, object]] = {}
for record in records:
param_value = record.param_value
if isinstance(param_value, str):
@ -107,7 +140,9 @@ class ConfigRepository:
else:
d[k] = v
def _decrypt_env_variables(self, env_vars: dict[str, Any], return_original_value: bool = True) -> dict[str, str]:
def _decrypt_env_variables(
self, env_vars: Mapping[str, object], return_original_value: bool = True
) -> dict[str, str]:
"""Decrypt environment variables from database."""
decrypted: Final[dict[str, str]] = {}
for key, value in env_vars.items():

58
litellm/repositories/model_repository.py
View file

@ -3,7 +3,8 @@ Model repository for database operations on LiteLLM_ProxyModelTable.
"""
import json
from typing import Any, Final
from collections.abc import Awaitable, Mapping, Sequence
from typing import Any, Final, Protocol
from litellm.models.model import LiteLLM_ProxyModelTable
from litellm.proxy.common_utils.config_sync_pubsub import wrap_table_actions_for_config_sync
@ -11,28 +12,51 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import (
decrypt_value_helper,
encrypt_value_helper,
)
from litellm.repositories.base_repository import BaseRepository
from litellm.repositories.base_repository import BaseRepository, DbRecord
class _PrismaModelDb(Protocol):
litellm_proxymodeltable: object
class _PrismaClientView(Protocol):
db: _PrismaModelDb
class _ProxyModelActions(Protocol):
"""Prisma table actions used by :class:`ModelRepository`."""
def find_many(self, *, where: Mapping[str, object] | None = None) -> Awaitable[Sequence[DbRecord]]: ...
def create(self, *, data: Mapping[str, object]) -> Awaitable[DbRecord]: ...
def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> Awaitable[DbRecord | None]: ...
class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]):
"""Repository for proxy model database operations with encryption support."""
def __init__(self, prisma_client: Any, encryption_key: str | None = None):
def __init__(self, prisma_client: object, encryption_key: str | None = None):
super().__init__(prisma_client)
self._encryption_key = encryption_key
@property
def table(self) -> Any:
client: Final[_PrismaClientView] = self.prisma_client
return wrap_table_actions_for_config_sync(
actions=self.prisma_client.db.litellm_proxymodeltable,
actions=client.db.litellm_proxymodeltable,
table_name="litellm_proxymodeltable",
)
@property
def _model_table(self) -> _ProxyModelActions:
return self.table
@property
def model_class(self) -> type[LiteLLM_ProxyModelTable]:
return LiteLLM_ProxyModelTable
def _encrypt_litellm_params(self, litellm_params: dict[str, Any]) -> dict[str, Any]:
def _encrypt_litellm_params(self, litellm_params: Mapping[str, object]) -> Mapping[str, object]:
"""Encrypt sensitive values in litellm_params."""
encrypted: Final = {}
for key, value in litellm_params.items():
@ -42,7 +66,7 @@ class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]):
encrypted[key] = value
return encrypted
def _decrypt_litellm_params(self, litellm_params: dict[str, Any]) -> dict[str, Any]:
def _decrypt_litellm_params(self, litellm_params: Mapping[str, object]) -> Mapping[str, object]:
"""Decrypt sensitive values in litellm_params."""
decrypted: Final = {}
for key, value in litellm_params.items():
@ -76,17 +100,17 @@ class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]):
async def find_by_name(self, model_name: str) -> list[LiteLLM_ProxyModelTable]:
"""Find models by name."""
records: Final = await self.table.find_many(where={"model_name": model_name})
records: Final = await self._model_table.find_many(where={"model_name": model_name})
return self._to_model_list(records)
async def find_all(self) -> list[LiteLLM_ProxyModelTable]:
"""Find all models."""
records: Final = await self.table.find_many()
records: Final = await self._model_table.find_many()
return self._to_model_list(records)
async def find_unblocked(self) -> list[LiteLLM_ProxyModelTable]:
"""Find all models that are not blocked."""
records: Final = await self.table.find_many(where={"blocked": False})
records: Final = await self._model_table.find_many(where={"blocked": False})
return self._to_model_list(records)
async def find_by_team_id(self, team_id: str) -> list[LiteLLM_ProxyModelTable]:
@ -102,16 +126,16 @@ class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]):
async def create_model(
self,
model_name: str,
litellm_params: dict[str, Any],
litellm_params: Mapping[str, object],
created_by: str,
model_id: str | None = None,
model_info: dict[str, Any] | None = None,
model_info: Mapping[str, object] | None = None,
blocked: bool = False,
) -> LiteLLM_ProxyModelTable:
"""Create a new model with encryption."""
encrypted_params: Final = self._encrypt_litellm_params(litellm_params)
data: Final[dict[str, Any]] = {
data: Final[dict[str, str | bool]] = {
"model_name": model_name,
"litellm_params": json.dumps(encrypted_params),
"created_by": created_by,
@ -123,7 +147,7 @@ class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]):
if model_info is not None:
data["model_info"] = json.dumps(model_info)
record: Final = await self.table.create(data=data)
record: Final = await self._model_table.create(data=data)
model: Final = self._to_model(record)
assert model is not None
return model
@ -133,12 +157,12 @@ class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]):
model_id: str,
updated_by: str,
model_name: str | None = None,
litellm_params: dict[str, Any] | None = None,
model_info: dict[str, Any] | None = None,
litellm_params: Mapping[str, object] | None = None,
model_info: Mapping[str, object] | None = None,
blocked: bool | None = None,
) -> LiteLLM_ProxyModelTable | None:
"""Update a model with encryption."""
data: Final[dict[str, Any]] = {"updated_by": updated_by}
data: Final[dict[str, str | bool]] = {"updated_by": updated_by}
if model_name is not None:
data["model_name"] = model_name
if litellm_params is not None:
@ -149,7 +173,7 @@ class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]):
if blocked is not None:
data["blocked"] = blocked
record: Final = await self.table.update(where={"model_id": model_id}, data=data)
record: Final = await self._model_table.update(where={"model_id": model_id}, data=data)
return self._to_model(record)
async def delete_model(self, model_id: str) -> LiteLLM_ProxyModelTable | None:

2
litellm/responses/main.py
View file

@ -810,7 +810,7 @@ def _responses_try_dispatch_emulated_file_search(
extra_body: dict[str, object] | None,
timeout: float | httpx.Timeout | None,
custom_llm_provider: str | None,
kwargs: dict[str, Any],
kwargs: dict[str, object],
_is_async: bool,
) -> ResponsesAPIResponse | Coroutine[object, object, ResponsesAPIResponse] | None:
"""Return a response when emulated file_search handles the call; otherwise None."""

21
litellm/responses/streaming_iterator.py
View file

@ -79,6 +79,11 @@ def _is_str_mapping(value: object) -> TypeIs[dict[str, str]]: # guard-ok: verif
return _is_json_object(value) and all(isinstance(item, str) for item in value.values())
def _load_json_object(payload: str | bytes) -> dict[str, object]:
"""Parse a JSON payload that the caller consumes as an object."""
return json.loads(payload)
def _model_id_from_metadata(litellm_metadata: dict[str, object] | None) -> str | None:
model_info: Final = litellm_metadata.get("model_info") if litellm_metadata else None
model_id: Final = model_info.get("id") if _is_json_object(model_info) else None
@ -1474,7 +1479,7 @@ class ResponsesWebSocketStreaming:
event = event.decode("utf-8")
if isinstance(event, str):
try:
event_obj = json.loads(event)
event_obj = _load_json_object(event)
except (json.JSONDecodeError, TypeError):
return
else:
@ -1487,7 +1492,7 @@ class ResponsesWebSocketStreaming:
"""Extract user input content from response.create for logging."""
try:
if isinstance(message, str):
msg_obj = json.loads(message)
msg_obj = _load_json_object(message)
elif _is_json_object(message):
msg_obj = message
else:
@ -1557,7 +1562,7 @@ class ResponsesWebSocketStreaming:
# masked response.completed.
if self.output_guardrail_callbacks:
try:
_evt_payload: Mapping[str, object] = json.loads(response_str)
_evt_payload: Mapping[str, object] = _load_json_object(response_str)
_evt_type = _evt_payload.get("type")
except (json.JSONDecodeError, TypeError):
_evt_type = None
@ -1622,7 +1627,7 @@ class ResponsesWebSocketStreaming:
Non-``response.create`` messages are returned unchanged.
"""
try:
msg_obj: Final[dict[str, object]] = json.loads(message)
msg_obj: Final = _load_json_object(message)
except (json.JSONDecodeError, TypeError):
return message
@ -1751,7 +1756,7 @@ class ResponsesWebSocketStreaming:
return response_str
try:
evt_obj: Final[dict[str, object]] = json.loads(response_str)
evt_obj: Final = _load_json_object(response_str)
except (json.JSONDecodeError, TypeError):
return response_str
@ -1807,7 +1812,7 @@ class ResponsesWebSocketStreaming:
return response_str
try:
evt_obj: Final[Mapping[str, object]] = json.loads(response_str)
evt_obj: Final[Mapping[str, object]] = _load_json_object(response_str)
except (json.JSONDecodeError, TypeError):
return response_str
@ -2141,7 +2146,7 @@ class ManagedResponsesWebSocketHandler:
async def _parse_message(self, raw_message: str) -> dict[str, object] | None:
"""Parse raw WS text; return the message dict or None (JSON error / ignored type)."""
try:
msg_obj: Final[dict[str, object]] = json.loads(raw_message)
msg_obj: Final = _load_json_object(raw_message)
except json.JSONDecodeError:
await self._send_error("Invalid JSON in response.create event", "invalid_request_error")
return None
@ -2345,7 +2350,7 @@ class ManagedResponsesWebSocketHandler:
continue
if chunk_type == "response.completed" and completed_event is None:
try:
completed_event = json.loads(serialized)
completed_event = _load_json_object(serialized)
except Exception:
pass
try:

83
litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py
View file

@ -12,6 +12,7 @@ from __future__ import annotations
import contextlib
import contextvars
from collections.abc import Mapping, MutableMapping
from typing import TYPE_CHECKING, Any, Final
import httpx
@ -26,13 +27,13 @@ from litellm.utils import get_utc_datetime
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
Span = _Span | Any
Span = _Span
else:
Span = Any
RoutingArgsTTL: Final = 60
_io_token_rate_limit_request_kwargs: Final[contextvars.ContextVar[dict[str, Any] | None]] = contextvars.ContextVar(
_io_token_rate_limit_request_kwargs: Final[contextvars.ContextVar[dict[str, object] | None]] = contextvars.ContextVar(
"io_token_rate_limit_request_kwargs",
default=None,
)
@ -43,7 +44,7 @@ ITPM_CACHE_KEY: Final = "_litellm_itpm_cache_key"
OTPM_CACHE_KEY: Final = "_litellm_otpm_cache_key"
def set_io_token_rate_limit_request_kwargs(kwargs: dict[str, Any] | None, store_in_context: bool = True) -> None:
def set_io_token_rate_limit_request_kwargs(kwargs: dict[str, object] | None, store_in_context: bool = True) -> None:
# The reservation sentinels are server-only, but `metadata` is caller
# controlled on proxy requests. Strip any client-supplied copies here (this
# runs before the router stashes its own reservation) so a forged
@ -60,7 +61,7 @@ def set_io_token_rate_limit_request_kwargs(kwargs: dict[str, Any] | None, store_
_io_token_rate_limit_request_kwargs.set(kwargs if store_in_context else None)
def get_io_token_rate_limit_request_kwargs() -> dict[str, Any] | None:
def get_io_token_rate_limit_request_kwargs() -> dict[str, object] | None:
return _io_token_rate_limit_request_kwargs.get()
@ -151,14 +152,14 @@ def _resolve_max_tokens(request_kwargs: dict[str, Any] | None, deployment: dict)
return 4096
def _get_usage_tokens(usage: Any) -> tuple[int, int, int]:
def _get_usage_tokens(usage: object) -> tuple[int, int, int]:
if usage is None:
return 0, 0, 0
if hasattr(usage, "prompt_tokens") or hasattr(usage, "input_tokens"):
prompt = int(getattr(usage, "prompt_tokens", None) or getattr(usage, "input_tokens", 0) or 0)
completion = int(getattr(usage, "completion_tokens", None) or getattr(usage, "output_tokens", 0) or 0)
cached = 0
details = getattr(usage, "prompt_tokens_details", None)
details: object = getattr(usage, "prompt_tokens_details", None)
if details is not None:
cached = int(getattr(details, "cached_tokens", 0) or 0)
if not cached:
@ -175,13 +176,13 @@ def _get_usage_tokens(usage: Any) -> tuple[int, int, int]:
return 0, 0, 0
def _extract_response_usage(response_obj: Any) -> Any:
def _extract_response_usage(response_obj: object) -> object:
if isinstance(response_obj, dict):
return response_obj.get("usage")
return getattr(response_obj, "usage", None)
def _usage_is_present(usage: Any) -> bool:
def _usage_is_present(usage: object) -> bool:
"""
True only if usage carries an actual input/output breakdown.
@ -199,8 +200,8 @@ def _usage_is_present(usage: Any) -> bool:
def _resolve_reconcile_usage_tokens(
kwargs: Any,
response_obj: Any,
kwargs: Mapping[str, object] | None,
response_obj: object,
) -> tuple[int, int, bool]:
"""
Resolve billable input and output tokens for post-call reconcile.
@ -233,7 +234,7 @@ def _resolve_reconcile_usage_tokens(
def _stash_reservation_in_metadata(
request_kwargs: dict[str, Any] | None,
request_kwargs: dict[str, object] | None,
*,
itpm_reserved: int,
otpm_reserved: int,
@ -256,7 +257,7 @@ def _stash_reservation_in_metadata(
request_kwargs[channel] = dict(reservation)
def _extract_reservation(reservation: dict[str, Any]) -> tuple[int, int, str | None, str | None]:
def _extract_reservation(reservation: Mapping[str, int | str | None]) -> tuple[int, int, str | None, str | None]:
itpm_cache_key: Final = reservation.get(ITPM_CACHE_KEY)
otpm_cache_key: Final = reservation.get(OTPM_CACHE_KEY)
return (
@ -267,7 +268,12 @@ def _extract_reservation(reservation: dict[str, Any]) -> tuple[int, int, str | N
)
def _reservation_channels(kwargs: Any) -> tuple[Any, ...]:
def _as_mutable_mapping(value: object) -> MutableMapping[str, object] | None:
"""``value`` when it is a dict, else ``None``."""
return value if isinstance(value, dict) else None
def _reservation_channels(kwargs: Mapping[str, object] | None) -> tuple[object, ...]:
"""
Places a reservation may live, in priority order: the top-level metadata
channels win over litellm_params.metadata (so a top-level stash is never
@ -275,30 +281,29 @@ def _reservation_channels(kwargs: Any) -> tuple[Any, ...]:
"""
if not isinstance(kwargs, dict):
return ()
channels: Final = [kwargs.get("metadata"), kwargs.get("litellm_metadata")]
litellm_params: Final = kwargs.get("litellm_params")
if isinstance(litellm_params, dict):
channels.append(litellm_params.get("metadata"))
standard_logging_object: Final = kwargs.get("standard_logging_object")
if isinstance(standard_logging_object, dict):
channels.append(standard_logging_object.get("metadata"))
return tuple(channels)
top_level: Final = (kwargs.get("metadata"), kwargs.get("litellm_metadata"))
litellm_params: Final = _as_mutable_mapping(kwargs.get("litellm_params"))
from_params: Final = () if litellm_params is None else (litellm_params.get("metadata"),)
standard_logging_object: Final = _as_mutable_mapping(kwargs.get("standard_logging_object"))
from_logging_object: Final = () if standard_logging_object is None else (standard_logging_object.get("metadata"),)
return top_level + from_params + from_logging_object
def _read_reservation_from_kwargs(kwargs: Any) -> tuple[int, int, str | None, str | None]:
def _read_reservation_from_kwargs(kwargs: Mapping[str, object] | None) -> tuple[int, int, str | None, str | None]:
for channel_dict in _reservation_channels(kwargs):
if isinstance(channel_dict, dict) and ITPM_RESERVED_KEY in channel_dict:
return _extract_reservation(channel_dict)
return 0, 0, None, None
def _clear_reservation_from_kwargs(kwargs: Any) -> None:
def _clear_reservation_from_kwargs(kwargs: Mapping[str, object] | None) -> None:
"""
Remove the stashed reservation so a retry on a different (e.g. non-IO)
deployment does not re-process the already-reconciled/refunded reservation.
"""
for channel_dict in _reservation_channels(kwargs):
if isinstance(channel_dict, dict):
for channel in _reservation_channels(kwargs):
channel_dict = _as_mutable_mapping(channel)
if channel_dict is not None:
for key in (ITPM_RESERVED_KEY, OTPM_RESERVED_KEY, ITPM_CACHE_KEY, OTPM_CACHE_KEY):
channel_dict.pop(key, None)
@ -524,11 +529,13 @@ def io_token_reconcile_success(
kwargs: Any,
response_obj: Any,
) -> None:
itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(kwargs)
request_kwargs: Final[Mapping[str, object] | None] = kwargs
response: Final[object] = response_obj
itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(request_kwargs)
if itpm_key is None and otpm_key is None:
return
billable_input, completion_tokens, usage_resolved = _resolve_reconcile_usage_tokens(kwargs, response_obj)
billable_input, completion_tokens, usage_resolved = _resolve_reconcile_usage_tokens(request_kwargs, response)
try:
if usage_resolved:
@ -556,7 +563,7 @@ def io_token_reconcile_success(
otpm_reserved,
)
finally:
_clear_reservation_from_kwargs(kwargs)
_clear_reservation_from_kwargs(request_kwargs)
verbose_router_logger.debug(
"[IO TOKEN LIMIT] reconciled (usage_resolved=%s, itpm_reserved=%s, billable_input=%s, otpm_reserved=%s, output=%s)",
@ -575,11 +582,13 @@ async def async_io_token_reconcile_success(
*,
parent_otel_span: Span | None = None,
) -> None:
itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(kwargs)
request_kwargs: Final[Mapping[str, object] | None] = kwargs
response: Final[object] = response_obj
itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(request_kwargs)
if itpm_key is None and otpm_key is None:
return
billable_input, completion_tokens, usage_resolved = _resolve_reconcile_usage_tokens(kwargs, response_obj)
billable_input, completion_tokens, usage_resolved = _resolve_reconcile_usage_tokens(request_kwargs, response)
# Reconcile against the exact key that held the reservation (which encodes
# the reservation's minute), not a key recomputed at response time. This
@ -615,7 +624,7 @@ async def async_io_token_reconcile_success(
otpm_reserved,
)
finally:
_clear_reservation_from_kwargs(kwargs)
_clear_reservation_from_kwargs(request_kwargs)
verbose_router_logger.debug(
"[IO TOKEN LIMIT] reconciled (usage_resolved=%s, itpm_reserved=%s, billable_input=%s, otpm_reserved=%s, output=%s)",
@ -631,7 +640,8 @@ def io_token_refund_failure(
dual_cache: DualCache,
kwargs: Any,
) -> None:
itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(kwargs)
request_kwargs: Final[Mapping[str, object] | None] = kwargs
itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(request_kwargs)
if itpm_key is None and otpm_key is None:
return
if itpm_key is not None and itpm_reserved > 0:
@ -646,11 +656,11 @@ def io_token_refund_failure(
value=-otpm_reserved,
ttl=RoutingArgsTTL,
)
_clear_reservation_from_kwargs(kwargs)
_clear_reservation_from_kwargs(request_kwargs)
verbose_router_logger.debug("[IO TOKEN LIMIT] refunded ITPM=%s OTPM=%s", itpm_reserved, otpm_reserved)
def refund_stale_reservation_before_retry(dual_cache: DualCache, kwargs: dict[str, Any] | None) -> None:
def refund_stale_reservation_before_retry(dual_cache: DualCache, kwargs: Mapping[str, object] | None) -> None:
"""
Synchronously refund and clear any reservation a previous deployment
attempt stashed in ``kwargs``, before it's overwritten for the next
@ -683,7 +693,8 @@ async def async_io_token_refund_failure(
*,
parent_otel_span: Span | None = None,
) -> None:
itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(kwargs)
request_kwargs: Final[Mapping[str, object] | None] = kwargs
itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(request_kwargs)
if itpm_key is None and otpm_key is None:
return
if itpm_key is not None and itpm_reserved > 0:
@ -700,7 +711,7 @@ async def async_io_token_refund_failure(
ttl=RoutingArgsTTL,
parent_otel_span=parent_otel_span,
)
_clear_reservation_from_kwargs(kwargs)
_clear_reservation_from_kwargs(request_kwargs)
verbose_router_logger.debug("[IO TOKEN LIMIT] refunded ITPM=%s OTPM=%s", itpm_reserved, otpm_reserved)

10
ruff-strict-budget.json
View file

@ -1,6 +1,6 @@
{
"ANN001": {
"limit": 3026
"limit": 3020
},
"ANN002": {
"limit": 71
@ -12,19 +12,19 @@
"limit": 2017
},
"ANN202": {
"limit": 854
"limit": 852
},
"ANN204": {
"limit": 711
},
"ANN205": {
"limit": 114
"limit": 113
},
"ANN206": {
"limit": 133
},
"ANN401": {
"limit": 1290
"limit": 1188
},
"ASYNC230": {
"limit": 11
@ -234,7 +234,7 @@
"limit": 5
},
"TID251": {
"limit": 1216
"limit": 1212
},
"TRY002": {
"limit": 524

10
type-discipline-budget.json
View file

@ -1,9 +1,9 @@
{
"LIT001": {
"limit": 22892
"limit": 22809
},
"LIT002": {
"limit": 26886
"limit": 26878
},
"LIT003": {
"limit": 269
@ -15,7 +15,7 @@
"limit": 0
},
"LIT006": {
"limit": 1071
"limit": 1069
},
"LIT007": {
"limit": 0
@ -27,10 +27,10 @@
"limit": 0
},
"LIT010": {
"limit": 16699
"limit": 16695
},
"LIT011": {
"limit": 5590
"limit": 5588
},
"LIT012": {
"limit": 4519