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

refactor(lint): apply every safe ruff autofix and zero 28 strict-rule budgets

Browse source 
About 35,000 fixes ruff marks safe across 32 rules (UP006/UP045/UP007
modern annotations, UP032 f-strings, SIM114/SIM118, RET501, and
friends), removal of the 1,296 typing imports the rewrite orphaned, and
hand fixes for what the fixers could not see: five star-import
freeloaders of typing names, two F823 late-import annotations, the
/get/config/list introspection crash on types.UnionType, redundant
function-local RoleMappings imports in ui_sso.py that shadowed the
module-level name once the annotation lost its quotes, and one FURB168
tautology.

B009/B010/PIE804/RUF019 are excluded on purpose: their safe fixes
rewrite getattr/setattr/**-splat/key-in-dict escape hatches into forms
basedpyright then rejects (283 new errors measured), so their budgets
stay at base values.

ruff-strict-budget.json drops by 39,579 this commit (39,968 across the
branch) with 28 rules at an actual 0 and 9 more sharply down.
type-discipline-budget.json ratchets LIT002/LIT006/LIT009 down; LIT001
moves to the now-honest total: the checker matches the spelling `set`
but not the alias `Set`, so the 160 typing.Set annotations rewritten to
set[...] were always mutable-set annotations and only now count.
This commit is contained in:
mateo-berri 2026-08-01 15:43:29 -07:00
parent 397e8e4918
commit b604e2b20c
1457 changed files with 29765 additions and 32318 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
litellm/_lazy_imports.py
View file

@ -18,7 +18,7 @@ until they're actually needed.
import importlib
import sys
from collections.abc import Callable
from typing import Any, Optional, cast
from typing import Any, cast
# Import all the data structures that define what can be lazy-loaded
# These are just lists of names and maps of where to find them
@ -78,7 +78,7 @@ def _get_utils_globals() -> dict:
# They're separate from the main lazy import system because they have specific use cases
# Lazy loader for default encoding - avoids importing heavy tiktoken library at startup
_default_encoding: Optional[Any] = None
_default_encoding: Any | None = None
def _get_default_encoding() -> Any:
@ -100,7 +100,7 @@ def _get_default_encoding() -> Any:
# Lazy loader for get_modified_max_tokens to avoid importing token_counter at module import time
_get_modified_max_tokens_func: Optional[Any] = None
_get_modified_max_tokens_func: Any | None = None
def _get_modified_max_tokens() -> Any:
@ -124,7 +124,7 @@ def _get_modified_max_tokens() -> Any:
# Lazy loader for token_counter to avoid importing token_counter module at module import time
_token_counter_new_func: Optional[Any] = None
_token_counter_new_func: Any | None = None
def _get_token_counter_new() -> Any:
@ -154,7 +154,7 @@ def _get_token_counter_new() -> Any:
# This registry maps attribute names (like "ModelResponse") to handler functions
# It's built once the first time someone accesses a lazy-loaded attribute
# Example: {"ModelResponse": _lazy_import_utils, "Cache": _lazy_import_caching, ...}
_LAZY_IMPORT_REGISTRY: Optional[dict[str, Callable[[str], Any]]] = None
_LAZY_IMPORT_REGISTRY: dict[str, Callable[[str], Any]] | None = None
def _get_lazy_import_registry() -> dict[str, Callable[[str], Any]]:

12
litellm/_logging.py
View file

@ -4,11 +4,11 @@ import os
import sys
from datetime import datetime
from logging import Formatter
from typing import Any, Dict, Optional
from typing import Any
from litellm.litellm_core_utils.secret_redaction import redact_string
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.litellm_core_utils.secret_redaction import redact_string
set_verbose = False
@ -86,7 +86,7 @@ handler.setLevel(numeric_level)
handler.addFilter(_secret_filter)
def _try_parse_json_message(message: str) -> Optional[Dict[str, Any]]:
def _try_parse_json_message(message: str) -> dict[str, Any] | None:
"""
Try to parse a log message as JSON. Returns parsed dict if valid, else None.
Handles messages that are entirely valid JSON (e.g. json.dumps output).
@ -103,7 +103,7 @@ def _try_parse_json_message(message: str) -> Optional[Dict[str, Any]]:
return parsed
def _try_parse_embedded_python_dict(message: str) -> Optional[Dict[str, Any]]:
def _try_parse_embedded_python_dict(message: str) -> dict[str, Any] | None:
"""
Try to find and parse a Python dict repr (e.g. str(d) or repr(d)) embedded in
the message. Handles patterns like:
@ -149,7 +149,7 @@ _STANDARD_RECORD_ATTRS = _get_standard_record_attrs()
class JsonFormatter(Formatter):
def __init__(self):
super(JsonFormatter, self).__init__()
super().__init__()
def formatTime(self, record, datefmt=None):
# Use datetime to format the timestamp in ISO 8601 format
@ -158,7 +158,7 @@ class JsonFormatter(Formatter):
def format(self, record):
message_str = record.getMessage()
json_record: Dict[str, Any] = {
json_record: dict[str, Any] = {
"message": message_str,
"level": record.levelname,
"timestamp": self.formatTime(record),

56
litellm/_redis.py
View file

@ -13,7 +13,6 @@ import json
# s/o [@Frank Colson](https://www.linkedin.com/in/frank-colson-422b9b183/) for this redis implementation
import os
from collections.abc import Callable
from typing import List, Optional, Union
import redis # type: ignore
import redis.asyncio as async_redis # type: ignore
@ -77,7 +76,7 @@ def _init_arg_names(cls: type) -> frozenset[str]:
)
def _get_redis_url_kwargs(client: Optional[type] = None) -> tuple[str, ...]:
def _get_redis_url_kwargs(client: type | None = None) -> tuple[str, ...]:
"""Connection kwargs that redis-py forwards from ``from_url`` down to the connection.
``from_url`` is declared as ``(cls, url, **kwargs)``, so introspecting it yields no
@ -161,7 +160,7 @@ def _redis_kwargs_from_environment():
def create_gcp_iam_redis_connect_func(
service_account: str,
ssl_ca_certs: Optional[str] = None,
ssl_ca_certs: str | None = None,
) -> Callable:
"""
Creates a custom Redis connection function for GCP IAM authentication.
@ -204,9 +203,9 @@ def create_gcp_iam_redis_connect_func(
def _build_azure_credential(
azure_client_id: Optional[str] = None,
azure_tenant_id: Optional[str] = None,
azure_client_secret: Optional[str] = None,
azure_client_id: str | None = None,
azure_tenant_id: str | None = None,
azure_client_secret: str | None = None,
):
"""
Build a long-lived Azure credential object.
@ -242,9 +241,9 @@ def _build_azure_credential(
def _generate_azure_ad_redis_token(
azure_client_id: Optional[str] = None,
azure_tenant_id: Optional[str] = None,
azure_client_secret: Optional[str] = None,
azure_client_id: str | None = None,
azure_tenant_id: str | None = None,
azure_client_secret: str | None = None,
) -> str:
"""
One-shot helper that builds a credential and fetches a single Azure AD
@ -264,9 +263,9 @@ def _generate_azure_ad_redis_token(
def create_azure_ad_redis_connect_func(
azure_client_id: Optional[str] = None,
azure_tenant_id: Optional[str] = None,
azure_client_secret: Optional[str] = None,
azure_client_id: str | None = None,
azure_tenant_id: str | None = None,
azure_client_secret: str | None = None,
) -> Callable:
"""
Creates a custom Redis connection function for Azure AD authentication.
@ -370,7 +369,7 @@ def _get_redis_client_logic(**env_overrides):
**env_overrides,
}
_startup_nodes: Optional[Union[str, list]] = redis_kwargs.get("startup_nodes", None) or get_secret( # type: ignore
_startup_nodes: str | list | None = redis_kwargs.get("startup_nodes", None) or get_secret( # type: ignore
"REDIS_CLUSTER_NODES"
)
@ -381,21 +380,21 @@ def _get_redis_client_logic(**env_overrides):
elif _startup_nodes is None:
redis_kwargs.pop("startup_nodes", None)
_sentinel_nodes: Optional[Union[str, list]] = redis_kwargs.get("sentinel_nodes", None) or get_secret( # type: ignore
_sentinel_nodes: str | list | None = redis_kwargs.get("sentinel_nodes", None) or get_secret( # type: ignore
"REDIS_SENTINEL_NODES"
)
if _sentinel_nodes is not None and isinstance(_sentinel_nodes, str):
redis_kwargs["sentinel_nodes"] = json.loads(_sentinel_nodes)
_sentinel_password: Optional[str] = redis_kwargs.get("sentinel_password", None) or get_secret_str(
_sentinel_password: str | None = redis_kwargs.get("sentinel_password", None) or get_secret_str(
"REDIS_SENTINEL_PASSWORD"
)
if _sentinel_password is not None:
redis_kwargs["sentinel_password"] = _sentinel_password
_service_name: Optional[str] = redis_kwargs.get("service_name", None) or get_secret( # type: ignore
_service_name: str | None = redis_kwargs.get("service_name", None) or get_secret( # type: ignore
"REDIS_SERVICE_NAME"
)
@ -466,9 +465,12 @@ def _get_redis_client_logic(**env_overrides):
redis_kwargs.pop("port", None)
redis_kwargs.pop("db", None)
redis_kwargs.pop("password", None)
elif "startup_nodes" in redis_kwargs and redis_kwargs["startup_nodes"] is not None:
pass
elif "sentinel_nodes" in redis_kwargs and redis_kwargs["sentinel_nodes"] is not None:
elif (
"startup_nodes" in redis_kwargs
and redis_kwargs["startup_nodes"] is not None
or "sentinel_nodes" in redis_kwargs
and redis_kwargs["sentinel_nodes"] is not None
):
pass
elif "host" not in redis_kwargs or redis_kwargs["host"] is None:
raise ValueError("Either 'host' or 'url' must be specified for redis.")
@ -478,7 +480,7 @@ def _get_redis_client_logic(**env_overrides):
def init_redis_cluster(redis_kwargs) -> redis.RedisCluster:
_redis_cluster_nodes_in_env: Optional[str] = get_secret("REDIS_CLUSTER_NODES") # type: ignore
_redis_cluster_nodes_in_env: str | None = get_secret("REDIS_CLUSTER_NODES") # type: ignore
if _redis_cluster_nodes_in_env is not None:
try:
redis_kwargs["startup_nodes"] = json.loads(_redis_cluster_nodes_in_env)
@ -496,7 +498,7 @@ def init_redis_cluster(redis_kwargs) -> redis.RedisCluster:
if arg in args:
cluster_kwargs[arg] = redis_kwargs[arg]
new_startup_nodes: List[ClusterNode] = []
new_startup_nodes: list[ClusterNode] = []
for item in redis_kwargs["startup_nodes"]:
new_startup_nodes.append(ClusterNode(**item))
@ -588,9 +590,9 @@ def get_redis_client(**env_overrides):
def get_redis_async_client(
connection_pool: Optional[async_redis.BlockingConnectionPool] = None,
connection_pool: async_redis.BlockingConnectionPool | None = None,
**env_overrides,
) -> Union[async_redis.Redis, async_redis.RedisCluster]:
) -> async_redis.Redis | async_redis.RedisCluster:
redis_kwargs = _get_redis_client_logic(**env_overrides)
if "startup_nodes" in redis_kwargs:
@ -619,7 +621,7 @@ def get_redis_async_client(
username=os.environ.get("REDIS_USERNAME") or None,
)
new_startup_nodes: List[ClusterNode] = []
new_startup_nodes: list[ClusterNode] = []
for item in redis_kwargs["startup_nodes"]:
new_startup_nodes.append(ClusterNode(**item))
@ -649,9 +651,7 @@ def get_redis_async_client(
if arg in args:
url_kwargs[arg] = redis_kwargs[arg]
else:
verbose_logger.debug(
"REDIS: ignoring argument: {}. Not an allowed async_redis.Redis.from_url arg.".format(arg)
)
verbose_logger.debug(f"REDIS: ignoring argument: {arg}. Not an allowed async_redis.Redis.from_url arg.")
return async_redis.Redis.from_url(**url_kwargs)
# Check for Redis Sentinel
@ -683,7 +683,7 @@ def get_redis_async_client(
def get_redis_connection_pool(
**env_overrides,
) -> Optional[async_redis.BlockingConnectionPool]:
) -> async_redis.BlockingConnectionPool | None:
redis_kwargs = _get_redis_client_logic(**env_overrides)
verbose_logger.debug("get_redis_connection_pool: redis_kwargs", redis_kwargs)

14
litellm/_redis_credential_provider.py
View file

@ -1,7 +1,7 @@
import asyncio
import threading
import time
from typing import Any, Dict, Optional, Tuple, Union
from typing import Any
from redis.credentials import CredentialProvider # type: ignore[attr-defined]
@ -14,7 +14,7 @@ _GCP_IAM_TOKEN_TTL_SECONDS = 3300
# Module-level cache shared across all GCPIAMCredentialProvider instances for the
# same service account, so multiple Redis connections on the same pod share one token.
# Keyed by service_account → (token, expiry_monotonic_timestamp).
_token_cache: Dict[str, Tuple[str, float]] = {}
_token_cache: dict[str, tuple[str, float]] = {}
_token_cache_lock = threading.Lock()
@ -95,11 +95,11 @@ class GCPIAMCredentialProvider(CredentialProvider):
def __init__(self, gcp_service_account: str) -> None:
self._gcp_service_account = gcp_service_account
def get_credentials(self) -> Tuple[str]:
def get_credentials(self) -> tuple[str]:
token = _get_cached_gcp_iam_token(self._gcp_service_account)
return (token,)
async def get_credentials_async(self) -> Tuple[str]:
async def get_credentials_async(self) -> tuple[str]:
token = await asyncio.to_thread(_get_cached_gcp_iam_token, self._gcp_service_account)
return (token,)
@ -115,17 +115,17 @@ class AzureADCredentialProvider(CredentialProvider):
fail authentication after the initial token expired (~1 hour TTL).
"""
def __init__(self, credential: Any, username: Optional[str] = None) -> None:
def __init__(self, credential: Any, username: str | None = None) -> None:
self._credential = credential
self._username = username
def get_credentials(self) -> Union[Tuple[str], Tuple[str, str]]:
def get_credentials(self) -> tuple[str] | tuple[str, str]:
token = self._credential.get_token(AZURE_REDIS_SCOPE).token
if self._username:
return (self._username, token)
return (token,)
async def get_credentials_async(self) -> Union[Tuple[str], Tuple[str, str]]:
async def get_credentials_async(self) -> tuple[str] | tuple[str, str]:
token_obj = await asyncio.to_thread(self._credential.get_token, AZURE_REDIS_SCOPE)
if self._username:
return (self._username, token_obj.token)

38
litellm/_service_logger.py
View file

@ -1,6 +1,6 @@
import asyncio
from datetime import datetime, timedelta
from typing import TYPE_CHECKING, Any, Optional, Union
from typing import TYPE_CHECKING, Any, Union
import litellm
from litellm._logging import verbose_logger
@ -24,7 +24,7 @@ else:
UserAPIKeyAuth = Any
def _get_otel_v2_class() -> Optional[type]:
def _get_otel_v2_class() -> type | None:
"""Return the ``OpenTelemetryV2`` class, or ``None`` if the OTel SDK is absent.
Imported lazily: ``litellm.integrations.otel.logger`` imports the OpenTelemetry
@ -54,7 +54,7 @@ class ServiceLogging(CustomLogger):
if "prometheus_system" in litellm.service_callback:
self.prometheusServicesLogger = PrometheusServicesLogger()
def _resolve_otel_service_logger(self, callback: Any) -> Optional[Any]:
def _resolve_otel_service_logger(self, callback: Any) -> Any | None:
"""Resolve the OTel logger (legacy or V2) to emit a service span on.
Returns the logger instance whose ``async_service_*_hook`` should fire for
@ -88,9 +88,9 @@ class ServiceLogging(CustomLogger):
service: ServiceTypes,
duration: float,
call_type: str,
parent_otel_span: Optional[Span] = None,
start_time: Optional[Union[datetime, float]] = None,
end_time: Optional[Union[float, datetime]] = None,
parent_otel_span: Span | None = None,
start_time: datetime | float | None = None,
end_time: float | datetime | None = None,
):
"""
Handles both sync and async monitoring by checking for existing event loop.
@ -152,10 +152,10 @@ class ServiceLogging(CustomLogger):
service: ServiceTypes,
call_type: str,
duration: float,
parent_otel_span: Optional[Span] = None,
start_time: Optional[Union[datetime, float]] = None,
end_time: Optional[Union[datetime, float]] = None,
event_metadata: Optional[dict] = None,
parent_otel_span: Span | None = None,
start_time: datetime | float | None = None,
end_time: datetime | float | None = None,
event_metadata: dict | None = None,
):
"""
- For counting if the redis, postgres call is successful
@ -218,7 +218,6 @@ class ServiceLogging(CustomLogger):
self.prometheusServicesLogger = PrometheusServicesLogger()
elif self.prometheusServicesLogger is None:
self.prometheusServicesLogger = self.prometheusServicesLogger()
return
async def init_datadog_logger_if_none(self):
"""
@ -230,8 +229,6 @@ class ServiceLogging(CustomLogger):
if not hasattr(self, "dd_logger"):
self.dd_logger: DataDogLogger = DataDogLogger()
return
async def init_otel_logger_if_none(self):
"""
initializes otel_logger if it is None or no attribute exists on ServiceLogging Object
@ -246,18 +243,17 @@ class ServiceLogging(CustomLogger):
verbose_logger.warning(
"ServiceLogger: open_telemetry_logger is None or not an instance of OpenTelemetry"
)
return
async def async_service_failure_hook(
self,
service: ServiceTypes,
duration: float,
error: Union[str, Exception],
error: str | Exception,
call_type: str,
parent_otel_span: Optional[Span] = None,
start_time: Optional[Union[datetime, float]] = None,
end_time: Optional[Union[float, datetime]] = None,
event_metadata: Optional[dict] = None,
parent_otel_span: Span | None = None,
start_time: datetime | float | None = None,
end_time: float | datetime | None = None,
event_metadata: dict | None = None,
):
"""
- For counting if the redis, postgres call is unsuccessful
@ -324,7 +320,7 @@ class ServiceLogging(CustomLogger):
request_data: dict,
original_exception: Exception,
user_api_key_dict: UserAPIKeyAuth,
traceback_str: Optional[str] = None,
traceback_str: str | None = None,
):
"""
Hook to track failed litellm-service calls
@ -347,7 +343,7 @@ class ServiceLogging(CustomLogger):
pass
else:
raise Exception(
"Duration={} is not a float or timedelta object. type={}".format(_duration, type(_duration))
f"Duration={_duration} is not a float or timedelta object. type={type(_duration)}"
) # invalid _duration value
# Batch polling callbacks (check_batch_cost) don't include call_type in kwargs.
# Use .get() to avoid KeyError.

4
litellm/a2a_protocol/card_resolver.py
View file

@ -4,7 +4,7 @@ Custom A2A Card Resolver for LiteLLM.
Extends the A2A SDK's card resolver to support multiple well-known paths.
"""
from typing import TYPE_CHECKING, Any, Dict
from typing import TYPE_CHECKING, Any
from litellm._logging import verbose_logger
from litellm.constants import LOCALHOST_URL_PATTERNS
@ -114,7 +114,7 @@ class LiteLLMA2ACardResolver(_A2ACardResolver): # type: ignore[misc]
async def get_agent_card(
self,
relative_card_path: str | None = None,
http_kwargs: Dict[str, Any] | None = None,
http_kwargs: dict[str, Any] | None = None,
) -> "AgentCard":
"""
Fetch the agent card, trying multiple well-known paths.

6
litellm/a2a_protocol/client.py
View file

@ -5,7 +5,7 @@ Provides a class-based interface for A2A agent invocation.
"""
from collections.abc import AsyncIterator
from typing import TYPE_CHECKING, Dict, Optional
from typing import TYPE_CHECKING
from litellm.types.agents import LiteLLMSendMessageResponse
@ -51,7 +51,7 @@ class A2AClient:
self,
base_url: str,
timeout: float = 60.0,
extra_headers: Optional[Dict[str, str]] = None,
extra_headers: dict[str, str] | None = None,
):
"""
Initialize the A2A client wrapper.
@ -64,7 +64,7 @@ class A2AClient:
self.base_url = base_url
self.timeout = timeout
self.extra_headers = extra_headers
self._a2a_client: Optional["A2AClientType"] = None
self._a2a_client: A2AClientType | None = None
async def _get_client(self) -> "A2AClientType":
"""Get or create the underlying A2A client."""

8
litellm/a2a_protocol/cost_calculator.py
View file

@ -5,7 +5,7 @@ Supports dynamic cost parameters that allow platform owners
to define custom costs per agent query or per token.
"""
from typing import TYPE_CHECKING, Any, Optional
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import (
@ -18,7 +18,7 @@ else:
class A2ACostCalculator:
@staticmethod
def calculate_a2a_cost(
litellm_logging_obj: Optional[LitellmLoggingObject],
litellm_logging_obj: LitellmLoggingObject | None,
) -> float:
"""
Calculate the cost of an A2A send_message call.
@ -73,8 +73,8 @@ class A2ACostCalculator:
@staticmethod
def _calculate_token_based_cost(
model_call_details: dict,
input_cost_per_token: Optional[float],
output_cost_per_token: Optional[float],
input_cost_per_token: float | None,
output_cost_per_token: float | None,
) -> float:
"""
Calculate cost based on token usage and per-token pricing.

10
litellm/a2a_protocol/exception_mapping_utils.py
View file

@ -4,7 +4,7 @@ A2A Protocol Exception Mapping Utils.
Maps A2A SDK exceptions to LiteLLM A2A exception types.
"""
from typing import TYPE_CHECKING, Any, Optional
from typing import TYPE_CHECKING, Any
from litellm._logging import verbose_logger
from litellm.a2a_protocol.card_resolver import (
@ -57,7 +57,7 @@ class A2AExceptionCheckers:
return any(pattern in error_str_lower for pattern in CONNECTION_ERROR_PATTERNS)
@staticmethod
def is_localhost_url(url: Optional[str]) -> bool:
def is_localhost_url(url: str | None) -> bool:
"""
Check if a URL is a localhost/internal URL.
@ -96,9 +96,9 @@ class A2AExceptionCheckers:
def map_a2a_exception(
original_exception: Exception,
card_url: Optional[str] = None,
api_base: Optional[str] = None,
model: Optional[str] = None,
card_url: str | None = None,
api_base: str | None = None,
model: str | None = None,
) -> Exception:
"""
Map an A2A SDK exception to a LiteLLM A2A exception type.

36
litellm/a2a_protocol/exceptions.py
View file

@ -4,8 +4,6 @@ A2A Protocol Exceptions.
Custom exception types for A2A protocol operations, following LiteLLM's exception pattern.
"""
from typing import Optional
import httpx
@ -21,11 +19,11 @@ class A2AError(Exception):
message: str,
status_code: int = 500,
llm_provider: str = "a2a_agent",
model: Optional[str] = None,
response: Optional[httpx.Response] = None,
litellm_debug_info: Optional[str] = None,
max_retries: Optional[int] = None,
num_retries: Optional[int] = None,
model: str | None = None,
response: httpx.Response | None = None,
litellm_debug_info: str | None = None,
max_retries: int | None = None,
num_retries: int | None = None,
):
self.status_code = status_code
self.message = f"litellm.A2AError: {message}"
@ -65,12 +63,12 @@ class A2AConnectionError(A2AError):
def __init__(
self,
message: str,
url: Optional[str] = None,
model: Optional[str] = None,
response: Optional[httpx.Response] = None,
litellm_debug_info: Optional[str] = None,
max_retries: Optional[int] = None,
num_retries: Optional[int] = None,
url: str | None = None,
model: str | None = None,
response: httpx.Response | None = None,
litellm_debug_info: str | None = None,
max_retries: int | None = None,
num_retries: int | None = None,
):
self.url = url
super().__init__(
@ -98,10 +96,10 @@ class A2AAgentCardError(A2AError):
def __init__(
self,
message: str,
url: Optional[str] = None,
model: Optional[str] = None,
response: Optional[httpx.Response] = None,
litellm_debug_info: Optional[str] = None,
url: str | None = None,
model: str | None = None,
response: httpx.Response | None = None,
litellm_debug_info: str | None = None,
):
self.url = url
super().__init__(
@ -132,8 +130,8 @@ class A2ALocalhostURLError(A2AConnectionError):
self,
localhost_url: str,
base_url: str,
original_error: Optional[Exception] = None,
model: Optional[str] = None,
original_error: Exception | None = None,
model: str | None = None,
):
self.localhost_url = localhost_url
self.base_url = base_url

2
litellm/a2a_protocol/litellm_completion_bridge/__init__.py
View file

@ -16,8 +16,8 @@ from litellm.a2a_protocol.litellm_completion_bridge.transformation import (
)
__all__ = [
"A2ACompletionBridgeTransformation",
"A2ACompletionBridgeHandler",
"A2ACompletionBridgeTransformation",
"handle_a2a_completion",
"handle_a2a_completion_streaming",
]

46
litellm/a2a_protocol/litellm_completion_bridge/handler.py
View file

@ -11,7 +11,7 @@ A2A Streaming Events (in order):
"""
from collections.abc import AsyncIterator
from typing import Any, Dict, Optional
from typing import Any
import litellm
from litellm._logging import verbose_logger
@ -47,13 +47,13 @@ class A2ACompletionBridgeHandler:
@staticmethod
async def handle_non_streaming(
request_id: str,
params: Dict[str, Any],
litellm_params: Dict[str, Any],
api_base: Optional[str] = None,
agent_extra_headers: Optional[Dict[str, str]] = None,
params: dict[str, Any],
litellm_params: dict[str, Any],
api_base: str | None = None,
agent_extra_headers: dict[str, str] | None = None,
*,
_skip_a2a_provider_routing: bool = False,
) -> Dict[str, Any]:
) -> dict[str, Any]:
"""
Handle non-streaming A2A request via litellm.acompletion.
@ -106,7 +106,7 @@ class A2ACompletionBridgeHandler:
verbose_logger.info(f"A2A completion bridge: model={full_model}, api_base={api_base}")
# Build completion params dict
completion_params: Dict[str, Any] = {
completion_params: dict[str, Any] = {
"model": full_model,
"messages": openai_messages,
"api_base": api_base,
@ -150,13 +150,13 @@ class A2ACompletionBridgeHandler:
@staticmethod
async def handle_streaming(
request_id: str,
params: Dict[str, Any],
litellm_params: Dict[str, Any],
api_base: Optional[str] = None,
agent_extra_headers: Optional[Dict[str, str]] = None,
params: dict[str, Any],
litellm_params: dict[str, Any],
api_base: str | None = None,
agent_extra_headers: dict[str, str] | None = None,
*,
_skip_a2a_provider_routing: bool = False,
) -> AsyncIterator[Dict[str, Any]]:
) -> AsyncIterator[dict[str, Any]]:
"""
Handle streaming A2A request via litellm.acompletion with stream=True.
@ -224,7 +224,7 @@ class A2ACompletionBridgeHandler:
verbose_logger.info(f"A2A completion bridge streaming: model={full_model}, api_base={api_base}")
# Build completion params dict
completion_params: Dict[str, Any] = {
completion_params: dict[str, Any] = {
"model": full_model,
"messages": openai_messages,
"api_base": api_base,
@ -306,11 +306,11 @@ class A2ACompletionBridgeHandler:
# Convenience functions that delegate to the class methods
async def handle_a2a_completion(
request_id: str,
params: Dict[str, Any],
litellm_params: Dict[str, Any],
api_base: Optional[str] = None,
agent_extra_headers: Optional[Dict[str, str]] = None,
) -> Dict[str, Any]:
params: dict[str, Any],
litellm_params: dict[str, Any],
api_base: str | None = None,
agent_extra_headers: dict[str, str] | None = None,
) -> dict[str, Any]:
"""Convenience function for non-streaming A2A completion."""
return await A2ACompletionBridgeHandler.handle_non_streaming(
request_id=request_id,
@ -323,11 +323,11 @@ async def handle_a2a_completion(
async def handle_a2a_completion_streaming(
request_id: str,
params: Dict[str, Any],
litellm_params: Dict[str, Any],
api_base: Optional[str] = None,
agent_extra_headers: Optional[Dict[str, str]] = None,
) -> AsyncIterator[Dict[str, Any]]:
params: dict[str, Any],
litellm_params: dict[str, Any],
api_base: str | None = None,
agent_extra_headers: dict[str, str] | None = None,
) -> AsyncIterator[dict[str, Any]]:
"""Convenience function for streaming A2A completion."""
async for chunk in A2ACompletionBridgeHandler.handle_streaming(
request_id=request_id,

46
litellm/a2a_protocol/litellm_completion_bridge/transformation.py
View file

@ -18,7 +18,7 @@ A2A Streaming Events:
"""
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional
from typing import Any
from uuid import uuid4
from litellm._logging import verbose_logger
@ -30,7 +30,7 @@ class A2AStreamingContext:
Tracks task_id, context_id, and message accumulation.
"""
def __init__(self, request_id: str, input_message: Dict[str, Any]):
def __init__(self, request_id: str, input_message: dict[str, Any]):
self.request_id = request_id
self.task_id = str(uuid4())
self.context_id = str(uuid4())
@ -46,9 +46,9 @@ class A2ACompletionBridgeTransformation:
"""
@staticmethod
def _extract_text_from_a2a_parts(parts: List[Dict[str, Any]]) -> str:
def _extract_text_from_a2a_parts(parts: list[dict[str, Any]]) -> str:
"""Extract text from A2A parts (with or without explicit ``kind``)."""
content_parts: List[str] = []
content_parts: list[str] = []
for part in parts:
if not isinstance(part, dict):
continue
@ -62,16 +62,16 @@ class A2ACompletionBridgeTransformation:
@staticmethod
def get_forward_metadata(
a2a_message: Dict[str, Any],
params: Optional[Dict[str, Any]] = None,
) -> Optional[Dict[str, Any]]:
a2a_message: dict[str, Any],
params: dict[str, Any] | None = None,
) -> dict[str, Any] | None:
"""
Merge A2A metadata from MessageSendParams and the message for downstream providers.
Forwarded once on the LangGraph run payload (``metadata``), not duplicated on
each input message — see ``apply_forward_metadata_to_completion_params``.
"""
merged: Dict[str, Any] = {}
merged: dict[str, Any] = {}
if params and isinstance(params.get("metadata"), dict):
merged.update(params["metadata"])
message_metadata = a2a_message.get("metadata")
@ -81,9 +81,9 @@ class A2ACompletionBridgeTransformation:
@staticmethod
def apply_forward_metadata_to_completion_params(
completion_params: Dict[str, Any],
a2a_message: Dict[str, Any],
params: Optional[Dict[str, Any]] = None,
completion_params: dict[str, Any],
a2a_message: dict[str, Any],
params: dict[str, Any] | None = None,
) -> None:
"""
Attach A2A metadata to completion kwargs for provider bridges (e.g. LangGraph).
@ -104,8 +104,8 @@ class A2ACompletionBridgeTransformation:
# ``extra_body.metadata`` so the configured keys remain authoritative
# and an A2A caller cannot overwrite server-set run metadata.
existing_metadata = extra_body.get("metadata")
existing_dict: Dict[str, Any] = existing_metadata if isinstance(existing_metadata, dict) else {}
merged_metadata: Dict[str, Any] = {**forward_metadata, **existing_dict}
existing_dict: dict[str, Any] = existing_metadata if isinstance(existing_metadata, dict) else {}
merged_metadata: dict[str, Any] = {**forward_metadata, **existing_dict}
extra_body = {**extra_body, "metadata": merged_metadata}
completion_params["extra_body"] = extra_body
@ -113,8 +113,8 @@ class A2ACompletionBridgeTransformation:
@staticmethod
def a2a_message_to_openai_messages(
a2a_message: Dict[str, Any],
) -> List[Dict[str, Any]]:
a2a_message: dict[str, Any],
) -> list[dict[str, Any]]:
"""
Transform an A2A message to OpenAI message format.
@ -143,7 +143,7 @@ class A2ACompletionBridgeTransformation:
# Do not attach A2A message.metadata here — the completion bridge forwards it
# once at run level via extra_body.metadata (LangGraph POST /runs/wait shape).
openai_message: Dict[str, Any] = {"role": openai_role, "content": content}
openai_message: dict[str, Any] = {"role": openai_role, "content": content}
verbose_logger.debug(f"A2A -> OpenAI transform: role={role} -> {openai_role}, content_length={len(content)}")
@ -152,8 +152,8 @@ class A2ACompletionBridgeTransformation:
@staticmethod
def openai_response_to_a2a_response(
response: Any,
request_id: Optional[str] = None,
) -> Dict[str, Any]:
request_id: str | None = None,
) -> dict[str, Any]:
"""
Transform a LiteLLM ModelResponse to A2A SendMessageResponse format.
@ -198,7 +198,7 @@ class A2ACompletionBridgeTransformation:
@staticmethod
def create_task_event(
ctx: A2AStreamingContext,
) -> Dict[str, Any]:
) -> dict[str, Any]:
"""
Create the initial task event with status 'submitted'.
@ -232,8 +232,8 @@ class A2ACompletionBridgeTransformation:
ctx: A2AStreamingContext,
state: str,
final: bool = False,
message_text: Optional[str] = None,
) -> Dict[str, Any]:
message_text: str | None = None,
) -> dict[str, Any]:
"""
Create a status update event.
@ -243,7 +243,7 @@ class A2ACompletionBridgeTransformation:
final: Whether this is the final event
message_text: Optional message text for 'working' status
"""
status: Dict[str, Any] = {
status: dict[str, Any] = {
"state": state,
"timestamp": A2ACompletionBridgeTransformation._get_timestamp(),
}
@ -275,7 +275,7 @@ class A2ACompletionBridgeTransformation:
def create_artifact_update_event(
ctx: A2AStreamingContext,
text: str,
) -> Dict[str, Any]:
) -> dict[str, Any]:
"""
Create an artifact update event with content.

38
litellm/a2a_protocol/main.py
View file

@ -16,9 +16,7 @@ from collections.abc import AsyncIterator, Coroutine
from typing import (
TYPE_CHECKING,
Any,
Dict,
Optional,
Union,
cast,
)
@ -86,7 +84,7 @@ A2ACardResolver = LiteLLMA2ACardResolver
def _set_usage_on_logging_obj(
kwargs: Dict[str, Any],
kwargs: dict[str, Any],
prompt_tokens: int,
completion_tokens: int,
) -> None:
@ -109,7 +107,7 @@ def _set_usage_on_logging_obj(
def _set_agent_id_on_logging_obj(
kwargs: Dict[str, Any],
kwargs: dict[str, Any],
agent_id: str | None,
) -> None:
"""
@ -155,7 +153,7 @@ def _set_litellm_params_on_logging_obj(
logging_obj.model_call_details["litellm_params"] = {**existing, **cost_params}
def _get_a2a_model_info(a2a_client: Any, kwargs: Dict[str, Any]) -> str:
def _get_a2a_model_info(a2a_client: Any, kwargs: dict[str, Any]) -> str:
"""
Extract agent info and set model/custom_llm_provider for cost tracking.
@ -198,8 +196,8 @@ async def _send_message_via_completion_bridge(
request: "SendMessageRequest",
custom_llm_provider: str,
api_base: str | None,
litellm_params: Dict[str, Any],
agent_extra_headers: Dict[str, str] | None = None,
litellm_params: dict[str, Any],
agent_extra_headers: dict[str, str] | None = None,
) -> LiteLLMSendMessageResponse:
"""
Route a send_message through the LiteLLM completion bridge (e.g. LangGraph, Bedrock AgentCore).
@ -369,9 +367,9 @@ async def asend_message(
a2a_client: Optional["A2AClientType"] = None,
request: Optional["SendMessageRequest"] = None,
api_base: str | None = None,
litellm_params: Dict[str, Any] | None = None,
litellm_params: dict[str, Any] | None = None,
agent_id: str | None = None,
agent_extra_headers: Dict[str, str] | None = None,
agent_extra_headers: dict[str, str] | None = None,
**kwargs: Any,
) -> LiteLLMSendMessageResponse:
"""
@ -452,7 +450,7 @@ async def asend_message(
if api_base is None:
raise ValueError("Either a2a_client or api_base is required for standard A2A flow")
trace_id = trace_id or str(uuid.uuid4())
extra_headers: Dict[str, str] = {"X-LiteLLM-Trace-Id": trace_id}
extra_headers: dict[str, str] = {"X-LiteLLM-Trace-Id": trace_id}
if agent_id:
extra_headers["X-LiteLLM-Agent-Id"] = agent_id
# Overlay agent-level headers (agent headers take precedence over LiteLLM internal ones)
@ -517,7 +515,7 @@ def send_message(
a2a_client: "A2AClientType",
request: "SendMessageRequest",
**kwargs: Any,
) -> Union[LiteLLMSendMessageResponse, Coroutine[Any, Any, LiteLLMSendMessageResponse]]:
) -> LiteLLMSendMessageResponse | Coroutine[Any, Any, LiteLLMSendMessageResponse]:
"""
Sync: Send a message to an A2A agent.
@ -546,9 +544,9 @@ def _build_streaming_logging_obj(
request: "SendStreamingMessageRequest",
agent_name: str,
agent_id: str | None,
litellm_params: Dict[str, Any] | None,
metadata: Dict[str, Any] | None,
proxy_server_request: Dict[str, Any] | None,
litellm_params: dict[str, Any] | None,
metadata: dict[str, Any] | None,
proxy_server_request: dict[str, Any] | None,
) -> Logging:
"""Build logging object for streaming A2A requests."""
start_time = datetime.datetime.now()
@ -589,11 +587,11 @@ async def asend_message_streaming(
a2a_client: Optional["A2AClientType"] = None,
request: Optional["SendStreamingMessageRequest"] = None,
api_base: str | None = None,
litellm_params: Dict[str, Any] | None = None,
litellm_params: dict[str, Any] | None = None,
agent_id: str | None = None,
metadata: Dict[str, Any] | None = None,
proxy_server_request: Dict[str, Any] | None = None,
agent_extra_headers: Dict[str, str] | None = None,
metadata: dict[str, Any] | None = None,
proxy_server_request: dict[str, Any] | None = None,
agent_extra_headers: dict[str, str] | None = None,
**kwargs: object,
) -> AsyncIterator[Any]:
"""
@ -727,7 +725,7 @@ async def asend_message_streaming(
async def create_a2a_client(
base_url: str,
timeout: float = DEFAULT_A2A_AGENT_TIMEOUT,
extra_headers: Dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
streaming: bool = False,
) -> "A2AClientType":
"""
@ -808,7 +806,7 @@ async def create_a2a_client(
async def aget_agent_card(
base_url: str,
timeout: float = DEFAULT_A2A_AGENT_TIMEOUT,
extra_headers: Dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
) -> "AgentCard":
"""
Fetch the agent card from an A2A agent.

2
litellm/a2a_protocol/providers/__init__.py
View file

@ -7,4 +7,4 @@ This module contains provider-specific implementations for the A2A protocol.
from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig
from litellm.a2a_protocol.providers.config_manager import A2AProviderConfigManager
__all__ = ["BaseA2AProviderConfig", "A2AProviderConfigManager"]
__all__ = ["A2AProviderConfigManager", "BaseA2AProviderConfig"]

15
litellm/a2a_protocol/providers/base.py
View file

@ -4,7 +4,7 @@ Base configuration for A2A protocol providers.
from abc import ABC, abstractmethod
from collections.abc import AsyncIterator
from typing import Any, Dict, Optional
from typing import Any
class BaseA2AProviderConfig(ABC):
@ -19,10 +19,10 @@ class BaseA2AProviderConfig(ABC):
async def handle_non_streaming(
self,
request_id: str,
params: Dict[str, Any],
api_base: Optional[str] = None,
params: dict[str, Any],
api_base: str | None = None,
**kwargs,
) -> Dict[str, Any]:
) -> dict[str, Any]:
"""
Handle non-streaming A2A request.
@ -35,16 +35,15 @@ class BaseA2AProviderConfig(ABC):
Returns:
A2A SendMessageResponse dict
"""
pass
@abstractmethod
async def handle_streaming(
self,
request_id: str,
params: Dict[str, Any],
api_base: Optional[str] = None,
params: dict[str, Any],
api_base: str | None = None,
**kwargs,
) -> AsyncIterator[Dict[str, Any]]:
) -> AsyncIterator[dict[str, Any]]:
"""
Handle streaming A2A request.

14
litellm/a2a_protocol/providers/bedrock_agentcore/config.py
View file

@ -3,7 +3,7 @@ Bedrock AgentCore A2A provider configuration.
"""
from collections.abc import AsyncIterator
from typing import Any, Dict, Optional
from typing import Any
from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig
from litellm.a2a_protocol.providers.bedrock_agentcore.handler import (
@ -23,10 +23,10 @@ class BedrockAgentCoreA2AConfig(BaseA2AProviderConfig):
async def handle_non_streaming(
self,
request_id: str,
params: Dict[str, Any],
api_base: Optional[str] = None,
params: dict[str, Any],
api_base: str | None = None,
**kwargs,
) -> Dict[str, Any]:
) -> dict[str, Any]:
"""Handle non-streaming request to AgentCore A2A agent."""
litellm_params = kwargs.get("litellm_params")
if not litellm_params:
@ -43,10 +43,10 @@ class BedrockAgentCoreA2AConfig(BaseA2AProviderConfig):
async def handle_streaming(
self,
request_id: str,
params: Dict[str, Any],
api_base: Optional[str] = None,
params: dict[str, Any],
api_base: str | None = None,
**kwargs,
) -> AsyncIterator[Dict[str, Any]]:
) -> AsyncIterator[dict[str, Any]]:
"""Handle streaming request to AgentCore A2A agent."""
litellm_params = kwargs.get("litellm_params")
if not litellm_params:

18
litellm/a2a_protocol/providers/bedrock_agentcore/handler.py
View file

@ -7,7 +7,7 @@ completion bridge that would otherwise strip the envelope.
import json
from collections.abc import AsyncIterator
from typing import Any, Dict, Optional, cast
from typing import Any, cast
from litellm._logging import verbose_logger
from litellm.a2a_protocol.providers.bedrock_agentcore.transformation import (
@ -28,10 +28,10 @@ class BedrockAgentCoreA2AHandler:
@staticmethod
async def handle_non_streaming(
request_id: str,
params: Dict[str, Any],
litellm_params: Dict[str, Any],
agent_extra_headers: Optional[Dict[str, str]] = None,
) -> Dict[str, Any]:
params: dict[str, Any],
litellm_params: dict[str, Any],
agent_extra_headers: dict[str, str] | None = None,
) -> dict[str, Any]:
"""
Handle non-streaming A2A request to AgentCore.
@ -74,10 +74,10 @@ class BedrockAgentCoreA2AHandler:
@staticmethod
async def handle_streaming(
request_id: str,
params: Dict[str, Any],
litellm_params: Dict[str, Any],
agent_extra_headers: Optional[Dict[str, str]] = None,
) -> AsyncIterator[Dict[str, Any]]:
params: dict[str, Any],
litellm_params: dict[str, Any],
agent_extra_headers: dict[str, str] | None = None,
) -> AsyncIterator[dict[str, Any]]:
"""
Handle streaming A2A request to AgentCore.

20
litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py
View file

@ -7,7 +7,7 @@ and signs requests via AmazonAgentCoreConfig (SigV4 or JWT).
import json
from collections.abc import AsyncIterator, Mapping
from typing import Any, Dict, Optional, Tuple
from typing import Any
from litellm._logging import verbose_logger
from litellm.llms.bedrock.chat.agentcore.transformation import AmazonAgentCoreConfig
@ -29,15 +29,15 @@ _RESERVED_EXACT_HEADERS = frozenset(
"host",
}
)
_RESERVED_PREFIX_HEADERS: Tuple[str, ...] = (
_RESERVED_PREFIX_HEADERS: tuple[str, ...] = (
"x-amzn-bedrock-agentcore-runtime-",
"x-amz-",
)
def _filter_reserved_headers(
agent_extra_headers: Optional[Mapping[str, str]],
) -> Optional[Dict[str, str]]:
agent_extra_headers: Mapping[str, str] | None,
) -> dict[str, str] | None:
"""
Strip reserved AWS / AgentCore headers from caller-supplied
``agent_extra_headers`` before they are merged into the signed request.
@ -47,7 +47,7 @@ def _filter_reserved_headers(
if not agent_extra_headers:
return None
filtered: Dict[str, str] = {}
filtered: dict[str, str] = {}
dropped: list = []
for k, v in agent_extra_headers.items():
k_lower = k.lower()
@ -77,12 +77,12 @@ class BedrockAgentCoreA2ATransformation:
@staticmethod
def get_url_and_signed_request(
request_id: str,
params: Dict[str, Any],
litellm_params: Dict[str, Any],
params: dict[str, Any],
litellm_params: dict[str, Any],
method: str = "message/send",
stream: bool = False,
agent_extra_headers: Optional[Dict[str, str]] = None,
) -> Tuple[str, dict, bytes]:
agent_extra_headers: dict[str, str] | None = None,
) -> tuple[str, dict, bytes]:
"""
Build the AgentCore URL, construct a JSON-RPC envelope, and sign the request.
@ -170,7 +170,7 @@ class BedrockAgentCoreA2ATransformation:
return url, signed_headers, signed_body
@staticmethod
async def parse_sse_events(response: Any) -> AsyncIterator[Dict[str, Any]]:
async def parse_sse_events(response: Any) -> AsyncIterator[dict[str, Any]]:
"""
Parse SSE events from an httpx streaming response.

8
litellm/a2a_protocol/providers/config_manager.py
View file

@ -4,8 +4,6 @@ A2A Provider Config Manager.
Manages provider-specific configurations for A2A protocol.
"""
from typing import Optional
from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig
@ -18,9 +16,9 @@ class A2AProviderConfigManager:
@staticmethod
def get_provider_config(
custom_llm_provider: Optional[str],
model: Optional[str] = None,
) -> Optional[BaseA2AProviderConfig]:
custom_llm_provider: str | None,
model: str | None = None,
) -> BaseA2AProviderConfig | None:
"""
Get the provider configuration for a given custom_llm_provider.

14
litellm/a2a_protocol/providers/langflow/config.py
View file

@ -1,5 +1,5 @@
from collections.abc import AsyncIterator
from typing import Any, Dict, Optional
from typing import Any
from litellm.a2a_protocol.litellm_completion_bridge.handler import (
A2A_USER_API_KEY_HASH_PARAM,
@ -16,10 +16,10 @@ class LangFlowA2AConfig(BaseA2AProviderConfig):
async def handle_non_streaming(
self,
request_id: str,
params: Dict[str, Any],
api_base: Optional[str] = None,
params: dict[str, Any],
api_base: str | None = None,
**kwargs,
) -> Dict[str, Any]:
) -> dict[str, Any]:
litellm_params = kwargs.get("litellm_params")
if not litellm_params:
raise ValueError(
@ -39,10 +39,10 @@ class LangFlowA2AConfig(BaseA2AProviderConfig):
async def handle_streaming(
self,
request_id: str,
params: Dict[str, Any],
api_base: Optional[str] = None,
params: dict[str, Any],
api_base: str | None = None,
**kwargs,
) -> AsyncIterator[Dict[str, Any]]:
) -> AsyncIterator[dict[str, Any]]:
litellm_params = kwargs.get("litellm_params")
if not litellm_params:
raise ValueError(

2
litellm/a2a_protocol/providers/pydantic_ai_agents/__init__.py
View file

@ -13,4 +13,4 @@ from litellm.a2a_protocol.providers.pydantic_ai_agents.transformation import (
PydanticAITransformation,
)
__all__ = ["PydanticAIHandler", "PydanticAITransformation", "PydanticAIProviderConfig"]
__all__ = ["PydanticAIHandler", "PydanticAIProviderConfig", "PydanticAITransformation"]

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

@ -3,7 +3,7 @@ Pydantic AI provider configuration.
"""
from collections.abc import AsyncIterator
from typing import Any, Dict, Optional
from typing import Any
from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig
from litellm.a2a_protocol.providers.pydantic_ai_agents.handler import PydanticAIHandler
@ -20,10 +20,10 @@ class PydanticAIProviderConfig(BaseA2AProviderConfig):
async def handle_non_streaming(
self,
request_id: str,
params: Dict[str, Any],
api_base: Optional[str] = None,
params: dict[str, Any],
api_base: str | None = None,
**kwargs: Any,
) -> Dict[str, Any]:
) -> dict[str, Any]:
"""Handle non-streaming request to Pydantic AI agent."""
if api_base is None:
raise ValueError("api_base is required for PydanticAIProviderConfig")
@ -38,10 +38,10 @@ class PydanticAIProviderConfig(BaseA2AProviderConfig):
async def handle_streaming(
self,
request_id: str,
params: Dict[str, Any],
api_base: Optional[str] = None,
params: dict[str, Any],
api_base: str | None = None,
**kwargs,
) -> AsyncIterator[Dict[str, Any]]:
) -> AsyncIterator[dict[str, Any]]:
"""Handle streaming request with fake streaming."""
if not api_base:
raise ValueError("api_base is required for Pydantic AI agents")

18
litellm/a2a_protocol/providers/pydantic_ai_agents/handler.py
View file

@ -6,7 +6,7 @@ This handler provides fake streaming by converting non-streaming responses into
"""
from collections.abc import AsyncIterator
from typing import Any, Dict, Optional
from typing import Any
from litellm._logging import verbose_logger
from litellm.a2a_protocol.providers.pydantic_ai_agents.transformation import (
@ -26,11 +26,11 @@ class PydanticAIHandler:
@staticmethod
async def handle_non_streaming(
request_id: str,
params: Dict[str, Any],
api_base: Optional[str] = None,
params: dict[str, Any],
api_base: str | None = None,
timeout: float = 60.0,
agent_extra_headers: Optional[Dict[str, str]] = None,
) -> Dict[str, Any]:
agent_extra_headers: dict[str, str] | None = None,
) -> dict[str, Any]:
"""
Handle non-streaming request to Pydantic AI agent.
@ -63,13 +63,13 @@ class PydanticAIHandler:
@staticmethod
async def handle_streaming(
request_id: str,
params: Dict[str, Any],
api_base: Optional[str] = None,
params: dict[str, Any],
api_base: str | None = None,
timeout: float = 60.0,
chunk_size: int = 50,
delay_ms: int = 10,
agent_extra_headers: Optional[Dict[str, str]] = None,
) -> AsyncIterator[Dict[str, Any]]:
agent_extra_headers: dict[str, str] | None = None,
) -> AsyncIterator[dict[str, Any]]:
"""
Handle streaming request to Pydantic AI agent with fake streaming.

30
litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py
View file

@ -7,7 +7,7 @@ This module provides fake streaming by converting non-streaming responses into s
import asyncio
from collections.abc import AsyncIterator
from typing import Any, Dict, Optional, cast
from typing import Any, cast
from uuid import uuid4
from litellm._logging import verbose_logger
@ -49,7 +49,7 @@ class PydanticAITransformation:
return obj
@staticmethod
def _params_to_dict(params: Any) -> Dict[str, Any]:
def _params_to_dict(params: Any) -> dict[str, Any]:
"""
Convert params to a dict, handling Pydantic models.
@ -79,8 +79,8 @@ class PydanticAITransformation:
request_id: str,
max_attempts: int = 30,
poll_interval: float = 0.5,
agent_extra_headers: Optional[Dict[str, str]] = None,
) -> Dict[str, Any]:
agent_extra_headers: dict[str, str] | None = None,
) -> dict[str, Any]:
"""
Poll for task completion using tasks/get method.
@ -135,8 +135,8 @@ class PydanticAITransformation:
request_id: str,
params: Any,
timeout: float = 60.0,
agent_extra_headers: Optional[Dict[str, str]] = None,
) -> Dict[str, Any]:
agent_extra_headers: dict[str, str] | None = None,
) -> dict[str, Any]:
"""
Send a request to Pydantic AI agent and return the raw task response.
@ -219,8 +219,8 @@ class PydanticAITransformation:
request_id: str,
params: Any,
timeout: float = 60.0,
agent_extra_headers: Optional[Dict[str, str]] = None,
) -> Dict[str, Any]:
agent_extra_headers: dict[str, str] | None = None,
) -> dict[str, Any]:
"""
Send a non-streaming A2A request to Pydantic AI agent and wait for completion.
@ -255,8 +255,8 @@ class PydanticAITransformation:
request_id: str,
params: Any,
timeout: float = 60.0,
agent_extra_headers: Optional[Dict[str, str]] = None,
) -> Dict[str, Any]:
agent_extra_headers: dict[str, str] | None = None,
) -> dict[str, Any]:
"""
Send a request to Pydantic AI agent and return the raw task response.
@ -282,9 +282,9 @@ class PydanticAITransformation:
@staticmethod
def _transform_to_a2a_response(
response_data: Dict[str, Any],
response_data: dict[str, Any],
request_id: str,
) -> Dict[str, Any]:
) -> dict[str, Any]:
"""
Transform Pydantic AI task response to standard A2A non-streaming format.
@ -328,7 +328,7 @@ class PydanticAITransformation:
}
@staticmethod
def _extract_response_text(response_data: Dict[str, Any]) -> tuple[str, str, list]:
def _extract_response_text(response_data: dict[str, Any]) -> tuple[str, str, list]:
"""
Extract response text from completed task response.
@ -383,11 +383,11 @@ class PydanticAITransformation:
@staticmethod
async def fake_streaming_from_response(
response_data: Dict[str, Any],
response_data: dict[str, Any],
request_id: str,
chunk_size: int = 50,
delay_ms: int = 10,
) -> AsyncIterator[Dict[str, Any]]:
) -> AsyncIterator[dict[str, Any]]:
"""
Convert a non-streaming A2A response into fake streaming chunks.

14
litellm/a2a_protocol/providers/watsonx_orchestrate/config.py
View file

@ -3,7 +3,7 @@ A2A provider configuration for IBM watsonx Orchestrate (WXO).
"""
from collections.abc import AsyncIterator
from typing import Any, Dict, Optional
from typing import Any
from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig
from litellm.a2a_protocol.providers.watsonx_orchestrate.handler import (
@ -17,10 +17,10 @@ class WatsonxOrchestrateA2AConfig(BaseA2AProviderConfig):
async def handle_non_streaming(
self,
request_id: str,
params: Dict[str, Any],
api_base: Optional[str] = None,
params: dict[str, Any],
api_base: str | None = None,
**kwargs: Any,
) -> Dict[str, Any]:
) -> dict[str, Any]:
"""Handle a non-streaming A2A request via WXO runs API."""
litellm_params = kwargs.get("litellm_params")
if not litellm_params:
@ -37,10 +37,10 @@ class WatsonxOrchestrateA2AConfig(BaseA2AProviderConfig):
async def handle_streaming(
self,
request_id: str,
params: Dict[str, Any],
api_base: Optional[str] = None,
params: dict[str, Any],
api_base: str | None = None,
**kwargs: Any,
) -> AsyncIterator[Dict[str, Any]]:
) -> AsyncIterator[dict[str, Any]]:
"""Handle a streaming A2A request via WXO streaming runs API."""
litellm_params = kwargs.get("litellm_params")
if not litellm_params:

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

@ -7,7 +7,7 @@ import hashlib
import json
import time
from collections.abc import AsyncIterator
from typing import Any, Dict, NamedTuple, Optional, Tuple, cast
from typing import Any, NamedTuple, cast
import httpx
@ -25,7 +25,7 @@ _IBM_CLOUD_IAM_URL = "https://iam.cloud.ibm.com/identity/token"
_POLL_INTERVAL_S = 2.0
_MAX_POLL_ATTEMPTS = 90
_TOKEN_CACHE_TTL_BUFFER_S = 60
_token_cache: Dict[str, Tuple[str, float]] = {}
_token_cache: dict[str, tuple[str, float]] = {}
class WXORequestParams(NamedTuple):
@ -33,9 +33,9 @@ class WXORequestParams(NamedTuple):
instance_id: str
wxo_agent_id: str
api_key: str
username: Optional[str]
username: str | None
auth_mode: str
thread_id: Optional[str]
thread_id: str | None
class WatsonxOrchestrateHandler:
@ -51,13 +51,13 @@ class WatsonxOrchestrateHandler:
auth_mode: str,
cp4d_host: str,
api_key: str,
username: Optional[str],
username: str | None,
) -> str:
material = f"{auth_mode}:{cp4d_host}:{username or ''}:{api_key}"
return hashlib.sha256(material.encode()).hexdigest()
@staticmethod
def _cp4d_token_ttl_seconds(expiration: Any, now_wall: Optional[float] = None) -> int:
def _cp4d_token_ttl_seconds(expiration: Any, now_wall: float | None = None) -> int:
# CP4D returns expiration as absolute Unix epoch seconds, not a duration.
expires_at = int(expiration)
wall = now_wall if now_wall is not None else time.time()
@ -68,8 +68,8 @@ class WatsonxOrchestrateHandler:
cp4d_host: str,
auth_mode: str,
api_key: str,
username: Optional[str] = None,
client: Optional[AsyncHTTPHandler] = None,
username: str | None = None,
client: AsyncHTTPHandler | None = None,
) -> str:
cache_key = WatsonxOrchestrateHandler._token_cache_key(auth_mode, cp4d_host, api_key, username)
now = time.monotonic()
@ -122,18 +122,18 @@ class WatsonxOrchestrateHandler:
async def _poll_run(
base_url: str,
run_id: str,
auth_headers: Dict[str, str],
auth_headers: dict[str, str],
client: AsyncHTTPHandler,
max_attempts: int = _MAX_POLL_ATTEMPTS,
interval_s: float = _POLL_INTERVAL_S,
) -> Dict[str, Any]:
) -> dict[str, Any]:
url = 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: dict[str, Any] = response.json()
status = result.get("status", "")
verbose_logger.debug(f"WXO: Poll {attempt + 1}/{max_attempts} run='{run_id}' status='{status}'")
if status in WatsonxOrchestrateTransformation.TERMINAL_STATES:
@ -145,11 +145,11 @@ class WatsonxOrchestrateHandler:
@staticmethod
async def _get_successful_run_data(
run_data: Dict[str, Any],
run_data: dict[str, Any],
base_url: str,
auth_headers: Dict[str, str],
auth_headers: dict[str, str],
client: AsyncHTTPHandler,
) -> Dict[str, Any]:
) -> dict[str, Any]:
status = run_data.get("status", "")
if status not in WatsonxOrchestrateTransformation.TERMINAL_STATES:
run_id = run_data.get("run_id") or run_data.get("id") or ""
@ -187,7 +187,7 @@ class WatsonxOrchestrateHandler:
return accumulated_text
@staticmethod
def _extract_litellm_params(litellm_params: Dict[str, Any]) -> WXORequestParams:
def _extract_litellm_params(litellm_params: dict[str, Any]) -> WXORequestParams:
cp4d_host = litellm_params.get("cp4d_host") or ""
instance_id = litellm_params.get("instance_id") or ""
wxo_agent_id = litellm_params.get("wxo_agent_id") or ""
@ -215,9 +215,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, Any],
litellm_params: dict[str, Any],
) -> dict[str, Any]:
wxo = WatsonxOrchestrateHandler._extract_litellm_params(litellm_params)
client = WatsonxOrchestrateHandler._http_client(timeout=90.0)
@ -246,7 +246,7 @@ class WatsonxOrchestrateHandler:
headers=auth_headers,
)
run_response.raise_for_status()
run_data: Dict[str, Any] = run_response.json()
run_data: dict[str, Any] = run_response.json()
run_data = await WatsonxOrchestrateHandler._get_successful_run_data(
run_data=run_data,
@ -261,11 +261,11 @@ class WatsonxOrchestrateHandler:
@staticmethod
async def handle_streaming(
request_id: str,
params: Dict[str, Any],
litellm_params: Dict[str, Any],
params: dict[str, Any],
litellm_params: dict[str, Any],
chunk_size: int = 50,
delay_ms: int = 10,
) -> AsyncIterator[Dict[str, Any]]:
) -> AsyncIterator[dict[str, Any]]:
wxo = WatsonxOrchestrateHandler._extract_litellm_params(litellm_params)
client = WatsonxOrchestrateHandler._http_client(timeout=120.0)

16
litellm/a2a_protocol/providers/watsonx_orchestrate/transformation.py
View file

@ -9,7 +9,7 @@ WXO uses a REST API (not A2A/JSON-RPC) with an async-poll execution model:
import asyncio
from collections.abc import AsyncIterator
from typing import Any, Dict, Optional
from typing import Any
from uuid import uuid4
from litellm._logging import verbose_logger
@ -29,7 +29,7 @@ class WatsonxOrchestrateTransformation:
return f"{cp4d_host.rstrip('/')}/orchestrate/cpd/instances/{instance_id}"
@staticmethod
def extract_text_from_a2a_params(params: Dict[str, Any]) -> str:
def extract_text_from_a2a_params(params: dict[str, Any]) -> str:
"""
Extract user message text from A2A MessageSendParams.
@ -50,10 +50,10 @@ class WatsonxOrchestrateTransformation:
def build_wxo_run_body(
wxo_agent_id: str,
text: str,
thread_id: Optional[str] = None,
) -> Dict[str, Any]:
thread_id: str | None = None,
) -> dict[str, Any]:
"""Build the WXO POST /v1/orchestrate/runs request body."""
body: Dict[str, Any] = {
body: dict[str, Any] = {
"agent_id": wxo_agent_id,
"message": {
"role": "user",
@ -103,7 +103,7 @@ class WatsonxOrchestrateTransformation:
return ""
@staticmethod
def extract_text_from_a2a_message_response(a2a_response: Dict[str, Any]) -> str:
def extract_text_from_a2a_message_response(a2a_response: dict[str, Any]) -> str:
result = a2a_response.get("result")
if not isinstance(result, dict):
verbose_logger.warning("WXO: A2A response missing result object")
@ -119,7 +119,7 @@ class WatsonxOrchestrateTransformation:
return ""
@staticmethod
def build_a2a_message_response(request_id: str, text: str) -> Dict[str, Any]:
def build_a2a_message_response(request_id: str, text: str) -> dict[str, Any]:
"""
Build a standard A2A non-streaming SendMessageResponse (kind=message).
"""
@ -140,7 +140,7 @@ class WatsonxOrchestrateTransformation:
request_id: str,
chunk_size: int = 50,
delay_ms: int = 10,
) -> AsyncIterator[Dict[str, Any]]:
) -> AsyncIterator[dict[str, Any]]:
"""
Emit standard A2A streaming events from a completed text response.

12
litellm/a2a_protocol/streaming_iterator.py
View file

@ -5,7 +5,7 @@ A2A Streaming Iterator with token tracking and logging support.
import asyncio
from collections.abc import AsyncIterator
from datetime import datetime
from typing import TYPE_CHECKING, Any, Dict, List, Optional
from typing import TYPE_CHECKING, Any
import litellm
from litellm._logging import verbose_logger
@ -38,9 +38,9 @@ class A2AStreamingIterator:
self.start_time = datetime.now()
# Collect chunks for token counting
self.chunks: List[Any] = []
self.collected_text_parts: List[str] = []
self.final_chunk: Optional[Any] = None
self.chunks: list[Any] = []
self.collected_text_parts: list[str] = []
self.final_chunk: Any | None = None
def __aiter__(self):
return self
@ -146,9 +146,9 @@ class A2AStreamingIterator:
except Exception as e:
verbose_logger.debug(f"Error in A2A streaming completion handler: {e}")
def _build_logging_result(self, usage: litellm.Usage) -> Dict[str, Any]:
def _build_logging_result(self, usage: litellm.Usage) -> dict[str, Any]:
"""Build a result dict for logging."""
result: Dict[str, Any] = {
result: dict[str, Any] = {
"id": getattr(self.request, "id", "unknown"),
"jsonrpc": "2.0",
"usage": (usage.model_dump() if hasattr(usage, "model_dump") else dict(usage)),

16
litellm/a2a_protocol/utils.py
View file

@ -2,7 +2,7 @@
Utility functions for A2A protocol.
"""
from typing import TYPE_CHECKING, Any, Dict, List, Tuple, Union
from typing import TYPE_CHECKING, Any
import litellm
from litellm._logging import verbose_logger
@ -34,7 +34,7 @@ class A2ARequestUtils:
else:
parts = getattr(message, "parts", []) or []
text_parts: List[str] = []
text_parts: list[str] = []
for part in parts:
if isinstance(part, dict):
if part.get("kind") == "text":
@ -46,7 +46,7 @@ class A2ARequestUtils:
return " ".join(text_parts)
@staticmethod
def extract_text_from_response(response_dict: Dict[str, Any]) -> str:
def extract_text_from_response(response_dict: dict[str, Any]) -> str:
"""
Extract text content from A2A response result.
@ -71,7 +71,7 @@ class A2ARequestUtils:
@staticmethod
def get_input_message_from_request(
request: "Union[SendMessageRequest, SendStreamingMessageRequest]",
request: "SendMessageRequest | SendStreamingMessageRequest",
) -> Any:
"""
Extract the input message from an A2A request.
@ -108,9 +108,9 @@ class A2ARequestUtils:
@staticmethod
def calculate_usage_from_request_response(
request: "Union[SendMessageRequest, SendStreamingMessageRequest]",
response_dict: Dict[str, Any],
) -> Tuple[int, int, int]:
request: "SendMessageRequest | SendStreamingMessageRequest",
response_dict: dict[str, Any],
) -> tuple[int, int, int]:
"""
Calculate token usage from A2A request and response.
@ -145,5 +145,5 @@ def extract_text_from_a2a_message(message: Any) -> str:
return A2ARequestUtils.extract_text_from_message(message)
def extract_text_from_a2a_response(response_dict: Dict[str, Any]) -> str:
def extract_text_from_a2a_response(response_dict: dict[str, Any]) -> str:
return A2ARequestUtils.extract_text_from_response(response_dict)

19
litellm/anthropic_beta_headers_manager.py
View file

@ -25,14 +25,13 @@ Environment Variables:
import json
import os
from importlib.resources import files
from typing import Dict, List, Optional, Set
import httpx
from litellm.litellm_core_utils.litellm_logging import verbose_logger
# Cache for the loaded configuration
_BETA_HEADERS_CONFIG: Optional[Dict] = None
_BETA_HEADERS_CONFIG: dict | None = None
class GetAnthropicBetaHeadersConfig:
@ -44,7 +43,7 @@ class GetAnthropicBetaHeadersConfig:
"""
@staticmethod
def load_local_beta_headers_config() -> Dict:
def load_local_beta_headers_config() -> dict:
"""Load the local backup beta headers config bundled with the package."""
try:
content = json.loads(
@ -159,7 +158,7 @@ def get_beta_headers_config(url: str) -> dict:
return content
def _load_beta_headers_config() -> Dict:
def _load_beta_headers_config() -> dict:
"""
Load the beta headers configuration.
Uses caching to avoid repeated fetches/file reads.
@ -183,7 +182,7 @@ def _load_beta_headers_config() -> Dict:
return _BETA_HEADERS_CONFIG
def reload_beta_headers_config() -> Dict:
def reload_beta_headers_config() -> dict:
"""
Force reload the beta headers configuration from source (remote or local).
Clears the cache and fetches fresh configuration.
@ -213,9 +212,9 @@ def get_provider_name(provider: str) -> str:
def filter_and_transform_beta_headers(
beta_headers: List[str],
beta_headers: list[str],
provider: str,
) -> List[str]:
) -> list[str]:
"""
Filter and transform beta headers based on provider's mapping configuration.
@ -240,7 +239,7 @@ def filter_and_transform_beta_headers(
# Get the header mapping for this provider
provider_mapping = config.get(provider, {})
filtered_headers: Set[str] = set()
filtered_headers: set[str] = set()
for header in beta_headers:
header = header.strip()
@ -289,7 +288,7 @@ def is_beta_header_supported(
def get_provider_beta_header(
anthropic_beta_header: str,
provider: str,
) -> Optional[str]:
) -> str | None:
"""
Get the provider-specific beta header name for a given Anthropic beta header.
@ -390,7 +389,7 @@ def update_request_with_filtered_beta(
return headers, request_data
def get_unsupported_headers(provider: str) -> List[str]:
def get_unsupported_headers(provider: str) -> list[str]:
"""
Get all beta headers that are unsupported by a provider (have null values in mapping).

4
litellm/anthropic_interface/exceptions/__init__.py
View file

@ -11,9 +11,9 @@ from .exceptions import (
)
__all__ = [
"AnthropicErrorType",
"ANTHROPIC_ERROR_TYPE_MAP",
"AnthropicErrorDetail",
"AnthropicErrorResponse",
"ANTHROPIC_ERROR_TYPE_MAP",
"AnthropicErrorType",
"AnthropicExceptionMapping",
]

9
litellm/anthropic_interface/exceptions/exception_mapping_utils.py
View file

@ -5,13 +5,12 @@ Similar to litellm/litellm_core_utils/exception_mapping_utils.py but for Anthrop
"""
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from typing import Dict, Optional
from .exceptions import AnthropicErrorResponse, AnthropicErrorType
# HTTP status code -> Anthropic error type
# Source: https://docs.anthropic.com/en/api/errors
ANTHROPIC_ERROR_TYPE_MAP: Dict[int, AnthropicErrorType] = {
ANTHROPIC_ERROR_TYPE_MAP: dict[int, AnthropicErrorType] = {
400: "invalid_request_error",
401: "authentication_error",
403: "permission_error",
@ -39,7 +38,7 @@ class AnthropicExceptionMapping:
def create_error_response(
status_code: int,
message: str,
request_id: Optional[str] = None,
request_id: str | None = None,
) -> AnthropicErrorResponse:
"""
Create an Anthropic-formatted error response dict.
@ -124,7 +123,7 @@ class AnthropicExceptionMapping:
def transform_to_anthropic_error(
status_code: int,
raw_message: str,
request_id: Optional[str] = None,
request_id: str | None = None,
) -> AnthropicErrorResponse:
"""
Transform an error message to Anthropic format.
@ -143,7 +142,7 @@ class AnthropicExceptionMapping:
AnthropicErrorResponse dict
"""
# Try to parse as JSON once
parsed: Optional[dict] = safe_json_loads(raw_message)
parsed: dict | None = safe_json_loads(raw_message)
if not isinstance(parsed, dict):
parsed = None

64
litellm/anthropic_interface/messages/__init__.py
View file

@ -11,7 +11,7 @@ This is an __init__.py file to allow the following interface
"""
from collections.abc import AsyncIterator, Coroutine, Iterator
from typing import Any, Dict, List, Optional, Union
from typing import Any
from litellm.llms.anthropic.experimental_pass_through.messages.handler import (
anthropic_messages as _async_anthropic_messages,
@ -26,21 +26,21 @@ from litellm.types.llms.anthropic_messages.anthropic_response import (
async def acreate(
max_tokens: int,
messages: List[Dict],
messages: list[dict],
model: str,
metadata: Optional[Dict] = None,
stop_sequences: Optional[List[str]] = None,
stream: Optional[bool] = False,
system: Optional[str] = None,
temperature: Optional[float] = None,
thinking: Optional[Dict] = None,
tool_choice: Optional[Dict] = None,
tools: Optional[List[Dict]] = None,
top_k: Optional[int] = None,
top_p: Optional[float] = None,
container: Optional[Dict] = None,
metadata: dict | None = None,
stop_sequences: list[str] | None = None,
stream: bool | None = False,
system: str | None = None,
temperature: float | None = None,
thinking: dict | None = None,
tool_choice: dict | None = None,
tools: list[dict] | None = None,
top_k: int | None = None,
top_p: float | None = None,
container: dict | None = None,
**kwargs,
) -> Union[AnthropicMessagesResponse, AsyncIterator]:
) -> AnthropicMessagesResponse | AsyncIterator:
"""
Async wrapper for Anthropic's messages API
@ -85,26 +85,26 @@ async def acreate(
def create(
max_tokens: int,
messages: List[Dict],
messages: list[dict],
model: str,
metadata: Optional[Dict] = None,
stop_sequences: Optional[List[str]] = None,
stream: Optional[bool] = False,
system: Optional[str] = None,
temperature: Optional[float] = None,
thinking: Optional[Dict] = None,
tool_choice: Optional[Dict] = None,
tools: Optional[List[Dict]] = None,
top_k: Optional[int] = None,
top_p: Optional[float] = None,
container: Optional[Dict] = None,
metadata: dict | None = None,
stop_sequences: list[str] | None = None,
stream: bool | None = False,
system: str | None = None,
temperature: float | None = None,
thinking: dict | None = None,
tool_choice: dict | None = None,
tools: list[dict] | None = None,
top_k: int | None = None,
top_p: float | None = None,
container: dict | None = None,
**kwargs,
) -> Union[
AnthropicMessagesResponse,
Iterator[bytes],
AsyncIterator[Any],
Coroutine[Any, Any, Union[AnthropicMessagesResponse, AsyncIterator[Any], Iterator[bytes]]],
]:
) -> (
AnthropicMessagesResponse
| Iterator[bytes]
| AsyncIterator[Any]
| Coroutine[Any, Any, AnthropicMessagesResponse | AsyncIterator[Any] | Iterator[bytes]]
):
"""
Async wrapper for Anthropic's messages API

214
litellm/assistants/main.py
View file

@ -5,7 +5,7 @@ import contextvars
import os
from collections.abc import Coroutine, Iterable
from functools import partial
from typing import Any, Dict, List, Literal, Optional, Union
from typing import Any, Literal
import httpx
from openai import AsyncOpenAI, OpenAI
@ -37,7 +37,7 @@ azure_assistants_api = AzureAssistantsAPI()
async def aget_assistants(
custom_llm_provider: Literal["openai", "azure"],
client: Optional[AsyncOpenAI] = None,
client: AsyncOpenAI | None = None,
**kwargs,
) -> AsyncCursorPage[Assistant]:
loop = asyncio.get_event_loop()
@ -74,13 +74,13 @@ async def aget_assistants(
def get_assistants(
custom_llm_provider: Literal["openai", "azure"],
client: Optional[Any] = None,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
api_version: Optional[str] = None,
client: Any | None = None,
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
**kwargs,
) -> SyncCursorPage[Assistant]:
aget_assistants: Optional[bool] = kwargs.pop("aget_assistants", None)
aget_assistants: bool | None = kwargs.pop("aget_assistants", None)
if aget_assistants is not None and not isinstance(aget_assistants, bool):
raise Exception("Invalid value passed in for aget_assistants. Only bool or None allowed")
optional_params = GenericLiteLLMParams(api_key=api_key, api_base=api_base, api_version=api_version, **kwargs)
@ -102,7 +102,7 @@ def get_assistants(
elif timeout is None:
timeout = 600.0
response: Optional[SyncCursorPage[Assistant]] = None
response: SyncCursorPage[Assistant] | None = None
if custom_llm_provider == "openai":
api_base = (
optional_params.api_base # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there
@ -148,7 +148,7 @@ def get_assistants(
) # type: ignore
extra_body = optional_params.get("extra_body", {})
azure_ad_token: Optional[str] = None
azure_ad_token: str | None = None
if extra_body is not None:
azure_ad_token = extra_body.pop("azure_ad_token", None)
else:
@ -167,9 +167,7 @@ def get_assistants(
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'get_assistants'. Only 'openai' is supported.".format(
custom_llm_provider
),
message=f"LiteLLM doesn't support {custom_llm_provider} for 'get_assistants'. Only 'openai' is supported.",
model="n/a",
llm_provider=custom_llm_provider,
response=httpx.Response(
@ -181,9 +179,7 @@ def get_assistants(
if response is None:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'get_assistants'. Only 'openai' is supported.".format(
custom_llm_provider
),
message=f"LiteLLM doesn't support {custom_llm_provider} for 'get_assistants'. Only 'openai' is supported.",
model="n/a",
llm_provider=custom_llm_provider,
response=httpx.Response(
@ -198,7 +194,7 @@ def get_assistants(
async def acreate_assistants(
custom_llm_provider: Literal["openai", "azure"],
client: Optional[AsyncOpenAI] = None,
client: AsyncOpenAI | None = None,
**kwargs,
) -> Assistant:
loop = asyncio.get_event_loop()
@ -238,22 +234,22 @@ async def acreate_assistants(
def create_assistants(
custom_llm_provider: Literal["openai", "azure"],
model: str,
name: Optional[str] = None,
description: Optional[str] = None,
instructions: Optional[str] = None,
tools: Optional[List[Dict[str, Any]]] = None,
tool_resources: Optional[Dict[str, Any]] = None,
metadata: Optional[Dict[str, str]] = None,
temperature: Optional[float] = None,
top_p: Optional[float] = None,
response_format: Optional[Union[str, Dict[str, str]]] = None,
client: Optional[Any] = None,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
api_version: Optional[str] = None,
name: str | None = None,
description: str | None = None,
instructions: str | None = None,
tools: list[dict[str, Any]] | None = None,
tool_resources: dict[str, Any] | None = None,
metadata: dict[str, str] | None = None,
temperature: float | None = None,
top_p: float | None = None,
response_format: str | dict[str, str] | None = None,
client: Any | None = None,
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
**kwargs,
) -> Union[Assistant, Coroutine[Any, Any, Assistant]]:
async_create_assistants: Optional[bool] = kwargs.pop("async_create_assistants", None)
) -> Assistant | Coroutine[Any, Any, Assistant]:
async_create_assistants: bool | None = kwargs.pop("async_create_assistants", None)
if async_create_assistants is not None and not isinstance(async_create_assistants, bool):
raise ValueError("Invalid value passed in for async_create_assistants. Only bool or None allowed")
optional_params = GenericLiteLLMParams(api_key=api_key, api_base=api_base, api_version=api_version, **kwargs)
@ -291,7 +287,7 @@ def create_assistants(
# only send params that are not None
create_assistant_data = {k: v for k, v in create_assistant_data.items() if v is not None}
response: Optional[Union[Coroutine[Any, Any, Assistant], Assistant]] = None
response: Coroutine[Any, Any, Assistant] | Assistant | None = None
if custom_llm_provider == "openai":
api_base = (
optional_params.api_base # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there
@ -338,7 +334,7 @@ def create_assistants(
) # type: ignore
extra_body = optional_params.get("extra_body", {})
azure_ad_token: Optional[str] = None
azure_ad_token: str | None = None
if extra_body is not None:
azure_ad_token = extra_body.pop("azure_ad_token", None)
else:
@ -361,9 +357,7 @@ def create_assistants(
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'create_assistants'. Only 'openai' is supported.".format(
custom_llm_provider
),
message=f"LiteLLM doesn't support {custom_llm_provider} for 'create_assistants'. Only 'openai' is supported.",
model="n/a",
llm_provider=custom_llm_provider,
response=httpx.Response(
@ -383,7 +377,7 @@ def create_assistants(
async def adelete_assistant(
custom_llm_provider: Literal["openai", "azure"],
client: Optional[AsyncOpenAI] = None,
client: AsyncOpenAI | None = None,
**kwargs,
) -> AssistantDeleted:
loop = asyncio.get_event_loop()
@ -422,17 +416,17 @@ async def adelete_assistant(
def delete_assistant(
custom_llm_provider: Literal["openai", "azure"],
assistant_id: str,
client: Optional[Any] = None,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
api_version: Optional[str] = None,
client: Any | None = None,
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
**kwargs,
) -> Union[AssistantDeleted, Coroutine[Any, Any, AssistantDeleted]]:
) -> AssistantDeleted | Coroutine[Any, Any, AssistantDeleted]:
optional_params = GenericLiteLLMParams(api_key=api_key, api_base=api_base, api_version=api_version, **kwargs)
litellm_params_dict = get_litellm_params(**kwargs)
async_delete_assistants: Optional[bool] = kwargs.pop("async_delete_assistants", None)
async_delete_assistants: bool | None = kwargs.pop("async_delete_assistants", None)
if async_delete_assistants is not None and not isinstance(async_delete_assistants, bool):
raise ValueError("Invalid value passed in for async_delete_assistants. Only bool or None allowed")
@ -452,7 +446,7 @@ def delete_assistant(
elif timeout is None:
timeout = 600.0
response: Optional[Union[AssistantDeleted, Coroutine[Any, Any, AssistantDeleted]]] = None
response: AssistantDeleted | Coroutine[Any, Any, AssistantDeleted] | None = None
if custom_llm_provider == "openai":
api_base = (
optional_params.api_base
@ -491,7 +485,7 @@ def delete_assistant(
) # type: ignore
extra_body = optional_params.get("extra_body", {})
azure_ad_token: Optional[str] = None
azure_ad_token: str | None = None
if extra_body is not None:
azure_ad_token = extra_body.pop("azure_ad_token", None)
else:
@ -514,9 +508,7 @@ def delete_assistant(
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'delete_assistant'. Only 'openai' is supported.".format(
custom_llm_provider
),
message=f"LiteLLM doesn't support {custom_llm_provider} for 'delete_assistant'. Only 'openai' is supported.",
model="n/a",
llm_provider=custom_llm_provider,
response=httpx.Response(
@ -572,10 +564,10 @@ async def acreate_thread(custom_llm_provider: Literal["openai", "azure"], **kwar
def create_thread(
custom_llm_provider: Literal["openai", "azure"],
messages: Optional[Iterable[OpenAICreateThreadParamsMessage]] = None,
metadata: Optional[dict] = None,
tool_resources: Optional[OpenAICreateThreadParamsToolResources] = None,
client: Optional[OpenAI] = None,
messages: Iterable[OpenAICreateThreadParamsMessage] | None = None,
metadata: dict | None = None,
tool_resources: OpenAICreateThreadParamsToolResources | None = None,
client: OpenAI | None = None,
**kwargs,
) -> Thread:
"""
@ -620,10 +612,10 @@ def create_thread(
elif timeout is None:
timeout = 600.0
api_base: Optional[str] = None
api_key: Optional[str] = None
api_base: str | None = None
api_key: str | None = None
response: Optional[Thread] = None
response: Thread | None = None
if custom_llm_provider == "openai":
api_base = (
optional_params.api_base # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there
@ -667,12 +659,10 @@ def create_thread(
or get_secret("AZURE_API_KEY")
) # type: ignore
api_version: Optional[str] = (
optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION")
) # type: ignore
api_version: str | None = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION") # type: ignore
extra_body = optional_params.get("extra_body", {})
azure_ad_token: Optional[str] = None
azure_ad_token: str | None = None
if extra_body is not None:
azure_ad_token = extra_body.pop("azure_ad_token", None)
else:
@ -696,9 +686,7 @@ def create_thread(
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'create_thread'. Only 'openai' is supported.".format(
custom_llm_provider
),
message=f"LiteLLM doesn't support {custom_llm_provider} for 'create_thread'. Only 'openai' is supported.",
model="n/a",
llm_provider=custom_llm_provider,
response=httpx.Response(
@ -713,7 +701,7 @@ def create_thread(
async def aget_thread(
custom_llm_provider: Literal["openai", "azure"],
thread_id: str,
client: Optional[AsyncOpenAI] = None,
client: AsyncOpenAI | None = None,
**kwargs,
) -> Thread:
loop = asyncio.get_event_loop()
@ -773,9 +761,9 @@ def get_thread(
timeout = float(timeout) # type: ignore
elif timeout is None:
timeout = 600.0
api_base: Optional[str] = None
api_key: Optional[str] = None
response: Optional[Thread] = None
api_base: str | None = None
api_key: str | None = None
response: Thread | None = None
if custom_llm_provider == "openai":
api_base = (
optional_params.api_base # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there
@ -811,9 +799,7 @@ def get_thread(
elif custom_llm_provider == "azure":
api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore
api_version: Optional[str] = (
optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION")
) # type: ignore
api_version: str | None = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION") # type: ignore
api_key = (
optional_params.api_key
@ -824,7 +810,7 @@ def get_thread(
) # type: ignore
extra_body = optional_params.get("extra_body", {})
azure_ad_token: Optional[str] = None
azure_ad_token: str | None = None
if extra_body is not None:
azure_ad_token = extra_body.pop("azure_ad_token", None)
else:
@ -847,9 +833,7 @@ def get_thread(
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'get_thread'. Only 'openai' is supported.".format(
custom_llm_provider
),
message=f"LiteLLM doesn't support {custom_llm_provider} for 'get_thread'. Only 'openai' is supported.",
model="n/a",
llm_provider=custom_llm_provider,
response=httpx.Response(
@ -869,8 +853,8 @@ async def a_add_message(
thread_id: str,
role: Literal["user", "assistant"],
content: str,
attachments: Optional[List[Attachment]] = None,
metadata: Optional[dict] = None,
attachments: list[Attachment] | None = None,
metadata: dict | None = None,
client=None,
**kwargs,
) -> OpenAIMessage:
@ -922,8 +906,8 @@ def add_message(
thread_id: str,
role: Literal["user", "assistant"],
content: str,
attachments: Optional[List[Attachment]] = None,
metadata: Optional[dict] = None,
attachments: list[Attachment] | None = None,
metadata: dict | None = None,
client=None,
**kwargs,
) -> OpenAIMessage:
@ -956,9 +940,9 @@ def add_message(
timeout = float(timeout) # type: ignore
elif timeout is None:
timeout = 600.0
api_key: Optional[str] = None
api_base: Optional[str] = None
response: Optional[OpenAIMessage] = None
api_key: str | None = None
api_base: str | None = None
response: OpenAIMessage | None = None
if custom_llm_provider == "openai":
api_base = (
optional_params.api_base # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there
@ -994,9 +978,7 @@ def add_message(
elif custom_llm_provider == "azure":
api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore
api_version: Optional[str] = (
optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION")
) # type: ignore
api_version: str | None = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION") # type: ignore
api_key = (
optional_params.api_key
@ -1007,7 +989,7 @@ def add_message(
) # type: ignore
extra_body = optional_params.get("extra_body", {})
azure_ad_token: Optional[str] = None
azure_ad_token: str | None = None
if extra_body is not None:
azure_ad_token = extra_body.pop("azure_ad_token", None)
else:
@ -1028,9 +1010,7 @@ def add_message(
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'create_thread'. Only 'openai' is supported.".format(
custom_llm_provider
),
message=f"LiteLLM doesn't support {custom_llm_provider} for 'create_thread'. Only 'openai' is supported.",
model="n/a",
llm_provider=custom_llm_provider,
response=httpx.Response(
@ -1046,7 +1026,7 @@ def add_message(
async def aget_messages(
custom_llm_provider: Literal["openai", "azure"],
thread_id: str,
client: Optional[AsyncOpenAI] = None,
client: AsyncOpenAI | None = None,
**kwargs,
) -> AsyncCursorPage[OpenAIMessage]:
loop = asyncio.get_event_loop()
@ -1091,7 +1071,7 @@ async def aget_messages(
def get_messages(
custom_llm_provider: Literal["openai", "azure"],
thread_id: str,
client: Optional[Any] = None,
client: Any | None = None,
**kwargs,
) -> SyncCursorPage[OpenAIMessage]:
aget_messages = kwargs.pop("aget_messages", None)
@ -1114,9 +1094,9 @@ def get_messages(
elif timeout is None:
timeout = 600.0
response: Optional[SyncCursorPage[OpenAIMessage]] = None
api_key: Optional[str] = None
api_base: Optional[str] = None
response: SyncCursorPage[OpenAIMessage] | None = None
api_key: str | None = None
api_base: str | None = None
if custom_llm_provider == "openai":
api_base = (
optional_params.api_base # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there
@ -1151,9 +1131,7 @@ def get_messages(
elif custom_llm_provider == "azure":
api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore
api_version: Optional[str] = (
optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION")
) # type: ignore
api_version: str | None = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION") # type: ignore
api_key = (
optional_params.api_key
@ -1164,7 +1142,7 @@ def get_messages(
) # type: ignore
extra_body = optional_params.get("extra_body", {})
azure_ad_token: Optional[str] = None
azure_ad_token: str | None = None
if extra_body is not None:
azure_ad_token = extra_body.pop("azure_ad_token", None)
else:
@ -1184,9 +1162,7 @@ def get_messages(
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'get_messages'. Only 'openai' is supported.".format(
custom_llm_provider
),
message=f"LiteLLM doesn't support {custom_llm_provider} for 'get_messages'. Only 'openai' is supported.",
model="n/a",
llm_provider=custom_llm_provider,
response=httpx.Response(
@ -1202,7 +1178,7 @@ def get_messages(
### RUNS ###
def arun_thread_stream(
*,
event_handler: Optional[AssistantEventHandler] = None,
event_handler: AssistantEventHandler | None = None,
**kwargs,
) -> AsyncAssistantStreamManager[AsyncAssistantEventHandler]:
kwargs["arun_thread"] = True
@ -1213,13 +1189,13 @@ async def arun_thread(
custom_llm_provider: Literal["openai", "azure"],
thread_id: str,
assistant_id: str,
additional_instructions: Optional[str] = None,
instructions: Optional[str] = None,
metadata: Optional[dict] = None,
model: Optional[str] = None,
stream: Optional[bool] = None,
tools: Optional[Iterable[AssistantToolParam]] = None,
client: Optional[Any] = None,
additional_instructions: str | None = None,
instructions: str | None = None,
metadata: dict | None = None,
model: str | None = None,
stream: bool | None = None,
tools: Iterable[AssistantToolParam] | None = None,
client: Any | None = None,
**kwargs,
) -> Run:
loop = asyncio.get_event_loop()
@ -1270,7 +1246,7 @@ async def arun_thread(
def run_thread_stream(
*,
event_handler: Optional[AssistantEventHandler] = None,
event_handler: AssistantEventHandler | None = None,
**kwargs,
) -> AssistantStreamManager[AssistantEventHandler]:
return run_thread(stream=True, event_handler=event_handler, **kwargs) # type: ignore
@ -1280,14 +1256,14 @@ def run_thread(
custom_llm_provider: Literal["openai", "azure"],
thread_id: str,
assistant_id: str,
additional_instructions: Optional[str] = None,
instructions: Optional[str] = None,
metadata: Optional[dict] = None,
model: Optional[str] = None,
stream: Optional[bool] = None,
tools: Optional[Iterable[AssistantToolParam]] = None,
client: Optional[Any] = None,
event_handler: Optional[AssistantEventHandler] = None, # for stream=True calls
additional_instructions: str | None = None,
instructions: str | None = None,
metadata: dict | None = None,
model: str | None = None,
stream: bool | None = None,
tools: Iterable[AssistantToolParam] | None = None,
client: Any | None = None,
event_handler: AssistantEventHandler | None = None, # for stream=True calls
**kwargs,
) -> Run:
"""Run a given thread + assistant."""
@ -1311,7 +1287,7 @@ def run_thread(
elif timeout is None:
timeout = 600.0
response: Optional[Run] = None
response: Run | None = None
if custom_llm_provider == "openai":
api_base = (
optional_params.api_base # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there
@ -1393,9 +1369,7 @@ def run_thread(
) # type: ignore
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'run_thread'. Only 'openai' is supported.".format(
custom_llm_provider
),
message=f"LiteLLM doesn't support {custom_llm_provider} for 'run_thread'. Only 'openai' is supported.",
model="n/a",
llm_provider=custom_llm_provider,
response=httpx.Response(

43
litellm/assistants/utils.py
View file

@ -1,5 +1,3 @@
from typing import Optional, Union
import litellm
from ..exceptions import UnsupportedParamsError
@ -7,21 +5,10 @@ from ..types.llms.openai import *
def get_optional_params_add_message(
role: Optional[str],
content: Optional[
Union[
str,
List[
Union[
MessageContentTextObject,
MessageContentImageFileObject,
MessageContentImageURLObject,
]
],
]
],
attachments: Optional[List[Attachment]],
metadata: Optional[dict],
role: str | None,
content: str | List[MessageContentTextObject | MessageContentImageFileObject | MessageContentImageURLObject] | None,
attachments: List[Attachment] | None,
metadata: dict | None,
custom_llm_provider: str,
**kwargs,
):
@ -56,9 +43,7 @@ def get_optional_params_add_message(
elif k not in supported_params:
raise litellm.utils.UnsupportedParamsError(
status_code=500,
message="k={}, not supported by {}. Supported params={}. To drop it from the call, set `litellm.drop_params = True`.".format(
k, custom_llm_provider, supported_params
),
message=f"k={k}, not supported by {custom_llm_provider}. Supported params={supported_params}. To drop it from the call, set `litellm.drop_params = True`.",
)
return non_default_params
@ -71,19 +56,19 @@ def get_optional_params_add_message(
non_default_params=non_default_params, optional_params=optional_params
)
for k in passed_params.keys():
if k not in default_params.keys():
if k not in default_params:
optional_params[k] = passed_params[k]
return optional_params
def get_optional_params_image_gen(
n: Optional[int] = None,
quality: Optional[str] = None,
response_format: Optional[str] = None,
size: Optional[str] = None,
style: Optional[str] = None,
user: Optional[str] = None,
custom_llm_provider: Optional[str] = None,
n: int | None = None,
quality: str | None = None,
response_format: str | None = None,
size: str | None = None,
style: str | None = None,
user: str | None = None,
custom_llm_provider: str | None = None,
**kwargs,
):
# retrieve all parameters passed to the function
@ -142,6 +127,6 @@ def get_optional_params_image_gen(
optional_params["sampleCount"] = int(n)
for k in passed_params.keys():
if k not in default_params.keys():
if k not in default_params:
optional_params[k] = passed_params[k]
return optional_params

35
litellm/batch_completion/main.py
View file

@ -1,5 +1,4 @@
from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait
from typing import List, Optional
import litellm
from litellm._logging import print_verbose
@ -11,23 +10,23 @@ from ..llms.vllm.completion import handler as vllm_handler
def batch_completion(
model: str,
# Optional OpenAI params: see https://platform.openai.com/docs/api-reference/chat/create
messages: List = [],
functions: Optional[List] = None,
function_call: Optional[str] = None,
temperature: Optional[float] = None,
top_p: Optional[float] = None,
n: Optional[int] = None,
stream: Optional[bool] = None,
messages: list = [],
functions: list | None = None,
function_call: str | None = None,
temperature: float | None = None,
top_p: float | None = None,
n: int | None = None,
stream: bool | None = None,
stop=None,
max_tokens: Optional[int] = None,
presence_penalty: Optional[float] = None,
frequency_penalty: Optional[float] = None,
logit_bias: Optional[dict] = None,
user: Optional[str] = None,
max_tokens: int | None = None,
presence_penalty: float | None = None,
frequency_penalty: float | None = None,
logit_bias: dict | None = None,
user: str | None = None,
deployment_id=None,
request_timeout: Optional[int] = None,
timeout: Optional[int] = 600,
max_workers: Optional[int] = 100,
request_timeout: int | None = None,
timeout: int | None = 600,
max_workers: int | None = 100,
# Optional liteLLM function params
**kwargs,
):
@ -164,7 +163,7 @@ def batch_completion_models(*args, **kwargs):
futures = {}
with ThreadPoolExecutor(max_workers=len(deployments)) as executor:
for deployment in deployments:
for key in kwargs.keys():
for key in kwargs:
if key not in deployment: # don't override deployment values e.g. model name, api base, etc.
deployment[key] = kwargs[key]
kwargs = {**deployment, **nested_kwargs}
@ -250,7 +249,7 @@ def batch_completion_models_all_responses(*args, **kwargs):
if result is not None:
responses.append(result)
except Exception as e:
print_verbose(f"batch_completion_models_all_responses: model request failed: {str(e)}")
print_verbose(f"batch_completion_models_all_responses: model request failed: {e!s}")
continue
return responses

42
litellm/batches/batch_utils.py
View file

@ -1,7 +1,7 @@
import json
from collections.abc import Iterable, Iterator
from dataclasses import dataclass
from typing import Any, List, Literal, Optional, Tuple
from typing import Any, Literal
import litellm
from litellm._logging import verbose_logger
@ -12,11 +12,11 @@ from litellm.utils import token_counter
async def calculate_batch_cost_and_usage(
file_content_dictionary: List[dict],
file_content_dictionary: list[dict],
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"],
model_name: Optional[str] = None,
model_info: Optional[ModelInfo] = None,
) -> Tuple[float, Usage, List[str]]:
model_name: str | None = None,
model_info: ModelInfo | None = None,
) -> tuple[float, Usage, list[str]]:
"""
Calculate the cost and usage of a batch.
@ -45,9 +45,9 @@ async def calculate_batch_cost_and_usage(
async def _handle_completed_batch(
batch: Batch,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"],
model_name: Optional[str] = None,
litellm_params: Optional[dict] = None,
) -> Tuple[float, Usage, List[str]]:
model_name: str | None = None,
litellm_params: dict | None = None,
) -> tuple[float, Usage, list[str]]:
"""Fetch a completed batch's output file and aggregate its cost, usage, and
models in a single pass over the JSONL lines, so the parsed file content is
never materialized in memory.
@ -85,14 +85,14 @@ class _BatchOutputLineStats:
total_tokens: int
cache_read_tokens: int
cache_creation_tokens: int
model: Optional[str]
model: str | None
def _iter_successful_output_line_stats(
entries: Iterable[dict],
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
model_name: Optional[str],
model_info: Optional[ModelInfo],
model_name: str | None,
model_info: ModelInfo | None,
) -> Iterator[_BatchOutputLineStats]:
from litellm.cost_calculator import batch_cost_calculator
@ -136,9 +136,9 @@ def _iter_successful_output_line_stats(
def _aggregate_batch_cost_usage_models(
entries: Iterable[dict],
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
model_name: Optional[str] = None,
model_info: Optional[ModelInfo] = None,
) -> Tuple[float, Usage, List[str]]:
model_name: str | None = None,
model_info: ModelInfo | None = None,
) -> tuple[float, Usage, list[str]]:
"""Aggregate cost, usage, and models from batch output entries in a single
pass, holding one small stats record per line instead of the parsed file."""
line_stats = tuple(_iter_successful_output_line_stats(entries, custom_llm_provider, model_name, model_info))
@ -164,9 +164,9 @@ def _aggregate_batch_cost_usage_models(
def calculate_vertex_ai_batch_cost_and_usage(
vertex_ai_batch_responses: List[dict],
model_name: Optional[str] = None,
) -> Tuple[float, Usage]:
vertex_ai_batch_responses: list[dict],
model_name: str | None = None,
) -> tuple[float, Usage]:
"""
Calculate both cost and usage from raw Vertex AI batch responses.
@ -234,7 +234,7 @@ def calculate_vertex_ai_batch_cost_and_usage(
async def _fetch_batch_output_file_content(
batch: Batch,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai",
litellm_params: Optional[dict] = None,
litellm_params: dict | None = None,
) -> bytes:
"""
Fetch the batch output file and return its raw JSONL bytes
@ -278,7 +278,7 @@ async def _fetch_batch_output_file_content(
return _file_content.content
def _extract_file_access_credentials(litellm_params: Optional[dict]) -> dict:
def _extract_file_access_credentials(litellm_params: dict | None) -> dict:
"""
Extract credentials from litellm_params for file access operations.
@ -317,7 +317,7 @@ def _extract_file_access_credentials(litellm_params: Optional[dict]) -> dict:
return credentials
def _get_file_content_as_dictionary(file_content: bytes) -> List[dict]:
def _get_file_content_as_dictionary(file_content: bytes) -> list[dict]:
"""
Get the file content as a list of dictionaries from JSON Lines format
"""
@ -367,7 +367,7 @@ def _estimate_batch_entry_tokens(raw_line: bytes) -> int:
def _count_entry_tokens(
entry: dict,
model_name: Optional[str] = None,
model_name: str | None = None,
) -> int:
"""Token-count a single batch input entry's body (chat / text / embedding)."""
body = entry.get("body", {}) or {}

104
litellm/batches/main.py
View file

@ -15,7 +15,7 @@ import contextvars
import os
from collections.abc import Coroutine
from functools import partial
from typing import Any, Dict, Literal, Optional, Union, cast
from typing import Any, Literal, cast
import httpx
from openai.types.batch import BatchRequestCounts
@ -64,7 +64,7 @@ base_llm_http_handler = BaseLLMHTTPHandler()
def _resolve_timeout(
optional_params: GenericLiteLLMParams,
kwargs: Dict[str, Any],
kwargs: dict[str, Any],
custom_llm_provider: str,
default_timeout: float = 600.0,
) -> float:
@ -107,10 +107,10 @@ async def acreate_batch(
endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions"],
input_file_id: str,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm"] = "openai",
metadata: Optional[Dict[str, str]] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
output_expires_after: Optional[Dict[str, Any]] = None,
metadata: dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
output_expires_after: dict[str, Any] | None = None,
**kwargs,
) -> LiteLLMBatch:
"""
@ -157,12 +157,12 @@ def create_batch(
endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions"],
input_file_id: str,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm"] = "openai",
metadata: Optional[Dict[str, str]] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
output_expires_after: Optional[Dict[str, Any]] = None,
metadata: dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
output_expires_after: dict[str, Any] | None = None,
**kwargs,
) -> Union[LiteLLMBatch, Coroutine[Any, Any, LiteLLMBatch]]:
) -> LiteLLMBatch | Coroutine[Any, Any, LiteLLMBatch]:
"""
Creates and executes a batch from an uploaded file of request
@ -173,7 +173,7 @@ def create_batch(
litellm_call_id = kwargs.get("litellm_call_id", None)
proxy_server_request = kwargs.get("proxy_server_request", None)
model_info = kwargs.get("model_info", None)
model: Optional[str] = kwargs.get("model", None)
model: str | None = kwargs.get("model", None)
try:
if model is not None:
model, _, _, _ = get_llm_provider(
@ -182,7 +182,7 @@ def create_batch(
)
except Exception as e:
verbose_logger.exception(
f"litellm.batches.main.py::create_batch() - Error inferring custom_llm_provider - {str(e)}"
f"litellm.batches.main.py::create_batch() - Error inferring custom_llm_provider - {e!s}"
)
_is_async = kwargs.pop("acreate_batch", False) is True
@ -238,7 +238,7 @@ def create_batch(
model=model,
)
return response
api_base: Optional[str] = None
api_base: str | None = None
if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS:
# for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there
api_base = (
@ -321,7 +321,7 @@ def create_batch(
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support custom_llm_provider={} for 'create_batch'".format(custom_llm_provider),
message=f"LiteLLM doesn't support custom_llm_provider={custom_llm_provider} for 'create_batch'",
model="n/a",
llm_provider=custom_llm_provider,
response=httpx.Response(
@ -339,9 +339,9 @@ def create_batch(
async def aretrieve_batch(
batch_id: str,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic"] = "openai",
metadata: Optional[Dict[str, str]] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
metadata: dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
**kwargs,
) -> LiteLLMBatch:
"""
@ -380,14 +380,14 @@ async def aretrieve_batch(
def _handle_retrieve_batch_providers_without_provider_config(
batch_id: str,
optional_params: GenericLiteLLMParams,
timeout: Union[float, httpx.Timeout],
timeout: float | httpx.Timeout,
litellm_params: dict,
_retrieve_batch_request: RetrieveBatchRequest,
_is_async: bool,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic"] = "openai",
logging_obj: Optional[Any] = None,
logging_obj: Any | None = None,
):
api_base: Optional[str] = None
api_base: str | None = None
if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS:
# for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there
api_base = (
@ -489,10 +489,10 @@ def _handle_retrieve_batch_providers_without_provider_config(
else:
raise litellm.exceptions.BadRequestError(
message=(
"LiteLLM doesn't support custom_llm_provider={} for 'retrieve_batch' without a `model` kwarg. "
f"LiteLLM doesn't support custom_llm_provider={custom_llm_provider} for 'retrieve_batch' without a `model` kwarg. "
"Supported via this path: 'openai', 'azure', 'vertex_ai', 'anthropic'. "
"'bedrock' is supported but requires `model` to be passed so the provider config can be loaded."
).format(custom_llm_provider),
),
model="n/a",
llm_provider=custom_llm_provider,
response=httpx.Response(
@ -508,11 +508,11 @@ def _handle_retrieve_batch_providers_without_provider_config(
def retrieve_batch(
batch_id: str,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic"] = "openai",
metadata: Optional[Dict[str, str]] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
metadata: dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
**kwargs,
) -> Union[LiteLLMBatch, Coroutine[Any, Any, LiteLLMBatch]]:
) -> LiteLLMBatch | Coroutine[Any, Any, LiteLLMBatch]:
"""
Retrieves a batch.
@ -520,7 +520,7 @@ def retrieve_batch(
"""
try:
optional_params = GenericLiteLLMParams(**kwargs)
litellm_logging_obj: Optional[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj", None)
litellm_logging_obj: LiteLLMLoggingObj | None = kwargs.get("litellm_logging_obj", None)
### TIMEOUT LOGIC ###
timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600
litellm_params = get_litellm_params(
@ -589,7 +589,7 @@ def retrieve_batch(
)
# Try to use provider config first (for providers like bedrock)
model: Optional[str] = kwargs.get("model", None)
model: str | None = kwargs.get("model", None)
if model is not None:
provider_config = ProviderConfigManager.get_provider_batches_config(
model=model,
@ -643,12 +643,12 @@ def retrieve_batch(
@client
async def alist_batches(
after: Optional[str] = None,
limit: Optional[int] = None,
after: str | None = None,
limit: int | None = None,
custom_llm_provider: ListBatchesSupportedProvider = "openai",
metadata: Optional[Dict[str, str]] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
metadata: dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
**kwargs,
):
"""
@ -686,11 +686,11 @@ async def alist_batches(
@client
def list_batches(
after: Optional[str] = None,
limit: Optional[int] = None,
after: str | None = None,
limit: int | None = None,
custom_llm_provider: ListBatchesSupportedProvider = "openai",
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
**kwargs,
):
"""
@ -823,11 +823,11 @@ def list_batches(
async def acancel_batch(
batch_id: str,
model: Optional[str] = None,
model: str | None = None,
custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai",
metadata: Optional[Dict[str, str]] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
metadata: dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
**kwargs,
) -> LiteLLMBatch:
"""
@ -869,13 +869,13 @@ async def acancel_batch(
def cancel_batch(
batch_id: str,
model: Optional[str] = None,
custom_llm_provider: Union[Literal["openai", "azure", "vertex_ai"], str] = "openai",
metadata: Optional[Dict[str, str]] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
model: str | None = None,
custom_llm_provider: Literal["openai", "azure", "vertex_ai"] | str = "openai",
metadata: dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
**kwargs,
) -> Union[LiteLLMBatch, Coroutine[Any, Any, LiteLLMBatch]]:
) -> LiteLLMBatch | Coroutine[Any, Any, LiteLLMBatch]:
"""
Cancels a batch.
@ -890,7 +890,7 @@ def cancel_batch(
)
except Exception as e:
verbose_logger.exception(
f"litellm.batches.main.py::cancel_batch() - Error inferring custom_llm_provider - {str(e)}"
f"litellm.batches.main.py::cancel_batch() - Error inferring custom_llm_provider - {e!s}"
)
optional_params = GenericLiteLLMParams(**kwargs)
litellm_params = get_litellm_params(
@ -920,7 +920,7 @@ def cancel_batch(
)
_is_async = kwargs.pop("acancel_batch", False) is True
api_base: Optional[str] = None
api_base: str | None = None
if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS:
api_base = (
optional_params.api_base
@ -993,9 +993,7 @@ def cancel_batch(
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'cancel_batch'. Only 'openai', 'azure', and 'vertex_ai' are supported.".format(
custom_llm_provider
),
message=f"LiteLLM doesn't support {custom_llm_provider} for 'cancel_batch'. Only 'openai', 'azure', and 'vertex_ai' are supported.",
model="n/a",
llm_provider=custom_llm_provider,
response=httpx.Response(

16
litellm/budget_manager.py
View file

@ -11,7 +11,7 @@ import json
import os
import threading
import time
from typing import Literal, Optional
from typing import Literal
import litellm
from litellm.constants import (
@ -28,8 +28,8 @@ class BudgetManager:
self,
project_name: str,
client_type: str = "local",
api_base: Optional[str] = None,
headers: Optional[dict] = None,
api_base: str | None = None,
headers: dict | None = None,
):
self.client_type = client_type
self.project_name = project_name
@ -73,7 +73,7 @@ class BudgetManager:
self,
total_budget: float,
user: str,
duration: Optional[Literal["daily", "weekly", "monthly", "yearly"]] = None,
duration: Literal["daily", "weekly", "monthly", "yearly"] | None = None,
created_at: float = time.time(),
):
self.user_dict[user] = {"total_budget": total_budget}
@ -113,10 +113,10 @@ class BudgetManager:
def update_cost(
self,
user: str,
completion_obj: Optional[ModelResponse] = None,
model: Optional[str] = None,
input_text: Optional[str] = None,
output_text: Optional[str] = None,
completion_obj: ModelResponse | None = None,
model: str | None = None,
input_text: str | None = None,
output_text: str | None = None,
):
if model and input_text and output_text:
prompt_tokens = litellm.token_counter(model=model, messages=[{"role": "user", "content": input_text}])

2
litellm/caching/__init__.py
View file

@ -2,10 +2,10 @@ from .azure_blob_cache import AzureBlobCache
from .caching import Cache, LiteLLMCacheType
from .disk_cache import DiskCache
from .dual_cache import DualCache
from .gcs_cache import GCSCache
from .in_memory_cache import InMemoryCache
from .qdrant_semantic_cache import QdrantSemanticCache
from .redis_cache import RedisCache
from .redis_cluster_cache import RedisClusterCache
from .redis_semantic_cache import RedisSemanticCache
from .s3_cache import S3Cache
from .gcs_cache import GCSCache

4
litellm/caching/_internal_lru_cache.py
View file

@ -1,12 +1,12 @@
from collections.abc import Callable
from functools import lru_cache
from typing import Optional, TypeVar
from typing import TypeVar
T = TypeVar("T")
def lru_cache_wrapper(
maxsize: Optional[int] = None,
maxsize: int | None = None,
) -> Callable[[Callable[..., T]], Callable[..., T]]:
"""
Wrapper for lru_cache that caches success and exceptions

2
litellm/caching/azure_blob_cache.py
View file

@ -19,12 +19,12 @@ from .base_cache import BaseCache
class AzureBlobCache(BaseCache):
def __init__(self, account_url, container) -> None:
from azure.storage.blob import BlobServiceClient
from azure.core.exceptions import ResourceExistsError
from azure.identity import DefaultAzureCredential
from azure.identity.aio import (
DefaultAzureCredential as AsyncDefaultAzureCredential,
)
from azure.storage.blob import BlobServiceClient
from azure.storage.blob.aio import BlobServiceClient as AsyncBlobServiceClient
self.container_client = BlobServiceClient(

6
litellm/caching/base_cache.py
View file

@ -9,7 +9,7 @@ Has 4 methods:
"""
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Any, Optional, Union
from typing import TYPE_CHECKING, Any, Union
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
@ -23,8 +23,8 @@ class BaseCache(ABC):
def __init__(self, default_ttl: int = 60):
self.default_ttl = default_ttl
def get_ttl(self, **kwargs) -> Optional[int]:
kwargs_ttl: Optional[int] = kwargs.get("ttl")
def get_ttl(self, **kwargs) -> int | None:
kwargs_ttl: int | None = kwargs.get("ttl")
if kwargs_ttl is not None:
try:
return int(kwargs_ttl)

159
litellm/caching/caching.py
View file

@ -13,7 +13,7 @@ import json
import time
import traceback
from enum import Enum
from typing import Any, Dict, List, Optional, Tuple, Union
from typing import Any
from pydantic import BaseModel
@ -55,19 +55,18 @@ class CacheMode(str, Enum):
class Cache:
def __init__(
self,
type: Optional[LiteLLMCacheType] = LiteLLMCacheType.LOCAL,
mode: Optional[
CacheMode
] = CacheMode.default_on, # when default_on cache is always on, when default_off cache is opt in
host: Optional[str] = None,
port: Optional[str] = None,
password: Optional[str] = None,
namespace: Optional[str] = None,
ttl: Optional[float] = None,
default_in_memory_ttl: Optional[float] = None,
default_in_redis_ttl: Optional[float] = None,
similarity_threshold: Optional[float] = None,
supported_call_types: Optional[List[CachingSupportedCallTypes]] = [
type: LiteLLMCacheType | None = LiteLLMCacheType.LOCAL,
mode: CacheMode
| None = CacheMode.default_on, # when default_on cache is always on, when default_off cache is opt in
host: str | None = None,
port: str | None = None,
password: str | None = None,
namespace: str | None = None,
ttl: float | None = None,
default_in_memory_ttl: float | None = None,
default_in_redis_ttl: float | None = None,
similarity_threshold: float | None = None,
supported_call_types: list[CachingSupportedCallTypes] | None = [
"completion",
"acompletion",
"embedding",
@ -82,38 +81,38 @@ class Cache:
"aresponses",
],
# s3 Bucket, boto3 configuration
azure_account_url: Optional[str] = None,
azure_blob_container: Optional[str] = None,
s3_bucket_name: Optional[str] = None,
s3_region_name: Optional[str] = None,
s3_api_version: Optional[str] = None,
s3_use_ssl: Optional[bool] = True,
s3_verify: Optional[Union[bool, str]] = None,
s3_endpoint_url: Optional[str] = None,
s3_aws_access_key_id: Optional[str] = None,
s3_aws_secret_access_key: Optional[str] = None,
s3_aws_session_token: Optional[str] = None,
s3_config: Optional[Any] = None,
s3_path: Optional[str] = None,
gcs_bucket_name: Optional[str] = None,
gcs_path_service_account: Optional[str] = None,
gcs_path: Optional[str] = None,
azure_account_url: str | None = None,
azure_blob_container: str | None = None,
s3_bucket_name: str | None = None,
s3_region_name: str | None = None,
s3_api_version: str | None = None,
s3_use_ssl: bool | None = True,
s3_verify: bool | str | None = None,
s3_endpoint_url: str | None = None,
s3_aws_access_key_id: str | None = None,
s3_aws_secret_access_key: str | None = None,
s3_aws_session_token: str | None = None,
s3_config: Any | None = None,
s3_path: str | None = None,
gcs_bucket_name: str | None = None,
gcs_path_service_account: str | None = None,
gcs_path: str | None = None,
redis_semantic_cache_embedding_model: str = "text-embedding-ada-002",
redis_semantic_cache_index_name: Optional[str] = None,
redis_semantic_cache_index_name: str | None = None,
valkey_semantic_cache_embedding_model: str = "text-embedding-ada-002",
valkey_semantic_cache_index_name: str | None = None,
redis_flush_size: Optional[int] = None,
redis_startup_nodes: Optional[List] = None,
disk_cache_dir: Optional[str] = None,
qdrant_api_base: Optional[str] = None,
qdrant_api_key: Optional[str] = None,
qdrant_collection_name: Optional[str] = None,
qdrant_quantization_config: Optional[str] = None,
redis_flush_size: int | None = None,
redis_startup_nodes: list | None = None,
disk_cache_dir: str | None = None,
qdrant_api_base: str | None = None,
qdrant_api_key: str | None = None,
qdrant_collection_name: str | None = None,
qdrant_quantization_config: str | None = None,
qdrant_semantic_cache_embedding_model: str = "text-embedding-ada-002",
qdrant_semantic_cache_vector_size: Optional[int] = None,
qdrant_semantic_cache_vector_size: int | None = None,
# GCP IAM authentication parameters
gcp_service_account: Optional[str] = None,
gcp_ssl_ca_certs: Optional[str] = None,
gcp_service_account: str | None = None,
gcp_ssl_ca_certs: str | None = None,
**kwargs,
):
"""
@ -352,15 +351,15 @@ class Cache:
if param in scope_excluded_params:
continue
if param in combined_kwargs:
param_value: Optional[str] = self._get_param_value(param, kwargs)
param_value: str | None = self._get_param_value(param, kwargs)
if param_value is not None:
cache_key += f"{str(param)}: {str(param_value)}"
cache_key += f"{param!s}: {param_value!s}"
elif param not in litellm_param_kwargs: # check if user passed in optional param - e.g. top_k
if litellm.enable_caching_on_provider_specific_optional_params is True: # feature flagged for now
if kwargs[param] is None:
continue # ignore None params
param_value = kwargs[param]
cache_key += f"{str(param)}: {str(param_value)}"
cache_key += f"{param!s}: {param_value!s}"
if is_semantic_cache:
cache_key += self._get_semantic_cache_tenant_scope(kwargs)
@ -382,7 +381,7 @@ class Cache:
self,
param: str,
kwargs: dict,
) -> Optional[str]:
) -> str | None:
"""
Get the value for the given param from kwargs
"""
@ -400,15 +399,15 @@ class Cache:
2. Else if a model_group is set, then return the model_group as the model. This is used for all requests sent through the litellm.Router()
3. Else use the `model` passed in kwargs
"""
metadata: Dict = kwargs.get("metadata", {}) or {}
litellm_params: Dict = kwargs.get("litellm_params", {}) or {}
metadata_in_litellm_params: Dict = litellm_params.get("metadata", {}) or {}
model_group: Optional[str] = metadata.get("model_group") or metadata_in_litellm_params.get("model_group")
metadata: dict = kwargs.get("metadata", {}) or {}
litellm_params: dict = kwargs.get("litellm_params", {}) or {}
metadata_in_litellm_params: dict = litellm_params.get("metadata", {}) or {}
model_group: str | None = metadata.get("model_group") or metadata_in_litellm_params.get("model_group")
caching_group = self._get_caching_group(metadata, model_group)
return caching_group or model_group or kwargs["model"]
def _get_caching_group(self, metadata: dict, model_group: Optional[str]) -> Optional[str]:
caching_groups: Optional[List] = metadata.get("caching_groups", [])
def _get_caching_group(self, metadata: dict, model_group: str | None) -> str | None:
caching_groups: list | None = metadata.get("caching_groups", [])
if caching_groups:
for group in caching_groups:
if model_group in group:
@ -429,7 +428,7 @@ class Cache:
or litellm_params.get("file_name")
)
def _get_preset_cache_key_from_kwargs(self, **kwargs) -> Optional[str]:
def _get_preset_cache_key_from_kwargs(self, **kwargs) -> str | None:
"""
Get the preset cache key from kwargs["litellm_params"]
@ -510,8 +509,8 @@ class Cache:
def _get_cache_logic(
self,
cached_result: Optional[Any],
max_age: Optional[float],
cached_result: Any | None,
max_age: float | None,
):
"""
Common get cache logic across sync + async implementations
@ -544,8 +543,8 @@ class Cache:
return cached_result
@staticmethod
def _get_safe_cache_lookup_kwargs(kwargs: Dict[str, Any]) -> Dict[str, Any]:
cache_lookup_kwargs: Dict[str, Any] = {}
def _get_safe_cache_lookup_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]:
cache_lookup_kwargs: dict[str, Any] = {}
for prompt_kwarg in ("messages", "input"):
if prompt_kwarg in kwargs:
cache_lookup_kwargs[prompt_kwarg] = kwargs[prompt_kwarg]
@ -558,7 +557,7 @@ class Cache:
@staticmethod
def _update_metadata_from_cache_lookup_kwargs(
original_kwargs: Dict[str, Any], cache_lookup_kwargs: Dict[str, Any]
original_kwargs: dict[str, Any], cache_lookup_kwargs: dict[str, Any]
) -> None:
original_metadata = original_kwargs.get("metadata")
cache_lookup_metadata = cache_lookup_kwargs.get("metadata")
@ -568,7 +567,7 @@ class Cache:
if "semantic-similarity" in cache_lookup_metadata:
original_metadata["semantic-similarity"] = cache_lookup_metadata["semantic-similarity"]
def get_cache(self, dynamic_cache_object: Optional[BaseCache] = None, **kwargs):
def get_cache(self, dynamic_cache_object: BaseCache | None = None, **kwargs):
"""
Retrieves the cached result for the given arguments.
@ -603,7 +602,7 @@ class Cache:
print_verbose(f"An exception occurred: {traceback.format_exc()}")
return None
async def async_get_cache(self, dynamic_cache_object: Optional[BaseCache] = None, **kwargs):
async def async_get_cache(self, dynamic_cache_object: BaseCache | None = None, **kwargs):
"""
Async get cache implementation.
@ -677,9 +676,9 @@ class Cache:
cache_key, cached_data, kwargs = self._add_cache_logic(result=result, **kwargs)
self.cache.set_cache(cache_key, cached_data, **kwargs)
except Exception as e:
verbose_logger.exception(f"LiteLLM Cache: Excepton add_cache: {str(e)}")
verbose_logger.exception(f"LiteLLM Cache: Excepton add_cache: {e!s}")
async def async_add_cache(self, result, dynamic_cache_object: Optional[BaseCache] = None, **kwargs):
async def async_add_cache(self, result, dynamic_cache_object: BaseCache | None = None, **kwargs):
"""
Async implementation of add_cache
"""
@ -696,14 +695,14 @@ class Cache:
else:
await self.cache.async_set_cache(cache_key, cached_data, **kwargs)
except Exception as e:
verbose_logger.exception(f"LiteLLM Cache: Excepton add_cache: {str(e)}")
verbose_logger.exception(f"LiteLLM Cache: Excepton add_cache: {e!s}")
def _convert_to_cached_embedding(
self,
embedding_response: Any,
model: Optional[str],
prompt_tokens: Optional[int] = None,
prompt_tokens_details: Optional[dict] = None,
model: str | None,
prompt_tokens: int | None = None,
prompt_tokens_details: dict | None = None,
) -> CachedEmbedding:
"""
Convert any embedding response into the standardized CachedEmbedding TypedDict format.
@ -745,7 +744,7 @@ class Cache:
self,
result: EmbeddingResponse,
idx_in_result_data: int,
) -> Optional[dict]:
) -> dict | None:
"""
Extract per-item prompt_tokens_details from a response for caching.
@ -788,7 +787,7 @@ class Cache:
self,
result: EmbeddingResponse,
idx_in_result_data: int,
) -> Optional[int]:
) -> int | None:
"""
Extract the per-item prompt_tokens from a response for caching.
@ -813,7 +812,7 @@ class Cache:
input: str,
kwargs: dict,
idx_in_result_data: int = 0,
) -> Tuple[str, dict, dict]:
) -> tuple[str, dict, dict]:
preset_cache_key = self.get_cache_key(**{**kwargs, "input": input})
kwargs["cache_key"] = preset_cache_key
embedding_response = result.data[idx_in_result_data]
@ -843,7 +842,7 @@ class Cache:
)
return cache_key, cached_data, kwargs
async def async_add_cache_pipeline(self, result, dynamic_cache_object: Optional[BaseCache] = None, **kwargs):
async def async_add_cache_pipeline(self, result, dynamic_cache_object: BaseCache | None = None, **kwargs):
"""
Async implementation of add_cache for Embedding calls
@ -875,7 +874,7 @@ class Cache:
else:
await self.cache.async_set_cache_pipeline(cache_list=cache_list, **kwargs)
except Exception as e:
verbose_logger.exception(f"LiteLLM Cache: Excepton add_cache: {str(e)}")
verbose_logger.exception(f"LiteLLM Cache: Excepton add_cache: {e!s}")
def should_use_cache(self, **kwargs):
"""
@ -926,11 +925,11 @@ class Cache:
def enable_cache(
type: Optional[LiteLLMCacheType] = LiteLLMCacheType.LOCAL,
host: Optional[str] = None,
port: Optional[str] = None,
password: Optional[str] = None,
supported_call_types: Optional[List[CachingSupportedCallTypes]] = [
type: LiteLLMCacheType | None = LiteLLMCacheType.LOCAL,
host: str | None = None,
port: str | None = None,
password: str | None = None,
supported_call_types: list[CachingSupportedCallTypes] | None = [
"completion",
"acompletion",
"embedding",
@ -986,11 +985,11 @@ def enable_cache(
def update_cache(
type: Optional[LiteLLMCacheType] = LiteLLMCacheType.LOCAL,
host: Optional[str] = None,
port: Optional[str] = None,
password: Optional[str] = None,
supported_call_types: Optional[List[CachingSupportedCallTypes]] = [
type: LiteLLMCacheType | None = LiteLLMCacheType.LOCAL,
host: str | None = None,
port: str | None = None,
password: str | None = None,
supported_call_types: list[CachingSupportedCallTypes] | None = [
"completion",
"acompletion",
"embedding",

116
litellm/caching/caching_handler.py
View file

@ -22,11 +22,7 @@ from collections.abc import AsyncGenerator, Callable, Generator
from typing import (
TYPE_CHECKING,
Any,
Dict,
List,
Optional,
Tuple,
Union,
)
from pydantic import BaseModel
@ -75,8 +71,8 @@ class CachingHandlerResponse(BaseModel):
For embeddings there can be a cache hit for some of the inputs in the list and a cache miss for others
"""
cached_result: Optional[Any] = None
final_embedding_cached_response: Optional[EmbeddingResponse] = None
cached_result: Any | None = None
final_embedding_cached_response: EmbeddingResponse | None = None
embedding_all_elements_cache_hit: bool = False # this is set to True when all elements in the list have a cache hit in the embedding cache, if true return the final_embedding_cached_response no need to make an API call
@ -109,7 +105,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, Any]) -> bool:
"""
When stream=True, do not run success callbacks at cache-hit time.
@ -125,25 +121,24 @@ class LLMCachingHandler:
def __init__(
self,
original_function: Callable,
request_kwargs: Dict[str, Any],
request_kwargs: dict[str, Any],
start_time: datetime.datetime,
):
from litellm.caching import DualCache, RedisCache
self.async_streaming_chunks: List[ModelResponse] = []
self.sync_streaming_chunks: List[ModelResponse] = []
self.async_streaming_chunks: list[ModelResponse] = []
self.sync_streaming_chunks: list[ModelResponse] = []
self.request_kwargs = _drop_logging_obj_from_kwargs(request_kwargs)
self.preset_cache_key: Optional[str] = None
self.preset_cache_key: str | None = None
self.original_function = original_function
self.start_time = start_time
if litellm.cache is not None and isinstance(litellm.cache.cache, RedisCache):
self.dual_cache: Optional[DualCache] = DualCache(
self.dual_cache: DualCache | None = DualCache(
redis_cache=litellm.cache.cache,
in_memory_cache=in_memory_cache_obj,
)
else:
self.dual_cache = None
pass
async def _async_get_cache(
self,
@ -152,9 +147,9 @@ class LLMCachingHandler:
logging_obj: LiteLLMLoggingObj,
start_time: datetime.datetime,
call_type: str,
kwargs: Dict[str, Any],
args: Optional[Tuple[Any, ...]] = None,
) -> Optional[CachingHandlerResponse]:
kwargs: dict[str, Any],
args: tuple[Any, ...] | None = None,
) -> CachingHandlerResponse | None:
"""
Internal method to get from the cache.
Handles different call types (embeddings, chat/completions, text_completion, transcription)
@ -182,15 +177,15 @@ class LLMCachingHandler:
kwargs.get("cache", {}).get("no-cache", False) is not True
): # allow users to control returning cached responses from the completion function
args = args or ()
final_embedding_cached_response: Optional[EmbeddingResponse] = None
final_embedding_cached_response: EmbeddingResponse | None = None
embedding_all_elements_cache_hit: bool = False
cached_result: Optional[Any] = None
cached_result: Any | None = None
kwargs = kwargs.copy()
#########################################################
# Init cache timing metrics
#########################################################
cache_check_start_time = time.perf_counter()
cache_check_end_time: Optional[float] = None
cache_check_end_time: float | None = None
#########################################################
parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs)
kwargs["parent_otel_span"] = parent_otel_span
@ -291,10 +286,10 @@ class LLMCachingHandler:
logging_obj: LiteLLMLoggingObj,
start_time: datetime.datetime,
call_type: str,
kwargs: Dict[str, Any],
args: Optional[Tuple[Any, ...]] = None,
kwargs: dict[str, Any],
args: tuple[Any, ...] | None = None,
) -> CachingHandlerResponse:
cached_result: Optional[Any] = None
cached_result: Any | None = None
# Check if caching should be performed BEFORE doing expensive kwargs copy
if litellm.cache is not None and self._is_call_type_supported_by_cache(original_function=original_function):
@ -369,7 +364,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, Any]) -> list[str]:
"""
Handles the input of kwargs['input'] being a list or a string
"""
@ -380,7 +375,7 @@ class LLMCachingHandler:
else:
raise ValueError("input must be a string or a list")
def _extract_model_from_cached_results(self, non_null_list: List[Tuple[int, CachedEmbedding]]) -> Optional[str]:
def _extract_model_from_cached_results(self, non_null_list: list[tuple[int, CachedEmbedding]]) -> str | None:
"""
Helper method to extract the model name from cached results.
@ -397,13 +392,13 @@ class LLMCachingHandler:
def _process_async_embedding_cached_response(
self,
final_embedding_cached_response: Optional[EmbeddingResponse],
cached_result: List[Optional[CachedEmbedding]],
kwargs: Dict[str, Any],
final_embedding_cached_response: EmbeddingResponse | None,
cached_result: list[CachedEmbedding | None],
kwargs: dict[str, Any],
logging_obj: LiteLLMLoggingObj,
start_time: datetime.datetime,
model: str,
) -> Tuple[Optional[EmbeddingResponse], bool]:
) -> tuple[EmbeddingResponse | None, bool]:
"""
Returns the final embedding cached response and a boolean indicating if all elements in the list have a cache hit
@ -446,7 +441,7 @@ class LLMCachingHandler:
final_embedding_cached_response._hidden_params["cache_hit"] = True
prompt_tokens = 0
aggregated_details: Optional[dict] = None
aggregated_details: dict | None = None
for val in non_null_list:
idx, cr = val # (idx, cr) tuple
if cr is not None:
@ -476,10 +471,10 @@ class LLMCachingHandler:
aggregated_details[key] = value
## USAGE
prompt_tokens_details: Optional["PromptTokensDetailsWrapper"] = None
if aggregated_details:
from litellm.types.utils import PromptTokensDetailsWrapper
from litellm.types.utils import PromptTokensDetailsWrapper
prompt_tokens_details: PromptTokensDetailsWrapper | None = None
if aggregated_details:
try:
prompt_tokens_details = PromptTokensDetailsWrapper(**aggregated_details)
except Exception:
@ -674,9 +669,7 @@ class LLMCachingHandler:
cache_hit=cache_hit,
)
async def _retrieve_from_cache(
self, call_type: str, kwargs: Dict[str, Any], args: Tuple[Any, ...]
) -> Optional[Any]:
async def _retrieve_from_cache(self, call_type: str, kwargs: dict[str, Any], args: tuple[Any, ...]) -> Any | None:
"""
Internal method to
- get cache key
@ -709,7 +702,7 @@ class LLMCachingHandler:
if new_kwargs.get("stream") is True and "cache_key" not in new_kwargs:
new_kwargs["cache_key"] = litellm.cache.get_cache_key(**new_kwargs)
self.request_kwargs = _drop_logging_obj_from_kwargs(new_kwargs)
cached_result: Optional[Any] = None
cached_result: Any | None = None
if call_type == CallTypes.aembedding.value:
if isinstance(new_kwargs["input"], str):
new_kwargs["input"] = [new_kwargs["input"]]
@ -754,21 +747,20 @@ class LLMCachingHandler:
self,
cached_result: Any,
call_type: str,
kwargs: Dict[str, Any],
kwargs: dict[str, Any],
logging_obj: LiteLLMLoggingObj,
model: str,
args: Tuple[Any, ...],
custom_llm_provider: Optional[str] = None,
) -> Optional[
Union[
ModelResponse,
TextCompletionResponse,
EmbeddingResponse,
RerankResponse,
TranscriptionResponse,
CustomStreamWrapper,
]
]:
args: tuple[Any, ...],
custom_llm_provider: str | None = None,
) -> (
ModelResponse
| TextCompletionResponse
| EmbeddingResponse
| RerankResponse
| TranscriptionResponse
| CustomStreamWrapper
| None
):
"""
Internal method to process the cached result
@ -921,7 +913,7 @@ class LLMCachingHandler:
convert_to_streaming_response_async,
)
_stream_cached_result: Union[AsyncGenerator, Generator]
_stream_cached_result: AsyncGenerator | Generator
if call_type == CallTypes.acompletion.value or call_type == CallTypes.atext_completion.value:
_stream_cached_result = convert_to_streaming_response_async(
response_object=cached_result,
@ -941,8 +933,8 @@ class LLMCachingHandler:
self,
result: Any,
original_function: Callable,
kwargs: Dict[str, Any],
args: Optional[Tuple[Any, ...]] = None,
kwargs: dict[str, Any],
args: tuple[Any, ...] | None = None,
):
"""
Internal method to check the type of the result & cache used and adds the result to the cache accordingly
@ -1007,8 +999,8 @@ class LLMCachingHandler:
def sync_set_cache(
self,
result: Any,
kwargs: Dict[str, Any],
args: Optional[Tuple[Any, ...]] = None,
kwargs: dict[str, Any],
args: tuple[Any, ...] | None = None,
):
"""
Sync internal method to add the result to the cache
@ -1029,7 +1021,7 @@ class LLMCachingHandler:
return
def _should_store_result_in_cache(self, original_function: Callable, kwargs: Dict[str, Any]) -> bool:
def _should_store_result_in_cache(self, original_function: Callable, kwargs: dict[str, Any]) -> bool:
"""
Helper function to determine if the result should be stored in the cache.
@ -1075,7 +1067,7 @@ class LLMCachingHandler:
"""
complete_streaming_response: Optional[Union[ModelResponse, TextCompletionResponse]] = (
complete_streaming_response: ModelResponse | TextCompletionResponse | None = (
_assemble_complete_response_from_streaming_chunks(
result=processed_chunk,
start_time=self.start_time,
@ -1097,7 +1089,7 @@ class LLMCachingHandler:
"""
Sync internal method to add the streaming response to the cache
"""
complete_streaming_response: Optional[Union[ModelResponse, TextCompletionResponse]] = (
complete_streaming_response: ModelResponse | TextCompletionResponse | None = (
_assemble_complete_response_from_streaming_chunks(
result=processed_chunk,
start_time=self.start_time,
@ -1119,12 +1111,12 @@ class LLMCachingHandler:
self,
logging_obj: LiteLLMLoggingObj,
model: str,
kwargs: Dict[str, Any],
kwargs: dict[str, Any],
cached_result: Any,
is_async: bool,
is_embedding: bool = False,
custom_llm_provider: Optional[str] = None,
cache_duration_ms: Optional[float] = None,
custom_llm_provider: str | None = None,
cache_duration_ms: float | None = None,
):
"""
Helper function to update the LiteLLMLoggingObj environment variables.
@ -1178,8 +1170,8 @@ class LLMCachingHandler:
def convert_args_to_kwargs(
original_function: Callable,
args: Optional[Tuple[Any, ...]] = None,
) -> Dict[str, Any]:
args: tuple[Any, ...] | None = None,
) -> dict[str, Any]:
# Get the signature of the original function
signature = inspect.signature(original_function)

4
litellm/caching/disk_cache.py
View file

@ -1,5 +1,5 @@
import json
from typing import TYPE_CHECKING, Any, Optional, Union
from typing import TYPE_CHECKING, Any, Union
from .base_cache import BaseCache
@ -12,7 +12,7 @@ else:
class DiskCache(BaseCache):
def __init__(self, disk_cache_dir: Optional[str] = None):
def __init__(self, disk_cache_dir: str | None = None):
try:
import diskcache as dc
except ModuleNotFoundError as e:

64
litellm/caching/dual_cache.py
View file

@ -13,7 +13,7 @@ import time
import traceback
from concurrent.futures import ThreadPoolExecutor
from threading import Lock
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
from typing import TYPE_CHECKING, Any, Union
if TYPE_CHECKING:
from litellm.types.caching import RedisPipelineIncrementOperation
@ -57,11 +57,11 @@ class DualCache(BaseCache):
def __init__(
self,
in_memory_cache: Optional[InMemoryCache] = None,
redis_cache: Optional[RedisCache] = None,
default_in_memory_ttl: Optional[float] = None,
default_redis_ttl: Optional[float] = None,
default_redis_batch_cache_expiry: Optional[float] = None,
in_memory_cache: InMemoryCache | None = None,
redis_cache: RedisCache | None = None,
default_in_memory_ttl: float | None = None,
default_redis_ttl: float | None = None,
default_redis_batch_cache_expiry: float | None = None,
default_max_redis_batch_cache_size: int = DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE,
) -> None:
super().__init__()
@ -77,7 +77,7 @@ class DualCache(BaseCache):
self.default_in_memory_ttl = default_in_memory_ttl or litellm.default_in_memory_ttl
self.default_redis_ttl = default_redis_ttl or litellm.default_redis_ttl
def update_cache_ttl(self, default_in_memory_ttl: Optional[float], default_redis_ttl: Optional[float]):
def update_cache_ttl(self, default_in_memory_ttl: float | None, default_redis_ttl: float | None):
if default_in_memory_ttl is not None:
self.default_in_memory_ttl = default_in_memory_ttl
@ -86,9 +86,9 @@ class DualCache(BaseCache):
def attach_redis_cache(
self,
redis_cache: Optional[RedisCache] = None,
redis_cache: RedisCache | None = None,
*,
default_redis_ttl: Optional[float] = None,
default_redis_ttl: float | None = None,
) -> None:
"""
Attach a Redis backend if this DualCache does not already have one.
@ -147,13 +147,13 @@ class DualCache(BaseCache):
return result
except Exception as e:
verbose_logger.error(f"LiteLLM Cache: Excepton async add_cache: {str(e)}")
verbose_logger.error(f"LiteLLM Cache: Excepton async add_cache: {e!s}")
raise e
def get_cache(
self,
key,
parent_otel_span: Optional[Span] = None,
parent_otel_span: Span | None = None,
local_only: bool = False,
**kwargs,
):
@ -184,7 +184,7 @@ class DualCache(BaseCache):
def batch_get_cache(
self,
keys: list,
parent_otel_span: Optional[Span] = None,
parent_otel_span: Span | None = None,
local_only: bool = False,
**kwargs,
):
@ -217,7 +217,7 @@ class DualCache(BaseCache):
async def async_get_cache(
self,
key,
parent_otel_span: Optional[Span] = None,
parent_otel_span: Span | None = None,
local_only: bool = False,
**kwargs,
):
@ -250,15 +250,15 @@ class DualCache(BaseCache):
def _reserve_redis_batch_keys(
self,
current_time: float,
keys: List[str],
result: List[Any],
) -> Tuple[List[str], Dict[str, Optional[float]]]:
keys: list[str],
result: list[Any],
) -> tuple[list[str], dict[str, float | None]]:
"""
Atomically choose keys to fetch from Redis and reserve their access time.
This prevents check-then-act races under concurrent async callers.
"""
sublist_keys: List[str] = []
previous_access_times: Dict[str, Optional[float]] = {}
sublist_keys: list[str] = []
previous_access_times: dict[str, float | None] = {}
with self._last_redis_batch_access_time_lock:
for key, value in zip(keys, result):
@ -275,7 +275,7 @@ class DualCache(BaseCache):
return sublist_keys, previous_access_times
def _rollback_redis_batch_key_reservations(self, previous_access_times: Dict[str, Optional[float]]) -> None:
def _rollback_redis_batch_key_reservations(self, previous_access_times: dict[str, float | None]) -> None:
with self._last_redis_batch_access_time_lock:
for key, previous_time in previous_access_times.items():
if previous_time is None:
@ -286,7 +286,7 @@ class DualCache(BaseCache):
async def async_batch_get_cache(
self,
keys: list,
parent_otel_span: Optional[Span] = None,
parent_otel_span: Span | None = None,
local_only: bool = False,
**kwargs,
):
@ -347,7 +347,7 @@ class DualCache(BaseCache):
if self.redis_cache is not None and local_only is False:
await self.redis_cache.async_set_cache(key, value, **kwargs)
except Exception as e:
verbose_logger.exception(f"LiteLLM Cache: Excepton async add_cache: {str(e)}")
verbose_logger.exception(f"LiteLLM Cache: Excepton async add_cache: {e!s}")
# async_batch_set_cache
async def async_set_cache_pipeline(self, cache_list: list, local_only: bool = False, **kwargs):
@ -366,17 +366,17 @@ class DualCache(BaseCache):
cache_list=cache_list, ttl=kwargs.pop("ttl", None), **kwargs
)
except Exception as e:
verbose_logger.exception(f"LiteLLM Cache: Excepton async add_cache: {str(e)}")
verbose_logger.exception(f"LiteLLM Cache: Excepton async add_cache: {e!s}")
async def async_increment_cache(
self,
key,
value: float,
parent_otel_span: Optional[Span] = None,
parent_otel_span: Span | None = None,
local_only: bool = False,
refresh_ttl: bool = False,
**kwargs,
) -> Optional[float]:
) -> float | None:
"""
Key - the key in cache
@ -388,7 +388,7 @@ class DualCache(BaseCache):
Returns - the incremented value, or None if no cache backend is
available (in_memory_cache is None and Redis failed/is absent).
"""
result: Optional[float] = None
result: float | None = None
try:
if self.in_memory_cache is not None:
result = await self.in_memory_cache.async_increment(key, value, **kwargs)
@ -412,12 +412,12 @@ class DualCache(BaseCache):
async def async_increment_cache_pipeline(
self,
increment_list: List["RedisPipelineIncrementOperation"],
increment_list: list["RedisPipelineIncrementOperation"],
local_only: bool = False,
parent_otel_span: Optional[Span] = None,
parent_otel_span: Span | None = None,
**kwargs,
) -> Optional[List[float]]:
result: Optional[List[float]] = None
) -> list[float] | None:
result: list[float] | None = None
try:
if self.in_memory_cache is not None:
result = await self.in_memory_cache.async_increment_pipeline(
@ -439,7 +439,7 @@ class DualCache(BaseCache):
)
return result
async def async_set_cache_sadd(self, key, value: List, local_only: bool = False, **kwargs) -> None:
async def async_set_cache_sadd(self, key, value: list, local_only: bool = False, **kwargs) -> None:
"""
Add value to a set
@ -456,7 +456,7 @@ class DualCache(BaseCache):
if self.redis_cache is not None and local_only is False:
_ = await self.redis_cache.async_set_cache_sadd(key, value, ttl=kwargs.get("ttl", None))
return None
return
except Exception as e:
raise e # don't log, if exception is raised
@ -484,7 +484,7 @@ class DualCache(BaseCache):
if self.redis_cache is not None:
await self.redis_cache.async_delete_cache(key)
async def async_get_ttl(self, key: str) -> Optional[int]:
async def async_get_ttl(self, key: str) -> int | None:
"""
Get the remaining TTL of a key in in-memory cache or redis
"""

12
litellm/caching/gcs_cache.py
View file

@ -2,27 +2,27 @@
Supports syncing responses to Google Cloud Storage Buckets using HTTP requests.
"""
import json
import asyncio
from typing import Optional
import json
from urllib.parse import quote
from litellm._logging import print_verbose, verbose_logger
from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
_get_httpx_client,
get_async_httpx_client,
httpxSpecialProvider,
)
from .base_cache import BaseCache
class GCSCache(BaseCache):
def __init__(
self,
bucket_name: Optional[str] = None,
path_service_account: Optional[str] = None,
gcs_path: Optional[str] = None,
bucket_name: str | None = None,
path_service_account: str | None = None,
gcs_path: str | None = None,
) -> None:
super().__init__()
self.bucket_name = bucket_name or GCSBucketBase(bucket_name=None).BUCKET_NAME

29
litellm/caching/in_memory_cache.py
View file

@ -8,12 +8,12 @@ Has 4 methods:
- async_get_cache
"""
import heapq
import json
import sys
import time
import heapq
import threading
from typing import TYPE_CHECKING, Any, List, Optional
import time
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from litellm.types.caching import RedisPipelineIncrementOperation
@ -28,11 +28,10 @@ from .base_cache import BaseCache
class InMemoryCache(BaseCache):
def __init__(
self,
max_size_in_memory: Optional[int] = 200,
default_ttl: Optional[
int
] = 600, # default ttl is 10 minutes. At maximum litellm rate limiting logic requires objects to be in memory for 1 minute
max_size_per_item: Optional[int] = 1024, # 1MB = 1024KB
max_size_in_memory: int | None = 200,
default_ttl: int
| None = 600, # default ttl is 10 minutes. At maximum litellm rate limiting logic requires objects to be in memory for 1 minute
max_size_per_item: int | None = 1024, # 1MB = 1024KB
):
"""
max_size_in_memory [int]: Maximum number of items in cache. done to prevent memory leaks. Use 200 items as a default
@ -146,9 +145,7 @@ class InMemoryCache(BaseCache):
Check if ttl is set for a key
"""
ttl_time = self.ttl_dict.get(key)
if ttl_time is None: # if ttl is not set, allow override
return True
elif float(ttl_time) < time.time(): # if ttl is expired, allow override
if ttl_time is None or float(ttl_time) < time.time(): # if ttl is not set, allow override
return True
else:
return False
@ -184,7 +181,7 @@ class InMemoryCache(BaseCache):
else:
self.set_cache(key=cache_key, value=cache_value)
async def async_set_cache_sadd(self, key, value: List, ttl: Optional[float]):
async def async_set_cache_sadd(self, key, value: list, ttl: float | None):
"""
Add value to set
"""
@ -247,8 +244,8 @@ class InMemoryCache(BaseCache):
return self.increment_cache(key=key, value=value, **kwargs)
async def async_increment_pipeline(
self, increment_list: List["RedisPipelineIncrementOperation"], **kwargs
) -> Optional[List[float]]:
self, increment_list: list["RedisPipelineIncrementOperation"], **kwargs
) -> list[float] | None:
results = []
for increment in increment_list:
result = await self.async_increment(increment["key"], increment["increment_value"], **kwargs)
@ -266,13 +263,13 @@ class InMemoryCache(BaseCache):
def delete_cache(self, key):
self._remove_key(key)
async def async_get_ttl(self, key: str) -> Optional[int]:
async def async_get_ttl(self, key: str) -> int | None:
"""
Get the remaining TTL of a key in in-memory cache
"""
return self.ttl_dict.get(key, None)
async def async_get_oldest_n_keys(self, n: int) -> List[str]:
async def async_get_oldest_n_keys(self, n: int) -> list[str]:
"""
Get the oldest n keys in the cache
"""

14
litellm/caching/qdrant_semantic_cache.py
View file

@ -12,7 +12,7 @@ import ast
import asyncio
import json
import os
from typing import Any, Dict, cast
from typing import Any, cast
import litellm
from litellm._logging import print_verbose
@ -104,7 +104,7 @@ class QdrantSemanticCache(BaseCache):
print_verbose(f"Collection already exists.\nCollection details:{self.collection_info}")
self._ensure_cache_key_payload_index()
else:
quantization_params: Dict[str, Any]
quantization_params: dict[str, Any]
if quantization_config is None or quantization_config == "binary":
quantization_params = {
"binary": {
@ -178,7 +178,7 @@ class QdrantSemanticCache(BaseCache):
if response.status_code not in (200, 201):
print_verbose(f"Qdrant semantic-cache could not create cache-key payload index: {response.text}")
except Exception as exc:
print_verbose(f"Qdrant semantic-cache could not create cache-key payload index: {str(exc)}")
print_verbose(f"Qdrant semantic-cache could not create cache-key payload index: {exc!s}")
def _payload_matches_cache_key(self, payload: dict, key: str) -> bool:
# Pre-isolation points stored only prompt + response with no cache-key
@ -188,7 +188,7 @@ class QdrantSemanticCache(BaseCache):
cached_key = payload.get(self.CACHE_KEY_FIELD_NAME)
return cached_key is not None and str(cached_key) == str(key)
def _get_embedding(self, prompt: str, metadata: Dict[str, Any] | None = None) -> EmbeddingResponse:
def _get_embedding(self, prompt: str, metadata: dict[str, Any] | None = None) -> EmbeddingResponse:
"""Embed via the proxy Router when it serves the model, else direct."""
try:
from litellm.proxy.proxy_server import llm_model_list, llm_router
@ -210,7 +210,7 @@ class QdrantSemanticCache(BaseCache):
cache={"no-store": True, "no-cache": True},
)
async def _get_async_embedding(self, prompt: str, metadata: Dict[str, Any] | None = None) -> EmbeddingResponse:
async def _get_async_embedding(self, prompt: str, metadata: dict[str, Any] | None = None) -> EmbeddingResponse:
try:
from litellm.proxy.proxy_server import llm_model_list, llm_router
except ImportError:
@ -270,7 +270,6 @@ class QdrantSemanticCache(BaseCache):
headers=self.headers,
json=data,
)
return
def get_cache(self, key, **kwargs):
print_verbose(f"sync qdrant semantic-cache get_cache, kwargs: {kwargs}")
@ -344,7 +343,6 @@ class QdrantSemanticCache(BaseCache):
else:
# cache miss !
return None
pass
async def async_set_cache(self, key, value, **kwargs):
from litellm._uuid import uuid
@ -381,7 +379,6 @@ class QdrantSemanticCache(BaseCache):
headers=self.headers,
json=data,
)
return
async def async_get_cache(self, key, **kwargs):
print_verbose(f"async qdrant semantic-cache get_cache, kwargs: {kwargs}")
@ -452,7 +449,6 @@ class QdrantSemanticCache(BaseCache):
else:
# cache miss !
return None
pass
async def _collection_info(self):
return self.collection_info

136
litellm/caching/redis_cache.py
View file

@ -16,9 +16,9 @@ import inspect
import json
import time
from collections.abc import Awaitable, Callable, Sequence
from datetime import timedelta
from contextvars import ContextVar
from typing import TYPE_CHECKING, Any, List, Optional, Tuple, TypeVar, Union, cast
from datetime import timedelta
from typing import TYPE_CHECKING, Any, TypeVar, Union, cast
import litellm
from litellm._logging import print_verbose, verbose_logger
@ -127,7 +127,7 @@ class RedisCircuitBreaker:
self.recovery_timeout = recovery_timeout
self.enabled = enabled
self._failure_count = 0
self._opened_at: Optional[float] = None
self._opened_at: float | None = None
self._state = self.CLOSED
def is_open(self) -> bool:
@ -272,10 +272,10 @@ class RedisCache(BaseCache):
host=None,
port=None,
password=None,
redis_flush_size: Optional[int] = 100,
namespace: Optional[str] = None,
startup_nodes: Optional[List] = None, # for redis-cluster
socket_timeout: Optional[float] = 5.0, # default 5 second timeout
redis_flush_size: int | None = 100,
namespace: str | None = None,
startup_nodes: list | None = None, # for redis-cluster
socket_timeout: float | None = 5.0, # default 5 second timeout
**kwargs,
):
from litellm._service_logger import ServiceLogging
@ -304,7 +304,7 @@ class RedisCache(BaseCache):
redis_kwargs.update(kwargs)
self.redis_client = get_redis_client(**redis_kwargs)
self.redis_async_client: Optional[Union[async_redis_client, async_redis_cluster_client]] = None
self.redis_async_client: async_redis_client | async_redis_cluster_client | None = None
self.redis_kwargs = redis_kwargs
self.async_redis_conn_pool = get_redis_connection_pool(**redis_kwargs)
@ -346,7 +346,7 @@ class RedisCache(BaseCache):
verbose_logger.debug("Ignoring async redis ping. No running event loop.")
else:
verbose_logger.error(
"Error connecting to Async Redis client - {}".format(str(e)),
f"Error connecting to Async Redis client - {e!s}",
extra={"error": str(e)},
)
self._handle_async_ping_error(e)
@ -407,7 +407,7 @@ class RedisCache(BaseCache):
def init_async_client(
self,
) -> Union[async_redis_client, async_redis_cluster_client]:
) -> async_redis_client | async_redis_cluster_client:
from litellm import in_memory_llm_clients_cache
from .._redis import get_redis_async_client, get_redis_connection_pool
@ -415,7 +415,7 @@ class RedisCache(BaseCache):
cache_key = self._get_async_client_cache_key()
cached_client = in_memory_llm_clients_cache.get_cache(key=cache_key)
if cached_client is not None:
redis_async_client = cast(Union[async_redis_client, async_redis_cluster_client], cached_client)
redis_async_client = cast(async_redis_client | async_redis_cluster_client, cached_client)
else:
# Create new connection pool and client for current event loop
self.async_redis_conn_pool = get_redis_connection_pool(**self.redis_kwargs)
@ -483,9 +483,9 @@ class RedisCache(BaseCache):
)
except Exception as e:
# NON blocking - notify users Redis is throwing an exception
print_verbose(f"litellm.caching.caching: set() - Got exception from REDIS : {str(e)}")
print_verbose(f"litellm.caching.caching: set() - Got exception from REDIS : {e!s}")
def increment_cache(self, key, value: int, ttl: Optional[float] = None, **kwargs) -> int:
def increment_cache(self, key, value: int, ttl: float | None = None, **kwargs) -> int:
_redis_client = self.redis_client
start_time = time.time()
set_ttl = self.get_ttl(ttl=ttl)
@ -626,7 +626,7 @@ class RedisCache(BaseCache):
async def run_script(keys: Sequence[str], args: Sequence[Any], client: Any = None) -> Any:
async def execute() -> object:
executor: Optional[Callable[..., Awaitable[Any]]] = litellm.in_memory_llm_clients_cache.get_cache(
executor: Callable[..., Awaitable[Any]] | None = litellm.in_memory_llm_clients_cache.get_cache(
key=script_cache_key
)
if executor is None:
@ -755,10 +755,10 @@ class RedisCache(BaseCache):
async def _pipeline_helper(
self,
pipe: Union[pipeline, cluster_pipeline],
cache_list: List[Tuple[Any, Any]],
ttl: Optional[float],
) -> List:
pipe: pipeline | cluster_pipeline,
cache_list: list[tuple[Any, Any]],
ttl: float | None,
) -> list:
"""
Helper function for executing a pipeline of set operations on Redis
"""
@ -769,7 +769,7 @@ class RedisCache(BaseCache):
print_verbose(f"Set ASYNC Redis Cache PIPELINE: key: {cache_key}\nValue {cache_value}\nttl={ttl}")
json_cache_value = json.dumps(cache_value)
# Set the value with a TTL if it's provided.
_td: Optional[timedelta] = None
_td: timedelta | None = None
if ttl is not None:
_td = timedelta(seconds=ttl)
pipe.set( # type: ignore
@ -782,7 +782,7 @@ class RedisCache(BaseCache):
return results
@_redis_circuit_breaker_guard
async def async_set_cache_pipeline(self, cache_list: List[Tuple[Any, Any]], ttl: Optional[float] = None, **kwargs):
async def async_set_cache_pipeline(self, cache_list: list[tuple[Any, Any]], ttl: float | None = None, **kwargs):
"""
Use Redis Pipelines for bulk write operations
"""
@ -814,7 +814,7 @@ class RedisCache(BaseCache):
parent_otel_span=_get_parent_otel_span_from_kwargs(kwargs),
)
)
return None
return
except Exception as e:
## LOGGING ##
end_time = time.time()
@ -842,8 +842,8 @@ class RedisCache(BaseCache):
self,
redis_client: async_redis_client,
key: str,
value: List,
ttl: Optional[float],
value: list,
ttl: float | None,
) -> None:
"""Helper function for async_set_cache_sadd. Separated for testing."""
ttl = self.get_ttl(ttl=ttl)
@ -856,7 +856,7 @@ class RedisCache(BaseCache):
raise
@_redis_circuit_breaker_guard
async def async_set_cache_sadd(self, key, value: List, ttl: Optional[float], **kwargs):
async def async_set_cache_sadd(self, key, value: list, ttl: float | None, **kwargs):
from redis.asyncio import Redis
start_time = time.time()
@ -938,8 +938,8 @@ class RedisCache(BaseCache):
self,
key,
value: float,
ttl: Optional[int] = None,
parent_otel_span: Optional[Span] = None,
ttl: int | None = None,
parent_otel_span: Span | None = None,
refresh_ttl: bool = False,
) -> float:
from redis.asyncio import Redis
@ -1051,7 +1051,7 @@ class RedisCache(BaseCache):
cached_response = ast.literal_eval(cached_response)
return cached_response
def get_cache(self, key, parent_otel_span: Optional[Span] = None, **kwargs):
def get_cache(self, key, parent_otel_span: Span | None = None, **kwargs):
try:
key = self.check_and_fix_namespace(key=key)
print_verbose(f"Get Redis Cache: key: {key}")
@ -1073,7 +1073,7 @@ class RedisCache(BaseCache):
# NON blocking - notify users Redis is throwing an exception
verbose_logger.error("litellm.caching.caching: get() - Got exception from REDIS: ", e)
def _run_redis_mget_operation(self, keys: List[str]) -> List[Any]:
def _run_redis_mget_operation(self, keys: list[str]) -> list[Any]:
"""
Wrapper to call `mget` on the redis client
@ -1081,7 +1081,7 @@ class RedisCache(BaseCache):
"""
return self.redis_client.mget(keys=keys) # type: ignore
async def _async_run_redis_mget_operation(self, keys: List[str]) -> List[Any]:
async def _async_run_redis_mget_operation(self, keys: list[str]) -> list[Any]:
"""
Wrapper to call `mget` on the redis client
@ -1092,8 +1092,8 @@ class RedisCache(BaseCache):
def batch_get_cache(
self,
key_list: Union[List[str], List[Optional[str]]],
parent_otel_span: Optional[Span] = None,
key_list: list[str] | list[str | None],
parent_otel_span: Span | None = None,
) -> dict:
"""
Use Redis for bulk read operations
@ -1114,7 +1114,7 @@ class RedisCache(BaseCache):
cache_key = self.check_and_fix_namespace(key=cache_key or "")
_keys.append(cache_key)
start_time = time.time()
results: List = self._run_redis_mget_operation(keys=_keys)
results: list = self._run_redis_mget_operation(keys=_keys)
end_time = time.time()
_duration = end_time - start_time
self.service_logger_obj.service_success_hook(
@ -1139,11 +1139,11 @@ class RedisCache(BaseCache):
return decoded_results
except Exception as e:
verbose_logger.error(f"Error occurred in batch get cache - {str(e)}")
verbose_logger.error(f"Error occurred in batch get cache - {e!s}")
return key_value_dict
@_redis_circuit_breaker_guard
async def async_get_cache(self, key, parent_otel_span: Optional[Span] = None, **kwargs):
async def async_get_cache(self, key, parent_otel_span: Span | None = None, **kwargs):
from redis.asyncio import Redis
_redis_client: Redis = self.init_async_client() # type: ignore
@ -1185,14 +1185,14 @@ class RedisCache(BaseCache):
event_metadata={"key": key},
)
)
print_verbose(f"litellm.caching.caching: async get() - Got exception from REDIS: {str(e)}")
print_verbose(f"litellm.caching.caching: async get() - Got exception from REDIS: {e!s}")
_record_swallowed_redis_failure(self._circuit_breaker, e)
@_redis_circuit_breaker_guard
async def async_batch_get_cache(
self,
key_list: Union[List[str], List[Optional[str]]],
parent_otel_span: Optional[Span] = None,
key_list: list[str] | list[str | None],
parent_otel_span: Span | None = None,
) -> dict:
"""
Use Redis for bulk read operations
@ -1257,7 +1257,7 @@ class RedisCache(BaseCache):
parent_otel_span=parent_otel_span,
)
)
verbose_logger.error(f"Error occurred in async batch get cache - {str(e)}")
verbose_logger.error(f"Error occurred in async batch get cache - {e!s}")
_record_swallowed_redis_failure(self._circuit_breaker, e)
return key_value_dict
@ -1292,7 +1292,7 @@ class RedisCache(BaseCache):
error=e,
call_type=f"sync_ping <- {_get_call_stack_info()}",
)
verbose_logger.error(f"LiteLLM Redis Cache PING: - Got exception from REDIS : {str(e)}")
verbose_logger.error(f"LiteLLM Redis Cache PING: - Got exception from REDIS : {e!s}")
raise e
async def ping(self) -> bool:
@ -1326,7 +1326,7 @@ class RedisCache(BaseCache):
call_type=f"async_ping <- {_get_call_stack_info()}",
)
)
verbose_logger.error(f"LiteLLM Redis Cache PING: - Got exception from REDIS : {str(e)}")
verbose_logger.error(f"LiteLLM Redis Cache PING: - Got exception from REDIS : {e!s}")
raise e
@_redis_circuit_breaker_guard
@ -1337,8 +1337,8 @@ class RedisCache(BaseCache):
# keys is a list, unpack it so it gets passed as individual elements to delete
await _redis_client.delete(*keys)
def client_list(self) -> List:
client_list: List = self.redis_client.client_list() # type: ignore
def client_list(self) -> list:
client_list: list = self.redis_client.client_list() # type: ignore
return client_list
def info(self):
@ -1388,10 +1388,10 @@ class RedisCache(BaseCache):
else:
return {"status": "failed", "message": "Redis ping returned False"}
except Exception as e:
verbose_logger.error(f"Redis connection test failed: {str(e)}")
verbose_logger.error(f"Redis connection test failed: {e!s}")
return {
"status": "failed",
"message": f"Redis connection failed: {str(e)}",
"message": f"Redis connection failed: {e!s}",
"error": str(e),
}
@ -1410,8 +1410,8 @@ class RedisCache(BaseCache):
async def _pipeline_increment_helper(
self,
pipe: pipeline,
increment_list: List[RedisPipelineIncrementOperation],
) -> Optional[List[float]]:
increment_list: list[RedisPipelineIncrementOperation],
) -> list[float] | None:
"""Helper function for pipeline increment operations"""
# Iterate through each increment operation and add commands to pipeline
for increment_op in increment_list:
@ -1431,8 +1431,8 @@ class RedisCache(BaseCache):
@_redis_circuit_breaker_guard
async def async_increment_pipeline(
self, increment_list: List[RedisPipelineIncrementOperation], **kwargs
) -> Optional[List[float]]:
self, increment_list: list[RedisPipelineIncrementOperation], **kwargs
) -> list[float] | None:
"""
Use Redis Pipelines for bulk increment operations
Args:
@ -1492,7 +1492,7 @@ class RedisCache(BaseCache):
raise e
@_redis_circuit_breaker_guard
async def async_get_ttl(self, key: str) -> Optional[int]:
async def async_get_ttl(self, key: str) -> int | None:
"""
Get the remaining TTL of a key in Redis
@ -1521,8 +1521,8 @@ class RedisCache(BaseCache):
async def async_rpush(
self,
key: str,
values: List[Any],
parent_otel_span: Optional[Span] = None,
values: list[Any],
parent_otel_span: Span | None = None,
**kwargs,
) -> int:
"""
@ -1565,14 +1565,14 @@ class RedisCache(BaseCache):
call_type=f"async_rpush <- {_get_call_stack_info()}",
)
)
verbose_logger.error(f"LiteLLM Redis Cache RPUSH: - Got exception from REDIS : {str(e)}")
verbose_logger.error(f"LiteLLM Redis Cache RPUSH: - Got exception from REDIS : {e!s}")
raise e
async def _pipeline_rpush_helper(
self,
pipe: pipeline,
rpush_list: List[RedisPipelineRpushOperation],
) -> List[int]:
rpush_list: list[RedisPipelineRpushOperation],
) -> list[int]:
"""Helper function for pipeline rpush operations"""
for rpush_op in rpush_list:
key = self.check_and_fix_namespace(key=rpush_op["key"])
@ -1587,8 +1587,8 @@ class RedisCache(BaseCache):
@_redis_circuit_breaker_guard
async def async_rpush_pipeline(
self,
rpush_list: List[RedisPipelineRpushOperation],
) -> List[int]:
rpush_list: list[RedisPipelineRpushOperation],
) -> list[int]:
"""
Use Redis Pipelines for bulk RPUSH operations
@ -1639,8 +1639,8 @@ class RedisCache(BaseCache):
)
raise e
async def handle_lpop_count_for_older_redis_versions(self, pipe: pipeline, key: str, count: int) -> List[bytes]:
result: List[bytes] = []
async def handle_lpop_count_for_older_redis_versions(self, pipe: pipeline, key: str, count: int) -> list[bytes]:
result: list[bytes] = []
for _ in range(count):
pipe.lpop(key)
results = await pipe.execute()
@ -1656,10 +1656,10 @@ class RedisCache(BaseCache):
async def async_lpop(
self,
key: str,
count: Optional[int] = None,
parent_otel_span: Optional[Span] = None,
count: int | None = None,
parent_otel_span: Span | None = None,
**kwargs,
) -> Union[Any, List[Any]]:
) -> Any | list[Any]:
_redis_client: Any = self.init_async_client()
key = self.check_and_fix_namespace(key=key)
start_time = time.time()
@ -1711,14 +1711,14 @@ class RedisCache(BaseCache):
call_type=f"async_lpop <- {_get_call_stack_info()}",
)
)
verbose_logger.error(f"LiteLLM Redis Cache LPOP: - Got exception from REDIS : {str(e)}")
verbose_logger.error(f"LiteLLM Redis Cache LPOP: - Got exception from REDIS : {e!s}")
raise e
async def _pipeline_lpop_helper(
self,
pipe: pipeline,
lpop_list: List[RedisPipelineLpopOperation],
) -> List[Optional[List[str]]]:
lpop_list: list[RedisPipelineLpopOperation],
) -> list[list[str] | None]:
"""Helper function for pipeline lpop operations.
For Redis >= 7, queues one LPOP(key, count) per operation.
@ -1734,7 +1734,7 @@ class RedisCache(BaseCache):
else:
# For Redis < 7, LPOP doesn't support count param.
# Issue `count` individual LPOP commands per key, all in one pipeline.
counts: List[int] = []
counts: list[int] = []
for lpop_op in lpop_list:
key = self.check_and_fix_namespace(key=lpop_op["key"])
count = lpop_op["count"] or 1
@ -1757,7 +1757,7 @@ class RedisCache(BaseCache):
raise r
# Decode bytes -> str for each result set
decoded_results: List[Optional[List[str]]] = []
decoded_results: list[list[str] | None] = []
for r in raw_results:
if r is None:
decoded_results.append(None)
@ -1776,8 +1776,8 @@ class RedisCache(BaseCache):
@_redis_circuit_breaker_guard
async def async_lpop_pipeline(
self,
lpop_list: List[RedisPipelineLpopOperation],
) -> List[Optional[List[str]]]:
lpop_list: list[RedisPipelineLpopOperation],
) -> list[list[str] | None]:
"""
Use Redis Pipelines for bulk LPOP operations

16
litellm/caching/redis_cluster_cache.py
View file

@ -5,7 +5,7 @@ Key differences:
- RedisClient NEEDs to be re-used across requests, adds 3000ms latency if it's re-created
"""
from typing import TYPE_CHECKING, Any, List, Optional, Union
from typing import TYPE_CHECKING, Any, Union
from litellm.caching.redis_cache import RedisCache
@ -26,8 +26,8 @@ else:
class RedisClusterCache(RedisCache):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.redis_async_redis_cluster_client: Optional[RedisCluster] = None
self.redis_sync_redis_cluster_client: Optional[RedisCluster] = None
self.redis_async_redis_cluster_client: RedisCluster | None = None
self.redis_sync_redis_cluster_client: RedisCluster | None = None
def init_async_client(self):
from redis.asyncio import RedisCluster
@ -43,13 +43,13 @@ class RedisClusterCache(RedisCache):
return _redis_client
def _run_redis_mget_operation(self, keys: List[str]) -> List[Any]:
def _run_redis_mget_operation(self, keys: list[str]) -> list[Any]:
"""
Overrides `_run_redis_mget_operation` in redis_cache.py
"""
return self.redis_client.mget_nonatomic(keys=keys) # type: ignore
async def _async_run_redis_mget_operation(self, keys: List[str]) -> List[Any]:
async def _async_run_redis_mget_operation(self, keys: list[str]) -> list[Any]:
"""
Overrides `_async_run_redis_mget_operation` in redis_cache.py
"""
@ -71,7 +71,7 @@ class RedisClusterCache(RedisCache):
cluster_kwargs = self.redis_kwargs.copy()
startup_nodes = cluster_kwargs.pop("startup_nodes", [])
new_startup_nodes: List[ClusterNode] = []
new_startup_nodes: list[ClusterNode] = []
for item in startup_nodes:
new_startup_nodes.append(ClusterNode(**item))
@ -100,9 +100,9 @@ class RedisClusterCache(RedisCache):
except Exception as e:
from litellm._logging import verbose_logger
verbose_logger.error(f"Redis Cluster connection test failed: {str(e)}")
verbose_logger.error(f"Redis Cluster connection test failed: {e!s}")
return {
"status": "failed",
"message": f"Redis Cluster connection failed: {str(e)}",
"message": f"Redis Cluster connection failed: {e!s}",
"error": str(e),
}

54
litellm/caching/redis_semantic_cache.py
View file

@ -13,7 +13,7 @@ import ast
import asyncio
import json
import os
from typing import Any, Dict, List, Optional, Tuple, cast
from typing import Any, cast
import litellm
from litellm._logging import print_verbose, verbose_logger
@ -40,13 +40,13 @@ class RedisSemanticCache(BaseCache):
def __init__(
self,
host: Optional[str] = None,
port: Optional[str] = None,
password: Optional[str] = None,
redis_url: Optional[str] = None,
similarity_threshold: Optional[float] = None,
host: str | None = None,
port: str | None = None,
password: str | None = None,
redis_url: str | None = None,
similarity_threshold: float | None = None,
embedding_model: str = "text-embedding-ada-002",
index_name: Optional[str] = None,
index_name: str | None = None,
**kwargs,
):
"""
@ -142,7 +142,7 @@ class RedisSemanticCache(BaseCache):
raise
@classmethod
def _cache_key_filterable_field(cls) -> Dict[str, str]:
def _cache_key_filterable_field(cls) -> dict[str, str]:
return {
"name": cls.CACHE_KEY_FIELD_NAME,
"type": "tag",
@ -203,7 +203,7 @@ class RedisSemanticCache(BaseCache):
overwrite=True,
)
def _get_cache_filters(self, key: str) -> Dict[str, str]:
def _get_cache_filters(self, key: str) -> dict[str, str]:
return {self.CACHE_KEY_FIELD_NAME: str(key)}
def _get_cache_key_filter_expression(self, key: str) -> Any:
@ -211,7 +211,7 @@ class RedisSemanticCache(BaseCache):
return Tag(self.CACHE_KEY_FIELD_NAME) == str(key)
def _cache_hit_matches_key(self, cache_hit: Dict[str, Any], key: str) -> bool:
def _cache_hit_matches_key(self, cache_hit: dict[str, Any], key: str) -> bool:
# Pre-isolation entries with no ``litellm_cache_key`` field cannot be
# safely reassigned to a caller's scope and are treated as misses.
cached_key = cache_hit.get(self.CACHE_KEY_FIELD_NAME)
@ -219,7 +219,7 @@ class RedisSemanticCache(BaseCache):
cached_key = cached_key.decode("utf-8")
return cached_key is not None and str(cached_key) == str(key)
def _get_ttl(self, **kwargs) -> Optional[int]:
def _get_ttl(self, **kwargs) -> int | None:
"""
Get the TTL (time-to-live) value for cache entries.
@ -235,7 +235,7 @@ class RedisSemanticCache(BaseCache):
return ttl
@classmethod
def _get_prompt_from_kwargs(cls, **kwargs) -> Optional[str]:
def _get_prompt_from_kwargs(cls, **kwargs) -> str | None:
"""
Extract a semantic-cache prompt from chat or Responses API request kwargs.
"""
@ -246,13 +246,13 @@ class RedisSemanticCache(BaseCache):
if "input" not in kwargs:
return None
prompt_parts: List[str] = []
prompt_parts: list[str] = []
cls._collect_responses_input_text(kwargs.get("input"), prompt_parts)
prompt = "\n".join(prompt_parts).strip()
return prompt or None
@classmethod
def _collect_responses_input_text(cls, value: Any, prompt_parts: List[str]) -> None:
def _collect_responses_input_text(cls, value: Any, prompt_parts: list[str]) -> None:
value = cls._coerce_response_input_value(value)
if value is None:
return
@ -306,7 +306,7 @@ class RedisSemanticCache(BaseCache):
return dict_method()
return value
def _get_embedding(self, prompt: str, metadata: Dict[str, Any] | None = None) -> List[float]:
def _get_embedding(self, prompt: str, metadata: dict[str, Any] | None = None) -> list[float]:
"""
Routes through the proxy Router when the embedding model is a Router
deployment so per-deployment auth (e.g. Bedrock aws_role_name) applies,
@ -364,7 +364,7 @@ class RedisSemanticCache(BaseCache):
try:
cached_response = ast.literal_eval(cached_response)
except (ValueError, SyntaxError) as e:
print_verbose(f"Error parsing cached response: {str(e)}")
print_verbose(f"Error parsing cached response: {e!s}")
return None
return cached_response
@ -381,7 +381,7 @@ class RedisSemanticCache(BaseCache):
"""
print_verbose(f"Redis semantic-cache set_cache, kwargs: {kwargs}")
value_str: Optional[str] = None
value_str: str | None = None
try:
prompt = self._get_prompt_from_kwargs(**kwargs)
if prompt is None:
@ -403,7 +403,7 @@ class RedisSemanticCache(BaseCache):
store_kwargs["ttl"] = int(ttl)
self.llmcache.store(prompt, value_str, **store_kwargs)
except Exception as e:
print_verbose(f"Error setting {value_str or value} in the Redis semantic cache: {str(e)}")
print_verbose(f"Error setting {value_str or value} in the Redis semantic cache: {e!s}")
def get_cache(self, key: str, **kwargs) -> Any:
"""
@ -468,10 +468,10 @@ class RedisSemanticCache(BaseCache):
return self._get_cache_logic(cached_response=cached_response)
except Exception as e:
print_verbose(f"Error retrieving from Redis semantic cache: {str(e)}")
print_verbose(f"Error retrieving from Redis semantic cache: {e!s}")
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
async def _get_async_embedding(self, prompt: str, metadata: Dict[str, Any] | None = None) -> List[float]:
async def _get_async_embedding(self, prompt: str, metadata: dict[str, Any] | None = None) -> list[float]:
"""
Asynchronously generate an embedding for the given prompt.
@ -505,8 +505,8 @@ class RedisSemanticCache(BaseCache):
)
return embedding_response["data"][0]["embedding"]
except Exception as e:
print_verbose(f"Error generating async embedding: {str(e)}")
raise ValueError(f"Failed to generate embedding: {str(e)}") from e
print_verbose(f"Error generating async embedding: {e!s}")
raise ValueError(f"Failed to generate embedding: {e!s}") from e
async def async_set_cache(self, key: str, value: Any, **kwargs) -> None:
"""
@ -546,7 +546,7 @@ class RedisSemanticCache(BaseCache):
**store_kwargs,
)
except Exception as e:
print_verbose(f"Error in async_set_cache: {str(e)}")
print_verbose(f"Error in async_set_cache: {e!s}")
async def async_get_cache(self, key: str, **kwargs) -> Any:
"""
@ -612,10 +612,10 @@ class RedisSemanticCache(BaseCache):
return self._get_cache_logic(cached_response=cached_response)
except Exception as e:
print_verbose(f"Error in async_get_cache: {str(e)}")
print_verbose(f"Error in async_get_cache: {e!s}")
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
async def _index_info(self) -> Dict[str, Any]:
async def _index_info(self) -> dict[str, Any]:
"""
Get information about the Redis index.
@ -625,7 +625,7 @@ class RedisSemanticCache(BaseCache):
aindex = await self.llmcache._get_async_index()
return await aindex.info()
async def async_set_cache_pipeline(self, cache_list: List[Tuple[str, Any]], **kwargs) -> None:
async def async_set_cache_pipeline(self, cache_list: list[tuple[str, Any]], **kwargs) -> None:
"""
Asynchronously store multiple values in the semantic cache.
@ -639,4 +639,4 @@ class RedisSemanticCache(BaseCache):
tasks.append(self.async_set_cache(val[0], val[1], **kwargs))
await asyncio.gather(*tasks)
except Exception as e:
print_verbose(f"Error in async_set_cache_pipeline: {str(e)}")
print_verbose(f"Error in async_set_cache_pipeline: {e!s}")

5
litellm/caching/s3_cache.py
View file

@ -11,9 +11,8 @@ Has 4 methods:
import ast
import asyncio
import json
from datetime import datetime, timedelta, timezone
from functools import partial
from typing import Optional
from datetime import datetime, timezone, timedelta
from litellm._logging import print_verbose, verbose_logger
@ -26,7 +25,7 @@ class S3Cache(BaseCache):
s3_bucket_name,
s3_region_name=None,
s3_api_version=None,
s3_use_ssl: Optional[bool] = True,
s3_use_ssl: bool | None = True,
s3_verify=None,
s3_endpoint_url=None,
s3_aws_access_key_id=None,

10
litellm/caching/valkey_semantic_cache.py
View file

@ -249,7 +249,7 @@ class ValkeySemanticCache(RedisSemanticCache):
if ttl is not None:
self.sync_client.expire(doc_key, ttl)
except Exception as e:
print_verbose(f"Error in Valkey semantic-cache set_cache: {str(e)}")
print_verbose(f"Error in Valkey semantic-cache set_cache: {e!s}")
def get_cache(self, key: str, **kwargs: Any) -> Any:
print_verbose(f"Valkey semantic-cache get_cache, kwargs: {kwargs}")
@ -268,7 +268,7 @@ class ValkeySemanticCache(RedisSemanticCache):
)
return self._resolve_hit(self._first_hit(search_result), key, **kwargs)
except Exception as e:
print_verbose(f"Error in Valkey semantic-cache get_cache: {str(e)}")
print_verbose(f"Error in Valkey semantic-cache get_cache: {e!s}")
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
async def async_set_cache(self, key: str, value: Any, **kwargs: Any) -> None:
@ -288,7 +288,7 @@ class ValkeySemanticCache(RedisSemanticCache):
if ttl is not None:
await self.async_client.expire(doc_key, ttl)
except Exception as e:
print_verbose(f"Error in async Valkey semantic-cache set_cache: {str(e)}")
print_verbose(f"Error in async Valkey semantic-cache set_cache: {e!s}")
async def async_get_cache(self, key: str, **kwargs: Any) -> Any:
print_verbose(f"Async Valkey semantic-cache get_cache, kwargs: {kwargs}")
@ -307,14 +307,14 @@ class ValkeySemanticCache(RedisSemanticCache):
)
return self._resolve_hit(self._first_hit(search_result), key, **kwargs)
except Exception as e:
print_verbose(f"Error in async Valkey semantic-cache get_cache: {str(e)}")
print_verbose(f"Error in async Valkey semantic-cache get_cache: {e!s}")
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
async def async_set_cache_pipeline(self, cache_list: list[tuple[str, Any]], **kwargs: Any) -> None:
try:
await asyncio.gather(*[self.async_set_cache(key, value, **kwargs) for key, value in cache_list])
except Exception as e:
print_verbose(f"Error in Valkey semantic-cache async_set_cache_pipeline: {str(e)}")
print_verbose(f"Error in Valkey semantic-cache async_set_cache_pipeline: {e!s}")
async def _index_info(self) -> dict:
return await self.async_client.ft(self.index_name).info()

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

@ -3,7 +3,7 @@ Handler for transforming /chat/completions api requests to litellm.responses req
"""
from collections.abc import Coroutine
from typing import TYPE_CHECKING, Any, Optional, Union
from typing import TYPE_CHECKING, Any, Union
from typing_extensions import TypedDict
@ -47,7 +47,7 @@ class ResponsesToCompletionBridgeHandler:
@staticmethod
def _coerce_response_object(
response_obj: Any,
hidden_params: Optional[dict],
hidden_params: dict | None,
) -> "ResponsesAPIResponse":
if isinstance(response_obj, ResponsesAPIResponse):
response = response_obj

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

@ -8,11 +8,7 @@ from collections.abc import AsyncIterator, Callable, Iterable, Iterator
from typing import (
TYPE_CHECKING,
Any,
Dict,
List,
Literal,
Optional,
Tuple,
Union,
cast,
)
@ -61,7 +57,7 @@ if TYPE_CHECKING:
def _get_reasoning_items(
msg: "AllMessageValues",
) -> List[ChatCompletionReasoningItem]:
) -> list[ChatCompletionReasoningItem]:
"""Extract reasoning_items from a message dict with proper typing."""
items = msg.get("reasoning_items") # type: ignore[union-attr]
if items:
@ -71,14 +67,14 @@ def _get_reasoning_items(
def _build_reasoning_item(
item_id: str,
encrypted_content: Optional[str],
encrypted_content: str | None,
summary_raw: Any,
) -> Dict[str, Any]:
) -> dict[str, Any]:
"""Build a ChatCompletionReasoningItem-shaped dict from raw response data.
Handles both pydantic objects (attribute access) and plain dicts.
"""
summary: List[Dict[str, Any]] = []
summary: list[dict[str, Any]] = []
for s in summary_raw or []:
if isinstance(s, dict):
summary.append({"type": s.get("type", "summary_text"), "text": s.get("text", "")})
@ -98,10 +94,10 @@ def _build_reasoning_item(
def _reasoning_item_to_response_input(
r_item: Union[ChatCompletionReasoningItem, Dict[str, Any]],
) -> Dict[str, Any]:
r_item: ChatCompletionReasoningItem | dict[str, Any],
) -> dict[str, Any]:
"""Convert a stored ChatCompletionReasoningItem back to a Responses API input item."""
r_input: Dict[str, Any] = {
r_input: dict[str, Any] = {
"type": "reasoning",
"id": r_item.get("id") or f"rs_{id(r_item)}",
# summary is always required by the Responses API, even when empty
@ -134,7 +130,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
return {"type": "function", "name": fn_name}
return tool_choice
def _handle_raw_dict_response_item(self, item: Dict[str, Any], index: int) -> Tuple[Optional[Any], int]:
def _handle_raw_dict_response_item(self, item: dict[str, Any], index: int) -> tuple[Any | None, int]:
"""
Handle raw dict response items from Responses API (e.g., GPT-5 Codex format).
@ -208,10 +204,10 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
return None, index
def convert_chat_completion_messages_to_responses_api(
self, messages: List["AllMessageValues"]
) -> Tuple[List[Any], Optional[str]]:
input_items: List[Any] = []
instructions: Optional[str] = None
self, messages: list["AllMessageValues"]
) -> tuple[list[Any], str | None]:
input_items: list[Any] = []
instructions: str | None = None
for msg in messages:
role = msg.get("role")
@ -242,7 +238,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
# Convert tool message to function call output format
# The Responses API expects 'output' to be a list with input_text/input_image types
# Using list format for consistency across text and multimodal content
tool_output: List[Dict[str, Any]]
tool_output: list[dict[str, Any]]
if content is None:
tool_output = []
elif isinstance(content, str):
@ -270,7 +266,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
for tool_call in tool_calls:
function = tool_call.get("function")
if function:
input_tool_call: Dict[str, Any] = {
input_tool_call: dict[str, Any] = {
"type": "function_call",
"call_id": tool_call["id"],
}
@ -308,7 +304,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
responses_api_request["max_output_tokens"] = value
elif key == "tools" and value is not None:
responses_api_request["tools"] = self._convert_tools_to_responses_format(
cast(List[Dict[str, Any]], value)
cast(list[dict[str, Any]], value)
)
elif key == "response_format":
text_format = self._transform_response_format_to_text_format(value)
@ -331,15 +327,15 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
elif key == "web_search_options":
self._add_web_search_tool(responses_api_request, value)
def _build_sanitized_litellm_params(self, litellm_params: dict) -> Dict[str, Any]:
def _build_sanitized_litellm_params(self, litellm_params: dict) -> dict[str, Any]:
"""Build sanitized litellm_params with merged metadata."""
responses_optional_param_keys = set(ResponsesAPIOptionalRequestParams.__annotations__.keys())
sanitized: Dict[str, Any] = {
sanitized: dict[str, Any] = {
key: value for key, value in litellm_params.items() if key not in responses_optional_param_keys
}
legacy_metadata = litellm_params.get("metadata")
existing_litellm_metadata = litellm_params.get("litellm_metadata")
merged_litellm_metadata: Dict[str, Any] = {}
merged_litellm_metadata: dict[str, Any] = {}
if isinstance(legacy_metadata, dict):
merged_litellm_metadata.update(legacy_metadata)
if isinstance(existing_litellm_metadata, dict):
@ -352,9 +348,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
def _merge_responses_api_request_into_request_data(
self,
request_data: Dict[str, Any],
request_data: dict[str, Any],
responses_api_request: "ResponsesAPIOptionalRequestParams",
instructions: Optional[str],
instructions: str | None,
) -> None:
"""Add non-None values from responses_api_request into request_data."""
for key, value in responses_api_request.items():
@ -374,12 +370,12 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
def transform_request(
self,
model: str,
messages: List["AllMessageValues"],
messages: list["AllMessageValues"],
optional_params: dict,
litellm_params: dict,
headers: dict,
litellm_logging_obj: "LiteLLMLoggingObj",
client: Optional[Any] = None,
client: Any | None = None,
) -> dict:
(
input_items,
@ -453,9 +449,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
@staticmethod
def _convert_response_output_to_choices(
output_items: List[Any],
handle_raw_dict_callback: Optional[Callable] = None,
) -> List[Any]:
output_items: list[Any],
handle_raw_dict_callback: Callable | None = None,
) -> list[Any]:
"""
Convert Responses API output items to chat completion choices.
@ -481,14 +477,14 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
from litellm.types.utils import Choices, Message
choices: List[Choices] = []
choices: list[Choices] = []
index = 0
reasoning_content: Optional[str] = None
pending_reasoning_item: Optional[Dict[str, Any]] = None
reasoning_content: str | None = None
pending_reasoning_item: dict[str, Any] | None = None
# Collect all tool calls to put them in a single choice
# (Chat Completions API expects all tool calls in one message)
accumulated_tool_calls: List[Dict[str, Any]] = []
accumulated_tool_calls: list[dict[str, Any]] = []
tool_call_index = 0
for item in output_items:
@ -514,7 +510,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
reasoning_content=reasoning_content,
annotations=annotations,
reasoning_items=cast(
Optional[List[ChatCompletionReasoningItem]],
list[ChatCompletionReasoningItem] | None,
([pending_reasoning_item] if pending_reasoning_item is not None else None),
),
)
@ -574,7 +570,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
tool_calls=accumulated_tool_calls,
reasoning_content=reasoning_content,
reasoning_items=cast(
Optional[List[ChatCompletionReasoningItem]],
list[ChatCompletionReasoningItem] | None,
([pending_reasoning_item] if pending_reasoning_item is not None else None),
),
)
@ -585,22 +581,22 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
return choices
@classmethod
def _extract_output_from_completed_event(cls, parsed_chunk: Dict[str, Any]) -> Optional[List[Dict[str, Any]]]:
def _extract_output_from_completed_event(cls, parsed_chunk: dict[str, Any]) -> list[dict[str, Any]] | None:
response_payload = parsed_chunk.get("response")
if not isinstance(response_payload, dict):
return None
response_output = response_payload.get("output")
if not isinstance(response_output, list) or len(response_output) == 0:
return None
return cast(List[Dict[str, Any]], response_output)
return cast(list[dict[str, Any]], response_output)
@classmethod
def _recover_output_items_from_raw_sse(cls, raw_sse: Optional[str]) -> List[Dict[str, Any]]:
def _recover_output_items_from_raw_sse(cls, raw_sse: str | None) -> list[dict[str, Any]]:
if not raw_sse or not isinstance(raw_sse, str):
return []
recovered_output_items: Dict[int, Dict[str, Any]] = {}
recovered_text_only_items: Dict[int, Dict[str, Any]] = {}
recovered_output_items: dict[int, dict[str, Any]] = {}
recovered_text_only_items: dict[int, dict[str, Any]] = {}
for chunk in raw_sse.splitlines():
parsed_chunk = parse_sse_json_chunk(chunk)
@ -635,7 +631,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
# but text-only items at indices without a matching OUTPUT_ITEM_DONE
# must still be preserved (e.g. multi-output responses where some
# indices only emitted OUTPUT_TEXT_DONE).
merged_items: Dict[int, Dict[str, Any]] = {**recovered_text_only_items}
merged_items: dict[int, dict[str, Any]] = {**recovered_text_only_items}
merged_items.update(recovered_output_items)
if merged_items:
@ -644,7 +640,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
return []
@classmethod
def _recover_output_items_from_logging(cls, logging_obj: "LiteLLMLoggingObj") -> List[Dict[str, Any]]:
def _recover_output_items_from_logging(cls, logging_obj: "LiteLLMLoggingObj") -> list[dict[str, Any]]:
model_call_details = getattr(logging_obj, "model_call_details", {}) or {}
original_response = model_call_details.get("original_response")
return cls._recover_output_items_from_raw_sse(original_response)
@ -656,12 +652,12 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
model_response: "ModelResponse",
logging_obj: "LiteLLMLoggingObj",
request_data: dict,
messages: List["AllMessageValues"],
messages: list["AllMessageValues"],
optional_params: dict,
litellm_params: dict,
encoding: Any,
api_key: Optional[str] = None,
json_mode: Optional[bool] = None,
api_key: str | None = None,
json_mode: bool | None = None,
) -> "ModelResponse":
"""Transform Responses API response to chat completion response"""
from litellm.responses.utils import ResponseAPILoggingUtils
@ -729,11 +725,11 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
self,
streaming_response: Union[Iterator[str], AsyncIterator[str], "ModelResponse", "BaseModel"],
sync_stream: bool,
json_mode: Optional[bool] = False,
json_mode: bool | None = False,
) -> BaseModelResponseIterator:
return OpenAiResponsesToChatCompletionStreamIterator(streaming_response, sync_stream, json_mode)
def _convert_content_str_to_input_text(self, content: str, role: str) -> Dict[str, Any]:
def _convert_content_str_to_input_text(self, content: str, role: str) -> dict[str, Any]:
if role == "user" or role == "system" or role == "tool":
return {"type": "input_text", "text": content}
else:
@ -745,15 +741,15 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
from openai.types.responses import ResponseInputImageParam
content_image_url = content.get("image_url")
actual_image_url: Optional[str] = None
detail: Optional[Literal["low", "high", "auto"]] = None
actual_image_url: str | None = None
detail: Literal["low", "high", "auto"] | None = None
if isinstance(content_image_url, str):
actual_image_url = content_image_url
elif isinstance(content_image_url, dict):
actual_image_url = content_image_url.get("url")
detail = cast(
Optional[Literal["low", "high", "auto"]],
Literal["low", "high", "auto"] | None,
content_image_url.get("detail"),
)
@ -769,21 +765,14 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
def _convert_content_to_responses_format(
self,
content: Optional[
Union[
str,
List[Any],
Iterable[
Union[
"OpenAIMessageContentListBlock",
"ChatCompletionThinkingBlock",
"ChatCompletionRedactedThinkingBlock",
]
],
]
],
content: str
| list[Any]
| Iterable[
Union["OpenAIMessageContentListBlock", "ChatCompletionThinkingBlock", "ChatCompletionRedactedThinkingBlock"]
]
| None,
role: str,
) -> List[Dict[str, Any]]:
) -> list[dict[str, Any]]:
"""Convert chat completion content to responses API format"""
from litellm.types.llms.openai import ChatCompletionImageObject
@ -863,9 +852,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
verbose_logger.debug(f"Chat provider: Other content type -> {result}")
return result
def _convert_tools_to_responses_format(self, tools: List[Dict[str, Any]]) -> List["ALL_RESPONSES_API_TOOL_PARAMS"]:
def _convert_tools_to_responses_format(self, tools: list[dict[str, Any]]) -> list["ALL_RESPONSES_API_TOOL_PARAMS"]:
"""Convert chat completion tools to responses API tools format"""
responses_tools: List["ALL_RESPONSES_API_TOOL_PARAMS"] = []
responses_tools: list[ALL_RESPONSES_API_TOOL_PARAMS] = []
for tool in tools:
# convert function tool from chat completion to responses API format
if tool.get("type") == "function":
@ -882,7 +871,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
else:
responses_tools.append(tool) # type: ignore
return cast(List["ALL_RESPONSES_API_TOOL_PARAMS"], responses_tools)
return cast(list["ALL_RESPONSES_API_TOOL_PARAMS"], responses_tools)
def _extract_extra_body_params(self, optional_params: dict):
"""
@ -913,7 +902,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
return optional_params
def _map_reasoning_effort(self, reasoning_effort: Union[str, Dict[str, Any]]) -> Optional[Reasoning]:
def _map_reasoning_effort(self, reasoning_effort: str | dict[str, Any]) -> Reasoning | None:
# If dict is passed, convert it directly to Reasoning object
if isinstance(reasoning_effort, dict):
return Reasoning(**reasoning_effort) # type: ignore[typeddict-item]
@ -964,16 +953,14 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
tools = []
responses_api_request["tools"] = tools
web_search_tool: Dict[str, Any] = {"type": "web_search"}
web_search_tool: dict[str, Any] = {"type": "web_search"}
if isinstance(web_search_options, dict):
web_search_tool.update(web_search_options)
# Cast to Any to match the expected union type for tools list items
tools.append(cast(Any, web_search_tool))
def _transform_response_format_to_text_format(
self, response_format: Union[Dict[str, Any], Any]
) -> Optional[Dict[str, Any]]:
def _transform_response_format_to_text_format(self, response_format: dict[str, Any] | Any) -> dict[str, Any] | None:
"""
Transform Chat Completion response_format parameter to Responses API text.format parameter.
@ -1022,8 +1009,8 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
@staticmethod
def _convert_annotations_to_chat_format(
annotations: Optional[List[Any]],
) -> Optional[List[ChatCompletionAnnotation]]:
annotations: list[Any] | None,
) -> list[ChatCompletionAnnotation] | None:
"""
Convert annotations from Responses API to Chat Completions format.
@ -1033,7 +1020,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
if not annotations:
return None
result: List[ChatCompletionAnnotation] = []
result: list[ChatCompletionAnnotation] = []
for annotation in annotations:
try:
# Convert Pydantic models to dicts (handles both v1 and v2)
@ -1056,7 +1043,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
return result if result else None
def _map_responses_status_to_finish_reason(self, status: Optional[str]) -> str:
def _map_responses_status_to_finish_reason(self, status: str | None) -> str:
"""Map responses API status to chat completion finish_reason"""
if not status:
return "stop"
@ -1072,7 +1059,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
def __init__(self, streaming_response, sync_stream: bool, json_mode: Optional[bool] = False):
def __init__(self, streaming_response, sync_stream: bool, json_mode: bool | None = False):
super().__init__(streaming_response, sync_stream, json_mode)
self._chat_completion_id: str | None = None
@ -1095,7 +1082,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
@staticmethod
def translate_responses_chunk_to_openai_stream(
parsed_chunk: Union[dict, BaseModel],
parsed_chunk: dict | BaseModel,
) -> "ModelResponseStream":
"""
Translate a Responses API streaming chunk to OpenAI chat completion streaming format.
@ -1196,7 +1183,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
]
)
elif event_type == "response.function_call_arguments.delta":
content_part: Optional[str] = parsed_chunk.get("delta", None)
content_part: str | None = parsed_chunk.get("delta", None)
if content_part:
tool_call_index = parsed_chunk.get("output_index", 0)
return ModelResponseStream(
@ -1319,7 +1306,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
finish_reason = "tool_calls" if has_function_calls else "stop"
# Extract reasoning items with encrypted_content for round-tripping
completed_reasoning_items: Optional[List[Dict[str, Any]]] = None
completed_reasoning_items: list[dict[str, Any]] | None = None
for item in output_items:
if not isinstance(item, dict) or item.get("type") != "reasoning":
continue
@ -1333,7 +1320,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
)
)
completed_reasoning_items_typed = cast(
Optional[List[ChatCompletionReasoningItem]],
list[ChatCompletionReasoningItem] | None,
completed_reasoning_items,
)

92
litellm/compression/compress.py
View file

@ -4,7 +4,7 @@ scoring, message stubbing, and retrieval tool injection.
"""
from collections.abc import Mapping, Sequence
from typing import Any, Dict, List, Optional, Set, Tuple, Union, cast
from typing import Any, cast
from litellm.caching.dual_cache import DualCache
from litellm.compression.message_stubbing import (
@ -33,7 +33,7 @@ _SUPPORTED_CALL_TYPES = frozenset(
)
def _normalize_call_type(call_type: Union[CallTypes, str]) -> str:
def _normalize_call_type(call_type: CallTypes | str) -> str:
"""Return the string value for a ``CallTypes`` enum or a raw string."""
if isinstance(call_type, CallTypes):
return call_type.value
@ -44,7 +44,7 @@ def _is_anthropic_call_type(call_type: str) -> bool:
return call_type in _ANTHROPIC_CALL_TYPES
def _build_retrieval_tools(keys: List[str], call_type: str) -> List[dict]:
def _build_retrieval_tools(keys: list[str], call_type: str) -> list[dict]:
"""
Build retrieval tool definitions in the target request schema.
@ -63,7 +63,7 @@ def _build_retrieval_tools(keys: List[str], call_type: str) -> List[dict]:
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
anthropic_tools, _mcp_servers = AnthropicConfig()._map_tools(openai_tools)
return cast(List[dict], anthropic_tools)
return cast(list[dict], anthropic_tools)
def _content_to_text(content: Any) -> str:
@ -77,8 +77,8 @@ def _content_to_text(content: Any) -> str:
Implemented iteratively (stack-based) to avoid unbounded recursion.
"""
parts: List[str] = []
stack: List[Any] = [content]
parts: list[str] = []
stack: list[Any] = [content]
while stack:
item = stack.pop()
if isinstance(item, str):
@ -97,9 +97,9 @@ def _content_to_text(content: Any) -> str:
def _normalize_messages_for_compression(
messages: List[dict],
messages: list[dict],
call_type: str,
) -> Tuple[List[dict], List[dict]]:
) -> tuple[list[dict], list[dict]]:
"""
Normalize each original message to a text-surrogate content for scoring.
@ -111,9 +111,9 @@ def _normalize_messages_for_compression(
f"Unsupported call_type={call_type!r} for compression. Expected one of: {sorted(_SUPPORTED_CALL_TYPES)}."
)
original_messages: List[Dict[str, Any]] = [dict(m) for m in messages]
original_messages: list[dict[str, Any]] = [dict(m) for m in messages]
normalized_messages: List[dict] = []
normalized_messages: list[dict] = []
for msg in original_messages:
normalized_messages.append(
{
@ -124,7 +124,7 @@ def _normalize_messages_for_compression(
return normalized_messages, original_messages
def _extract_last_user_message(messages: List[dict]) -> str:
def _extract_last_user_message(messages: list[dict]) -> str:
"""Return the text content of the last user message."""
for msg in reversed(messages):
if msg.get("role") == "user":
@ -132,10 +132,10 @@ def _extract_last_user_message(messages: List[dict]) -> str:
return ""
def _extract_tool_use_ids(content: Any) -> List[str]:
def _extract_tool_use_ids(content: Any) -> list[str]:
if not isinstance(content, list):
return []
tool_use_ids: List[str] = []
tool_use_ids: list[str] = []
for part in content:
if not isinstance(part, dict):
continue
@ -147,10 +147,10 @@ def _extract_tool_use_ids(content: Any) -> List[str]:
return tool_use_ids
def _extract_tool_result_ids(content: Any) -> Set[str]:
def _extract_tool_result_ids(content: Any) -> set[str]:
if not isinstance(content, list):
return set()
tool_result_ids: Set[str] = set()
tool_result_ids: set[str] = set()
for part in content:
if not isinstance(part, dict):
continue
@ -163,15 +163,15 @@ def _extract_tool_result_ids(content: Any) -> Set[str]:
def _extract_anthropic_tool_exchange_spans(
messages: List[dict],
) -> Tuple[List[Set[int]], Optional[str]]:
messages: list[dict],
) -> tuple[list[set[int]], str | None]:
"""
Return atomic 2-message spans for Anthropic tool exchanges.
Each assistant message containing `tool_use` must be immediately followed by a
user message containing matching `tool_result` blocks for all tool_use ids.
"""
spans: List[Set[int]] = []
spans: list[set[int]] = []
i = 0
while i < len(messages):
current = messages[i]
@ -223,13 +223,13 @@ def get_protected_indices(messages: Sequence[Mapping[str, object]]) -> tuple[int
def _combine_scores(
bm25_scores: List[float],
emb_scores: List[float],
bm25_scores: list[float],
emb_scores: list[float],
bm25_weight: float = 0.4,
) -> List[float]:
) -> list[float]:
"""Weighted average of BM25 and embedding scores, with min-max normalization."""
def _normalize(scores: List[float]) -> List[float]:
def _normalize(scores: list[float]) -> list[float]:
min_s = min(scores) if scores else 0.0
max_s = max(scores) if scores else 0.0
rng = max_s - min_s
@ -245,14 +245,14 @@ def _combine_scores(
def _select_kept_indices_for_budget(
normalized_messages: List[dict],
original_messages: List[dict],
combined_scores: List[float],
normalized_messages: list[dict],
original_messages: list[dict],
combined_scores: list[float],
compression_target: int,
model: str,
initial_kept_indices: Set[int],
tool_exchange_spans: List[Set[int]],
) -> Tuple[Set[int], Dict[int, dict]]:
initial_kept_indices: set[int],
tool_exchange_spans: list[set[int]],
) -> tuple[set[int], dict[int, dict]]:
kept_indices = set(initial_kept_indices)
current_tokens = 0
for i in kept_indices:
@ -265,14 +265,14 @@ def _select_kept_indices_for_budget(
# A unit is either:
# 1) a single message index, or
# 2) an Anthropic tool-exchange span that must be kept/dropped atomically.
truncated_overrides: Dict[int, dict] = {} # idx -> truncated message dict
span_id_by_index: Dict[int, int] = {}
truncated_overrides: dict[int, dict] = {} # idx -> truncated message dict
span_id_by_index: dict[int, int] = {}
for span_id, span in enumerate(tool_exchange_spans):
for idx in span:
span_id_by_index[idx] = span_id
# Build single-message candidate units (non-span messages).
candidate_units: List[Tuple[float, Tuple[int, ...], bool]] = []
candidate_units: list[tuple[float, tuple[int, ...], bool]] = []
for idx in range(len(normalized_messages)):
if idx in span_id_by_index or idx in kept_indices:
continue
@ -322,8 +322,8 @@ def _select_kept_indices_for_budget(
return kept_indices, truncated_overrides
def _get_dropped_tool_span_indices(kept_indices: Set[int], tool_exchange_spans: List[Set[int]]) -> Set[int]:
dropped_tool_span_indices: Set[int] = set()
def _get_dropped_tool_span_indices(kept_indices: set[int], tool_exchange_spans: list[set[int]]) -> set[int]:
dropped_tool_span_indices: set[int] = set()
for span in tool_exchange_spans:
if not any(idx in kept_indices for idx in span):
dropped_tool_span_indices.update(span)
@ -331,14 +331,14 @@ def _get_dropped_tool_span_indices(kept_indices: Set[int], tool_exchange_spans:
def compress(
messages: List[dict],
messages: list[dict],
model: str,
call_type: Union[CallTypes, str] = CallTypes.completion,
call_type: CallTypes | str = CallTypes.completion,
compression_trigger: int = 200_000,
compression_target: Optional[int] = None,
embedding_model: Optional[str] = None,
embedding_model_params: Optional[Dict[str, Any]] = None,
compression_cache: Optional[DualCache] = None,
compression_target: int | None = None,
embedding_model: str | None = None,
embedding_model_params: dict[str, Any] | None = None,
compression_cache: DualCache | None = None,
) -> CompressedResult:
"""
Compress a list of messages by replacing low-relevance content with stubs.
@ -383,7 +383,7 @@ def compress(
original_tokens = token_counter(
model=model,
messages=cast(List[Any], original_messages),
messages=cast(list[Any], original_messages),
)
# Pass through if below trigger
@ -422,9 +422,9 @@ def compress(
# Protected messages are never compressed
protected_indices = get_protected_indices(normalized_messages)
kept_indices: Set[int] = set(protected_indices)
kept_indices: set[int] = set(protected_indices)
tool_exchange_spans: List[Set[int]] = []
tool_exchange_spans: list[set[int]] = []
if _is_anthropic_call_type(call_type_str):
tool_exchange_spans, tool_sequence_error = _extract_anthropic_tool_exchange_spans(original_messages)
if tool_sequence_error is not None:
@ -454,9 +454,9 @@ def compress(
)
# Build compressed messages and cache
compressed_messages: List[dict] = []
cache: Dict[str, str] = {}
used_keys: Set[str] = set()
compressed_messages: list[dict] = []
cache: dict[str, str] = {}
used_keys: set[str] = set()
dropped_tool_span_indices = _get_dropped_tool_span_indices(
kept_indices=kept_indices, tool_exchange_spans=tool_exchange_spans
)
@ -478,7 +478,7 @@ def compress(
compressed_tokens = token_counter(
model=model,
messages=cast(List[Any], compressed_messages),
messages=cast(list[Any], compressed_messages),
)
return CompressedResult(

3
litellm/compression/message_stubbing.py
View file

@ -3,7 +3,6 @@ Replace messages with compact stubs and extract human-readable keys.
"""
import re
from typing import Set
from litellm.compression.content_detection import detect_content_type
@ -17,7 +16,7 @@ _FILE_PATH_PATTERNS = [
]
def extract_key(message: dict, fallback_index: int, used_keys: Set[str]) -> str:
def extract_key(message: dict, fallback_index: int, used_keys: set[str]) -> str:
"""
Extract a human-readable key for the message.

4
litellm/compression/retrieval_tool.py
View file

@ -2,10 +2,8 @@
Build the litellm_content_retrieve tool definition for the LLM.
"""
from typing import List
def build_retrieval_tool(available_keys: List[str]) -> dict:
def build_retrieval_tool(available_keys: list[str]) -> dict:
"""
Return an OpenAI-format tool definition that lets the model
retrieve the full content of a compressed message.

15
litellm/compression/scoring/bm25.py
View file

@ -7,10 +7,9 @@ No external dependencies — uses only stdlib.
import math
import re
from collections import Counter
from typing import Dict, List
def _tokenize(text: str) -> List[str]:
def _tokenize(text: str) -> list[str]:
"""Split text into lowercase tokens on word boundaries."""
return re.findall(r"[a-z0-9_]+", text.lower())
@ -33,10 +32,10 @@ def _extract_content(message: dict) -> str:
def bm25_score_messages(
query: str,
messages: List[dict],
messages: list[dict],
k1: float = 1.5,
b: float = 0.75,
) -> List[float]:
) -> list[float]:
"""
Score each message's relevance to the query using BM25 (Okapi BM25).
@ -54,7 +53,7 @@ def bm25_score_messages(
return [0.0] * len(messages)
# Tokenize all documents
doc_tokens: List[List[str]] = []
doc_tokens: list[list[str]] = []
for msg in messages:
doc_tokens.append(_tokenize(_extract_content(msg)))
@ -67,14 +66,14 @@ def bm25_score_messages(
avgdl = sum(doc_lengths) / n if n > 0 else 1.0
# Document frequency for each term
df: Dict[str, int] = {}
df: dict[str, int] = {}
for dt in doc_tokens:
seen = set(dt)
for term in seen:
df[term] = df.get(term, 0) + 1
# IDF for query terms
idf: Dict[str, float] = {}
idf: dict[str, float] = {}
for term in set(query_terms):
term_df = df.get(term, 0)
# Standard BM25 IDF: log((N - df + 0.5) / (df + 0.5) + 1)
@ -94,7 +93,7 @@ def bm25_score_messages(
return sum(count for token, count in tf_counts.items() if token != query_term and token.startswith(query_term))
# Score each document
scores: List[float] = []
scores: list[float] = []
for i, dt in enumerate(doc_tokens):
if not dt:
scores.append(0.0)

16
litellm/compression/scoring/embedding_scorer.py
View file

@ -5,7 +5,7 @@ Computes cosine similarity between the query embedding and each message embeddin
"""
import math
from typing import Any, Dict, List, Optional
from typing import Any
from litellm.caching.dual_cache import DualCache
@ -34,7 +34,7 @@ def _truncate_text(text: str, max_chars: int = 30000) -> str:
return text[:half] + "\n...\n" + text[-half:]
def _cosine_similarity(a: List[float], b: List[float]) -> float:
def _cosine_similarity(a: list[float], b: list[float]) -> float:
"""Compute cosine similarity between two vectors."""
dot = sum(x * y for x, y in zip(a, b))
norm_a = math.sqrt(sum(x * x for x in a))
@ -46,11 +46,11 @@ def _cosine_similarity(a: List[float], b: List[float]) -> float:
def embedding_score_messages(
query: str,
messages: List[dict],
messages: list[dict],
model: str,
cache: Optional[DualCache] = None,
embedding_model_params: Optional[Dict[str, Any]] = None,
) -> List[float]:
cache: DualCache | None = None,
embedding_model_params: dict[str, Any] | None = None,
) -> list[float]:
"""
Score each message's semantic similarity to the query using embeddings.
@ -74,7 +74,7 @@ def embedding_score_messages(
# Filter out empty texts — replace with a placeholder to maintain indexing
processed_texts = [t if t.strip() else "empty" for t in texts]
kwargs: Dict[str, Any] = {
kwargs: dict[str, Any] = {
"model": model,
"input": processed_texts,
"caching": cache is not None,
@ -88,7 +88,7 @@ def embedding_score_messages(
embeddings = [item["embedding"] for item in response.data]
query_embedding = embeddings[0]
scores: List[float] = []
scores: list[float] = []
for i in range(1, len(embeddings)):
scores.append(_cosine_similarity(query_embedding, embeddings[i]))

18
litellm/constants.py
View file

@ -1,6 +1,6 @@
import os
import sys
from typing import List, Literal, Optional
from typing import Literal
from litellm.litellm_core_utils.env_utils import get_env_int, get_env_int_or_none
@ -396,14 +396,14 @@ DEFAULT_A2A_AGENT_TIMEOUT: float = float(os.getenv("DEFAULT_A2A_AGENT_TIMEOUT",
# Patterns that indicate a localhost/internal URL in A2A agent cards that should be
# replaced with the original base_url. This is a common misconfiguration where
# developers deploy agents with development URLs in their agent cards.
LOCALHOST_URL_PATTERNS: List[str] = [
LOCALHOST_URL_PATTERNS: list[str] = [
"localhost",
"127.0.0.1",
"0.0.0.0",
"[::1]", # IPv6 localhost
]
# Patterns in error messages that indicate a connection failure
CONNECTION_ERROR_PATTERNS: List[str] = [
CONNECTION_ERROR_PATTERNS: list[str] = [
"connect",
"connection",
"network",
@ -685,7 +685,7 @@ DEFAULT_CHAT_COMPLETION_PARAM_VALUES = {
"context_management": None,
}
openai_compatible_endpoints: List = [
openai_compatible_endpoints: list = [
"api.perplexity.ai",
"api.endpoints.anyscale.com/v1",
"api.deepinfra.com/v1/openai",
@ -731,7 +731,7 @@ openai_compatible_endpoints: List = [
]
openai_compatible_providers: List = [
openai_compatible_providers: list = [
"anyscale",
"groq",
"nvidia_nim",
@ -796,7 +796,7 @@ openai_compatible_providers: List = [
"darkbloom",
"meta", # Meta Model API (Muse Spark) - JSON-configured provider
]
openai_text_completion_compatible_providers: List = [ # providers that support `/v1/completions`
openai_text_completion_compatible_providers: list = [ # providers that support `/v1/completions`
"together_ai",
"fireworks_ai",
"hosted_vllm",
@ -819,7 +819,7 @@ openai_text_completion_compatible_providers: List = [ # providers that support
"hyperbolic",
"wandb",
]
_openai_like_providers: List = [
_openai_like_providers: list = [
"predibase",
"databricks",
"lemonade",
@ -1362,7 +1362,7 @@ try:
_raw_background_health_check_max_tokens = (
_background_health_check_max_tokens_env.strip() if _background_health_check_max_tokens_env is not None else ""
)
BACKGROUND_HEALTH_CHECK_MAX_TOKENS: Optional[int] = (
BACKGROUND_HEALTH_CHECK_MAX_TOKENS: int | None = (
int(_raw_background_health_check_max_tokens) if _raw_background_health_check_max_tokens else None
)
except (ValueError, TypeError):
@ -1376,7 +1376,7 @@ try:
if _background_health_check_max_tokens_reasoning_env is not None
else ""
)
BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING: Optional[int] = (
BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING: int | None = (
int(_raw_background_health_check_max_tokens_reasoning)
if _raw_background_health_check_max_tokens_reasoning
else None

36
litellm/containers/endpoint_factory.py
View file

@ -11,7 +11,7 @@ import json
from collections.abc import Callable
from functools import partial
from pathlib import Path
from typing import Any, Dict, List, Literal, Optional, Type
from typing import Any, Literal
import litellm
from litellm.constants import request_timeout as DEFAULT_REQUEST_TIMEOUT
@ -28,21 +28,21 @@ from litellm.types.router import GenericLiteLLMParams
from litellm.utils import ProviderConfigManager, client
# Response type mapping
RESPONSE_TYPES: Dict[str, Type] = {
RESPONSE_TYPES: dict[str, type] = {
"ContainerFileListResponse": ContainerFileListResponse,
"ContainerFileObject": ContainerFileObject,
"DeleteContainerFileResponse": DeleteContainerFileResponse,
}
def _load_endpoints_config() -> Dict:
def _load_endpoints_config() -> dict:
"""Load the endpoints configuration from JSON file."""
config_path = Path(__file__).parent / "endpoints.json"
with open(config_path) as f:
return json.load(f)
def create_sync_endpoint_function(endpoint_config: Dict) -> Callable:
def create_sync_endpoint_function(endpoint_config: dict) -> Callable:
"""
Create a sync SDK function from endpoint config.
@ -56,16 +56,16 @@ def create_sync_endpoint_function(endpoint_config: Dict) -> Callable:
def endpoint_func(
timeout: int = 600,
custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai",
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
**kwargs,
):
local_vars = locals()
try:
resolved_custom_llm_provider: str = custom_llm_provider
litellm_logging_obj: LiteLLMLoggingObj = kwargs.pop("litellm_logging_obj")
litellm_call_id: Optional[str] = kwargs.get("litellm_call_id")
litellm_call_id: str | None = kwargs.get("litellm_call_id")
_is_async = kwargs.pop("async_call", False) is True
# Check for mock response
@ -91,10 +91,8 @@ def create_sync_endpoint_function(endpoint_config: Dict) -> Callable:
custom_llm_provider=resolved_custom_llm_provider,
litellm_params=litellm_params,
)
container_provider_config: Optional[BaseContainerConfig] = (
ProviderConfigManager.get_provider_container_config(
provider=litellm.LlmProviders(resolved_custom_llm_provider),
)
container_provider_config: BaseContainerConfig | None = ProviderConfigManager.get_provider_container_config(
provider=litellm.LlmProviders(resolved_custom_llm_provider),
)
if container_provider_config is None:
@ -139,7 +137,7 @@ def create_sync_endpoint_function(endpoint_config: Dict) -> Callable:
def create_async_endpoint_function(
sync_func: Callable,
endpoint_config: Dict,
endpoint_config: dict,
) -> Callable:
"""Create an async SDK function that wraps the sync function."""
@ -147,9 +145,9 @@ def create_async_endpoint_function(
async def async_endpoint_func(
timeout: int = 600,
custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai",
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
**kwargs,
):
local_vars = locals()
@ -189,7 +187,7 @@ def create_async_endpoint_function(
return async_endpoint_func
def generate_container_endpoints() -> Dict[str, Callable]:
def generate_container_endpoints() -> dict[str, Callable]:
"""
Generate all container endpoint functions from the JSON config.
@ -210,7 +208,7 @@ def generate_container_endpoints() -> Dict[str, Callable]:
return endpoints
def get_all_endpoint_names() -> List[str]:
def get_all_endpoint_names() -> list[str]:
"""Get all endpoint names (sync and async) from config."""
config = _load_endpoints_config()
names = []
@ -220,7 +218,7 @@ def get_all_endpoint_names() -> List[str]:
return names
def get_async_endpoint_names() -> List[str]:
def get_async_endpoint_names() -> list[str]:
"""Get all async endpoint names for router registration."""
config = _load_endpoints_config()
return [endpoint["async_name"] for endpoint in config["endpoints"]]

302
litellm/containers/main.py
View file

@ -3,7 +3,7 @@ import contextvars
import json
from collections.abc import Coroutine
from functools import partial
from typing import Any, Dict, List, Literal, Optional, Union, overload
from typing import Any, Literal, overload
import litellm
from litellm.constants import request_timeout as DEFAULT_REQUEST_TIMEOUT
@ -48,16 +48,16 @@ __all__ = [
@client
async def acreate_container(
name: str,
expires_after: Optional[Dict[str, Any]] = None,
file_ids: Optional[List[str]] = None,
expires_after: dict[str, Any] | None = None,
file_ids: list[str] | None = None,
timeout=600, # default to 10 minutes
# LiteLLM specific params,
custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai",
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
# The extra values given here take precedence over values defined on the client or passed to this method.
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
**kwargs,
) -> ContainerObject:
"""Asynchronously calls the `create_container` function with the given arguments and keyword arguments.
@ -120,12 +120,12 @@ async def acreate_container(
@overload
def create_container(
name: str,
expires_after: Optional[Dict[str, Any]] = None,
file_ids: Optional[List[str]] = None,
expires_after: dict[str, Any] | None = None,
file_ids: list[str] | None = None,
timeout=600, # default to 10 minutes
api_key: Optional[str] = None,
api_base: Optional[str] = None,
api_version: Optional[str] = None,
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai",
*,
acreate_container: Literal[True],
@ -137,12 +137,12 @@ def create_container(
@overload
def create_container(
name: str,
expires_after: Optional[Dict[str, Any]] = None,
file_ids: Optional[List[str]] = None,
expires_after: dict[str, Any] | None = None,
file_ids: list[str] | None = None,
timeout=600, # default to 10 minutes
api_key: Optional[str] = None,
api_base: Optional[str] = None,
api_version: Optional[str] = None,
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai",
*,
acreate_container: Literal[False] = False,
@ -156,23 +156,20 @@ def create_container(
@client
def create_container(
name: str,
expires_after: Optional[Dict[str, Any]] = None,
file_ids: Optional[List[str]] = None,
expires_after: dict[str, Any] | None = None,
file_ids: list[str] | None = None,
timeout=600, # default to 10 minutes
api_key: Optional[str] = None,
api_base: Optional[str] = None,
api_version: Optional[str] = None,
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai",
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
# The extra values given here take precedence over values defined on the client or passed to this method.
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
**kwargs,
) -> Union[
ContainerObject,
Coroutine[Any, Any, ContainerObject],
]:
) -> ContainerObject | Coroutine[Any, Any, ContainerObject]:
"""Create a container using the OpenAI Container API.
Currently supports OpenAI
@ -191,7 +188,7 @@ def create_container(
local_vars = locals()
try:
litellm_logging_obj: LiteLLMLoggingObj = kwargs.pop("litellm_logging_obj") # type: ignore
litellm_call_id: Optional[str] = kwargs.get("litellm_call_id")
litellm_call_id: str | None = kwargs.get("litellm_call_id")
_is_async = kwargs.pop("async_call", False) is True
# Check for mock response first
@ -212,7 +209,7 @@ def create_container(
**kwargs,
)
# get provider config
container_provider_config: Optional[BaseContainerConfig] = ProviderConfigManager.get_provider_container_config(
container_provider_config: BaseContainerConfig | None = ProviderConfigManager.get_provider_container_config(
provider=litellm.LlmProviders(custom_llm_provider),
)
@ -226,7 +223,7 @@ def create_container(
)
# Get optional parameters for the container API
container_create_request_params: Dict = ContainerRequestUtils.get_optional_params_container_create(
container_create_request_params: dict = ContainerRequestUtils.get_optional_params_container_create(
container_provider_config=container_provider_config,
container_create_optional_params=container_create_optional_params,
)
@ -281,16 +278,16 @@ def create_container(
##### Container List #######################
@client
async def alist_containers(
after: Optional[str] = None,
limit: Optional[int] = None,
order: Optional[str] = None,
after: str | None = None,
limit: int | None = None,
order: str | None = None,
timeout=600, # default to 10 minutes
custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai",
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
# The extra values given here take precedence over values defined on the client or passed to this method.
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
**kwargs,
) -> ContainerListResponse:
"""Asynchronously list containers.
@ -351,13 +348,13 @@ async def alist_containers(
@overload
def list_containers(
after: Optional[str] = None,
limit: Optional[int] = None,
order: Optional[str] = None,
after: str | None = None,
limit: int | None = None,
order: str | None = None,
timeout=600, # default to 10 minutes
api_key: Optional[str] = None,
api_base: Optional[str] = None,
api_version: Optional[str] = None,
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai",
*,
alist_containers: Literal[True],
@ -368,13 +365,13 @@ def list_containers(
@overload
def list_containers(
after: Optional[str] = None,
limit: Optional[int] = None,
order: Optional[str] = None,
after: str | None = None,
limit: int | None = None,
order: str | None = None,
timeout=600, # default to 10 minutes
api_key: Optional[str] = None,
api_base: Optional[str] = None,
api_version: Optional[str] = None,
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai",
*,
alist_containers: Literal[False] = False,
@ -387,24 +384,21 @@ def list_containers(
@client
def list_containers(
after: Optional[str] = None,
limit: Optional[int] = None,
order: Optional[str] = None,
after: str | None = None,
limit: int | None = None,
order: str | None = None,
timeout=600, # default to 10 minutes
api_key: Optional[str] = None,
api_base: Optional[str] = None,
api_version: Optional[str] = None,
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai",
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
# The extra values given here take precedence over values defined on the client or passed to this method.
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
**kwargs,
) -> Union[
ContainerListResponse,
Coroutine[Any, Any, ContainerListResponse],
]:
) -> ContainerListResponse | Coroutine[Any, Any, ContainerListResponse]:
"""List containers using the OpenAI Container API.
Currently supports OpenAI
@ -412,7 +406,7 @@ def list_containers(
local_vars = locals()
try:
litellm_logging_obj: LiteLLMLoggingObj = kwargs.pop("litellm_logging_obj") # type: ignore
litellm_call_id: Optional[str] = kwargs.get("litellm_call_id")
litellm_call_id: str | None = kwargs.get("litellm_call_id")
_is_async = kwargs.pop("async_call", False) is True
# Check for mock response first
@ -433,7 +427,7 @@ def list_containers(
**kwargs,
)
# get provider config
container_provider_config: Optional[BaseContainerConfig] = ProviderConfigManager.get_provider_container_config(
container_provider_config: BaseContainerConfig | None = ProviderConfigManager.get_provider_container_config(
provider=litellm.LlmProviders(custom_llm_provider),
)
@ -491,9 +485,9 @@ async def aretrieve_container(
custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai",
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
# The extra values given here take precedence over values defined on the client or passed to this method.
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
**kwargs,
) -> ContainerObject:
"""Asynchronously retrieve a container.
@ -552,9 +546,9 @@ async def aretrieve_container(
def retrieve_container(
container_id: str,
timeout=600, # default to 10 minutes
api_key: Optional[str] = None,
api_base: Optional[str] = None,
api_version: Optional[str] = None,
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai",
*,
aretrieve_container: Literal[True],
@ -567,9 +561,9 @@ def retrieve_container(
def retrieve_container(
container_id: str,
timeout=600, # default to 10 minutes
api_key: Optional[str] = None,
api_base: Optional[str] = None,
api_version: Optional[str] = None,
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai",
*,
aretrieve_container: Literal[False] = False,
@ -584,20 +578,17 @@ def retrieve_container(
def retrieve_container(
container_id: str,
timeout=600, # default to 10 minutes
api_key: Optional[str] = None,
api_base: Optional[str] = None,
api_version: Optional[str] = None,
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai",
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
# The extra values given here take precedence over values defined on the client or passed to this method.
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
**kwargs,
) -> Union[
ContainerObject,
Coroutine[Any, Any, ContainerObject],
]:
) -> ContainerObject | Coroutine[Any, Any, ContainerObject]:
"""Retrieve a container using the OpenAI Container API.
Currently supports OpenAI
@ -606,7 +597,7 @@ def retrieve_container(
try:
resolved_custom_llm_provider: str = custom_llm_provider
litellm_logging_obj: LiteLLMLoggingObj = kwargs.pop("litellm_logging_obj") # type: ignore
litellm_call_id: Optional[str] = kwargs.get("litellm_call_id")
litellm_call_id: str | None = kwargs.get("litellm_call_id")
_is_async = kwargs.pop("async_call", False) is True
# Check for mock response first
@ -637,7 +628,7 @@ def retrieve_container(
was_encoded = original_container_id != container_id
# get provider config
container_provider_config: Optional[BaseContainerConfig] = ProviderConfigManager.get_provider_container_config(
container_provider_config: BaseContainerConfig | None = ProviderConfigManager.get_provider_container_config(
provider=litellm.LlmProviders(resolved_custom_llm_provider),
)
@ -709,9 +700,9 @@ async def adelete_container(
custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai",
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
# The extra values given here take precedence over values defined on the client or passed to this method.
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
**kwargs,
) -> DeleteContainerResult:
"""Asynchronously delete a container.
@ -770,9 +761,9 @@ async def adelete_container(
def delete_container(
container_id: str,
timeout=600, # default to 10 minutes
api_key: Optional[str] = None,
api_base: Optional[str] = None,
api_version: Optional[str] = None,
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai",
*,
adelete_container: Literal[True],
@ -785,9 +776,9 @@ def delete_container(
def delete_container(
container_id: str,
timeout=600, # default to 10 minutes
api_key: Optional[str] = None,
api_base: Optional[str] = None,
api_version: Optional[str] = None,
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai",
*,
adelete_container: Literal[False] = False,
@ -802,20 +793,17 @@ def delete_container(
def delete_container(
container_id: str,
timeout=600, # default to 10 minutes
api_key: Optional[str] = None,
api_base: Optional[str] = None,
api_version: Optional[str] = None,
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai",
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
# The extra values given here take precedence over values defined on the client or passed to this method.
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
**kwargs,
) -> Union[
DeleteContainerResult,
Coroutine[Any, Any, DeleteContainerResult],
]:
) -> DeleteContainerResult | Coroutine[Any, Any, DeleteContainerResult]:
"""Delete a container using the OpenAI Container API.
Currently supports OpenAI
@ -824,7 +812,7 @@ def delete_container(
try:
resolved_custom_llm_provider: str = custom_llm_provider
litellm_logging_obj: LiteLLMLoggingObj = kwargs.pop("litellm_logging_obj") # type: ignore
litellm_call_id: Optional[str] = kwargs.get("litellm_call_id")
litellm_call_id: str | None = kwargs.get("litellm_call_id")
_is_async = kwargs.pop("async_call", False) is True
# Check for mock response first
@ -855,7 +843,7 @@ def delete_container(
was_encoded = original_container_id != container_id
# get provider config
container_provider_config: Optional[BaseContainerConfig] = ProviderConfigManager.get_provider_container_config(
container_provider_config: BaseContainerConfig | None = ProviderConfigManager.get_provider_container_config(
provider=litellm.LlmProviders(resolved_custom_llm_provider),
)
@ -923,14 +911,14 @@ def delete_container(
@client
async def alist_container_files(
container_id: str,
after: Optional[str] = None,
limit: Optional[int] = None,
order: Optional[str] = None,
after: str | None = None,
limit: int | None = None,
order: str | None = None,
timeout=600, # default to 10 minutes
custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai",
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
**kwargs,
) -> ContainerFileListResponse:
"""Asynchronously list files in a container.
@ -994,13 +982,13 @@ async def alist_container_files(
@overload
def list_container_files(
container_id: str,
after: Optional[str] = None,
limit: Optional[int] = None,
order: Optional[str] = None,
after: str | None = None,
limit: int | None = None,
order: str | None = None,
timeout=600,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
api_version: Optional[str] = None,
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai",
*,
alist_container_files: Literal[True],
@ -1012,13 +1000,13 @@ def list_container_files(
@overload
def list_container_files(
container_id: str,
after: Optional[str] = None,
limit: Optional[int] = None,
order: Optional[str] = None,
after: str | None = None,
limit: int | None = None,
order: str | None = None,
timeout=600,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
api_version: Optional[str] = None,
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai",
*,
alist_container_files: Literal[False] = False,
@ -1032,22 +1020,19 @@ def list_container_files(
@client
def list_container_files(
container_id: str,
after: Optional[str] = None,
limit: Optional[int] = None,
order: Optional[str] = None,
after: str | None = None,
limit: int | None = None,
order: str | None = None,
timeout=600, # default to 10 minutes
api_key: Optional[str] = None,
api_base: Optional[str] = None,
api_version: Optional[str] = None,
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai",
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
**kwargs,
) -> Union[
ContainerFileListResponse,
Coroutine[Any, Any, ContainerFileListResponse],
]:
) -> ContainerFileListResponse | Coroutine[Any, Any, ContainerFileListResponse]:
"""List files in a container using the OpenAI Container API.
Currently supports OpenAI
@ -1056,7 +1041,7 @@ def list_container_files(
try:
resolved_custom_llm_provider: str = custom_llm_provider
litellm_logging_obj: LiteLLMLoggingObj = kwargs.pop("litellm_logging_obj") # type: ignore
litellm_call_id: Optional[str] = kwargs.get("litellm_call_id")
litellm_call_id: str | None = kwargs.get("litellm_call_id")
_is_async = kwargs.pop("async_call", False) is True
# Check for mock response first
@ -1085,7 +1070,7 @@ def list_container_files(
)
# get provider config
container_provider_config: Optional[BaseContainerConfig] = ProviderConfigManager.get_provider_container_config(
container_provider_config: BaseContainerConfig | None = ProviderConfigManager.get_provider_container_config(
provider=litellm.LlmProviders(resolved_custom_llm_provider),
)
@ -1142,9 +1127,9 @@ async def aupload_container_file(
file: FileTypes,
timeout=600, # default to 10 minutes
custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai",
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
**kwargs,
) -> ContainerFileObject:
"""Asynchronously upload a file to a container.
@ -1227,9 +1212,9 @@ def upload_container_file(
container_id: str,
file: FileTypes,
timeout=600,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
api_version: Optional[str] = None,
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai",
*,
aupload_container_file: Literal[True],
@ -1243,9 +1228,9 @@ def upload_container_file(
container_id: str,
file: FileTypes,
timeout=600,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
api_version: Optional[str] = None,
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai",
*,
aupload_container_file: Literal[False] = False,
@ -1261,18 +1246,15 @@ def upload_container_file(
container_id: str,
file: FileTypes,
timeout=600, # default to 10 minutes
api_key: Optional[str] = None,
api_base: Optional[str] = None,
api_version: Optional[str] = None,
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai",
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
**kwargs,
) -> Union[
ContainerFileObject,
Coroutine[Any, Any, ContainerFileObject],
]:
) -> ContainerFileObject | Coroutine[Any, Any, ContainerFileObject]:
"""Upload a file to a container using the OpenAI Container API.
This endpoint allows uploading files directly to a container session,
@ -1310,7 +1292,7 @@ def upload_container_file(
try:
resolved_custom_llm_provider: str = custom_llm_provider
litellm_logging_obj: LiteLLMLoggingObj = kwargs.pop("litellm_logging_obj") # type: ignore
litellm_call_id: Optional[str] = kwargs.get("litellm_call_id")
litellm_call_id: str | None = kwargs.get("litellm_call_id")
_is_async = kwargs.pop("async_call", False) is True
# Check for mock response first
@ -1339,7 +1321,7 @@ def upload_container_file(
)
# get provider config
container_provider_config: Optional[BaseContainerConfig] = ProviderConfigManager.get_provider_container_config(
container_provider_config: BaseContainerConfig | None = ProviderConfigManager.get_provider_container_config(
provider=litellm.LlmProviders(resolved_custom_llm_provider),
)

12
litellm/containers/utils.py
View file

@ -1,4 +1,4 @@
from typing import Any, Dict, Optional, TypeVar
from typing import Any, TypeVar
from litellm.llms.base_llm.containers.transformation import BaseContainerConfig
from litellm.responses.utils import ResponsesAPIRequestUtils
@ -61,7 +61,7 @@ class ContainerRequestUtils:
def get_optional_params_container_create(
container_provider_config: BaseContainerConfig,
container_create_optional_params: ContainerCreateOptionalRequestParams,
) -> Dict:
) -> dict:
"""Get the optional parameters for container creation."""
supported_params = container_provider_config.get_supported_openai_params()
@ -97,9 +97,9 @@ class ContainerRequestUtils:
@staticmethod
def encode_container_id_in_response(
response_obj: T,
custom_llm_provider: Optional[str],
litellm_metadata: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
custom_llm_provider: str | None,
litellm_metadata: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
) -> T:
"""
Encode container_id in response object with provider/model metadata for routing.
@ -124,7 +124,7 @@ class ContainerRequestUtils:
"""
# Extract model_id from litellm_metadata
litellm_metadata = litellm_metadata or {}
model_info: Dict[str, Any] = litellm_metadata.get("model_info", {}) or {}
model_info: dict[str, Any] = litellm_metadata.get("model_info", {}) or {}
model_id = model_info.get("id")
# Check if we should encode based on routing metadata

381
litellm/cost_calculator.py
View file

@ -3,7 +3,7 @@
import logging
import time
from functools import lru_cache
from typing import TYPE_CHECKING, Any, List, Literal, Optional, Tuple, Union, cast
from typing import TYPE_CHECKING, Any, Literal, cast
from httpx import Response
from pydantic import BaseModel
@ -29,8 +29,8 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import (
_parse_prompt_tokens_details,
calculate_cost_component,
generic_cost_per_token,
get_token_type_cost_breakdown,
get_billable_input_tokens,
get_token_type_cost_breakdown,
select_cost_metric_for_model,
)
from litellm.llms.anthropic.cost_calculation import (
@ -52,9 +52,6 @@ from litellm.llms.databricks.cost_calculator import (
from litellm.llms.deepseek.cost_calculator import (
cost_per_token as deepseek_cost_per_token,
)
from litellm.llms.tencent.cost_calculator import (
cost_per_token as tencent_cost_per_token,
)
from litellm.llms.fireworks_ai.cost_calculator import (
cost_per_token as fireworks_ai_cost_per_token,
)
@ -64,12 +61,19 @@ from litellm.llms.lemonade.cost_calculator import (
)
from litellm.llms.openai.cost_calculation import (
_video_output_cost_per_second,
)
from litellm.llms.openai.cost_calculation import (
cost_per_second as openai_cost_per_second,
)
from litellm.llms.openai.cost_calculation import (
cost_per_token as openai_cost_per_token,
)
from litellm.llms.perplexity.cost_calculator import (
cost_per_token as perplexity_cost_per_token,
)
from litellm.llms.tencent.cost_calculator import (
cost_per_token as tencent_cost_per_token,
)
from litellm.llms.together_ai.cost_calculator import get_model_params_and_category
from litellm.llms.vertex_ai.cost_calculator import (
cost_per_character as google_cost_per_character,
@ -180,13 +184,13 @@ _MCP_CALL_TYPE = CallTypes.call_mcp_tool.value
def _cost_per_token_custom_pricing_helper(
prompt_tokens: float = 0,
completion_tokens: float = 0,
response_time_ms: Optional[float] = 0.0,
response_time_ms: float | None = 0.0,
cached_tokens: float = 0,
cache_creation_tokens: float = 0,
### CUSTOM PRICING ###
custom_cost_per_token: Optional[CostPerToken] = None,
custom_cost_per_second: Optional[float] = None,
) -> Optional[Tuple[float, float]]:
custom_cost_per_token: CostPerToken | None = None,
custom_cost_per_second: float | None = None,
) -> tuple[float, float] | None:
"""Internal helper function for calculating cost, if custom pricing given.
prompt_tokens is assumed to include both cached_tokens and cache_creation_tokens
@ -230,10 +234,10 @@ def _cost_per_token_custom_pricing_helper(
def _get_additional_costs(
model: str,
custom_llm_provider: Optional[str],
custom_llm_provider: str | None,
prompt_tokens: int,
completion_tokens: int,
) -> Optional[dict]:
) -> dict | None:
"""
Calculate additional costs beyond standard token costs.
@ -275,7 +279,7 @@ def _get_additional_costs(
def _transcription_usage_has_token_details(
usage_block: Optional[Usage],
usage_block: Usage | None,
) -> bool:
if usage_block is None:
return False
@ -297,35 +301,35 @@ def cost_per_token(
model: str = "",
prompt_tokens: int = 0,
completion_tokens: int = 0,
response_time_ms: Optional[float] = 0.0,
custom_llm_provider: Optional[str] = None,
response_time_ms: float | None = 0.0,
custom_llm_provider: str | None = None,
region_name=None,
### CHARACTER PRICING ###
prompt_characters: Optional[int] = None,
completion_characters: Optional[int] = None,
prompt_characters: int | None = None,
completion_characters: int | None = None,
### PROMPT CACHING PRICING ### - used for anthropic
cache_creation_input_tokens: Optional[int] = 0,
cache_read_input_tokens: Optional[int] = 0,
cache_creation_input_tokens: int | None = 0,
cache_read_input_tokens: int | None = 0,
### CUSTOM PRICING ###
custom_cost_per_token: Optional[CostPerToken] = None,
custom_cost_per_second: Optional[float] = None,
custom_cost_per_token: CostPerToken | None = None,
custom_cost_per_second: float | None = None,
### NUMBER OF QUERIES ###
number_of_queries: Optional[int] = None,
number_of_queries: int | None = None,
### USAGE OBJECT ###
usage_object: Optional[Usage] = None, # just read the usage object if provided
usage_object: Usage | None = None, # just read the usage object if provided
### BILLED UNITS ###
rerank_billed_units: Optional[RerankBilledUnits] = None,
rerank_billed_units: RerankBilledUnits | None = None,
### CALL TYPE ###
call_type: CallTypesLiteral = "completion",
audio_transcription_file_duration: float = 0.0, # for audio transcription calls - the file time in seconds
### SERVICE TIER ###
service_tier: Optional[str] = None, # for OpenAI service tier pricing
service_tier: str | None = None, # for OpenAI service tier pricing
### DATA RESIDENCY ###
data_residency: Optional[str] = None, # for OpenAI regional-processing uplift (e.g. "eu", "us")
response: Optional[Any] = None,
data_residency: str | None = None, # for OpenAI regional-processing uplift (e.g. "eu", "us")
response: Any | None = None,
### REQUEST MODEL ###
request_model: Optional[str] = None, # original request model for router detection
) -> Tuple[float, float]: # type: ignore
request_model: str | None = None, # original request model for router detection
) -> tuple[float, float]: # type: ignore
"""
Calculates the cost per token for a given model, prompt tokens, and completion tokens.
@ -489,12 +493,7 @@ def cost_per_token(
if cost_metric == "cost_per_character":
if prompt_characters is None:
raise ValueError(
"prompt_characters must be provided for tts calls. prompt_characters={}, model={}, custom_llm_provider={}, call_type={}".format(
prompt_characters,
model,
custom_llm_provider,
call_type,
)
f"prompt_characters must be provided for tts calls. prompt_characters={prompt_characters}, model={model}, custom_llm_provider={custom_llm_provider}, call_type={call_type}"
)
_prompt_cost, _completion_cost = _generic_cost_per_character(
model=model_without_prefix,
@ -506,14 +505,7 @@ def cost_per_token(
)
if _prompt_cost is None or _completion_cost is None:
raise ValueError(
"cost for tts call is None. prompt_cost={}, completion_cost={}, model={}, custom_llm_provider={}, prompt_characters={}, completion_characters={}".format(
_prompt_cost,
_completion_cost,
model_without_prefix,
custom_llm_provider,
prompt_characters,
completion_characters,
)
f"cost for tts call is None. prompt_cost={_prompt_cost}, completion_cost={_completion_cost}, model={model_without_prefix}, custom_llm_provider={custom_llm_provider}, prompt_characters={prompt_characters}, completion_characters={completion_characters}"
)
prompt_cost = _prompt_cost
completion_cost = _completion_cost
@ -712,9 +704,9 @@ def has_hidden_params(obj: Any) -> bool:
def _get_provider_for_cost_calc(
model: Optional[str],
custom_llm_provider: Optional[str] = None,
) -> Optional[str]:
model: str | None,
custom_llm_provider: str | None = None,
) -> str | None:
if custom_llm_provider is not None:
return custom_llm_provider
if model is None:
@ -723,7 +715,7 @@ def _get_provider_for_cost_calc(
_, custom_llm_provider, _, _ = litellm.get_llm_provider(model=model)
except Exception as e:
verbose_logger.debug(
f"litellm.cost_calculator.py::_get_provider_for_cost_calc() - Error inferring custom_llm_provider - {str(e)}"
f"litellm.cost_calculator.py::_get_provider_for_cost_calc() - Error inferring custom_llm_provider - {e!s}"
)
return None
@ -731,13 +723,13 @@ def _get_provider_for_cost_calc(
def _select_model_name_for_cost_calc(
model: Optional[str],
completion_response: Optional[Any],
base_model: Optional[str] = None,
custom_pricing: Optional[bool] = None,
custom_llm_provider: Optional[str] = None,
router_model_id: Optional[str] = None,
) -> Optional[str]:
model: str | None,
completion_response: Any | None,
base_model: str | None = None,
custom_pricing: bool | None = None,
custom_llm_provider: str | None = None,
router_model_id: str | None = None,
) -> str | None:
"""
1. If custom pricing is true, return received model name
2. If base_model is set (e.g. for azure models), return that
@ -745,17 +737,17 @@ def _select_model_name_for_cost_calc(
4. Check if model is passed in return that
"""
return_model: Optional[str] = None
region_name: Optional[str] = None
return_model: str | None = None
region_name: str | None = None
custom_llm_provider = _get_provider_for_cost_calc(model=model, custom_llm_provider=custom_llm_provider)
completion_response_model: Optional[str] = None
completion_response_model: str | None = None
if completion_response is not None:
if isinstance(completion_response, BaseModel):
completion_response_model = getattr(completion_response, "model", None)
elif isinstance(completion_response, dict):
completion_response_model = completion_response.get("model", None)
hidden_params: Optional[dict] = getattr(completion_response, "_hidden_params", None)
hidden_params: dict | None = getattr(completion_response, "_hidden_params", None)
if custom_pricing is True:
if router_model_id is not None and router_model_id in litellm.model_cost:
@ -808,7 +800,7 @@ def _model_contains_known_llm_provider(model: str) -> bool:
return _provider_prefix in LlmProvidersSet
def _get_response_model(completion_response: Any) -> Optional[str]:
def _get_response_model(completion_response: Any) -> str | None:
"""
Extract the model name from a completion response object.
@ -837,7 +829,7 @@ _GEMINI_TRAFFIC_TYPE_TO_SERVICE_TIER: dict = {
}
def _map_traffic_type_to_service_tier(traffic_type: Optional[str]) -> Optional[str]:
def _map_traffic_type_to_service_tier(traffic_type: str | None) -> str | None:
"""
Map a Gemini usageMetadata.trafficType value to a LiteLLM service_tier string.
@ -872,9 +864,9 @@ def _normalize_service_tier(service_tier: object) -> str | None:
def _get_usage_object(
completion_response: Any,
) -> Optional[Usage]:
) -> Usage | None:
usage_obj = cast(
Union[Usage, ResponseAPIUsage, dict, BaseModel],
Usage | ResponseAPIUsage | dict | BaseModel,
(
completion_response.get("usage")
if isinstance(completion_response, dict)
@ -895,7 +887,7 @@ def _get_usage_object(
elif TranscriptionUsageObjectTransformation.is_transcription_usage_object(usage_obj):
return TranscriptionUsageObjectTransformation.transform_transcription_usage_object(
cast(
Union[TranscriptionUsageDurationObject, TranscriptionUsageTokensObject],
TranscriptionUsageDurationObject | TranscriptionUsageTokensObject,
usage_obj,
)
)
@ -917,7 +909,7 @@ def _is_known_usage_objects(usage_obj):
)
def _infer_call_type(call_type: Optional[CallTypesLiteral], completion_response: Any) -> Optional[CallTypesLiteral]:
def _infer_call_type(call_type: CallTypesLiteral | None, completion_response: Any) -> CallTypesLiteral | None:
if call_type is not None:
return call_type
@ -946,8 +938,8 @@ def _infer_call_type(call_type: Optional[CallTypesLiteral], completion_response:
def _apply_cost_discount(
base_cost: float,
custom_llm_provider: Optional[str],
) -> Tuple[float, float, float]:
custom_llm_provider: str | None,
) -> tuple[float, float, float]:
"""
Apply provider-specific cost discount from module-level config.
@ -980,8 +972,8 @@ def _apply_cost_discount(
def _apply_cost_margin(
base_cost: float,
custom_llm_provider: Optional[str],
) -> Tuple[float, float, float, float]:
custom_llm_provider: str | None,
) -> tuple[float, float, float, float]:
"""
Apply provider-specific or global cost margin from module-level config.
@ -1044,21 +1036,21 @@ def _apply_cost_margin(
def _store_cost_breakdown_in_logging_obj(
litellm_logging_obj: Optional[LitellmLoggingObject],
litellm_logging_obj: LitellmLoggingObject | None,
prompt_tokens_cost_usd_dollar: float,
completion_tokens_cost_usd_dollar: float,
cost_for_built_in_tools_cost_usd_dollar: float,
total_cost_usd_dollar: float,
additional_costs: Optional[dict] = None,
original_cost: Optional[float] = None,
discount_percent: Optional[float] = None,
discount_amount: Optional[float] = None,
margin_percent: Optional[float] = None,
margin_fixed_amount: Optional[float] = None,
margin_total_amount: Optional[float] = None,
cache_read_cost: Optional[float] = None,
cache_creation_cost: Optional[float] = None,
reasoning_cost: Optional[float] = None,
additional_costs: dict | None = None,
original_cost: float | None = None,
discount_percent: float | None = None,
discount_amount: float | None = None,
margin_percent: float | None = None,
margin_fixed_amount: float | None = None,
margin_total_amount: float | None = None,
cache_read_cost: float | None = None,
cache_creation_cost: float | None = None,
reasoning_cost: float | None = None,
) -> None:
"""
Helper function to store cost breakdown in the logging object.
@ -1100,40 +1092,39 @@ def _store_cost_breakdown_in_logging_obj(
)
except Exception as breakdown_error:
verbose_logger.debug(f"Error storing cost breakdown: {str(breakdown_error)}")
verbose_logger.debug(f"Error storing cost breakdown: {breakdown_error!s}")
# Don't fail the main cost calculation if breakdown storage fails
pass
def completion_cost(
completion_response=None,
model: Optional[str] = None,
model: str | None = None,
prompt="",
messages: List = [],
messages: list = [],
completion="",
total_time: Optional[float] = 0.0, # used for replicate, sagemaker
call_type: Optional[CallTypesLiteral] = None,
total_time: float | None = 0.0, # used for replicate, sagemaker
call_type: CallTypesLiteral | None = None,
### REGION ###
custom_llm_provider=None,
region_name=None, # used for bedrock pricing
### IMAGE GEN ###
size: Optional[str] = None,
quality: Optional[str] = None,
n: Optional[int] = None, # number of images
size: str | None = None,
quality: str | None = None,
n: int | None = None, # number of images
### CUSTOM PRICING ###
custom_cost_per_token: Optional[CostPerToken] = None,
custom_cost_per_second: Optional[float] = None,
optional_params: Optional[dict] = None,
custom_pricing: Optional[bool] = None,
base_model: Optional[str] = None,
standard_built_in_tools_params: Optional[StandardBuiltInToolsParams] = None,
litellm_model_name: Optional[str] = None,
router_model_id: Optional[str] = None,
litellm_logging_obj: Optional[LitellmLoggingObject] = None,
custom_cost_per_token: CostPerToken | None = None,
custom_cost_per_second: float | None = None,
optional_params: dict | None = None,
custom_pricing: bool | None = None,
base_model: str | None = None,
standard_built_in_tools_params: StandardBuiltInToolsParams | None = None,
litellm_model_name: str | None = None,
router_model_id: str | None = None,
litellm_logging_obj: LitellmLoggingObject | None = None,
### SERVICE TIER ###
service_tier: Optional[str] = None, # for OpenAI service tier pricing
service_tier: str | None = None, # for OpenAI service tier pricing
### DATA RESIDENCY ###
data_residency: Optional[str] = None, # for OpenAI regional-processing uplift (e.g. "eu", "us")
data_residency: str | None = None, # for OpenAI regional-processing uplift (e.g. "eu", "us")
) -> float:
"""
Calculate the cost of a given completion call fot GPT-3.5-turbo, llama2, any litellm supported llm.
@ -1176,14 +1167,14 @@ def completion_cost(
model = "dall-e-2" # for dall-e-2, azure expects an empty model name
# Handle Inputs to completion_cost
prompt_tokens = 0
prompt_characters: Optional[int] = None
prompt_characters: int | None = None
completion_tokens = 0
completion_characters: Optional[int] = None
cache_creation_input_tokens: Optional[int] = None
cache_read_input_tokens: Optional[int] = None
completion_characters: int | None = None
cache_creation_input_tokens: int | None = None
cache_read_input_tokens: int | None = None
audio_transcription_file_duration: float = 0.0
cost_per_token_usage_object: Optional[Usage] = _get_usage_object(completion_response=completion_response)
rerank_billed_units: Optional[RerankBilledUnits] = None
cost_per_token_usage_object: Usage | None = _get_usage_object(completion_response=completion_response)
rerank_billed_units: RerankBilledUnits | None = None
# Extract service_tier from optional_params if not provided directly
if service_tier is None and optional_params is not None:
@ -1234,7 +1225,7 @@ def completion_cost(
isinstance(completion_response, BaseModel) or isinstance(completion_response, dict)
): # tts returns a custom class
if isinstance(completion_response, dict):
usage_obj: Optional[Union[dict, Usage]] = completion_response.get("usage", {})
usage_obj: dict | Usage | None = completion_response.get("usage", {})
else:
usage_obj = getattr(completion_response, "usage", {})
if isinstance(usage_obj, BaseModel) and not _is_known_usage_objects(usage_obj=usage_obj):
@ -1258,10 +1249,7 @@ def completion_cost(
elif TranscriptionUsageObjectTransformation.is_transcription_usage_object(_usage):
tr_usage = TranscriptionUsageObjectTransformation.transform_transcription_usage_object(
cast(
Union[
TranscriptionUsageDurationObject,
TranscriptionUsageTokensObject,
],
TranscriptionUsageDurationObject | TranscriptionUsageTokensObject,
_usage,
)
)
@ -1327,9 +1315,7 @@ def completion_cost(
) # strip the llm provider from the model name -> for image gen cost calculation
except Exception as e:
verbose_logger.debug(
"litellm.cost_calculator.py::completion_cost() - Error inferring custom_llm_provider - {}".format(
str(e)
)
f"litellm.cost_calculator.py::completion_cost() - Error inferring custom_llm_provider - {e!s}"
)
if CostCalculatorUtils._call_type_has_image_response(call_type) and isinstance(
completion_response, ImageResponse
@ -1348,7 +1334,7 @@ def completion_cost(
elif call_type in _VIDEO_CALL_TYPES:
### VIDEO GENERATION COST CALCULATION ###
# Extract custom model_info for deployment-specific pricing
_video_model_info: Optional[ModelInfo] = None
_video_model_info: ModelInfo | None = None
if custom_pricing and litellm_logging_obj is not None:
_litellm_params = getattr(litellm_logging_obj, "litellm_params", None)
if _litellm_params is not None:
@ -1356,8 +1342,8 @@ def completion_cost(
_video_model_info = _metadata.get("model_info", None)
usage_obj = getattr(completion_response, "usage", None)
duration_seconds: Optional[float] = None
video_resolution: Optional[str] = None
duration_seconds: float | None = None
video_resolution: str | None = None
if completion_response is not None and usage_obj:
# Handle both dict and Pydantic Usage object
if isinstance(usage_obj, dict):
@ -1491,10 +1477,7 @@ def completion_cost(
):
if cost_per_token_usage_object is None or custom_llm_provider is None:
raise ValueError(
"usage object and custom_llm_provider must be provided for realtime stream cost calculation. Got cost_per_token_usage_object={}, custom_llm_provider={}".format(
cost_per_token_usage_object,
custom_llm_provider,
)
f"usage object and custom_llm_provider must be provided for realtime stream cost calculation. Got cost_per_token_usage_object={cost_per_token_usage_object}, custom_llm_provider={custom_llm_provider}"
)
return handle_realtime_stream_cost_calculation(
results=completion_response.results,
@ -1577,13 +1560,15 @@ def completion_cost(
if completion_response is not None:
hidden_params = getattr(completion_response, "_hidden_params", None) or {}
hidden_model = hidden_params.get("model") or hidden_params.get("litellm_model_name")
if hidden_model and (
"model_router" in (hidden_model or "").lower()
or "model-router" in (hidden_model or "").lower()
if (
hidden_model
and (
"model_router" in (hidden_model or "").lower()
or "model-router" in (hidden_model or "").lower()
)
or model_for_additional_costs is None
):
model_for_additional_costs = hidden_model
elif model_for_additional_costs is None:
model_for_additional_costs = hidden_model
if model_for_additional_costs is None:
model_for_additional_costs = model
additional_costs = _get_additional_costs(
@ -1639,11 +1624,11 @@ def completion_cost(
# Store cost breakdown in logging object if available
if litellm_logging_obj is not None:
_reasoning_cost: Optional[float] = None
_cache_read_cost: Optional[float] = None
_cache_creation_cost: Optional[float] = None
_reasoning_cost: float | None = None
_cache_read_cost: float | None = None
_cache_creation_cost: float | None = None
if cost_per_token_usage_object is not None and model:
_breakdown_provider: Optional[str] = (
_breakdown_provider: str | None = (
custom_llm_provider if isinstance(custom_llm_provider, str) else None
)
_token_type_breakdown = get_token_type_cost_breakdown(
@ -1677,20 +1662,18 @@ def completion_cost(
return _final_cost
except Exception as e:
verbose_logger.debug(
"litellm.cost_calculator.py::completion_cost() - Error calculating cost for model={} - {}".format(
model, str(e)
)
f"litellm.cost_calculator.py::completion_cost() - Error calculating cost for model={model} - {e!s}"
)
if idx == len(potential_model_names) - 1:
raise e
raise Exception("Unable to calculat cost for received potential model names - {}".format(potential_model_names))
raise Exception(f"Unable to calculat cost for received potential model names - {potential_model_names}")
except Exception as e:
raise e
def get_response_cost_from_hidden_params(
hidden_params: Union[dict, BaseModel],
) -> Optional[float]:
hidden_params: dict | BaseModel,
) -> float | None:
if isinstance(hidden_params, BaseModel):
_hidden_params_dict = cast(BaseModel, hidden_params).model_dump()
else:
@ -1706,22 +1689,20 @@ def get_response_cost_from_hidden_params(
def response_cost_calculator(
response_object: Union[
ModelResponse,
EmbeddingResponse,
ImageResponse,
TranscriptionResponse,
TextCompletionResponse,
HttpxBinaryResponseContent,
RerankResponse,
ResponsesAPIResponse,
LiteLLMRealtimeStreamLoggingObject,
OpenAIModerationResponse,
Response,
SearchResponse,
],
response_object: ModelResponse
| EmbeddingResponse
| ImageResponse
| TranscriptionResponse
| TextCompletionResponse
| HttpxBinaryResponseContent
| RerankResponse
| ResponsesAPIResponse
| LiteLLMRealtimeStreamLoggingObject
| OpenAIModerationResponse
| Response
| SearchResponse,
model: str,
custom_llm_provider: Optional[str],
custom_llm_provider: str | None,
call_type: Literal[
"embedding",
"aembedding",
@ -1743,18 +1724,18 @@ def response_cost_calculator(
"asearch",
],
optional_params: dict,
cache_hit: Optional[bool] = None,
base_model: Optional[str] = None,
custom_pricing: Optional[bool] = None,
cache_hit: bool | None = None,
base_model: str | None = None,
custom_pricing: bool | None = None,
prompt: str = "",
standard_built_in_tools_params: Optional[StandardBuiltInToolsParams] = None,
litellm_model_name: Optional[str] = None,
router_model_id: Optional[str] = None,
litellm_logging_obj: Optional[LitellmLoggingObject] = None,
standard_built_in_tools_params: StandardBuiltInToolsParams | None = None,
litellm_model_name: str | None = None,
router_model_id: str | None = None,
litellm_logging_obj: LitellmLoggingObject | None = None,
### SERVICE TIER ###
service_tier: Optional[str] = None, # for OpenAI service tier pricing
service_tier: str | None = None, # for OpenAI service tier pricing
### DATA RESIDENCY ###
data_residency: Optional[str] = None, # for OpenAI regional-processing uplift (e.g. "eu", "us")
data_residency: str | None = None, # for OpenAI regional-processing uplift (e.g. "eu", "us")
) -> float:
"""
Returns
@ -1795,9 +1776,9 @@ def response_cost_calculator(
def ocr_cost(
model: str,
custom_llm_provider: Optional[str],
response: Optional[Any] = None,
) -> Tuple[float, float]:
custom_llm_provider: str | None,
response: Any | None = None,
) -> tuple[float, float]:
"""
Args:
model: str - model name
@ -1821,7 +1802,7 @@ def ocr_cost(
raise ValueError("OCR response usage_info is None")
try:
model_info: Optional[ModelInfo] = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider)
model_info: ModelInfo | None = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider)
except Exception:
model_info = None
@ -1832,7 +1813,7 @@ def ocr_cost(
if credits is not None and cost_per_credit is not None:
return cost_per_credit * credits, 0.0
ocr_cost_per_page: Optional[float] = None
ocr_cost_per_page: float | None = None
if model_info is not None:
ocr_cost_per_page = model_info.get("ocr_cost_per_page")
@ -1874,15 +1855,15 @@ def ocr_cost(
def vector_store_search_cost(
model: Optional[str],
model: str | None,
custom_llm_provider: str,
response: VectorStoreSearchResponse,
) -> Tuple[float, float]:
) -> tuple[float, float]:
"""
Returns
- float or None: cost of vector store search
"""
api_type: Optional[str] = None
api_type: str | None = None
if custom_llm_provider is None:
custom_llm_provider = "openai"
@ -1907,9 +1888,9 @@ def vector_store_search_cost(
def rerank_cost(
model: str,
custom_llm_provider: Optional[str],
billed_units: Optional[RerankBilledUnits] = None,
) -> Tuple[float, float]:
custom_llm_provider: str | None,
billed_units: RerankBilledUnits | None = None,
) -> tuple[float, float]:
"""
Returns
- float or None: cost of response OR none if error.
@ -1925,9 +1906,7 @@ def rerank_cost(
)
try:
model_info: Optional[ModelInfo] = litellm.get_model_info(
model=model, custom_llm_provider=custom_llm_provider
)
model_info: ModelInfo | None = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider)
except Exception:
model_info = None
@ -1941,17 +1920,17 @@ def rerank_cost(
raise e
def transcription_cost(model: str, custom_llm_provider: Optional[str], duration: float) -> Tuple[float, float]:
def transcription_cost(model: str, custom_llm_provider: str | None, duration: float) -> tuple[float, float]:
return openai_cost_per_second(model=model, custom_llm_provider=custom_llm_provider, duration=duration)
def default_image_cost_calculator(
model: str,
custom_llm_provider: Optional[str] = None,
quality: Optional[str] = None,
n: Optional[int] = 1, # Default to 1 image
size: Optional[str] = "1024-x-1024", # OpenAI default
optional_params: Optional[dict] = None,
custom_llm_provider: str | None = None,
quality: str | None = None,
n: int | None = 1, # Default to 1 image
size: str | None = "1024-x-1024", # OpenAI default
optional_params: dict | None = None,
) -> float:
"""
Default image cost calculator for image generation
@ -1978,7 +1957,7 @@ def default_image_cost_calculator(
# Build model names for cost lookup
base_model_name = f"{size_str}/{model}"
model_name_without_custom_llm_provider: Optional[str] = None
model_name_without_custom_llm_provider: str | None = None
if custom_llm_provider and model.startswith(f"{custom_llm_provider}/"):
model_name_without_custom_llm_provider = model.replace(f"{custom_llm_provider}/", "")
base_model_name = f"{custom_llm_provider}/{size_str}/{model_name_without_custom_llm_provider}"
@ -1993,8 +1972,8 @@ def default_image_cost_calculator(
model_with_quality_without_provider = f"{quality}/{model_without_provider}" if quality else model_without_provider
# Try model with quality first, fall back to base model name
cost_info: Optional[dict] = None
models_to_check: List[Optional[str]] = [
cost_info: dict | None = None
models_to_check: list[str | None] = [
model_name_with_quality,
base_model_name,
model_name_with_v2_quality,
@ -2023,9 +2002,9 @@ def default_image_cost_calculator(
def default_video_cost_calculator(
model: str,
duration_seconds: float,
custom_llm_provider: Optional[str] = None,
model_info: Optional[ModelInfo] = None,
video_resolution: Optional[str] = None,
custom_llm_provider: str | None = None,
model_info: ModelInfo | None = None,
video_resolution: str | None = None,
) -> float:
"""
Default video cost calculator for video generation
@ -2046,13 +2025,13 @@ def default_video_cost_calculator(
Exception: If model pricing not found in cost map
"""
# Use custom model_info pricing if provided (deployment-specific pricing)
cost_info: Optional[dict] = None
cost_info: dict | None = None
if model_info is not None:
cost_info = dict(model_info)
else:
# Build model names for cost lookup
base_model_name = model
model_name_without_custom_llm_provider: Optional[str] = None
model_name_without_custom_llm_provider: str | None = None
if custom_llm_provider and model.startswith(f"{custom_llm_provider}/"):
model_name_without_custom_llm_provider = model.replace(f"{custom_llm_provider}/", "")
base_model_name = f"{custom_llm_provider}/{model_name_without_custom_llm_provider}"
@ -2062,7 +2041,7 @@ def default_video_cost_calculator(
model_without_provider = model.split("/")[-1]
# Try model with provider first, fall back to base model name
models_to_check: List[Optional[str]] = [
models_to_check: list[str | None] = [
base_model_name,
model,
model_without_provider,
@ -2101,10 +2080,10 @@ def default_video_cost_calculator(
def batch_cost_calculator(
usage: Usage,
model: str,
custom_llm_provider: Optional[str] = None,
model_info: Optional[ModelInfo] = None,
data_residency: Optional[str] = None,
) -> Tuple[float, float]:
custom_llm_provider: str | None = None,
model_info: ModelInfo | None = None,
data_residency: str | None = None,
) -> tuple[float, float]:
"""
Calculate the cost of a batch job.
@ -2191,7 +2170,7 @@ def batch_cost_calculator(
return total_prompt_cost, total_completion_cost
def _summable_prompt_token_fields(prompt_tokens_details: BaseModel) -> List[str]:
def _summable_prompt_token_fields(prompt_tokens_details: BaseModel) -> list[str]:
field_names = list(type(prompt_tokens_details).model_fields)
if getattr(prompt_tokens_details, "cache_write_tokens", None) is None:
return field_names
@ -2200,7 +2179,7 @@ def _summable_prompt_token_fields(prompt_tokens_details: BaseModel) -> List[str]
class BaseTokenUsageProcessor:
@staticmethod
def combine_usage_objects(usage_objects: List[Usage]) -> Usage:
def combine_usage_objects(usage_objects: list[Usage]) -> Usage:
"""
Combine multiple Usage objects into a single Usage object, checking model keys for nested values.
"""
@ -2272,15 +2251,15 @@ class RealtimeAPITokenUsageProcessor(BaseTokenUsageProcessor):
@staticmethod
def collect_usage_from_realtime_stream_results(
results: OpenAIRealtimeStreamList,
) -> List[Usage]:
) -> list[Usage]:
"""
Collect usage from realtime stream results
"""
response_done_events: List[OpenAIRealtimeStreamResponseBaseObject] = cast(
List[OpenAIRealtimeStreamResponseBaseObject],
response_done_events: list[OpenAIRealtimeStreamResponseBaseObject] = cast(
list[OpenAIRealtimeStreamResponseBaseObject],
[result for result in results if result["type"] == "response.done"],
)
usage_objects: List[Usage] = []
usage_objects: list[Usage] = []
for result in response_done_events:
usage_object = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(
result["response"].get("usage", {})
@ -2317,8 +2296,8 @@ def handle_realtime_stream_cost_calculation(
combined_usage_object: Usage,
custom_llm_provider: str,
litellm_model_name: str,
data_residency: Optional[str] = None,
litellm_logging_obj: Optional[LitellmLoggingObject] = None,
data_residency: str | None = None,
litellm_logging_obj: LitellmLoggingObject | None = None,
) -> float:
"""
Handles the cost calculation for realtime stream responses.
@ -2412,7 +2391,7 @@ def handle_realtime_transcription_cost_calculation(
def _get_transcription_model_name_from_results(
results: OpenAIRealtimeStreamList,
) -> Optional[str]:
) -> str | None:
"""Resolve the ASR model from a transcription_session.* / session.* event."""
for result in results:
if result.get("type") in (
@ -2431,7 +2410,7 @@ def _get_transcription_model_name_from_results(
return None
def _transcription_usage_cost(usage: dict, model_info: Optional[ModelInfo]) -> float:
def _transcription_usage_cost(usage: dict, model_info: ModelInfo | None) -> float:
if model_info is None:
return 0.0
usage_type = usage.get("type")

8
litellm/endpoints/speech/speech_to_completion_bridge/handler.py
View file

@ -2,7 +2,7 @@
Handler for transforming /chat/completions api requests to litellm.responses requests
"""
from typing import TYPE_CHECKING, Optional, Union
from typing import TYPE_CHECKING
from typing_extensions import TypedDict
@ -14,7 +14,7 @@ if TYPE_CHECKING:
class SpeechToCompletionBridgeHandlerInputKwargs(TypedDict):
model: str
input: str
voice: Optional[Union[str, dict]]
voice: str | dict | None
optional_params: dict
litellm_params: dict
logging_obj: "LiteLLMLoggingObj"
@ -79,7 +79,7 @@ class SpeechToCompletionBridgeHandler:
self,
model: str,
input: str,
voice: Optional[Union[str, dict]],
voice: str | dict | None,
optional_params: dict,
litellm_params: dict,
headers: dict,
@ -120,7 +120,7 @@ class SpeechToCompletionBridgeHandler:
model_response=result,
)
else:
raise Exception("Unmapped response type. Got type: {}".format(type(result)))
raise Exception(f"Unmapped response type. Got type: {type(result)}")
speech_to_completion_bridge_handler = SpeechToCompletionBridgeHandler()

4
litellm/endpoints/speech/speech_to_completion_bridge/transformation.py
View file

@ -1,4 +1,4 @@
from typing import TYPE_CHECKING, Optional, Union, cast
from typing import TYPE_CHECKING, cast
from litellm.constants import OPENAI_CHAT_COMPLETION_PARAMS
@ -13,7 +13,7 @@ class SpeechToCompletionBridgeTransformationHandler:
self,
model: str,
input: str,
voice: Optional[Union[str, dict]],
voice: str | dict | None,
optional_params: dict,
litellm_params: dict,
headers: dict,

20
litellm/evals/__init__.py
View file

@ -18,16 +18,16 @@ from .main import (
)
__all__ = [
"acreate_eval",
"alist_evals",
"aget_eval",
"aupdate_eval",
"adelete_eval",
"acancel_eval",
"create_eval",
"list_evals",
"get_eval",
"update_eval",
"delete_eval",
"acreate_eval",
"adelete_eval",
"aget_eval",
"alist_evals",
"aupdate_eval",
"cancel_eval",
"create_eval",
"delete_eval",
"get_eval",
"list_evals",
"update_eval",
]

328
litellm/evals/main.py
View file

@ -7,7 +7,7 @@ import asyncio
import contextvars
from collections.abc import Coroutine
from functools import partial
from typing import Any, Dict, List, Optional, Union
from typing import Any
import httpx
@ -41,15 +41,15 @@ DEFAULT_OPENAI_API_BASE = "https://api.openai.com"
@client
async def acreate_eval(
data_source_config: Dict[str, Any],
testing_criteria: List[Dict[str, Any]],
name: Optional[str] = None,
metadata: Optional[Dict[str, Any]] = None,
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
custom_llm_provider: Optional[str] = None,
data_source_config: dict[str, Any],
testing_criteria: list[dict[str, Any]],
name: str | None = None,
metadata: dict[str, Any] | None = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> Eval:
"""
@ -110,17 +110,17 @@ async def acreate_eval(
@client
def create_eval(
data_source_config: Dict[str, Any],
testing_criteria: List[Dict[str, Any]],
name: Optional[str] = None,
metadata: Optional[Dict[str, Any]] = None,
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
custom_llm_provider: Optional[str] = None,
data_source_config: dict[str, Any],
testing_criteria: list[dict[str, Any]],
name: str | None = None,
metadata: dict[str, Any] | None = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> Union[Eval, Coroutine[Any, Any, Eval]]:
) -> Eval | Coroutine[Any, Any, Eval]:
"""
Create a new evaluation
@ -142,7 +142,7 @@ def create_eval(
local_vars = locals()
try:
litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore
litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None)
litellm_call_id: str | None = kwargs.get("litellm_call_id", None)
_is_async = kwargs.pop("acreate_eval", False) is True
# Get LiteLLM parameters
@ -153,7 +153,7 @@ def create_eval(
custom_llm_provider = "openai"
# Get provider config
evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
evals_api_provider_config: BaseEvalsAPIConfig | None = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
provider=litellm.LlmProviders(custom_llm_provider),
)
@ -226,15 +226,15 @@ def create_eval(
@client
async def alist_evals(
limit: Optional[int] = None,
after: Optional[str] = None,
before: Optional[str] = None,
order: Optional[str] = None,
order_by: Optional[str] = None,
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
custom_llm_provider: Optional[str] = None,
limit: int | None = None,
after: str | None = None,
before: str | None = None,
order: str | None = None,
order_by: str | None = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> ListEvalsResponse:
"""
@ -295,17 +295,17 @@ async def alist_evals(
@client
def list_evals(
limit: Optional[int] = None,
after: Optional[str] = None,
before: Optional[str] = None,
order: Optional[str] = None,
order_by: Optional[str] = None,
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
custom_llm_provider: Optional[str] = None,
limit: int | None = None,
after: str | None = None,
before: str | None = None,
order: str | None = None,
order_by: str | None = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> Union[ListEvalsResponse, Coroutine[Any, Any, ListEvalsResponse]]:
) -> ListEvalsResponse | Coroutine[Any, Any, ListEvalsResponse]:
"""
List all evaluations
@ -327,7 +327,7 @@ def list_evals(
local_vars = locals()
try:
litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore
litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None)
litellm_call_id: str | None = kwargs.get("litellm_call_id", None)
_is_async = kwargs.pop("alist_evals", False) is True
# Get LiteLLM parameters
@ -338,7 +338,7 @@ def list_evals(
custom_llm_provider = "openai"
# Get provider config
evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
evals_api_provider_config: BaseEvalsAPIConfig | None = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
provider=litellm.LlmProviders(custom_llm_provider),
)
@ -413,10 +413,10 @@ def list_evals(
@client
async def aget_eval(
eval_id: str,
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
custom_llm_provider: Optional[str] = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> Eval:
"""
@ -470,12 +470,12 @@ async def aget_eval(
@client
def get_eval(
eval_id: str,
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
custom_llm_provider: Optional[str] = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> Union[Eval, Coroutine[Any, Any, Eval]]:
) -> Eval | Coroutine[Any, Any, Eval]:
"""
Get an evaluation by ID
@ -493,7 +493,7 @@ def get_eval(
local_vars = locals()
try:
litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore
litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None)
litellm_call_id: str | None = kwargs.get("litellm_call_id", None)
_is_async = kwargs.pop("aget_eval", False) is True
# Get LiteLLM parameters
@ -504,7 +504,7 @@ def get_eval(
custom_llm_provider = "openai"
# Get provider config
evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
evals_api_provider_config: BaseEvalsAPIConfig | None = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
provider=litellm.LlmProviders(custom_llm_provider),
)
@ -563,13 +563,13 @@ def get_eval(
@client
async def aupdate_eval(
eval_id: str,
name: Optional[str] = None,
metadata: Optional[Dict[str, Any]] = None,
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
custom_llm_provider: Optional[str] = None,
name: str | None = None,
metadata: dict[str, Any] | None = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> Eval:
"""
@ -629,15 +629,15 @@ async def aupdate_eval(
@client
def update_eval(
eval_id: str,
name: Optional[str] = None,
metadata: Optional[Dict[str, Any]] = None,
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
custom_llm_provider: Optional[str] = None,
name: str | None = None,
metadata: dict[str, Any] | None = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> Union[Eval, Coroutine[Any, Any, Eval]]:
) -> Eval | Coroutine[Any, Any, Eval]:
"""
Update an evaluation
@ -658,7 +658,7 @@ def update_eval(
local_vars = locals()
try:
litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore
litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None)
litellm_call_id: str | None = kwargs.get("litellm_call_id", None)
_is_async = kwargs.pop("aupdate_eval", False) is True
# Get LiteLLM parameters
@ -669,7 +669,7 @@ def update_eval(
custom_llm_provider = "openai"
# Get provider config
evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
evals_api_provider_config: BaseEvalsAPIConfig | None = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
provider=litellm.LlmProviders(custom_llm_provider),
)
@ -783,10 +783,10 @@ def update_eval(
@client
async def adelete_eval(
eval_id: str,
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
custom_llm_provider: Optional[str] = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> DeleteEvalResponse:
"""
@ -840,12 +840,12 @@ async def adelete_eval(
@client
def delete_eval(
eval_id: str,
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
custom_llm_provider: Optional[str] = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> Union[DeleteEvalResponse, Coroutine[Any, Any, DeleteEvalResponse]]:
) -> DeleteEvalResponse | Coroutine[Any, Any, DeleteEvalResponse]:
"""
Delete an evaluation
@ -863,7 +863,7 @@ def delete_eval(
local_vars = locals()
try:
litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore
litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None)
litellm_call_id: str | None = kwargs.get("litellm_call_id", None)
_is_async = kwargs.pop("adelete_eval", False) is True
# Get LiteLLM parameters
@ -874,7 +874,7 @@ def delete_eval(
custom_llm_provider = "openai"
# Get provider config
evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
evals_api_provider_config: BaseEvalsAPIConfig | None = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
provider=litellm.LlmProviders(custom_llm_provider),
)
@ -933,10 +933,10 @@ def delete_eval(
@client
async def acancel_eval(
eval_id: str,
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
custom_llm_provider: Optional[str] = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> CancelEvalResponse:
"""
@ -990,12 +990,12 @@ async def acancel_eval(
@client
def cancel_eval(
eval_id: str,
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
custom_llm_provider: Optional[str] = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> Union[CancelEvalResponse, Coroutine[Any, Any, CancelEvalResponse]]:
) -> CancelEvalResponse | Coroutine[Any, Any, CancelEvalResponse]:
"""
Cancel a running evaluation
@ -1013,7 +1013,7 @@ def cancel_eval(
local_vars = locals()
try:
litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore
litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None)
litellm_call_id: str | None = kwargs.get("litellm_call_id", None)
_is_async = kwargs.pop("acancel_eval", False) is True
# Get LiteLLM parameters
@ -1024,7 +1024,7 @@ def cancel_eval(
custom_llm_provider = "openai"
# Get provider config
evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
evals_api_provider_config: BaseEvalsAPIConfig | None = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
provider=litellm.LlmProviders(custom_llm_provider),
)
@ -1092,14 +1092,14 @@ def cancel_eval(
@client
async def acreate_run(
eval_id: str,
data_source: Dict[str, Any],
name: Optional[str] = None,
metadata: Optional[Dict[str, Any]] = None,
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
custom_llm_provider: Optional[str] = None,
data_source: dict[str, Any],
name: str | None = None,
metadata: dict[str, Any] | None = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> Run:
"""
@ -1161,16 +1161,16 @@ async def acreate_run(
@client
def create_run(
eval_id: str,
data_source: Dict[str, Any],
name: Optional[str] = None,
metadata: Optional[Dict[str, Any]] = None,
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
custom_llm_provider: Optional[str] = None,
data_source: dict[str, Any],
name: str | None = None,
metadata: dict[str, Any] | None = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> Union[Run, Coroutine[Any, Any, Run]]:
) -> Run | Coroutine[Any, Any, Run]:
"""
Create a new run for an evaluation
@ -1192,7 +1192,7 @@ def create_run(
local_vars = locals()
try:
litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore
litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None)
litellm_call_id: str | None = kwargs.get("litellm_call_id", None)
_is_async = kwargs.pop("acreate_run", False) is True
# Get LiteLLM parameters
@ -1203,7 +1203,7 @@ def create_run(
custom_llm_provider = "openai"
# Get provider config
evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
evals_api_provider_config: BaseEvalsAPIConfig | None = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
provider=litellm.LlmProviders(custom_llm_provider),
)
@ -1276,14 +1276,14 @@ def create_run(
@client
async def alist_runs(
eval_id: str,
limit: Optional[int] = None,
after: Optional[str] = None,
before: Optional[str] = None,
order: Optional[str] = None,
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
custom_llm_provider: Optional[str] = None,
limit: int | None = None,
after: str | None = None,
before: str | None = None,
order: str | None = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> ListRunsResponse:
"""
@ -1345,16 +1345,16 @@ async def alist_runs(
@client
def list_runs(
eval_id: str,
limit: Optional[int] = None,
after: Optional[str] = None,
before: Optional[str] = None,
order: Optional[str] = None,
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
custom_llm_provider: Optional[str] = None,
limit: int | None = None,
after: str | None = None,
before: str | None = None,
order: str | None = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> Union[ListRunsResponse, Coroutine[Any, Any, ListRunsResponse]]:
) -> ListRunsResponse | Coroutine[Any, Any, ListRunsResponse]:
"""
List all runs for an evaluation
@ -1376,7 +1376,7 @@ def list_runs(
local_vars = locals()
try:
litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore
litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None)
litellm_call_id: str | None = kwargs.get("litellm_call_id", None)
_is_async = kwargs.pop("alist_runs", False) is True
# Get LiteLLM parameters
@ -1387,7 +1387,7 @@ def list_runs(
custom_llm_provider = "openai"
# Get provider config
evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
evals_api_provider_config: BaseEvalsAPIConfig | None = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
provider=litellm.LlmProviders(custom_llm_provider),
)
@ -1462,10 +1462,10 @@ def list_runs(
async def aget_run(
eval_id: str,
run_id: str,
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
custom_llm_provider: Optional[str] = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> Run:
"""
@ -1522,12 +1522,12 @@ async def aget_run(
def get_run(
eval_id: str,
run_id: str,
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
custom_llm_provider: Optional[str] = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> Union[Run, Coroutine[Any, Any, Run]]:
) -> Run | Coroutine[Any, Any, Run]:
"""
Get a specific run
@ -1546,7 +1546,7 @@ def get_run(
local_vars = locals()
try:
litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore
litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None)
litellm_call_id: str | None = kwargs.get("litellm_call_id", None)
_is_async = kwargs.pop("aget_run", False) is True
# Get LiteLLM parameters
@ -1557,7 +1557,7 @@ def get_run(
custom_llm_provider = "openai"
# Get provider config
evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
evals_api_provider_config: BaseEvalsAPIConfig | None = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
provider=litellm.LlmProviders(custom_llm_provider),
)
@ -1618,10 +1618,10 @@ def get_run(
async def acancel_run(
eval_id: str,
run_id: str,
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
custom_llm_provider: Optional[str] = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> CancelRunResponse:
"""
@ -1678,12 +1678,12 @@ async def acancel_run(
def cancel_run(
eval_id: str,
run_id: str,
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
custom_llm_provider: Optional[str] = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> Union[CancelRunResponse, Coroutine[Any, Any, CancelRunResponse]]:
) -> CancelRunResponse | Coroutine[Any, Any, CancelRunResponse]:
"""
Cancel a running run
@ -1702,7 +1702,7 @@ def cancel_run(
local_vars = locals()
try:
litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore
litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None)
litellm_call_id: str | None = kwargs.get("litellm_call_id", None)
_is_async = kwargs.pop("acancel_run", False) is True
# Get LiteLLM parameters
@ -1713,7 +1713,7 @@ def cancel_run(
custom_llm_provider = "openai"
# Get provider config
evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
evals_api_provider_config: BaseEvalsAPIConfig | None = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
provider=litellm.LlmProviders(custom_llm_provider),
)
@ -1783,10 +1783,10 @@ def cancel_run(
async def adelete_run(
eval_id: str,
run_id: str,
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
custom_llm_provider: Optional[str] = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> RunDeleteResponse:
"""
@ -1843,12 +1843,12 @@ async def adelete_run(
def delete_run(
eval_id: str,
run_id: str,
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
custom_llm_provider: Optional[str] = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> Union[RunDeleteResponse, Coroutine[Any, Any, RunDeleteResponse]]:
) -> RunDeleteResponse | Coroutine[Any, Any, RunDeleteResponse]:
"""
Delete a run
@ -1867,7 +1867,7 @@ def delete_run(
local_vars = locals()
try:
litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore
litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None)
litellm_call_id: str | None = kwargs.get("litellm_call_id", None)
_is_async = kwargs.pop("adelete_run", False) is True
# Get LiteLLM parameters
@ -1878,7 +1878,7 @@ def delete_run(
custom_llm_provider = "openai"
# Get provider config
evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
evals_api_provider_config: BaseEvalsAPIConfig | None = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
provider=litellm.LlmProviders(custom_llm_provider),
)

244
litellm/exceptions.py
View file

@ -10,7 +10,7 @@
## LiteLLM versions of the OpenAI Exception Types
import enum
from typing import Any, Dict, Optional, Union
from typing import Any
import httpx
import openai
@ -85,7 +85,7 @@ _RATE_LIMIT_CATEGORY_VALUES = frozenset(c.value for c in RateLimitErrorCategory)
_RATE_LIMIT_TYPE_VALUES = frozenset(t.value for t in RateLimitType)
def validate_rate_limit_category(value: Any) -> Optional[str]:
def validate_rate_limit_category(value: Any) -> str | None:
"""Return ``value`` only if it matches a known :class:`RateLimitErrorCategory`.
Used at duck-typed read sites (StandardLoggingPayload extraction, Prometheus
@ -100,7 +100,7 @@ def validate_rate_limit_category(value: Any) -> Optional[str]:
return None
def validate_rate_limit_type(value: Any) -> Optional[str]:
def validate_rate_limit_type(value: Any) -> str | None:
"""Return ``value`` only if it matches a known :class:`RateLimitType`.
See :func:`validate_rate_limit_category` for the rationale.
@ -112,7 +112,7 @@ def validate_rate_limit_type(value: Any) -> Optional[str]:
return None
_MINIMAL_ERROR_RESPONSE: Optional[httpx.Response] = None
_MINIMAL_ERROR_RESPONSE: httpx.Response | None = None
def _get_minimal_error_response() -> httpx.Response:
@ -132,13 +132,13 @@ class AuthenticationError(openai.AuthenticationError): # type: ignore
message,
llm_provider,
model,
response: Optional[httpx.Response] = None,
litellm_debug_info: Optional[str] = None,
max_retries: Optional[int] = None,
num_retries: Optional[int] = None,
response: httpx.Response | None = None,
litellm_debug_info: str | None = None,
max_retries: int | None = None,
num_retries: int | None = None,
):
self.status_code = 401
self.message = "litellm.AuthenticationError: {}".format(message)
self.message = f"litellm.AuthenticationError: {message}"
self.llm_provider = llm_provider
self.model = model
self.litellm_debug_info = litellm_debug_info
@ -176,13 +176,13 @@ class NotFoundError(openai.NotFoundError): # type: ignore
message,
model,
llm_provider,
response: Optional[httpx.Response] = None,
litellm_debug_info: Optional[str] = None,
max_retries: Optional[int] = None,
num_retries: Optional[int] = None,
response: httpx.Response | None = None,
litellm_debug_info: str | None = None,
max_retries: int | None = None,
num_retries: int | None = None,
):
self.status_code = 404
self.message = "litellm.NotFoundError: {}".format(message)
self.message = f"litellm.NotFoundError: {message}"
self.model = model
self.llm_provider = llm_provider
self.litellm_debug_info = litellm_debug_info
@ -219,14 +219,14 @@ class BadRequestError(openai.BadRequestError): # type: ignore
message,
model,
llm_provider,
response: Optional[httpx.Response] = None,
litellm_debug_info: Optional[str] = None,
max_retries: Optional[int] = None,
num_retries: Optional[int] = None,
body: Optional[dict] = None,
response: httpx.Response | None = None,
litellm_debug_info: str | None = None,
max_retries: int | None = None,
num_retries: int | None = None,
body: dict | None = None,
):
self.status_code = 400
self.message = "litellm.BadRequestError: {}".format(message)
self.message = f"litellm.BadRequestError: {message}"
self.model = model
self.llm_provider = llm_provider
self.litellm_debug_info = litellm_debug_info
@ -270,11 +270,11 @@ class ImageFetchError(BadRequestError):
message,
model=None,
llm_provider=None,
response: Optional[httpx.Response] = None,
litellm_debug_info: Optional[str] = None,
max_retries: Optional[int] = None,
num_retries: Optional[int] = None,
body: Optional[dict] = None,
response: httpx.Response | None = None,
litellm_debug_info: str | None = None,
max_retries: int | None = None,
num_retries: int | None = None,
body: dict | None = None,
):
super().__init__(
message=message,
@ -295,12 +295,12 @@ class UnprocessableEntityError(openai.UnprocessableEntityError): # type: ignore
model,
llm_provider,
response: httpx.Response,
litellm_debug_info: Optional[str] = None,
max_retries: Optional[int] = None,
num_retries: Optional[int] = None,
litellm_debug_info: str | None = None,
max_retries: int | None = None,
num_retries: int | None = None,
):
self.status_code = 422
self.message = "litellm.UnprocessableEntityError: {}".format(message)
self.message = f"litellm.UnprocessableEntityError: {message}"
self.model = model
self.llm_provider = llm_provider
self.litellm_debug_info = litellm_debug_info
@ -333,11 +333,11 @@ class Timeout(openai.APITimeoutError): # type: ignore
message,
model,
llm_provider,
litellm_debug_info: Optional[str] = None,
max_retries: Optional[int] = None,
num_retries: Optional[int] = None,
headers: Optional[dict] = None,
exception_status_code: Optional[int] = None,
litellm_debug_info: str | None = None,
max_retries: int | None = None,
num_retries: int | None = None,
headers: dict | None = None,
exception_status_code: int | None = None,
):
request = httpx.Request(
method="POST",
@ -345,7 +345,7 @@ class Timeout(openai.APITimeoutError): # type: ignore
)
super().__init__(request=request) # Call the base class constructor with the parameters it needs
self.status_code = exception_status_code or 408
self.message = "litellm.Timeout: {}".format(message)
self.message = f"litellm.Timeout: {message}"
self.model = model
self.llm_provider = llm_provider
self.litellm_debug_info = litellm_debug_info
@ -378,12 +378,12 @@ class PermissionDeniedError(openai.PermissionDeniedError): # type: ignore
llm_provider,
model,
response: httpx.Response,
litellm_debug_info: Optional[str] = None,
max_retries: Optional[int] = None,
num_retries: Optional[int] = None,
litellm_debug_info: str | None = None,
max_retries: int | None = None,
num_retries: int | None = None,
):
self.status_code = 403
self.message = "litellm.PermissionDeniedError: {}".format(message)
self.message = f"litellm.PermissionDeniedError: {message}"
self.llm_provider = llm_provider
self.model = model
self.litellm_debug_info = litellm_debug_info
@ -428,17 +428,17 @@ class RateLimitError(openai.RateLimitError): # type: ignore
message,
llm_provider,
model,
response: Optional[httpx.Response] = None,
litellm_debug_info: Optional[str] = None,
max_retries: Optional[int] = None,
num_retries: Optional[int] = None,
category: Union[str, RateLimitErrorCategory] = (RateLimitErrorCategory.VENDOR_RATE_LIMIT),
rate_limit_type: Optional[Union[str, RateLimitType]] = None,
headers: Optional[Dict[str, str]] = None,
response: httpx.Response | None = None,
litellm_debug_info: str | None = None,
max_retries: int | None = None,
num_retries: int | None = None,
category: str | RateLimitErrorCategory = (RateLimitErrorCategory.VENDOR_RATE_LIMIT),
rate_limit_type: str | RateLimitType | None = None,
headers: dict[str, str] | None = None,
detail: Any = None,
):
self.status_code = 429
self.message = "litellm.RateLimitError: {}".format(message)
self.message = f"litellm.RateLimitError: {message}"
self.llm_provider = llm_provider
self.model = model
self.litellm_debug_info = litellm_debug_info
@ -448,7 +448,7 @@ class RateLimitError(openai.RateLimitError): # type: ignore
# Which dimension was exceeded — request count, token count, parallel
# requests, budget, max iterations. None when the source didn't
# classify the failure (e.g. legacy vendor 429 with no header hints).
self.rate_limit_type: Optional[str] = (
self.rate_limit_type: str | None = (
rate_limit_type.value if isinstance(rate_limit_type, RateLimitType) else rate_limit_type
)
# Headers explicitly attached to the error (e.g. retry-after,
@ -465,7 +465,7 @@ class RateLimitError(openai.RateLimitError): # type: ignore
# explicitly want them; only the proxy-supplied `headers=` kwarg
# makes it onto `self.headers`.
_response_headers = getattr(response, "headers", None) if response is not None else None
self.headers: Optional[Dict[str, str]] = {k: str(v) for k, v in headers.items()} if headers else None
self.headers: dict[str, str] | None = {k: str(v) for k, v in headers.items()} if headers else None
# Mirrors FastAPI HTTPException.detail so the same instance can be
# serialized through both the ProxyException and HTTPException paths.
self.detail = detail if detail is not None else self.message
@ -507,8 +507,8 @@ class ContextWindowExceededError(BadRequestError): # type: ignore
message,
model,
llm_provider,
response: Optional[httpx.Response] = None,
litellm_debug_info: Optional[str] = None,
response: httpx.Response | None = None,
litellm_debug_info: str | None = None,
):
self.status_code = 400
self.model = model
@ -523,7 +523,7 @@ class ContextWindowExceededError(BadRequestError): # type: ignore
) # Call the base class constructor with the parameters it needs
# set after, to make it clear the raised error is a context window exceeded error
self.message = "litellm.ContextWindowExceededError: {}".format(self.message)
self.message = f"litellm.ContextWindowExceededError: {self.message}"
def __str__(self):
_message = self.message
@ -550,10 +550,10 @@ class RejectedRequestError(BadRequestError): # type: ignore
model,
llm_provider,
request_data: dict,
litellm_debug_info: Optional[str] = None,
litellm_debug_info: str | None = None,
):
self.status_code = 400
self.message = "litellm.RejectedRequestError: {}".format(message)
self.message = f"litellm.RejectedRequestError: {message}"
self.model = model
self.llm_provider = llm_provider
self.litellm_debug_info = litellm_debug_info
@ -592,13 +592,13 @@ class ContentPolicyViolationError(BadRequestError): # type: ignore
message,
model,
llm_provider,
response: Optional[httpx.Response] = None,
litellm_debug_info: Optional[str] = None,
provider_specific_fields: Optional[dict] = None,
body: Optional[dict] = None,
response: httpx.Response | None = None,
litellm_debug_info: str | None = None,
provider_specific_fields: dict | None = None,
body: dict | None = None,
):
self.status_code = 400
self.message = "litellm.ContentPolicyViolationError: {}".format(message)
self.message = f"litellm.ContentPolicyViolationError: {message}"
self.model = model
self.llm_provider = llm_provider
self.litellm_debug_info = litellm_debug_info
@ -636,13 +636,13 @@ class ServiceUnavailableError(openai.APIStatusError): # type: ignore
message,
llm_provider,
model,
response: Optional[httpx.Response] = None,
litellm_debug_info: Optional[str] = None,
max_retries: Optional[int] = None,
num_retries: Optional[int] = None,
response: httpx.Response | None = None,
litellm_debug_info: str | None = None,
max_retries: int | None = None,
num_retries: int | None = None,
):
self.status_code = 503
self.message = "litellm.ServiceUnavailableError: {}".format(message)
self.message = f"litellm.ServiceUnavailableError: {message}"
self.llm_provider = llm_provider
self.model = model
self.litellm_debug_info = litellm_debug_info
@ -684,13 +684,13 @@ class BadGatewayError(openai.APIStatusError): # type: ignore
message,
llm_provider,
model,
response: Optional[httpx.Response] = None,
litellm_debug_info: Optional[str] = None,
max_retries: Optional[int] = None,
num_retries: Optional[int] = None,
response: httpx.Response | None = None,
litellm_debug_info: str | None = None,
max_retries: int | None = None,
num_retries: int | None = None,
):
self.status_code = 502
self.message = "litellm.BadGatewayError: {}".format(message)
self.message = f"litellm.BadGatewayError: {message}"
self.llm_provider = llm_provider
self.model = model
self.litellm_debug_info = litellm_debug_info
@ -732,13 +732,13 @@ class InternalServerError(openai.InternalServerError): # type: ignore
message,
llm_provider,
model,
response: Optional[httpx.Response] = None,
litellm_debug_info: Optional[str] = None,
max_retries: Optional[int] = None,
num_retries: Optional[int] = None,
response: httpx.Response | None = None,
litellm_debug_info: str | None = None,
max_retries: int | None = None,
num_retries: int | None = None,
):
self.status_code = 500
self.message = "litellm.InternalServerError: {}".format(message)
self.message = f"litellm.InternalServerError: {message}"
self.llm_provider = llm_provider
self.model = model
self.litellm_debug_info = litellm_debug_info
@ -782,13 +782,13 @@ class APIError(openai.APIError): # type: ignore
message,
llm_provider,
model,
request: Optional[httpx.Request] = None,
litellm_debug_info: Optional[str] = None,
max_retries: Optional[int] = None,
num_retries: Optional[int] = None,
request: httpx.Request | None = None,
litellm_debug_info: str | None = None,
max_retries: int | None = None,
num_retries: int | None = None,
):
self.status_code = status_code
self.message = "litellm.APIError: {}".format(message)
self.message = f"litellm.APIError: {message}"
self.llm_provider = llm_provider
self.model = model
self.litellm_debug_info = litellm_debug_info
@ -822,12 +822,12 @@ class APIConnectionError(openai.APIConnectionError): # type: ignore
message,
llm_provider,
model,
request: Optional[httpx.Request] = None,
litellm_debug_info: Optional[str] = None,
max_retries: Optional[int] = None,
num_retries: Optional[int] = None,
request: httpx.Request | None = None,
litellm_debug_info: str | None = None,
max_retries: int | None = None,
num_retries: int | None = None,
):
self.message = "litellm.APIConnectionError: {}".format(message)
self.message = f"litellm.APIConnectionError: {message}"
self.llm_provider = llm_provider
self.model = model
self.status_code = 500
@ -861,11 +861,11 @@ class APIResponseValidationError(openai.APIResponseValidationError): # type: ig
message,
llm_provider,
model,
litellm_debug_info: Optional[str] = None,
max_retries: Optional[int] = None,
num_retries: Optional[int] = None,
litellm_debug_info: str | None = None,
max_retries: int | None = None,
num_retries: int | None = None,
):
self.message = "litellm.APIResponseValidationError: {}".format(message)
self.message = f"litellm.APIResponseValidationError: {message}"
self.llm_provider = llm_provider
self.model = model
request = httpx.Request(method="POST", url="https://api.openai.com/v1")
@ -897,9 +897,7 @@ class JSONSchemaValidationError(APIResponseValidationError):
self.raw_response = raw_response
self.schema = schema
self.model = model
message = "litellm.JSONSchemaValidationError: model={}, returned an invalid response={}, for schema={}.\nAccess raw response with `e.raw_response`".format(
model, raw_response, schema
)
message = f"litellm.JSONSchemaValidationError: model={model}, returned an invalid response={raw_response}, for schema={schema}.\nAccess raw response with `e.raw_response`"
self.message = message
super().__init__(model=model, message=message, llm_provider=llm_provider)
@ -914,16 +912,16 @@ class UnsupportedParamsError(BadRequestError):
def __init__(
self,
message,
llm_provider: Optional[str] = None,
model: Optional[str] = None,
llm_provider: str | None = None,
model: str | None = None,
status_code: int = 400,
response: Optional[httpx.Response] = None,
litellm_debug_info: Optional[str] = None,
max_retries: Optional[int] = None,
num_retries: Optional[int] = None,
response: httpx.Response | None = None,
litellm_debug_info: str | None = None,
max_retries: int | None = None,
num_retries: int | None = None,
):
self.status_code = 400
self.message = "litellm.UnsupportedParamsError: {}".format(message)
self.message = f"litellm.UnsupportedParamsError: {message}"
self.model = model
self.llm_provider = llm_provider
self.litellm_debug_info = litellm_debug_info
@ -964,10 +962,10 @@ class BudgetExceededError(Exception):
self,
current_cost: float,
max_budget: float,
message: Optional[str] = None,
llm_provider: Optional[str] = None,
entity_type: Optional[str] = None,
entity_id: Optional[str] = None,
message: str | None = None,
llm_provider: str | None = None,
entity_type: str | None = None,
entity_id: str | None = None,
):
self.current_cost = current_cost
self.max_budget = max_budget
@ -1012,13 +1010,13 @@ class MockException(openai.APIError):
message,
llm_provider,
model,
request: Optional[httpx.Request] = None,
litellm_debug_info: Optional[str] = None,
max_retries: Optional[int] = None,
num_retries: Optional[int] = None,
request: httpx.Request | None = None,
litellm_debug_info: str | None = None,
max_retries: int | None = None,
num_retries: int | None = None,
):
self.status_code = status_code
self.message = "litellm.MockException: {}".format(message)
self.message = f"litellm.MockException: {message}"
self.llm_provider = llm_provider
self.model = model
self.litellm_debug_info = litellm_debug_info
@ -1030,7 +1028,7 @@ class MockException(openai.APIError):
class LiteLLMUnknownProvider(BadRequestError):
def __init__(self, model: str, custom_llm_provider: Optional[str] = None):
def __init__(self, model: str, custom_llm_provider: str | None = None):
self.message = LiteLLMCommonStrings.llm_provider_not_provided.value.format(
model=model, custom_llm_provider=custom_llm_provider
)
@ -1043,7 +1041,7 @@ class LiteLLMUnknownProvider(BadRequestError):
class GuardrailRaisedException(Exception):
def __init__(
self,
guardrail_name: Optional[str] = None,
guardrail_name: str | None = None,
message: str = "",
should_wrap_with_default_message: bool = True,
status_code: int = 400,
@ -1059,7 +1057,7 @@ class BlockedPiiEntityError(Exception):
def __init__(
self,
entity_type: str,
guardrail_name: Optional[str] = None,
guardrail_name: str | None = None,
status_code: int = 400,
):
"""
@ -1078,11 +1076,11 @@ class MidStreamFallbackError(ServiceUnavailableError): # type: ignore
message: str,
model: str,
llm_provider: str,
original_exception: Optional[Exception] = None,
response: Optional[httpx.Response] = None,
litellm_debug_info: Optional[str] = None,
max_retries: Optional[int] = None,
num_retries: Optional[int] = None,
original_exception: Exception | None = None,
response: httpx.Response | None = None,
litellm_debug_info: str | None = None,
max_retries: int | None = None,
num_retries: int | None = None,
generated_content: str = "",
is_pre_first_chunk: bool = False,
):
@ -1142,7 +1140,7 @@ class MidStreamFallbackError(ServiceUnavailableError): # type: ignore
if self.max_retries:
_message += f", LiteLLM Max Retries: {self.max_retries}"
if self.original_exception:
_message += f" Original exception: {type(self.original_exception).__name__}: {str(self.original_exception)}"
_message += f" Original exception: {type(self.original_exception).__name__}: {self.original_exception!s}"
return _message
def __repr__(self):
@ -1166,10 +1164,10 @@ class ModifyResponseException(Exception):
self,
message: str,
model: str,
request_data: Dict[str, Any],
guardrail_name: Optional[str] = None,
detection_info: Optional[Dict[str, Any]] = None,
original_response: Optional[Any] = None,
request_data: dict[str, Any],
guardrail_name: str | None = None,
detection_info: dict[str, Any] | None = None,
original_response: Any | None = None,
):
self.message = message
self.model = model
@ -1201,9 +1199,9 @@ class SensitiveDataRouteException(Exception):
self,
route_to_model: str,
session_id: str,
guardrail_name: Optional[str] = None,
detection_info: Optional[Dict[str, Any]] = None,
message: Optional[str] = None,
guardrail_name: str | None = None,
detection_info: dict[str, Any] | None = None,
message: str | None = None,
sticky_session_routing: bool = True,
):
self.route_to_model = route_to_model

2
litellm/experimental_mcp_client/__init__.py
View file

@ -1,3 +1,3 @@
from .tools import call_openai_tool, load_mcp_tools
__all__ = ["load_mcp_tools", "call_openai_tool"]
__all__ = ["call_openai_tool", "load_mcp_tools"]

117
litellm/experimental_mcp_client/client.py
View file

@ -8,12 +8,7 @@ import os
from collections.abc import Awaitable, Callable, Generator
from typing import (
Any,
Dict,
List,
Optional,
Tuple,
TypeVar,
Union,
)
import httpx
@ -21,7 +16,7 @@ from mcp import ClientSession, ReadResourceResult, Resource, StdioServerParamete
from mcp.client.sse import sse_client
from mcp.client.stdio import stdio_client
streamable_http_client: Optional[Any] = None
streamable_http_client: Any | None = None
try:
import mcp.client.streamable_http as streamable_http_module # type: ignore
@ -58,15 +53,15 @@ def to_basic_auth(auth_value: str) -> str:
return base64.b64encode(auth_value.encode("utf-8")).decode()
def _strip_header_whitespace(headers: Dict[str, str]) -> Dict[str, str]:
def _strip_header_whitespace(headers: dict[str, str]) -> dict[str, str]:
return {
(key.strip() if isinstance(key, str) else key): (value.strip() if isinstance(value, str) else value)
for key, value in headers.items()
}
def _first_non_cancelled_cause(exc: BaseException) -> Optional[BaseException]:
queue: List[BaseException] = [exc]
def _first_non_cancelled_cause(exc: BaseException) -> BaseException | None:
queue: list[BaseException] = [exc]
while queue:
current = queue.pop(0)
nested = getattr(current, "exceptions", None)
@ -92,13 +87,13 @@ class MCPSigV4Auth(httpx.Auth):
def __init__(
self,
aws_access_key_id: Optional[str] = None,
aws_secret_access_key: Optional[str] = None,
aws_session_token: Optional[str] = None,
aws_region_name: Optional[str] = None,
aws_service_name: Optional[str] = None,
aws_role_name: Optional[str] = None,
aws_session_name: Optional[str] = None,
aws_access_key_id: str | None = None,
aws_secret_access_key: str | None = None,
aws_session_token: str | None = None,
aws_region_name: str | None = None,
aws_service_name: str | None = None,
aws_role_name: str | None = None,
aws_session_name: str | None = None,
):
try:
from botocore.credentials import Credentials
@ -140,10 +135,10 @@ class MCPSigV4Auth(httpx.Auth):
@staticmethod
def _assume_role(
aws_role_name: str,
aws_session_name: Optional[str],
aws_access_key_id: Optional[str],
aws_secret_access_key: Optional[str],
aws_session_token: Optional[str],
aws_session_name: str | None,
aws_access_key_id: str | None,
aws_secret_access_key: str | None,
aws_session_token: str | None,
aws_region_name: str,
):
"""Call STS AssumeRole and return temporary credentials."""
@ -207,47 +202,47 @@ class MCPClient:
server_url: str = "",
transport_type: MCPTransportType = MCPTransport.http,
auth_type: MCPAuthType = None,
auth_value: Optional[Union[str, Dict[str, str]]] = None,
timeout: Optional[float] = None,
stdio_config: Optional[MCPStdioConfig] = None,
extra_headers: Optional[Dict[str, str]] = None,
ssl_verify: Optional[VerifyTypes] = None,
aws_auth: Optional[httpx.Auth] = None,
resolved_auth: Optional[httpx.Auth] = None,
sampling_callback: Optional[Callable] = None,
elicitation_callback: Optional[Callable] = None,
logging_callback: Optional[Callable] = None,
auth_value: str | dict[str, str] | None = None,
timeout: float | None = None,
stdio_config: MCPStdioConfig | None = None,
extra_headers: dict[str, str] | None = None,
ssl_verify: VerifyTypes | None = None,
aws_auth: httpx.Auth | None = None,
resolved_auth: httpx.Auth | None = None,
sampling_callback: Callable | None = None,
elicitation_callback: Callable | None = None,
logging_callback: Callable | None = None,
):
self.server_url: str = server_url
self.transport_type: MCPTransport = transport_type
self.auth_type: MCPAuthType = auth_type
self.timeout: float = timeout if timeout is not None else MCP_CLIENT_TIMEOUT
self._mcp_auth_value: Optional[Union[str, Dict[str, str]]] = None
self.stdio_config: Optional[MCPStdioConfig] = stdio_config
self.extra_headers: Optional[Dict[str, str]] = extra_headers
self.ssl_verify: Optional[VerifyTypes] = ssl_verify
self._aws_auth: Optional[httpx.Auth] = aws_auth
self._mcp_auth_value: str | dict[str, str] | None = None
self.stdio_config: MCPStdioConfig | None = stdio_config
self.extra_headers: dict[str, str] | None = extra_headers
self.ssl_verify: VerifyTypes | None = ssl_verify
self._aws_auth: httpx.Auth | None = aws_auth
# A pre-resolved httpx.Auth (e.g. from the v2 credential resolver) attached to the
# upstream client's auth= slot, taking precedence over the SigV4 aws_auth.
self._resolved_auth: Optional[httpx.Auth] = resolved_auth
self._last_initialize_instructions: Optional[str] = None
self._sampling_callback: Optional[Callable] = sampling_callback
self._elicitation_callback: Optional[Callable] = elicitation_callback
self._logging_callback: Optional[Callable] = logging_callback
self._resolved_auth: httpx.Auth | None = resolved_auth
self._last_initialize_instructions: str | None = None
self._sampling_callback: Callable | None = sampling_callback
self._elicitation_callback: Callable | None = elicitation_callback
self._logging_callback: Callable | None = logging_callback
# handle the basic auth value if provided
if auth_value:
self.update_auth_value(auth_value)
def _create_transport_context(
self,
) -> Tuple[Any, Optional[httpx.AsyncClient]]:
) -> tuple[Any, httpx.AsyncClient | None]:
"""
Create the appropriate transport context based on transport type.
Returns:
Tuple of (transport_context, http_client).
http_client is only set for HTTP transport and needs cleanup.
"""
http_client: Optional[httpx.AsyncClient] = None
http_client: httpx.AsyncClient | None = None
if self.transport_type == MCPTransport.stdio:
if not self.stdio_config:
raise ValueError("stdio_config is required for stdio transport")
@ -285,7 +280,7 @@ class MCPClient:
)
return transport_ctx, http_client
def _get_safe_stdio_env(self, provided_env: Optional[Dict[str, str]]) -> Optional[Dict[str, str]]:
def _get_safe_stdio_env(self, provided_env: dict[str, str] | None) -> dict[str, str] | None:
"""
Return a safe environment for the stdio subprocess.
@ -344,11 +339,11 @@ class MCPClient:
user input (elicitation), or send log messages.
"""
transport = await transport_ctx.__aenter__()
in_flight_error: Optional[BaseException] = None
in_flight_error: BaseException | None = None
try:
read_stream, write_stream = transport[0], transport[1]
# Build session kwargs with optional callbacks
session_kwargs: Dict[str, Any] = {}
session_kwargs: dict[str, Any] = {}
if self._sampling_callback is not None:
session_kwargs["sampling_callback"] = self._sampling_callback
if self._elicitation_callback is not None:
@ -393,7 +388,7 @@ class MCPClient:
quiet_on_error demotes the failure line to debug for callers that own the exception
(call_tool / list_tools under raise_on_error), so an expected pass-through re-auth does
not emit a warning per call; every other caller keeps the operator-visible warning."""
http_client: Optional[httpx.AsyncClient] = None
http_client: httpx.AsyncClient | None = None
try:
self._last_initialize_instructions = None
transport_ctx, http_client = self._create_transport_context()
@ -409,7 +404,7 @@ class MCPClient:
except BaseException as e:
verbose_logger.debug(f"Error during http_client cleanup: {e}")
def update_auth_value(self, mcp_auth_value: Union[str, Dict[str, str]]):
def update_auth_value(self, mcp_auth_value: str | dict[str, str]):
"""
Set the authentication header for the MCP client.
"""
@ -462,9 +457,9 @@ class MCPClient:
def factory(
*,
headers: Optional[Dict[str, str]] = None,
timeout: Optional[httpx.Timeout] = None,
auth: Optional[httpx.Auth] = None,
headers: dict[str, str] | None = None,
timeout: httpx.Timeout | None = None,
auth: httpx.Auth | None = None,
) -> httpx.AsyncClient:
"""Create an httpx.AsyncClient with LiteLLM's SSL configuration."""
# Get unified SSL configuration using the same logic as http_handler.py
@ -485,7 +480,7 @@ class MCPClient:
return factory
async def list_tools(self, raise_on_error: bool = False) -> List[MCPTool]:
async def list_tools(self, raise_on_error: bool = False) -> list[MCPTool]:
"""List available tools from the server.
Args:
@ -520,7 +515,7 @@ class MCPClient:
_log(
f"MCP client list_tools failed - "
f"Error Type: {error_type}, "
f"Error: {str(e)}, "
f"Error: {e!s}, "
f"Server: {self.server_url or 'stdio'}, "
f"Transport: {self.transport_type}"
)
@ -541,14 +536,14 @@ class MCPClient:
def error_tool_result(exc: Exception) -> MCPCallToolResult:
"""The error result ``call_tool`` returns when it swallows a failure (no re-execution)."""
return MCPCallToolResult(
content=[TextContent(type="text", text=f"{type(exc).__name__}: {str(exc)}")],
content=[TextContent(type="text", text=f"{type(exc).__name__}: {exc!s}")],
isError=True,
)
async def call_tool(
self,
call_tool_request_params: MCPCallToolRequestParams,
host_progress_callback: Optional[Callable] = None,
host_progress_callback: Callable | None = None,
raise_on_error: bool = False,
) -> MCPCallToolResult:
"""
@ -606,7 +601,7 @@ class MCPClient:
_log(
f"MCP client call_tool failed - "
f"Error Type: {error_type}, "
f"Error: {str(e)}, "
f"Error: {e!s}, "
f"Tool: {call_tool_request_params.name}, "
f"Server: {self.server_url or 'stdio'}, "
f"Transport: {self.transport_type}"
@ -622,7 +617,7 @@ class MCPClient:
# Return a default error result instead of raising
return self.error_tool_result(e)
async def list_prompts(self) -> List[Prompt]:
async def list_prompts(self) -> list[Prompt]:
"""List available prompts from the server."""
verbose_logger.debug(f"MCP client listing tools from {self.server_url or 'stdio'}")
@ -645,7 +640,7 @@ class MCPClient:
verbose_logger.error(
f"MCP client list_prompts failed - "
f"Error Type: {error_type}, "
f"Error: {str(e)}, "
f"Error: {e!s}, "
f"Server: {self.server_url or 'stdio'}, "
f"Transport: {self.transport_type}"
)
@ -686,7 +681,7 @@ class MCPClient:
verbose_logger.error(
f"MCP client get_prompt failed - "
f"Error Type: {error_type}, "
f"Error: {str(e)}, "
f"Error: {e!s}, "
f"Prompt: {get_prompt_request_params.name}, "
f"Server: {self.server_url or 'stdio'}, "
f"Transport: {self.transport_type}"
@ -722,7 +717,7 @@ class MCPClient:
verbose_logger.error(
f"MCP client list_resources failed - "
f"Error Type: {error_type}, "
f"Error: {str(e)}, "
f"Error: {e!s}, "
f"Server: {self.server_url or 'stdio'}, "
f"Transport: {self.transport_type}"
)
@ -758,7 +753,7 @@ class MCPClient:
verbose_logger.error(
f"MCP client list_resource_templates failed - "
f"Error Type: {error_type}, "
f"Error: {str(e)}, "
f"Error: {e!s}, "
f"Server: {self.server_url or 'stdio'}, "
f"Transport: {self.transport_type}"
)
@ -796,7 +791,7 @@ class MCPClient:
verbose_logger.error(
f"MCP client read_resource failed - "
f"Error Type: {error_type}, "
f"Error: {str(e)}, "
f"Error: {e!s}, "
f"Url: {url}, "
f"Server: {self.server_url or 'stdio'}, "
f"Transport: {self.transport_type}"

6
litellm/experimental_mcp_client/tools.py
View file

@ -1,5 +1,5 @@
import json
from typing import Dict, List, Literal, Union
from typing import Literal
from mcp import ClientSession
from mcp.types import CallToolRequestParams as MCPCallToolRequestParams
@ -92,7 +92,7 @@ def transform_mcp_tool_to_anthropic_tool(mcp_tool: MCPTool) -> AnthropicMessages
async def load_mcp_tools(
session: ClientSession, format: Literal["mcp", "openai"] = "mcp"
) -> Union[List[MCPTool], List[ChatCompletionToolParam]]:
) -> list[MCPTool] | list[ChatCompletionToolParam]:
"""
Load all available MCP tools
@ -138,7 +138,7 @@ def _get_function_arguments(function: FunctionDefinition) -> dict:
def transform_openai_tool_call_request_to_mcp_tool_call_request(
openai_tool: Union[ChatCompletionMessageToolCall, Dict],
openai_tool: ChatCompletionMessageToolCall | dict,
) -> MCPCallToolRequestParams:
"""Convert an OpenAI ChatCompletionMessageToolCall to an MCP CallToolRequestParams."""
function = openai_tool["function"]

129
litellm/files/main.py
View file

@ -12,7 +12,7 @@ import uuid as uuid_module
from collections.abc import Coroutine
from functools import partial
from types import MappingProxyType
from typing import Any, Dict, Literal, Optional, Union, cast
from typing import Any, Literal, cast
import httpx
@ -70,7 +70,7 @@ base_llm_http_handler = BaseLLMHTTPHandler()
def _should_sdk_support_streaming(
custom_llm_provider: Optional[Union[FileContentProvider, str]],
custom_llm_provider: FileContentProvider | str | None,
) -> bool:
"""
Return whether file content streaming is supported for the provider.
@ -86,7 +86,7 @@ bedrock_files_instance = BedrockFilesHandler()
def _add_trusted_model_credentials_to_litellm_params(
litellm_params_dict: Dict[str, Any], kwargs: Dict[str, Any]
litellm_params_dict: dict[str, Any], kwargs: dict[str, Any]
) -> None:
trusted_model_credentials = kwargs.get("_litellm_internal_model_credentials")
if isinstance(trusted_model_credentials, type(MappingProxyType({}))):
@ -97,10 +97,10 @@ def _add_trusted_model_credentials_to_litellm_params(
async def acreate_file(
file: FileTypes,
purpose: Literal["assistants", "batch", "fine-tune", "messages"],
expires_after: Optional[FileExpiresAfter] = None,
expires_after: FileExpiresAfter | None = None,
custom_llm_provider: FileCreateProvider = "openai",
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
**kwargs,
) -> OpenAIFileObject:
"""
@ -142,12 +142,12 @@ async def acreate_file(
def create_file(
file: FileTypes,
purpose: Literal["assistants", "batch", "fine-tune", "messages"],
expires_after: Optional[FileExpiresAfter] = None,
custom_llm_provider: Optional[FileCreateProvider] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
expires_after: FileExpiresAfter | None = None,
custom_llm_provider: FileCreateProvider | None = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
**kwargs,
) -> Union[OpenAIFileObject, Coroutine[Any, Any, OpenAIFileObject]]:
) -> OpenAIFileObject | Coroutine[Any, Any, OpenAIFileObject]:
"""
Files are used to upload documents that can be used with features like Assistants, Fine-tuning, and Batch API.
@ -159,7 +159,7 @@ def create_file(
_is_async = kwargs.pop("acreate_file", False) is True
optional_params = GenericLiteLLMParams(**kwargs)
litellm_params_dict = dict(**kwargs)
logging_obj = cast(Optional[LiteLLMLoggingObj], kwargs.get("litellm_logging_obj"))
logging_obj = cast(LiteLLMLoggingObj | None, kwargs.get("litellm_logging_obj"))
if logging_obj is None:
raise ValueError("logging_obj is required")
client = kwargs.get("client")
@ -246,9 +246,7 @@ def create_file(
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'create_file'. Only ['openai', 'azure', 'vertex_ai', 'manus', 'anthropic'] are supported.".format(
custom_llm_provider
),
message=f"LiteLLM doesn't support {custom_llm_provider} for 'create_file'. Only ['openai', 'azure', 'vertex_ai', 'manus', 'anthropic'] are supported.",
model="n/a",
llm_provider=custom_llm_provider,
response=httpx.Response(
@ -266,8 +264,8 @@ def create_file(
async def afile_retrieve(
file_id: str,
custom_llm_provider: FileRetrieveProvider = "openai",
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
**kwargs,
) -> OpenAIFileObject:
"""
@ -307,8 +305,8 @@ async def afile_retrieve(
def file_retrieve(
file_id: str,
custom_llm_provider: FileRetrieveProvider = "openai",
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
**kwargs,
) -> FileObject:
"""
@ -412,9 +410,7 @@ def file_retrieve(
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'file_retrieve'. Only 'openai', 'azure', 'manus', and 'anthropic' are supported.".format(
custom_llm_provider
),
message=f"LiteLLM doesn't support {custom_llm_provider} for 'file_retrieve'. Only 'openai', 'azure', 'manus', and 'anthropic' are supported.",
model="n/a",
llm_provider=custom_llm_provider,
response=httpx.Response(
@ -437,8 +433,8 @@ def file_retrieve(
async def afile_delete(
file_id: str,
custom_llm_provider: FileDeleteProvider = "openai",
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
**kwargs,
) -> Coroutine[Any, Any, FileObject]:
"""
@ -479,10 +475,10 @@ async def afile_delete(
@client
def file_delete(
file_id: str,
model: Optional[str] = None,
custom_llm_provider: Union[FileDeleteProvider, str] = "openai",
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
model: str | None = None,
custom_llm_provider: FileDeleteProvider | str = "openai",
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
**kwargs,
) -> FileDeleted:
"""
@ -591,9 +587,7 @@ def file_delete(
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'file_delete'. Only 'openai', 'azure', 'gemini', 'manus', and 'anthropic' are supported.".format(
custom_llm_provider
),
message=f"LiteLLM doesn't support {custom_llm_provider} for 'file_delete'. Only 'openai', 'azure', 'gemini', 'manus', and 'anthropic' are supported.",
model="n/a",
llm_provider=custom_llm_provider,
response=httpx.Response(
@ -614,9 +608,9 @@ def file_delete(
@client
async def afile_list(
custom_llm_provider: FileListProvider = "openai",
purpose: Optional[str] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
purpose: str | None = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
**kwargs,
):
"""
@ -655,9 +649,9 @@ async def afile_list(
@client
def file_list(
custom_llm_provider: FileListProvider = "openai",
purpose: Optional[str] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
purpose: str | None = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
**kwargs,
):
"""
@ -755,9 +749,7 @@ def file_list(
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'file_list'. Only 'openai', 'azure', 'manus', and 'anthropic' are supported.".format(
custom_llm_provider
),
message=f"LiteLLM doesn't support {custom_llm_provider} for 'file_list'. Only 'openai', 'azure', 'manus', and 'anthropic' are supported.",
model="n/a",
llm_provider=custom_llm_provider,
response=httpx.Response(
@ -775,12 +767,12 @@ def file_list(
async def afile_content(
file_id: str,
custom_llm_provider: FileContentProvider = "openai",
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
chunk_size: int = 1024 * 1024,
stream: bool = False,
**kwargs,
) -> Union[HttpxBinaryResponseContent, FileContentStreamingResult]:
) -> HttpxBinaryResponseContent | FileContentStreamingResult:
"""
Async: Get file contents
@ -821,19 +813,19 @@ async def afile_content(
@client
def file_content(
file_id: str,
model: Optional[str] = None,
custom_llm_provider: Optional[Union[FileContentProvider, str]] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
model: str | None = None,
custom_llm_provider: FileContentProvider | str | None = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
chunk_size: int = 1024 * 1024,
stream: bool = False,
**kwargs,
) -> Union[
HttpxBinaryResponseContent,
FileContentStreamingResult,
Coroutine[Any, Any, HttpxBinaryResponseContent],
Coroutine[Any, Any, FileContentStreamingResult],
]:
) -> (
HttpxBinaryResponseContent
| FileContentStreamingResult
| Coroutine[Any, Any, HttpxBinaryResponseContent]
| Coroutine[Any, Any, FileContentStreamingResult]
):
"""
Returns the contents of the specified file.
@ -887,7 +879,7 @@ def file_content(
chunk_size=chunk_size,
optional_params=optional_params,
timeout=timeout,
logging_obj=cast(Optional[LiteLLMLoggingObj], kwargs.get("litellm_logging_obj")),
logging_obj=cast(LiteLLMLoggingObj | None, kwargs.get("litellm_logging_obj")),
_is_async=_is_async,
client=client,
)
@ -989,9 +981,7 @@ def file_content(
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'file_content'. Supported providers are 'openai', 'azure', 'vertex_ai', 'bedrock', 'manus', 'anthropic'.".format(
custom_llm_provider
),
message=f"LiteLLM doesn't support {custom_llm_provider} for 'file_content'. Supported providers are 'openai', 'azure', 'vertex_ai', 'bedrock', 'manus', 'anthropic'.",
model="n/a",
llm_provider=custom_llm_provider,
response=httpx.Response(
@ -1008,17 +998,17 @@ def file_content(
def file_content_streaming(
*,
file_id: str,
model: Optional[str],
custom_llm_provider: Optional[Union[FileContentProvider, str]],
extra_headers: Optional[Dict[str, str]],
extra_body: Optional[Dict[str, str]],
model: str | None,
custom_llm_provider: FileContentProvider | str | None,
extra_headers: dict[str, str] | None,
extra_body: dict[str, str] | None,
chunk_size: int,
optional_params: GenericLiteLLMParams,
timeout: Union[float, httpx.Timeout],
logging_obj: Optional[LiteLLMLoggingObj],
timeout: float | httpx.Timeout,
logging_obj: LiteLLMLoggingObj | None,
_is_async: bool,
client: Optional[Any],
) -> Union[FileContentStreamingResult, Coroutine[Any, Any, FileContentStreamingResult]]:
client: Any | None,
) -> FileContentStreamingResult | Coroutine[Any, Any, FileContentStreamingResult]:
if logging_obj is not None:
logging_obj.model = model or ""
logging_obj.model_call_details["model"] = model or ""
@ -1043,8 +1033,8 @@ def file_content_streaming(
headers=response.headers,
)
response: Union[FileContentStreamingResult, Coroutine[Any, Any, FileContentStreamingResult]] = (
FileContentStreamingResult(stream_iterator=iter(()), headers={})
response: FileContentStreamingResult | Coroutine[Any, Any, FileContentStreamingResult] = FileContentStreamingResult(
stream_iterator=iter(()), headers={}
)
if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS:
openai_creds = get_openai_credentials(
@ -1069,10 +1059,7 @@ def file_content_streaming(
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for streaming 'file_content'. Supported providers are {}.".format(
custom_llm_provider,
sorted(OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS),
),
message=f"LiteLLM doesn't support {custom_llm_provider} for streaming 'file_content'. Supported providers are {sorted(OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS)}.",
model="n/a",
llm_provider=custom_llm_provider,
response=httpx.Response(

20
litellm/files/streaming.py
View file

@ -4,9 +4,7 @@ from collections.abc import AsyncIterator, Iterator
from typing import (
TYPE_CHECKING,
Any,
Dict,
Optional,
Union,
cast,
)
@ -29,10 +27,10 @@ class FileContentStreamingResponse:
def __init__(
self,
stream_iterator: Union[Iterator[bytes], AsyncIterator[bytes]],
stream_iterator: Iterator[bytes] | AsyncIterator[bytes],
file_id: str,
model: Optional[str],
custom_llm_provider: Optional[Union[FileContentProvider, str]],
model: str | None,
custom_llm_provider: FileContentProvider | str | None,
logging_obj: Optional["LiteLLMLoggingObj"],
) -> None:
self.stream_iterator = stream_iterator
@ -40,8 +38,8 @@ class FileContentStreamingResponse:
self.model = model
self.custom_llm_provider = custom_llm_provider
self.logging_obj = logging_obj
self.standard_logging_object: Optional["StandardLoggingPayload"] = None
self._hidden_params: Dict[str, Any] = {}
self.standard_logging_object: StandardLoggingPayload | None = None
self._hidden_params: dict[str, Any] = {}
self._logging_completed = False
self._close_completed = False
self._start_time = (
@ -94,7 +92,7 @@ class FileContentStreamingResponse:
self._close_completed = True
self._logging_completed = True
stream_to_close = self.stream_iterator
self.stream_iterator = cast(Union[Iterator[bytes], AsyncIterator[bytes]], iter(()))
self.stream_iterator = cast(Iterator[bytes] | AsyncIterator[bytes], iter(()))
# Shield cleanup from request cancellation so upstream HTTP connections
# are released promptly on client disconnects.
@ -113,12 +111,12 @@ class FileContentStreamingResponse:
self._close_completed = True
self._logging_completed = True
stream_to_close = self.stream_iterator
self.stream_iterator = cast(Union[Iterator[bytes], AsyncIterator[bytes]], iter(()))
self.stream_iterator = cast(Iterator[bytes] | AsyncIterator[bytes], iter(()))
if hasattr(stream_to_close, "close"):
cast(Iterator[bytes], stream_to_close).close() # type: ignore[attr-defined]
def _build_logging_response(self) -> Dict[str, str]:
def _build_logging_response(self) -> dict[str, str]:
response = {
"id": self.file_id,
"object": "file.content",
@ -170,7 +168,7 @@ class FileContentStreamingResponse:
merged_hidden_params = cast(
"StandardLoggingHiddenParams",
{
**cast(Dict[str, Any], payload.get("hidden_params") or {}),
**cast(dict[str, Any], payload.get("hidden_params") or {}),
**self._hidden_params,
},
)

6
litellm/files/types.py
View file

@ -1,9 +1,9 @@
from collections.abc import AsyncIterator, Iterator
from typing import Dict, Literal, NamedTuple, Union
from typing import Literal, NamedTuple
FileContentProvider = Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic", "manus"]
class FileContentStreamingResult(NamedTuple):
stream_iterator: Union[Iterator[bytes], AsyncIterator[bytes]]
headers: Dict[str, str]
stream_iterator: Iterator[bytes] | AsyncIterator[bytes]
headers: dict[str, str]

6
litellm/files/utils.py
View file

@ -1,5 +1,3 @@
from typing import Optional
from litellm.types.llms.openai import CreateFileRequest
from litellm.types.utils import ExtractedFileData
@ -37,7 +35,7 @@ class FilesAPIUtils:
)
@staticmethod
def is_batch_jsonl_request(create_file_data: CreateFileRequest, content_type: Optional[str]) -> bool:
def is_batch_jsonl_request(create_file_data: CreateFileRequest, content_type: str | None) -> bool:
"""
Batch-jsonl check from metadata only, so the body can stay a streamable
Path/handle instead of being read into memory.
@ -49,7 +47,7 @@ class FilesAPIUtils:
)
@staticmethod
def valid_content_type(content_type: Optional[str]) -> bool:
def valid_content_type(content_type: str | None) -> bool:
"""
Whether the upload's MIME type is one a batch JSONL file is plausibly
sent as (see ``_BATCH_JSONL_CONTENT_TYPES``).

94
litellm/fine_tuning/main.py
View file

@ -13,7 +13,7 @@ import contextvars
import os
from collections.abc import Coroutine
from functools import partial
from typing import Any, Dict, Literal, Optional, Union
from typing import Any, Literal
import httpx
@ -36,10 +36,10 @@ vertex_fine_tuning_apis_instance = VertexFineTuningAPI()
def _prepare_azure_extra_body(
extra_body: Optional[Dict[str, Any]],
kwargs: Dict[str, Any],
azure_specific_hyperparams: Dict[str, Any],
) -> Dict[str, Any]:
extra_body: dict[str, Any] | None,
kwargs: dict[str, Any],
azure_specific_hyperparams: dict[str, Any],
) -> dict[str, Any]:
"""
Prepare extra_body for Azure fine-tuning API by combining Azure-specific parameters.
@ -77,14 +77,14 @@ def _prepare_azure_extra_body(
async def acreate_fine_tuning_job(
model: str,
training_file: str,
hyperparameters: Optional[dict] = {},
suffix: Optional[str] = None,
validation_file: Optional[str] = None,
integrations: Optional[List[str]] = None,
seed: Optional[int] = None,
hyperparameters: dict | None = {},
suffix: str | None = None,
validation_file: str | None = None,
integrations: List[str] | None = None,
seed: int | None = None,
custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai",
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
**kwargs,
) -> LiteLLMFineTuningJob:
"""
@ -140,7 +140,7 @@ def _build_fine_tuning_job_data(model, training_file, hyperparameters, suffix, v
def _resolve_fine_tuning_timeout(
timeout: Any,
custom_llm_provider: str,
) -> Union[float, httpx.Timeout]:
) -> float | httpx.Timeout:
"""Normalise a raw timeout value to a float (seconds) or httpx.Timeout for fine-tuning calls."""
timeout = timeout or 600.0
if isinstance(timeout, httpx.Timeout):
@ -154,16 +154,16 @@ def _resolve_fine_tuning_timeout(
def create_fine_tuning_job(
model: str,
training_file: str,
hyperparameters: Optional[dict] = {},
suffix: Optional[str] = None,
validation_file: Optional[str] = None,
integrations: Optional[List[str]] = None,
seed: Optional[int] = None,
hyperparameters: dict | None = {},
suffix: str | None = None,
validation_file: str | None = None,
integrations: List[str] | None = None,
seed: int | None = None,
custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai",
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
**kwargs,
) -> Union[LiteLLMFineTuningJob, Coroutine[Any, Any, LiteLLMFineTuningJob]]:
) -> LiteLLMFineTuningJob | Coroutine[Any, Any, LiteLLMFineTuningJob]:
"""
Creates a fine-tuning job which begins the process of creating a new model from a given dataset.
@ -315,9 +315,7 @@ def create_fine_tuning_job(
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'create_batch'. Only 'openai' is supported.".format(
custom_llm_provider
),
message=f"LiteLLM doesn't support {custom_llm_provider} for 'create_batch'. Only 'openai' is supported.",
model="n/a",
llm_provider=custom_llm_provider,
response=httpx.Response(
@ -336,8 +334,8 @@ def create_fine_tuning_job(
async def acancel_fine_tuning_job(
fine_tuning_job_id: str,
custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai",
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
**kwargs,
) -> LiteLLMFineTuningJob:
"""
@ -374,10 +372,10 @@ async def acancel_fine_tuning_job(
def cancel_fine_tuning_job(
fine_tuning_job_id: str,
custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai",
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
**kwargs,
) -> Union[LiteLLMFineTuningJob, Coroutine[Any, Any, LiteLLMFineTuningJob]]:
) -> LiteLLMFineTuningJob | Coroutine[Any, Any, LiteLLMFineTuningJob]:
"""
Immediately cancel a fine-tune job.
@ -469,9 +467,7 @@ def cancel_fine_tuning_job(
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'create_batch'. Only 'openai' is supported.".format(
custom_llm_provider
),
message=f"LiteLLM doesn't support {custom_llm_provider} for 'create_batch'. Only 'openai' is supported.",
model="n/a",
llm_provider=custom_llm_provider,
response=httpx.Response(
@ -486,11 +482,11 @@ def cancel_fine_tuning_job(
async def alist_fine_tuning_jobs(
after: Optional[str] = None,
limit: Optional[int] = None,
after: str | None = None,
limit: int | None = None,
custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai",
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
**kwargs,
):
"""
@ -525,11 +521,11 @@ async def alist_fine_tuning_jobs(
def list_fine_tuning_jobs(
after: Optional[str] = None,
limit: Optional[int] = None,
after: str | None = None,
limit: int | None = None,
custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai",
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
**kwargs,
):
"""
@ -627,9 +623,7 @@ def list_fine_tuning_jobs(
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'create_batch'. Only 'openai' is supported.".format(
custom_llm_provider
),
message=f"LiteLLM doesn't support {custom_llm_provider} for 'create_batch'. Only 'openai' is supported.",
model="n/a",
llm_provider=custom_llm_provider,
response=httpx.Response(
@ -647,8 +641,8 @@ def list_fine_tuning_jobs(
async def aretrieve_fine_tuning_job(
fine_tuning_job_id: str,
custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai",
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
**kwargs,
) -> LiteLLMFineTuningJob:
"""
@ -685,10 +679,10 @@ async def aretrieve_fine_tuning_job(
def retrieve_fine_tuning_job(
fine_tuning_job_id: str,
custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai",
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
**kwargs,
) -> Union[LiteLLMFineTuningJob, Coroutine[Any, Any, LiteLLMFineTuningJob]]:
) -> LiteLLMFineTuningJob | Coroutine[Any, Any, LiteLLMFineTuningJob]:
"""
Get info about a fine-tuning job.
"""
@ -767,9 +761,7 @@ def retrieve_fine_tuning_job(
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'retrieve_fine_tuning_job'. Only 'openai' and 'azure' are supported.".format(
custom_llm_provider
),
message=f"LiteLLM doesn't support {custom_llm_provider} for 'retrieve_fine_tuning_job'. Only 'openai' and 'azure' are supported.",
model="n/a",
llm_provider=custom_llm_provider,
response=httpx.Response(

4
litellm/google_genai/__init__.py
View file

@ -12,8 +12,8 @@ from .main import (
)
__all__ = [
"generate_content",
"agenerate_content",
"generate_content_stream",
"agenerate_content_stream",
"generate_content",
"generate_content_stream",
]

2
litellm/google_genai/adapters/__init__.py
View file

@ -13,7 +13,7 @@ from .handler import GenerateContentToCompletionHandler
from .transformation import GoogleGenAIAdapter, GoogleGenAIStreamWrapper
__all__ = [
"GenerateContentToCompletionHandler",
"GoogleGenAIAdapter",
"GoogleGenAIStreamWrapper",
"GenerateContentToCompletionHandler",
]

34
litellm/google_genai/adapters/handler.py
View file

@ -1,5 +1,5 @@
from collections.abc import AsyncIterator, Coroutine
from typing import Any, Dict, List, Optional, Union, cast
from typing import Any, cast
import litellm
from litellm.types.router import GenericLiteLLMParams
@ -17,12 +17,12 @@ class GenerateContentToCompletionHandler:
@staticmethod
def _prepare_completion_kwargs(
model: str,
contents: Union[List[Dict[str, Any]], Dict[str, Any]],
config: Optional[Dict[str, Any]] = None,
contents: list[dict[str, Any]] | dict[str, Any],
config: dict[str, Any] | None = None,
stream: bool = False,
litellm_params: Optional[GenericLiteLLMParams] = None,
extra_kwargs: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
litellm_params: GenericLiteLLMParams | None = None,
extra_kwargs: dict[str, Any] | None = None,
) -> dict[str, Any]:
"""Prepare kwargs for litellm.completion/acompletion"""
# Transform generate_content request to completion format
@ -34,7 +34,7 @@ class GenerateContentToCompletionHandler:
**(extra_kwargs or {}),
)
completion_kwargs: Dict[str, Any] = dict(completion_request)
completion_kwargs: dict[str, Any] = dict(completion_request)
# Forward extra_kwargs that should be passed to completion call
if extra_kwargs is not None:
@ -53,12 +53,12 @@ class GenerateContentToCompletionHandler:
@staticmethod
async def async_generate_content_handler(
model: str,
contents: Union[List[Dict[str, Any]], Dict[str, Any]],
contents: list[dict[str, Any]] | dict[str, Any],
litellm_params: GenericLiteLLMParams,
config: Optional[Dict[str, Any]] = None,
config: dict[str, Any] | None = None,
stream: bool = False,
**kwargs,
) -> Union[Dict[str, Any], AsyncIterator[bytes]]:
) -> dict[str, Any] | AsyncIterator[bytes]:
"""Handle generate_content call asynchronously using completion adapter"""
completion_kwargs = GenerateContentToCompletionHandler._prepare_completion_kwargs(
@ -98,22 +98,18 @@ class GenerateContentToCompletionHandler:
return generate_content_response
except Exception as e:
raise ValueError(f"Error calling litellm.acompletion for generate_content: {str(e)}")
raise ValueError(f"Error calling litellm.acompletion for generate_content: {e!s}")
@staticmethod
def generate_content_handler(
model: str,
contents: Union[List[Dict[str, Any]], Dict[str, Any]],
contents: list[dict[str, Any]] | dict[str, Any],
litellm_params: GenericLiteLLMParams,
config: Optional[Dict[str, Any]] = None,
config: dict[str, Any] | None = None,
stream: bool = False,
_is_async: bool = False,
**kwargs,
) -> Union[
Dict[str, Any],
AsyncIterator[bytes],
Coroutine[Any, Any, Union[Dict[str, Any], AsyncIterator[bytes]]],
]:
) -> dict[str, Any] | AsyncIterator[bytes] | Coroutine[Any, Any, dict[str, Any] | AsyncIterator[bytes]]:
"""Handle generate_content call using completion adapter"""
if _is_async:
@ -163,4 +159,4 @@ class GenerateContentToCompletionHandler:
return generate_content_response
except Exception as e:
raise ValueError(f"Error calling litellm.completion for generate_content: {str(e)}")
raise ValueError(f"Error calling litellm.completion for generate_content: {e!s}")

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

@ -1,6 +1,6 @@
import json
from collections.abc import AsyncIterator, Iterator
from typing import Any, Dict, List, Optional, Union, cast
from typing import Any, cast
from litellm import verbose_logger
from litellm.litellm_core_utils.json_validation_rule import normalize_tool_schema
@ -36,7 +36,7 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
sent_first_chunk: bool = False
# State tracking for accumulating partial tool calls
accumulated_tool_calls: Dict[str, Dict[str, Any]]
accumulated_tool_calls: dict[str, dict[str, Any]]
def __init__(self, completion_stream: Any):
self.sent_first_chunk = False
@ -108,7 +108,6 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
f"Name: {tool_call_data['name']}. "
f"Partial args: {tool_call_data['arguments']}"
)
pass
if parts:
final_chunk = {
"candidates": [
@ -178,11 +177,11 @@ class GoogleGenAIAdapter:
def translate_generate_content_to_completion(
self,
model: str,
contents: Union[List[Dict[str, Any]], Dict[str, Any]],
config: Optional[Dict[str, Any]] = None,
litellm_params: Optional[GenericLiteLLMParams] = None,
contents: list[dict[str, Any]] | dict[str, Any],
config: dict[str, Any] | None = None,
litellm_params: GenericLiteLLMParams | None = None,
**kwargs,
) -> Dict[str, Any]:
) -> dict[str, Any]:
"""
Transform generate_content request to litellm completion format
@ -273,8 +272,8 @@ class GoogleGenAIAdapter:
def _add_generic_litellm_params_to_request(
self,
completion_request_dict: Dict[str, Any],
litellm_params: Optional[GenericLiteLLMParams] = None,
completion_request_dict: dict[str, Any],
litellm_params: GenericLiteLLMParams | None = None,
) -> dict:
"""Add generic litellm params to request. e.g add api_base, api_key, api_version, etc.
@ -296,7 +295,7 @@ class GoogleGenAIAdapter:
def translate_completion_output_params_streaming(
self,
completion_stream: Any,
) -> Union[AsyncIterator[bytes], None]:
) -> AsyncIterator[bytes] | None:
"""Transform streaming completion output to Google GenAI format"""
google_genai_wrapper = GoogleGenAIStreamWrapper(completion_stream=completion_stream)
# Return the SSE-wrapped version for proper event formatting
@ -304,15 +303,15 @@ class GoogleGenAIAdapter:
def _transform_google_genai_tools_to_openai(
self,
tools: List[Dict[str, Any]],
) -> List[ChatCompletionToolParam]:
tools: list[dict[str, Any]],
) -> list[ChatCompletionToolParam]:
"""Transform Google GenAI tools to OpenAI tools format"""
openai_tools: List[Dict[str, Any]] = []
openai_tools: list[dict[str, Any]] = []
for tool in tools:
if "functionDeclarations" in tool:
for func_decl in tool["functionDeclarations"]:
function_chunk: Dict[str, Any] = {
function_chunk: dict[str, Any] = {
"name": func_decl.get("name", ""),
}
@ -327,12 +326,12 @@ class GoogleGenAIAdapter:
# normalize the tool schemas
normalized_tools = [normalize_tool_schema(tool) for tool in openai_tools]
return cast(List[ChatCompletionToolParam], normalized_tools)
return cast(list[ChatCompletionToolParam], normalized_tools)
def _transform_google_genai_tool_config_to_openai(
self,
tool_config: Dict[str, Any],
) -> Optional[ChatCompletionToolChoiceValues]:
tool_config: dict[str, Any],
) -> ChatCompletionToolChoiceValues | None:
"""Transform Google GenAI tool_config to OpenAI tool_choice"""
function_calling_config = tool_config.get("functionCallingConfig", {})
mode = function_calling_config.get("mode", "AUTO")
@ -344,11 +343,11 @@ class GoogleGenAIAdapter:
def _transform_contents_to_messages(
self,
contents: List[Dict[str, Any]],
system_instruction: Optional[Dict[str, Any]] = None,
) -> List[AllMessageValues]:
contents: list[dict[str, Any]],
system_instruction: dict[str, Any] | None = None,
) -> list[AllMessageValues]:
"""Transform Google GenAI contents to OpenAI messages format"""
messages: List[AllMessageValues] = []
messages: list[AllMessageValues] = []
# Handle system instruction
if system_instruction:
@ -362,8 +361,8 @@ class GoogleGenAIAdapter:
if role == "user":
# Handle user messages with potential function responses
content_parts: List[Union[ChatCompletionTextObject, ChatCompletionImageObject]] = []
tool_messages: List[ChatCompletionToolMessage] = []
content_parts: list[ChatCompletionTextObject | ChatCompletionImageObject] = []
tool_messages: list[ChatCompletionToolMessage] = []
for part in parts:
if isinstance(part, dict):
@ -420,7 +419,7 @@ class GoogleGenAIAdapter:
elif role == "model":
# Handle assistant messages with potential function calls
combined_text = ""
tool_calls: List[ChatCompletionAssistantToolCall] = []
tool_calls: list[ChatCompletionAssistantToolCall] = []
for part in parts:
if isinstance(part, dict):
@ -461,7 +460,7 @@ class GoogleGenAIAdapter:
def translate_completion_to_generate_content(
self,
response: ModelResponse,
) -> Dict[str, Any]:
) -> dict[str, Any]:
"""
Transform litellm completion response to Google GenAI generate_content format
@ -490,7 +489,7 @@ class GoogleGenAIAdapter:
parts = [{"text": message_content}] if message_content else []
# Create Google GenAI format response
generate_content_response: Dict[str, Any] = {
generate_content_response: dict[str, Any] = {
"candidates": [
{
"content": {"parts": parts, "role": "model"},
@ -522,9 +521,9 @@ class GoogleGenAIAdapter:
def translate_streaming_completion_to_generate_content(
self,
response: Union[ModelResponse, ModelResponseStream],
response: ModelResponse | ModelResponseStream,
wrapper: GoogleGenAIStreamWrapper,
) -> Optional[Dict[str, Any]]:
) -> dict[str, Any] | None:
"""
Transform streaming litellm completion chunk to Google GenAI generate_content format
@ -560,7 +559,7 @@ class GoogleGenAIAdapter:
return None
# Create Google GenAI streaming format response
streaming_chunk: Dict[str, Any] = {
streaming_chunk: dict[str, Any] = {
"candidates": [
{
"content": {"parts": parts, "role": "model"},
@ -597,9 +596,9 @@ class GoogleGenAIAdapter:
def _transform_openai_message_to_google_genai_parts(
self,
message: Any,
) -> List[Dict[str, Any]]:
) -> list[dict[str, Any]]:
"""Transform OpenAI message to Google GenAI parts format"""
parts: List[Dict[str, Any]] = []
parts: list[dict[str, Any]] = []
# Add text content if present
if hasattr(message, "content") and message.content:
@ -626,14 +625,14 @@ class GoogleGenAIAdapter:
def _transform_openai_delta_to_google_genai_parts_with_accumulation(
self, delta: Any, wrapper: GoogleGenAIStreamWrapper
) -> List[Dict[str, Any]]:
) -> list[dict[str, Any]]:
"""Transforms OpenAI delta to Google GenAI parts, accumulating streaming tool calls."""
# 1. Initialize wrapper state if it doesn't exist
if not hasattr(wrapper, "accumulated_tool_calls"):
wrapper.accumulated_tool_calls = {}
parts: List[Dict[str, Any]] = []
parts: list[dict[str, Any]] = []
if hasattr(delta, "content") and delta.content:
parts.append({"text": delta.content})
@ -699,7 +698,7 @@ class GoogleGenAIAdapter:
return parts
def _map_finish_reason(self, finish_reason: Optional[str]) -> str:
def _map_finish_reason(self, finish_reason: str | None) -> str:
"""Map OpenAI finish reasons to Google GenAI finish reasons"""
if not finish_reason:
return "STOP"
@ -714,7 +713,7 @@ class GoogleGenAIAdapter:
return mapping.get(finish_reason, "STOP")
def _map_usage(self, usage: Any) -> Dict[str, int]:
def _map_usage(self, usage: Any) -> dict[str, int]:
"""Map OpenAI usage to Google GenAI usage format"""
return {
"promptTokenCount": getattr(usage, "prompt_tokens", 0) or 0,

80
litellm/google_genai/main.py
View file

@ -2,7 +2,7 @@ import asyncio
import contextvars
from collections.abc import Iterator
from functools import partial
from typing import TYPE_CHECKING, Any, ClassVar, Dict, Optional, Union
from typing import TYPE_CHECKING, Any, ClassVar
import httpx
from pydantic import BaseModel, ConfigDict
@ -52,14 +52,14 @@ class GenerateContentSetupResult(BaseModel):
model_config: ClassVar[ConfigDict] = ConfigDict(arbitrary_types_allowed=True)
model: str
request_body: Dict[str, Any]
request_body: dict[str, Any]
custom_llm_provider: str
generate_content_provider_config: Optional[BaseGoogleGenAIGenerateContentConfig]
generate_content_config_dict: Dict[str, Any]
generate_content_provider_config: BaseGoogleGenAIGenerateContentConfig | None
generate_content_config_dict: dict[str, Any]
native_request_fields: dict[str, object]
litellm_params: GenericLiteLLMParams
litellm_logging_obj: LiteLLMLoggingObj
litellm_call_id: Optional[str]
litellm_call_id: str | None
class GenerateContentHelper:
@ -68,7 +68,7 @@ class GenerateContentHelper:
@staticmethod
def mock_generate_content_response(
mock_response: str = "This is a mock response from Google GenAI generate_content.",
) -> Dict[str, Any]:
) -> dict[str, Any]:
"""Mock response for generate_content for testing purposes"""
return {
"text": mock_response,
@ -91,9 +91,9 @@ class GenerateContentHelper:
def setup_generate_content_call(
model: str,
contents: GenerateContentContentListUnionDict,
config: Optional[GenerateContentConfigDict] = None,
custom_llm_provider: Optional[str] = None,
tools: Optional[ToolConfigDict] = None,
config: GenerateContentConfigDict | None = None,
custom_llm_provider: str | None = None,
tools: ToolConfigDict | None = None,
**kwargs,
) -> GenerateContentSetupResult:
"""
@ -110,8 +110,8 @@ class GenerateContentHelper:
Returns:
GenerateContentSetupResult containing all setup information
"""
litellm_logging_obj: Optional[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj")
litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None)
litellm_logging_obj: LiteLLMLoggingObj | None = kwargs.get("litellm_logging_obj")
litellm_call_id: str | None = kwargs.get("litellm_call_id", None)
# get llm provider logic
litellm_params = GenericLiteLLMParams(**kwargs)
@ -140,7 +140,7 @@ class GenerateContentHelper:
litellm_params.custom_llm_provider = custom_llm_provider
# get provider config
generate_content_provider_config: Optional[BaseGoogleGenAIGenerateContentConfig] = (
generate_content_provider_config: BaseGoogleGenAIGenerateContentConfig | None = (
ProviderConfigManager.get_provider_google_genai_generate_content_config(
model=model,
provider=litellm.LlmProviders(custom_llm_provider),
@ -235,16 +235,16 @@ def _merge_native_request_fields(
async def agenerate_content(
model: str,
contents: GenerateContentContentListUnionDict,
config: Optional[GenerateContentConfigDict] = None,
tools: Optional[ToolConfigDict] = None,
config: GenerateContentConfigDict | None = None,
tools: ToolConfigDict | None = None,
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
# The extra values given here take precedence over values defined on the client or passed to this method.
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
timeout: float | httpx.Timeout | None = None,
# LiteLLM specific params,
custom_llm_provider: Optional[str] = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> Any:
"""
@ -303,16 +303,16 @@ async def agenerate_content(
def generate_content(
model: str,
contents: GenerateContentContentListUnionDict,
config: Optional[GenerateContentConfigDict] = None,
tools: Optional[ToolConfigDict] = None,
config: GenerateContentConfigDict | None = None,
tools: ToolConfigDict | None = None,
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
# The extra values given here take precedence over values defined on the client or passed to this method.
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
timeout: float | httpx.Timeout | None = None,
# LiteLLM specific params,
custom_llm_provider: Optional[str] = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> Any:
"""
@ -393,16 +393,16 @@ def generate_content(
async def agenerate_content_stream(
model: str,
contents: GenerateContentContentListUnionDict,
config: Optional[GenerateContentConfigDict] = None,
tools: Optional[ToolConfigDict] = None,
config: GenerateContentConfigDict | None = None,
tools: ToolConfigDict | None = None,
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
# The extra values given here take precedence over values defined on the client or passed to this method.
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
timeout: float | httpx.Timeout | None = None,
# LiteLLM specific params,
custom_llm_provider: Optional[str] = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> Any:
"""
@ -488,16 +488,16 @@ async def agenerate_content_stream(
def generate_content_stream(
model: str,
contents: GenerateContentContentListUnionDict,
config: Optional[GenerateContentConfigDict] = None,
tools: Optional[ToolConfigDict] = None,
config: GenerateContentConfigDict | None = None,
tools: ToolConfigDict | None = None,
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
# The extra values given here take precedence over values defined on the client or passed to this method.
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
timeout: float | httpx.Timeout | None = None,
# LiteLLM specific params,
custom_llm_provider: Optional[str] = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> Iterator[Any]:
"""

22
litellm/google_genai/streaming_iterator.py
View file

@ -1,6 +1,6 @@
import asyncio
from datetime import datetime
from typing import TYPE_CHECKING, Any, Dict, List, Optional
from typing import TYPE_CHECKING, Any
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy.pass_through_endpoints.success_handler import (
@ -18,12 +18,12 @@ else:
GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ = PassThroughEndpointLogging()
def _encode_google_genai_sse_event(event_lines: List[str]) -> bytes:
def _encode_google_genai_sse_event(event_lines: list[str]) -> bytes:
return ("\n".join(event_lines) + "\n\n").encode("utf-8")
def _next_google_genai_sse_chunk(line_iter) -> bytes:
event_lines: List[str] = []
event_lines: list[str] = []
while True:
try:
line = next(line_iter)
@ -39,7 +39,7 @@ def _next_google_genai_sse_chunk(line_iter) -> bytes:
async def _anext_google_genai_sse_chunk(line_iter) -> bytes:
event_lines: List[str] = []
event_lines: list[str] = []
while True:
try:
line = await line_iter.__anext__()
@ -65,14 +65,14 @@ class BaseGoogleGenAIGenerateContentStreamingIterator:
litellm_logging_obj: LiteLLMLoggingObj,
request_body: dict,
model: str,
hidden_params: Optional[Dict[str, Any]] = None,
hidden_params: dict[str, Any] | None = None,
):
self.litellm_logging_obj = litellm_logging_obj
self.request_body = request_body
self.start_time = datetime.now()
self.collected_chunks: List[bytes] = []
self.collected_chunks: list[bytes] = []
self.model = model
self._hidden_params: Dict[str, Any] = hidden_params or {}
self._hidden_params: dict[str, Any] = hidden_params or {}
async def _handle_async_streaming_logging(
self,
@ -111,8 +111,8 @@ class GoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateContent
generate_content_provider_config: BaseGoogleGenAIGenerateContentConfig,
litellm_metadata: dict,
custom_llm_provider: str,
request_body: Optional[dict] = None,
hidden_params: Optional[Dict[str, Any]] = None,
request_body: dict | None = None,
hidden_params: dict[str, Any] | None = None,
):
super().__init__(
litellm_logging_obj=logging_obj,
@ -162,8 +162,8 @@ class AsyncGoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateCo
generate_content_provider_config: BaseGoogleGenAIGenerateContentConfig,
litellm_metadata: dict,
custom_llm_provider: str,
request_body: Optional[dict] = None,
hidden_params: Optional[Dict[str, Any]] = None,
request_body: dict | None = None,
hidden_params: dict[str, Any] | None = None,
):
super().__init__(
litellm_logging_obj=logging_obj,

164
litellm/images/main.py
View file

@ -6,11 +6,8 @@ from functools import partial
from typing import (
TYPE_CHECKING,
Any,
Dict,
List,
Literal,
Optional,
Union,
cast,
overload,
)
@ -116,7 +113,7 @@ async def aimage_generation(*args, **kwargs) -> ImageResponse:
# Await normally
init_response = await loop.run_in_executor(None, func_with_context)
response: Optional[ImageResponse] = None
response: ImageResponse | None = None
if isinstance(init_response, dict):
response = ImageResponse(**init_response)
elif isinstance(init_response, ImageResponse): ## CACHING SCENARIO
@ -145,17 +142,17 @@ async def aimage_generation(*args, **kwargs) -> ImageResponse:
@overload
def image_generation(
prompt: str,
model: Optional[str] = None,
n: Optional[int] = None,
quality: Optional[Union[str, ImageGenerationRequestQuality]] = None,
response_format: Optional[str] = None,
size: Optional[str] = None,
style: Optional[str] = None,
user: Optional[str] = None,
model: str | None = None,
n: int | None = None,
quality: str | ImageGenerationRequestQuality | None = None,
response_format: str | None = None,
size: str | None = None,
style: str | None = None,
user: str | None = None,
timeout=600, # default to 10 minutes
api_key: Optional[str] = None,
api_base: Optional[str] = None,
api_version: Optional[str] = None,
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
custom_llm_provider=None,
*,
aimg_generation: Literal[True],
@ -169,17 +166,17 @@ def image_generation(
@overload
def image_generation(
prompt: str,
model: Optional[str] = None,
n: Optional[int] = None,
quality: Optional[Union[str, ImageGenerationRequestQuality]] = None,
response_format: Optional[str] = None,
size: Optional[str] = None,
style: Optional[str] = None,
user: Optional[str] = None,
model: str | None = None,
n: int | None = None,
quality: str | ImageGenerationRequestQuality | None = None,
response_format: str | None = None,
size: str | None = None,
style: str | None = None,
user: str | None = None,
timeout=600, # default to 10 minutes
api_key: Optional[str] = None,
api_base: Optional[str] = None,
api_version: Optional[str] = None,
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
custom_llm_provider=None,
*,
aimg_generation: Literal[False] = False,
@ -193,23 +190,20 @@ def image_generation(
@client
def image_generation(
prompt: str,
model: Optional[str] = None,
n: Optional[int] = None,
quality: Optional[Union[str, ImageGenerationRequestQuality]] = None,
response_format: Optional[str] = None,
size: Optional[str] = None,
style: Optional[str] = None,
user: Optional[str] = None,
model: str | None = None,
n: int | None = None,
quality: str | ImageGenerationRequestQuality | None = None,
response_format: str | None = None,
size: str | None = None,
style: str | None = None,
user: str | None = None,
timeout=600, # default to 10 minutes
api_key: Optional[str] = None,
api_base: Optional[str] = None,
api_version: Optional[str] = None,
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
custom_llm_provider=None,
**kwargs,
) -> Union[
ImageResponse,
Coroutine[Any, Any, ImageResponse],
]:
) -> ImageResponse | Coroutine[Any, Any, ImageResponse]:
"""
Maps the https://api.openai.com/v1/images/generations endpoint.
@ -220,7 +214,7 @@ def image_generation(
aimg_generation = kwargs.get("aimg_generation", False)
litellm_call_id = kwargs.get("litellm_call_id", None)
logger_fn = kwargs.get("logger_fn", None)
mock_response: Optional[str] = kwargs.get("mock_response", None) # type: ignore
mock_response: str | None = kwargs.get("mock_response", None) # type: ignore
proxy_server_request = kwargs.get("proxy_server_request", None)
azure_ad_token_provider = kwargs.get("azure_ad_token_provider", None)
model_info = kwargs.get("model_info", None)
@ -233,7 +227,7 @@ def image_generation(
if extra_headers is not None:
headers.update(extra_headers)
model_response: ImageResponse = litellm.utils.ImageResponse()
dynamic_api_key: Optional[str] = None
dynamic_api_key: str | None = None
if model is not None or custom_llm_provider is not None:
model, custom_llm_provider, dynamic_api_key, api_base = get_llm_provider(
model=model, # type: ignore
@ -267,7 +261,7 @@ def image_generation(
k: v for k, v in kwargs.items() if k not in default_params
} # model-specific params - pass them straight to the model/provider
image_generation_config: Optional[BaseImageGenerationConfig] = None
image_generation_config: BaseImageGenerationConfig | None = None
if custom_llm_provider is not None and custom_llm_provider in LlmProviders._member_map_.values():
image_generation_config = ProviderConfigManager.get_provider_image_generation_config(
model=base_model or model,
@ -474,7 +468,7 @@ def image_generation(
if extra_headers is not None:
optional_params["extra_headers"] = extra_headers
# Forward OpenAI organization if present (set by proxy pre-call utils)
organization: Optional[str] = kwargs.get("organization", None)
organization: str | None = kwargs.get("organization", None)
model_response = openai_chat_completions.image_generation(
model=model,
prompt=prompt,
@ -506,7 +500,7 @@ def image_generation(
)
elif custom_llm_provider in litellm._custom_providers: # Assume custom LLM provider
# Get the Custom Handler
custom_handler: Optional[CustomLLM] = None
custom_handler: CustomLLM | None = None
for item in litellm.custom_provider_map:
if item["provider"] == custom_llm_provider:
custom_handler = item["custom_handler"]
@ -516,7 +510,7 @@ def image_generation(
## ROUTE LLM CALL ##
if aimg_generation is True:
async_custom_client: Optional[AsyncHTTPHandler] = None
async_custom_client: AsyncHTTPHandler | None = None
if client is not None and isinstance(client, AsyncHTTPHandler):
async_custom_client = client
@ -533,7 +527,7 @@ def image_generation(
client=async_custom_client,
)
else:
custom_client: Optional[HTTPHandler] = None
custom_client: HTTPHandler | None = None
if client is not None and isinstance(client, HTTPHandler):
custom_client = client
@ -619,8 +613,8 @@ def image_variation(
model: str = "dall-e-2", # set to dall-e-2 by default - like OpenAI.
n: int = 1,
response_format: Literal["url", "b64_json"] = "url",
size: Optional[str] = None,
user: Optional[str] = None,
size: str | None = None,
user: str | None = None,
**kwargs,
) -> ImageResponse:
# get non-default params
@ -648,7 +642,7 @@ def image_variation(
)
model_response = ImageResponse()
response: Optional[ImageResponse] = None
response: ImageResponse | None = None
provider_config = ProviderConfigManager.get_provider_model_info(
model=model or "", # openai defaults to dall-e-2
@ -711,25 +705,25 @@ def image_variation(
@client
def image_edit(
image: Optional[Union[FileTypes, List[FileTypes]]] = None,
prompt: Optional[str] = None,
model: Optional[str] = None,
mask: Optional[str] = None,
n: Optional[int] = None,
quality: Optional[Union[str, ImageGenerationRequestQuality]] = None,
response_format: Optional[str] = None,
size: Optional[str] = None,
user: Optional[str] = None,
image: FileTypes | list[FileTypes] | None = None,
prompt: str | None = None,
model: str | None = None,
mask: str | None = None,
n: int | None = None,
quality: str | ImageGenerationRequestQuality | None = None,
response_format: str | None = None,
size: str | None = None,
user: str | None = None,
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
# The extra values given here take precedence over values defined on the client or passed to this method.
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
timeout: float | httpx.Timeout | None = None,
# LiteLLM specific params,
custom_llm_provider: Optional[str] = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> Union[ImageResponse, Coroutine[Any, Any, ImageResponse]]:
) -> ImageResponse | Coroutine[Any, Any, ImageResponse]:
"""
Maps the image edit functionality, similar to OpenAI's images/edits endpoint.
"""
@ -759,7 +753,7 @@ def image_edit(
k: v for k, v in kwargs.items() if k not in default_params
} # model-specific params - pass them straight to the model/provider
litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore
litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None)
litellm_call_id: str | None = kwargs.get("litellm_call_id", None)
model_info = kwargs.get("model_info", None)
metadata = kwargs.get("metadata", {})
_is_async = kwargs.pop("async_call", False) is True
@ -768,7 +762,7 @@ def image_edit(
images = image if isinstance(image, list) else ([image] if image is not None else [])
headers_from_kwargs = kwargs.get("headers")
merged_extra_headers: Dict[str, Any] = {}
merged_extra_headers: dict[str, Any] = {}
if isinstance(headers_from_kwargs, dict):
merged_extra_headers.update(headers_from_kwargs)
if isinstance(extra_headers, dict):
@ -786,7 +780,7 @@ def image_edit(
# Check for custom provider
if custom_llm_provider in litellm._custom_providers:
custom_handler: Optional[CustomLLM] = None
custom_handler: CustomLLM | None = None
for item in litellm.custom_provider_map:
if item["provider"] == custom_llm_provider:
custom_handler = item["custom_handler"]
@ -797,7 +791,7 @@ def image_edit(
model_response = ImageResponse()
if _is_async:
async_custom_client: Optional[AsyncHTTPHandler] = None
async_custom_client: AsyncHTTPHandler | None = None
if kwargs.get("client") is not None and isinstance(kwargs.get("client"), AsyncHTTPHandler):
async_custom_client = kwargs.get("client")
@ -814,7 +808,7 @@ def image_edit(
client=async_custom_client,
)
else:
custom_client: Optional[HTTPHandler] = None
custom_client: HTTPHandler | None = None
if kwargs.get("client") is not None and isinstance(kwargs.get("client"), HTTPHandler):
custom_client = kwargs.get("client")
@ -832,11 +826,9 @@ def image_edit(
)
# get provider config
image_edit_provider_config: Optional[BaseImageEditConfig] = (
ProviderConfigManager.get_provider_image_edit_config(
model=model,
provider=litellm.LlmProviders(custom_llm_provider),
)
image_edit_provider_config: BaseImageEditConfig | None = ProviderConfigManager.get_provider_image_edit_config(
model=model,
provider=litellm.LlmProviders(custom_llm_provider),
)
if image_edit_provider_config is None:
@ -848,7 +840,7 @@ def image_edit(
_get_ImageEditRequestUtils().get_requested_image_edit_optional_param(local_vars)
)
# Get optional parameters for the responses API
image_edit_request_params: Dict = _get_ImageEditRequestUtils().get_optional_params_image_edit(
image_edit_request_params: dict = _get_ImageEditRequestUtils().get_optional_params_image_edit(
model=model,
image_edit_provider_config=image_edit_provider_config,
image_edit_optional_params=image_edit_optional_params,
@ -952,23 +944,23 @@ def image_edit(
@client
async def aimage_edit(
image: Union[FileTypes, List[FileTypes]],
image: FileTypes | list[FileTypes],
model: str,
prompt: str,
mask: Optional[str] = None,
n: Optional[int] = None,
quality: Optional[Union[str, ImageGenerationRequestQuality]] = None,
response_format: Optional[str] = None,
size: Optional[str] = None,
user: Optional[str] = None,
mask: str | None = None,
n: int | None = None,
quality: str | ImageGenerationRequestQuality | None = None,
response_format: str | None = None,
size: str | None = None,
user: str | None = None,
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
# The extra values given here take precedence over values defined on the client or passed to this method.
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
timeout: float | httpx.Timeout | None = None,
# LiteLLM specific params,
custom_llm_provider: Optional[str] = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> ImageResponse:
"""

10
litellm/images/utils.py
View file

@ -1,5 +1,5 @@
from io import BufferedReader, BytesIO
from typing import Any, Dict, List, Optional, cast, get_type_hints
from typing import Any, cast, get_type_hints
import litellm
from litellm.litellm_core_utils.token_counter import get_image_type
@ -14,9 +14,9 @@ class ImageEditRequestUtils:
model: str,
image_edit_provider_config: BaseImageEditConfig,
image_edit_optional_params: ImageEditOptionalRequestParams,
drop_params: Optional[bool] = None,
additional_drop_params: Optional[List[str]] = None,
) -> Dict:
drop_params: bool | None = None,
additional_drop_params: list[str] | None = None,
) -> dict:
"""
Get optional parameters for the image edit API.
@ -61,7 +61,7 @@ class ImageEditRequestUtils:
@staticmethod
def get_requested_image_edit_optional_param(
params: Dict[str, Any],
params: dict[str, Any],
) -> ImageEditOptionalRequestParams:
"""
Filter parameters to only include those defined in ImageEditOptionalRequestParams.

2
litellm/integrations/SlackAlerting/batching_handler.py
View file

@ -70,6 +70,6 @@ async def send_to_webhook(slackAlertingInstance: SlackAlertingType, item, count)
if response.status_code != 200:
verbose_proxy_logger.debug(f"Error sending slack alert to url={item['url']}. Error={response.text}")
except Exception as e:
verbose_proxy_logger.debug(f"Error sending slack alert: {str(e)}")
verbose_proxy_logger.debug(f"Error sending slack alert: {e!s}")
finally:
_print_alerting_payload_warning(payload, slackAlertingInstance=slackAlertingInstance)

2
litellm/integrations/SlackAlerting/budget_alert_types.py
View file

@ -10,12 +10,10 @@ class BaseBudgetAlertType(ABC):
@abstractmethod
def get_event_message(self) -> str:
"""Return the event message for this alert type"""
pass
@abstractmethod
def get_id(self, user_info: CallInfo) -> str:
"""Return the ID to use for caching/tracking this alert"""
pass
class ProxyBudgetAlert(BaseBudgetAlertType):

10
litellm/integrations/SlackAlerting/hanging_request_check.py
View file

@ -9,7 +9,7 @@ Notes:
import asyncio
import time
from typing import TYPE_CHECKING, Any, Optional
from typing import TYPE_CHECKING, Any
import litellm
from litellm._logging import verbose_proxy_logger
@ -49,7 +49,7 @@ class AlertingHangingRequestCheck:
async def add_request_to_hanging_request_check(
self,
request_data: Optional[dict] = None,
request_data: dict | None = None,
):
"""
Add a request to the hanging request cache. This is the list of request_ids that gets periodicall checked for hanging requests
@ -59,7 +59,7 @@ class AlertingHangingRequestCheck:
request_metadata = get_litellm_metadata_from_kwargs(kwargs=request_data)
model = request_data.get("model", "")
api_base: Optional[str] = None
api_base: str | None = None
if request_data.get("deployment", None) is not None and isinstance(request_data["deployment"], dict):
api_base = litellm.get_api_base(
@ -101,7 +101,7 @@ class AlertingHangingRequestCheck:
)
for request_id in hanging_requests:
hanging_request_data: Optional[HangingRequestData] = await self.hanging_request_cache.async_get_cache(
hanging_request_data: HangingRequestData | None = await self.hanging_request_cache.async_get_cache(
key=request_id,
)
@ -112,7 +112,7 @@ class AlertingHangingRequestCheck:
continue
request_status = await proxy_logging_obj.internal_usage_cache.async_get_cache(
key="request_status:{}".format(hanging_request_data.request_id),
key=f"request_status:{hanging_request_data.request_id}",
litellm_parent_otel_span=None,
local_only=True,
)

159
litellm/integrations/SlackAlerting/slack_alerting.py
View file

@ -6,7 +6,7 @@ import os
import random
import time
from datetime import timedelta
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union
from typing import TYPE_CHECKING, Any, Literal
from openai import APIError
@ -61,16 +61,15 @@ class SlackAlerting(CustomBatchLogger):
# Class variables or attributes
def __init__(
self,
internal_usage_cache: Optional[DualCache] = None,
alerting_threshold: Optional[float] = None, # threshold for slow / hanging llm responses (in seconds)
alerting: Optional[List] = [],
alert_types: List[AlertType] = DEFAULT_ALERT_TYPES,
alert_to_webhook_url: Optional[
Dict[AlertType, Union[List[str], str]]
] = None, # if user wants to separate alerts to diff channels
internal_usage_cache: DualCache | None = None,
alerting_threshold: float | None = None, # threshold for slow / hanging llm responses (in seconds)
alerting: list | None = [],
alert_types: list[AlertType] = DEFAULT_ALERT_TYPES,
alert_to_webhook_url: dict[AlertType, list[str] | str]
| None = None, # if user wants to separate alerts to diff channels
alerting_args={},
default_webhook_url: Optional[str] = None,
alert_type_config: Optional[Dict[str, dict]] = None,
default_webhook_url: str | None = None,
alert_type_config: dict[str, dict] | None = None,
**kwargs,
):
if alerting_threshold is None:
@ -89,23 +88,23 @@ class SlackAlerting(CustomBatchLogger):
self.hanging_request_check = AlertingHangingRequestCheck(
slack_alerting_object=self,
)
self.alert_type_config: Dict[str, AlertTypeConfig] = {}
self.alert_type_config: dict[str, AlertTypeConfig] = {}
if alert_type_config:
for key, val in alert_type_config.items():
self.alert_type_config[key] = AlertTypeConfig(**val) if isinstance(val, dict) else val
self.digest_buckets: Dict[str, DigestEntry] = {}
self.digest_buckets: dict[str, DigestEntry] = {}
self.digest_lock = asyncio.Lock()
super().__init__(**kwargs, flush_lock=self.flush_lock)
def update_values(
self,
alerting: Optional[List] = None,
alerting_threshold: Optional[float] = None,
alert_types: Optional[List[AlertType]] = None,
alert_to_webhook_url: Optional[Dict[AlertType, Union[List[str], str]]] = None,
alerting_args: Optional[Dict] = None,
llm_router: Optional[Router] = None,
alert_type_config: Optional[Dict[str, dict]] = None,
alerting: list | None = None,
alerting_threshold: float | None = None,
alert_types: list[AlertType] | None = None,
alert_to_webhook_url: dict[AlertType, list[str] | str] | None = None,
alerting_args: dict | None = None,
llm_router: Router | None = None,
alert_type_config: dict[str, dict] | None = None,
):
if alerting is not None:
self.alerting = alerting
@ -134,9 +133,7 @@ class SlackAlerting(CustomBatchLogger):
if llm_router is not None:
self.llm_router = llm_router
def _prepare_outage_value_for_cache(
self, outage_value: Union[dict, ProviderRegionOutageModel, OutageModel]
) -> dict:
def _prepare_outage_value_for_cache(self, outage_value: dict | ProviderRegionOutageModel | OutageModel) -> dict:
"""
Helper method to prepare outage value for Redis caching.
Converts set objects to lists for JSON serialization.
@ -148,7 +145,7 @@ class SlackAlerting(CustomBatchLogger):
cache_value["deployment_ids"] = list(cache_value["deployment_ids"])
return cache_value
def _restore_outage_value_from_cache(self, outage_value: Optional[dict]) -> Optional[dict]:
def _restore_outage_value_from_cache(self, outage_value: dict | None) -> dict | None:
"""
Helper method to restore outage value after retrieving from cache.
Converts list objects back to sets for proper handling.
@ -210,7 +207,7 @@ class SlackAlerting(CustomBatchLogger):
_deployment_latencies = metadata["_latency_per_deployment"]
if len(_deployment_latencies) == 0:
return None
_deployment_latency_map: Optional[dict] = None
_deployment_latency_map: dict | None = None
try:
# try sorting deployments by latency
_deployment_latencies = sorted(_deployment_latencies.items(), key=lambda x: x[1])
@ -290,10 +287,7 @@ class SlackAlerting(CustomBatchLogger):
## FAILED REQUESTS ##
if deployment_metrics.failed_request:
await self.internal_usage_cache.async_increment_cache(
key="{}:{}".format(
deployment_metrics.id,
SlackAlertingCacheKeys.failed_requests_key.value,
),
key=f"{deployment_metrics.id}:{SlackAlertingCacheKeys.failed_requests_key.value}",
value=1,
parent_otel_span=None, # no attached request, this is a background operation
)
@ -303,7 +297,7 @@ class SlackAlerting(CustomBatchLogger):
## LATENCY ##
if deployment_metrics.latency_per_output_token is not None:
await self.internal_usage_cache.async_increment_cache(
key="{}:{}".format(deployment_metrics.id, SlackAlertingCacheKeys.latency_key.value),
key=f"{deployment_metrics.id}:{SlackAlertingCacheKeys.latency_key.value}",
value=deployment_metrics.latency_per_output_token,
parent_otel_span=None, # no attached request, this is a background operation
)
@ -333,8 +327,8 @@ class SlackAlerting(CustomBatchLogger):
ids = router.get_model_ids()
# get keys
failed_request_keys = ["{}:{}".format(id, SlackAlertingCacheKeys.failed_requests_key.value) for id in ids]
latency_keys = ["{}:{}".format(id, SlackAlertingCacheKeys.latency_key.value) for id in ids]
failed_request_keys = [f"{id}:{SlackAlertingCacheKeys.failed_requests_key.value}" for id in ids]
latency_keys = [f"{id}:{SlackAlertingCacheKeys.latency_key.value}" for id in ids]
combined_metrics_keys = failed_request_keys + latency_keys # reduce cache calls
@ -445,7 +439,7 @@ class SlackAlerting(CustomBatchLogger):
async def response_taking_too_long(
self,
request_data: Optional[dict] = None,
request_data: dict | None = None,
):
if self.alerting is None or self.alert_types is None:
return
@ -471,7 +465,7 @@ class SlackAlerting(CustomBatchLogger):
_cache: DualCache = self.internal_usage_cache
message = "Failed Tracking Cost for " + error_message
_cache_key = "budget_alerts:failed_tracking:{}".format(failing_model)
_cache_key = f"budget_alerts:failed_tracking:{failing_model}"
result = await _cache.async_get_cache(key=_cache_key)
if result is None:
await self.send_alert(
@ -528,16 +522,11 @@ class SlackAlerting(CustomBatchLogger):
event_message = budget_alert_class.get_event_message()
# Set default event unless we're in projected_limit_exceeded
event: Optional[
Literal[
"budget_crossed",
"threshold_crossed",
"projected_limit_exceeded",
"soft_budget_crossed",
]
] = "projected_limit_exceeded" if type == "projected_limit_exceeded" else None
event: (
Literal["budget_crossed", "threshold_crossed", "projected_limit_exceeded", "soft_budget_crossed"] | None
) = "projected_limit_exceeded" if type == "projected_limit_exceeded" else None
webhook_event: Optional[WebhookEvent] = None
webhook_event: WebhookEvent | None = None
# percent of max_budget left to spend
if user_info.max_budget is None and user_info.soft_budget is None:
@ -552,7 +541,7 @@ class SlackAlerting(CustomBatchLogger):
# send alert
if event is not None and user_info.event_group is not None:
_cache_key = "budget_alerts:{}:{}".format(event, _id)
_cache_key = f"budget_alerts:{event}:{_id}"
result = await _cache.async_get_cache(key=_cache_key)
if result is None:
webhook_event = WebhookEvent(
@ -579,24 +568,10 @@ class SlackAlerting(CustomBatchLogger):
def _get_event_and_event_message(
self,
user_info: CallInfo,
event: Optional[
Literal[
"budget_crossed",
"threshold_crossed",
"soft_budget_crossed",
"projected_limit_exceeded",
]
],
event: Literal["budget_crossed", "threshold_crossed", "soft_budget_crossed", "projected_limit_exceeded"] | None,
event_message: str,
) -> Tuple[
Optional[
Literal[
"budget_crossed",
"threshold_crossed",
"soft_budget_crossed",
"projected_limit_exceeded",
]
],
) -> tuple[
Literal["budget_crossed", "threshold_crossed", "soft_budget_crossed", "projected_limit_exceeded"] | None,
str,
]:
"""
@ -642,7 +617,7 @@ class SlackAlerting(CustomBatchLogger):
"""
percent_left: float = 0.0
current_spend: float = user_info.spend
max_budget: Optional[float] = user_info.max_budget
max_budget: float | None = user_info.max_budget
if max_budget is None:
return percent_left
if max_budget <= 0:
@ -666,11 +641,11 @@ class SlackAlerting(CustomBatchLogger):
async def customer_spend_alert(
self,
token: Optional[str],
key_alias: Optional[str],
end_user_id: Optional[str],
response_cost: Optional[float],
max_budget: Optional[float],
token: str | None,
key_alias: str | None,
end_user_id: str | None,
response_cost: float | None,
max_budget: float | None,
):
if (
self.alerting is not None
@ -693,12 +668,12 @@ class SlackAlerting(CustomBatchLogger):
projected_spend=None,
event="spend_tracked",
event_group=Litellm_EntityType.END_USER,
event_message="Customer spend tracked. Customer={}, spend={}".format(end_user_id, response_cost),
event_message=f"Customer spend tracked. Customer={end_user_id}, spend={response_cost}",
)
await self.send_webhook_alert(webhook_event=event)
def _count_outage_alerts(self, alerts: List[int]) -> str:
def _count_outage_alerts(self, alerts: list[int]) -> str:
"""
Parameters:
- alerts: List[int] -> list of error codes (either 408 or 500+)
@ -718,7 +693,7 @@ class SlackAlerting(CustomBatchLogger):
error_msg = ""
for key, value in error_breakdown.items():
if value > 0:
error_msg += "\n{}: {}\n".format(key, value)
error_msg += f"\n{key}: {value}\n"
return error_msg
@ -728,7 +703,7 @@ class SlackAlerting(CustomBatchLogger):
key: Literal["Model", "Region"],
key_val: str,
provider: str,
api_base: Optional[str],
api_base: str | None,
outage_value: BaseOutageModel,
) -> str:
"""Format an alert message for slack"""
@ -788,9 +763,7 @@ class SlackAlerting(CustomBatchLogger):
### UNIQUE CACHE KEY ###
cache_key = provider + region_name
outage_value: Optional[ProviderRegionOutageModel] = await self.internal_usage_cache.async_get_cache(
key=cache_key
)
outage_value: ProviderRegionOutageModel | None = await self.internal_usage_cache.async_get_cache(key=cache_key)
# Convert deployment_ids back to set if it was stored as a list
if outage_value is not None:
@ -911,7 +884,7 @@ class SlackAlerting(CustomBatchLogger):
max_alerts_size = 10
"""
try:
outage_value: Optional[OutageModel] = await self.internal_usage_cache.async_get_cache(key=deployment_id) # type: ignore
outage_value: OutageModel | None = await self.internal_usage_cache.async_get_cache(key=deployment_id) # type: ignore
if (
getattr(exception, "status_code", None) is None
or (
@ -1024,7 +997,7 @@ class SlackAlerting(CustomBatchLogger):
for k, v in model_info.items():
if k == "input_cost_per_token" or k == "output_cost_per_token":
# when converting to string it should not be 1.63e-06
v = "{:.8f}".format(v)
v = f"{v:.8f}"
model_info_str += f"{k}: {v}\n"
@ -1105,15 +1078,14 @@ Model Info:
async def _check_if_using_premium_email_feature(
self,
premium_user: bool,
email_logo_url: Optional[str] = None,
email_support_contact: Optional[str] = None,
email_logo_url: str | None = None,
email_support_contact: str | None = None,
):
from litellm.proxy.proxy_server import CommonProxyErrors, premium_user
if premium_user is not True:
if email_logo_url is not None or email_support_contact is not None:
raise ValueError(f"Trying to Customize Email Alerting\n {CommonProxyErrors.not_premium_user.value}")
return
async def send_key_created_or_user_invited_email(self, webhook_event: WebhookEvent) -> bool:
try:
@ -1274,9 +1246,9 @@ Model Info:
level: Literal["Low", "Medium", "High"],
alert_type: AlertType,
alerting_metadata: dict,
user_info: Optional[WebhookEvent] = None,
request_model: Optional[str] = None,
api_base: Optional[str] = None,
user_info: WebhookEvent | None = None,
request_model: str | None = None,
api_base: str | None = None,
**kwargs,
):
"""
@ -1323,7 +1295,7 @@ Model Info:
if _atc is not None and _atc.digest:
# Resolve webhook URL for this alert type (needed for digest entry)
if self.alert_to_webhook_url is not None and alert_type in self.alert_to_webhook_url:
_digest_webhook: Optional[Union[str, List[str]]] = self.alert_to_webhook_url[alert_type]
_digest_webhook: str | list[str] | None = self.alert_to_webhook_url[alert_type]
elif self.default_webhook_url is not None:
_digest_webhook = self.default_webhook_url
else:
@ -1376,7 +1348,7 @@ Model Info:
# check if we find the slack webhook url in self.alert_to_webhook_url
if self.alert_to_webhook_url is not None and alert_type in self.alert_to_webhook_url:
slack_webhook_url: Optional[Union[str, List[str]]] = self.alert_to_webhook_url[alert_type]
slack_webhook_url: str | list[str] | None = self.alert_to_webhook_url[alert_type]
elif self.default_webhook_url is not None:
slack_webhook_url = self.default_webhook_url
else:
@ -1431,7 +1403,7 @@ Model Info:
from datetime import datetime
now = datetime.now()
flushed_keys: List[str] = []
flushed_keys: list[str] = []
async with self.digest_lock:
for key, entry in self.digest_buckets.items():
@ -1495,7 +1467,7 @@ Model Info:
try:
await self._flush_digest_buckets()
except Exception as e:
verbose_proxy_logger.debug(f"Error flushing digest buckets: {str(e)}")
verbose_proxy_logger.debug(f"Error flushing digest buckets: {e!s}")
await self.flush_queue()
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
@ -1530,9 +1502,8 @@ Model Info:
)
except Exception as e:
verbose_proxy_logger.error(
f"[Non-Blocking Error] Slack Alerting: Got error in logging LLM deployment latency: {str(e)}"
f"[Non-Blocking Error] Slack Alerting: Got error in logging LLM deployment latency: {e!s}"
)
pass
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
"""Log failure + deployment latency"""
@ -1551,7 +1522,7 @@ Model Info:
)
)
except Exception as e:
verbose_logger.debug(f"Exception raises -{str(e)}")
verbose_logger.debug(f"Exception raises -{e!s}")
if isinstance(kwargs.get("exception", ""), APIError):
if "outage_alerts" in self.alert_types:
@ -1601,7 +1572,7 @@ Model Info:
return report_sent_bool
async def _run_scheduled_daily_report(self, llm_router: Optional[Any] = None):
async def _run_scheduled_daily_report(self, llm_router: Any | None = None):
"""
If 'daily_reports' enabled
@ -1785,8 +1756,6 @@ Model Info:
except Exception as e:
verbose_proxy_logger.error("Error sending weekly spend report %s", e)
pass
async def send_virtual_key_event_slack(
self,
key_event: VirtualKeyEvent,
@ -1830,9 +1799,7 @@ Model Info:
except Exception as e:
verbose_proxy_logger.error("Error sending send_virtual_key_event_slack %s", e)
return
async def _request_is_completed(self, request_data: Optional[dict]) -> bool:
async def _request_is_completed(self, request_data: dict | None) -> bool:
"""
Returns True if the request is completed - either as a success or failure
"""
@ -1842,8 +1809,8 @@ Model Info:
if request_data.get("litellm_status", "") != "success" and request_data.get("litellm_status", "") != "fail":
## CHECK IF CACHE IS UPDATED
litellm_call_id = request_data.get("litellm_call_id", "")
status: Optional[str] = await self.internal_usage_cache.async_get_cache(
key="request_status:{}".format(litellm_call_id), local_only=True
status: str | None = await self.internal_usage_cache.async_get_cache(
key=f"request_status:{litellm_call_id}", local_only=True
)
if status is not None and (status == "success" or status == "fail"):
return True

14
litellm/integrations/SlackAlerting/utils.py
View file

@ -3,7 +3,7 @@ Utils used for slack alerting
"""
import asyncio
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
from typing import TYPE_CHECKING, Any
import litellm
from litellm.proxy._types import AlertType
@ -18,8 +18,8 @@ else:
def process_slack_alerting_variables(
alert_to_webhook_url: Optional[Dict[AlertType, Union[List[str], str]]],
) -> Optional[Dict[AlertType, Union[List[str], str]]]:
alert_to_webhook_url: dict[AlertType, list[str] | str] | None,
) -> dict[AlertType, list[str] | str] | None:
"""
process alert_to_webhook_url
- check if any urls are set as os.environ/SLACK_WEBHOOK_URL_1 read env var and set the correct value
@ -29,7 +29,7 @@ def process_slack_alerting_variables(
for alert_type, webhook_urls in alert_to_webhook_url.items():
if isinstance(webhook_urls, list):
_webhook_values: List[str] = []
_webhook_values: list[str] = []
for webhook_url in webhook_urls:
if "os.environ/" in webhook_url:
_env_value = get_secret(secret_name=webhook_url)
@ -56,8 +56,8 @@ def process_slack_alerting_variables(
async def _add_langfuse_trace_id_to_alert(
request_data: Optional[dict] = None,
) -> Optional[str]:
request_data: dict | None = None,
) -> str | None:
"""
Returns langfuse trace url
@ -73,7 +73,7 @@ async def _add_langfuse_trace_id_to_alert(
#########################################################
if request_data is not None and request_data.get("litellm_logging_obj", None) is not None:
trace_id: Optional[str] = None
trace_id: str | None = None
litellm_logging_obj: Logging = request_data["litellm_logging_obj"]
for _ in range(3):

8
litellm/integrations/additional_logging_utils.py
View file

@ -7,7 +7,6 @@ Base class for Additional Logging Utils for CustomLoggers
from abc import ABC, abstractmethod
from datetime import datetime
from typing import Optional
from litellm.types.integrations.base_health_check import IntegrationHealthCheckStatus
@ -21,15 +20,14 @@ class AdditionalLoggingUtils(ABC):
"""
Check if the service is healthy
"""
pass
@abstractmethod
async def get_request_response_payload(
self,
request_id: str,
start_time_utc: Optional[datetime],
end_time_utc: Optional[datetime],
) -> Optional[dict]:
start_time_utc: datetime | None,
end_time_utc: datetime | None,
) -> dict | None:
"""
Get the request and response payload for a given `request_id`
"""

13
litellm/integrations/agentops/agentops.py
View file

@ -4,7 +4,8 @@ AgentOps integration for LiteLLM - Provides OpenTelemetry tracing for LLM calls
import os
from dataclasses import dataclass
from typing import Optional, Dict, Any
from typing import Any
from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig
from litellm.llms.custom_httpx.http_handler import _get_httpx_client
@ -12,9 +13,9 @@ from litellm.llms.custom_httpx.http_handler import _get_httpx_client
@dataclass
class AgentOpsConfig:
endpoint: str = "https://otlp.agentops.cloud/v1/traces"
api_key: Optional[str] = None
service_name: Optional[str] = None
deployment_environment: Optional[str] = None
api_key: str | None = None
service_name: str | None = None
deployment_environment: str | None = None
auth_endpoint: str = "https://api.agentops.ai/v3/auth/token"
@classmethod
@ -47,7 +48,7 @@ class AgentOps(OpenTelemetry):
def __init__(
self,
config: Optional[AgentOpsConfig] = None,
config: AgentOpsConfig | None = None,
):
if config is None:
config = AgentOpsConfig.from_env()
@ -82,7 +83,7 @@ class AgentOps(OpenTelemetry):
self.resource_attributes = resource_attrs
def _fetch_auth_token(self, api_key: str, auth_endpoint: str) -> Dict[str, Any]:
def _fetch_auth_token(self, api_key: str, auth_endpoint: str) -> dict[str, Any]:
"""
Fetch JWT authentication token from AgentOps API

112
litellm/integrations/anthropic_cache_control_hook.py
View file

@ -10,7 +10,7 @@ Supported for both `v1/chat/completions` (via the prompt-management hook) and
"""
import copy
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast
from typing import TYPE_CHECKING, Any, cast
from litellm._logging import verbose_logger
from litellm.integrations.custom_logger import CustomLogger
@ -39,17 +39,17 @@ class AnthropicCacheControlHook(CustomPromptManagement):
def get_chat_completion_prompt(
self,
model: str,
messages: List[AllMessageValues],
messages: list[AllMessageValues],
non_default_params: dict,
prompt_id: Optional[str],
prompt_variables: Optional[dict],
prompt_id: str | None,
prompt_variables: dict | None,
dynamic_callback_params: StandardCallbackDynamicParams,
prompt_spec: Optional[PromptSpec] = None,
prompt_label: Optional[str] = None,
prompt_version: Optional[int] = None,
ignore_prompt_manager_model: Optional[bool] = False,
ignore_prompt_manager_optional_params: Optional[bool] = False,
) -> Tuple[str, List[AllMessageValues], dict]:
prompt_spec: PromptSpec | None = None,
prompt_label: str | None = None,
prompt_version: int | None = None,
ignore_prompt_manager_model: bool | None = False,
ignore_prompt_manager_optional_params: bool | None = False,
) -> tuple[str, list[AllMessageValues], dict]:
"""
Apply cache control directives based on specified injection points.
@ -59,7 +59,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
- non_default_params: dict - params with any global cache controls
"""
# Extract cache control injection points
injection_points: List[CacheControlInjectionPoint] = non_default_params.pop(
injection_points: list[CacheControlInjectionPoint] = non_default_params.pop(
"cache_control_injection_points", []
)
if not injection_points:
@ -69,8 +69,8 @@ class AnthropicCacheControlHook(CustomPromptManagement):
processed_messages = copy.deepcopy(messages)
# Separate message-level and non-message-level injection points
message_points: List[CacheControlMessageInjectionPoint] = []
remaining_points: List[CacheControlInjectionPoint] = []
message_points: list[CacheControlMessageInjectionPoint] = []
remaining_points: list[CacheControlInjectionPoint] = []
for point in injection_points:
if point.get("location") == "message":
message_points.append(cast(CacheControlMessageInjectionPoint, point))
@ -99,10 +99,10 @@ class AnthropicCacheControlHook(CustomPromptManagement):
@staticmethod
def _apply_message_injections(
points: List[CacheControlMessageInjectionPoint],
messages: List[AllMessageValues],
points: list[CacheControlMessageInjectionPoint],
messages: list[AllMessageValues],
max_blocks: int,
) -> List[AllMessageValues]:
) -> list[AllMessageValues]:
"""Apply message-level cache control injection points in order.
Anthropic allows at most ``MAX_CACHE_CONTROL_BLOCKS`` cache_control
@ -151,11 +151,11 @@ class AnthropicCacheControlHook(CustomPromptManagement):
@staticmethod
def _resolve_target_indices(
point: CacheControlMessageInjectionPoint, messages: List[AllMessageValues]
) -> List[int]:
point: CacheControlMessageInjectionPoint, messages: list[AllMessageValues]
) -> list[int]:
"""Resolve which message indices an injection point targets."""
_targetted_index: Optional[Union[int, str]] = point.get("index", None)
targetted_index: Optional[int] = None
_targetted_index: int | str | None = point.get("index", None)
targetted_index: int | None = None
if isinstance(_targetted_index, str):
try:
targetted_index = int(_targetted_index)
@ -232,10 +232,10 @@ class AnthropicCacheControlHook(CustomPromptManagement):
@staticmethod
def apply_to_anthropic_messages_request(
messages: List[Dict],
messages: list[dict],
system: str | list | None,
injection_points: List[CacheControlInjectionPoint],
) -> Tuple[List[Dict], str | list | None, List[CacheControlInjectionPoint]]:
injection_points: list[CacheControlInjectionPoint],
) -> tuple[list[dict], str | list | None, list[CacheControlInjectionPoint]]:
"""Apply cache control injection for the Anthropic-native v1/messages endpoint.
Returns (messages, system, remaining_non_message_points).
@ -243,12 +243,12 @@ class AnthropicCacheControlHook(CustomPromptManagement):
if not injection_points:
return messages, system, []
processed_messages: List[Dict] = copy.deepcopy(messages)
processed_messages: list[dict] = copy.deepcopy(messages)
processed_system = copy.deepcopy(system) if system is not None else None
message_points: List[CacheControlMessageInjectionPoint] = []
system_points: List[CacheControlMessageInjectionPoint] = []
remaining_points: List[CacheControlInjectionPoint] = []
message_points: list[CacheControlMessageInjectionPoint] = []
system_points: list[CacheControlMessageInjectionPoint] = []
remaining_points: list[CacheControlInjectionPoint] = []
for point in injection_points:
if point.get("location") == "message":
@ -292,7 +292,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
processed_messages = AnthropicCacheControlHook._apply_message_injections(
points=message_points,
messages=cast(List[AllMessageValues], processed_messages),
messages=cast(list[AllMessageValues], processed_messages),
max_blocks=max_blocks - used_blocks,
)
@ -462,13 +462,13 @@ class AnthropicCacheControlHook(CustomPromptManagement):
@staticmethod
def maybe_inject_cache_control(
messages: List[Dict],
messages: list[dict],
system: str | list | None,
kwargs: Dict[str, Any],
kwargs: dict[str, Any],
model: str | None = None,
custom_llm_provider: str | None = None,
tools: list[dict] | None = None,
) -> Tuple[List[Dict], str | list | None]:
) -> tuple[list[dict], str | list | None]:
"""Extract cache_control_injection_points from kwargs and apply if present.
Configured points stand down entirely when the client already marked
@ -515,8 +515,8 @@ class AnthropicCacheControlHook(CustomPromptManagement):
def should_run_prompt_management(
self,
prompt_id: Optional[str],
prompt_spec: Optional[PromptSpec],
prompt_id: str | None,
prompt_spec: PromptSpec | None,
dynamic_callback_params: StandardCallbackDynamicParams,
) -> bool:
"""Always return False since this is not a true prompt management system."""
@ -524,12 +524,12 @@ class AnthropicCacheControlHook(CustomPromptManagement):
def _compile_prompt_helper(
self,
prompt_id: Optional[str],
prompt_spec: Optional[PromptSpec],
prompt_variables: Optional[dict],
prompt_id: str | None,
prompt_spec: PromptSpec | None,
prompt_variables: dict | None,
dynamic_callback_params: StandardCallbackDynamicParams,
prompt_label: Optional[str] = None,
prompt_version: Optional[int] = None,
prompt_label: str | None = None,
prompt_version: int | None = None,
) -> PromptManagementClient:
"""Not used - this hook only modifies messages, doesn't fetch prompts."""
return PromptManagementClient(
@ -542,12 +542,12 @@ class AnthropicCacheControlHook(CustomPromptManagement):
async def async_compile_prompt_helper(
self,
prompt_id: Optional[str],
prompt_variables: Optional[dict],
prompt_id: str | None,
prompt_variables: dict | None,
dynamic_callback_params: StandardCallbackDynamicParams,
prompt_spec: Optional[PromptSpec] = None,
prompt_label: Optional[str] = None,
prompt_version: Optional[int] = None,
prompt_spec: PromptSpec | None = None,
prompt_label: str | None = None,
prompt_version: int | None = None,
) -> PromptManagementClient:
"""Not used - this hook only modifies messages, doesn't fetch prompts."""
return self._compile_prompt_helper(
@ -562,19 +562,19 @@ class AnthropicCacheControlHook(CustomPromptManagement):
async def async_get_chat_completion_prompt(
self,
model: str,
messages: List[AllMessageValues],
messages: list[AllMessageValues],
non_default_params: dict,
prompt_id: Optional[str],
prompt_variables: Optional[dict],
prompt_id: str | None,
prompt_variables: dict | None,
dynamic_callback_params: StandardCallbackDynamicParams,
litellm_logging_obj: LiteLLMLoggingObj,
prompt_spec: Optional[PromptSpec] = None,
tools: Optional[List[Dict]] = None,
prompt_label: Optional[str] = None,
prompt_version: Optional[int] = None,
ignore_prompt_manager_model: Optional[bool] = False,
ignore_prompt_manager_optional_params: Optional[bool] = False,
) -> Tuple[str, List[AllMessageValues], dict]:
prompt_spec: PromptSpec | None = None,
tools: list[dict] | None = None,
prompt_label: str | None = None,
prompt_version: int | None = None,
ignore_prompt_manager_model: bool | None = False,
ignore_prompt_manager_optional_params: bool | None = False,
) -> tuple[str, list[AllMessageValues], dict]:
"""Async version - delegates to sync since no async operations needed."""
return self.get_chat_completion_prompt(
model=model,
@ -591,15 +591,15 @@ class AnthropicCacheControlHook(CustomPromptManagement):
)
@staticmethod
def should_use_anthropic_cache_control_hook(non_default_params: Dict) -> bool:
def should_use_anthropic_cache_control_hook(non_default_params: dict) -> bool:
if non_default_params.get("cache_control_injection_points", None):
return True
return False
@staticmethod
def get_custom_logger_for_anthropic_cache_control_hook(
non_default_params: Dict,
) -> Optional[CustomLogger]:
non_default_params: dict,
) -> CustomLogger | None:
from litellm.litellm_core_utils.litellm_logging import (
_init_custom_logger_compatible_class,
)

34
litellm/integrations/argilla.py
View file

@ -7,7 +7,7 @@ import json
import os
import random
import types
from typing import Any, Dict, List, Optional
from typing import Any
import httpx
from pydantic import BaseModel # type: ignore
@ -41,9 +41,9 @@ def is_serializable(value):
class ArgillaLogger(CustomBatchLogger):
def __init__(
self,
argilla_api_key: Optional[str] = None,
argilla_dataset_name: Optional[str] = None,
argilla_base_url: Optional[str] = None,
argilla_api_key: str | None = None,
argilla_dataset_name: str | None = None,
argilla_base_url: str | None = None,
**kwargs,
):
if litellm.argilla_transformation_object is None:
@ -69,7 +69,7 @@ class ArgillaLogger(CustomBatchLogger):
self.flush_lock = asyncio.Lock()
super().__init__(**kwargs, flush_lock=self.flush_lock)
def validate_argilla_transformation_object(self, argilla_transformation_object: Dict[str, Any]):
def validate_argilla_transformation_object(self, argilla_transformation_object: dict[str, Any]):
if not isinstance(argilla_transformation_object, dict):
raise Exception("'argilla_transformation_object' must be a dictionary, to log your payload to Argilla.")
@ -81,9 +81,9 @@ class ArgillaLogger(CustomBatchLogger):
def get_credentials_from_env(
self,
argilla_api_key: Optional[str],
argilla_dataset_name: Optional[str],
argilla_base_url: Optional[str],
argilla_api_key: str | None,
argilla_dataset_name: str | None,
argilla_base_url: str | None,
) -> ArgillaCredentialsObject:
_credentials_api_key = argilla_api_key or os.getenv("ARGILLA_API_KEY")
if _credentials_api_key is None:
@ -115,7 +115,7 @@ class ArgillaLogger(CustomBatchLogger):
ARGILLA_DATASET_NAME=_credentials_dataset_name,
)
def get_chat_messages(self, payload: StandardLoggingPayload) -> List[Dict[str, Any]]:
def get_chat_messages(self, payload: StandardLoggingPayload) -> list[dict[str, Any]]:
payload_messages = payload.get("messages", None)
if payload_messages is None:
@ -141,10 +141,10 @@ class ArgillaLogger(CustomBatchLogger):
else:
raise Exception(f"Invalid response format: {response}")
def _prepare_log_data(self, kwargs, response_obj, start_time, end_time) -> Optional[ArgillaItem]:
def _prepare_log_data(self, kwargs, response_obj, start_time, end_time) -> ArgillaItem | None:
try:
# Ensure everything in the payload is converted to str
payload: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object", None)
payload: StandardLoggingPayload | None = kwargs.get("standard_logging_object", None)
if payload is None:
raise Exception("Error logging request payload. Payload=none.")
@ -204,9 +204,7 @@ class ArgillaLogger(CustomBatchLogger):
random_sample = random.random()
if random_sample > sampling_rate:
verbose_logger.info(
"Skipping Langsmith logging. Sampling rate={}, random_sample={}".format(
sampling_rate, random_sample
)
f"Skipping Langsmith logging. Sampling rate={sampling_rate}, random_sample={random_sample}"
)
return # Skip logging
verbose_logger.debug(
@ -233,9 +231,7 @@ class ArgillaLogger(CustomBatchLogger):
random_sample = random.random()
if random_sample > sampling_rate:
verbose_logger.info(
"Skipping Langsmith logging. Sampling rate={}, random_sample={}".format(
sampling_rate, random_sample
)
f"Skipping Langsmith logging. Sampling rate={sampling_rate}, random_sample={random_sample}"
)
return # Skip logging
verbose_logger.debug(
@ -243,7 +239,7 @@ class ArgillaLogger(CustomBatchLogger):
kwargs,
response_obj,
)
payload: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object", None)
payload: StandardLoggingPayload | None = kwargs.get("standard_logging_object", None)
data = self._prepare_log_data(kwargs, response_obj, start_time, end_time)
@ -276,7 +272,7 @@ class ArgillaLogger(CustomBatchLogger):
random_sample = random.random()
if random_sample > sampling_rate:
verbose_logger.info(
"Skipping Langsmith logging. Sampling rate={}, random_sample={}".format(sampling_rate, random_sample)
f"Skipping Langsmith logging. Sampling rate={sampling_rate}, random_sample={random_sample}"
)
return # Skip logging
verbose_logger.info("Langsmith Failure Event Logging!")

6
litellm/integrations/arize/__init__.py
View file

@ -1,16 +1,16 @@
import os
from typing import TYPE_CHECKING, Optional
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from litellm.types.prompts.init_prompts import PromptLiteLLMParams, PromptSpec
from litellm.integrations.custom_prompt_management import CustomPromptManagement
from litellm.types.prompts.init_prompts import PromptLiteLLMParams, PromptSpec
from litellm.types.prompts.init_prompts import SupportedPromptIntegrations
from .arize_phoenix_prompt_manager import ArizePhoenixPromptManager
# Global instances
global_arize_config: Optional[dict] = None
global_arize_config: dict | None = None
def prompt_initializer(litellm_params: "PromptLiteLLMParams", prompt_spec: "PromptSpec") -> "CustomPromptManagement":

30
litellm/integrations/arize/_utils.py
View file

@ -1,5 +1,5 @@
import json
from typing import TYPE_CHECKING, Any, Dict, Optional, Type
from typing import TYPE_CHECKING, Any
from typing_extensions import override
@ -31,7 +31,7 @@ from litellm.integrations._types.open_inference import (
class ArizeOTELAttributes(BaseLLMObsOTELAttributes):
@staticmethod
@override
def set_messages(span: "Span", kwargs: Dict[str, Any]):
def set_messages(span: "Span", kwargs: dict[str, Any]):
messages = kwargs.get("messages")
# for /chat/completions
@ -302,7 +302,7 @@ def _set_usage_outputs(span: "Span", response_obj, span_attrs):
)
def _infer_open_inference_span_kind(call_type: Optional[str]) -> str:
def _infer_open_inference_span_kind(call_type: str | None) -> str:
"""
Map LiteLLM call types to OpenInference span kinds.
"""
@ -360,7 +360,7 @@ def _infer_open_inference_span_kind(call_type: Optional[str]) -> str:
return OpenInferenceSpanKindValues.UNKNOWN.value
def _set_tool_attributes(span: "Span", optional_tools: Optional[list], metadata_tools: Optional[list]):
def _set_tool_attributes(span: "Span", optional_tools: list | None, metadata_tools: list | None):
"""set tool attributes on span from optional_params or tool call metadata"""
if optional_tools:
for idx, tool in enumerate(optional_tools):
@ -408,7 +408,7 @@ def _set_tool_attributes(span: "Span", optional_tools: Optional[list], metadata_
)
def set_attributes(span: "Span", kwargs, response_obj, attributes: Type[BaseLLMObsOTELAttributes]):
def set_attributes(span: "Span", kwargs, response_obj, attributes: type[BaseLLMObsOTELAttributes]):
"""
Populates span with OpenInference-compliant LLM attributes for Arize and Phoenix tracing.
"""
@ -427,7 +427,7 @@ def set_attributes(span: "Span", kwargs, response_obj, attributes: Type[BaseLLMO
try:
optional_params = _sanitize_optional_params(kwargs.get("optional_params"))
litellm_params = kwargs.get("litellm_params", {}) or {}
standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object")
standard_logging_payload: StandardLoggingPayload | None = kwargs.get("standard_logging_object")
if standard_logging_payload is None:
raise ValueError("standard_logging_object not found in kwargs")
@ -482,19 +482,19 @@ def set_attributes(span: "Span", kwargs, response_obj, attributes: Type[BaseLLMO
)
def _sanitize_optional_params(optional_params: Optional[dict]) -> dict:
def _sanitize_optional_params(optional_params: dict | None) -> dict:
if not isinstance(optional_params, dict):
return {}
optional_params.pop("secret_fields", None)
return optional_params
def _set_metadata_attributes(span: "Span", metadata: Optional[Any], span_attrs) -> None:
def _set_metadata_attributes(span: "Span", metadata: Any | None, span_attrs) -> None:
if metadata is not None:
safe_set_attribute(span, span_attrs.METADATA, safe_dumps(metadata))
def _extract_metadata_tools(metadata: Optional[Any]) -> Optional[list]:
def _extract_metadata_tools(metadata: Any | None) -> list | None:
if not isinstance(metadata, dict):
return None
llm_obj = metadata.get("llm")
@ -503,7 +503,7 @@ def _extract_metadata_tools(metadata: Optional[Any]) -> Optional[list]:
return None
def _extract_optional_tools(optional_params: dict) -> Optional[list]:
def _extract_optional_tools(optional_params: dict) -> list | None:
return optional_params.get("tools") if isinstance(optional_params, dict) else None
@ -544,7 +544,7 @@ def _set_request_attributes(
safe_set_attribute(span, "llm.response.model", response_obj.get("model"))
def _set_model_params(span: "Span", model_params: Optional[dict], span_attrs) -> None:
def _set_model_params(span: "Span", model_params: dict | None, span_attrs) -> None:
if not model_params:
return
@ -606,7 +606,7 @@ def _coerce_response_obj_for_attrs(response_obj):
return response_obj
def _coerce_text(value) -> Optional[str]:
def _coerce_text(value) -> str | None:
"""Best-effort text extraction from a message-content value.
Returns None when no textual portion can be derived. Handles:
@ -650,7 +650,7 @@ def _to_plain_dict(value):
return value
def _get_tool_calls(message) -> Optional[list]:
def _get_tool_calls(message) -> list | None:
"""Return ``message.tool_calls`` only when it's a non-empty list.
Works for dicts and Pydantic message objects via ``_safe_get``.
@ -659,7 +659,7 @@ def _get_tool_calls(message) -> Optional[list]:
return tool_calls if isinstance(tool_calls, list) and tool_calls else None
def _normalize_tool_call(raw_tc) -> Optional[Dict[str, Any]]:
def _normalize_tool_call(raw_tc) -> dict[str, Any] | None:
"""Normalize a single tool_call (dict or Pydantic) into a stable shape:
{"id": str|None, "type": str, "function": {"name": str|None, "arguments": str|None}}
@ -879,7 +879,7 @@ def _set_response_cost_attr(span: "Span", standard_logging_payload) -> None:
safe_set_attribute(span, "llm.response.cost", cost_value)
def _is_passthrough_call_type(call_type: Optional[str]) -> bool:
def _is_passthrough_call_type(call_type: str | None) -> bool:
if not call_type:
return False
lowered = str(call_type).lower()

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