diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 43df27ea2e2..aa110d421e3 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,9 +1,9 @@ { "reportAny": { - "limit": 31256 + "limit": 29310 }, "reportArgumentType": { - "limit": 2645 + "limit": 2633 }, "reportAssignmentType": { "limit": 329 @@ -18,13 +18,13 @@ "limit": 59 }, "reportDeprecated": { - "limit": 325 + "limit": 322 }, "reportDuplicateImport": { - "limit": 42 + "limit": 39 }, "reportExplicitAny": { - "limit": 10208 + "limit": 8957 }, "reportFunctionMemberAccess": { "limit": 11 @@ -54,10 +54,10 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5869 + "limit": 5851 }, "reportMissingTypeArgument": { - "limit": 15861 + "limit": 15840 }, "reportMissingTypeStubs": { "limit": 41 @@ -99,19 +99,19 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 45357 + "limit": 44980 }, "reportUnknownLambdaType": { "limit": 113 }, "reportUnknownMemberType": { - "limit": 40477 + "limit": 40408 }, "reportUnknownParameterType": { - "limit": 20338 + "limit": 20308 }, "reportUnknownVariableType": { - "limit": 32047 + "limit": 31987 }, "reportUnnecessaryCast": { "limit": 177 @@ -138,7 +138,7 @@ "limit": 204 }, "reportUnusedImport": { - "limit": 1003 + "limit": 1000 }, "reportUnusedVariable": { "limit": 1297 diff --git a/litellm/_lazy_imports.py b/litellm/_lazy_imports.py index b04fae86e47..529287b0ca0 100644 --- a/litellm/_lazy_imports.py +++ b/litellm/_lazy_imports.py @@ -17,39 +17,43 @@ until they're actually needed. import importlib import sys -from typing import Any, Optional, cast, Callable +from collections.abc import Callable +from typing import TYPE_CHECKING, Any, Optional, cast + +if TYPE_CHECKING: + from tiktoken import Encoding # 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 from ._lazy_imports_registry import ( - # Name tuples - COST_CALCULATOR_NAMES, - LITELLM_LOGGING_NAMES, - UTILS_NAMES, - TOKEN_COUNTER_NAMES, - LLM_CLIENT_CACHE_NAMES, - BEDROCK_TYPES_NAMES, - TYPES_UTILS_NAMES, - CACHING_NAMES, - HTTP_HANDLER_NAMES, - DOTPROMPT_NAMES, - LLM_CONFIG_NAMES, - TYPES_NAMES, - LLM_PROVIDER_LOGIC_NAMES, - UTILS_MODULE_NAMES, - # Import maps - _UTILS_IMPORT_MAP, - _COST_CALCULATOR_IMPORT_MAP, - _TYPES_UTILS_IMPORT_MAP, - _TOKEN_COUNTER_IMPORT_MAP, _BEDROCK_TYPES_IMPORT_MAP, _CACHING_IMPORT_MAP, - _LITELLM_LOGGING_IMPORT_MAP, + _COST_CALCULATOR_IMPORT_MAP, _DOTPROMPT_IMPORT_MAP, - _TYPES_IMPORT_MAP, + _LITELLM_LOGGING_IMPORT_MAP, _LLM_CONFIGS_IMPORT_MAP, _LLM_PROVIDER_LOGIC_IMPORT_MAP, + _TOKEN_COUNTER_IMPORT_MAP, + _TYPES_IMPORT_MAP, + _TYPES_UTILS_IMPORT_MAP, + # Import maps + _UTILS_IMPORT_MAP, _UTILS_MODULE_IMPORT_MAP, + BEDROCK_TYPES_NAMES, + CACHING_NAMES, + # Name tuples + COST_CALCULATOR_NAMES, + DOTPROMPT_NAMES, + HTTP_HANDLER_NAMES, + LITELLM_LOGGING_NAMES, + LLM_CLIENT_CACHE_NAMES, + LLM_CONFIG_NAMES, + LLM_PROVIDER_LOGIC_NAMES, + TOKEN_COUNTER_NAMES, + TYPES_NAMES, + TYPES_UTILS_NAMES, + UTILS_MODULE_NAMES, + UTILS_NAMES, ) @@ -77,10 +81,10 @@ 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: Optional["Encoding"] = None -def _get_default_encoding() -> Any: +def _get_default_encoding() -> "Encoding": """ Lazily load and cache the default OpenAI encoding. @@ -99,7 +103,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: @@ -123,7 +127,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: @@ -153,7 +157,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]]: diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index c2dc7189934..70f5ea48919 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -10,18 +10,14 @@ import subprocess import sys import time import traceback +from collections.abc import Callable from datetime import datetime as dt_object from functools import lru_cache from typing import ( TYPE_CHECKING, Any, - Callable, - Dict, - List, Literal, Optional, - Tuple, - Type, Union, cast, ) @@ -37,11 +33,6 @@ from litellm import ( turn_off_message_logging, ) from litellm._logging import _is_debugging_on, _redact_string, verbose_logger -from litellm.exceptions import ( - BudgetExceededError, - validate_rate_limit_category, - validate_rate_limit_type, -) from litellm._uuid import uuid from litellm.batches.batch_utils import _handle_completed_batch from litellm.caching.caching import DualCache, InMemoryCache @@ -56,6 +47,11 @@ from litellm.cost_calculator import ( RealtimeAPITokenUsageProcessor, _select_model_name_for_cost_calc, ) +from litellm.exceptions import ( + BudgetExceededError, + validate_rate_limit_category, + validate_rate_limit_type, +) from litellm.integrations.agentops import AgentOps from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook from litellm.integrations.arize.arize import ArizeLogger @@ -176,7 +172,9 @@ from .initialize_dynamic_callback_params import ( from .specialty_caches.dynamic_logging_cache import DynamicLoggingCache if TYPE_CHECKING: + from litellm.integrations.otel.logger import OpenTelemetryV2 from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig + from litellm.router import Router try: from litellm_enterprise.enterprise_callbacks.callback_controls import ( EnterpriseCallbackControls, @@ -199,7 +197,7 @@ try: from litellm.integrations.generic_api.generic_api_callback import GenericAPILogger - EnterpriseStandardLoggingPayloadSetupVAR: Optional[Type[EnterpriseStandardLoggingPayloadSetup]] = ( + EnterpriseStandardLoggingPayloadSetupVAR: type[EnterpriseStandardLoggingPayloadSetup] | None = ( EnterpriseStandardLoggingPayloadSetup ) except Exception as e: @@ -211,7 +209,16 @@ except Exception as e: PagerDutyAlerting = CustomLogger # type: ignore EnterpriseCallbackControls = None # type: ignore EnterpriseStandardLoggingPayloadSetupVAR = None -_in_memory_loggers: List[Any] = [] +_in_memory_loggers: list[CustomLogger] = [] + + +def _erase_static_type(value: object) -> object: + return value + + +def _as_custom_logger(value: Any) -> CustomLogger: # any-ok: unresolvable enterprise import fallback + return value + _STANDARD_LOGGING_METADATA_KEYS: frozenset = frozenset(StandardLoggingMetadata.__annotations__.keys()) @@ -242,10 +249,10 @@ greenscaleLogger = None lunaryLogger = None supabaseClient = None deepevalLogger = None -callback_list: Optional[List[str]] = [] +callback_list: list[str] | None = [] user_logger_fn = None -additional_details: Optional[Dict[str, str]] = {} -local_cache: Optional[Dict[str, str]] = {} +additional_details: dict[str, str] | None = {} +local_cache: dict[str, str] | None = {} last_fetched_at = None last_fetched_at_keys = None @@ -255,7 +262,7 @@ class ServiceTraceIDCache: def __init__(self) -> None: self.cache = InMemoryCache() - def get_cache(self, litellm_call_id: str, service_name: str) -> Optional[str]: + def get_cache(self, litellm_call_id: str, service_name: str) -> str | None: key_name = "{}:{}".format(service_name, litellm_call_id) response = self.cache.get_cache(key=key_name) return response @@ -313,17 +320,17 @@ class Logging(LiteLLMLoggingBaseClass): start_time, litellm_call_id: str, function_id: str, - litellm_trace_id: Optional[str] = None, - dynamic_input_callbacks: Optional[List[Union[str, Callable, CustomLogger]]] = None, - dynamic_success_callbacks: Optional[List[Union[str, Callable, CustomLogger]]] = None, - dynamic_async_success_callbacks: Optional[List[Union[str, Callable, CustomLogger]]] = None, - dynamic_failure_callbacks: Optional[List[Union[str, Callable, CustomLogger]]] = None, - dynamic_async_failure_callbacks: Optional[List[Union[str, Callable, CustomLogger]]] = None, - applied_guardrails: Optional[List[str]] = None, - kwargs: Optional[Dict] = None, + litellm_trace_id: str | None = None, + dynamic_input_callbacks: list[str | Callable | CustomLogger] | None = None, + dynamic_success_callbacks: list[str | Callable | CustomLogger] | None = None, + dynamic_async_success_callbacks: list[str | Callable | CustomLogger] | None = None, + dynamic_failure_callbacks: list[str | Callable | CustomLogger] | None = None, + dynamic_async_failure_callbacks: list[str | Callable | CustomLogger] | None = None, + applied_guardrails: list[str] | None = None, + kwargs: dict | None = None, log_raw_request_response: bool = False, ): - _input: Optional[str] = messages # save original value of messages + _input: str | None = messages # save original value of messages if messages is not None: if isinstance(messages, str): messages = [ @@ -348,18 +355,18 @@ class Logging(LiteLLMLoggingBaseClass): self.litellm_call_id = litellm_call_id self.litellm_trace_id: str = litellm_trace_id if litellm_trace_id else str(uuid.uuid4()) self.function_id = function_id - self.streaming_chunks: List[Any] = [] # for generating complete stream response - self.sync_streaming_chunks: List[Any] = [] # for generating complete stream response + self.streaming_chunks: list[Any] = [] # for generating complete stream response + self.sync_streaming_chunks: list[Any] = [] # for generating complete stream response self.log_raw_request_response = log_raw_request_response # Initialize dynamic callbacks - self.dynamic_input_callbacks: Optional[List[Union[str, Callable, CustomLogger]]] = dynamic_input_callbacks - self.dynamic_success_callbacks: Optional[List[Union[str, Callable, CustomLogger]]] = dynamic_success_callbacks - self.dynamic_async_success_callbacks: Optional[List[Union[str, Callable, CustomLogger]]] = ( + self.dynamic_input_callbacks: list[str | Callable | CustomLogger] | None = dynamic_input_callbacks + self.dynamic_success_callbacks: list[str | Callable | CustomLogger] | None = dynamic_success_callbacks + self.dynamic_async_success_callbacks: list[str | Callable | CustomLogger] | None = ( dynamic_async_success_callbacks ) - self.dynamic_failure_callbacks: Optional[List[Union[str, Callable, CustomLogger]]] = dynamic_failure_callbacks - self.dynamic_async_failure_callbacks: Optional[List[Union[str, Callable, CustomLogger]]] = ( + self.dynamic_failure_callbacks: list[str | Callable | CustomLogger] | None = dynamic_failure_callbacks + self.dynamic_async_failure_callbacks: list[str | Callable | CustomLogger] | None = ( dynamic_async_failure_callbacks ) @@ -375,8 +382,8 @@ class Logging(LiteLLMLoggingBaseClass): self.initialize_standard_built_in_tools_params(kwargs) ) ## TIME TO FIRST TOKEN LOGGING ## - self.completion_start_time: Optional[datetime.datetime] = None - self._llm_caching_handler: Optional[LLMCachingHandler] = None + self.completion_start_time: datetime.datetime | None = None + self._llm_caching_handler: LLMCachingHandler | None = None # INITIAL LITELLM_PARAMS litellm_params = {} @@ -387,15 +394,15 @@ class Logging(LiteLLMLoggingBaseClass): self.litellm_params = litellm_params # Initialize cost breakdown field - self.cost_breakdown: Optional[CostBreakdown] = None + self.cost_breakdown: CostBreakdown | None = None # Init Caching related details - self.caching_details: Optional[CachingDetails] = None + self.caching_details: CachingDetails | None = None # Passthrough endpoint guardrails config for field targeting - self.passthrough_guardrails_config: Optional[Dict[str, Any]] = None + self.passthrough_guardrails_config: dict[str, Any] | None = None - self.model_call_details: Dict[str, Any] = { + self.model_call_details: dict[str, Any] = { "litellm_trace_id": self.litellm_trace_id, "litellm_call_id": litellm_call_id, "input": _input, @@ -408,7 +415,7 @@ class Logging(LiteLLMLoggingBaseClass): # post_call guardrails have run; the @client decorator then stores the # enqueue closure here instead of firing it immediately. self._defer_async_logging: bool = False - self._enqueue_deferred_logging: Optional[Callable[[], None]] = None + self._enqueue_deferred_logging: Callable[[], None] | None = None def process_dynamic_callbacks(self): """ @@ -443,9 +450,9 @@ class Logging(LiteLLMLoggingBaseClass): def _process_dynamic_callback_list( self, - callback_list: Optional[List[Union[str, Callable, CustomLogger]]], + callback_list: list[str | Callable | CustomLogger] | None, dynamic_callbacks_type: Literal["input", "success", "failure", "async_success", "async_failure"], - ) -> Optional[List[Union[str, Callable, CustomLogger]]]: + ) -> list[str | Callable | CustomLogger] | None: """ Helper function to initialize CustomLogger compatible callbacks in self.dynamic_* callbacks @@ -457,12 +464,12 @@ class Logging(LiteLLMLoggingBaseClass): if callback_list is None: return None - processed_list: List[Union[str, Callable, CustomLogger]] = [] + processed_list: list[str | Callable | CustomLogger] = [] for callback in callback_list: if isinstance(callback, str) and callback in litellm._known_custom_logger_compatible_callbacks: # For callbacks that support team-scoped credentials (e.g. datadog), # pass only the relevant dynamic params as custom_logger_init_args. - _custom_logger_init_args: Optional[dict] = None + _custom_logger_init_args: dict | None = None if callback == "datadog": _custom_logger_init_args = { k: v for k, v in self.standard_callback_dynamic_params.items() if k.startswith("dd_") @@ -490,9 +497,7 @@ class Logging(LiteLLMLoggingBaseClass): processed_list.append(callback) return processed_list - def initialize_standard_callback_dynamic_params( - self, kwargs: Optional[Dict] = None - ) -> StandardCallbackDynamicParams: + def initialize_standard_callback_dynamic_params(self, kwargs: dict | None = None) -> StandardCallbackDynamicParams: """ Initialize the standard callback dynamic params from the kwargs @@ -501,7 +506,7 @@ class Logging(LiteLLMLoggingBaseClass): return _initialize_standard_callback_dynamic_params(kwargs) - def initialize_standard_built_in_tools_params(self, kwargs: Optional[Dict] = None) -> StandardBuiltInToolsParams: + def initialize_standard_built_in_tools_params(self, kwargs: dict | None = None) -> StandardBuiltInToolsParams: """ Initialize the standard built-in tools params from the kwargs @@ -512,7 +517,7 @@ class Logging(LiteLLMLoggingBaseClass): file_search=StandardBuiltInToolCostTracking._get_file_search_tool_call(kwargs or {}), ) - def get_router_model_id(self) -> Optional[str]: + def get_router_model_id(self) -> str | None: """Extract the router deployment model_id from litellm_params. Checks both litellm_metadata and metadata for model_info.id. @@ -531,10 +536,10 @@ class Logging(LiteLLMLoggingBaseClass): def update_environment_variables( self, - litellm_params: Dict, - optional_params: Dict, - model: Optional[str] = None, - user: Optional[str] = None, + litellm_params: dict, + optional_params: dict, + model: str | None = None, + user: str | None = None, **additional_params, ): self.optional_params = optional_params @@ -580,11 +585,11 @@ class Logging(LiteLLMLoggingBaseClass): def update_from_kwargs( self, - kwargs: Dict, - litellm_params: Optional[Dict] = None, - optional_params: Optional[Dict] = None, - model: Optional[str] = None, - user: Optional[str] = None, + kwargs: dict, + litellm_params: dict | None = None, + optional_params: dict | None = None, + model: str | None = None, + user: str | None = None, **additional_params, ): """ @@ -592,7 +597,7 @@ class Logging(LiteLLMLoggingBaseClass): automatically extracts metadata/litellm_metadata from kwargs, so callers don't need to manually plumb them into litellm_params. """ - base_litellm_params: Dict[str, Any] = {} + base_litellm_params: dict[str, Any] = {} if "metadata" in kwargs: base_litellm_params["metadata"] = kwargs["metadata"] @@ -622,7 +627,7 @@ class Logging(LiteLLMLoggingBaseClass): **additional_params, ) - def update_messages(self, messages: List[AllMessageValues]): + def update_messages(self, messages: list[AllMessageValues]): """ Update the logged value of the messages in the model_call_details @@ -633,9 +638,9 @@ class Logging(LiteLLMLoggingBaseClass): def should_run_prompt_management_hooks( self, - non_default_params: Dict, - prompt_id: Optional[str] = None, - tools: Optional[List[Dict]] = None, + non_default_params: dict, + prompt_id: str | None = None, + tools: list[dict] | None = None, ) -> bool: """ Return True if prompt management hooks should be run @@ -658,8 +663,8 @@ class Logging(LiteLLMLoggingBaseClass): def _should_run_prompt_management_hooks_without_prompt_id( self, - non_default_params: Dict, - tools: Optional[List[Dict]] = None, + non_default_params: dict, + tools: list[dict] | None = None, ) -> bool: """ Certain prompt management hooks don't need a `prompt_id` to be passed in, they are triggered by dynamic params @@ -683,15 +688,15 @@ class Logging(LiteLLMLoggingBaseClass): def get_chat_completion_prompt( self, model: str, - messages: List[AllMessageValues], - non_default_params: Dict, - prompt_variables: Optional[dict], - prompt_id: Optional[str] = None, - prompt_spec: Optional[PromptSpec] = None, - prompt_management_logger: Optional[CustomLogger] = None, - prompt_label: Optional[str] = None, - prompt_version: Optional[int] = None, - ) -> Tuple[str, List[AllMessageValues], dict]: + messages: list[AllMessageValues], + non_default_params: dict, + prompt_variables: dict | None, + prompt_id: str | None = None, + prompt_spec: PromptSpec | None = None, + prompt_management_logger: CustomLogger | None = None, + prompt_label: str | None = None, + prompt_version: int | None = None, + ) -> tuple[str, list[AllMessageValues], dict]: custom_logger = prompt_management_logger or self.get_custom_logger_for_prompt_management( model=model, non_default_params=non_default_params, @@ -722,16 +727,16 @@ class Logging(LiteLLMLoggingBaseClass): async def async_get_chat_completion_prompt( self, model: str, - messages: List[AllMessageValues], - non_default_params: Dict, - prompt_variables: Optional[dict], - prompt_id: Optional[str] = None, - prompt_spec: Optional[PromptSpec] = None, - prompt_management_logger: Optional[CustomLogger] = None, - tools: Optional[List[Dict]] = None, - prompt_label: Optional[str] = None, - prompt_version: Optional[int] = None, - ) -> Tuple[str, List[AllMessageValues], dict]: + messages: list[AllMessageValues], + non_default_params: dict, + prompt_variables: dict | None, + prompt_id: str | None = None, + prompt_spec: PromptSpec | None = None, + prompt_management_logger: CustomLogger | None = None, + tools: list[dict] | None = None, + prompt_label: str | None = None, + prompt_version: int | None = None, + ) -> tuple[str, list[AllMessageValues], dict]: custom_logger = prompt_management_logger or self.get_custom_logger_for_prompt_management( model=model, tools=tools, @@ -765,9 +770,9 @@ class Logging(LiteLLMLoggingBaseClass): def _auto_detect_prompt_management_logger( self, prompt_id: str, - prompt_spec: Optional[PromptSpec], + prompt_spec: PromptSpec | None, dynamic_callback_params: StandardCallbackDynamicParams, - ) -> Optional[CustomLogger]: + ) -> CustomLogger | None: """ Auto-detect which prompt management system owns the given prompt_id. @@ -803,12 +808,12 @@ class Logging(LiteLLMLoggingBaseClass): def get_custom_logger_for_prompt_management( self, model: str, - non_default_params: Dict, - tools: Optional[List[Dict]] = None, - prompt_id: Optional[str] = None, - prompt_spec: Optional[PromptSpec] = None, - dynamic_callback_params: Optional[StandardCallbackDynamicParams] = None, - ) -> Optional[CustomLogger]: + non_default_params: dict, + tools: list[dict] | None = None, + prompt_id: str | None = None, + prompt_spec: PromptSpec | None = None, + dynamic_callback_params: StandardCallbackDynamicParams | None = None, + ) -> CustomLogger | None: """ Get a custom logger for prompt management based on model name or available callbacks. @@ -879,7 +884,7 @@ class Logging(LiteLLMLoggingBaseClass): return None - def get_custom_logger_for_anthropic_cache_control_hook(self, non_default_params: Dict) -> Optional[CustomLogger]: + def get_custom_logger_for_anthropic_cache_control_hook(self, non_default_params: dict) -> CustomLogger | None: if non_default_params.get("cache_control_injection_points", None): custom_logger = _init_custom_logger_compatible_class( logging_integration="anthropic_cache_control_hook", @@ -889,7 +894,7 @@ class Logging(LiteLLMLoggingBaseClass): return custom_logger return None - def _get_raw_request_body(self, data: Optional[Union[dict, str]]) -> dict: + def _get_raw_request_body(self, data: dict | str | None) -> dict: if data is None: return {"error": "Received empty dictionary for raw request body"} if isinstance(data, str): @@ -1103,9 +1108,7 @@ class Logging(LiteLLMLoggingBaseClass): def _get_request_body(self, data: dict) -> str: return str(data) - def _get_request_curl_command( - self, api_base: str, headers: Optional[dict], additional_args: dict, data: dict - ) -> str: + def _get_request_curl_command(self, api_base: str, headers: dict | None, additional_args: dict, data: dict) -> str: masked_api_base = self._get_masked_api_base(api_base) if headers is None: headers = {} @@ -1246,7 +1249,7 @@ class Logging(LiteLLMLoggingBaseClass): for callback in callbacks: try: if isinstance(callback, CustomLogger): - response: Optional[MCPPostCallResponseObject] = await callback.async_post_mcp_tool_call_hook( + response: MCPPostCallResponseObject | None = await callback.async_post_mcp_tool_call_hook( kwargs=kwargs, response_obj=post_mcp_tool_call_response_obj, start_time=start_time, @@ -1264,7 +1267,7 @@ class Logging(LiteLLMLoggingBaseClass): ) return response_obj - def _parse_post_mcp_call_hook_response(self, response: Optional[MCPPostCallResponseObject]) -> Any: + def _parse_post_mcp_call_hook_response(self, response: MCPPostCallResponseObject | None) -> Any: """ Parse the response from the post_mcp_tool_call_hook @@ -1288,16 +1291,16 @@ class Logging(LiteLLMLoggingBaseClass): output_cost: float, total_cost: float, cost_for_built_in_tools_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 method to store cost breakdown in the logging object. @@ -1371,10 +1374,10 @@ class Logging(LiteLLMLoggingBaseClass): dict, list, ], - cache_hit: Optional[bool] = None, - litellm_model_name: Optional[str] = None, - router_model_id: Optional[str] = None, - ) -> Optional[float]: + cache_hit: bool | None = None, + litellm_model_name: str | None = None, + router_model_id: str | None = None, + ) -> float | None: """ Calculate response cost using result + logging object variables. @@ -1473,7 +1476,7 @@ class Logging(LiteLLMLoggingBaseClass): return None - def _generate_content_result_as_model_response(self, result: object) -> Optional[ModelResponse]: + def _generate_content_result_as_model_response(self, result: object) -> ModelResponse | None: """ Native Google :generateContent bodies report token usage under ``usageMetadata``, which the cost calculator does not read, so a raw body @@ -1508,20 +1511,18 @@ class Logging(LiteLLMLoggingBaseClass): async def _response_cost_calculator_async( self, - result: Union[ - ModelResponse, - ModelResponseStream, - EmbeddingResponse, - ImageResponse, - TranscriptionResponse, - TextCompletionResponse, - HttpxBinaryResponseContent, - RerankResponse, - Batch, - FineTuningJob, - ], - cache_hit: Optional[bool] = None, - ) -> Optional[float]: + result: ModelResponse + | ModelResponseStream + | EmbeddingResponse + | ImageResponse + | TranscriptionResponse + | TextCompletionResponse + | HttpxBinaryResponseContent + | RerankResponse + | Batch + | FineTuningJob, + cache_hit: bool | None = None, + ) -> float | None: return self._response_cost_calculator(result=result, cache_hit=cache_hit) @staticmethod @@ -1839,7 +1840,7 @@ class Logging(LiteLLMLoggingBaseClass): start_time=None, end_time=None, cache_hit=None, - standard_logging_object: Optional[StandardLoggingPayload] = None, + standard_logging_object: StandardLoggingPayload | None = None, ): try: if start_time is None: @@ -1948,7 +1949,7 @@ class Logging(LiteLLMLoggingBaseClass): def _flush_passthrough_collected_chunks_helper( self, - raw_bytes: List[bytes], + raw_bytes: list[bytes], provider_config: "BasePassthroughConfig", ) -> Optional["CostResponseTypes"]: all_chunks = provider_config._convert_raw_bytes_to_str_lines(raw_bytes) @@ -1963,7 +1964,7 @@ class Logging(LiteLLMLoggingBaseClass): def flush_passthrough_collected_chunks( self, - raw_bytes: List[bytes], + raw_bytes: list[bytes], provider_config: "BasePassthroughConfig", ): """ @@ -1986,7 +1987,7 @@ class Logging(LiteLLMLoggingBaseClass): async def async_flush_passthrough_collected_chunks( self, - raw_bytes: List[bytes], + raw_bytes: list[bytes], provider_config: "BasePassthroughConfig", ): complete_streaming_response = self._flush_passthrough_collected_chunks_helper( @@ -2013,9 +2014,7 @@ class Logging(LiteLLMLoggingBaseClass): is_sync_request = self._is_sync_litellm_request(litellm_params) try: ## BUILD COMPLETE STREAMED RESPONSE - complete_streaming_response: Optional[ - Union[ModelResponse, TextCompletionResponse, ResponsesAPIResponse] - ] = None + complete_streaming_response: ModelResponse | TextCompletionResponse | ResponsesAPIResponse | None = None if "complete_streaming_response" in self.model_call_details: return # break out of this. complete_streaming_response = self._get_assembled_streaming_response( @@ -2469,7 +2468,7 @@ class Logging(LiteLLMLoggingBaseClass): ## BUILD COMPLETE STREAMED RESPONSE if "async_complete_streaming_response" in self.model_call_details: return # break out of this. - complete_streaming_response: Optional[Union[ModelResponse, TextCompletionResponse, ResponsesAPIResponse]] = ( + complete_streaming_response: ModelResponse | TextCompletionResponse | ResponsesAPIResponse | None = ( self._get_assembled_streaming_response( result=result, start_time=start_time, @@ -2612,7 +2611,7 @@ class Logging(LiteLLMLoggingBaseClass): ) if isinstance(callback, CustomLogger): # custom logger class - model_call_details: Dict = self.model_call_details + model_call_details: dict = self.model_call_details ################################## # call redaction hook for custom logger model_call_details = callback.redact_standard_logging_payload_from_model_call_details( @@ -3024,7 +3023,7 @@ class Logging(LiteLLMLoggingBaseClass): # Track callback logging failures in Prometheus self._handle_callback_failure(callback=callback) - def _get_trace_id(self, service_name: Literal["langfuse"]) -> Optional[str]: + def _get_trace_id(self, service_name: Literal["langfuse"]) -> str | None: """ For the given service (e.g. langfuse), return the trace_id actually logged. @@ -3034,7 +3033,7 @@ class Logging(LiteLLMLoggingBaseClass): - str: The logged trace id - None: If trace id not yet emitted. """ - trace_id: Optional[str] = None + trace_id: str | None = None if service_name == "langfuse": trace_id = in_memory_trace_id_cache.get_cache( litellm_call_id=self.litellm_call_id, service_name=service_name @@ -3042,7 +3041,7 @@ class Logging(LiteLLMLoggingBaseClass): return trace_id - def _get_callback_object(self, service_name: Literal["langfuse"]) -> Optional[Any]: + def _get_callback_object(self, service_name: Literal["langfuse"]) -> LangFuseLogger | None: """ Return dynamic callback object. @@ -3081,7 +3080,7 @@ class Logging(LiteLLMLoggingBaseClass): result: Any, start_time: datetime.datetime, end_time: datetime.datetime, - cache_hit: Optional[Any] = None, + cache_hit: bool | None = None, ) -> None: """ Handles calling success callbacks for Async calls. @@ -3130,12 +3129,12 @@ class Logging(LiteLLMLoggingBaseClass): _filtered_failure_callbacks = self._remove_internal_litellm_callbacks(_filtered_failure_callbacks) return len(_filtered_failure_callbacks) > 0 - def get_combined_callback_list(self, dynamic_success_callbacks: Optional[List], global_callbacks: List) -> List: + def get_combined_callback_list(self, dynamic_success_callbacks: list | None, global_callbacks: list) -> list: if dynamic_success_callbacks is None: return list(global_callbacks) return list(dict.fromkeys(dynamic_success_callbacks + global_callbacks)) - def _remove_internal_litellm_callbacks(self, callbacks: List) -> List: + def _remove_internal_litellm_callbacks(self, callbacks: list) -> list: """ Creates a filtered list of callbacks, excluding internal LiteLLM callbacks. @@ -3186,7 +3185,7 @@ class Logging(LiteLLMLoggingBaseClass): cb_name = self._get_callback_name(cb) return any(prefix in cb_name for prefix in INTERNAL_PREFIXES) - def _remove_internal_custom_logger_callbacks(self, callbacks: List) -> List: + def _remove_internal_custom_logger_callbacks(self, callbacks: list) -> list: """ Removes internal custom logger callbacks from the list. """ @@ -3201,18 +3200,12 @@ class Logging(LiteLLMLoggingBaseClass): def _get_assembled_streaming_response( self, - result: Union[ - ModelResponse, - TextCompletionResponse, - ModelResponseStream, - ResponseCompletedEvent, - Any, - ], + result: ModelResponse | TextCompletionResponse | ModelResponseStream | ResponseCompletedEvent | Any, start_time: datetime.datetime, end_time: datetime.datetime, is_async: bool, - streaming_chunks: List[Any], - ) -> Optional[Union[ModelResponse, TextCompletionResponse, ResponsesAPIResponse]]: + streaming_chunks: list[object], + ) -> ModelResponse | TextCompletionResponse | ResponsesAPIResponse | None: if self.stream is not True: return None if isinstance(result, ModelResponse): @@ -3396,7 +3389,7 @@ def _get_masked_values( ignore_sensitive_values: bool = False, mask_all_values: bool = False, unmasked_length: int = 4, - number_of_asterisks: Optional[int] = 4, + number_of_asterisks: int | None = 4, _depth: int = 0, _max_depth: int = 20, ) -> dict: @@ -3563,10 +3556,10 @@ def set_callbacks(callback_list, function_id=None): def _init_custom_logger_compatible_class( logging_integration: _custom_logger_compatible_callbacks_literal, - internal_usage_cache: Optional[DualCache], - llm_router: Optional[Any], # expect litellm.Router, but typing errors due to circular import - custom_logger_init_args: Optional[dict] = {}, -) -> Optional[CustomLogger]: + internal_usage_cache: DualCache | None, + llm_router: Optional["Router"], + custom_logger_init_args: dict | None = {}, +) -> CustomLogger | None: """ Initialize a custom logger compatible class """ @@ -3954,7 +3947,7 @@ def _init_custom_logger_compatible_class( dynamic_rate_limiter_obj = _PROXY_DynamicRateLimitHandler(internal_usage_cache=internal_usage_cache) - if llm_router is not None and isinstance(llm_router, litellm.Router): + if llm_router is not None: dynamic_rate_limiter_obj.update_variables(llm_router=llm_router) _in_memory_loggers.append(dynamic_rate_limiter_obj) return dynamic_rate_limiter_obj # type: ignore @@ -3974,7 +3967,7 @@ def _init_custom_logger_compatible_class( dynamic_rate_limiter_obj_v3 = _PROXY_DynamicRateLimitHandlerV3(internal_usage_cache=internal_usage_cache) - if llm_router is not None and isinstance(llm_router, litellm.Router): + if llm_router is not None: dynamic_rate_limiter_obj_v3.update_variables(llm_router=llm_router) _in_memory_loggers.append(dynamic_rate_limiter_obj_v3) return dynamic_rate_limiter_obj_v3 # type: ignore @@ -4061,7 +4054,7 @@ def _init_custom_logger_compatible_class( if isinstance(callback, PagerDutyAlerting): return callback pagerduty_logger = PagerDutyAlerting(**custom_logger_init_args) - _in_memory_loggers.append(pagerduty_logger) + _in_memory_loggers.append(_as_custom_logger(pagerduty_logger)) return pagerduty_logger # type: ignore elif logging_integration == "anthropic_cache_control_hook": for callback in _in_memory_loggers: @@ -4090,7 +4083,7 @@ def _init_custom_logger_compatible_class( return _gcs_pubsub_logger # type: ignore elif logging_integration == "generic_api": for callback in _in_memory_loggers: - if isinstance(callback, GenericAPILogger): + if isinstance(_erase_static_type(callback), GenericAPILogger): return callback generic_api_logger = GenericAPILogger() _in_memory_loggers.append(generic_api_logger) @@ -4100,21 +4093,21 @@ def _init_custom_logger_compatible_class( if isinstance(callback, ResendEmailLogger): return callback resend_email_logger = ResendEmailLogger() - _in_memory_loggers.append(resend_email_logger) + _in_memory_loggers.append(_as_custom_logger(resend_email_logger)) return resend_email_logger # type: ignore elif logging_integration == "sendgrid_email": for callback in _in_memory_loggers: if isinstance(callback, SendGridEmailLogger): return callback sendgrid_email_logger = SendGridEmailLogger() - _in_memory_loggers.append(sendgrid_email_logger) + _in_memory_loggers.append(_as_custom_logger(sendgrid_email_logger)) return sendgrid_email_logger # type: ignore elif logging_integration == "smtp_email": for callback in _in_memory_loggers: if isinstance(callback, SMTPEmailLogger): return callback smtp_email_logger = SMTPEmailLogger() - _in_memory_loggers.append(smtp_email_logger) + _in_memory_loggers.append(_as_custom_logger(smtp_email_logger)) return smtp_email_logger # type: ignore elif logging_integration == "humanloop": for callback in _in_memory_loggers: @@ -4142,7 +4135,7 @@ def _init_custom_logger_compatible_class( return callback # Get global BitBucket config - bitbucket_config = getattr(litellm, "global_bitbucket_config", None) + bitbucket_config = litellm.global_bitbucket_config if bitbucket_config is None: raise ValueError("BitBucket configuration not found. Please set litellm.global_bitbucket_config first.") @@ -4159,7 +4152,7 @@ def _init_custom_logger_compatible_class( return callback # Get global BitBucket config - gitlab_config = getattr(litellm, "global_gitlab_config", None) + gitlab_config = litellm.global_gitlab_config if gitlab_config is None: raise ValueError("Gitlab configuration not found. Please set litellm.global_gitlab_config first.") @@ -4180,7 +4173,7 @@ def _init_custom_logger_compatible_class( return None -def _maybe_construct_otel_v2(callback_name: str, _in_memory_loggers: list) -> Optional[Any]: +def _maybe_construct_otel_v2(callback_name: str, _in_memory_loggers: list[CustomLogger]) -> Optional["OpenTelemetryV2"]: """If ``LITELLM_OTEL_V2`` is on, build (or reuse) a single ``OpenTelemetryV2`` instance configured via the preset for ``callback_name``. @@ -4198,7 +4191,7 @@ def _maybe_construct_otel_v2(callback_name: str, _in_memory_loggers: list) -> Op if preset_fn is None: return None for callback in _in_memory_loggers: - if isinstance(callback, OpenTelemetryV2) and getattr(callback, "callback_name", None) == callback_name: + if isinstance(callback, OpenTelemetryV2) and callback.callback_name == callback_name: return callback try: config = preset_fn() @@ -4211,7 +4204,7 @@ def _maybe_construct_otel_v2(callback_name: str, _in_memory_loggers: list) -> Op return v2_logger -def _maybe_auto_initialize_arize_phoenix(_in_memory_loggers: list) -> None: +def _maybe_auto_initialize_arize_phoenix(_in_memory_loggers: list[CustomLogger]) -> None: """ Auto-initialize ArizePhoenixLogger when Phoenix env vars are detected. @@ -4256,7 +4249,7 @@ def _maybe_auto_initialize_arize_phoenix(_in_memory_loggers: list) -> None: def get_custom_logger_compatible_class( logging_integration: _custom_logger_compatible_callbacks_literal, -) -> Optional[CustomLogger]: +) -> CustomLogger | None: try: if logging_integration == "lago": for callback in _in_memory_loggers: @@ -4438,7 +4431,7 @@ def get_custom_logger_compatible_class( return callback elif logging_integration == "generic_api": for callback in _in_memory_loggers: - if isinstance(callback, GenericAPILogger): + if isinstance(_erase_static_type(callback), GenericAPILogger): return callback elif logging_integration == "resend_email": for callback in _in_memory_loggers: @@ -4463,7 +4456,7 @@ def get_custom_logger_compatible_class( return None -def _get_custom_logger_settings_from_proxy_server(callback_name: str) -> Dict: +def _get_custom_logger_settings_from_proxy_server(callback_name: str) -> dict: """ Get the settings for a custom logger from the proxy server config.yaml @@ -4478,7 +4471,7 @@ def _get_custom_logger_settings_from_proxy_server(callback_name: str) -> Dict: return {} -def use_custom_pricing_for_model(litellm_params: Optional[dict]) -> bool: +def use_custom_pricing_for_model(litellm_params: dict | None) -> bool: """ Check if the model uses custom pricing @@ -4516,10 +4509,10 @@ def is_valid_sha256_hash(value: str) -> bool: class StandardLoggingPayloadSetup: @staticmethod def cleanup_timestamps( - start_time: Union[dt_object, float], - end_time: Union[dt_object, float], - completion_start_time: Union[dt_object, float], - ) -> Tuple[float, float, float]: + start_time: dt_object | float, + end_time: dt_object | float, + completion_start_time: dt_object | float, + ) -> tuple[float, float, float]: """ Convert datetime objects to floats @@ -4556,7 +4549,7 @@ class StandardLoggingPayloadSetup: return start_time_float, end_time_float, completion_start_time_float @staticmethod - def append_system_prompt_messages(kwargs: Optional[Dict] = None, messages: Optional[Any] = None): + def append_system_prompt_messages(kwargs: dict | None = None, messages: Any | None = None): """ Append system prompt messages to the messages """ @@ -4614,16 +4607,16 @@ class StandardLoggingPayloadSetup: @staticmethod def get_standard_logging_metadata( - metadata: Optional[Dict[str, Any]], - litellm_params: Optional[dict] = None, - prompt_integration: Optional[str] = None, - applied_guardrails: Optional[List[str]] = None, - mcp_tool_call_metadata: Optional[StandardLoggingMCPToolCall] = None, - vector_store_request_metadata: Optional[List[StandardLoggingVectorStoreRequest]] = None, - usage_object: Optional[dict] = None, - proxy_server_request: Optional[dict] = None, - start_time: Optional[dt_object] = None, - response_id: Optional[str] = None, + metadata: dict[str, Any] | None, + litellm_params: dict | None = None, + prompt_integration: str | None = None, + applied_guardrails: list[str] | None = None, + mcp_tool_call_metadata: StandardLoggingMCPToolCall | None = None, + vector_store_request_metadata: list[StandardLoggingVectorStoreRequest] | None = None, + usage_object: dict | None = None, + proxy_server_request: dict | None = None, + start_time: dt_object | None = None, + response_id: str | None = None, ) -> StandardLoggingMetadata: """ Clean and filter the metadata dictionary to include only the specified keys in StandardLoggingMetadata. @@ -4639,10 +4632,10 @@ class StandardLoggingPayloadSetup: - If 'user_api_key' is present in metadata and is a valid SHA256 hash, it's stored as 'user_api_key_hash'. """ - prompt_management_metadata: Optional[StandardLoggingPromptManagementMetadata] = None + prompt_management_metadata: StandardLoggingPromptManagementMetadata | None = None if litellm_params is not None: - prompt_id = cast(Optional[str], litellm_params.get("prompt_id", None)) - prompt_variables = cast(Optional[dict], litellm_params.get("prompt_variables", None)) + prompt_id = cast(str | None, litellm_params.get("prompt_id", None)) + prompt_variables = cast(dict | None, litellm_params.get("prompt_variables", None)) if prompt_id is not None and prompt_integration is not None: prompt_management_metadata = StandardLoggingPromptManagementMetadata( @@ -4724,9 +4717,7 @@ class StandardLoggingPayloadSetup: return clean_metadata @staticmethod - def get_usage_from_response_obj( - response_obj: Optional[dict], combined_usage_object: Optional[Usage] = None - ) -> Usage: + def get_usage_from_response_obj(response_obj: dict | None, combined_usage_object: Usage | None = None) -> Usage: ## BASE CASE ## if combined_usage_object is not None: return combined_usage_object @@ -4757,8 +4748,8 @@ class StandardLoggingPayloadSetup: @staticmethod def get_usage_as_dict( - response_obj: Optional[dict], - combined_usage_object: Optional[Usage] = None, + response_obj: dict | None, + combined_usage_object: Usage | None = None, ) -> dict: """ Like get_usage_from_response_obj but returns a plain dict, skipping @@ -4784,11 +4775,11 @@ class StandardLoggingPayloadSetup: @staticmethod def get_model_cost_information( - base_model: Optional[str], - custom_pricing: Optional[bool], - custom_llm_provider: Optional[str], - init_response_obj: Union[Any, BaseModel, dict], - api_base: Optional[str] = None, + base_model: str | None, + custom_pricing: bool | None, + custom_llm_provider: str | None, + init_response_obj: Any | BaseModel | dict, + api_base: str | None = None, ) -> StandardLoggingModelInformation: model_cost_name = _select_model_name_for_cost_calc( model=base_model if custom_pricing else None, @@ -4822,13 +4813,13 @@ class StandardLoggingPayloadSetup: @staticmethod def get_final_response_obj( - response_obj: dict, init_response_obj: Union[Any, BaseModel, dict], kwargs: dict - ) -> Optional[Union[dict, str, list]]: + response_obj: dict, init_response_obj: Any | BaseModel | dict, kwargs: dict + ) -> dict | str | list | None: """ Get final response object after redacting the message input/output from logging """ if response_obj: - final_response_obj: Optional[Union[dict, str, list]] = response_obj + final_response_obj: dict | str | list | None = response_obj elif isinstance(init_response_obj, list) or isinstance(init_response_obj, str): final_response_obj = init_response_obj else: @@ -4848,8 +4839,8 @@ class StandardLoggingPayloadSetup: @staticmethod def get_additional_headers( - additiona_headers: Optional[dict], - ) -> Optional[StandardLoggingAdditionalHeaders]: + additiona_headers: dict | None, + ) -> StandardLoggingAdditionalHeaders | None: if additiona_headers is None: return None @@ -4875,7 +4866,7 @@ class StandardLoggingPayloadSetup: @staticmethod def get_hidden_params( - hidden_params: Optional[dict], + hidden_params: dict | None, ) -> StandardLoggingHiddenParams: clean_hidden_params = StandardLoggingHiddenParams( model_id=None, @@ -4900,7 +4891,7 @@ class StandardLoggingPayloadSetup: return clean_hidden_params @staticmethod - def strip_trailing_slash(api_base: Optional[str]) -> Optional[str]: + def strip_trailing_slash(api_base: str | None) -> str | None: if api_base: if api_base.endswith("//"): return api_base.rstrip("/") @@ -4912,8 +4903,8 @@ class StandardLoggingPayloadSetup: def _generate_cold_storage_object_key( start_time: dt_object, response_id: str, - team_alias: Optional[str] = None, - ) -> Optional[str]: + team_alias: str | None = None, + ) -> str | None: """ Generate cold storage object key in the same format as S3Logger. @@ -4965,8 +4956,8 @@ class StandardLoggingPayloadSetup: @staticmethod def get_error_information( - original_exception: Optional[Exception], - traceback_str: Optional[str] = None, + original_exception: Exception | None, + traceback_str: str | None = None, ) -> StandardLoggingPayloadErrorInformation: from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG @@ -5110,13 +5101,13 @@ class StandardLoggingPayloadSetup: return logging_obj.litellm_trace_id @staticmethod - def _get_user_agent_tags(proxy_server_request: dict) -> Optional[List[str]]: + def _get_user_agent_tags(proxy_server_request: dict) -> list[str] | None: """ Return the user agent tags from the proxy server request for spend tracking """ if litellm.disable_add_user_agent_to_request_tags is True: return None - user_agent_tags: Optional[List[str]] = None + user_agent_tags: list[str] | None = None headers = proxy_server_request.get("headers", {}) if headers is not None and isinstance(headers, dict): if "user-agent" in headers: @@ -5124,7 +5115,7 @@ class StandardLoggingPayloadSetup: if user_agent is not None: if user_agent_tags is None: user_agent_tags = [] - user_agent_part: Optional[str] = None + user_agent_part: str | None = None if "/" in user_agent: user_agent_part = user_agent.split("/")[0] if user_agent_part is not None: @@ -5134,11 +5125,11 @@ class StandardLoggingPayloadSetup: return user_agent_tags @staticmethod - def _get_extra_header_tags(proxy_server_request: dict) -> Optional[List[str]]: + def _get_extra_header_tags(proxy_server_request: dict) -> list[str] | None: """ Extract additional header tags for spend tracking based on config. """ - extra_headers: List[str] = getattr(litellm, "extra_spend_tag_headers", None) or [] + extra_headers: list[str] = getattr(litellm, "extra_spend_tag_headers", None) or [] if not extra_headers: return None @@ -5155,7 +5146,7 @@ class StandardLoggingPayloadSetup: return header_tags if header_tags else None @staticmethod - def _get_request_tags(litellm_params: dict, proxy_server_request: dict) -> List[str]: + def _get_request_tags(litellm_params: dict, proxy_server_request: dict) -> list[str]: # check for 'tags' in both 'metadata' and 'litellm_metadata' metadata = litellm_params.get("metadata") or {} litellm_metadata = litellm_params.get("litellm_metadata") or {} @@ -5176,8 +5167,8 @@ class StandardLoggingPayloadSetup: def _get_status_fields( status: StandardLoggingPayloadStatus, - guardrail_information: Optional[List[dict]], - error_str: Optional[str], + guardrail_information: list[dict] | None, + error_str: str | None, ) -> "StandardLoggingPayloadStatusFields": """ Determine status fields based on request status and guardrail information. @@ -5191,7 +5182,7 @@ def _get_status_fields( StandardLoggingPayloadStatusFields with llm_api_status and guardrail_status """ # Mapping for legacy guardrail status values to new GuardrailStatus values - GUARDRAIL_STATUS_MAP: Dict[str, GuardrailStatus] = { + GUARDRAIL_STATUS_MAP: dict[str, GuardrailStatus] = { "success": "success", "blocked": "guardrail_intervened", # legacy "guardrail_intervened": "guardrail_intervened", # direct @@ -5219,11 +5210,11 @@ def _get_status_fields( def _extract_response_obj_and_hidden_params( - init_response_obj: Union[Any, BaseModel, dict], - original_exception: Optional[Exception], -) -> Tuple[dict, Optional[dict]]: + init_response_obj: Any | BaseModel | dict, + original_exception: Exception | None, +) -> tuple[dict, dict | None]: """Extract response_obj and hidden_params from init_response_obj.""" - hidden_params: Optional[dict] = None + hidden_params: dict | None = None if init_response_obj is None: response_obj = {} elif isinstance(init_response_obj, BaseModel): @@ -5255,16 +5246,16 @@ def _extract_response_obj_and_hidden_params( def get_standard_logging_object_payload( - kwargs: Optional[dict], - init_response_obj: Union[Any, BaseModel, dict], + kwargs: dict | None, + init_response_obj: Any | BaseModel | dict, start_time: dt_object, end_time: dt_object, logging_obj: Logging, status: StandardLoggingPayloadStatus, - error_str: Optional[str] = None, - original_exception: Optional[Exception] = None, - standard_built_in_tools_params: Optional[StandardBuiltInToolsParams] = None, -) -> Optional[StandardLoggingPayload]: + error_str: str | None = None, + original_exception: Exception | None = None, + standard_built_in_tools_params: StandardBuiltInToolsParams | None = None, +) -> StandardLoggingPayload | None: try: kwargs = kwargs or {} @@ -5283,7 +5274,7 @@ def get_standard_logging_object_payload( # Extract usage as a plain dict, avoiding Pydantic round-trip raw_usage_dict = StandardLoggingPayloadSetup.get_usage_as_dict( response_obj=response_obj, - combined_usage_object=cast(Optional[Usage], kwargs.get("combined_usage_object")), + combined_usage_object=cast(Usage | None, kwargs.get("combined_usage_object")), ) usage_dict = ( {**raw_usage_dict, "output_image_count": len(init_response_obj.data)} @@ -5382,7 +5373,7 @@ def get_standard_logging_object_payload( kwargs=kwargs, ) - stream: Optional[bool] = None + stream: bool | None = None if ( kwargs.get("complete_streaming_response") is not None or kwargs.get("async_complete_streaming_response") is not None @@ -5392,9 +5383,9 @@ def get_standard_logging_object_payload( # Reconstruct full model name with provider prefix for logging # This ensures Bedrock models like "us.anthropic.claude-3-5-sonnet-20240620-v1:0" # are logged as "bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0" - custom_llm_provider = cast(Optional[str], kwargs.get("custom_llm_provider")) + custom_llm_provider = cast(str | None, kwargs.get("custom_llm_provider")) model_name = reconstruct_model_name(kwargs.get("model", "") or "", custom_llm_provider, metadata) - response_model_name: Optional[str] = None + response_model_name: str | None = None if isinstance(final_response_obj, dict): response_model_name = final_response_obj.get("model") @@ -5477,7 +5468,7 @@ def emit_standard_logging_payload(payload: StandardLoggingPayload): def get_standard_logging_metadata( - metadata: Optional[Dict[str, Any]], + metadata: dict[str, Any] | None, ) -> StandardLoggingMetadata: """ Clean and filter the metadata dictionary to include only the specified keys in StandardLoggingMetadata. @@ -5541,7 +5532,7 @@ def get_standard_logging_metadata( return clean_metadata -def scrub_sensitive_keys_in_metadata(litellm_params: Optional[dict]): +def scrub_sensitive_keys_in_metadata(litellm_params: dict | None): if litellm_params is None: litellm_params = {} @@ -5578,7 +5569,7 @@ def _get_traceback_str_for_error(error_str: str) -> str: from decimal import Decimal # used for unit testing -from typing import Any, Dict, List, Optional, Union +from typing import Any, Optional, Union def create_dummy_standard_logging_payload() -> StandardLoggingPayload: @@ -5622,8 +5613,8 @@ def create_dummy_standard_logging_payload() -> StandardLoggingPayload: saved_cache_cost = Decimal("0.0") # Create messages and response with proper typing - messages: List[Dict[str, str]] = [{"role": "user", "content": "Hello, world!"}] - response: Dict[str, List[Dict[str, Dict[str, str]]]] = {"choices": [{"message": {"content": "Hi there!"}}]} + messages: list[dict[str, str]] = [{"role": "user", "content": "Hello, world!"}] + response: dict[str, list[dict[str, dict[str, str]]]] = {"choices": [{"message": {"content": "Hi there!"}}]} # Main payload initialization return StandardLoggingPayload( # type: ignore diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 0ac22b5bb1e..2b603f353ff 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -6,16 +6,11 @@ import logging import threading import time import traceback +from collections.abc import AsyncIterator, Callable, Iterator from dataclasses import dataclass from typing import ( Any, - AsyncIterator, - Callable, - Dict, - Iterator, - List, NoReturn, - Optional, Union, cast, ) @@ -36,15 +31,13 @@ from litellm.types.llms.openai import OpenAIChatCompletionChunk from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( Delta, -) -from litellm.types.utils import GenericStreamingChunk as GChunk -from litellm.types.utils import ( LlmProviders, ModelResponse, ModelResponseStream, StreamingChoices, Usage, ) +from litellm.types.utils import GenericStreamingChunk as GChunk from ..exceptions import OpenAIError from .core_helpers import map_finish_reason, process_response_headers @@ -105,7 +98,7 @@ class _ProviderChunkParsed: @dataclass(frozen=True, slots=True) class _ProviderChunkEarlyReturn: - value: Any + value: ModelResponseStream | None _ProviderChunkResult = Union[_ProviderChunkParsed, _ProviderChunkEarlyReturn] @@ -116,11 +109,11 @@ class CustomStreamWrapper: self, completion_stream, model, - logging_obj: Any, - custom_llm_provider: Optional[str] = None, + logging_obj: LiteLLMLoggingObject, + custom_llm_provider: str | None = None, stream_options=None, - make_call: Optional[Callable] = None, - _response_headers: Optional[dict] = None, + make_call: Callable | None = None, + _response_headers: dict | None = None, ): self.model = model self.make_call = make_call @@ -139,9 +132,9 @@ class CustomStreamWrapper: self.sent_last_thinking_block = False self.thinking_content = "" - self.system_fingerprint: Optional[str] = None - self.received_finish_reason: Optional[str] = None - self.intermittent_finish_reason: Optional[str] = None # finish reasons that show up mid-stream + self.system_fingerprint: str | None = None + self.received_finish_reason: str | None = None + self.intermittent_finish_reason: str | None = None # finish reasons that show up mid-stream self.special_tokens = [ "<|assistant|>", "<|system|>", @@ -154,7 +147,7 @@ class CustomStreamWrapper: self.holding_chunk = "" self.complete_response = "" self.response_uptil_now = "" - _model_info: Dict = litellm_params.model_info or {} + _model_info: dict = litellm_params.model_info or {} _api_base = get_api_base( model=model or "", @@ -171,7 +164,7 @@ class CustomStreamWrapper: ) # GUARANTEE OPENAI HEADERS IN RESPONSE self._response_headers = _response_headers - self.response_id: Optional[str] = None + self.response_id: str | None = None self.logging_loop = None self.rules = Rules() self.stream_options = stream_options or getattr(logging_obj, "stream_options", None) @@ -179,14 +172,14 @@ class CustomStreamWrapper: self.sent_stream_usage = False self.send_stream_usage = True if self.check_send_stream_usage(self.stream_options) else False self.tool_call = False - self.chunks: List = [] # keep track of the returned chunks - used for calculating the input/output tokens for stream options + self.chunks: list = [] # keep track of the returned chunks - used for calculating the input/output tokens for stream options self._repeated_messages_count = 1 self.is_function_call = self.check_is_function_call(logging_obj=logging_obj) - self.created: Optional[int] = None - self._last_returned_hidden_params: Optional[dict] = None + self.created: int | None = None + self._last_returned_hidden_params: dict | None = None _cached_logging_provider = self.logging_obj.model_call_details.get("custom_llm_provider", None) - self._cached_logging_llm_provider: Optional[str] = _cached_logging_provider + self._cached_logging_llm_provider: str | None = _cached_logging_provider _effective_model = model or "" if custom_llm_provider == "openai" and custom_llm_provider != _cached_logging_provider: _effective_model = "{}/{}".format(_cached_logging_provider, _effective_model) @@ -195,12 +188,12 @@ class CustomStreamWrapper: # Snapshot assumes self._hidden_params is populated from litellm_params # at init and never mutated during the stream. If that ever changes, # this cache must be removed. - self._base_hidden_params: Dict[str, Any] = { + self._base_hidden_params: dict[str, Any] = { **self._hidden_params, "response_cost": None, } - self._post_streaming_hooks: Optional[List] = None + self._post_streaming_hooks: list | None = None def _check_max_streaming_duration(self) -> None: """Raise litellm.Timeout if the stream has exceeded LITELLM_MAX_STREAMING_DURATION_SECONDS.""" @@ -243,7 +236,7 @@ class CustomStreamWrapper: e, ) - def check_send_stream_usage(self, stream_options: Optional[dict]): + def check_send_stream_usage(self, stream_options: dict | None): return stream_options is not None and stream_options.get("include_usage", False) is True def check_is_function_call(self, logging_obj) -> bool: @@ -314,7 +307,7 @@ class CustomStreamWrapper: llm_provider="", ) - def check_special_tokens(self, chunk: str, finish_reason: Optional[str]): + def check_special_tokens(self, chunk: str, finish_reason: str | None): """ Output parse / special tokens for sagemaker + hf streaming. """ @@ -596,7 +589,7 @@ class CustomStreamWrapper: except Exception as e: raise e - def handle_baseten_chunk(self, chunk): + def handle_baseten_chunk(self, chunk) -> str: try: chunk = chunk.decode("utf-8") if len(chunk) > 0: @@ -665,12 +658,12 @@ class CustomStreamWrapper: except Exception as e: raise e - def model_response_creator(self, chunk: Optional[dict] = None, hidden_params: Optional[dict] = None): + def model_response_creator(self, chunk: dict | None = None, hidden_params: dict | None = None): _model = self._cached_model_name _logging_obj_llm_provider = self._cached_logging_llm_provider if chunk is None: - args: Dict[str, Any] = {"model": _model} + args: dict[str, Any] = {"model": _model} else: chunk.pop("model", None) args = {"model": _model} @@ -744,7 +737,7 @@ class CustomStreamWrapper: def copy_model_response_level_provider_specific_fields( self, - original_chunk: Union[ModelResponseStream, OpenAIChatCompletionChunk], + original_chunk: ModelResponseStream | OpenAIChatCompletionChunk, model_response: ModelResponseStream, ) -> ModelResponseStream: """ @@ -759,9 +752,9 @@ class CustomStreamWrapper: def is_chunk_non_empty( self, - completion_obj: Dict[str, Any], + completion_obj: dict[str, Any], model_response: ModelResponseStream, - response_obj: Dict[str, Any], + response_obj: dict[str, Any], ) -> bool: if ( "content" in completion_obj @@ -885,9 +878,9 @@ class CustomStreamWrapper: def return_processed_chunk_logic( # noqa: C901 self, - completion_obj: Dict[str, Any], + completion_obj: dict[str, Any], model_response: ModelResponseStream, - response_obj: Dict[str, Any], + response_obj: dict[str, Any], ): from litellm.litellm_core_utils.core_helpers import ( preserve_upstream_non_openai_attributes, @@ -947,7 +940,7 @@ class CustomStreamWrapper: if response_obj.get("provider_specific_fields") is not None: completion_obj["provider_specific_fields"] = response_obj["provider_specific_fields"] model_response.choices[0].delta = Delta(**completion_obj) - _index: Optional[int] = completion_obj.get("index") + _index: int | None = completion_obj.get("index") if _index is not None: model_response.choices[0].index = _index @@ -1728,7 +1721,7 @@ class CustomStreamWrapper: print_verbose( f"PROCESSED CHUNK PRE CHUNK CREATOR: {chunk.decode('utf-8', errors='replace') if isinstance(chunk, bytes) else chunk}; custom_llm_provider: {self.custom_llm_provider}" ) - response: Optional[ModelResponseStream] = self.chunk_creator(chunk=chunk) + response: ModelResponseStream | None = self.chunk_creator(chunk=chunk) print_verbose(f"PROCESSED CHUNK POST CHUNK CREATOR: {response}") if response is None: @@ -1916,7 +1909,7 @@ class CustomStreamWrapper: elif self.custom_llm_provider == "gemini" and hasattr(chunk, "parts") and len(chunk.parts) == 0: continue - processed_chunk: Optional[ModelResponseStream] = self.chunk_creator(chunk=chunk) + processed_chunk: ModelResponseStream | None = self.chunk_creator(chunk=chunk) if processed_chunk is None: continue @@ -2137,7 +2130,7 @@ class CustomStreamWrapper: return try: partial_response = litellm.stream_chunk_builder(chunks=self.chunks) - usage = cast(Optional[Usage], getattr(partial_response, "usage", None)) + usage = cast(Usage | None, getattr(partial_response, "usage", None)) if usage is None: return self.logging_obj.model_call_details["combined_usage_object"] = usage @@ -2178,7 +2171,7 @@ class CustomStreamWrapper: except Exception as mapping_error: mapped_exception = mapping_error - def _normalize_status_code(exc: Exception) -> Optional[int]: + def _normalize_status_code(exc: Exception) -> int | None: """Best-effort status_code extraction.""" try: code = getattr(exc, "status_code", None) @@ -2218,7 +2211,7 @@ class CustomStreamWrapper: ) @staticmethod - def _strip_sse_data_from_chunk(chunk: Optional[str]) -> Optional[str]: + def _strip_sse_data_from_chunk(chunk: str | None) -> str | None: """ Strips the 'data: ' prefix from Server-Sent Events (SSE) chunks. @@ -2254,7 +2247,7 @@ class CustomStreamWrapper: return chunk -def calculate_total_usage(chunks: List[ModelResponse]) -> Usage: +def calculate_total_usage(chunks: list[ModelResponse]) -> Usage: """Assume most recent usage chunk has total usage uptil then.""" prompt_tokens: int = 0 completion_tokens: int = 0 diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index d6acbaae434..4ed7121c6c7 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2,19 +2,15 @@ import asyncio import json import os import ssl +from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping from contextlib import asynccontextmanager from functools import lru_cache from typing import ( TYPE_CHECKING, Any, - AsyncIterator, - Coroutine, - Dict, - Iterator, - List, Literal, Optional, - Tuple, + TypeVar, Union, cast, get_type_hints, @@ -23,6 +19,7 @@ from urllib.parse import parse_qs, urlencode, urlparse, urlunparse import httpx # type: ignore from openai.types.file_deleted import FileDeleted +from tiktoken import Encoding import litellm import litellm.litellm_core_utils @@ -149,6 +146,8 @@ from litellm.utils import ( async_pre_call_deployment_hook, ) +FakeStreamWrappableResponseT = TypeVar("FakeStreamWrappableResponseT") + def _rust_responses_websocket_enabled( custom_llm_provider: str | None, @@ -162,7 +161,11 @@ from .http_handler import get_shared_realtime_ssl_context if TYPE_CHECKING: from aiohttp import ClientSession + from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import ( + FakeAnthropicMessagesStreamIterator, + ) from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( AnthropicMessagesStreamingResponse, ) @@ -189,11 +192,11 @@ def _google_genai_streaming_hidden_params( litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, response_headers: httpx.Headers, -) -> Dict[str, Any]: +) -> dict[str, object]: """Pre-stream metadata for proxy response headers (mirrors CustomStreamWrapper._hidden_params).""" from litellm.litellm_core_utils.core_helpers import process_response_headers - _model_info: Dict[str, Any] = dict(getattr(litellm_params, "model_info", None) or {}) + _model_info: dict[str, object] = dict(getattr(litellm_params, "model_info", None) or {}) _raw_id = _model_info.get("id") or logging_obj.get_router_model_id() or "" _model_id = _raw_id if isinstance(_raw_id, str) else str(_raw_id) return { @@ -210,7 +213,7 @@ def _responses_api_optional_request_param_names() -> frozenset[str]: return frozenset(get_type_hints(ResponsesAPIOptionalRequestParams).keys()) -def _custom_logger_callbacks(logging_obj: Any) -> list[Any]: +def _custom_logger_callbacks(logging_obj: LiteLLMLoggingObj) -> list["CustomLogger"]: from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import ( get_custom_logger_compatible_class, @@ -221,7 +224,7 @@ def _custom_logger_callbacks(logging_obj: Any) -> list[Any]: if isinstance(dynamic_success_callbacks, (list, tuple)): callbacks.extend(dynamic_success_callbacks) - custom_loggers: list[Any] = [] + custom_loggers: list[CustomLogger] = [] for cb in callbacks: if isinstance(cb, str): resolved = get_custom_logger_compatible_class(cb) # type: ignore[arg-type] @@ -233,7 +236,7 @@ def _custom_logger_callbacks(logging_obj: Any) -> list[Any]: return custom_loggers -def _has_pre_call_deployment_hook(logging_obj: Any) -> bool: +def _has_pre_call_deployment_hook(logging_obj: LiteLLMLoggingObj) -> bool: from litellm.integrations.custom_logger import CustomLogger base_func = CustomLogger.async_pre_call_deployment_hook @@ -252,16 +255,16 @@ class BaseLLMHTTPHandler: api_base: str, headers: dict, data: dict, - timeout: Union[float, httpx.Timeout], + timeout: float | httpx.Timeout, litellm_params: dict, logging_obj: LiteLLMLoggingObj, stream: bool = False, - signed_json_body: Optional[bytes] = None, + signed_json_body: bytes | None = None, ) -> httpx.Response: """Common implementation across stream + non-stream calls. Meant to ensure consistent error-handling.""" max_retry_on_unprocessable_entity_error = provider_config.max_retry_on_unprocessable_entity_error - response: Optional[httpx.Response] = None + response: httpx.Response | None = None for i in range(max(max_retry_on_unprocessable_entity_error, 1)): try: response = await async_httpx_client.post( @@ -302,15 +305,15 @@ class BaseLLMHTTPHandler: api_base: str, headers: dict, data: dict, - timeout: Union[float, httpx.Timeout], + timeout: float | httpx.Timeout, litellm_params: dict, logging_obj: LiteLLMLoggingObj, stream: bool = False, - signed_json_body: Optional[bytes] = None, + signed_json_body: bytes | None = None, ) -> httpx.Response: max_retry_on_unprocessable_entity_error = provider_config.max_retry_on_unprocessable_entity_error - response: Optional[httpx.Response] = None + response: httpx.Response | None = None for i in range(max(max_retry_on_unprocessable_entity_error, 1)): try: @@ -352,18 +355,18 @@ class BaseLLMHTTPHandler: api_base: str, headers: dict, data: dict, - timeout: Union[float, httpx.Timeout], + timeout: float | httpx.Timeout, model: str, model_response: ModelResponse, logging_obj: LiteLLMLoggingObj, messages: list, optional_params: dict, litellm_params: dict, - encoding: Any, - api_key: Optional[str] = None, - client: Optional[AsyncHTTPHandler] = None, + encoding: Encoding, + api_key: str | None = None, + client: AsyncHTTPHandler | None = None, json_mode: bool = False, - signed_json_body: Optional[bytes] = None, + signed_json_body: bytes | None = None, shared_session: Optional["ClientSession"] = None, ): if client is None: @@ -422,25 +425,25 @@ class BaseLLMHTTPHandler: self, model: str, messages: list, - api_base: Optional[str], + api_base: str | None, custom_llm_provider: str, model_response: ModelResponse, - encoding, + encoding: Encoding, logging_obj: LiteLLMLoggingObj, optional_params: dict, - timeout: Union[float, httpx.Timeout], + timeout: float | httpx.Timeout, litellm_params: dict, acompletion: bool, - stream: Optional[bool] = False, + stream: bool | None = False, fake_stream: bool = False, - api_key: Optional[str] = None, - headers: Optional[Dict[str, Any]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - provider_config: Optional[BaseConfig] = None, + api_key: str | None = None, + headers: dict[str, Any] | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, + provider_config: BaseConfig | None = None, shared_session: Optional["ClientSession"] = None, ): json_mode: bool = optional_params.pop("json_mode", False) - extra_body: Optional[dict] = optional_params.pop("extra_body", None) + extra_body: dict | None = optional_params.pop("extra_body", None) provider_config = provider_config or ProviderConfigManager.get_provider_chat_config( model=model, provider=litellm.LlmProviders(custom_llm_provider) @@ -640,18 +643,18 @@ class BaseLLMHTTPHandler: api_base: str, headers: dict, data: dict, - signed_json_body: Optional[bytes], + signed_json_body: bytes | None, original_data: dict, model: str, messages: list, logging_obj, optional_params: dict, litellm_params: dict, - timeout: Union[float, httpx.Timeout], + timeout: float | httpx.Timeout, fake_stream: bool = False, - client: Optional[HTTPHandler] = None, + client: HTTPHandler | None = None, json_mode: bool = False, - ) -> Tuple[Any, dict]: + ) -> tuple[MockResponseIterator | BaseModelResponseIterator, dict]: if client is None or not isinstance(client, HTTPHandler): sync_httpx_client = _get_httpx_client( { @@ -691,7 +694,9 @@ class BaseLLMHTTPHandler: json_mode=json_mode, ) - completion_stream: Any = MockResponseIterator(model_response=model_response, json_mode=json_mode) + completion_stream: MockResponseIterator | BaseModelResponseIterator = MockResponseIterator( + model_response=model_response, json_mode=json_mode + ) else: completion_stream = provider_config.get_model_response_iterator( streaming_response=response.iter_lines(), @@ -717,15 +722,15 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, headers: dict, provider_config: BaseConfig, - timeout: Union[float, httpx.Timeout], + timeout: float | httpx.Timeout, logging_obj: LiteLLMLoggingObj, data: dict, litellm_params: dict, optional_params: dict, fake_stream: bool = False, - client: Optional[AsyncHTTPHandler] = None, - json_mode: Optional[bool] = None, - signed_json_body: Optional[bytes] = None, + client: AsyncHTTPHandler | None = None, + json_mode: bool | None = None, + signed_json_body: bytes | None = None, ): if provider_config.has_custom_stream_wrapper is True: return await provider_config.get_async_custom_stream_wrapper( @@ -776,14 +781,14 @@ class BaseLLMHTTPHandler: data: dict, messages: list, logging_obj: LiteLLMLoggingObj, - timeout: Union[float, httpx.Timeout], + timeout: float | httpx.Timeout, litellm_params: dict, optional_params: dict, fake_stream: bool = False, - client: Optional[AsyncHTTPHandler] = None, - json_mode: Optional[bool] = None, - signed_json_body: Optional[bytes] = None, - ) -> Tuple[Any, httpx.Headers]: + client: AsyncHTTPHandler | None = None, + json_mode: bool | None = None, + signed_json_body: bytes | None = None, + ) -> tuple[MockResponseIterator | BaseModelResponseIterator, httpx.Headers]: """ Helper function for making an async call with stream. @@ -827,7 +832,9 @@ class BaseLLMHTTPHandler: json_mode=json_mode, ) - completion_stream: Any = MockResponseIterator(model_response=model_response, json_mode=json_mode) + completion_stream: MockResponseIterator | BaseModelResponseIterator = MockResponseIterator( + model_response=model_response, json_mode=json_mode + ) else: completion_stream = provider_config.get_model_response_iterator( streaming_response=response.aiter_lines(), sync_stream=False @@ -870,14 +877,14 @@ class BaseLLMHTTPHandler: timeout: float, custom_llm_provider: str, logging_obj: LiteLLMLoggingObj, - api_base: Optional[str], + api_base: str | None, optional_params: dict, litellm_params: dict, model_response: EmbeddingResponse, - api_key: Optional[str] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - aembedding: Optional[bool] = False, - headers: Optional[Dict[str, Any]] = None, + api_key: str | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, + aembedding: bool | None = False, + headers: dict[str, Any] | None = None, ) -> EmbeddingResponse: provider_config = ProviderConfigManager.get_provider_embedding_config( model=model, provider=litellm.LlmProviders(custom_llm_provider) @@ -999,10 +1006,10 @@ class BaseLLMHTTPHandler: logging_obj: LiteLLMLoggingObj, optional_params: dict, litellm_params: dict, - api_key: Optional[str] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - signed_body: Optional[bytes] = None, + api_key: str | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, + signed_body: bytes | None = None, ) -> EmbeddingResponse: if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( @@ -1047,15 +1054,15 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, logging_obj: LiteLLMLoggingObj, provider_config: BaseRerankConfig, - optional_rerank_params: Dict, - timeout: Optional[Union[float, httpx.Timeout]], + optional_rerank_params: dict, + timeout: float | httpx.Timeout | None, model_response: RerankResponse, _is_async: bool = False, - headers: Optional[Dict[str, Any]] = None, - api_key: Optional[str] = None, - api_base: Optional[str] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - litellm_params: Optional[Dict[str, Any]] = None, + headers: dict[str, str] | None = None, + api_key: str | None = None, + api_base: str | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, + litellm_params: dict[str, Any] | None = None, ) -> RerankResponse: # get config from model, custom llm provider headers = provider_config.validate_environment( @@ -1141,9 +1148,9 @@ class BaseLLMHTTPHandler: model_response: RerankResponse, api_base: str, headers: dict, - api_key: Optional[str] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + api_key: str | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, ) -> RerankResponse: if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client(llm_provider=litellm.LlmProviders(custom_llm_provider)) @@ -1175,11 +1182,11 @@ class BaseLLMHTTPHandler: optional_params: dict, litellm_params: dict, logging_obj: LiteLLMLoggingObj, - api_key: Optional[str], - api_base: Optional[str], - headers: Optional[Dict[str, Any]], + api_key: str | None, + api_base: str | None, + headers: dict[str, str] | None, provider_config: BaseAudioTranscriptionConfig, - ) -> Tuple[dict, str, Union[dict, bytes, None], Optional[dict]]: + ) -> tuple[dict, str, dict | bytes | None, dict | None]: """ Shared logic for preparing audio transcription requests. Returns: (headers, complete_url, data, files) @@ -1244,7 +1251,7 @@ class BaseLLMHTTPHandler: model_response: TranscriptionResponse, logging_obj: LiteLLMLoggingObj, optional_params: dict, - api_key: Optional[str], + api_key: str | None, ) -> TranscriptionResponse: """Shared logic for transforming audio transcription responses.""" return provider_config.transform_audio_transcription_response( @@ -1261,15 +1268,15 @@ class BaseLLMHTTPHandler: timeout: float, max_retries: int, logging_obj: LiteLLMLoggingObj, - api_key: Optional[str], - api_base: Optional[str], + api_key: str | None, + api_base: str | None, custom_llm_provider: str, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, atranscription: bool = False, - headers: Optional[Dict[str, Any]] = None, - provider_config: Optional[BaseAudioTranscriptionConfig] = None, + headers: dict[str, str] | None = None, + provider_config: BaseAudioTranscriptionConfig | None = None, shared_session: Optional["ClientSession"] = None, - ) -> Union[TranscriptionResponse, Coroutine[Any, Any, TranscriptionResponse]]: + ) -> TranscriptionResponse | Coroutine[None, None, TranscriptionResponse]: if provider_config is None: raise ValueError(f"No provider config found for model: {model} and provider: {custom_llm_provider}") @@ -1347,12 +1354,12 @@ class BaseLLMHTTPHandler: timeout: float, max_retries: int, logging_obj: LiteLLMLoggingObj, - api_key: Optional[str], - api_base: Optional[str], + api_key: str | None, + api_base: str | None, custom_llm_provider: str, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - headers: Optional[Dict[str, Any]] = None, - provider_config: Optional[BaseAudioTranscriptionConfig] = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, + headers: dict[str, str] | None = None, + provider_config: BaseAudioTranscriptionConfig | None = None, shared_session: Optional["ClientSession"] = None, ) -> TranscriptionResponse: if provider_config is None: @@ -1412,15 +1419,15 @@ class BaseLLMHTTPHandler: def _prepare_ocr_request( self, model: str, - document: Dict[str, str], + document: dict[str, str], optional_params: dict, logging_obj: LiteLLMLoggingObj, - api_key: Optional[str], - api_base: Optional[str], - headers: Optional[Dict[str, Any]], + api_key: str | None, + api_base: str | None, + headers: dict[str, object] | None, provider_config: BaseOCRConfig, litellm_params: dict, - ) -> Tuple[Dict[str, Any], str, Dict[str, Any], None]: + ) -> tuple[dict[str, Any], str, dict[str, Any], None]: """ Shared logic for preparing OCR requests. Returns: (headers, complete_url, data, files) @@ -1478,15 +1485,15 @@ class BaseLLMHTTPHandler: async def _async_prepare_ocr_request( self, model: str, - document: Dict[str, str], + document: dict[str, str], optional_params: dict, logging_obj: LiteLLMLoggingObj, - api_key: Optional[str], - api_base: Optional[str], - headers: Optional[Dict[str, Any]], + api_key: str | None, + api_base: str | None, + headers: dict[str, object] | None, provider_config: BaseOCRConfig, litellm_params: dict, - ) -> Tuple[Dict[str, Any], str, Dict[str, Any], None]: + ) -> tuple[dict[str, Any], str, dict[str, Any], None]: """ Async version of _prepare_ocr_request for providers that need async transforms. Returns: (headers, complete_url, data, files) @@ -1558,19 +1565,19 @@ class BaseLLMHTTPHandler: def ocr( self, model: str, - document: Dict[str, str], + document: dict[str, str], optional_params: dict, - timeout: Union[float, httpx.Timeout], + timeout: float | httpx.Timeout, logging_obj: LiteLLMLoggingObj, - api_key: Optional[str], - api_base: Optional[str], + api_key: str | None, + api_base: str | None, custom_llm_provider: str, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, aocr: bool = False, - headers: Optional[Dict[str, Any]] = None, - provider_config: Optional[BaseOCRConfig] = None, - litellm_params: Optional[dict] = None, - ) -> Union[OCRResponse, Coroutine[Any, Any, OCRResponse]]: + headers: dict[str, object] | None = None, + provider_config: BaseOCRConfig | None = None, + litellm_params: dict | None = None, + ) -> OCRResponse | Coroutine[object, object, OCRResponse]: """ Sync OCR handler. """ @@ -1633,17 +1640,17 @@ class BaseLLMHTTPHandler: async def async_ocr( self, model: str, - document: Dict[str, str], + document: dict[str, str], optional_params: dict, - timeout: Union[float, httpx.Timeout], + timeout: float | httpx.Timeout, logging_obj: LiteLLMLoggingObj, - api_key: Optional[str], - api_base: Optional[str], + api_key: str | None, + api_base: str | None, custom_llm_provider: str, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - headers: Optional[Dict[str, Any]] = None, - provider_config: Optional[BaseOCRConfig] = None, - litellm_params: Optional[dict] = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, + headers: dict[str, object] | None = None, + provider_config: BaseOCRConfig | None = None, + litellm_params: dict | None = None, ) -> OCRResponse: """ Async OCR handler. @@ -1694,18 +1701,18 @@ class BaseLLMHTTPHandler: def search( self, - query: Union[str, List[str]], + query: str | list[str], optional_params: dict, - timeout: Union[float, httpx.Timeout], + timeout: float | httpx.Timeout, logging_obj: LiteLLMLoggingObj, - api_key: Optional[str], - api_base: Optional[str], + api_key: str | None, + api_base: str | None, custom_llm_provider: str, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, asearch: bool = False, - headers: Optional[Dict[str, Any]] = None, - provider_config: Optional[BaseSearchConfig] = None, - ) -> Union[SearchResponse, Coroutine[Any, Any, SearchResponse]]: + headers: dict[str, str] | None = None, + provider_config: BaseSearchConfig | None = None, + ) -> SearchResponse | Coroutine[None, None, SearchResponse]: """ Sync Search handler. """ @@ -1790,16 +1797,16 @@ class BaseLLMHTTPHandler: async def async_search( self, - query: Union[str, List[str]], + query: str | list[str], optional_params: dict, - timeout: Union[float, httpx.Timeout], + timeout: float | httpx.Timeout, logging_obj: LiteLLMLoggingObj, - api_key: Optional[str], - api_base: Optional[str], + api_key: str | None, + api_base: str | None, custom_llm_provider: str, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - headers: Optional[Dict[str, Any]] = None, - provider_config: Optional[BaseSearchConfig] = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, + headers: dict[str, str] | None = None, + provider_config: BaseSearchConfig | None = None, ) -> SearchResponse: """ Async Search handler. @@ -1884,15 +1891,15 @@ class BaseLLMHTTPHandler: headers: dict, # str when the caller passes a pre-serialized (unsigned) body to avoid # re-dumping; bytes when a provider signed the request (e.g. Bedrock). - signed_json_body: Optional[Union[str, bytes]], + signed_json_body: str | bytes | None, request_body: dict, stream: bool, logging_obj: LiteLLMLoggingObj, provider_config: BaseAnthropicMessagesConfig, litellm_params: GenericLiteLLMParams, - api_key: Optional[str], + api_key: str | None, model: str, - timeout: Optional[Union[float, httpx.Timeout]] = None, + timeout: float | httpx.Timeout | None = None, ) -> httpx.Response: max_attempts = max(provider_config.max_retry_on_anthropic_messages_http_error, 1) litellm_params_dict = dict(litellm_params) @@ -1945,7 +1952,7 @@ class BaseLLMHTTPHandler: litellm_params: GenericLiteLLMParams, stream: bool, custom_llm_provider: str, - ) -> Optional[Union[float, httpx.Timeout]]: + ) -> float | httpx.Timeout | None: from litellm.litellm_core_utils.completion_timeout import CompletionTimeout from litellm.litellm_core_utils.request_timeout_resolver import ( get_configured_request_timeout, @@ -1969,19 +1976,19 @@ class BaseLLMHTTPHandler: async def async_anthropic_messages_handler( self, model: str, - messages: List[Dict], + messages: list[dict], anthropic_messages_provider_config: BaseAnthropicMessagesConfig, - anthropic_messages_optional_request_params: Dict, + anthropic_messages_optional_request_params: dict, custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - client: Optional[AsyncHTTPHandler] = None, - extra_headers: Optional[Dict[str, Any]] = None, - api_key: Optional[str] = None, - api_base: Optional[str] = None, - stream: Optional[bool] = False, - kwargs: Optional[Dict[str, Any]] = None, - ) -> Union[AnthropicMessagesResponse, AsyncIterator]: + client: AsyncHTTPHandler | None = None, + extra_headers: dict[str, str] | None = None, + api_key: str | None = None, + api_base: str | None = None, + stream: bool | None = False, + kwargs: dict[str, object] | None = None, + ) -> AnthropicMessagesResponse | AsyncIterator: from litellm.litellm_core_utils.get_provider_specific_headers import ( ProviderSpecificHeaderUtils, ) @@ -1994,7 +2001,7 @@ class BaseLLMHTTPHandler: # Prepare headers kwargs = kwargs or {} provider_specific_header = cast( - Optional[litellm.types.utils.ProviderSpecificHeader], + litellm.types.utils.ProviderSpecificHeader | None, kwargs.get("provider_specific_header", None), ) provider_specific_headers = ProviderSpecificHeaderUtils.get_provider_specific_headers( @@ -2006,9 +2013,9 @@ class BaseLLMHTTPHandler: extra_headers_from_kwargs = kwargs.get("extra_headers", None) # Merge all header sources: forwarded < extra_headers < provider_specific merged_headers = {} - if forwarded_headers: + if isinstance(forwarded_headers, dict): merged_headers.update(forwarded_headers) - if extra_headers_from_kwargs: + if isinstance(extra_headers_from_kwargs, dict): merged_headers.update(extra_headers_from_kwargs) if provider_specific_headers: merged_headers.update(provider_specific_headers) @@ -2156,7 +2163,7 @@ class BaseLLMHTTPHandler: # used for logging + cost tracking logging_obj.model_call_details["httpx_response"] = response - initial_response: Union[AsyncIterator, AnthropicMessagesResponse] + initial_response: AsyncIterator | AnthropicMessagesResponse if stream: from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( AnthropicMessagesStreamingResponse, @@ -2328,22 +2335,19 @@ class BaseLLMHTTPHandler: def anthropic_messages_handler( self, model: str, - messages: List[Dict], + messages: list[dict], anthropic_messages_provider_config: BaseAnthropicMessagesConfig, - anthropic_messages_optional_request_params: Dict, + anthropic_messages_optional_request_params: dict, custom_llm_provider: str, _is_async: bool, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - api_key: Optional[str] = None, - api_base: Optional[str] = None, - stream: Optional[bool] = False, - kwargs: Optional[Dict[str, Any]] = None, - ) -> Union[ - AnthropicMessagesResponse, - Coroutine[Any, Any, Union[AnthropicMessagesResponse, AsyncIterator]], - ]: + client: HTTPHandler | AsyncHTTPHandler | None = None, + api_key: str | None = None, + api_base: str | None = None, + stream: bool | None = False, + kwargs: dict[str, object] | None = None, + ) -> AnthropicMessagesResponse | Coroutine[None, None, AnthropicMessagesResponse | AsyncIterator]: """ LLM HTTP Handler for Anthropic Messages """ @@ -2369,14 +2373,14 @@ class BaseLLMHTTPHandler: self, *, model: str, - input: Union[str, ResponseInputParam], + input: str | ResponseInputParam, custom_llm_provider: str, response_api_optional_request_params: dict[str, Any], litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, ) -> tuple[ str, - Union[str, ResponseInputParam], + str | ResponseInputParam, str, dict[str, Any], GenericLiteLLMParams, @@ -2428,7 +2432,7 @@ class BaseLLMHTTPHandler: return ( str(modified_kwargs["model"]) if "model" in modified_kwargs else model, cast( - Union[str, ResponseInputParam], + str | ResponseInputParam, modified_kwargs["input"] if "input" in modified_kwargs else input, ), ( @@ -2443,25 +2447,25 @@ class BaseLLMHTTPHandler: def response_api_handler( self, model: str, - input: Union[str, ResponseInputParam], + input: str | ResponseInputParam, responses_api_provider_config: BaseResponsesAPIConfig, response_api_optional_request_params: dict[str, Any], custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, fake_stream: bool = False, - litellm_metadata: Optional[Dict[str, Any]] = None, + litellm_metadata: dict[str, object] | None = None, shared_session: Optional["ClientSession"] = None, - ) -> Union[ - ResponsesAPIResponse, - BaseResponsesAPIStreamingIterator, - Coroutine[Any, Any, Union[ResponsesAPIResponse, BaseResponsesAPIStreamingIterator]], - ]: + ) -> ( + ResponsesAPIResponse + | BaseResponsesAPIStreamingIterator + | Coroutine[None, None, ResponsesAPIResponse | BaseResponsesAPIStreamingIterator] + ): """ Handles responses API requests. When _is_async=True, returns a coroutine instead of making the call directly. @@ -2543,7 +2547,7 @@ class BaseLLMHTTPHandler: # Preserve the OpenAI-style request context (not sent to the provider) for streaming # hooks/metadata; the streaming iterator now consumes this to run deployment hooks # with the same info as chat, including litellm_params. - request_context: Dict[str, Any] = {"input": input} + request_context: dict[str, Any] = {"input": input} try: request_context.update(response_api_optional_request_params) except Exception: @@ -2573,7 +2577,7 @@ class BaseLLMHTTPHandler: stream=stream, fake_stream=fake_stream, ) - body_kwargs: Dict[str, Any] = {"data": signed_body} if signed_body is not None else {"json": data} + body_kwargs: dict[str, Any] = {"data": signed_body} if signed_body is not None else {"json": data} ## LOGGING logging_obj.pre_call( @@ -2657,20 +2661,20 @@ class BaseLLMHTTPHandler: async def async_response_api_handler( self, model: str, - input: Union[str, ResponseInputParam], + input: str | ResponseInputParam, responses_api_provider_config: BaseResponsesAPIConfig, - response_api_optional_request_params: Dict, + response_api_optional_request_params: dict, custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, fake_stream: bool = False, - litellm_metadata: Optional[Dict[str, Any]] = None, + litellm_metadata: dict[str, object] | None = None, shared_session: Optional["ClientSession"] = None, - ) -> Union[ResponsesAPIResponse, BaseResponsesAPIStreamingIterator]: + ) -> ResponsesAPIResponse | BaseResponsesAPIStreamingIterator: """ Async version of the responses API handler. Uses async HTTP client to make requests. @@ -2720,7 +2724,7 @@ class BaseLLMHTTPHandler: # Preserve the OpenAI-style request context (not sent to the provider) for streaming # hooks/metadata; the streaming iterator now consumes this to run deployment hooks # with the same info as chat, including litellm_params. - request_context: Dict[str, Any] = {"input": input} + request_context: dict[str, Any] = {"input": input} try: request_context.update(response_api_optional_request_params) except Exception: @@ -2747,7 +2751,7 @@ class BaseLLMHTTPHandler: stream=stream, fake_stream=fake_stream, ) - body_kwargs: Dict[str, Any] = {"data": signed_body} if signed_body is not None else {"json": data} + body_kwargs: dict[str, Any] = {"data": signed_body} if signed_body is not None else {"json": data} ## LOGGING logging_obj.pre_call( @@ -2846,11 +2850,11 @@ class BaseLLMHTTPHandler: responses_api_provider_config: BaseResponsesAPIConfig, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - custom_llm_provider: Optional[str], - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + custom_llm_provider: str | None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, shared_session: Optional["ClientSession"] = None, ) -> DeleteResponseResult: @@ -2902,7 +2906,7 @@ class BaseLLMHTTPHandler: }, ) - delete_kwargs: Dict[str, Any] = { + delete_kwargs: dict[str, Any] = { "url": url, "headers": headers, "timeout": timeout, @@ -2930,14 +2934,14 @@ class BaseLLMHTTPHandler: responses_api_provider_config: BaseResponsesAPIConfig, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - custom_llm_provider: Optional[str], - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + custom_llm_provider: str | None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, shared_session: Optional["ClientSession"] = None, - ) -> Union[DeleteResponseResult, Coroutine[Any, Any, DeleteResponseResult]]: + ) -> DeleteResponseResult | Coroutine[None, None, DeleteResponseResult]: """ Async version of the responses API handler. Uses async HTTP client to make requests. @@ -2992,7 +2996,7 @@ class BaseLLMHTTPHandler: }, ) - delete_kwargs: Dict[str, Any] = { + delete_kwargs: dict[str, Any] = { "url": url, "headers": headers, "timeout": timeout, @@ -3020,14 +3024,14 @@ class BaseLLMHTTPHandler: responses_api_provider_config: BaseResponsesAPIConfig, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - custom_llm_provider: Optional[str] = None, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + custom_llm_provider: str | None = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, shared_session: Optional["ClientSession"] = None, - ) -> Union[ResponsesAPIResponse, Coroutine[Any, Any, ResponsesAPIResponse]]: + ) -> ResponsesAPIResponse | Coroutine[None, None, ResponsesAPIResponse]: """ Get a response by ID Uses GET /v1/responses/{response_id} endpoint in the responses API @@ -3101,11 +3105,11 @@ class BaseLLMHTTPHandler: responses_api_provider_config: BaseResponsesAPIConfig, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - custom_llm_provider: Optional[str] = None, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + custom_llm_provider: str | None = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, shared_session: Optional["ClientSession"] = None, ) -> ResponsesAPIResponse: """ @@ -3177,18 +3181,18 @@ class BaseLLMHTTPHandler: responses_api_provider_config: BaseResponsesAPIConfig, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - custom_llm_provider: Optional[str] = None, - after: Optional[str] = None, - before: Optional[str] = None, - include: Optional[List[str]] = None, + custom_llm_provider: str | None = None, + after: str | None = None, + before: str | None = None, + include: list[str] | None = None, limit: int = 20, order: Literal["asc", "desc"] = "desc", - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, shared_session: Optional["ClientSession"] = None, - ) -> Union[Dict, Coroutine[Any, Any, Dict]]: + ) -> dict | Coroutine[None, None, dict]: if _is_async: return self.async_list_responses_input_items( response_id=response_id, @@ -3263,17 +3267,17 @@ class BaseLLMHTTPHandler: responses_api_provider_config: BaseResponsesAPIConfig, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - custom_llm_provider: Optional[str] = None, - after: Optional[str] = None, - before: Optional[str] = None, - include: Optional[List[str]] = None, + custom_llm_provider: str | None = None, + after: str | None = None, + before: str | None = None, + include: list[str] | None = None, limit: int = 20, order: Literal["asc", "desc"] = "desc", - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, shared_session: Optional["ClientSession"] = None, - ) -> Dict: + ) -> dict: if client is None or not isinstance(client, AsyncHTTPHandler): verbose_logger.debug( f"Creating HTTP client for list_input_items with shared_session: {id(shared_session) if shared_session else None}" @@ -3336,7 +3340,7 @@ class BaseLLMHTTPHandler: response: httpx.Response, upload_url_location: str, upload_url_key: str = "upload_url", - ) -> tuple[Optional[str], Optional[dict]]: + ) -> tuple[str | None, dict | None]: """ Extract upload URL from initial file creation response. @@ -3352,13 +3356,15 @@ class BaseLLMHTTPHandler: """ if upload_url_location == "headers": # Google Cloud Storage style - URL in X-Goog-Upload-URL header - upload_url = response.headers.get("X-Goog-Upload-URL") - return upload_url, None + raw_header_upload_url = response.headers.get("X-Goog-Upload-URL") + header_upload_url = raw_header_upload_url if isinstance(raw_header_upload_url, str) else None + return header_upload_url, None else: # Response body style (e.g., Manus, S3 presigned URLs) try: response_data = response.json() - upload_url = response_data.get(upload_url_key) + raw_upload_url = response_data.get(upload_url_key) + upload_url = raw_upload_url if isinstance(raw_upload_url, str) else None return upload_url, response_data if upload_url else None except Exception: return None, None @@ -3369,13 +3375,13 @@ class BaseLLMHTTPHandler: litellm_params: dict, provider_config: BaseFilesConfig, headers: dict, - api_base: Optional[str], - api_key: Optional[str], + api_base: str | None, + api_key: str | None, logging_obj: LiteLLMLoggingObj, _is_async: bool = False, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - ) -> Union[OpenAIFileObject, Coroutine[Any, Any, OpenAIFileObject]]: + client: HTTPHandler | AsyncHTTPHandler | None = None, + timeout: float | httpx.Timeout | None = None, + ) -> OpenAIFileObject | Coroutine[None, None, OpenAIFileObject]: """ Creates a file using Gemini's two-step upload process """ @@ -3477,7 +3483,7 @@ class BaseLLMHTTPHandler: ): # Handle pre-signed requests (e.g., from Bedrock S3 uploads) # Type narrowing: this is a plain dict, not TwoStepFileUploadConfig - presigned_request = cast(Dict[str, Any], transformed_request) + presigned_request = cast(dict[str, Any], transformed_request) upload_response = getattr(sync_httpx_client, presigned_request["method"].lower())( url=presigned_request["url"], headers=presigned_request["headers"], @@ -3523,7 +3529,7 @@ class BaseLLMHTTPHandler: elif isinstance(transformed_request, dict) and "file" in transformed_request: # Handle multipart form-data uploads (e.g., Anthropic Files API) # The dict contains tuples suitable for httpx's `files` parameter - file_request = cast(Dict[str, Any], transformed_request) + file_request = cast(dict[str, Any], transformed_request) upload_response = sync_httpx_client.post( url=api_base, headers=headers, @@ -3555,8 +3561,8 @@ class BaseLLMHTTPHandler: headers: dict, api_base: str, logging_obj: LiteLLMLoggingObj, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, + timeout: float | httpx.Timeout | None = None, ): """ Creates a file using Gemini's two-step upload process @@ -3639,7 +3645,7 @@ class BaseLLMHTTPHandler: ): # Handle pre-signed requests (e.g., from Bedrock S3 uploads) # Type narrowing: this is a plain dict, not TwoStepFileUploadConfig - presigned_request = cast(Dict[str, Any], transformed_request) + presigned_request = cast(dict[str, Any], transformed_request) upload_response = await getattr(async_httpx_client, presigned_request["method"].lower())( url=presigned_request["url"], headers=presigned_request["headers"], @@ -3725,19 +3731,18 @@ class BaseLLMHTTPHandler: *, client: HTTPHandler, url: str, - base_headers: Dict[str, str], + base_headers: dict[str, str], body_stream: BaseFileUploadStream, content_type: str, - timeout: Optional[Union[float, httpx.Timeout]], + timeout: float | httpx.Timeout | None, ) -> httpx.Response: headers = {**base_headers, "Content-Type": content_type} - kwargs: Dict[str, Any] = { - "headers": headers, - "content": self._iter_in_blocks(body_stream.iter_bytes(), self._MEDIA_UPLOAD_BLOCK_SIZE), - } - if timeout is not None: - kwargs["timeout"] = timeout - resp = client.client.post(url, **kwargs) + content = self._iter_in_blocks(body_stream.iter_bytes(), self._MEDIA_UPLOAD_BLOCK_SIZE) + resp = ( + client.client.post(url, headers=headers, content=content, timeout=timeout) + if timeout is not None + else client.client.post(url, headers=headers, content=content) + ) self._check_media_upload_response(resp) return resp @@ -3746,10 +3751,10 @@ class BaseLLMHTTPHandler: *, client: AsyncHTTPHandler, url: str, - base_headers: Dict[str, str], + base_headers: dict[str, str], body_stream: BaseFileUploadStream, content_type: str, - timeout: Optional[Union[float, httpx.Timeout]], + timeout: float | httpx.Timeout | None, ) -> httpx.Response: """Stream the transformed body straight to a single media upload. Each block is produced on a worker thread (the transform never runs on the @@ -3768,10 +3773,11 @@ class BaseLLMHTTPHandler: break yield cast(bytes, block) - kwargs: Dict[str, Any] = {"headers": headers, "content": _abody()} - if timeout is not None: - kwargs["timeout"] = timeout - resp = await client.client.post(url, **kwargs) + resp = ( + await client.client.post(url, headers=headers, content=_abody(), timeout=timeout) + if timeout is not None + else await client.client.post(url, headers=headers, content=_abody()) + ) await resp.aread() self._check_media_upload_response(resp) return resp @@ -3782,14 +3788,14 @@ class BaseLLMHTTPHandler: litellm_params: dict, provider_config: "BaseBatchesConfig", headers: dict, - api_base: Optional[str], - api_key: Optional[str], + api_base: str | None, + api_key: str | None, logging_obj: "LiteLLMLoggingObj", _is_async: bool = False, - client: Optional[Union["HTTPHandler", "AsyncHTTPHandler"]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - model: Optional[str] = None, - ) -> Union["LiteLLMBatch", Coroutine[Any, Any, "LiteLLMBatch"]]: + client: Union["HTTPHandler", "AsyncHTTPHandler"] | None = None, + timeout: float | httpx.Timeout | None = None, + model: str | None = None, + ) -> Union["LiteLLMBatch", Coroutine[None, None, "LiteLLMBatch"]]: """ Creates a batch using provider-specific batch creation process """ @@ -3894,14 +3900,14 @@ class BaseLLMHTTPHandler: litellm_params: dict, provider_config: "BaseBatchesConfig", headers: dict, - api_base: Optional[str], - api_key: Optional[str], + api_base: str | None, + api_key: str | None, logging_obj: "LiteLLMLoggingObj", _is_async: bool = False, - client: Optional[Union["HTTPHandler", "AsyncHTTPHandler"]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - model: Optional[str] = None, - ) -> Union["LiteLLMBatch", Coroutine[Any, Any, "LiteLLMBatch"]]: + client: Union["HTTPHandler", "AsyncHTTPHandler"] | None = None, + timeout: float | httpx.Timeout | None = None, + model: str | None = None, + ) -> Union["LiteLLMBatch", Coroutine[None, None, "LiteLLMBatch"]]: """ Retrieve a batch using provider-specific configuration. """ @@ -3976,16 +3982,16 @@ class BaseLLMHTTPHandler: async def async_create_batch( self, - transformed_request: Union[bytes, str, dict], + transformed_request: bytes | str | dict, litellm_params: dict, provider_config: "BaseBatchesConfig", headers: dict, api_base: str, logging_obj: "LiteLLMLoggingObj", - client: Optional[Union["HTTPHandler", "AsyncHTTPHandler"]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, + client: Union["HTTPHandler", "AsyncHTTPHandler"] | None = None, + timeout: float | httpx.Timeout | None = None, create_batch_data: Optional["CreateBatchRequest"] = None, - model: Optional[str] = None, + model: str | None = None, ): """ Async version of create_batch @@ -4055,16 +4061,16 @@ class BaseLLMHTTPHandler: async def async_retrieve_batch( self, - transformed_request: Union[bytes, str, dict], + transformed_request: bytes | str | dict, litellm_params: dict, provider_config: "BaseBatchesConfig", headers: dict, - api_base: Optional[str], + api_base: str | None, logging_obj: "LiteLLMLoggingObj", - client: Optional[Union["HTTPHandler", "AsyncHTTPHandler"]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - batch_id: Optional[str] = None, - model: Optional[str] = None, + client: Union["HTTPHandler", "AsyncHTTPHandler"] | None = None, + timeout: float | httpx.Timeout | None = None, + batch_id: str | None = None, + model: str | None = None, ): """ Async version of retrieve_batch @@ -4137,14 +4143,14 @@ class BaseLLMHTTPHandler: responses_api_provider_config: BaseResponsesAPIConfig, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - custom_llm_provider: Optional[str], - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + custom_llm_provider: str | None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, shared_session: Optional["ClientSession"] = None, - ) -> Union[ResponsesAPIResponse, Coroutine[Any, Any, ResponsesAPIResponse]]: + ) -> ResponsesAPIResponse | Coroutine[None, None, ResponsesAPIResponse]: """ Async version of the responses API handler. Uses async HTTP client to make requests. @@ -4217,11 +4223,11 @@ class BaseLLMHTTPHandler: responses_api_provider_config: BaseResponsesAPIConfig, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - custom_llm_provider: Optional[str], - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + custom_llm_provider: str | None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, shared_session: Optional["ClientSession"] = None, ) -> ResponsesAPIResponse: @@ -4290,17 +4296,17 @@ class BaseLLMHTTPHandler: model: str, input: Union[str, "ResponseInputParam"], responses_api_provider_config: BaseResponsesAPIConfig, - response_api_optional_request_params: Dict, + response_api_optional_request_params: dict, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - custom_llm_provider: Optional[str], - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + custom_llm_provider: str | None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, shared_session: Optional["ClientSession"] = None, - ) -> Union[ResponsesAPIResponse, Coroutine[Any, Any, ResponsesAPIResponse]]: + ) -> ResponsesAPIResponse | Coroutine[None, None, ResponsesAPIResponse]: """ Handler for the compact responses API. """ @@ -4357,7 +4363,7 @@ class BaseLLMHTTPHandler: api_key=litellm_params.api_key, model=model, ) - body_kwargs: Dict[str, Any] = {"data": signed_body} if signed_body is not None else {"json": data} + body_kwargs: dict[str, Any] = {"data": signed_body} if signed_body is not None else {"json": data} ## LOGGING logging_obj.pre_call( @@ -4389,14 +4395,14 @@ class BaseLLMHTTPHandler: model: str, input: Union[str, "ResponseInputParam"], responses_api_provider_config: BaseResponsesAPIConfig, - response_api_optional_request_params: Dict, + response_api_optional_request_params: dict, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - custom_llm_provider: Optional[str], - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + custom_llm_provider: str | None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, shared_session: Optional["ClientSession"] = None, ) -> ResponsesAPIResponse: @@ -4448,7 +4454,7 @@ class BaseLLMHTTPHandler: api_key=litellm_params.api_key, model=model, ) - body_kwargs: Dict[str, Any] = {"data": signed_body} if signed_body is not None else {"json": data} + body_kwargs: dict[str, Any] = {"data": signed_body} if signed_body is not None else {"json": data} ## LOGGING logging_obj.pre_call( @@ -4483,9 +4489,9 @@ class BaseLLMHTTPHandler: headers: dict, logging_obj: LiteLLMLoggingObj, _is_async: bool = False, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - ) -> Union[OpenAIFileObject, Coroutine[Any, Any, OpenAIFileObject]]: + client: HTTPHandler | AsyncHTTPHandler | None = None, + timeout: float | httpx.Timeout | None = None, + ) -> OpenAIFileObject | Coroutine[None, None, OpenAIFileObject]: """ Retrieve file metadata by ID """ @@ -4550,8 +4556,8 @@ class BaseLLMHTTPHandler: litellm_params: dict, headers: dict, logging_obj: LiteLLMLoggingObj, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, + timeout: float | httpx.Timeout | None = None, ) -> OpenAIFileObject: """ Async retrieve file metadata by ID @@ -4607,9 +4613,9 @@ class BaseLLMHTTPHandler: headers: dict, logging_obj: LiteLLMLoggingObj, _is_async: bool = False, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - ) -> Union["FileDeleted", Coroutine[Any, Any, "FileDeleted"]]: + client: HTTPHandler | AsyncHTTPHandler | None = None, + timeout: float | httpx.Timeout | None = None, + ) -> Union["FileDeleted", Coroutine[None, None, "FileDeleted"]]: """ Delete a file by ID """ @@ -4674,8 +4680,8 @@ class BaseLLMHTTPHandler: litellm_params: dict, headers: dict, logging_obj: LiteLLMLoggingObj, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, + timeout: float | httpx.Timeout | None = None, ) -> "FileDeleted": """ Async delete a file by ID @@ -4725,15 +4731,15 @@ class BaseLLMHTTPHandler: def list_files( self, - purpose: Optional[str], + purpose: str | None, provider_config: BaseFilesConfig, litellm_params: dict, headers: dict, logging_obj: LiteLLMLoggingObj, _is_async: bool = False, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - ) -> Union[List[OpenAIFileObject], Coroutine[Any, Any, List[OpenAIFileObject]]]: + client: HTTPHandler | AsyncHTTPHandler | None = None, + timeout: float | httpx.Timeout | None = None, + ) -> list[OpenAIFileObject] | Coroutine[None, None, list[OpenAIFileObject]]: """ List all files """ @@ -4793,14 +4799,14 @@ class BaseLLMHTTPHandler: async def async_list_files( self, - purpose: Optional[str], + purpose: str | None, provider_config: BaseFilesConfig, litellm_params: dict, headers: dict, logging_obj: LiteLLMLoggingObj, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - ) -> List[OpenAIFileObject]: + client: HTTPHandler | AsyncHTTPHandler | None = None, + timeout: float | httpx.Timeout | None = None, + ) -> list[OpenAIFileObject]: """ Async list all files """ @@ -4855,9 +4861,9 @@ class BaseLLMHTTPHandler: headers: dict, logging_obj: LiteLLMLoggingObj, _is_async: bool = False, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - ) -> Union["HttpxBinaryResponseContent", Coroutine[Any, Any, "HttpxBinaryResponseContent"]]: + client: HTTPHandler | AsyncHTTPHandler | None = None, + timeout: float | httpx.Timeout | None = None, + ) -> Union["HttpxBinaryResponseContent", Coroutine[None, None, "HttpxBinaryResponseContent"]]: """ Retrieve file content by ID """ @@ -4929,8 +4935,8 @@ class BaseLLMHTTPHandler: litellm_params: dict, headers: dict, logging_obj: LiteLLMLoggingObj, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, + timeout: float | httpx.Timeout | None = None, ) -> "HttpxBinaryResponseContent": """ Async retrieve file content by ID @@ -4990,7 +4996,7 @@ class BaseLLMHTTPHandler: stream: bool, data: dict, fake_stream: bool, - ) -> Tuple[bool, dict]: + ) -> tuple[bool, dict]: """ Handles preparing a request when `fake_stream` is True. """ @@ -5001,14 +5007,14 @@ class BaseLLMHTTPHandler: return stream, data @staticmethod - def _get_agentic_loop_settings(kwargs: Dict) -> Tuple[int, int, List[str]]: + def _get_agentic_loop_settings(kwargs: dict) -> tuple[int, int, list[str]]: depth = int(kwargs.get("_agentic_loop_depth", 0) or 0) max_loops = int(kwargs.get("max_agentic_loops", 3) or 3) fingerprints = list(kwargs.get("_agentic_loop_fingerprints", []) or []) return depth, max(max_loops, 1), fingerprints @staticmethod - def _has_agentic_completion_hook(logging_obj: Any) -> bool: + def _has_agentic_completion_hook(logging_obj: LiteLLMLoggingObj) -> bool: """ True if any registered callback actually overrides ``async_should_run_agentic_loop`` (the gate every agentic hook goes @@ -5039,8 +5045,8 @@ class BaseLLMHTTPHandler: @staticmethod def _check_agentic_loop_safety( - tool_calls: Any, - fingerprints: List[str], + tool_calls: Mapping[str, object] | None, + fingerprints: list[str], depth: int, max_loops: int, model: str, @@ -5062,7 +5068,7 @@ class BaseLLMHTTPHandler: return fingerprint @staticmethod - def _fingerprint_agentic_tools(tools: Dict) -> str: + def _fingerprint_agentic_tools(tools: Mapping[str, object] | None) -> str: try: return json.dumps(tools, sort_keys=True, default=str) except Exception: @@ -5072,16 +5078,16 @@ class BaseLLMHTTPHandler: self, plan: AgenticLoopPlan, model: str, - messages: List[Dict], - anthropic_messages_optional_request_params: Dict, + messages: list[dict], + anthropic_messages_optional_request_params: dict, logging_obj: "LiteLLMLoggingObj", - kwargs: Dict, + kwargs: dict, depth: int, max_loops: int, - fingerprints: List[str], + fingerprints: list[str], fingerprint: str, stream: bool = False, - callback: Optional[Any] = None, + callback: Any | None = None, ) -> Any: from litellm.anthropic_interface import messages as anthropic_messages @@ -5101,7 +5107,7 @@ class BaseLLMHTTPHandler: max_tokens = patch.max_tokens if max_tokens is None: - max_tokens = cast(Optional[int], optional_params.pop("max_tokens", None)) + max_tokens = cast(int | None, optional_params.pop("max_tokens", None)) else: optional_params.pop("max_tokens", None) if max_tokens is None: @@ -5248,10 +5254,10 @@ class BaseLLMHTTPHandler: self, result: Any, model: str, - responses_api_provider_config: Any, + responses_api_provider_config: BaseResponsesAPIConfig, logging_obj: "LiteLLMLoggingObj", custom_llm_provider: str, - ) -> Any: + ) -> MockResponsesAPIStreamingIterator: """ Wrap a completed responses result as a synthetic stream. @@ -5260,10 +5266,6 @@ class BaseLLMHTTPHandler: """ import httpx - from litellm.responses.streaming_iterator import ( - MockResponsesAPIStreamingIterator, - ) - payload = result.model_dump() if hasattr(result, "model_dump") else result raw_response = httpx.Response(status_code=200, json=payload) return MockResponsesAPIStreamingIterator( @@ -5278,13 +5280,13 @@ class BaseLLMHTTPHandler: self, plan: AgenticLoopPlan, model: str, - messages: List[Dict], - optional_params: Dict, - kwargs: Dict, + messages: list[dict], + optional_params: dict, + kwargs: dict, custom_llm_provider: str, depth: int, max_loops: int, - fingerprints: List[str], + fingerprints: list[str], fingerprint: str, ) -> Any: patch = plan.request_patch or AgenticLoopRequestPatch() @@ -5330,10 +5332,10 @@ class BaseLLMHTTPHandler: def _maybe_wrap_in_fake_stream( self, - response: Any, + response: FakeStreamWrappableResponseT, logging_obj: Optional["LiteLLMLoggingObj"], api_surface: str, - ) -> Any: + ) -> Union[FakeStreamWrappableResponseT, "FakeAnthropicMessagesStreamIterator"]: """ If the original request was streaming but converted to non-streaming for WebSearch interception, wrap the dict response in a FakeAnthropicMessagesStreamIterator. @@ -5371,15 +5373,15 @@ class BaseLLMHTTPHandler: self, response: Any, model: str, - messages: List[Dict], + messages: list[dict], anthropic_messages_provider_config: "BaseAnthropicMessagesConfig", - anthropic_messages_optional_request_params: Dict, + anthropic_messages_optional_request_params: dict, logging_obj: "LiteLLMLoggingObj", stream: bool, custom_llm_provider: str, - kwargs: Dict, + kwargs: dict, api_surface: str = "anthropic_messages", - ) -> Optional[Any]: + ) -> Any | None: """ Call agentic completion hooks for all custom loggers (Anthropic Messages API). @@ -5402,7 +5404,7 @@ class BaseLLMHTTPHandler: continue should_run: bool = False - tool_calls: Any = None + tool_calls: Mapping[str, object] | None = None try: # First: Check if agentic loop should run. Wrap in try/except # to shield from buggy user callbacks — a callback crash should @@ -5542,13 +5544,13 @@ class BaseLLMHTTPHandler: self, response: Any, model: str, - messages: List[Dict], - optional_params: Dict, + messages: list[dict], + optional_params: dict, logging_obj: "LiteLLMLoggingObj", stream: bool, custom_llm_provider: str, - kwargs: Dict, - ) -> Optional[Any]: + kwargs: dict, + ) -> Any | None: """ Call agentic chat completion hooks for all custom loggers (Chat Completions API). @@ -5571,7 +5573,7 @@ class BaseLLMHTTPHandler: continue should_run: bool = False - tool_calls: Any = None + tool_calls: Mapping[str, object] | None = None try: ( should_run, @@ -5750,7 +5752,7 @@ class BaseLLMHTTPHandler: ) @staticmethod - def _append_query_params(url: str, query_params: Optional[RealtimeQueryParams]) -> str: + def _append_query_params(url: str, query_params: RealtimeQueryParams | None) -> str: """Append query_params to url, skipping keys already present in the URL.""" if not query_params: return url @@ -5795,7 +5797,7 @@ class BaseLLMHTTPHandler: ) if exc is not None ) - last_exc: Optional[BaseException] = None + last_exc: BaseException | None = None for _ in range(max_attempts): try: return await websockets_module.connect( @@ -5823,13 +5825,13 @@ class BaseLLMHTTPHandler: logging_obj: LiteLLMLoggingObj, provider_config: BaseRealtimeConfig, headers: dict, - api_base: Optional[str] = None, - api_key: Optional[str] = None, - client: Optional[Any] = None, - timeout: Optional[float] = None, - user_api_key_dict: Optional[Any] = None, - litellm_metadata: Optional[Dict[str, Any]] = None, - query_params: Optional[RealtimeQueryParams] = None, + api_base: str | None = None, + api_key: str | None = None, + client: Any | None = None, + timeout: float | None = None, + user_api_key_dict: Any | None = None, + litellm_metadata: dict[str, object] | None = None, + query_params: RealtimeQueryParams | None = None, ): import websockets from websockets.asyncio.client import ClientConnection @@ -5850,7 +5852,7 @@ class BaseLLMHTTPHandler: ssl_context.verify_mode = ssl.CERT_NONE backend_ws = await self._open_realtime_backend_ws(websockets, url, headers, ssl_context) async with backend_ws: - _request_data: Dict[str, Any] = {} + _request_data: dict[str, Any] = {} if litellm_metadata: _request_data["litellm_metadata"] = litellm_metadata realtime_streaming = RealTimeStreaming( @@ -5872,7 +5874,7 @@ class BaseLLMHTTPHandler: # auto-response disable can be folded into this one setup: Gemini # rejects a second setup, so a follow-up disable would be dropped # and the guardrail bypassed. - _session_config: Optional[str] = None + _session_config: str | None = None if provider_config.requires_session_configuration(): _session_config = provider_config.session_configuration_request(model) if _session_config: @@ -5922,14 +5924,14 @@ class BaseLLMHTTPHandler: self, api_base: str, api_key: str, - request_data: Dict[str, Any], + request_data: dict[str, Any], logging_obj: LiteLLMLoggingObj, - timeout: Union[float, httpx.Timeout], - provider_config: Optional[Any] = None, - model: Optional[str] = None, - extra_headers: Optional[Dict[str, Any]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - api_version: Optional[str] = None, + timeout: float | httpx.Timeout, + provider_config: Any | None = None, + model: str | None = None, + extra_headers: dict[str, str] | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, + api_version: str | None = None, ) -> httpx.Response: """ Forward POST /v1/realtime/client_secrets to upstream provider. @@ -5955,14 +5957,14 @@ class BaseLLMHTTPHandler: self, api_base: str, api_key: str, - request_data: Dict[str, Any], + request_data: dict[str, Any], logging_obj: LiteLLMLoggingObj, - timeout: Union[float, httpx.Timeout], - provider_config: Optional[Any] = None, - model: Optional[str] = None, - extra_headers: Optional[Dict[str, Any]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - api_version: Optional[str] = None, + timeout: float | httpx.Timeout, + provider_config: Any | None = None, + model: str | None = None, + extra_headers: dict[str, str] | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, + api_version: str | None = None, ) -> httpx.Response: """Forward POST /v1/realtime/transcription_sessions to upstream provider.""" return await self._async_realtime_session_post( @@ -5984,14 +5986,14 @@ class BaseLLMHTTPHandler: endpoint: Literal["client_secrets", "transcription_sessions"], api_base: str, api_key: str, - request_data: Dict[str, Any], + request_data: dict[str, Any], logging_obj: LiteLLMLoggingObj, - timeout: Union[float, httpx.Timeout], - provider_config: Optional[Any] = None, - model: Optional[str] = None, - extra_headers: Optional[Dict[str, Any]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - api_version: Optional[str] = None, + timeout: float | httpx.Timeout, + provider_config: Any | None = None, + model: str | None = None, + extra_headers: dict[str, str] | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, + api_version: str | None = None, ) -> httpx.Response: """ Shared POST flow for the realtime HTTP session endpoints @@ -6014,7 +6016,7 @@ class BaseLLMHTTPHandler: ) else: url = provider_config.get_complete_url(api_base=api_base, model=model or "", api_version=api_version) - headers: Dict[str, Any] = provider_config.validate_environment( + headers: dict[str, Any] = provider_config.validate_environment( headers={}, model=model or "", api_key=api_key ) else: @@ -6058,13 +6060,13 @@ class BaseLLMHTTPHandler: openai_ephemeral_key: str, sdp_body: bytes, logging_obj: LiteLLMLoggingObj, - timeout: Union[float, httpx.Timeout], - provider_config: Optional[Any] = None, - model: Optional[str] = None, - session_config: Optional[Dict[str, Any]] = None, - extra_headers: Optional[Dict[str, Any]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - api_version: Optional[str] = None, + timeout: float | httpx.Timeout, + provider_config: Any | None = None, + model: str | None = None, + session_config: dict[str, object] | None = None, + extra_headers: dict[str, str] | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, + api_version: str | None = None, ) -> httpx.Response: """ Forward POST /v1/realtime/calls (SDP exchange) to upstream provider. @@ -6085,7 +6087,7 @@ class BaseLLMHTTPHandler: if provider_config is not None: url = provider_config.get_realtime_calls_url(api_base=api_base, model=model or "", api_version=api_version) - headers: Dict[str, Any] = provider_config.get_realtime_calls_headers(ephemeral_key=openai_ephemeral_key) + headers: dict[str, Any] = provider_config.get_realtime_calls_headers(ephemeral_key=openai_ephemeral_key) else: url = f"{api_base.rstrip('/')}/v1/realtime/calls" headers = { @@ -6139,14 +6141,14 @@ class BaseLLMHTTPHandler: model: str, websocket: Any, logging_obj: LiteLLMLoggingObj, - responses_api_provider_config: Optional[BaseResponsesAPIConfig], - api_base: Optional[str] = None, - api_key: Optional[str] = None, - timeout: Optional[float] = None, - user_api_key_dict: Optional[Any] = None, - litellm_metadata: Optional[Dict[str, Any]] = None, - custom_llm_provider: Optional[str] = None, - first_message: Optional[str] = None, + responses_api_provider_config: BaseResponsesAPIConfig | None, + api_base: str | None = None, + api_key: str | None = None, + timeout: float | None = None, + user_api_key_dict: Any | None = None, + litellm_metadata: dict[str, object] | None = None, + custom_llm_provider: str | None = None, + first_message: str | None = None, **kwargs: Any, ): """ @@ -6253,7 +6255,7 @@ class BaseLLMHTTPHandler: yield backend async with _backend_connection() as backend_ws: - _request_data: Dict[str, Any] = {} + _request_data: dict[str, Any] = {} if litellm_metadata: _request_data["litellm_metadata"] = litellm_metadata @@ -6316,24 +6318,21 @@ class BaseLLMHTTPHandler: def image_edit_handler( self, model: str, - image: Any, - prompt: Optional[str], + image: Any, # any-ok: real callers pass either a single FileTypes or a List[FileTypes] + prompt: str | None, image_edit_provider_config: BaseImageEditConfig, - image_edit_optional_request_params: Dict, + image_edit_optional_request_params: dict, custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - timeout: Union[float, httpx.Timeout], - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + timeout: float | httpx.Timeout, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, fake_stream: bool = False, - litellm_metadata: Optional[Dict[str, Any]] = None, - ) -> Union[ - ImageResponse, - Coroutine[Any, Any, ImageResponse], - ]: + litellm_metadata: dict[str, object] | None = None, + ) -> ImageResponse | Coroutine[None, None, ImageResponse]: """ Handles image edit requests. @@ -6436,19 +6435,19 @@ class BaseLLMHTTPHandler: async def async_image_edit_handler( self, model: str, - image: FileTypes, - prompt: Optional[str], + image: Any, # any-ok: real callers pass either a single FileTypes or a List[FileTypes] + prompt: str | None, image_edit_provider_config: BaseImageEditConfig, - image_edit_optional_request_params: Dict, + image_edit_optional_request_params: dict, custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - timeout: Union[float, httpx.Timeout], - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + timeout: float | httpx.Timeout, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, fake_stream: bool = False, - litellm_metadata: Optional[Dict[str, Any]] = None, + litellm_metadata: dict[str, object] | None = None, ) -> ImageResponse: """ Async version of the image edit handler. @@ -6537,22 +6536,19 @@ class BaseLLMHTTPHandler: model: str, prompt: str, image_generation_provider_config: BaseImageGenerationConfig, - image_generation_optional_request_params: Dict, + image_generation_optional_request_params: dict, custom_llm_provider: str, - litellm_params: Dict, + litellm_params: dict, logging_obj: LiteLLMLoggingObj, - timeout: Union[float, httpx.Timeout], - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + timeout: float | httpx.Timeout, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, fake_stream: bool = False, - litellm_metadata: Optional[Dict[str, Any]] = None, - api_key: Optional[str] = None, - ) -> Union[ - ImageResponse, - Coroutine[Any, Any, ImageResponse], - ]: + litellm_metadata: dict[str, object] | None = None, + api_key: str | None = None, + ) -> ImageResponse | Coroutine[None, None, ImageResponse]: """ Handles image generation requests. When _is_async=True, returns a coroutine instead of making the call directly. @@ -6664,17 +6660,17 @@ class BaseLLMHTTPHandler: model: str, prompt: str, image_generation_provider_config: BaseImageGenerationConfig, - image_generation_optional_request_params: Dict, + image_generation_optional_request_params: dict, custom_llm_provider: str, - litellm_params: Dict, + litellm_params: dict, logging_obj: LiteLLMLoggingObj, - timeout: Union[float, httpx.Timeout], - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + timeout: float | httpx.Timeout, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, fake_stream: bool = False, - litellm_metadata: Optional[Dict[str, Any]] = None, - api_key: Optional[str] = None, + litellm_metadata: dict[str, object] | None = None, + api_key: str | None = None, ) -> ImageResponse: """ Async version of the image generation handler. @@ -6772,22 +6768,19 @@ class BaseLLMHTTPHandler: model: str, prompt: str, video_generation_provider_config: BaseVideoConfig, - video_generation_optional_request_params: Dict, + video_generation_optional_request_params: dict, custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - timeout: Union[float, httpx.Timeout], - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + timeout: float | httpx.Timeout, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, fake_stream: bool = False, - litellm_metadata: Optional[Dict[str, Any]] = None, - api_key: Optional[str] = None, - ) -> Union[ - VideoObject, - Coroutine[Any, Any, VideoObject], - ]: + litellm_metadata: dict[str, object] | None = None, + api_key: str | None = None, + ) -> VideoObject | Coroutine[None, None, VideoObject]: """ Handles video generation requests. When _is_async=True, returns a coroutine instead of making the call directly. @@ -6896,17 +6889,17 @@ class BaseLLMHTTPHandler: model: str, prompt: str, video_generation_provider_config: "BaseVideoConfig", - video_generation_optional_request_params: Dict, + video_generation_optional_request_params: dict, custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - timeout: Union[float, httpx.Timeout], - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + timeout: float | httpx.Timeout, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, fake_stream: bool = False, - litellm_metadata: Optional[Dict[str, Any]] = None, - api_key: Optional[str] = None, + litellm_metadata: dict[str, object] | None = None, + api_key: str | None = None, ) -> VideoObject: """ Async version of the video generation handler. @@ -7000,13 +6993,13 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - timeout: Union[float, httpx.Timeout], - extra_headers: Optional[Dict[str, Any]] = None, - api_key: Optional[str] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + timeout: float | httpx.Timeout, + extra_headers: dict[str, str] | None = None, + api_key: str | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, - variant: Optional[str] = None, - ) -> Union[bytes, Coroutine[Any, Any, bytes]]: + variant: str | None = None, + ) -> bytes | Coroutine[None, None, bytes]: """ Handle video content download requests. """ @@ -7090,11 +7083,11 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - timeout: Union[float, httpx.Timeout], - extra_headers: Optional[Dict[str, Any]] = None, - api_key: Optional[str] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - variant: Optional[str] = None, + timeout: float | httpx.Timeout, + extra_headers: dict[str, str] | None = None, + api_key: str | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, + variant: str | None = None, ) -> bytes: """ Async version of the video content download handler. @@ -7169,12 +7162,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params, logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[float] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + timeout: float | None = None, _is_async: bool = False, client=None, - api_key: Optional[str] = None, + api_key: str | None = None, ): """ Handler for video remix requests. @@ -7268,11 +7261,11 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params, logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[float] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + timeout: float | None = None, client=None, - api_key: Optional[str] = None, + api_key: str | None = None, ): """ Async version of the video remix handler. @@ -7351,11 +7344,11 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params, logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[float] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | None = None, _is_async: bool = False, client=None, - api_key: Optional[str] = None, + api_key: str | None = None, ): if _is_async: return self.async_video_create_character_handler( @@ -7435,10 +7428,10 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params, logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[float] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | None = None, client=None, - api_key: Optional[str] = None, + api_key: str | None = None, ): if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( @@ -7506,11 +7499,11 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params, logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[float] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | None = None, _is_async: bool = False, client=None, - api_key: Optional[str] = None, + api_key: str | None = None, ): if _is_async: return self.async_video_get_character_handler( @@ -7575,10 +7568,10 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params, logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[float] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | None = None, client=None, - api_key: Optional[str] = None, + api_key: str | None = None, ): if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( @@ -7634,12 +7627,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params, logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[float] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + timeout: float | None = None, _is_async: bool = False, client=None, - api_key: Optional[str] = None, + api_key: str | None = None, ): if _is_async: return self.async_video_edit_handler( @@ -7743,11 +7736,11 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params, logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[float] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + timeout: float | None = None, client=None, - api_key: Optional[str] = None, + api_key: str | None = None, ): if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( @@ -7840,12 +7833,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params, logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[float] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + timeout: float | None = None, _is_async: bool = False, client=None, - api_key: Optional[str] = None, + api_key: str | None = None, ): if _is_async: return self.async_video_extension_handler( @@ -7929,11 +7922,11 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params, logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[float] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + timeout: float | None = None, client=None, - api_key: Optional[str] = None, + api_key: str | None = None, ): if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( @@ -7997,19 +7990,19 @@ class BaseLLMHTTPHandler: def video_list_handler( self, - after: Optional[str], - limit: Optional[int], - order: Optional[str], + after: str | None, + limit: int | None, + order: str | None, video_list_provider_config, custom_llm_provider: str, litellm_params, logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - timeout: Optional[float] = None, + extra_headers: dict[str, str] | None = None, + extra_query: dict[str, str] | None = None, + timeout: float | None = None, _is_async: bool = False, client=None, - api_key: Optional[str] = None, + api_key: str | None = None, ): """ Handler for video list requests. @@ -8051,18 +8044,18 @@ class BaseLLMHTTPHandler: async def async_video_list_handler( self, - after: Optional[str], - limit: Optional[int], - order: Optional[str], + after: str | None, + limit: int | None, + order: str | None, video_list_provider_config: BaseVideoConfig, custom_llm_provider: str, litellm_params, logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - timeout: Optional[float] = None, + extra_headers: dict[str, str] | None = None, + extra_query: dict[str, str] | None = None, + timeout: float | None = None, client=None, - api_key: Optional[str] = None, + api_key: str | None = None, ): """ Async version of the video list handler. @@ -8139,10 +8132,10 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params, logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[float] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | None = None, client=None, - api_key: Optional[str] = None, + api_key: str | None = None, ): """ Async version of the video delete handler. @@ -8215,12 +8208,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params, logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[float] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + timeout: float | None = None, _is_async: bool = False, client=None, - api_key: Optional[str] = None, + api_key: str | None = None, ): """ Handler for video status requests. @@ -8320,11 +8313,11 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params, logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[float] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + timeout: float | None = None, client=None, - api_key: Optional[str] = None, + api_key: str | None = None, ): """ Async version of the video status handler. @@ -8406,15 +8399,15 @@ class BaseLLMHTTPHandler: def container_create_handler( self, name: str, - container_create_request_params: Dict, + container_create_request_params: dict, container_provider_config: "BaseContainerConfig", litellm_params: GenericLiteLLMParams, logging_obj: "LiteLLMLoggingObj", - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Union[float, httpx.Timeout] = 600, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout = 600, _is_async: bool = False, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - ) -> Union["ContainerObject", Coroutine[Any, Any, "ContainerObject"]]: + client: HTTPHandler | AsyncHTTPHandler | None = None, + ) -> Union["ContainerObject", Coroutine[None, None, "ContainerObject"]]: if _is_async: # Return the async coroutine if called with _is_async=True return self.async_container_create_handler( @@ -8493,13 +8486,13 @@ class BaseLLMHTTPHandler: async def async_container_create_handler( self, name: str, - container_create_request_params: Dict, + container_create_request_params: dict, container_provider_config: "BaseContainerConfig", litellm_params: GenericLiteLLMParams, logging_obj: "LiteLLMLoggingObj", - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Union[float, httpx.Timeout] = 600, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout = 600, + client: HTTPHandler | AsyncHTTPHandler | None = None, ) -> "ContainerObject": # For async calls, use async HTTP client if client is None or not isinstance(client, AsyncHTTPHandler): @@ -8571,15 +8564,15 @@ class BaseLLMHTTPHandler: container_provider_config: "BaseContainerConfig", litellm_params: GenericLiteLLMParams, logging_obj: "LiteLLMLoggingObj", - after: Optional[str] = None, - limit: Optional[int] = None, - order: Optional[str] = None, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - timeout: Union[float, httpx.Timeout] = 600, + after: str | None = None, + limit: int | None = None, + order: str | None = None, + extra_headers: dict[str, str] | None = None, + extra_query: dict[str, str] | None = None, + timeout: float | httpx.Timeout = 600, _is_async: bool = False, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - ) -> Union["ContainerListResponse", Coroutine[Any, Any, "ContainerListResponse"]]: + client: HTTPHandler | AsyncHTTPHandler | None = None, + ) -> Union["ContainerListResponse", Coroutine[None, None, "ContainerListResponse"]]: if _is_async: # Return the async coroutine if called with _is_async=True return self.async_container_list_handler( @@ -8661,13 +8654,13 @@ class BaseLLMHTTPHandler: container_provider_config: "BaseContainerConfig", litellm_params: GenericLiteLLMParams, logging_obj: "LiteLLMLoggingObj", - after: Optional[str] = None, - limit: Optional[int] = None, - order: Optional[str] = None, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - timeout: Union[float, httpx.Timeout] = 600, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + after: str | None = None, + limit: int | None = None, + order: str | None = None, + extra_headers: dict[str, str] | None = None, + extra_query: dict[str, str] | None = None, + timeout: float | httpx.Timeout = 600, + client: HTTPHandler | AsyncHTTPHandler | None = None, ) -> "ContainerListResponse": # For async calls, use async HTTP client if client is None or not isinstance(client, AsyncHTTPHandler): @@ -8739,12 +8732,12 @@ class BaseLLMHTTPHandler: container_provider_config: "BaseContainerConfig", litellm_params: GenericLiteLLMParams, logging_obj: "LiteLLMLoggingObj", - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - timeout: Union[float, httpx.Timeout] = 600, + extra_headers: dict[str, str] | None = None, + extra_query: dict[str, str] | None = None, + timeout: float | httpx.Timeout = 600, _is_async: bool = False, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - ) -> Union["ContainerObject", Coroutine[Any, Any, "ContainerObject"]]: + client: HTTPHandler | AsyncHTTPHandler | None = None, + ) -> Union["ContainerObject", Coroutine[None, None, "ContainerObject"]]: if _is_async: # Return the async coroutine if called with _is_async=True return self.async_container_retrieve_handler( @@ -8827,10 +8820,10 @@ class BaseLLMHTTPHandler: container_provider_config: "BaseContainerConfig", litellm_params: GenericLiteLLMParams, logging_obj: "LiteLLMLoggingObj", - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - timeout: Union[float, httpx.Timeout] = 600, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + extra_query: dict[str, str] | None = None, + timeout: float | httpx.Timeout = 600, + client: HTTPHandler | AsyncHTTPHandler | None = None, ) -> "ContainerObject": # For async calls, use async HTTP client if client is None or not isinstance(client, AsyncHTTPHandler): @@ -8904,12 +8897,12 @@ class BaseLLMHTTPHandler: container_provider_config: "BaseContainerConfig", litellm_params: GenericLiteLLMParams, logging_obj: "LiteLLMLoggingObj", - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - timeout: Union[float, httpx.Timeout] = 600, + extra_headers: dict[str, str] | None = None, + extra_query: dict[str, str] | None = None, + timeout: float | httpx.Timeout = 600, _is_async: bool = False, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - ) -> Union["DeleteContainerResult", Coroutine[Any, Any, "DeleteContainerResult"]]: + client: HTTPHandler | AsyncHTTPHandler | None = None, + ) -> Union["DeleteContainerResult", Coroutine[None, None, "DeleteContainerResult"]]: if _is_async: # Return the async coroutine if called with _is_async=True return self.async_container_delete_handler( @@ -8992,10 +8985,10 @@ class BaseLLMHTTPHandler: container_provider_config: "BaseContainerConfig", litellm_params: GenericLiteLLMParams, logging_obj: "LiteLLMLoggingObj", - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - timeout: Union[float, httpx.Timeout] = 600, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + extra_query: dict[str, str] | None = None, + timeout: float | httpx.Timeout = 600, + client: HTTPHandler | AsyncHTTPHandler | None = None, ) -> "DeleteContainerResult": # For async calls, use async HTTP client if client is None or not isinstance(client, AsyncHTTPHandler): @@ -9069,15 +9062,15 @@ class BaseLLMHTTPHandler: container_provider_config: "BaseContainerConfig", litellm_params: GenericLiteLLMParams, logging_obj: "LiteLLMLoggingObj", - after: Optional[str] = None, - limit: Optional[int] = None, - order: Optional[str] = None, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - timeout: Union[float, httpx.Timeout] = 600, + after: str | None = None, + limit: int | None = None, + order: str | None = None, + extra_headers: dict[str, str] | None = None, + extra_query: dict[str, str] | None = None, + timeout: float | httpx.Timeout = 600, _is_async: bool = False, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - ) -> Union["ContainerFileListResponse", Coroutine[Any, Any, "ContainerFileListResponse"]]: + client: HTTPHandler | AsyncHTTPHandler | None = None, + ) -> Union["ContainerFileListResponse", Coroutine[None, None, "ContainerFileListResponse"]]: if _is_async: return self.async_container_file_list_handler( container_id=container_id, @@ -9161,13 +9154,13 @@ class BaseLLMHTTPHandler: container_provider_config: "BaseContainerConfig", litellm_params: GenericLiteLLMParams, logging_obj: "LiteLLMLoggingObj", - after: Optional[str] = None, - limit: Optional[int] = None, - order: Optional[str] = None, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - timeout: Union[float, httpx.Timeout] = 600, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + after: str | None = None, + limit: int | None = None, + order: str | None = None, + extra_headers: dict[str, str] | None = None, + extra_query: dict[str, str] | None = None, + timeout: float | httpx.Timeout = 600, + client: HTTPHandler | AsyncHTTPHandler | None = None, ) -> "ContainerFileListResponse": # For async calls, use async HTTP client if client is None or not isinstance(client, AsyncHTTPHandler): @@ -9241,11 +9234,11 @@ class BaseLLMHTTPHandler: container_provider_config: "BaseContainerConfig", litellm_params: GenericLiteLLMParams, logging_obj: "LiteLLMLoggingObj", - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Union[float, httpx.Timeout] = 600, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout = 600, _is_async: bool = False, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - ) -> Union[bytes, Coroutine[Any, Any, bytes]]: + client: HTTPHandler | AsyncHTTPHandler | None = None, + ) -> bytes | Coroutine[None, None, bytes]: if _is_async: return self.async_container_file_content_handler( container_id=container_id, @@ -9327,9 +9320,9 @@ class BaseLLMHTTPHandler: container_provider_config: "BaseContainerConfig", litellm_params: GenericLiteLLMParams, logging_obj: "LiteLLMLoggingObj", - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Union[float, httpx.Timeout] = 600, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout = 600, + client: HTTPHandler | AsyncHTTPHandler | None = None, ) -> bytes: # For async calls, use async HTTP client if client is None or not isinstance(client, AsyncHTTPHandler): @@ -9400,16 +9393,16 @@ class BaseLLMHTTPHandler: async def async_vector_store_search_handler( self, vector_store_id: str, - query: Union[str, List[str]], + query: str | list[str], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, vector_store_provider_config: BaseVectorStoreConfig, custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, ) -> VectorStoreSearchResponse: if client is None or not isinstance(client, AsyncHTTPHandler): @@ -9459,7 +9452,7 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), extra_body=extra_body, ) - all_optional_params: Dict[str, Any] = dict(litellm_params) + all_optional_params: dict[str, Any] = dict(litellm_params) all_optional_params.update(vector_store_search_optional_params or {}) headers, signed_json_body = vector_store_provider_config.sign_request( headers=headers, @@ -9498,18 +9491,18 @@ class BaseLLMHTTPHandler: def vector_store_search_handler( self, vector_store_id: str, - query: Union[str, List[str]], + query: str | list[str], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, vector_store_provider_config: BaseVectorStoreConfig, custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, - ) -> Union[VectorStoreSearchResponse, Coroutine[Any, Any, VectorStoreSearchResponse]]: + ) -> VectorStoreSearchResponse | Coroutine[None, None, VectorStoreSearchResponse]: if _is_async: return self.async_vector_store_search_handler( vector_store_id=vector_store_id, @@ -9555,7 +9548,7 @@ class BaseLLMHTTPHandler: extra_body=extra_body, ) - all_optional_params: Dict[str, Any] = dict(litellm_params) + all_optional_params: dict[str, Any] = dict(litellm_params) all_optional_params.update(vector_store_search_optional_params or {}) headers, signed_json_body = vector_store_provider_config.sign_request( @@ -9598,10 +9591,10 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, ) -> VectorStoreCreateResponse: if client is None or not isinstance(client, AsyncHTTPHandler): @@ -9658,12 +9651,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, - ) -> Union[VectorStoreCreateResponse, Coroutine[Any, Any, VectorStoreCreateResponse]]: + ) -> VectorStoreCreateResponse | Coroutine[None, None, VectorStoreCreateResponse]: if _is_async: return self.async_vector_store_create_handler( vector_store_create_optional_params=vector_store_create_optional_params, @@ -9728,10 +9721,10 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, ) -> VectorStoreCreateResponse: if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( @@ -9781,12 +9774,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, - ) -> Union[VectorStoreCreateResponse, Coroutine[Any, Any, VectorStoreCreateResponse]]: + ) -> VectorStoreCreateResponse | Coroutine[None, None, VectorStoreCreateResponse]: if _is_async: return self.async_vector_store_retrieve_handler( vector_store_id=vector_store_id, @@ -9840,18 +9833,18 @@ class BaseLLMHTTPHandler: async def async_vector_store_list_handler( self, - after: Optional[str], - before: Optional[str], - limit: Optional[int], - order: Optional[str], + after: str | None, + before: str | None, + limit: int | None, + order: str | None, vector_store_provider_config: BaseVectorStoreConfig, custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, ): if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( @@ -9875,7 +9868,7 @@ class BaseLLMHTTPHandler: url = api_base - params: Dict[str, Any] = {} + params: dict[str, Any] = {} if after is not None: params["after"] = after if before is not None: @@ -9904,18 +9897,18 @@ class BaseLLMHTTPHandler: def vector_store_list_handler( self, - after: Optional[str], - before: Optional[str], - limit: Optional[int], - order: Optional[str], + after: str | None, + before: str | None, + limit: int | None, + order: str | None, vector_store_provider_config: BaseVectorStoreConfig, custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, ): if _is_async: @@ -9953,7 +9946,7 @@ class BaseLLMHTTPHandler: url = api_base - params: Dict[str, Any] = {} + params: dict[str, Any] = {} if after is not None: params["after"] = after if before is not None: @@ -9988,10 +9981,10 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, ) -> VectorStoreCreateResponse: if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( @@ -10016,7 +10009,7 @@ class BaseLLMHTTPHandler: encoded_vector_store_id = encode_url_path_segment(vector_store_id, field_name="vector_store_id") url = f"{api_base}/{encoded_vector_store_id}" - request_body: Dict[str, Any] = dict(vector_store_update_optional_params) + request_body: dict[str, Any] = dict(vector_store_update_optional_params) # Clean metadata to only include string values (OpenAI requirement) if "metadata" in request_body and request_body["metadata"] is not None: @@ -10054,12 +10047,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, - ) -> Union[VectorStoreCreateResponse, Coroutine[Any, Any, VectorStoreCreateResponse]]: + ) -> VectorStoreCreateResponse | Coroutine[None, None, VectorStoreCreateResponse]: if _is_async: return self.async_vector_store_update_handler( vector_store_id=vector_store_id, @@ -10094,7 +10087,7 @@ class BaseLLMHTTPHandler: encoded_vector_store_id = encode_url_path_segment(vector_store_id, field_name="vector_store_id") url = f"{api_base}/{encoded_vector_store_id}" - request_body: Dict[str, Any] = dict(vector_store_update_optional_params) + request_body: dict[str, Any] = dict(vector_store_update_optional_params) # Clean metadata to only include string values (OpenAI requirement) if "metadata" in request_body and request_body["metadata"] is not None: @@ -10131,10 +10124,10 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, ): if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( @@ -10182,10 +10175,10 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, ): if _is_async: @@ -10249,10 +10242,10 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, ) -> VectorStoreFileObject: if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( @@ -10314,12 +10307,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, - ) -> Union[VectorStoreFileObject, Coroutine[Any, Any, VectorStoreFileObject]]: + ) -> VectorStoreFileObject | Coroutine[None, None, VectorStoreFileObject]: if _is_async: return self.async_vector_store_file_create_handler( vector_store_id=vector_store_id, @@ -10391,10 +10384,10 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + extra_query: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, ) -> VectorStoreFileListResponse: if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( @@ -10455,12 +10448,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + extra_query: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, - ) -> Union[VectorStoreFileListResponse, Coroutine[Any, Any, VectorStoreFileListResponse]]: + ) -> VectorStoreFileListResponse | Coroutine[None, None, VectorStoreFileListResponse]: if _is_async: return self.async_vector_store_file_list_handler( vector_store_id=vector_store_id, @@ -10531,9 +10524,9 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, ) -> VectorStoreFileObject: if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( @@ -10590,11 +10583,11 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, - ) -> Union[VectorStoreFileObject, Coroutine[Any, Any, VectorStoreFileObject]]: + ) -> VectorStoreFileObject | Coroutine[None, None, VectorStoreFileObject]: if _is_async: return self.async_vector_store_file_retrieve_handler( vector_store_id=vector_store_id, @@ -10660,9 +10653,9 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, ) -> VectorStoreFileContentResponse: if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( @@ -10721,14 +10714,11 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, - ) -> Union[ - VectorStoreFileContentResponse, - Coroutine[Any, Any, VectorStoreFileContentResponse], - ]: + ) -> VectorStoreFileContentResponse | Coroutine[None, None, VectorStoreFileContentResponse]: if _is_async: return self.async_vector_store_file_content_handler( vector_store_id=vector_store_id, @@ -10797,10 +10787,10 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, ) -> VectorStoreFileObject: if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( @@ -10863,12 +10853,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, - ) -> Union[VectorStoreFileObject, Coroutine[Any, Any, VectorStoreFileObject]]: + ) -> VectorStoreFileObject | Coroutine[None, None, VectorStoreFileObject]: if _is_async: return self.async_vector_store_file_update_handler( vector_store_id=vector_store_id, @@ -10941,9 +10931,9 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, ) -> VectorStoreFileDeleteResponse: if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( @@ -11000,14 +10990,11 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, - ) -> Union[ - VectorStoreFileDeleteResponse, - Coroutine[Any, Any, VectorStoreFileDeleteResponse], - ]: + ) -> VectorStoreFileDeleteResponse | Coroutine[None, None, VectorStoreFileDeleteResponse]: if _is_async: return self.async_vector_store_file_delete_handler( vector_store_id=vector_store_id, @@ -11072,19 +11059,19 @@ class BaseLLMHTTPHandler: model: str, contents: Any, generate_content_provider_config: BaseGoogleGenAIGenerateContentConfig, - generate_content_config_dict: Dict, + generate_content_config_dict: dict, tools: Any, custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, _is_async: bool = False, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, stream: bool = False, - litellm_metadata: Optional[Dict[str, Any]] = None, - system_instruction: Optional[Any] = None, + litellm_metadata: dict[str, object] | None = None, + system_instruction: Any | None = None, ) -> Any: """ Handles Google GenAI generate content requests. @@ -11204,18 +11191,18 @@ class BaseLLMHTTPHandler: model: str, contents: Any, generate_content_provider_config: BaseGoogleGenAIGenerateContentConfig, - generate_content_config_dict: Dict, + generate_content_config_dict: dict, tools: Any, custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[AsyncHTTPHandler] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, + client: AsyncHTTPHandler | None = None, stream: bool = False, - litellm_metadata: Optional[Dict[str, Any]] = None, - system_instruction: Optional[Any] = None, + litellm_metadata: dict[str, object] | None = None, + system_instruction: Any | None = None, ) -> Any: """ Async version of the generate content handler. @@ -11321,19 +11308,19 @@ class BaseLLMHTTPHandler: self, model: str, input: str, - voice: Optional[str], + voice: str | None, text_to_speech_provider_config: BaseTextToSpeechConfig, - text_to_speech_optional_params: Dict, + text_to_speech_optional_params: dict, custom_llm_provider: str, - litellm_params: Dict, + litellm_params: dict, logging_obj: LiteLLMLoggingObj, - timeout: Union[float, httpx.Timeout], - extra_headers: Optional[Dict[str, Any]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + timeout: float | httpx.Timeout, + extra_headers: dict[str, str] | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, ) -> Union[ "HttpxBinaryResponseContent", - Coroutine[Any, Any, "HttpxBinaryResponseContent"], + Coroutine[None, None, "HttpxBinaryResponseContent"], ]: """ Handles text-to-speech requests. @@ -11436,15 +11423,15 @@ class BaseLLMHTTPHandler: self, model: str, input: str, - voice: Optional[str], + voice: str | None, text_to_speech_provider_config: BaseTextToSpeechConfig, - text_to_speech_optional_params: Dict, + text_to_speech_optional_params: dict, custom_llm_provider: str, - litellm_params: Dict, + litellm_params: dict, logging_obj: LiteLLMLoggingObj, - timeout: Union[float, httpx.Timeout], - extra_headers: Optional[Dict[str, Any]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + timeout: float | httpx.Timeout, + extra_headers: dict[str, str] | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, ) -> "HttpxBinaryResponseContent": """ Async version of the text-to-speech handler. @@ -11537,9 +11524,9 @@ class BaseLLMHTTPHandler: def _prepare_skill_multipart_request( self, - request_body: Dict, + request_body: dict, headers: dict, - ) -> tuple[Optional[Dict], Optional[list]]: + ) -> tuple[dict | None, list | None]: """ Helper to prepare multipart/form-data request for skills API. @@ -11570,17 +11557,17 @@ class BaseLLMHTTPHandler: def create_skill_handler( self, url: str, - request_body: Dict, + request_body: dict, skills_api_provider_config: "BaseSkillsAPIConfig", custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, shared_session: Optional["ClientSession"] = None, - ) -> Union["Skill", Coroutine[Any, Any, "Skill"]]: + ) -> Union["Skill", Coroutine[None, None, "Skill"]]: """Create a skill""" if _is_async: return self.async_create_skill_handler( @@ -11636,14 +11623,14 @@ class BaseLLMHTTPHandler: async def async_create_skill_handler( self, url: str, - request_body: Dict, + request_body: dict, skills_api_provider_config: "BaseSkillsAPIConfig", custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, shared_session: Optional["ClientSession"] = None, ) -> "Skill": """Async create a skill""" @@ -11692,17 +11679,17 @@ class BaseLLMHTTPHandler: def list_skills_handler( self, url: str, - query_params: Dict, + query_params: dict, skills_api_provider_config: "BaseSkillsAPIConfig", custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, shared_session: Optional["ClientSession"] = None, - ) -> Union["ListSkillsResponse", Coroutine[Any, Any, "ListSkillsResponse"]]: + ) -> Union["ListSkillsResponse", Coroutine[None, None, "ListSkillsResponse"]]: """List skills""" if _is_async: return self.async_list_skills_handler( @@ -11751,14 +11738,14 @@ class BaseLLMHTTPHandler: async def async_list_skills_handler( self, url: str, - query_params: Dict, + query_params: dict, skills_api_provider_config: "BaseSkillsAPIConfig", custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, shared_session: Optional["ClientSession"] = None, ) -> "ListSkillsResponse": """Async list skills""" @@ -11802,12 +11789,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, shared_session: Optional["ClientSession"] = None, - ) -> Union["Skill", Coroutine[Any, Any, "Skill"]]: + ) -> Union["Skill", Coroutine[None, None, "Skill"]]: """Get a skill""" if _is_async: return self.async_get_skill_handler( @@ -11858,9 +11845,9 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, shared_session: Optional["ClientSession"] = None, ) -> "Skill": """Async get a skill""" @@ -11903,12 +11890,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, shared_session: Optional["ClientSession"] = None, - ) -> Union["DeleteSkillResponse", Coroutine[Any, Any, "DeleteSkillResponse"]]: + ) -> Union["DeleteSkillResponse", Coroutine[None, None, "DeleteSkillResponse"]]: """Delete a skill""" if _is_async: return self.async_delete_skill_handler( @@ -11959,9 +11946,9 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, shared_session: Optional["ClientSession"] = None, ) -> "DeleteSkillResponse": """Async delete a skill""" @@ -12004,17 +11991,17 @@ class BaseLLMHTTPHandler: def create_eval_handler( self, url: str, - request_body: Dict, + request_body: dict, evals_api_provider_config: "BaseEvalsAPIConfig", custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, shared_session: Optional["ClientSession"] = None, - ) -> Union["Eval", Coroutine[Any, Any, "Eval"]]: + ) -> Union["Eval", Coroutine[None, None, "Eval"]]: """Create an eval""" if _is_async: return self.async_create_eval_handler( @@ -12063,14 +12050,14 @@ class BaseLLMHTTPHandler: async def async_create_eval_handler( self, url: str, - request_body: Dict, + request_body: dict, evals_api_provider_config: "BaseEvalsAPIConfig", custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, shared_session: Optional["ClientSession"] = None, ) -> "Eval": """Async create an eval""" @@ -12110,17 +12097,17 @@ class BaseLLMHTTPHandler: def list_evals_handler( self, url: str, - query_params: Dict, + query_params: dict, evals_api_provider_config: "BaseEvalsAPIConfig", custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, shared_session: Optional["ClientSession"] = None, - ) -> Union["ListEvalsResponse", Coroutine[Any, Any, "ListEvalsResponse"]]: + ) -> Union["ListEvalsResponse", Coroutine[None, None, "ListEvalsResponse"]]: """List evals""" if _is_async: return self.async_list_evals_handler( @@ -12169,14 +12156,14 @@ class BaseLLMHTTPHandler: async def async_list_evals_handler( self, url: str, - query_params: Dict, + query_params: dict, evals_api_provider_config: "BaseEvalsAPIConfig", custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, shared_session: Optional["ClientSession"] = None, ) -> "ListEvalsResponse": """Async list evals""" @@ -12220,12 +12207,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, shared_session: Optional["ClientSession"] = None, - ) -> Union["Eval", Coroutine[Any, Any, "Eval"]]: + ) -> Union["Eval", Coroutine[None, None, "Eval"]]: """Get an eval""" if _is_async: return self.async_get_eval_handler( @@ -12276,9 +12263,9 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, shared_session: Optional["ClientSession"] = None, ) -> "Eval": """Async get an eval""" @@ -12317,17 +12304,17 @@ class BaseLLMHTTPHandler: def update_eval_handler( self, url: str, - request_body: Dict, + request_body: dict, evals_api_provider_config: "BaseEvalsAPIConfig", custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, shared_session: Optional["ClientSession"] = None, - ) -> Union["Eval", Coroutine[Any, Any, "Eval"]]: + ) -> Union["Eval", Coroutine[None, None, "Eval"]]: """Update an eval""" if _is_async: return self.async_update_eval_handler( @@ -12376,14 +12363,14 @@ class BaseLLMHTTPHandler: async def async_update_eval_handler( self, url: str, - request_body: Dict, + request_body: dict, evals_api_provider_config: "BaseEvalsAPIConfig", custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, shared_session: Optional["ClientSession"] = None, ) -> "Eval": """Async update an eval""" @@ -12427,12 +12414,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, shared_session: Optional["ClientSession"] = None, - ) -> Union["DeleteEvalResponse", Coroutine[Any, Any, "DeleteEvalResponse"]]: + ) -> Union["DeleteEvalResponse", Coroutine[None, None, "DeleteEvalResponse"]]: """Delete an eval""" if _is_async: return self.async_delete_eval_handler( @@ -12483,9 +12470,9 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, shared_session: Optional["ClientSession"] = None, ) -> "DeleteEvalResponse": """Async delete an eval""" @@ -12528,12 +12515,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, shared_session: Optional["ClientSession"] = None, - ) -> Union["CancelEvalResponse", Coroutine[Any, Any, "CancelEvalResponse"]]: + ) -> Union["CancelEvalResponse", Coroutine[None, None, "CancelEvalResponse"]]: """Cancel an eval""" if _is_async: return self.async_cancel_eval_handler( @@ -12584,9 +12571,9 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, shared_session: Optional["ClientSession"] = None, ) -> "CancelEvalResponse": """Async cancel an eval""" @@ -12629,17 +12616,17 @@ class BaseLLMHTTPHandler: def create_run_handler( self, url: str, - request_body: Dict, + request_body: dict, evals_api_provider_config: "BaseEvalsAPIConfig", custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, shared_session: Optional["ClientSession"] = None, - ) -> Union["Run", Coroutine[Any, Any, "Run"]]: + ) -> Union["Run", Coroutine[None, None, "Run"]]: """Create a run""" if _is_async: return self.async_create_run_handler( @@ -12688,14 +12675,14 @@ class BaseLLMHTTPHandler: async def async_create_run_handler( self, url: str, - request_body: Dict, + request_body: dict, evals_api_provider_config: "BaseEvalsAPIConfig", custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, shared_session: Optional["ClientSession"] = None, ) -> "Run": """Async create a run""" @@ -12735,17 +12722,17 @@ class BaseLLMHTTPHandler: def list_runs_handler( self, url: str, - query_params: Dict, + query_params: dict, evals_api_provider_config: "BaseEvalsAPIConfig", custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, shared_session: Optional["ClientSession"] = None, - ) -> Union["ListRunsResponse", Coroutine[Any, Any, "ListRunsResponse"]]: + ) -> Union["ListRunsResponse", Coroutine[None, None, "ListRunsResponse"]]: """List runs""" if _is_async: return self.async_list_runs_handler( @@ -12794,14 +12781,14 @@ class BaseLLMHTTPHandler: async def async_list_runs_handler( self, url: str, - query_params: Dict, + query_params: dict, evals_api_provider_config: "BaseEvalsAPIConfig", custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, shared_session: Optional["ClientSession"] = None, ) -> "ListRunsResponse": """Async list runs""" @@ -12845,12 +12832,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, shared_session: Optional["ClientSession"] = None, - ) -> Union["Run", Coroutine[Any, Any, "Run"]]: + ) -> Union["Run", Coroutine[None, None, "Run"]]: """Get a run""" if _is_async: return self.async_get_run_handler( @@ -12901,9 +12888,9 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, shared_session: Optional["ClientSession"] = None, ) -> "Run": """Async get a run""" @@ -12946,12 +12933,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, shared_session: Optional["ClientSession"] = None, - ) -> Union["CancelRunResponse", Coroutine[Any, Any, "CancelRunResponse"]]: + ) -> Union["CancelRunResponse", Coroutine[None, None, "CancelRunResponse"]]: """Cancel a run""" if _is_async: return self.async_cancel_run_handler( @@ -13002,9 +12989,9 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, shared_session: Optional["ClientSession"] = None, ) -> "CancelRunResponse": """Async cancel a run""" @@ -13047,12 +13034,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, shared_session: Optional["ClientSession"] = None, - ) -> Union["RunDeleteResponse", Coroutine[Any, Any, "RunDeleteResponse"]]: + ) -> Union["RunDeleteResponse", Coroutine[None, None, "RunDeleteResponse"]]: """Delete a run""" if _is_async: return self.async_delete_run_handler( @@ -13103,9 +13090,9 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + extra_headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, shared_session: Optional["ClientSession"] = None, ) -> "RunDeleteResponse": """Async delete a run""" diff --git a/litellm/main.py b/litellm/main.py index acdec7385da..b5d16360c40 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -19,6 +19,7 @@ import random import sys import time import traceback +from collections.abc import AsyncIterator, Coroutine, Iterable, Mapping from concurrent import futures from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait from copy import deepcopy @@ -26,16 +27,8 @@ from functools import partial from typing import ( TYPE_CHECKING, Any, - AsyncIterator, - Coroutine, - Dict, - Iterable, - List, Literal, - Mapping, Optional, - Tuple, - Type, Union, cast, get_args, @@ -346,30 +339,28 @@ class LiteLLM: self, *, api_key=None, - organization: Optional[str] = None, - base_url: Optional[str] = None, - timeout: Optional[float] = 600, - max_retries: Optional[int] = litellm.num_retries, - default_headers: Optional[Mapping[str, str]] = None, + organization: str | None = None, + base_url: str | None = None, + timeout: float | None = 600, + max_retries: int | None = litellm.num_retries, + default_headers: Mapping[str, str] | None = None, ): self.params = locals() self.chat = Chat(self.params, router_obj=None) class Chat: - def __init__(self, params, router_obj: Optional[Any]): + def __init__(self, params, router_obj: Any | None): self.params = params if self.params.get("acompletion", False) is True: self.params.pop("acompletion") - self.completions: Union[AsyncCompletions, Completions] = AsyncCompletions( - self.params, router_obj=router_obj - ) + self.completions: AsyncCompletions | Completions = AsyncCompletions(self.params, router_obj=router_obj) else: self.completions = Completions(self.params, router_obj=router_obj) class Completions: - def __init__(self, params, router_obj: Optional[Any]): + def __init__(self, params, router_obj: Any | None): self.params = params self.router_obj = router_obj @@ -385,7 +376,7 @@ class Completions: class AsyncCompletions: - def __init__(self, params, router_obj: Optional[Any]): + def __init__(self, params, router_obj: Any | None): self.params = params self.router_obj = router_obj @@ -405,54 +396,54 @@ class AsyncCompletions: async def acompletion( 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, - timeout: Optional[Union[float, int]] = None, - temperature: Optional[float] = None, - top_p: Optional[float] = None, - n: Optional[int] = None, - stream: Optional[bool] = None, - stream_options: Optional[dict] = None, + messages: list = [], + functions: list | None = None, + function_call: str | None = None, + timeout: float | int | None = None, + temperature: float | None = None, + top_p: float | None = None, + n: int | None = None, + stream: bool | None = None, + stream_options: dict | None = None, stop=None, - max_tokens: Optional[int] = None, - max_completion_tokens: Optional[int] = None, - modalities: Optional[List[ChatCompletionModality]] = None, - prediction: Optional[ChatCompletionPredictionContentParam] = None, - audio: Optional[ChatCompletionAudioParam] = 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, + max_completion_tokens: int | None = None, + modalities: list[ChatCompletionModality] | None = None, + prediction: ChatCompletionPredictionContentParam | None = None, + audio: ChatCompletionAudioParam | None = None, + presence_penalty: float | None = None, + frequency_penalty: float | None = None, + logit_bias: dict | None = None, + user: str | None = None, # openai v1.0+ new params - response_format: Optional[Union[dict, Type[BaseModel]]] = None, - seed: Optional[int] = None, - tools: Optional[List] = None, - tool_choice: Optional[Union[str, dict]] = None, - parallel_tool_calls: Optional[bool] = None, - logprobs: Optional[bool] = None, - top_logprobs: Optional[int] = None, + response_format: dict | type[BaseModel] | None = None, + seed: int | None = None, + tools: list | None = None, + tool_choice: str | dict | None = None, + parallel_tool_calls: bool | None = None, + logprobs: bool | None = None, + top_logprobs: int | None = None, deployment_id=None, - reasoning_effort: Optional[Literal["none", "minimal", "low", "medium", "high", "xhigh", "default"]] = None, - verbosity: Optional[Literal["low", "medium", "high"]] = None, - safety_identifier: Optional[str] = None, - service_tier: Optional[str] = None, + reasoning_effort: Literal["none", "minimal", "low", "medium", "high", "xhigh", "default"] | None = None, + verbosity: Literal["low", "medium", "high"] | None = None, + safety_identifier: str | None = None, + service_tier: str | None = None, # set api_base, api_version, api_key - base_url: Optional[str] = None, - api_version: Optional[str] = None, - api_key: Optional[str] = None, - model_list: Optional[list] = None, # pass in a list of api_base,keys, etc. - extra_headers: Optional[dict] = None, + base_url: str | None = None, + api_version: str | None = None, + api_key: str | None = None, + model_list: list | None = None, # pass in a list of api_base,keys, etc. + extra_headers: dict | None = None, # Optional liteLLM function params - thinking: Optional[AnthropicThinkingParam] = None, - web_search_options: Optional[OpenAIWebSearchOptions] = None, - include_server_side_tool_invocations: Optional[bool] = None, + thinking: AnthropicThinkingParam | None = None, + web_search_options: OpenAIWebSearchOptions | None = None, + include_server_side_tool_invocations: bool | None = None, # Session management shared_session: Optional["ClientSession"] = None, # Per-request JSON schema validation (overrides litellm.enable_json_schema_validation) - enable_json_schema_validation: Optional[bool] = None, + enable_json_schema_validation: bool | None = None, **kwargs, -) -> Union[ModelResponse, CustomStreamWrapper]: +) -> ModelResponse | CustomStreamWrapper: """ Asynchronously executes a litellm.completion() call for any of litellm supported llms (example gpt-4, gpt-3.5-turbo, claude-2, command-nightly) @@ -519,7 +510,7 @@ async def acompletion( non_default_params=kwargs, messages=cast(list[AllMessageValues], messages), # cast-ok: acompletion types messages as a bare List model=model, - custom_llm_provider=cast(Optional[str], custom_llm_provider), # cast-ok: read from untyped kwargs + custom_llm_provider=cast(str | None, custom_llm_provider), # cast-ok: read from untyped kwargs tools=tools, ) @@ -733,9 +724,9 @@ async def _async_streaming(response, model, custom_llm_provider, args): def _handle_mock_potential_exceptions( - mock_response: Union[str, Exception], + mock_response: str | Exception, model: str, - custom_llm_provider: Optional[str] = None, + custom_llm_provider: str | None = None, ): if isinstance(mock_response, Exception): if isinstance(mock_response, openai.APIError): @@ -776,8 +767,8 @@ def _handle_mock_potential_exceptions( def _handle_mock_timeout( - mock_timeout: Optional[bool], - timeout: Optional[Union[float, str, httpx.Timeout]], + mock_timeout: bool | None, + timeout: float | str | httpx.Timeout | None, model: str, ): if mock_timeout is True and timeout is not None: @@ -790,8 +781,8 @@ def _handle_mock_timeout( async def _handle_mock_timeout_async( - mock_timeout: Optional[bool], - timeout: Optional[Union[float, str, httpx.Timeout]], + mock_timeout: bool | None, + timeout: float | str | httpx.Timeout | None, model: str, ): if mock_timeout is True and timeout is not None: @@ -803,7 +794,7 @@ async def _handle_mock_timeout_async( ) -def _sleep_for_timeout(timeout: Union[float, str, httpx.Timeout]): +def _sleep_for_timeout(timeout: float | str | httpx.Timeout): if isinstance(timeout, float): time.sleep(timeout) elif isinstance(timeout, str): @@ -812,7 +803,7 @@ def _sleep_for_timeout(timeout: Union[float, str, httpx.Timeout]): time.sleep(timeout.connect) -async def _sleep_for_timeout_async(timeout: Union[float, str, httpx.Timeout]): +async def _sleep_for_timeout_async(timeout: float | str | httpx.Timeout): if isinstance(timeout, float): await asyncio.sleep(timeout) elif isinstance(timeout, str): @@ -823,15 +814,15 @@ async def _sleep_for_timeout_async(timeout: Union[float, str, httpx.Timeout]): def mock_completion( model: str, - messages: List, - stream: Optional[bool] = False, - n: Optional[int] = None, - mock_response: Optional[MOCK_RESPONSE_TYPE] = "This is a mock request", - mock_tool_calls: Optional[List] = None, - mock_timeout: Optional[bool] = False, - logging=None, + messages: list, + stream: bool | None = False, + n: int | None = None, + mock_response: MOCK_RESPONSE_TYPE | None = "This is a mock request", + mock_tool_calls: list | None = None, + mock_timeout: bool | None = False, + logging: LiteLLMLoggingObj | None = None, custom_llm_provider=None, - timeout: Optional[Union[float, str, httpx.Timeout]] = None, + timeout: float | str | httpx.Timeout | None = None, **kwargs, ): """ @@ -879,7 +870,7 @@ def mock_completion( ) mock_response = cast( - Union[str, dict, ModelResponse, ModelResponseStream], mock_response + str | dict | ModelResponse | ModelResponseStream, mock_response ) # after this point, mock_response is a string, dict, ModelResponse, or ModelResponseStream if isinstance(mock_response, str) and mock_response.startswith("Exception: mock_streaming_error"): mock_response = litellm.MockException( @@ -901,10 +892,11 @@ def mock_completion( # convert to ModelResponseStream mock_response = convert_model_response_to_streaming(mock_response) # type: ignore - model_response: Union[ModelResponse, ModelResponseStream] = ModelResponse() + model_response: ModelResponse | ModelResponseStream = ModelResponse() if stream is True: model_response = ModelResponseStream() + assert logging is not None # don't try to access stream object, if kwargs.get("acompletion", False) is True: return CustomStreamWrapper( @@ -982,12 +974,12 @@ def mock_completion( def responses_api_bridge_check( model: str, custom_llm_provider: str, - web_search_options: Optional[OpenAIWebSearchOptions] = None, - tools: Optional[List[Any]] = None, - reasoning_effort: Optional[Any] = None, - reasoning_summary: Optional[Any] = None, -) -> Tuple[dict, str]: - model_info: Dict[str, Any] = {} + web_search_options: OpenAIWebSearchOptions | None = None, + tools: list[Any] | None = None, + reasoning_effort: Any | None = None, + reasoning_summary: Any | None = None, +) -> tuple[dict, str]: + model_info: dict[str, Any] = {} # Global flag: route ALL OpenAI chat completions through Responses API. # Returns early with minimal model_info; callers only inspect the "mode" key. @@ -1039,7 +1031,7 @@ def responses_api_bridge_check( return model_info, model -def _should_allow_input_examples(custom_llm_provider: Optional[str], model: str) -> bool: +def _should_allow_input_examples(custom_llm_provider: str | None, model: str) -> bool: if custom_llm_provider == "anthropic": return True if custom_llm_provider == "azure_ai" or custom_llm_provider == "bedrock" or custom_llm_provider == "vertex_ai": @@ -1059,11 +1051,11 @@ def _drop_input_examples_from_tool(tool: dict) -> dict: def _drop_input_examples_from_tools( - tools: Optional[List[dict]], -) -> Optional[List[dict]]: + tools: list[dict] | None, +) -> list[dict] | None: if tools is None: return None - cleaned_tools: List[dict] = [] + cleaned_tools: list[dict] = [] for tool in tools: if isinstance(tool, dict): cleaned_tools.append(_drop_input_examples_from_tool(tool)) @@ -1075,7 +1067,7 @@ def _drop_input_examples_from_tools( def _build_custom_pricing_entry( custom_llm_provider: str, kwargs: dict, - model_info: Optional[dict] = None, + model_info: dict | None = None, ) -> dict: """Build a complete model cost entry from kwargs and model_info. @@ -1098,7 +1090,7 @@ def _build_custom_pricing_entry( return entry -def _get_router_deployment_id(kwargs: dict) -> Optional[str]: +def _get_router_deployment_id(kwargs: dict) -> str | None: for metadata_key in ("litellm_metadata", "metadata"): metadata = kwargs.get(metadata_key) or {} if not isinstance(metadata, dict): @@ -1116,7 +1108,7 @@ def _register_custom_pricing_for_request( model: str, custom_llm_provider: str, kwargs: dict, - model_info: Optional[dict], + model_info: dict | None, ) -> None: """Register per-request custom pricing in litellm.model_cost. @@ -2601,7 +2593,7 @@ def _complete_anthropic_text( api_key = api_key or litellm.anthropic_key or litellm.api_key or os.environ.get("ANTHROPIC_API_KEY") custom_prompt_dict = custom_prompt_dict or litellm.custom_prompt_dict api_base = cast( - Optional[str], + str | None, api_base or litellm.api_base or get_secret("ANTHROPIC_API_BASE") @@ -2657,7 +2649,7 @@ def _complete_anthropic(ctx: _CompletionDispatchContext) -> _CompletionDispatchR # call /messages # default route for all anthropic models api_base = cast( - Optional[str], + str | None, api_base or litellm.api_base or get_secret("ANTHROPIC_API_BASE") @@ -4047,7 +4039,7 @@ def _complete_watsonx_text( optional_params.pop("watsonx_credentials", None), # follow {provider}_credentials, same as vertex ai ) - token: Optional[str] = None + token: str | None = None if wx_credentials is not None: api_base = wx_credentials.get("url", api_base) api_key = wx_credentials.get("apikey", wx_credentials.get("api_key", api_key)) @@ -4665,7 +4657,7 @@ def _complete_custom_providers( stream = ctx.stream timeout = ctx.timeout - 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"] @@ -4811,55 +4803,55 @@ def _complete_langflow(ctx: _CompletionDispatchContext) -> _CompletionDispatchRe def completion( # type: ignore model: str, # Optional OpenAI params: see https://platform.openai.com/docs/api-reference/chat/create - messages: List = [], - timeout: Optional[Union[float, str, httpx.Timeout]] = None, - temperature: Optional[float] = None, - top_p: Optional[float] = None, - n: Optional[int] = None, - stream: Optional[bool] = None, - stream_options: Optional[dict] = None, + messages: list = [], + timeout: float | str | httpx.Timeout | None = None, + temperature: float | None = None, + top_p: float | None = None, + n: int | None = None, + stream: bool | None = None, + stream_options: dict | None = None, stop=None, - max_completion_tokens: Optional[int] = None, - max_tokens: Optional[int] = None, - modalities: Optional[List[ChatCompletionModality]] = None, - prediction: Optional[ChatCompletionPredictionContentParam] = None, - audio: Optional[ChatCompletionAudioParam] = None, - presence_penalty: Optional[float] = None, - frequency_penalty: Optional[float] = None, - logit_bias: Optional[dict] = None, - user: Optional[str] = None, + max_completion_tokens: int | None = None, + max_tokens: int | None = None, + modalities: list[ChatCompletionModality] | None = None, + prediction: ChatCompletionPredictionContentParam | None = None, + audio: ChatCompletionAudioParam | None = None, + presence_penalty: float | None = None, + frequency_penalty: float | None = None, + logit_bias: dict | None = None, + user: str | None = None, # openai v1.0+ new params - reasoning_effort: Optional[Literal["none", "minimal", "low", "medium", "high", "xhigh", "default"]] = None, - verbosity: Optional[Literal["low", "medium", "high"]] = None, - response_format: Optional[Union[dict, Type[BaseModel]]] = None, - seed: Optional[int] = None, - tools: Optional[List] = None, - tool_choice: Optional[Union[str, dict]] = None, - logprobs: Optional[bool] = None, - top_logprobs: Optional[int] = None, - parallel_tool_calls: Optional[bool] = None, - web_search_options: Optional[OpenAIWebSearchOptions] = None, - include_server_side_tool_invocations: Optional[bool] = None, + reasoning_effort: Literal["none", "minimal", "low", "medium", "high", "xhigh", "default"] | None = None, + verbosity: Literal["low", "medium", "high"] | None = None, + response_format: dict | type[BaseModel] | None = None, + seed: int | None = None, + tools: list | None = None, + tool_choice: str | dict | None = None, + logprobs: bool | None = None, + top_logprobs: int | None = None, + parallel_tool_calls: bool | None = None, + web_search_options: OpenAIWebSearchOptions | None = None, + include_server_side_tool_invocations: bool | None = None, deployment_id=None, - extra_headers: Optional[dict] = None, - safety_identifier: Optional[str] = None, - service_tier: Optional[str] = None, + extra_headers: dict | None = None, + safety_identifier: str | None = None, + service_tier: str | None = None, # soon to be deprecated params by OpenAI - functions: Optional[List] = None, - function_call: Optional[str] = None, + functions: list | None = None, + function_call: str | None = None, # set api_base, api_version, api_key - base_url: Optional[str] = None, - api_version: Optional[str] = None, - api_key: Optional[str] = None, - model_list: Optional[list] = None, # pass in a list of api_base,keys, etc. + base_url: str | None = None, + api_version: str | None = None, + api_key: str | None = None, + model_list: list | None = None, # pass in a list of api_base,keys, etc. # Optional liteLLM function params - thinking: Optional[AnthropicThinkingParam] = None, + thinking: AnthropicThinkingParam | None = None, # Session management shared_session: Optional["ClientSession"] = None, # Per-request JSON schema validation (overrides litellm.enable_json_schema_validation) - enable_json_schema_validation: Optional[bool] = None, + enable_json_schema_validation: bool | None = None, **kwargs, -) -> Union[ModelResponse, CustomStreamWrapper]: +) -> ModelResponse | CustomStreamWrapper: """ Perform a completion() using any of litellm supported llms (example gpt-4, gpt-3.5-turbo, claude-2, command-nightly) Parameters: @@ -4937,7 +4929,7 @@ def completion( # type: ignore # Check if MCP tools are present (following responses pattern) # Cast tools to Optional[Iterable[ToolParam]] for type checking - tools_for_mcp = cast(Optional[Iterable[ToolParam]], tools) + tools_for_mcp = cast(Iterable[ToolParam] | None, tools) if LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(tools=tools_for_mcp): return acompletion_with_mcp( # pyright: ignore[reportReturnType] # MCP path returns a coroutine that acompletion() awaits; completion()'s sync return type omits it model=model, @@ -4984,9 +4976,9 @@ def completion( # type: ignore **kwargs, ) api_base = kwargs.get("api_base", None) - mock_response: Optional[MOCK_RESPONSE_TYPE] = kwargs.get("mock_response", None) + mock_response: MOCK_RESPONSE_TYPE | None = kwargs.get("mock_response", None) mock_tool_calls = kwargs.get("mock_tool_calls", None) - mock_timeout = cast(Optional[bool], kwargs.get("mock_timeout", None)) + mock_timeout = cast(bool | None, kwargs.get("mock_timeout", None)) force_timeout = kwargs.get("force_timeout", 600) ## deprecated logger_fn = kwargs.get("logger_fn", None) verbose = kwargs.get("verbose", False) @@ -4997,14 +4989,12 @@ def completion( # type: ignore model_info = kwargs.get("model_info", None) proxy_server_request = kwargs.get("proxy_server_request", None) fallbacks = kwargs.get("fallbacks", None) - provider_specific_header = cast(Optional[ProviderSpecificHeader], kwargs.get("provider_specific_header", None)) + provider_specific_header = cast(ProviderSpecificHeader | None, kwargs.get("provider_specific_header", None)) headers = kwargs.get("headers", None) or extra_headers - ensure_alternating_roles: Optional[bool] = kwargs.get("ensure_alternating_roles", None) - user_continue_message: Optional[ChatCompletionUserMessage] = kwargs.get("user_continue_message", None) - assistant_continue_message: Optional[ChatCompletionAssistantMessage] = kwargs.get( - "assistant_continue_message", None - ) + ensure_alternating_roles: bool | None = kwargs.get("ensure_alternating_roles", None) + user_continue_message: ChatCompletionUserMessage | None = kwargs.get("user_continue_message", None) + assistant_continue_message: ChatCompletionAssistantMessage | None = kwargs.get("assistant_continue_message", None) if headers is None: headers = {} if extra_headers is not None: @@ -5053,8 +5043,8 @@ def completion( # type: ignore ### Admin Controls ### no_log = kwargs.get("no-log", False) ### PROMPT MANAGEMENT ### - prompt_id = cast(Optional[str], kwargs.get("prompt_id", None)) - prompt_variables = cast(Optional[dict], kwargs.get("prompt_variables", None)) + prompt_id = cast(str | None, kwargs.get("prompt_id", None)) + prompt_variables = cast(dict | None, kwargs.get("prompt_variables", None)) litellm_system_prompt = kwargs.get("litellm_system_prompt", None) ### COPY MESSAGES ### - related issue https://github.com/BerriAI/litellm/discussions/4489 messages = get_completion_messages( @@ -5077,7 +5067,7 @@ def completion( # type: ignore non_default_params=non_default_params, messages=cast(list[AllMessageValues], messages), # cast-ok: completion types messages as a bare List model=model, - custom_llm_provider=cast(Optional[str], kwargs.get("custom_llm_provider")), # cast-ok: untyped kwargs + custom_llm_provider=cast(str | None, kwargs.get("custom_llm_provider")), # cast-ok: untyped kwargs tools=tools, ) @@ -5213,12 +5203,12 @@ def completion( # type: ignore messages=messages, model_id=(kwargs.get("model_info") or {}).get("id", None), model_file_id_mapping=cast( - Dict[str, Dict[str, str]], + dict[str, dict[str, str]], kwargs.get("model_file_id_mapping") or {}, ), ) - provider_config: Optional[BaseConfig] = None + provider_config: BaseConfig | None = None if custom_llm_provider is not None and custom_llm_provider in [provider.value for provider in LlmProviders]: provider_config = ProviderConfigManager.get_provider_chat_config( model=model, @@ -5838,7 +5828,7 @@ async def aembedding(*args, **kwargs) -> EmbeddingResponse: # Await normally init_response = await loop.run_in_executor(None, func_with_context) - response: Optional[EmbeddingResponse] = None + response: EmbeddingResponse | None = None if isinstance(init_response, dict): response = EmbeddingResponse(**init_response) elif isinstance(init_response, EmbeddingResponse): ## CACHING SCENARIO @@ -5870,16 +5860,16 @@ def embedding( model, input=[], # Optional params - dimensions: Optional[int] = None, - encoding_format: Optional[str] = None, + dimensions: int | None = None, + encoding_format: str | None = None, timeout=600, # default to 10 minutes # set api_base, api_version, api_key - api_base: Optional[str] = None, - api_version: Optional[str] = None, - api_key: Optional[str] = None, - api_type: Optional[str] = None, + api_base: str | None = None, + api_version: str | None = None, + api_key: str | None = None, + api_type: str | None = None, caching: bool = False, - user: Optional[str] = None, + user: str | None = None, custom_llm_provider=None, litellm_call_id=None, logger_fn=None, @@ -5896,16 +5886,16 @@ def embedding( model, input=[], # Optional params - dimensions: Optional[int] = None, - encoding_format: Optional[str] = None, + dimensions: int | None = None, + encoding_format: str | None = None, timeout=600, # default to 10 minutes # set api_base, api_version, api_key - api_base: Optional[str] = None, - api_version: Optional[str] = None, - api_key: Optional[str] = None, - api_type: Optional[str] = None, + api_base: str | None = None, + api_version: str | None = None, + api_key: str | None = None, + api_type: str | None = None, caching: bool = False, - user: Optional[str] = None, + user: str | None = None, custom_llm_provider=None, litellm_call_id=None, logger_fn=None, @@ -5923,21 +5913,21 @@ def embedding( model, input=[], # Optional params - dimensions: Optional[int] = None, - encoding_format: Optional[str] = None, + dimensions: int | None = None, + encoding_format: str | None = None, timeout=600, # default to 10 minutes # set api_base, api_version, api_key - api_base: Optional[str] = None, - api_version: Optional[str] = None, - api_key: Optional[str] = None, - api_type: Optional[str] = None, + api_base: str | None = None, + api_version: str | None = None, + api_key: str | None = None, + api_type: str | None = None, caching: bool = False, - user: Optional[str] = None, + user: str | None = None, custom_llm_provider=None, litellm_call_id=None, logger_fn=None, **kwargs, -) -> Union[EmbeddingResponse, Coroutine[Any, Any, EmbeddingResponse]]: +) -> EmbeddingResponse | Coroutine[Any, Any, EmbeddingResponse]: """ Embedding function that calls an API to generate embeddings for the given input. @@ -5968,9 +5958,9 @@ def embedding( shared_session = kwargs.get("shared_session", None) max_retries = kwargs.get("max_retries", None) litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore - mock_response: Optional[List[float]] = kwargs.get("mock_response", None) # type: ignore + mock_response: list[float] | None = kwargs.get("mock_response", None) # type: ignore azure_ad_token_provider = kwargs.get("azure_ad_token_provider", None) - aembedding: Optional[bool] = kwargs.get("aembedding", None) + aembedding: bool | None = kwargs.get("aembedding", None) extra_headers = kwargs.get("extra_headers", None) headers = kwargs.get("headers", None) or extra_headers if headers is None: @@ -6023,7 +6013,7 @@ def embedding( if dynamic_api_key is not None: api_key = dynamic_api_key - allowed_openai_params: Optional[List[str]] = kwargs.get("allowed_openai_params", None) + allowed_openai_params: list[str] | None = kwargs.get("allowed_openai_params", None) optional_params = get_optional_params_embeddings( model=model, user=user, @@ -6057,7 +6047,7 @@ def embedding( if mock_response is not None: return mock_embedding(model=model, mock_response=mock_response) try: - response: Optional[Union[EmbeddingResponse, Coroutine[Any, Any, EmbeddingResponse]]] = None + response: EmbeddingResponse | Coroutine[Any, Any, EmbeddingResponse] | None = None if azure is True or custom_llm_provider == "azure": # azure configs @@ -6878,7 +6868,7 @@ def embedding( litellm_params={}, ) elif 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"] @@ -6979,7 +6969,7 @@ def embedding( ###### Text Completion ################ @client -async def atext_completion(*args, **kwargs) -> Union[TextCompletionResponse, TextCompletionStreamWrapper]: +async def atext_completion(*args, **kwargs) -> TextCompletionResponse | TextCompletionStreamWrapper: """ Implemented to handle async streaming for the text completion endpoint """ @@ -7050,36 +7040,31 @@ async def atext_completion(*args, **kwargs) -> Union[TextCompletionResponse, Tex @client def text_completion( - prompt: Union[ - str, List[Union[str, List[Union[str, List[int]]]]] - ], # Required: The prompt(s) to generate completions for. - model: Optional[str] = None, # Optional: either `model` or `engine` can be set - best_of: Optional[int] = None, # Optional: Generates best_of completions server-side. - echo: Optional[bool] = None, # Optional: Echo back the prompt in addition to the completion. - frequency_penalty: Optional[float] = None, # Optional: Penalize new tokens based on their existing frequency. - logit_bias: Optional[Dict[int, int]] = None, # Optional: Modify the likelihood of specified tokens. - logprobs: Optional[int] = None, # Optional: Include the log probabilities on the most likely tokens. - max_tokens: Optional[int] = None, # Optional: The maximum number of tokens to generate in the completion. - n: Optional[int] = None, # Optional: How many completions to generate for each prompt. - presence_penalty: Optional[ - float - ] = None, # Optional: Penalize new tokens based on whether they appear in the text so far. - stop: Optional[ - Union[str, List[str]] - ] = None, # Optional: Sequences where the API will stop generating further tokens. - stream: Optional[bool] = None, # Optional: Whether to stream back partial progress. - stream_options: Optional[dict] = None, - suffix: Optional[str] = None, # Optional: The suffix that comes after a completion of inserted text. - temperature: Optional[float] = None, # Optional: Sampling temperature to use. - top_p: Optional[float] = None, # Optional: Nucleus sampling parameter. - user: Optional[str] = None, # Optional: A unique identifier representing your end-user. + prompt: str | list[str | list[str | list[int]]], # Required: The prompt(s) to generate completions for. + model: str | None = None, # Optional: either `model` or `engine` can be set + best_of: int | None = None, # Optional: Generates best_of completions server-side. + echo: bool | None = None, # Optional: Echo back the prompt in addition to the completion. + frequency_penalty: float | None = None, # Optional: Penalize new tokens based on their existing frequency. + logit_bias: dict[int, int] | None = None, # Optional: Modify the likelihood of specified tokens. + logprobs: int | None = None, # Optional: Include the log probabilities on the most likely tokens. + max_tokens: int | None = None, # Optional: The maximum number of tokens to generate in the completion. + n: int | None = None, # Optional: How many completions to generate for each prompt. + presence_penalty: float + | None = None, # Optional: Penalize new tokens based on whether they appear in the text so far. + stop: str | list[str] | None = None, # Optional: Sequences where the API will stop generating further tokens. + stream: bool | None = None, # Optional: Whether to stream back partial progress. + stream_options: dict | None = None, + suffix: str | None = None, # Optional: The suffix that comes after a completion of inserted text. + temperature: float | None = None, # Optional: Sampling temperature to use. + top_p: float | None = None, # Optional: Nucleus sampling parameter. + user: str | None = None, # Optional: A unique identifier representing your end-user. # set api_base, api_version, api_key - api_base: Optional[str] = None, - api_version: Optional[str] = None, - api_key: Optional[str] = None, - model_list: Optional[list] = None, # pass in a list of api_base,keys, etc. + api_base: str | None = None, + api_version: str | None = None, + api_key: str | None = None, + model_list: list | None = None, # pass in a list of api_base,keys, etc. # Optional liteLLM function params - custom_llm_provider: Optional[str] = None, + custom_llm_provider: str | None = None, *args, **kwargs, ): @@ -7120,7 +7105,7 @@ def text_completion( text_completion_response = TextCompletionResponse() - optional_params: Dict[str, Any] = {} + optional_params: dict[str, Any] = {} # default values for all optional params are none, litellm only passes them to the llm when they are set to non None values if best_of is not None: optional_params["best_of"] = best_of @@ -7290,14 +7275,12 @@ def text_completion( ###### Adapter Completion ################ -async def aadapter_completion( - *, adapter_id: str, **kwargs -) -> Optional[Union[BaseModel, AdapterCompletionStreamWrapper]]: +async def aadapter_completion(*, adapter_id: str, **kwargs) -> BaseModel | AdapterCompletionStreamWrapper | None: """ Implemented to handle async calls for adapter_completion() """ try: - translation_obj: Optional[CustomLogger] = None + translation_obj: CustomLogger | None = None for item in litellm.adapters: if item["id"] == adapter_id: translation_obj = item["adapter"] @@ -7311,8 +7294,8 @@ async def aadapter_completion( new_kwargs = translation_obj.translate_completion_input_params(kwargs=kwargs) - response: Union[ModelResponse, CustomStreamWrapper] = await acompletion(**new_kwargs) # type: ignore - translated_response: Optional[Union[BaseModel, AdapterCompletionStreamWrapper]] = None + response: ModelResponse | CustomStreamWrapper = await acompletion(**new_kwargs) # type: ignore + translated_response: BaseModel | AdapterCompletionStreamWrapper | None = None if isinstance(response, ModelResponse): translated_response = translation_obj.translate_completion_output_params(response=response) if isinstance(response, CustomStreamWrapper): @@ -7327,18 +7310,18 @@ async def aadapter_completion( async def aadapter_generate_content( **kwargs, -) -> Union[Dict[str, Any], AsyncIterator[bytes]]: +) -> dict[str, Any] | AsyncIterator[bytes]: from litellm.google_genai.adapters.handler import GenerateContentToCompletionHandler coro = cast( - Coroutine[Any, Any, Union[Dict[str, Any], AsyncIterator[bytes]]], + Coroutine[Any, Any, dict[str, Any] | AsyncIterator[bytes]], GenerateContentToCompletionHandler.generate_content_handler(**kwargs, _is_async=True), ) return await coro -def adapter_completion(*, adapter_id: str, **kwargs) -> Optional[Union[BaseModel, AdapterCompletionStreamWrapper]]: - translation_obj: Optional[CustomLogger] = None +def adapter_completion(*, adapter_id: str, **kwargs) -> BaseModel | AdapterCompletionStreamWrapper | None: + translation_obj: CustomLogger | None = None for item in litellm.adapters: if item["id"] == adapter_id: translation_obj = item["adapter"] @@ -7352,8 +7335,8 @@ def adapter_completion(*, adapter_id: str, **kwargs) -> Optional[Union[BaseModel new_kwargs = translation_obj.translate_completion_input_params(kwargs=kwargs) - response: Union[ModelResponse, CustomStreamWrapper] = completion(**new_kwargs) # type: ignore - translated_response: Optional[Union[BaseModel, AdapterCompletionStreamWrapper]] = None + response: ModelResponse | CustomStreamWrapper = completion(**new_kwargs) # type: ignore + translated_response: BaseModel | AdapterCompletionStreamWrapper | None = None if isinstance(response, ModelResponse): translated_response = translation_obj.translate_completion_output_params(response=response) elif isinstance(response, CustomStreamWrapper) or inspect.isgenerator(response): @@ -7365,9 +7348,7 @@ def adapter_completion(*, adapter_id: str, **kwargs) -> Optional[Union[BaseModel ##### Moderation ####################### -def moderation( - input: str, model: Optional[str] = None, api_key: Optional[str] = None, **kwargs -) -> OpenAIModerationResponse: +def moderation(input: str, model: str | None = None, api_key: str | None = None, **kwargs) -> OpenAIModerationResponse: # only supports open ai for now api_key = api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY") @@ -7386,7 +7367,7 @@ def moderation( else: response = openai_client.moderations.create(input=input) - response_dict: Dict = response.model_dump() + response_dict: dict = response.model_dump() return litellm.utils.LiteLLMResponseObjectHandler.convert_to_moderation_response( response_object=response_dict, ) @@ -7395,9 +7376,9 @@ def moderation( @client async def amoderation( input: str, - model: Optional[str] = None, - api_key: Optional[str] = None, - custom_llm_provider: Optional[str] = None, + model: str | None = None, + api_key: str | None = None, + custom_llm_provider: str | None = None, **kwargs, ) -> OpenAIModerationResponse: from openai import AsyncOpenAI @@ -7405,7 +7386,7 @@ async def amoderation( # only supports open ai for now api_key = api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY") 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) _dynamic_api_base = None try: ( @@ -7452,7 +7433,7 @@ async def amoderation( response = await _openai_client.moderations.create(input=input, model=model) else: response = await _openai_client.moderations.create(input=input) - response_dict: Dict = response.model_dump() + response_dict: dict = response.model_dump() return litellm.utils.LiteLLMResponseObjectHandler.convert_to_moderation_response( response_object=response_dict, ) @@ -7528,21 +7509,21 @@ def transcription( model: str, file: FileTypes, ## OPTIONAL OPENAI PARAMS ## - language: Optional[str] = None, - prompt: Optional[str] = None, - response_format: Optional[Literal["json", "text", "srt", "verbose_json", "vtt"]] = None, - timestamp_granularities: Optional[List[Literal["word", "segment"]]] = None, - temperature: Optional[int] = None, # openai defaults this to 0 + language: str | None = None, + prompt: str | None = None, + response_format: Literal["json", "text", "srt", "verbose_json", "vtt"] | None = None, + timestamp_granularities: list[Literal["word", "segment"]] | None = None, + temperature: int | None = None, # openai defaults this to 0 ## LITELLM PARAMS ## - user: Optional[str] = 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, - max_retries: Optional[int] = None, + api_key: str | None = None, + api_base: str | None = None, + api_version: str | None = None, + max_retries: int | None = None, custom_llm_provider=None, **kwargs, -) -> Union[TranscriptionResponse, Coroutine[Any, Any, TranscriptionResponse]]: +) -> TranscriptionResponse | Coroutine[Any, Any, TranscriptionResponse]: """ Calls openai + azure whisper endpoints. @@ -7559,14 +7540,9 @@ def transcription( kwargs.pop("tags", []) non_default_params = get_non_default_transcription_params(kwargs) - client: Optional[ - Union[ - openai.AsyncOpenAI, - openai.OpenAI, - openai.AzureOpenAI, - openai.AsyncAzureOpenAI, - ] - ] = kwargs.pop("client", None) + client: openai.AsyncOpenAI | openai.OpenAI | openai.AzureOpenAI | openai.AsyncAzureOpenAI | None = kwargs.pop( + "client", None + ) if litellm_logging_obj: litellm_logging_obj.model_call_details["client"] = str(client) @@ -7614,7 +7590,7 @@ def transcription( custom_llm_provider=custom_llm_provider, ) - response: Optional[Union[TranscriptionResponse, Coroutine[Any, Any, TranscriptionResponse]]] = None + response: TranscriptionResponse | Coroutine[Any, Any, TranscriptionResponse] | None = None provider_config = ProviderConfigManager.get_provider_audio_transcription_config( model=model, @@ -7831,26 +7807,26 @@ async def aspeech(*args, **kwargs) -> HttpxBinaryResponseContent: def speech( model: str, input: str, - voice: Optional[Union[str, dict]] = None, - api_key: Optional[str] = None, - api_base: Optional[str] = None, - api_version: Optional[str] = None, - organization: Optional[str] = None, - project: Optional[str] = None, - max_retries: Optional[int] = None, - metadata: Optional[dict] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - response_format: Optional[str] = None, - speed: Optional[int] = None, - instructions: Optional[str] = None, + voice: str | dict | None = None, + api_key: str | None = None, + api_base: str | None = None, + api_version: str | None = None, + organization: str | None = None, + project: str | None = None, + max_retries: int | None = None, + metadata: dict | None = None, + timeout: float | httpx.Timeout | None = None, + response_format: str | None = None, + speed: int | None = None, + instructions: str | None = None, client=None, - headers: Optional[dict] = None, - custom_llm_provider: Optional[str] = None, - aspeech: Optional[bool] = None, + headers: dict | None = None, + custom_llm_provider: str | None = None, + aspeech: bool | None = None, **kwargs, -) -> Union[HttpxBinaryResponseContent, Coroutine[Any, Any, HttpxBinaryResponseContent]]: +) -> HttpxBinaryResponseContent | Coroutine[Any, Any, HttpxBinaryResponseContent]: user = kwargs.get("user", None) - litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None) + litellm_call_id: str | None = kwargs.get("litellm_call_id", None) proxy_server_request = kwargs.get("proxy_server_request", None) extra_headers = kwargs.get("extra_headers", None) model_info = kwargs.get("model_info", None) @@ -7907,11 +7883,7 @@ def speech( }, custom_llm_provider=custom_llm_provider, ) - response: Union[ - HttpxBinaryResponseContent, - Coroutine[Any, Any, HttpxBinaryResponseContent], - None, - ] = None + response: HttpxBinaryResponseContent | Coroutine[Any, Any, HttpxBinaryResponseContent] | None = None if custom_llm_provider == "openai" or custom_llm_provider in litellm.openai_compatible_providers: if voice is None or not (isinstance(voice, str)): raise litellm.BadRequestError( @@ -8018,7 +7990,7 @@ def speech( or get_secret("AZURE_API_KEY") ) # type: ignore - azure_ad_token: Optional[str] = optional_params.get("extra_body", {}).pop( # type: ignore + azure_ad_token: str | None = optional_params.get("extra_body", {}).pop( # type: ignore "azure_ad_token", None ) or get_secret("AZURE_AD_TOKEN") azure_ad_token_provider = kwargs.get("azure_ad_token_provider", None) @@ -8205,7 +8177,7 @@ def speech( litellm_params_dict["api_key"] = api_key # Convert voice to string if it's a dict (minimax handler expects Optional[str]) - voice_str: Optional[str] = None + voice_str: str | None = None if isinstance(voice, str): voice_str = voice elif isinstance(voice, dict): @@ -8269,8 +8241,8 @@ def speech( async def ahealth_check( model_params: dict, mode: str | None = "chat", - prompt: Optional[str] = None, - input: Optional[List] = None, + prompt: str | None = None, + input: list | None = None, ): """ Support health checks for different providers. Return remaining rate limit, etc. @@ -8308,7 +8280,7 @@ async def ahealth_check( ) ######################################################### try: - model: Optional[str] = model_params.get("model", None) + model: str | None = model_params.get("model", None) if model is None: raise Exception("model not set") @@ -8397,7 +8369,7 @@ def config_completion(**kwargs): ) -def stream_chunk_builder_text_completion(chunks: list, messages: Optional[List] = None) -> TextCompletionResponse: +def stream_chunk_builder_text_completion(chunks: list, messages: list | None = None) -> TextCompletionResponse: id = chunks[0]["id"] object = chunks[0]["object"] created = chunks[0]["created"] @@ -8453,11 +8425,11 @@ def stream_chunk_builder_text_completion(chunks: list, messages: Optional[List] def stream_chunk_builder( chunks: list, - messages: Optional[list] = None, + messages: list | None = None, start_time=None, end_time=None, logging_obj: Optional["Logging"] = None, -) -> Optional[Union[ModelResponse, TextCompletionResponse]]: +) -> ModelResponse | TextCompletionResponse | None: try: if chunks is None: raise litellm.APIError( @@ -8488,7 +8460,7 @@ def stream_chunk_builder( # Fast path for the common text-only streaming case: # avoid repeated multi-pass list scans over chunks. - simple_content_parts: List[str] = [] + simple_content_parts: list[str] = [] is_simple_text_stream = True for chunk in chunks: if len(chunk["choices"]) == 0: @@ -8499,7 +8471,7 @@ def stream_chunk_builder( if isinstance(delta_obj, dict): delta = delta_obj elif hasattr(delta_obj, "model_dump"): - delta = cast(Dict[str, Any], delta_obj.model_dump()) + delta = cast(dict[str, Any], delta_obj.model_dump()) else: delta = {} @@ -8675,7 +8647,7 @@ def stream_chunk_builder( ] if len(provider_specific_chunks) > 0: - combined_provider_fields: Dict[str, Any] = {} + combined_provider_fields: dict[str, Any] = {} for chunk in provider_specific_chunks: fields = chunk["choices"][0]["delta"]["provider_specific_fields"] if isinstance(fields, dict): @@ -8740,11 +8712,11 @@ def stream_chunk_builder( async def acount_tokens( model: str, - messages: Optional[List[Dict[str, Any]]] = None, - tools: Optional[List[Dict[str, Any]]] = None, - system: Optional[str] = None, - api_key: Optional[str] = None, - api_base: Optional[str] = None, + messages: list[dict[str, Any]] | None = None, + tools: list[dict[str, Any]] | None = None, + system: str | None = None, + api_key: str | None = None, + api_base: str | None = None, ) -> "TokenCountResponse": """ Count tokens for a given model and messages using provider-specific APIs. @@ -8786,7 +8758,7 @@ async def acount_tokens( api_base = dynamic_api_base # Build deployment dict for the token counter - deployment: Dict[str, Any] = { + deployment: dict[str, Any] = { "litellm_params": { "model": model, "api_key": api_key, @@ -8837,10 +8809,10 @@ async def acount_tokens( # Cache for encoding to avoid repeated __getattr__ calls -_encoding_cache: Optional[Any] = None +_encoding_cache: tiktoken.Encoding | None = None -def _get_encoding(): +def _get_encoding() -> tiktoken.Encoding: """Get encoding, loading it lazily if needed.""" global _encoding_cache if _encoding_cache is None: @@ -8848,6 +8820,7 @@ def _get_encoding(): # Access via module to trigger __getattr__ if not cached _encoding_cache = sys.modules[__name__].encoding + assert _encoding_cache is not None return _encoding_cache diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index a94a75fdfa3..d277a90835f 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -18,9 +18,17 @@ import os import re import secrets import traceback -from collections.abc import Mapping +from collections.abc import Callable, Mapping, Sequence from datetime import datetime, timedelta, timezone -from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, cast +from typing import ( + TYPE_CHECKING, + Any, + Literal, + Optional, + Protocol, + cast, + overload, +) import fastapi import yaml @@ -131,8 +139,100 @@ from litellm.types.utils import ( TeamUIKeyGenerationConfig, ) +if TYPE_CHECKING: + from prisma.models import LiteLLM_Config as PrismaConfig + from prisma.models import ( + LiteLLM_DeletedVerificationToken as PrismaDeletedVerificationToken, + ) + from prisma.models import ( + LiteLLM_DeprecatedVerificationToken as PrismaDeprecatedVerificationToken, + ) + from prisma.models import LiteLLM_UserTable as PrismaUserTable + from prisma.models import LiteLLM_VerificationToken as PrismaVerificationToken -async def _check_custom_key_allowed(custom_key_value: Optional[str]) -> None: + +class _ConfigTableClient(Protocol): + async def find_many(self) -> "Sequence[PrismaConfig]": ... + + async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> "PrismaConfig": ... + + +class _DeletedVerificationTokenTableClient(Protocol): + async def find_many( + self, + where: Mapping[str, object] | None = None, + skip: int | None = None, + take: int | None = None, + order: Mapping[str, str] | Sequence[Mapping[str, str]] | None = None, + ) -> "Sequence[PrismaDeletedVerificationToken]": ... + + async def count(self, where: Mapping[str, object] | None = None) -> int: ... + + async def create_many(self, data: Sequence[Mapping[str, object]]) -> int: ... + + +class _DeprecatedVerificationTokenTableClient(Protocol): + async def upsert( + self, where: Mapping[str, object], data: Mapping[str, object] + ) -> "PrismaDeprecatedVerificationToken": ... + + +class _UserTableClient(Protocol): + async def find_unique( + self, where: Mapping[str, object], include: Mapping[str, object] | None = None + ) -> "PrismaUserTable | None": ... + + +class _VerificationTokenTableClient(Protocol): + async def find_unique( + self, where: Mapping[str, object], include: Mapping[str, object] | None = None + ) -> "PrismaVerificationToken | None": ... + + async def find_many( + self, + where: Mapping[str, object] | None = None, + skip: int | None = None, + take: int | None = None, + order: Mapping[str, str] | Sequence[Mapping[str, str]] | None = None, + include: Mapping[str, object] | None = None, + ) -> "Sequence[PrismaVerificationToken]": ... + + async def find_first(self, where: Mapping[str, object]) -> "PrismaVerificationToken | None": ... + + async def update( + self, where: Mapping[str, object], data: Mapping[str, object] + ) -> "PrismaVerificationToken | None": ... + + +@overload +def _table(repository: ConfigRepository) -> _ConfigTableClient: ... +@overload +def _table(repository: DeletedVerificationTokenRepository) -> _DeletedVerificationTokenTableClient: ... +@overload +def _table(repository: DeprecatedVerificationTokenRepository) -> _DeprecatedVerificationTokenTableClient: ... +@overload +def _table(repository: UserRepository) -> _UserTableClient: ... +@overload +def _table(repository: VerificationTokenRepository) -> _VerificationTokenTableClient: ... +def _table( + repository: ConfigRepository + | DeletedVerificationTokenRepository + | DeprecatedVerificationTokenRepository + | UserRepository + | VerificationTokenRepository, +) -> object: + prisma_table: object = repository.table # any-ok: repository.table crosses into the untyped prisma client wrapper + return prisma_table + + +class _CustomAuthDecision(BaseModel): + """Validated shape of the dict returned by a user-supplied custom key generate/update hook.""" + + decision: bool = True + message: str = "Authentication Failed - Custom Auth Rule" + + +async def _check_custom_key_allowed(custom_key_value: str | None) -> None: """Raise 403 if custom API keys are disabled and a custom key was provided.""" if custom_key_value is None: return @@ -150,7 +250,7 @@ def _is_team_key(data: Union[GenerateKeyRequest, LiteLLM_VerificationToken]): return data.team_id is not None -def _get_user_in_team(team_table: LiteLLM_TeamTableCachedObj, user_id: Optional[str]) -> Optional[Member]: +def _get_user_in_team(team_table: LiteLLM_TeamTableCachedObj, user_id: str | None) -> Member | None: if user_id is None: return None for member in team_table.members_with_roles: @@ -178,8 +278,8 @@ def _calculate_key_rotation_time(rotation_interval: str) -> datetime: def _set_key_rotation_fields( data: dict, auto_rotate: bool, - rotation_interval: Optional[str], - existing_key_alias: Optional[str] = None, + rotation_interval: str | None, + existing_key_alias: str | None = None, ) -> None: """ Helper function to set rotation fields in key data if auto_rotate is enabled. @@ -214,8 +314,8 @@ def _set_key_rotation_fields( def _is_allowed_to_make_key_request( user_api_key_dict: UserAPIKeyAuth, - user_id: Optional[str], - team_id: Optional[str], + user_id: str | None, + team_id: str | None, ) -> bool: """ Assert user only creates/updates keys for themselves @@ -241,7 +341,7 @@ def _is_allowed_to_make_key_request( def _team_key_operation_team_member_check( - assigned_user_id: Optional[str], + assigned_user_id: str | None, team_table: LiteLLM_TeamTableCachedObj, user_api_key_dict: UserAPIKeyAuth, team_key_generation: TeamUIKeyGenerationConfig, @@ -286,7 +386,7 @@ def _team_key_operation_team_member_check( return True -def _key_generation_required_param_check(data: GenerateKeyRequest, required_params: Optional[List[str]]): +def _key_generation_required_param_check(data: GenerateKeyRequest, required_params: list[str] | None): if required_params is None: return True @@ -340,7 +440,7 @@ def _team_key_generation_check( def _personal_key_membership_check( user_api_key_dict: UserAPIKeyAuth, - personal_key_generation: Optional[PersonalUIKeyGenerationConfig], + personal_key_generation: PersonalUIKeyGenerationConfig | None, ): if personal_key_generation is None or "allowed_user_roles" not in personal_key_generation: return True @@ -355,8 +455,8 @@ def _personal_key_membership_check( def _object_permission_to_dict( - object_permission: Optional[LiteLLM_ObjectPermissionBase], -) -> Optional[ObjectPermissionDict]: + object_permission: LiteLLM_ObjectPermissionBase | None, +) -> ObjectPermissionDict | None: if object_permission is None: return None return cast(ObjectPermissionDict, object_permission.model_dump(exclude_unset=True)) @@ -391,7 +491,7 @@ def _personal_key_generation_check(user_api_key_dict: UserAPIKeyAuth, data: Gene def key_generation_check( - team_table: Optional[LiteLLM_TeamTableCachedObj], + team_table: LiteLLM_TeamTableCachedObj | None, user_api_key_dict: UserAPIKeyAuth, data: GenerateKeyRequest, route: KeyManagementRoutes, @@ -432,9 +532,9 @@ def key_generation_check( def common_key_access_checks( user_api_key_dict: UserAPIKeyAuth, data: Union[GenerateKeyRequest, UpdateKeyRequest], - llm_router: Optional[Router], + llm_router: Router | None, premium_user: bool, - user_id: Optional[str] = None, + user_id: str | None = None, ) -> Literal[True]: """ Check if user is allowed to make a key request, for this key @@ -489,8 +589,8 @@ _NON_ADMIN_SAFE_ALLOWED_ROUTES_PRESETS = frozenset({"llm_api_routes", "info_rout def _validate_caller_can_change_key_ownership( - data: Optional[BaseModel], - existing_key_row: Any, + data: BaseModel | None, + existing_key_row: LiteLLM_VerificationToken, user_api_key_dict: UserAPIKeyAuth, ) -> None: """ @@ -523,7 +623,7 @@ def _validate_caller_can_change_key_ownership( status_code=403, detail="Non-admin users cannot remove the user_id from a key.", ) - existing_user_id = getattr(existing_key_row, "user_id", None) + existing_user_id = existing_key_row.user_id if incoming_user_id != existing_user_id: raise HTTPException( status_code=403, @@ -535,7 +635,7 @@ def _validate_caller_can_change_key_ownership( def _check_allowed_routes_caller_permission( - allowed_routes: Optional[list], + allowed_routes: list | None, user_api_key_dict: UserAPIKeyAuth, *, allowed_routes_was_provided: bool = False, @@ -600,11 +700,11 @@ def _check_permissions_caller_permission( def _check_budget_limits_delegation_ceiling( - budget_limits: Optional[List[BudgetLimitEntry]], - delegation_ceiling: Optional[float], + budget_limits: list[BudgetLimitEntry] | None, + delegation_ceiling: float | None, user_api_key_dict: UserAPIKeyAuth, is_ui_session_team_key: bool, - team_table: Optional[LiteLLM_TeamTableCachedObj], + team_table: LiteLLM_TeamTableCachedObj | None, ) -> None: """ Enforce three invariants on `budget_limits`: @@ -651,8 +751,8 @@ def _check_budget_limits_delegation_ceiling( async def validate_team_id_used_in_service_account_request( - team_id: Optional[str], - prisma_client: Optional[PrismaClient], + team_id: str | None, + prisma_client: PrismaClient | None, ): """ Validate team_id is used in the request body for generating a service account key @@ -670,9 +770,7 @@ async def validate_team_id_used_in_service_account_request( ) # check if team_id exists in the database - team = await TeamRepository(prisma_client).table.find_unique( - where={"team_id": team_id}, - ) + team = await TeamRepository(prisma_client).find_by_id(team_id) if team is None: raise HTTPException( status_code=400, @@ -749,8 +847,8 @@ def _enforce_upperbound_key_params( async def _common_key_generation_helper( data: GenerateKeyRequest, user_api_key_dict: UserAPIKeyAuth, - litellm_changed_by: Optional[str], - team_table: Optional[LiteLLM_TeamTableCachedObj], + litellm_changed_by: str | None, + team_table: LiteLLM_TeamTableCachedObj | None, ) -> GenerateKeyResponse: from litellm.proxy.proxy_server import ( litellm_proxy_admin_name, @@ -889,14 +987,14 @@ async def _common_key_generation_helper( ) new_budget = prisma_client.jsonify_object(budget_row.json(exclude_none=True)) - _budget = await BudgetRepository(prisma_client).table.create( + _budget = await BudgetRepository(prisma_client).create( data={ **new_budget, # type: ignore "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name, } ) - _budget_id = getattr(_budget, "budget_id", None) + _budget_id = _budget.budget_id # ADD METADATA FIELDS # Set Management Endpoint Metadata Fields @@ -1095,12 +1193,12 @@ async def _common_key_generation_helper( def _check_key_model_specific_limits( - keys: List[LiteLLM_VerificationToken], + keys: list[LiteLLM_VerificationToken], data: Union[GenerateKeyRequest, UpdateKeyRequest], - entity_rpm_limit: Optional[int], - entity_tpm_limit: Optional[int], - entity_model_rpm_limit_dict: Dict[str, int], - entity_model_tpm_limit_dict: Dict[str, int], + entity_rpm_limit: int | None, + entity_tpm_limit: int | None, + entity_model_rpm_limit_dict: dict[str, int], + entity_model_tpm_limit_dict: dict[str, int], entity_type: str, # "team" or "organization" ) -> None: """ @@ -1117,8 +1215,8 @@ def _check_key_model_specific_limits( return # get total model specific tpm/rpm limit - model_specific_rpm_limit: Dict[str, int] = {} - model_specific_tpm_limit: Dict[str, int] = {} + model_specific_rpm_limit: dict[str, int] = {} + model_specific_tpm_limit: dict[str, int] = {} for key in keys: if key.metadata.get("model_rpm_limit", None) is not None: @@ -1166,10 +1264,10 @@ def _check_key_model_specific_limits( def _check_key_rpm_tpm_limits( - keys: List[LiteLLM_VerificationToken], + keys: list[LiteLLM_VerificationToken], data: Union[GenerateKeyRequest, UpdateKeyRequest], - entity_rpm_limit: Optional[int], - entity_tpm_limit: Optional[int], + entity_rpm_limit: int | None, + entity_tpm_limit: int | None, entity_type: str, # "team" or "organization" ) -> None: """ @@ -1204,7 +1302,7 @@ def _check_key_rpm_tpm_limits( def check_team_key_model_specific_limits( - keys: List[LiteLLM_VerificationToken], + keys: list[LiteLLM_VerificationToken], team_table: LiteLLM_TeamTableCachedObj, data: Union[GenerateKeyRequest, UpdateKeyRequest], ) -> None: @@ -1229,7 +1327,7 @@ def check_team_key_model_specific_limits( def check_team_key_rpm_tpm_limits( - keys: List[LiteLLM_VerificationToken], + keys: list[LiteLLM_VerificationToken], team_table: LiteLLM_TeamTableCachedObj, data: Union[GenerateKeyRequest, UpdateKeyRequest], ) -> None: @@ -1261,9 +1359,7 @@ async def _check_team_key_limits( # calculate allocated tpm/rpm limit # check if specified tpm/rpm limit is greater than allocated tpm/rpm limit - keys = await VerificationTokenRepository(prisma_client).table.find_many( - where={"team_id": team_table.team_id}, - ) + keys = await VerificationTokenRepository(prisma_client).find_by_team_id(team_table.team_id) # Exclude the key being updated to avoid double-counting its limits. # data.key may be a raw key (sk-...) or a pre-hashed token_id. if isinstance(data, UpdateKeyRequest) and data.key is not None: @@ -1331,7 +1427,7 @@ async def _check_project_key_limits( def check_org_key_model_specific_limits( - keys: List[LiteLLM_VerificationToken], + keys: list[LiteLLM_VerificationToken], org_table: LiteLLM_OrganizationTable, data: Union[GenerateKeyRequest, UpdateKeyRequest], ) -> None: @@ -1364,7 +1460,7 @@ def check_org_key_model_specific_limits( def check_org_key_rpm_tpm_limits( - keys: List[LiteLLM_VerificationToken], + keys: list[LiteLLM_VerificationToken], org_table: LiteLLM_OrganizationTable, data: Union[GenerateKeyRequest, UpdateKeyRequest], ) -> None: @@ -1405,14 +1501,12 @@ async def _validate_caller_can_assign_key_org( detail="Cannot assign a key to an organization without a user_id on the caller's token", ) - user_row = await UserRepository(prisma_client).table.find_unique( + user_row = await _table(UserRepository(prisma_client)).find_unique( where={"user_id": user_api_key_dict.user_id}, include={"organization_memberships": True}, ) - memberships = getattr(user_row, "organization_memberships", None) if user_row else None - member_org_ids = { - membership.organization_id for membership in (memberships or []) if membership.organization_id is not None - } + memberships = user_row.organization_memberships if user_row else None + member_org_ids = {membership.organization_id for membership in (memberships or [])} if organization_id not in member_org_ids: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, @@ -1443,7 +1537,7 @@ async def _check_org_key_limits( # get all organization keys # calculate allocated tpm/rpm limit # check if specified tpm/rpm limit is greater than allocated tpm/rpm limit - keys = await VerificationTokenRepository(prisma_client).table.find_many( + keys = await VerificationTokenRepository(prisma_client).find_many( where={"organization_id": org_table.organization_id}, ) # Exclude the key being updated to avoid double-counting its limits. @@ -1473,7 +1567,7 @@ async def _check_org_key_limits( async def generate_key_fn( data: GenerateKeyRequest, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - litellm_changed_by: Optional[str] = Header( + litellm_changed_by: str | None = Header( None, description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", ), @@ -1587,11 +1681,12 @@ async def generate_key_fn( if user_custom_key_generate is not None: if inspect.iscoroutinefunction(user_custom_key_generate): - result = await user_custom_key_generate(data) # type: ignore + raw_result = await user_custom_key_generate(data) # type: ignore else: raise ValueError("user_custom_key_generate must be a coroutine") - decision = result.get("decision", True) - message = result.get("message", "Authentication Failed - Custom Auth Rule") + result = _CustomAuthDecision.model_validate(raw_result) + decision = result.decision + message = result.message if not decision: raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=message) @@ -1618,7 +1713,7 @@ async def generate_key_fn( user_api_key_dict.user_id, ) - team_table: Optional[LiteLLM_TeamTableCachedObj] = None + team_table: LiteLLM_TeamTableCachedObj | None = None if data.team_id is not None: try: team_table = await get_team_object( @@ -1683,7 +1778,7 @@ async def generate_key_fn( async def generate_service_account_key_fn( data: GenerateKeyRequest, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - litellm_changed_by: Optional[str] = Header( + litellm_changed_by: str | None = Header( None, description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", ), @@ -1786,14 +1881,15 @@ async def generate_service_account_key_fn( if user_custom_key_generate is not None: if inspect.iscoroutinefunction(user_custom_key_generate): - result = await user_custom_key_generate(data) # type: ignore + raw_result = await user_custom_key_generate(data) # type: ignore else: raise ValueError("user_custom_key_generate must be a coroutine") - decision = result.get("decision", True) - message = result.get("message", "Authentication Failed - Custom Auth Rule") + result = _CustomAuthDecision.model_validate(raw_result) + decision = result.decision + message = result.message if not decision: raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=message) - team_table: Optional[LiteLLM_TeamTableCachedObj] = None + team_table: LiteLLM_TeamTableCachedObj | None = None if data.team_id is not None: try: team_table = await get_team_object( @@ -2005,7 +2101,7 @@ def is_different_team(data: UpdateKeyRequest, existing_key_row: LiteLLM_Verifica return data.team_id != existing_key_row.team_id -def _validate_max_budget(max_budget: Optional[float]) -> None: +def _validate_max_budget(max_budget: float | None) -> None: """ Validate that max_budget is not negative. @@ -2023,7 +2119,7 @@ def _validate_max_budget(max_budget: Optional[float]) -> None: async def _get_and_validate_existing_key( - token: str | None, prisma_client: Optional[PrismaClient], key_alias: str | None = None + token: str | None, prisma_client: PrismaClient | None, key_alias: str | None = None ) -> LiteLLM_VerificationToken: """ Get existing key from database and validate it exists. @@ -2048,9 +2144,7 @@ async def _get_and_validate_existing_key( if token is not None: hashed_token = _hash_token_if_needed(token=token) - existing_key_row: LiteLLM_VerificationToken | None = await VerificationTokenRepository( - prisma_client - ).table.find_unique(where={"token": hashed_token}) + existing_key_row = await VerificationTokenRepository(prisma_client).find_by_id(hashed_token) if existing_key_row is None: raise ProxyException( @@ -2070,9 +2164,7 @@ async def _get_and_validate_existing_key( code=status.HTTP_400_BAD_REQUEST, ) - rows: list[LiteLLM_VerificationToken] = await VerificationTokenRepository(prisma_client).table.find_many( - where={"key_alias": key_alias}, take=2 - ) + rows = await VerificationTokenRepository(prisma_client).find_many(where={"key_alias": key_alias}, take=2) if len(rows) == 0: raise ProxyException( @@ -2109,14 +2201,14 @@ def _resolve_token_to_update(data: UpdateKeyRequest, existing_key_row: LiteLLM_V async def _process_single_key_update( update_key_request: UpdateKeyRequest, user_api_key_dict: UserAPIKeyAuth, - litellm_changed_by: Optional[str], - prisma_client: Optional[PrismaClient], + litellm_changed_by: str | None, + prisma_client: PrismaClient | None, user_api_key_cache: UserApiKeyCache, - proxy_logging_obj: Any, - llm_router: Optional[Router], - user_custom_key_update: Optional[Callable] = None, - existing_key_row: Optional[LiteLLM_VerificationToken] = None, -) -> Dict[str, Any]: + proxy_logging_obj: ProxyLogging, + llm_router: Router | None, + user_custom_key_update: Callable | None = None, + existing_key_row: LiteLLM_VerificationToken | None = None, +) -> dict[str, object]: """ Process a single key update with all validations and checks. @@ -2167,11 +2259,12 @@ async def _process_single_key_update( # Custom key update hook if user_custom_key_update is not None: if inspect.iscoroutinefunction(user_custom_key_update): - result = await user_custom_key_update(update_key_request) + raw_result = await user_custom_key_update(update_key_request) else: raise ValueError("user_custom_key_update must be a coroutine") - decision = result.get("decision", True) - message = result.get("message", "Authentication Failed - Custom Auth Rule") + result = _CustomAuthDecision.model_validate(raw_result) + decision = result.decision + message = result.message if not decision: raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=message) @@ -2179,7 +2272,7 @@ async def _process_single_key_update( _enforce_upperbound_key_params(update_key_request, fill_defaults=False) # Get team object and check team limits if team_id is provided - team_obj: Optional[LiteLLM_TeamTableCachedObj] = None + team_obj: LiteLLM_TeamTableCachedObj | None = None if update_key_request.team_id is not None: team_obj = await get_team_object( team_id=update_key_request.team_id, @@ -2265,11 +2358,11 @@ async def _process_single_key_update( async def _validate_mcp_servers_for_key_update( data: "UpdateKeyRequest", team_obj: Optional["LiteLLM_TeamTableCachedObj"], - existing_key_row: Any, - prisma_client: Any, - user_api_key_cache: Any, + existing_key_row: LiteLLM_VerificationToken, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, is_proxy_admin: bool, -) -> Optional[ObjectPermissionDict]: +) -> ObjectPermissionDict | None: """Validate MCP servers in object_permission against the effective team.""" effective_team_obj = team_obj # If team_id isn't being changed, resolve the existing key's team @@ -2300,16 +2393,22 @@ async def _validate_mcp_servers_for_key_update( return normalized_object_permission -async def _validate_update_key_data( +async def _validate_update_key_data( # noqa: C901 # branchy authorization checks; splitting risks behavior drift in security-critical key-update permission logic data: UpdateKeyRequest, - existing_key_row: Any, + existing_key_row: LiteLLM_VerificationToken, user_api_key_dict: UserAPIKeyAuth, - llm_router: Any, + llm_router: Router | None, premium_user: bool, - prisma_client: Any, - user_api_key_cache: Any, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, ) -> None: """Validate permissions and constraints for key update.""" + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": "Database not connected"}, + ) + # Reject NaN/±inf spend before it can reach the DB / spend counter. validate_finite_spend(data.spend) @@ -2414,8 +2513,8 @@ async def _validate_update_key_data( # _check_key_admin_access that would otherwise require team/org admin status. _key_is_team_key = getattr(existing_key_row, "team_id", None) is not None can_skip_admin_check = (caller_is_creator or _key_is_team_key) and not _is_budget_change - if (not _is_proxy_admin) and prisma_client is not None and not can_skip_admin_check: - hashed_key = existing_key_row.token + if (not _is_proxy_admin) and not can_skip_admin_check: + hashed_key = existing_key_row.token or "" await _check_key_admin_access( user_api_key_dict=user_api_key_dict, hashed_token=hashed_key, @@ -2425,7 +2524,7 @@ async def _validate_update_key_data( ) # Check team limits if key has a team_id (from request or existing key) - team_obj: Optional[LiteLLM_TeamTableCachedObj] = None + team_obj: LiteLLM_TeamTableCachedObj | None = None _team_id_to_check = data.team_id or getattr(existing_key_row, "team_id", None) if _team_id_to_check is not None: team_obj = await get_team_object( @@ -2456,7 +2555,7 @@ async def _validate_update_key_data( ) # Validate key against project limits if project_id is being set - _project_id_to_check = getattr(data, "project_id", None) or getattr(existing_key_row, "project_id", None) + _project_id_to_check = existing_key_row.project_id if _project_id_to_check is not None and (data.models is not None or data.max_budget is not None): await _check_project_key_limits( project_id=_project_id_to_check, @@ -2546,7 +2645,7 @@ async def update_key_fn( request: Request, data: UpdateKeyRequest, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - litellm_changed_by: Optional[str] = Header( + litellm_changed_by: str | None = Header( None, description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", ), @@ -2661,11 +2760,12 @@ async def update_key_fn( # Custom key update hook if user_custom_key_update is not None: if inspect.iscoroutinefunction(user_custom_key_update): - result = await user_custom_key_update(data) + raw_result = await user_custom_key_update(data) else: raise ValueError("user_custom_key_update must be a coroutine") - decision = result.get("decision", True) - message = result.get("message", "Authentication Failed - Custom Auth Rule") + result = _CustomAuthDecision.model_validate(raw_result) + decision = result.decision + message = result.message if not decision: raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=message) @@ -2767,7 +2867,7 @@ async def update_key_fn( async def bulk_update_keys( data: BulkUpdateKeyRequest, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - litellm_changed_by: Optional[str] = Header( + litellm_changed_by: str | None = Header( None, description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", ), @@ -2846,8 +2946,8 @@ async def bulk_update_keys( detail={"error": f"Maximum {MAX_BATCH_SIZE} keys can be updated at once. Found {len(data.keys)} keys."}, ) - successful_updates: List[SuccessfulKeyUpdate] = [] - failed_updates: List[FailedKeyUpdate] = [] + successful_updates: list[SuccessfulKeyUpdate] = [] + failed_updates: list[FailedKeyUpdate] = [] for key_update_item in data.keys: try: @@ -2925,7 +3025,7 @@ async def bulk_update_keys( def _build_failed_team_key_update( token: str, exception: Exception, - existing_key_row: Optional[LiteLLM_VerificationToken], + existing_key_row: LiteLLM_VerificationToken | None, ) -> FailedKeyUpdate: """Normalize an exception from the per-key update loop into a FailedKeyUpdate.""" if isinstance(exception, HTTPException): @@ -2939,14 +3039,10 @@ def _build_failed_team_key_update( else: error_message = str(exception) - key_info: Optional[Dict[str, Any]] = None + key_info: dict[str, object] | None = None if existing_key_row is not None: - if hasattr(existing_key_row, "model_dump"): - key_info = existing_key_row.model_dump() - elif hasattr(existing_key_row, "dict"): - key_info = existing_key_row.dict() - if key_info: - key_info.pop("token", None) + key_info = existing_key_row.model_dump() + key_info.pop("token", None) return FailedKeyUpdate(key=token, key_info=key_info, failed_reason=error_message) @@ -2961,7 +3057,7 @@ def _build_failed_team_key_update( async def bulk_update_team_keys( data: BulkUpdateTeamKeysRequest, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - litellm_changed_by: Optional[str] = Header( + litellm_changed_by: str | None = Header( None, description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", ), @@ -3009,7 +3105,7 @@ async def bulk_update_team_keys( # `blocked` is Boolean? with no default; `/key/generate` writes NULL. Prisma's `NOT` # excludes NULLs, so explicitly OR `false` with `null` to include them. now = datetime.now(timezone.utc) - existing_keys = await VerificationTokenRepository(prisma_client).table.find_many( + existing_keys = await VerificationTokenRepository(prisma_client).find_many( where={ "team_id": data.team_id, "AND": [ @@ -3027,7 +3123,7 @@ async def bulk_update_team_keys( "error": f"Team {data.team_id} has more than {MAX_BATCH_SIZE} keys. Use `key_ids` to update in batches of {MAX_BATCH_SIZE}." }, ) - requested_tokens = [row.token for row in existing_keys] + requested_tokens = [row.token for row in existing_keys if row.token is not None] else: if data.key_ids is None or len(data.key_ids) == 0: raise HTTPException( @@ -3045,7 +3141,7 @@ async def bulk_update_team_keys( seen_hashes.add(h) requested_tokens.append(k) hashed_key_ids.append(h) - existing_keys = await VerificationTokenRepository(prisma_client).table.find_many( + existing_keys = await VerificationTokenRepository(prisma_client).find_many( where={"team_id": data.team_id, "token": {"in": hashed_key_ids}} ) @@ -3081,8 +3177,8 @@ async def bulk_update_team_keys( existing_by_token = {row.token: row for row in existing_keys} update_field_dict = data.update_fields.model_dump(exclude_unset=True) - successful_updates: List[SuccessfulKeyUpdate] = [] - failed_updates: List[FailedKeyUpdate] = [] + successful_updates: list[SuccessfulKeyUpdate] = [] + failed_updates: list[FailedKeyUpdate] = [] for token in requested_tokens: db_token = _hash_token_if_needed(token) @@ -3209,7 +3305,7 @@ async def validate_key_team_change( async def delete_key_fn( data: KeyRequest, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - litellm_changed_by: Optional[str] = Header( + litellm_changed_by: str | None = Header( None, description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", ), @@ -3380,7 +3476,7 @@ async def _build_model_max_budget_usage( include_in_schema=False, ) async def info_key_fn_v2( - data: Optional[KeyRequest] = None, + data: KeyRequest | None = None, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -3416,7 +3512,7 @@ async def info_key_fn_v2( # Resolve key_aliases to tokens so we never pass token=None (unbounded query) tokens_to_query = list(data.keys) if data.keys else [] if data.key_aliases: - alias_rows = await VerificationTokenRepository(prisma_client).table.find_many( + alias_rows = await _table(VerificationTokenRepository(prisma_client)).find_many( where={"key_alias": {"in": data.key_aliases}}, include={"litellm_budget_table": True}, ) @@ -3465,7 +3561,7 @@ async def info_key_fn_v2( @router.get("/key/info", tags=["key management"], dependencies=[Depends(user_api_key_auth)]) @management_endpoint_wrapper async def info_key_fn( - key: Optional[str] = fastapi.Query(default=None, description="Key in the request parameters"), + key: str | None = fastapi.Query(default=None, description="Key in the request parameters"), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -3498,11 +3594,11 @@ async def info_key_fn( # default to using Auth token if no key is passed in key = key or user_api_key_dict.api_key - hashed_key: Optional[str] = key + hashed_key: str | None = key if key is not None: hashed_key = _hash_token_if_needed(token=key) key_info = await VerificationTokenRepository(prisma_client).table.find_unique( - where={"token": hashed_key}, # type: ignore + where={"token": hashed_key}, include={"litellm_budget_table": True}, ) if key_info is None: @@ -3554,9 +3650,7 @@ async def info_key_fn( raise handle_exception_on_proxy(e) -def _check_model_access_group( - models: Optional[List[str]], llm_router: Optional[Router], premium_user: bool -) -> Literal[True]: +def _check_model_access_group(models: list[str] | None, llm_router: Router | None, premium_user: bool) -> Literal[True]: """ if is_model_access_group is True + is_wildcard_route is True, check if user is a premium user @@ -3582,63 +3676,62 @@ def _check_model_access_group( async def generate_key_helper_fn( request_type: Literal["user", "key"], # identifies if this request is from /user/new or /key/generate - duration: Optional[str] = None, + duration: str | None = None, models: list = [], aliases: dict = {}, config: dict = {}, spend: float = 0.0, - key_max_budget: Optional[float] = None, # key_max_budget is used to Budget Per key - key_budget_duration: Optional[str] = None, - budget_id: Optional[float] = None, # budget id <-> LiteLLM_BudgetTable - soft_budget: Optional[float] = None, # soft_budget is used to set soft Budgets Per user - max_budget: Optional[float] = None, # max_budget is used to Budget Per user - blocked: Optional[bool] = None, - budget_duration: Optional[str] = None, # max_budget is used to Budget Per user - token: Optional[str] = None, - key: Optional[ - str - ] = None, # dev-friendly alt param for 'token'. Exposed on `/key/generate` for setting key value yourself. - user_id: Optional[str] = None, - user_alias: Optional[str] = None, - team_id: Optional[str] = None, - agent_id: Optional[str] = None, - user_email: Optional[str] = None, - user_role: Optional[str] = None, - max_parallel_requests: Optional[int] = None, - metadata: Optional[dict] = {}, - tpm_limit: Optional[int] = None, - rpm_limit: Optional[int] = None, + key_max_budget: float | None = None, # key_max_budget is used to Budget Per key + key_budget_duration: str | None = None, + budget_id: float | None = None, # budget id <-> LiteLLM_BudgetTable + soft_budget: float | None = None, # soft_budget is used to set soft Budgets Per user + max_budget: float | None = None, # max_budget is used to Budget Per user + blocked: bool | None = None, + budget_duration: str | None = None, # max_budget is used to Budget Per user + token: str | None = None, + key: str + | None = None, # dev-friendly alt param for 'token'. Exposed on `/key/generate` for setting key value yourself. + user_id: str | None = None, + user_alias: str | None = None, + team_id: str | None = None, + agent_id: str | None = None, + user_email: str | None = None, + user_role: str | None = None, + max_parallel_requests: int | None = None, + metadata: dict | None = {}, + tpm_limit: int | None = None, + rpm_limit: int | None = None, query_type: Literal["insert_data", "update_data"] = "insert_data", - update_key_values: Optional[dict] = None, - key_alias: Optional[str] = None, - allowed_cache_controls: Optional[list] = [], - permissions: Optional[dict] = {}, - model_max_budget: Optional[dict] = {}, - budget_fallbacks: Optional[dict] = None, - model_rpm_limit: Optional[dict] = None, - model_tpm_limit: Optional[dict] = None, - mcp_rpm_limit: Optional[dict] = None, - tag_rpm_limit: Optional[dict] = None, - guardrails: Optional[list] = None, - policies: Optional[list] = None, - prompts: Optional[list] = None, - teams: Optional[list] = None, - organization_id: Optional[str] = None, - project_id: Optional[str] = None, - table_name: Optional[Literal["key", "user"]] = None, - send_invite_email: Optional[bool] = None, - created_by: Optional[str] = None, - updated_by: Optional[str] = None, - allowed_routes: Optional[list] = None, + update_key_values: dict | None = None, + key_alias: str | None = None, + allowed_cache_controls: list | None = [], + permissions: dict | None = {}, + model_max_budget: dict | None = {}, + budget_fallbacks: dict | None = None, + model_rpm_limit: dict | None = None, + model_tpm_limit: dict | None = None, + mcp_rpm_limit: dict | None = None, + tag_rpm_limit: dict | None = None, + guardrails: list | None = None, + policies: list | None = None, + prompts: list | None = None, + teams: list | None = None, + organization_id: str | None = None, + project_id: str | None = None, + table_name: Literal["key", "user"] | None = None, + send_invite_email: bool | None = None, + created_by: str | None = None, + updated_by: str | None = None, + allowed_routes: list | None = None, key_type: str | None = None, - sso_user_id: Optional[str] = None, - object_permission_id: Optional[str] = None, # object_permission_id <-> LiteLLM_ObjectPermissionTable - object_permission: Optional[LiteLLM_ObjectPermissionBase] = None, - auto_rotate: Optional[bool] = None, - rotation_interval: Optional[str] = None, - router_settings: Optional[dict] = None, - access_group_ids: Optional[list] = None, - budget_limits: Optional[list] = None, # multiple concurrent budget windows + sso_user_id: str | None = None, + object_permission_id: str | None = None, # object_permission_id <-> LiteLLM_ObjectPermissionTable + object_permission: LiteLLM_ObjectPermissionBase | None = None, + auto_rotate: bool | None = None, + rotation_interval: str | None = None, + router_settings: dict | None = None, + access_group_ids: list | None = None, + budget_limits: list | None = None, # multiple concurrent budget windows ): from litellm.proxy.proxy_server import premium_user, prisma_client @@ -3669,7 +3762,7 @@ async def generate_key_helper_fn( reset_at = get_budget_reset_time(budget_duration=budget_duration) # Initialize reset_at for each budget window - budget_limits_json: Optional[str] = None + budget_limits_json: str | None = None if budget_limits: initialized_windows = [] for window in budget_limits: @@ -3804,7 +3897,7 @@ async def generate_key_helper_fn( saved_token["permissions"] = json.loads(saved_token["permissions"]) if isinstance(saved_token["model_max_budget"], str): saved_token["model_max_budget"] = json.loads(saved_token["model_max_budget"]) - router_settings = cast(Optional[dict], saved_token.get("router_settings")) + router_settings = cast(dict | None, saved_token.get("router_settings")) if router_settings is not None and isinstance(router_settings, str): try: saved_token["router_settings"] = yaml.safe_load(router_settings) @@ -3989,11 +4082,11 @@ async def can_modify_verification_token( async def delete_verification_tokens( - tokens: List, + tokens: list, user_api_key_cache: UserApiKeyCache, user_api_key_dict: UserAPIKeyAuth, - litellm_changed_by: Optional[str] = None, -) -> Tuple[Optional[Dict], List[LiteLLM_VerificationToken]]: + litellm_changed_by: str | None = None, +) -> tuple[dict | None, list[LiteLLM_VerificationToken]]: """ Helper that deletes the list of tokens from the database @@ -4014,13 +4107,13 @@ async def delete_verification_tokens( """ from litellm.proxy.proxy_server import prisma_client - failed_tokens: List = [] + failed_tokens: list = [] try: if prisma_client: tokens = [_hash_token_if_needed(token=key) for key in tokens] - _keys_being_deleted: List[LiteLLM_VerificationToken] = await VerificationTokenRepository( - prisma_client - ).table.find_many(where={"token": {"in": tokens}}) + _keys_being_deleted = await VerificationTokenRepository(prisma_client).find_many( + where={"token": {"in": tokens}} + ) if len(_keys_being_deleted) == 0: raise HTTPException( @@ -4085,16 +4178,16 @@ async def delete_verification_tokens( def _transform_verification_tokens_to_deleted_records( - keys: List[LiteLLM_VerificationToken], + keys: list[LiteLLM_VerificationToken], user_api_key_dict: UserAPIKeyAuth, - litellm_changed_by: Optional[str] = None, -) -> List[Dict[str, Any]]: + litellm_changed_by: str | None = None, +) -> list[dict[str, object]]: """Transform verification tokens into deleted token records ready for persistence.""" if not keys: return [] deleted_at = datetime.now(timezone.utc) - records = [] + records: list[dict[str, object]] = [] for key in keys: key_payload = key.model_dump() deleted_record = LiteLLM_DeletedVerificationToken.model_validate( @@ -4106,7 +4199,7 @@ def _transform_verification_tokens_to_deleted_records( "litellm_changed_by": litellm_changed_by, } ) - record = deleted_record.model_dump() + record: dict[str, object] = deleted_record.model_dump() # Map org_id to organization_id (model uses org_id, but schema expects organization_id) org_id_value = record.pop("org_id", None) @@ -4141,20 +4234,20 @@ def _transform_verification_tokens_to_deleted_records( async def _save_deleted_verification_token_records( - records: List[Dict[str, Any]], + records: list[dict[str, object]], prisma_client: PrismaClient, ) -> None: """Save deleted verification token records to the database.""" if not records: return - await DeletedVerificationTokenRepository(prisma_client).table.create_many(data=records) + await _table(DeletedVerificationTokenRepository(prisma_client)).create_many(data=records) async def _persist_deleted_verification_tokens( - keys: List[LiteLLM_VerificationToken], + keys: list[LiteLLM_VerificationToken], prisma_client: PrismaClient, user_api_key_dict: UserAPIKeyAuth, - litellm_changed_by: Optional[str] = None, + litellm_changed_by: str | None = None, ) -> None: """Persist deleted verification token records by transforming and saving them.""" records = _transform_verification_tokens_to_deleted_records( @@ -4169,13 +4262,13 @@ async def _persist_deleted_verification_tokens( async def delete_key_aliases( - key_aliases: List[str], + key_aliases: list[str], user_api_key_cache: UserApiKeyCache, prisma_client: PrismaClient, user_api_key_dict: UserAPIKeyAuth, - litellm_changed_by: Optional[str] = None, -) -> Tuple[Optional[Dict], List[LiteLLM_VerificationToken]]: - _keys_being_deleted = await VerificationTokenRepository(prisma_client).table.find_many( + litellm_changed_by: str | None = None, +) -> tuple[dict | None, list[LiteLLM_VerificationToken]]: + _keys_being_deleted = await VerificationTokenRepository(prisma_client).find_many( where={"key_alias": {"in": key_aliases}} ) @@ -4212,7 +4305,7 @@ async def _rotate_master_key( from litellm.proxy.proxy_server import proxy_config try: - models: Optional[List] = await ModelRepository(prisma_client).table.find_many() + models = await ModelRepository(prisma_client).find_all() except Exception: models = None # 2. process model table @@ -4263,7 +4356,7 @@ async def _rotate_master_key( ) if encrypted_env_vars: - await ConfigRepository(prisma_client).table.update( + await _table(ConfigRepository(prisma_client)).update( where={"param_name": "environment_variables"}, data={"param_value": prisma.Json(encrypted_env_vars)}, # type: ignore[attr-defined] ) @@ -4424,7 +4517,7 @@ async def check_encryption_endpoint( return {"status": "success", "report": report.as_dict()} -async def get_new_token(data: Optional[RegenerateKeyRequest]) -> str: +async def get_new_token(data: RegenerateKeyRequest | None) -> str: if data and data.new_key is not None: # Reject custom key values if disabled by admin await _check_custom_key_allowed(data.new_key) @@ -4450,7 +4543,7 @@ async def _insert_deprecated_key( prisma_client: "PrismaClient", old_token_hash: str, new_token_hash: str, - grace_period: Optional[str], + grace_period: str | None, ) -> None: """ Insert old key into deprecated table so it remains valid during grace period. @@ -4481,7 +4574,7 @@ async def _insert_deprecated_key( try: revoke_at = datetime.now(timezone.utc) + timedelta(seconds=grace_seconds) - await DeprecatedVerificationTokenRepository(prisma_client).table.upsert( + await _table(DeprecatedVerificationTokenRepository(prisma_client)).upsert( where={"token": old_token_hash}, data={ "create": { @@ -4513,9 +4606,9 @@ async def _execute_virtual_key_regeneration( key_in_db: LiteLLM_VerificationToken, hashed_api_key: str, key: str, - data: Optional[RegenerateKeyRequest], + data: RegenerateKeyRequest | None, user_api_key_dict: UserAPIKeyAuth, - litellm_changed_by: Optional[str], + litellm_changed_by: str | None, user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, ) -> GenerateKeyResponse: @@ -4571,9 +4664,9 @@ async def _execute_virtual_key_regeneration( grace_period=data.grace_period if data else None, ) - updated_token = await VerificationTokenRepository(prisma_client).table.update( + updated_token = await _table(VerificationTokenRepository(prisma_client)).update( where={"token": hashed_api_key}, - data=update_data, # type: ignore + data=update_data, ) updated_token_dict = dict(updated_token) if updated_token is not None else {} updated_token_dict["key"] = new_token @@ -4611,14 +4704,14 @@ async def _execute_virtual_key_regeneration( ) @management_endpoint_wrapper async def regenerate_key_fn( - key: Optional[str] = None, - data: Optional[RegenerateKeyRequest] = None, + key: str | None = None, + data: RegenerateKeyRequest | None = None, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - litellm_changed_by: Optional[str] = Header( + litellm_changed_by: str | None = Header( None, description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", ), -) -> Optional[GenerateKeyResponse]: +) -> GenerateKeyResponse | None: """ Regenerate an existing API key while optionally updating its parameters. @@ -4772,9 +4865,7 @@ async def regenerate_key_fn( else: hashed_api_key = hash_token(key) - _key_in_db = await VerificationTokenRepository(prisma_client).table.find_unique( - where={"token": hashed_api_key}, - ) + _key_in_db = await VerificationTokenRepository(prisma_client).find_by_id(hashed_api_key) if _key_in_db is None: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, @@ -4803,7 +4894,7 @@ async def regenerate_key_fn( ) if data is not None and (data.access_group_ids or data.object_permission is not None): - regenerate_team_table: Optional[LiteLLM_TeamTableCachedObj] = None + regenerate_team_table: LiteLLM_TeamTableCachedObj | None = None if _key_in_db.team_id is not None: regenerate_team_table = await get_team_object( team_id=_key_in_db.team_id, @@ -4876,7 +4967,7 @@ async def regenerate_key_fn( async def _check_proxy_or_team_admin_for_key( - key_in_db: LiteLLM_VerificationToken, + key_in_db: "PrismaVerificationToken", user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient, user_api_key_cache: UserApiKeyCache, @@ -4904,7 +4995,7 @@ async def _check_proxy_or_team_admin_for_key( ) -def _validate_reset_spend_value(reset_to: object, key_in_db: LiteLLM_VerificationToken) -> float: +def _validate_reset_spend_value(reset_to: object, key_in_db: "PrismaVerificationToken") -> float: if not isinstance(reset_to, (int, float)): raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, @@ -4928,7 +5019,7 @@ def _validate_reset_spend_value(reset_to: object, key_in_db: LiteLLM_Verificatio max_budget = key_in_db.max_budget if key_in_db.litellm_budget_table is not None: - budget_max_budget = getattr(key_in_db.litellm_budget_table, "max_budget", None) + budget_max_budget = key_in_db.litellm_budget_table.max_budget if budget_max_budget is not None: if max_budget is None or budget_max_budget < max_budget: max_budget = budget_max_budget @@ -4952,11 +5043,11 @@ async def reset_key_spend_fn( key: str, data: ResetSpendRequest, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - litellm_changed_by: Optional[str] = Header( + litellm_changed_by: str | None = Header( None, description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", ), -) -> Dict[str, Any]: +) -> dict[str, object]: try: from litellm.proxy.proxy_server import ( hash_token, @@ -4976,7 +5067,7 @@ async def reset_key_spend_fn( else: hashed_api_key = hash_token(key) - _key_in_db = await VerificationTokenRepository(prisma_client).table.find_unique( + _key_in_db = await _table(VerificationTokenRepository(prisma_client)).find_unique( where={"token": hashed_api_key}, include={"litellm_budget_table": True}, ) @@ -4996,10 +5087,7 @@ async def reset_key_spend_fn( user_api_key_cache=user_api_key_cache, ) - updated_key = await VerificationTokenRepository(prisma_client).table.update( - where={"token": hashed_api_key}, - data={"spend": reset_to}, - ) + updated_key = await VerificationTokenRepository(prisma_client).update_spend(hashed_api_key, reset_to) if updated_key is None: raise HTTPException( @@ -5050,13 +5138,13 @@ async def reset_key_spend_fn( async def validate_key_list_check( user_api_key_dict: UserAPIKeyAuth, - user_id: Optional[str], - team_id: Optional[str], - organization_id: Optional[str], - key_alias: Optional[str], - key_hash: Optional[str], + user_id: str | None, + team_id: str | None, + organization_id: str | None, + key_alias: str | None, + key_hash: str | None, prisma_client: PrismaClient, -) -> Optional[LiteLLM_UserTable]: +) -> LiteLLM_UserTable | None: if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: return None @@ -5067,7 +5155,7 @@ async def validate_key_list_check( param="user_id", code=status.HTTP_403_FORBIDDEN, ) - complete_user_info_db_obj: Optional[BaseModel] = await UserRepository(prisma_client).table.find_unique( + complete_user_info_db_obj = await _table(UserRepository(prisma_client)).find_unique( where={"user_id": user_api_key_dict.user_id}, include={"organization_memberships": True}, ) @@ -5140,26 +5228,20 @@ async def validate_key_list_check( async def _fetch_user_team_objects( - complete_user_info: Optional[LiteLLM_UserTable], + complete_user_info: LiteLLM_UserTable | None, prisma_client: PrismaClient, -) -> List[LiteLLM_TeamTable]: +) -> list[LiteLLM_TeamTable]: """Fetch team objects for all teams a user belongs to (single DB query).""" if complete_user_info is None or not complete_user_info.teams: return [] - teams: Optional[List[BaseModel]] = await TeamRepository(prisma_client).table.find_many( - where={"team_id": {"in": complete_user_info.teams}} - ) - if teams is None: - return [] - - return [LiteLLM_TeamTable.model_validate(team.model_dump()) for team in teams] + return await TeamRepository(prisma_client).find_many(where={"team_id": {"in": complete_user_info.teams}}) def _get_admin_team_ids_from_objects( user_api_key_dict: UserAPIKeyAuth, - team_objects: List[LiteLLM_TeamTable], -) -> List[str]: + team_objects: list[LiteLLM_TeamTable], +) -> list[str]: """Filter team objects to those where the user is an admin.""" return [ team.team_id for team in team_objects if _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team) @@ -5168,8 +5250,8 @@ def _get_admin_team_ids_from_objects( def _get_team_ids_with_key_list_permission_from_objects( user_api_key_dict: UserAPIKeyAuth, - team_objects: List[LiteLLM_TeamTable], -) -> List[str]: + team_objects: list[LiteLLM_TeamTable], +) -> list[str]: """Filter team objects to non-admin teams where the caller has /key/list permission via team_member_permissions. These teams should grant the caller full key visibility (same as a team admin), so other members' @@ -5188,8 +5270,8 @@ def _get_team_ids_with_key_list_permission_from_objects( def _get_member_team_ids_from_objects( user_api_key_dict: UserAPIKeyAuth, - team_objects: List[LiteLLM_TeamTable], -) -> List[str]: + team_objects: list[LiteLLM_TeamTable], +) -> list[str]: """Filter team objects to those where the user is a member (any role).""" return [ team.team_id @@ -5202,20 +5284,20 @@ def _get_member_team_ids_from_objects( async def get_admin_team_ids( - complete_user_info: Optional[LiteLLM_UserTable], + complete_user_info: LiteLLM_UserTable | None, user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient, -) -> List[str]: +) -> list[str]: """Get all team IDs where the user is an admin.""" team_objects = await _fetch_user_team_objects(complete_user_info, prisma_client) return _get_admin_team_ids_from_objects(user_api_key_dict, team_objects) async def get_member_team_ids( - complete_user_info: Optional[LiteLLM_UserTable], + complete_user_info: LiteLLM_UserTable | None, user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient, -) -> List[str]: +) -> list[str]: """ Get all team IDs where the user is a member (any role, including admin). @@ -5240,30 +5322,30 @@ async def list_keys( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), page: int = Query(1, description="Page number", ge=1), size: int = Query(10, description="Page size", ge=1, le=100), - user_id: Optional[str] = Query( + user_id: str | None = Query( None, description="Filter keys by user ID. Exact match by default; set substring_matching=true (admin only) for case-insensitive substring matching.", ), - team_id: Optional[str] = Query(None, description="Filter keys by team ID"), - organization_id: Optional[str] = Query(None, description="Filter keys by organization ID"), - key_hash: Optional[str] = Query(None, description="Filter keys by key hash"), - key_alias: Optional[str] = Query( + team_id: str | None = Query(None, description="Filter keys by team ID"), + organization_id: str | None = Query(None, description="Filter keys by organization ID"), + key_hash: str | None = Query(None, description="Filter keys by key hash"), + key_alias: str | None = Query( None, description="Filter keys by key alias. Exact match by default; set substring_matching=true (admin only) for case-insensitive substring matching.", ), return_full_object: bool = Query(False, description="Return full key object"), include_team_keys: bool = Query(False, description="Include all keys for teams that user is an admin of."), include_created_by_keys: bool = Query(False, description="Include keys created by the user"), - sort_by: Optional[str] = Query( + sort_by: str | None = Query( default=None, description="Column to sort by (e.g. 'user_id', 'created_at', 'spend')", ), sort_order: str = Query(default="desc", description="Sort order ('asc' or 'desc')"), - expand: Optional[List[str]] = Query(None, description="Expand related objects (e.g. 'user')"), - status: Optional[str] = Query(None, description="Filter by status (e.g. 'deleted')"), - project_id: Optional[str] = Query(None, description="Filter keys by project ID"), - access_group_id: Optional[str] = Query(None, description="Filter keys by access group ID"), - agent_id: Optional[str] = Query(None, description="Filter keys by agent ID"), + expand: list[str] | None = Query(None, description="Expand related objects (e.g. 'user')"), + status: str | None = Query(None, description="Filter by status (e.g. 'deleted')"), + project_id: str | None = Query(None, description="Filter keys by project ID"), + access_group_id: str | None = Query(None, description="Filter keys by access group ID"), + agent_id: str | None = Query(None, description="Filter keys by agent ID"), substring_matching: bool = Query( False, description="If true (proxy admins only), match user_id/key_alias as case-insensitive substrings instead of exact values. Defaults to false: /key/list matched these exactly before substring search was added, and an exact user_id/key_alias filter must never return another user's keys.", @@ -5421,23 +5503,23 @@ async def list_keys( async def _apply_non_admin_alias_scope( user_api_key_dict: UserAPIKeyAuth, - prisma_client: Any, - query_params: List[Any], - where_parts: List[str], + prisma_client: PrismaClient, + query_params: list[str], + where_parts: list[str], ) -> None: """Append SQL scope conditions so non-admin users only see aliases for keys they own or keys belonging to teams they are members of.""" - scope_conditions: List[str] = [] + scope_conditions: list[str] = [] if user_api_key_dict.user_id: query_params.append(user_api_key_dict.user_id) scope_conditions.append(f"user_id = ${len(query_params)}") # Look up the user's teams from the user table - user_teams: List[str] = [] + user_teams: list[str] = [] if user_api_key_dict.user_id: - user_row = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_api_key_dict.user_id}) + user_row = await UserRepository(prisma_client).find_by_id(user_api_key_dict.user_id) if user_row is not None: - user_teams = getattr(user_row, "teams", []) or [] + user_teams = user_row.teams if user_teams: team_placeholders = ", ".join(f"${len(query_params) + i + 1}" for i in range(len(user_teams))) @@ -5461,9 +5543,9 @@ async def key_aliases( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), page: int = Query(1, ge=1, description="Page number"), size: int = Query(50, ge=1, le=100, description="Page size"), - search: Optional[str] = Query(None, description="Search key aliases (case-insensitive partial match)"), - team_id: Optional[str] = Query(None, description="Filter aliases to keys belonging to this team"), -) -> Dict[str, Any]: + search: str | None = Query(None, description="Search key aliases (case-insensitive partial match)"), + team_id: str | None = Query(None, description="Filter aliases to keys belonging to this team"), +) -> dict[str, object]: """ Lists key aliases with pagination and optional search. @@ -5493,7 +5575,7 @@ async def key_aliases( # support column-level SELECT projection on find_many. # # $1 is always UI_SESSION_TOKEN_TEAM_ID (filters out UI session tokens). - query_params: List[Any] = [UI_SESSION_TOKEN_TEAM_ID] + query_params: list[str] = [UI_SESSION_TOKEN_TEAM_ID] where_parts = [ "key_alias IS NOT NULL", "key_alias != ''", @@ -5534,7 +5616,7 @@ async def key_aliases( f" LIMIT ${limit_idx} OFFSET ${offset_idx}" ) alias_rows = await prisma_client.db.query_raw(aliases_sql, *aliases_params) - aliases: List[str] = [row["key_alias"] for row in alias_rows if row.get("key_alias")] + aliases: list[str] = [row["key_alias"] for row in alias_rows if row.get("key_alias")] total_pages = -(-total_count // size) if total_count > 0 else 0 verbose_proxy_logger.debug( @@ -5569,8 +5651,8 @@ async def key_aliases( ) -def _validate_sort_params(sort_by: Optional[str], sort_order: str) -> Optional[Dict[str, str]]: - order_by: Dict[str, str] = {} +def _validate_sort_params(sort_by: str | None, sort_order: str) -> dict[str, str] | None: + order_by: dict[str, str] = {} if sort_by is None: return None @@ -5601,28 +5683,28 @@ def _validate_sort_params(sort_by: Optional[str], sort_order: str) -> Optional[D return order_by -def _build_expires_where_clause(expires_filter: str, now: datetime) -> dict[str, Any]: +def _build_expires_where_clause(expires_filter: str, now: datetime) -> dict[str, object]: if expires_filter == "expired": return {"AND": [{"expires": {"not": None}}, {"expires": {"lt": now}}]} return {"OR": [{"expires": None}, {"expires": {"gte": now}}]} def _build_key_filter_conditions( - user_id: Optional[str], - team_id: Optional[str], - organization_id: Optional[str], - key_alias: Optional[str], - key_hash: Optional[str], - exclude_team_id: Optional[str], - admin_team_ids: Optional[List[str]], - member_team_ids: Optional[List[str]] = None, + user_id: str | None, + team_id: str | None, + organization_id: str | None, + key_alias: str | None, + key_hash: str | None, + exclude_team_id: str | None, + admin_team_ids: list[str] | None, + member_team_ids: list[str] | None = None, include_created_by_keys: bool = False, - project_id: Optional[str] = None, - access_group_id: Optional[str] = None, - agent_id: Optional[str] = None, + project_id: str | None = None, + access_group_id: str | None = None, + agent_id: str | None = None, use_substring_matching: bool = False, expires_filter: str | None = None, -) -> Dict[str, Union[str, Dict[str, Any], List[Dict[str, Any]]]]: +) -> dict[str, object]: """Build filter conditions for key listing. Visibility rules: @@ -5634,14 +5716,14 @@ def _build_key_filter_conditions( so former members cannot see service accounts they created after leaving. """ # Prepare filter conditions - where: Dict[str, Union[str, Dict[str, Any], List[Dict[str, Any]]]] = {} + where: dict[str, object] = {} where.update(_get_condition_to_filter_out_ui_session_tokens()) # Build the OR conditions for user's keys and admin team keys - or_conditions: List[Dict[str, Any]] = [] + or_conditions: list[dict[str, object]] = [] # Base conditions for user's own keys - user_condition: Dict[str, Any] = {} + user_condition: dict[str, object] = {} if user_id and isinstance(user_id, str): if use_substring_matching: user_condition["user_id"] = { @@ -5740,25 +5822,24 @@ async def _list_key_helper( prisma_client: PrismaClient, page: int, size: int, - user_id: Optional[str], - team_id: Optional[str], - organization_id: Optional[str], - key_alias: Optional[str], - key_hash: Optional[str], - exclude_team_id: Optional[str] = None, + user_id: str | None, + team_id: str | None, + organization_id: str | None, + key_alias: str | None, + key_hash: str | None, + exclude_team_id: str | None = None, return_full_object: bool = False, - admin_team_ids: Optional[List[str]] = None, # New parameter for teams where user is admin - member_team_ids: Optional[ - List[str] - ] = None, # Team IDs where user is a member (any role) - for service account visibility + admin_team_ids: list[str] | None = None, # New parameter for teams where user is admin + member_team_ids: list[str] + | None = None, # Team IDs where user is a member (any role) - for service account visibility include_created_by_keys: bool = False, - sort_by: Optional[str] = None, + sort_by: str | None = None, sort_order: str = "desc", - expand: Optional[List[str]] = None, - status: Optional[str] = None, - project_id: Optional[str] = None, - access_group_id: Optional[str] = None, - agent_id: Optional[str] = None, + expand: list[str] | None = None, + status: str | None = None, + project_id: str | None = None, + access_group_id: str | None = None, + agent_id: str | None = None, use_substring_matching: bool = False, expires_filter: str | None = None, ) -> KeyListResponseObject: @@ -5806,7 +5887,7 @@ async def _list_key_helper( verbose_proxy_logger.debug(f"Pagination: skip={skip}, take={size}") - order_by: Optional[Dict[str, str]] = ( + order_by: dict[str, str] | None = ( _validate_sort_params(sort_by, sort_order) if sort_by is not None and isinstance(sort_by, str) else None ) @@ -5815,10 +5896,10 @@ async def _list_key_helper( # Fetch keys with pagination if use_deleted_table: - keys = await DeletedVerificationTokenRepository(prisma_client).table.find_many( - where=where, # type: ignore - skip=skip, # type: ignore - take=size, # type: ignore + keys = await _table(DeletedVerificationTokenRepository(prisma_client)).find_many( + where=where, + skip=skip, + take=size, order=( order_by if order_by @@ -5829,10 +5910,10 @@ async def _list_key_helper( ), ) else: - keys = await VerificationTokenRepository(prisma_client).table.find_many( - where=where, # type: ignore - skip=skip, # type: ignore - take=size, # type: ignore + keys = await _table(VerificationTokenRepository(prisma_client)).find_many( + where=where, + skip=skip, + take=size, order=( order_by if order_by @@ -5848,13 +5929,9 @@ async def _list_key_helper( # Get total count of keys if use_deleted_table: - total_count = await DeletedVerificationTokenRepository(prisma_client).table.count( - where=where # type: ignore - ) + total_count = await _table(DeletedVerificationTokenRepository(prisma_client)).count(where=where) else: - total_count = await VerificationTokenRepository(prisma_client).table.count( - where=where # type: ignore - ) + total_count = await VerificationTokenRepository(prisma_client).count(where=where) verbose_proxy_logger.debug(f"Total count of keys: {total_count}") @@ -5868,18 +5945,13 @@ async def _list_key_helper( created_by_ids = [key.created_by for key in keys if key.created_by] all_ids = list(set(user_ids + created_by_ids)) # Remove duplicates if all_ids: - users = await UserRepository(prisma_client).table.find_many(where={"user_id": {"in": all_ids}}) + users = await UserRepository(prisma_client).find_many(where={"user_id": {"in": all_ids}}) user_map = {user.user_id: user for user in users} # Prepare response - key_list: List[Union[str, UserAPIKeyAuth, LiteLLM_DeletedVerificationToken]] = [] + key_list: list[Union[str, UserAPIKeyAuth, LiteLLM_DeletedVerificationToken]] = [] for key in keys: - # Convert Prisma model to dict (supports both Pydantic v1 and v2) - try: - key_dict = key.model_dump() - except Exception: - # Fallback for Pydantic v1 compatibility - key_dict = key.dict() + key_dict = key.model_dump() # Attach object_permission if object_permission_id is set (only for non-deleted keys) if not use_deleted_table: key_dict = await attach_object_permission_to_dict(key_dict, prisma_client) @@ -5887,10 +5959,7 @@ async def _list_key_helper( # Include user information if expand includes "user" if expand and "user" in expand: if key.user_id and key.user_id in user_map: - try: - key_dict["user"] = user_map[key.user_id].model_dump() - except Exception: - key_dict["user"] = user_map[key.user_id].dict() + key_dict["user"] = user_map[key.user_id].model_dump() if key.created_by and key.created_by in user_map: created_by_user = user_map[key.created_by] key_dict["created_by_user"] = { @@ -5917,7 +5986,7 @@ async def _list_key_helper( ) -def _get_condition_to_filter_out_ui_session_tokens() -> Dict[str, Any]: +def _get_condition_to_filter_out_ui_session_tokens() -> dict[str, object]: """ Condition to filter out UI session tokens """ @@ -5932,7 +6001,7 @@ def _get_condition_to_filter_out_ui_session_tokens() -> Dict[str, Any]: async def _check_key_admin_access( user_api_key_dict: UserAPIKeyAuth, hashed_token: str, - prisma_client: Any, + prisma_client: PrismaClient, user_api_key_cache: UserApiKeyCache, route: str, ) -> None: @@ -5951,7 +6020,7 @@ async def _check_key_admin_access( return # Look up the target key to find its team - target_key_row = await VerificationTokenRepository(prisma_client).table.find_unique(where={"token": hashed_token}) + target_key_row = await VerificationTokenRepository(prisma_client).find_by_id(hashed_token) if target_key_row is None: raise HTTPException( status_code=404, @@ -5987,11 +6056,11 @@ async def block_key( data: BlockKeyRequest, http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - litellm_changed_by: Optional[str] = Header( + litellm_changed_by: str | None = Header( None, description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", ), -) -> Optional[LiteLLM_VerificationToken]: +) -> LiteLLM_VerificationToken | None: """ Block an Virtual key from making any requests. @@ -6047,7 +6116,7 @@ async def block_key( ) # Check if the key exists before trying to block it - existing_record = await VerificationTokenRepository(prisma_client).table.find_unique(where={"token": hashed_token}) + existing_record = await VerificationTokenRepository(prisma_client).find_by_id(hashed_token) if existing_record is None: raise ProxyException( message="Key not found.", @@ -6077,10 +6146,7 @@ async def block_key( ) ) - record = await VerificationTokenRepository(prisma_client).table.update( - where={"token": hashed_token}, - data={"blocked": True}, # type: ignore - ) + record = await VerificationTokenRepository(prisma_client).block_token(hashed_token) ## UPDATE KEY CACHE - invalidate so next read re-fetches from DB await _delete_cache_key_object( @@ -6098,7 +6164,7 @@ async def unblock_key( data: BlockKeyRequest, http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - litellm_changed_by: Optional[str] = Header( + litellm_changed_by: str | None = Header( None, description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", ), @@ -6158,7 +6224,7 @@ async def unblock_key( ) # Check if the key exists before trying to unblock it - existing_record = await VerificationTokenRepository(prisma_client).table.find_unique(where={"token": hashed_token}) + existing_record = await VerificationTokenRepository(prisma_client).find_by_id(hashed_token) if existing_record is None: raise ProxyException( message="Key not found.", @@ -6188,10 +6254,7 @@ async def unblock_key( ) ) - record = await VerificationTokenRepository(prisma_client).table.update( - where={"token": hashed_token}, - data={"blocked": False}, # type: ignore - ) + record = await VerificationTokenRepository(prisma_client).unblock_token(hashed_token) ## UPDATE KEY CACHE - invalidate so next read re-fetches from DB await _delete_cache_key_object( @@ -6295,7 +6358,7 @@ async def key_health( async def _can_user_query_key_info( user_api_key_dict: UserAPIKeyAuth, - key: Optional[str], + key: str | None, key_info: LiteLLM_VerificationToken, ) -> bool: """ @@ -6322,7 +6385,7 @@ async def _can_user_query_key_info( async def test_key_logging( user_api_key_dict: UserAPIKeyAuth, request: Request, - key_logging: List[Dict[str, Any]], + key_logging: list[dict[str, Any]], ) -> LoggingCallbackStatus: """ Test the key-based logging @@ -6337,7 +6400,7 @@ async def test_key_logging( from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request from litellm.proxy.proxy_server import general_settings, proxy_config - logging_callbacks: List[str] = [] + logging_callbacks: list[str] = [] for callback in key_logging: if callback.get("callback_name") is not None: logging_callbacks.append(callback["callback_name"]) @@ -6398,7 +6461,7 @@ async def test_key_logging( _KEY_ALIAS_PATTERN = re.compile(r"^[a-zA-Z0-9][a-zA-Z0-9_\-/\.@]{0,253}[a-zA-Z0-9]$") -def _validate_key_alias_format(key_alias: Optional[str]) -> None: +def _validate_key_alias_format(key_alias: str | None) -> None: """ Validate the format of the key_alias. @@ -6442,16 +6505,16 @@ def _validate_key_alias_format(key_alias: Optional[str]) -> None: async def _enforce_unique_key_alias( - key_alias: Optional[str], - prisma_client: Any, - existing_key_token: Optional[str] = None, + key_alias: str | None, + prisma_client: PrismaClient | None, + existing_key_token: str | None = None, ) -> None: """ Helper to enforce unique key aliases across all keys. Args: key_alias (Optional[str]): The key alias to check - prisma_client (Any): Prisma client instance + prisma_client (Optional[PrismaClient]): Prisma client instance existing_key_token (Optional[str]): ID of existing key being updated, to exclude from uniqueness check (The Admin UI passes key_alias, in all Edit key requests. So we need to be sure that if we find a key with the same alias, it's not the same key we're updating) @@ -6459,12 +6522,12 @@ async def _enforce_unique_key_alias( ProxyException: If key alias already exists on a different key """ if key_alias is not None and prisma_client is not None: - where_clause: dict[str, Any] = {"key_alias": key_alias} + where_clause: dict[str, object] = {"key_alias": key_alias} if existing_key_token: # Exclude the current key from the uniqueness check where_clause["NOT"] = {"token": existing_key_token} - existing_key = await VerificationTokenRepository(prisma_client).table.find_first(where=where_clause) + existing_key = await _table(VerificationTokenRepository(prisma_client)).find_first(where=where_clause) if existing_key is not None: raise ProxyException( message=f"Key with alias '{key_alias}' already exists. Unique key aliases across all keys are required.", @@ -6474,7 +6537,7 @@ async def _enforce_unique_key_alias( ) -def validate_model_max_budget(model_max_budget: Optional[Dict]) -> None: +def validate_model_max_budget(model_max_budget: dict | None) -> None: """ Validate the model_max_budget is GenericBudgetConfigType + enforce user has an enterprise license diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index fa01f43d049..07cd5485f2e 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -13,8 +13,9 @@ import asyncio import json import math import traceback +from collections.abc import Mapping, Sequence from datetime import datetime, timezone -from typing import Annotated, Any, Dict, List, Mapping, Optional, Tuple, Union, cast +from typing import Annotated, Protocol, TypedDict, cast import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Request, status @@ -30,18 +31,20 @@ from litellm.proxy._types import ( BlockTeamRequest, CommonProxyErrors, DeleteTeamRequest, + LiteLLM_AccessGroupTable, LiteLLM_AuditLogs, + LiteLLM_BudgetTableFull, LiteLLM_DeletedTeamTable, LiteLLM_ManagementEndpoint_MetadataFields, LiteLLM_ManagementEndpoint_MetadataFields_Premium, LiteLLM_ModelTable, + LiteLLM_OrganizationMembershipTable, LiteLLM_OrganizationTable, LiteLLM_OrganizationTableWithMembers, LiteLLM_TeamMembership, LiteLLM_TeamTable, LiteLLM_TeamTableCachedObj, LiteLLM_UserTable, - LiteLLM_VerificationToken, LitellmTableNames, LitellmUserRoles, Member, @@ -78,6 +81,7 @@ from litellm.proxy.auth.auth_checks import ( from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.callback_utils import encrypt_callback_vars from litellm.proxy.common_utils.json_merge_patch import apply_json_merge_patch +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.management_endpoints.common_utils import ( _check_passthrough_routes_caller_permission, _is_user_org_admin_for_team, @@ -106,7 +110,7 @@ from litellm.proxy.management_helpers.utils import ( add_new_member, management_endpoint_wrapper, ) -from litellm.proxy.utils import PrismaClient, handle_exception_on_proxy +from litellm.proxy.utils import PrismaClient, ProxyLogging, handle_exception_on_proxy from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.table_repositories import ( @@ -152,9 +156,9 @@ def _sanitize_for_log(value: object) -> str: async def _refresh_cached_team( - team_row: Any, - user_api_key_cache: Any, - proxy_logging_obj: Any, + team_row: LiteLLM_TeamTable, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, ) -> None: """ Refresh the in-memory cached team object after a DB write. @@ -215,7 +219,7 @@ class TeamMemberBudgetHandler: SYSTEM_MANAGED_METADATA_KEYS = ("team_member_budget_id",) @staticmethod - def strip_system_managed_metadata_keys(metadata: Optional[dict]) -> None: + def strip_system_managed_metadata_keys(metadata: dict | None) -> None: """Remove server-owned metadata keys from a caller-supplied dict.""" if not isinstance(metadata, dict): return @@ -224,10 +228,10 @@ class TeamMemberBudgetHandler: @staticmethod def should_create_budget( - team_member_budget: Optional[float] = None, - team_member_rpm_limit: Optional[int] = None, - team_member_tpm_limit: Optional[int] = None, - team_member_budget_duration: Optional[str] = None, + team_member_budget: float | None = None, + team_member_rpm_limit: int | None = None, + team_member_tpm_limit: int | None = None, + team_member_budget_duration: str | None = None, ) -> bool: """Check if any team member limits are provided""" return any( @@ -241,13 +245,13 @@ class TeamMemberBudgetHandler: @staticmethod async def create_team_member_budget_table( - data: Union[NewTeamRequest, LiteLLM_TeamTable], + data: NewTeamRequest | LiteLLM_TeamTable, new_team_data_json: dict, user_api_key_dict: UserAPIKeyAuth, - team_member_budget: Optional[float] = None, - team_member_rpm_limit: Optional[int] = None, - team_member_tpm_limit: Optional[int] = None, - team_member_budget_duration: Optional[str] = None, + team_member_budget: float | None = None, + team_member_rpm_limit: int | None = None, + team_member_tpm_limit: int | None = None, + team_member_budget_duration: str | None = None, ) -> dict: """Create team member budget table with provided limits""" from litellm.proxy._types import BudgetNewRequest @@ -295,10 +299,10 @@ class TeamMemberBudgetHandler: team_table: LiteLLM_TeamTable, user_api_key_dict: UserAPIKeyAuth, updated_kv: dict, - team_member_budget: Optional[float] = None, - team_member_rpm_limit: Optional[int] = None, - team_member_tpm_limit: Optional[int] = None, - team_member_budget_duration: Optional[str] = None, + team_member_budget: float | None = None, + team_member_rpm_limit: int | None = None, + team_member_tpm_limit: int | None = None, + team_member_budget_duration: str | None = None, ) -> dict: """Upsert team member budget table with provided limits""" from litellm.proxy._types import BudgetNewRequest @@ -396,7 +400,7 @@ class TeamMemberBudgetHandler: @staticmethod async def backfill_team_member_budget_entries( team_id: str, - members_with_roles: List[Union[Member, dict]], + members_with_roles: Sequence[Member | dict], team_member_budget_id: str, prisma_client: PrismaClient, ) -> None: @@ -415,23 +419,20 @@ class TeamMemberBudgetHandler: return # Batch-fetch existing memberships for this team (avoids N+1 queries) - existing_memberships = await TeamMembershipRepository(prisma_client).table.find_many(where={"team_id": team_id}) + existing_memberships: Sequence[LiteLLM_TeamMembership] = await TeamMembershipRepository( + prisma_client + ).table.find_many(where={"team_id": team_id}) existing_user_ids = {m.user_id for m in existing_memberships} # Identify members with no existing membership row. # members_with_roles may contain Member instances or raw dicts depending # on how the team was fetched/deserialized. - missing = [] - for m in members_with_roles: - user_id = m.get("user_id") if isinstance(m, dict) else m.user_id - if user_id is not None and user_id not in existing_user_ids: - missing.append( - { - "team_id": team_id, - "user_id": user_id, - "budget_id": team_member_budget_id, - } - ) + member_user_ids = (m.get("user_id") if isinstance(m, dict) else m.user_id for m in members_with_roles) + missing: tuple[Mapping[str, str], ...] = tuple( + {"team_id": team_id, "user_id": user_id, "budget_id": team_member_budget_id} + for user_id in member_user_ids + if user_id is not None and user_id not in existing_user_ids + ) if missing: await TeamMembershipRepository(prisma_client).table.create_many( @@ -448,7 +449,7 @@ class TeamMemberBudgetHandler: # Heal existing membership rows that predate the team_member_budget # configuration: populate budget_id where it is currently NULL. # Rows with an explicit budget_id (per-member override) are left alone. - updated = await TeamMembershipRepository(prisma_client).table.update_many( + updated: int = await TeamMembershipRepository(prisma_client).table.update_many( where={"team_id": team_id, "budget_id": None}, data={"budget_id": team_member_budget_id}, ) @@ -461,7 +462,7 @@ class TeamMemberBudgetHandler: ) -def _get_default_team_param(field: str) -> Any: +def _get_default_team_param(field: str) -> float | str | int | tuple[str, ...] | None: """ Returns a default value for the given field from litellm.default_team_params config. Returns None if no default is configured. @@ -471,6 +472,7 @@ def _get_default_team_param(field: str) -> Any: default_params = litellm.default_team_params if default_params is None: return None + value: object if isinstance(default_params, dict): value = default_params.get(field) else: @@ -479,8 +481,10 @@ def _get_default_team_param(field: str) -> Any: return None # Convert enum values in lists to strings if isinstance(value, list): - return [v.value if hasattr(v, "value") else v for v in value] - return value + return tuple(v.value if hasattr(v, "value") else v for v in value) + if isinstance(value, (float, str, int)): + return value + return None def _is_available_team(team_id: str, user_api_key_dict: UserAPIKeyAuth) -> bool: @@ -492,11 +496,11 @@ def _is_available_team(team_id: str, user_api_key_dict: UserAPIKeyAuth) -> bool: async def get_all_team_memberships( - prisma_client: PrismaClient, team_ids: List[str], user_id: Optional[str] = None -) -> List[LiteLLM_TeamMembership]: + prisma_client: PrismaClient, team_ids: list[str], user_id: str | None = None +) -> list[LiteLLM_TeamMembership]: """Get all team memberships for a given user""" ## GET ALL MEMBERSHIPS ## - where_obj: Dict[str, Dict[str, List[str]]] = {"team_id": {"in": team_ids}} + where_obj: dict[str, dict[str, list[str]]] = {"team_id": {"in": team_ids}} if user_id is not None: where_obj["user_id"] = {"in": [user_id]} # if user_id is None: @@ -504,12 +508,12 @@ async def get_all_team_memberships( # else: # where_obj = {"user_id": str(user_id), "team_id": {"in": team_id}} - team_memberships = await TeamMembershipRepository(prisma_client).table.find_many( + team_memberships: list[LiteLLM_TeamMembership] = await TeamMembershipRepository(prisma_client).table.find_many( where=where_obj, include={"litellm_budget_table": True}, ) - returned_tm: List[LiteLLM_TeamMembership] = [] + returned_tm: list[LiteLLM_TeamMembership] = [] for tm in team_memberships: returned_tm.append(LiteLLM_TeamMembership.model_validate(tm.model_dump())) @@ -517,12 +521,12 @@ async def get_all_team_memberships( def _check_team_model_specific_limits( - teams: List[LiteLLM_TeamTable], - data: Union[NewTeamRequest, UpdateTeamRequest], - entity_rpm_limit: Optional[int], - entity_tpm_limit: Optional[int], - entity_model_rpm_limit_dict: Dict[str, int], - entity_model_tpm_limit_dict: Dict[str, int], + teams: list[LiteLLM_TeamTable], + data: NewTeamRequest | UpdateTeamRequest, + entity_rpm_limit: int | None, + entity_tpm_limit: int | None, + entity_model_rpm_limit_dict: dict[str, int], + entity_model_tpm_limit_dict: dict[str, int], entity_type: str, # "organization" ) -> None: """ @@ -539,8 +543,8 @@ def _check_team_model_specific_limits( return # get total model specific tpm/rpm limit - model_specific_rpm_limit: Dict[str, int] = {} - model_specific_tpm_limit: Dict[str, int] = {} + model_specific_rpm_limit: dict[str, int] = {} + model_specific_tpm_limit: dict[str, int] = {} for team in teams: if team.metadata and team.metadata.get("model_rpm_limit", None) is not None: @@ -588,10 +592,10 @@ def _check_team_model_specific_limits( def _check_team_rpm_tpm_limits( - teams: List[LiteLLM_TeamTable], - data: Union[NewTeamRequest, UpdateTeamRequest], - entity_rpm_limit: Optional[int], - entity_tpm_limit: Optional[int], + teams: list[LiteLLM_TeamTable], + data: NewTeamRequest | UpdateTeamRequest, + entity_rpm_limit: int | None, + entity_tpm_limit: int | None, entity_type: str, # "organization" ) -> None: """ @@ -626,9 +630,9 @@ def _check_team_rpm_tpm_limits( def check_org_team_model_specific_limits( - teams: List[LiteLLM_TeamTable], + teams: list[LiteLLM_TeamTable], org_table: LiteLLM_OrganizationTable, - data: Union[NewTeamRequest, UpdateTeamRequest], + data: NewTeamRequest | UpdateTeamRequest, ) -> None: """ Check if the organization team is allocating model specific limits. If so, raise an error if we're overallocating. @@ -660,9 +664,9 @@ def check_org_team_model_specific_limits( def check_org_team_rpm_tpm_limits( - teams: List[LiteLLM_TeamTable], + teams: list[LiteLLM_TeamTable], org_table: LiteLLM_OrganizationTable, - data: Union[NewTeamRequest, UpdateTeamRequest], + data: NewTeamRequest | UpdateTeamRequest, ) -> None: """ Check if the organization team is allocating rpm/tpm limits. If so, raise an error if we're overallocating. @@ -686,7 +690,7 @@ def check_org_team_rpm_tpm_limits( async def _check_org_team_limits( org_table: LiteLLM_OrganizationTable, - data: Union[NewTeamRequest, UpdateTeamRequest], + data: NewTeamRequest | UpdateTeamRequest, prisma_client: PrismaClient, ) -> None: """ @@ -766,15 +770,10 @@ async def _check_org_team_limits( # calculate allocated tpm/rpm limit # check if specified tpm/rpm limit is greater than allocated tpm/rpm limit - teams = await TeamRepository(prisma_client).table.find_many( + team_objs = await TeamRepository(prisma_client).find_many( where={"organization_id": org_table.organization_id}, ) - # Convert teams to LiteLLM_TeamTable objects - team_objs: List[LiteLLM_TeamTable] = [] - for team in teams: - team_objs.append(LiteLLM_TeamTable.model_validate(team.model_dump())) - check_org_team_model_specific_limits( teams=team_objs, org_table=org_table, @@ -788,10 +787,10 @@ async def _check_org_team_limits( async def _check_user_team_limits( - data: Union[NewTeamRequest, UpdateTeamRequest], + data: NewTeamRequest | UpdateTeamRequest, user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient, - user_api_key_cache: Any, + user_api_key_cache: UserApiKeyCache, ) -> None: """ Enforce the caller's personal limits when CREATING a standalone team. @@ -862,7 +861,7 @@ async def _check_user_team_limits( def _check_team_budget_update_authority( data: UpdateTeamRequest, user_api_key_dict: UserAPIKeyAuth, - existing_team_max_budget: Optional[float], + existing_team_max_budget: float | None, ) -> None: """ Restrict who can grow a standalone team's spend ceiling on /team/update. @@ -919,7 +918,7 @@ async def new_team( data: NewTeamRequest, http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - litellm_changed_by: Optional[str] = Header( + litellm_changed_by: str | None = Header( None, description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", ), @@ -1052,7 +1051,7 @@ async def new_team( ) # Check if license is over limit - total_teams = await TeamRepository(prisma_client).table.count() + total_teams = await TeamRepository(prisma_client).count() if total_teams and _license_check.is_team_count_over_limit(team_count=total_teams): raise HTTPException( status_code=403, @@ -1154,7 +1153,7 @@ async def new_team( created_by=user_api_key_dict.user_id or litellm_proxy_admin_name, updated_by=user_api_key_dict.user_id or litellm_proxy_admin_name, ) - model_dict = await ModelTableRepository(prisma_client).table.create( + model_dict: LiteLLM_ModelTable = await ModelTableRepository(prisma_client).table.create( {**litellm_modeltable.json(exclude_none=True)} # type: ignore ) # type: ignore @@ -1236,7 +1235,7 @@ async def new_team( complete_team_data.budget_limits = initialized_windows # type: ignore[assignment] ## Add Team Member Budget Table - members_with_roles: List[Member] = [] + members_with_roles: list[Member] = [] if complete_team_data.members_with_roles is not None: members_with_roles = complete_team_data.members_with_roles complete_team_data.members_with_roles = [] @@ -1311,7 +1310,7 @@ async def _create_team_update_audit_log( existing_team_row: LiteLLM_TeamTable, updated_kv: dict, team_id: str, - litellm_changed_by: Optional[str], + litellm_changed_by: str | None, user_api_key_dict: UserAPIKeyAuth, litellm_proxy_admin_name: str, ) -> None: @@ -1358,11 +1357,11 @@ async def _create_team_update_audit_log( async def _update_model_table( data: UpdateTeamRequest, - model_id: Optional[str], + model_id: int | None, prisma_client: PrismaClient, user_api_key_dict: UserAPIKeyAuth, litellm_proxy_admin_name: str, -) -> Optional[str]: +) -> int | None: """ Upsert model table and return the model id """ @@ -1374,6 +1373,7 @@ async def _update_model_table( created_by=user_api_key_dict.user_id or litellm_proxy_admin_name, updated_by=user_api_key_dict.user_id or litellm_proxy_admin_name, ) + model_dict: LiteLLM_ModelTable if model_id is None: model_dict = await ModelTableRepository(prisma_client).table.create( data={**litellm_modeltable.json(exclude_none=True)} # type: ignore @@ -1395,7 +1395,7 @@ async def _update_model_table( async def _auto_add_team_members_to_organization( team: LiteLLM_TeamTable, organization: LiteLLM_OrganizationTableWithMembers, - prisma_client: Any, + prisma_client: PrismaClient, ) -> None: """ When moving a team to an org, ensure all team members are also org members. @@ -1433,11 +1433,11 @@ async def _auto_add_team_members_to_organization( async def fetch_and_validate_organization( organization_id: str, - existing_team_row: Any, - llm_router: Optional[Router], - prisma_client: Any, - user_api_key_dict: Optional[UserAPIKeyAuth] = None, -) -> Any: + existing_team_row: LiteLLM_TeamTable, + llm_router: Router | None, + prisma_client: PrismaClient, + user_api_key_dict: UserAPIKeyAuth | None = None, +) -> LiteLLM_OrganizationTableWithMembers: """ Fetch and validate an organization for team update operations. @@ -1456,7 +1456,9 @@ async def fetch_and_validate_organization( if llm_router is None: raise HTTPException(status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value}) - organization_row = await OrganizationRepository(prisma_client).table.find_unique( + organization_row: LiteLLM_OrganizationTableWithMembers | None = await OrganizationRepository( + prisma_client + ).table.find_unique( where={"organization_id": organization_id}, include={"litellm_budget_table": True, "members": True, "teams": True}, ) @@ -1470,7 +1472,7 @@ async def fetch_and_validate_organization( is_proxy_admin = user_api_key_dict is not None and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN organization = LiteLLM_OrganizationTableWithMembers.model_validate(organization_row.model_dump()) validate_team_org_change( - team=LiteLLM_TeamTable.model_validate(existing_team_row.model_dump()), + team=existing_team_row, organization=organization, llm_router=llm_router, is_proxy_admin=is_proxy_admin, @@ -1478,7 +1480,7 @@ async def fetch_and_validate_organization( if is_proxy_admin: await _auto_add_team_members_to_organization( - team=LiteLLM_TeamTable.model_validate(existing_team_row.model_dump()), + team=existing_team_row, organization=organization, prisma_client=prisma_client, ) @@ -1595,7 +1597,7 @@ async def update_team( data: UpdateTeamRequest, http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - litellm_changed_by: Optional[str] = Header( + litellm_changed_by: str | None = Header( None, description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", ), @@ -1705,7 +1707,7 @@ async def update_team( detail={"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"}, ) - existing_team_row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": data.team_id}) + existing_team_row = await TeamRepository(prisma_client).find_by_id(data.team_id) if existing_team_row is None: raise HTTPException( @@ -1715,7 +1717,7 @@ async def update_team( # Verify caller has access to manage this team await _verify_team_access( - team_obj=LiteLLM_TeamTable.model_validate(existing_team_row.model_dump()), + team_obj=existing_team_row, user_api_key_dict=user_api_key_dict, ) @@ -1733,9 +1735,9 @@ async def update_team( ) if data.max_budget is not None: - existing_soft_budget = getattr(existing_team_row, "soft_budget", None) + existing_soft_budget = existing_team_row.soft_budget soft_budget_to_check = data.soft_budget if data.soft_budget is not None else existing_soft_budget - if soft_budget_to_check is not None and isinstance(soft_budget_to_check, (int, float)): + if soft_budget_to_check is not None: if data.max_budget <= soft_budget_to_check: raise HTTPException( status_code=400, @@ -1751,13 +1753,13 @@ async def update_team( # so without this gate an org-admin could hand their team to any # other org (or capture a team from another org they once # administered into a new destination). - current_org_id = getattr(existing_team_row, "organization_id", None) + current_org_id = existing_team_row.organization_id if ( data.organization_id != current_org_id and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value ): # Is the caller org_admin of the destination org? - caller_memberships = ( + caller_memberships: Sequence[LiteLLM_OrganizationMembershipTable] = ( await OrganizationMembershipRepository(prisma_client).table.find_many( where={ "user_id": user_api_key_dict.user_id, @@ -1795,7 +1797,7 @@ async def update_team( org_id_to_check = ( data.organization_id if data.organization_id is not None else existing_team_row.organization_id ) - if org_id_to_check is not None and isinstance(org_id_to_check, str) and prisma_client is not None: + if org_id_to_check is not None and prisma_client is not None: org_table = await get_org_object( org_id=org_id_to_check, user_api_key_cache=user_api_key_cache, @@ -1909,7 +1911,7 @@ async def update_team( updated_kv["router_settings"] = safe_dumps(updated_kv["router_settings"]) updated_kv = prisma_client.jsonify_team_object(db_data=updated_kv) - team_row: Optional[LiteLLM_TeamTable] = await TeamRepository(prisma_client).table.update( + team_row: LiteLLM_TeamTable | None = await TeamRepository(prisma_client).table.update( where={"team_id": data.team_id}, data=updated_kv, # `object_permission` is included so `_refresh_cached_team` @@ -1962,7 +1964,7 @@ async def patch_team( http_request: Request, user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], litellm_changed_by: Annotated[ - Optional[str], + str | None, Header( description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", ), @@ -2005,7 +2007,7 @@ async def patch_team( patch_fields = data.model_dump(exclude_unset=True, exclude={"team_id"}) if "metadata" in patch_fields: - existing_team_row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id}) + existing_team_row = await TeamRepository(prisma_client).find_by_id(team_id) if existing_team_row is None: raise HTTPException( status_code=404, @@ -2073,13 +2075,13 @@ async def handle_update_object_permission(data_json: dict, existing_team_row: Li def _check_team_member_admin_add( - member: Union[Member, List[Member]], + member: Member | list[Member], premium_user: bool, ): if isinstance(member, Member) and member.role == "admin": if premium_user is not True: raise ValueError(f"Assigning team admins is a premium feature. {CommonProxyErrors.not_premium_user.value}") - elif isinstance(member, List): + elif isinstance(member, list): for m in member: if m.role == "admin": if premium_user is not True: @@ -2089,7 +2091,7 @@ def _check_team_member_admin_add( def team_call_validation_checks( - prisma_client: Optional[PrismaClient], + prisma_client: PrismaClient | None, data: TeamMemberAddRequest, premium_user: bool, ): @@ -2138,7 +2140,7 @@ def team_member_add_duplication_check( # First, populate the invalid_team_members list by checking for duplicates if isinstance(data.member, Member): _check_member_duplication(data.member) - elif isinstance(data.member, List): + elif isinstance(data.member, list): for m in data.member: _check_member_duplication(m) @@ -2249,10 +2251,10 @@ async def _process_team_members( prisma_client: PrismaClient, user_api_key_dict: UserAPIKeyAuth, litellm_proxy_admin_name: str, -) -> Tuple[List[LiteLLM_UserTable], List[LiteLLM_TeamMembership]]: +) -> tuple[list[LiteLLM_UserTable], list[LiteLLM_TeamMembership]]: """Process and add new team members.""" - updated_users: List[LiteLLM_UserTable] = [] - updated_team_memberships: List[LiteLLM_TeamMembership] = [] + updated_users: list[LiteLLM_UserTable] = [] + updated_team_memberships: list[LiteLLM_TeamMembership] = [] default_team_budget_id = ( complete_team_data.metadata.get("team_member_budget_id") if complete_team_data.metadata is not None else None @@ -2289,7 +2291,7 @@ async def _process_team_members( updated_users.append(updated_user) if updated_tm is not None: updated_team_memberships.append(updated_tm) - elif isinstance(data.member, List): + elif isinstance(data.member, list): for m in data.member: try: updated_user, updated_tm = await add_new_member( @@ -2322,7 +2324,7 @@ async def _process_team_members( async def _update_team_members_list( data: TeamMemberAddRequest, complete_team_data: LiteLLM_TeamTable, - updated_users: List[LiteLLM_UserTable], + updated_users: list[LiteLLM_UserTable], ) -> None: """Update the team's members_with_roles list.""" if isinstance(data.member, Member): @@ -2346,7 +2348,7 @@ async def _update_team_members_list( if not member_already_exists: complete_team_data.members_with_roles.append(new_member) - elif isinstance(data.member, List): + elif isinstance(data.member, list): for nm in data.member: if nm.user_id is None and nm.user_email is not None: for user in updated_users: @@ -2372,7 +2374,7 @@ async def _add_team_members_to_team( prisma_client: PrismaClient, user_api_key_dict: UserAPIKeyAuth, litellm_proxy_admin_name: str, -) -> Tuple[LiteLLM_TeamTable, List[LiteLLM_UserTable], List[LiteLLM_TeamMembership]]: +) -> tuple[LiteLLM_TeamTable, list[LiteLLM_UserTable], list[LiteLLM_TeamMembership]]: """Add team members to the team. The members_with_roles reconciliation runs inside a transaction that locks @@ -2487,11 +2489,11 @@ async def _validate_and_populate_member_user_info( # Case 2: Only user_email provided - populate user_id from DB if member.user_email is not None and member.user_id is None: - user_by_email = await UserRepository(prisma_client).table.find_first( + user_by_email_match: LiteLLM_UserTable | None = await UserRepository(prisma_client).table.find_first( where={"user_email": {"equals": member.user_email, "mode": "insensitive"}} ) - if user_by_email is None: + if user_by_email_match is None: # User doesn't exist yet - this is fine, will be created later return member @@ -2511,12 +2513,12 @@ async def _validate_and_populate_member_user_info( ) # Populate user_id - member.user_id = user_by_email.user_id + member.user_id = user_by_email_match.user_id return member # Case 3: Only user_id provided - populate user_email from DB if user exists if member.user_id is not None and member.user_email is None: - user_by_id = await UserRepository(prisma_client).table.find_unique(where={"user_id": member.user_id}) + user_by_id = await UserRepository(prisma_client).find_by_id(member.user_id) if user_by_id is None: # User doesn't exist yet - allow it to pass with user_email as None @@ -2612,7 +2614,7 @@ async def team_member_add( member=data.member, prisma_client=prisma_client, ) - elif isinstance(data.member, List): + elif isinstance(data.member, list): for m in data.member: await _validate_and_populate_member_user_info( member=m, @@ -2649,10 +2651,10 @@ async def team_member_add( def _cleanup_members_with_roles( existing_team_row: LiteLLM_TeamTable, data: TeamMemberDeleteRequest, -) -> Tuple[bool, List[Member]]: +) -> tuple[bool, list[Member]]: """Cleanup members_with_roles list for a team.""" is_member_in_team = False - new_team_members: List[Member] = [] + new_team_members: list[Member] = [] for m in existing_team_row.members_with_roles: if data.user_id is not None and m.user_id is not None and data.user_id == m.user_id: is_member_in_team = True @@ -2707,14 +2709,13 @@ async def team_member_delete( detail={"error": "Either user_id or user_email needs to be passed in"}, ) - _existing_team_row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": data.team_id}) + existing_team_row = await TeamRepository(prisma_client).find_by_id(data.team_id) - if _existing_team_row is None: + if existing_team_row is None: raise HTTPException( status_code=400, detail={"error": "Team id={} does not exist in db".format(data.team_id)}, ) - existing_team_row = LiteLLM_TeamTable.model_validate(_existing_team_row.model_dump()) ## CHECK IF USER IS PROXY ADMIN OR TEAM ADMIN OR ORG ADMIN @@ -2743,9 +2744,9 @@ async def team_member_delete( existing_team_row.members_with_roles = new_team_members - _db_new_team_members: List[dict] = [m.model_dump() for m in new_team_members] + _db_new_team_members: list[dict] = [m.model_dump() for m in new_team_members] - _ = await TeamRepository(prisma_client).table.update( + await TeamRepository(prisma_client).table.update( where={ "team_id": data.team_id, }, @@ -2756,16 +2757,16 @@ async def team_member_delete( ## DELETE TEAM ID from USER ROW, IF EXISTS ## # get user row - key_val = {} - if data.user_id is not None: - key_val["user_id"] = data.user_id - elif data.user_email is not None: - key_val["user_email"] = data.user_email - existing_user_rows = await UserRepository(prisma_client).table.find_many( - where=key_val # type: ignore + key_val = ( + {"user_id": data.user_id} + if data.user_id is not None + else {"user_email": data.user_email} + if data.user_email is not None + else {} ) + existing_user_rows = await UserRepository(prisma_client).find_many(where=key_val) - if existing_user_rows is not None and (isinstance(existing_user_rows, list) and len(existing_user_rows) > 0): + if existing_user_rows: for existing_user in existing_user_rows: team_list = [] if data.team_id in existing_user.teams: @@ -2782,10 +2783,9 @@ async def team_member_delete( user_ids_to_delete = set() if data.user_id is not None: user_ids_to_delete.add(data.user_id) - if existing_user_rows is not None and isinstance(existing_user_rows, list): - for existing_user in existing_user_rows: - if getattr(existing_user, "user_id", None): - user_ids_to_delete.add(existing_user.user_id) + for existing_user in existing_user_rows: + if existing_user.user_id: + user_ids_to_delete.add(existing_user.user_id) for _uid in user_ids_to_delete: await TeamMembershipRepository(prisma_client).table.delete_many( @@ -2799,9 +2799,7 @@ async def team_member_delete( ) # Fetch keys before deletion to persist them - keys_to_delete: List[LiteLLM_VerificationToken] = await VerificationTokenRepository( - prisma_client - ).table.find_many( + keys_to_delete = await VerificationTokenRepository(prisma_client).find_many( where={ "user_id": {"in": list(user_ids_to_delete)}, "team_id": data.team_id, @@ -2835,7 +2833,7 @@ _MEMBER_BUDGET_PATCH_FIELDS = { } -def _build_member_budget_patch(data: TeamMemberUpdateRequest) -> Dict[str, Any]: +def _build_member_budget_patch(data: TeamMemberUpdateRequest) -> dict[str, object]: """Map the budget fields the request actually set (merge-patch: a sent value updates, an explicit null clears, an absent field is left untouched) to their budget-table columns.""" @@ -2847,7 +2845,7 @@ def _build_member_budget_patch(data: TeamMemberUpdateRequest) -> Dict[str, Any]: } -def _validate_budget_duration(budget_duration: Optional[str]) -> None: +def _validate_budget_duration(budget_duration: str | None) -> None: """Reject budget durations that can't be parsed, are non-positive, or overflow date math, so a bad value can't be persisted and later crash the budget reset job.""" @@ -2911,14 +2909,13 @@ async def team_member_update( _validate_budget_duration(data.budget_duration) - _existing_team_row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": data.team_id}) + existing_team_row = await TeamRepository(prisma_client).find_by_id(data.team_id) - if _existing_team_row is None: + if existing_team_row is None: raise HTTPException( status_code=400, detail={"error": "Team id={} does not exist in db".format(data.team_id)}, ) - existing_team_row = LiteLLM_TeamTable.model_validate(_existing_team_row.model_dump()) ## CHECK IF USER IS PROXY ADMIN OR TEAM ADMIN OR ORG ADMIN @@ -2946,7 +2943,7 @@ async def team_member_update( team_table = returned_team_info["team_info"] ## get user id - received_user_id: Optional[str] = None + received_user_id: str | None = None if data.user_id is not None: received_user_id = data.user_id elif data.user_email is not None: @@ -2961,7 +2958,7 @@ async def team_member_update( detail={"error": "User id doesn't exist in team table. Data={}".format(data)}, ) ## find the relevant team membership - identified_budget_id: Optional[str] = None + identified_budget_id: str | None = None for tm in returned_team_info["team_memberships"]: if tm.user_id == received_user_id: identified_budget_id = tm.budget_id @@ -2970,7 +2967,7 @@ async def team_member_update( # If this membership still points at the team's shared default member # budget, _upsert_budget_and_membership will clone-on-write so that the # update only touches this user (not every member sharing the default). - team_default_budget_id: Optional[str] = None + team_default_budget_id: str | None = None if team_table.metadata is not None: raw_default_budget_id = team_table.metadata.get("team_member_budget_id") if isinstance(raw_default_budget_id, str): @@ -2991,7 +2988,7 @@ async def team_member_update( ### update team member role if data.role is not None: - team_members: List[Member] = [] + team_members: list[Member] = [] for member in team_table.members_with_roles: if member.user_id == received_user_id: team_members.append( @@ -3006,7 +3003,7 @@ async def team_member_update( team_table.members_with_roles = team_members - _db_team_members: List[dict] = [m.model_dump() for m in team_members] + _db_team_members: list[dict] = [m.model_dump() for m in team_members] await TeamRepository(prisma_client).table.update( where={"team_id": data.team_id}, data={"members_with_roles": json.dumps(_db_team_members)}, # type: ignore @@ -3025,13 +3022,13 @@ async def team_member_update( def _create_results_from_response( - members: List[Member], + members: list[Member], response: TeamAddMemberResponse, -) -> List[TeamMemberAddResult]: +) -> list[TeamMemberAddResult]: """ Convert TeamAddMemberResponse into individual TeamMemberAddResult objects """ - results: List[TeamMemberAddResult] = [] + results: list[TeamMemberAddResult] = [] for member in members: # Find corresponding updated user @@ -3138,7 +3135,7 @@ async def bulk_team_member_add( }, ) # get all users from the database - all_users_in_db = await UserRepository(prisma_client).table.find_many(order={"created_at": "desc"}) + all_users_in_db = await UserRepository(prisma_client).find_many(order={"created_at": "desc"}) data.members = [ Member( user_id=user.user_id, @@ -3215,7 +3212,7 @@ async def delete_team( data: DeleteTeamRequest, http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - litellm_changed_by: Optional[str] = Header( + litellm_changed_by: str | None = Header( None, description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", ), @@ -3251,20 +3248,17 @@ async def delete_team( raise HTTPException(status_code=400, detail={"error": "No team id passed in"}) # check that all teams passed exist - team_rows: List[LiteLLM_TeamTable] = [] + team_rows: list[LiteLLM_TeamTable] = [] for team_id in data.team_ids: try: - team_row_base: Optional[BaseModel] = await TeamRepository(prisma_client).table.find_unique( - where={"team_id": team_id} - ) - if team_row_base is None: + team_row_pydantic = await TeamRepository(prisma_client).find_by_id(team_id) + if team_row_pydantic is None: raise Exception except Exception: raise HTTPException( status_code=404, detail={"error": f"Team not found, passed team_id={team_id}"}, ) - team_row_pydantic = LiteLLM_TeamTable.model_validate(team_row_base.model_dump()) # Verify caller has access to manage this team await _verify_team_access( @@ -3286,7 +3280,7 @@ async def delete_team( if litellm.store_audit_logs is True: # make an audit log for each team deleted for team_id in data.team_ids: - team_row: Optional[LiteLLM_TeamTable] = await prisma_client.get_data( # type: ignore + team_row: LiteLLM_TeamTable | None = await prisma_client.get_data( # type: ignore team_id=team_id, table_name="team", query_type="find_unique" ) @@ -3323,7 +3317,7 @@ async def delete_team( _persist_deleted_verification_tokens, ) - keys_to_delete: List[LiteLLM_VerificationToken] = await VerificationTokenRepository(prisma_client).table.find_many( + keys_to_delete = await VerificationTokenRepository(prisma_client).find_many( where={"team_id": {"in": data.team_ids}} ) @@ -3375,56 +3369,65 @@ async def delete_team( return deleted_teams -def _transform_teams_to_deleted_records( - teams: List[LiteLLM_TeamTable], +def _transform_single_team_to_deleted_record( + team: LiteLLM_TeamTable, user_api_key_dict: UserAPIKeyAuth, - litellm_changed_by: Optional[str] = None, -) -> List[Dict[str, Any]]: + deleted_at: datetime, + litellm_changed_by: str | None, +) -> Mapping[str, object]: + """Transform a single team into a deleted team record ready for persistence.""" + team_payload = team.model_dump() + deleted_record = LiteLLM_DeletedTeamTable.model_validate( + { + **team_payload, + "deleted_at": deleted_at, + "deleted_by": user_api_key_dict.user_id, + "deleted_by_api_key": user_api_key_dict.api_key, + "litellm_changed_by": litellm_changed_by, + } + ) + record = deleted_record.model_dump() + + for json_field in [ + "members_with_roles", + "metadata", + "model_spend", + "model_max_budget", + "router_settings", + ]: + if json_field in record and record[json_field] is not None: + record[json_field] = json.dumps(record[json_field]) + + for rel_key in ( + "litellm_model_table", + "object_permission", + "id", + "budget_limits", # not in LiteLLM_DeletedTeamTable schema + "default_team_member_models", # not in LiteLLM_DeletedTeamTable schema + ): + record.pop(rel_key, None) + + return record + + +def _transform_teams_to_deleted_records( + teams: list[LiteLLM_TeamTable], + user_api_key_dict: UserAPIKeyAuth, + litellm_changed_by: str | None = None, +) -> tuple[Mapping[str, object], ...]: """Transform teams into deleted team records ready for persistence.""" if not teams: - return [] + return () deleted_at = datetime.now(timezone.utc) - records = [] - for team in teams: - team_payload = team.model_dump() - deleted_record = LiteLLM_DeletedTeamTable.model_validate( - { - **team_payload, - "deleted_at": deleted_at, - "deleted_by": user_api_key_dict.user_id, - "deleted_by_api_key": user_api_key_dict.api_key, - "litellm_changed_by": litellm_changed_by, - } - ) - record = deleted_record.model_dump() - - for json_field in [ - "members_with_roles", - "metadata", - "model_spend", - "model_max_budget", - "router_settings", - ]: - if json_field in record and record[json_field] is not None: - record[json_field] = json.dumps(record[json_field]) - - for rel_key in ( - "litellm_model_table", - "object_permission", - "id", - "budget_limits", # not in LiteLLM_DeletedTeamTable schema - "default_team_member_models", # not in LiteLLM_DeletedTeamTable schema - ): - record.pop(rel_key, None) - - records.append(record) - - return records + return tuple( + _transform_single_team_to_deleted_record(team, user_api_key_dict, deleted_at, litellm_changed_by) + for team in teams + ) async def _save_deleted_team_records( - records: List[Dict[str, Any]], + records: tuple[Mapping[str, object], ...], prisma_client: PrismaClient, ) -> None: """Save deleted team records to the database.""" @@ -3434,10 +3437,10 @@ async def _save_deleted_team_records( async def _persist_deleted_team_records( - teams: List[LiteLLM_TeamTable], + teams: list[LiteLLM_TeamTable], prisma_client: PrismaClient, user_api_key_dict: UserAPIKeyAuth, - litellm_changed_by: Optional[str] = None, + litellm_changed_by: str | None = None, ) -> None: """Persist deleted team records by transforming and saving them.""" records = _transform_teams_to_deleted_records( @@ -3506,7 +3509,7 @@ async def _add_team_member_budget_table( team_info_response_object: TeamInfoResponseObjectTeamTable, ) -> TeamInfoResponseObjectTeamTable: try: - team_budget = await BudgetRepository(prisma_client).table.find_unique( + team_budget: LiteLLM_BudgetTableFull | None = await BudgetRepository(prisma_client).table.find_unique( where={"budget_id": team_member_budget_id} ) team_info_response_object.team_member_budget_table = team_budget @@ -3518,7 +3521,7 @@ async def _add_team_member_budget_table( return team_info_response_object -async def _resolve_team_access_group_resources(_team_info: Any) -> None: +async def _resolve_team_access_group_resources(_team_info: TeamInfoResponseObjectTeamTable) -> None: """Populate access_group_models / mcp_server_ids / agent_ids on the team info response by resolving inherited resources from its access groups.""" if not _team_info.access_group_ids: @@ -3572,7 +3575,7 @@ async def team_info( ) try: - team_info: Optional[BaseModel] = await TeamRepository(prisma_client).table.find_unique( + team_info: BaseModel | None = await TeamRepository(prisma_client).table.find_unique( where={"team_id": team_id}, include={"litellm_model_table": True, "object_permission": True}, ) @@ -3723,7 +3726,7 @@ async def team_member_me( ) caller_user_email = user_api_key_dict.user_email - member_role: Optional[str] = None + member_role: str | None = None for m in team_table.members_with_roles: # Match by user_id when present, else fall back to email — members # added by email may have user_id=None on the stored entry. @@ -3819,7 +3822,7 @@ async def block_team( if prisma_client is None: raise Exception("No DB Connected.") - existing_team = await TeamRepository(prisma_client).table.find_unique(where={"team_id": data.team_id}) + existing_team = await TeamRepository(prisma_client).find_by_id(data.team_id) if existing_team is None: raise HTTPException( status_code=404, @@ -3828,11 +3831,11 @@ async def block_team( # Verify caller has access to manage this team await _verify_team_access( - team_obj=LiteLLM_TeamTable.model_validate(existing_team.model_dump()), + team_obj=existing_team, user_api_key_dict=user_api_key_dict, ) - record = await TeamRepository(prisma_client).table.update( + record: LiteLLM_TeamTable = await TeamRepository(prisma_client).table.update( where={"team_id": data.team_id}, data={"blocked": True}, # type: ignore ) @@ -3868,7 +3871,7 @@ async def unblock_team( if prisma_client is None: raise Exception("No DB Connected.") - existing_team = await TeamRepository(prisma_client).table.find_unique(where={"team_id": data.team_id}) + existing_team = await TeamRepository(prisma_client).find_by_id(data.team_id) if existing_team is None: raise HTTPException( status_code=404, @@ -3877,11 +3880,11 @@ async def unblock_team( # Verify caller has access to manage this team await _verify_team_access( - team_obj=LiteLLM_TeamTable.model_validate(existing_team.model_dump()), + team_obj=existing_team, user_api_key_dict=user_api_key_dict, ) - record = await TeamRepository(prisma_client).table.update( + record: LiteLLM_TeamTable = await TeamRepository(prisma_client).table.update( where={"team_id": data.team_id}, data={"blocked": False}, # type: ignore ) @@ -3893,7 +3896,7 @@ async def unblock_team( async def list_available_teams( http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - response_model=List[LiteLLM_TeamTable], + response_model=list[LiteLLM_TeamTable], ): from litellm.proxy.proxy_server import prisma_client @@ -3904,7 +3907,7 @@ async def list_available_teams( ) available_teams = cast( - Optional[List[str]], + list[str] | None, ( litellm.default_internal_user_params.get("available_teams") if litellm.default_internal_user_params is not None @@ -3915,29 +3918,26 @@ async def list_available_teams( return [] # filter out teams that the user is already a member of - user_info = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_api_key_dict.user_id}) + user_info: LiteLLM_UserTable | None = await UserRepository(prisma_client).table.find_unique( + where={"user_id": user_api_key_dict.user_id} + ) if user_info is None: raise HTTPException( status_code=404, detail={"error": "User not found"}, ) - user_info_correct_type = LiteLLM_UserTable.model_validate(user_info.model_dump()) - available_teams = [team for team in available_teams if team not in user_info_correct_type.teams] + available_teams = [team for team in available_teams if team not in user_info.teams] - available_teams_db = await TeamRepository(prisma_client).table.find_many(where={"team_id": {"in": available_teams}}) - - available_teams_correct_type = [LiteLLM_TeamTable.model_validate(team.model_dump()) for team in available_teams_db] - - return available_teams_correct_type + return await TeamRepository(prisma_client).find_many(where={"team_id": {"in": available_teams}}) async def _get_org_admin_org_ids( user_id: str, - prisma_client: Any, - user_api_key_cache: Any, - proxy_logging_obj: Any, -) -> Optional[List[str]]: + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> list[str] | None: """ Return the list of organization IDs where the user is an org admin. Returns None if the user is not an org admin of any organization or if @@ -3968,24 +3968,24 @@ async def _get_org_admin_org_ids( async def _build_team_list_where_conditions( prisma_client: PrismaClient, - team_id: Optional[str], - team_alias: Optional[str], - organization_id: Optional[str], - user_id: Optional[str], + team_id: str | None, + team_alias: str | None, + organization_id: str | None, + user_id: str | None, use_deleted_table: bool, - search: Optional[str] = None, + search: str | None = None, search_team_id_match: TeamIdSearchMatch = "exact", - org_admin_org_ids: Optional[List[str]] = None, - user_api_key_cache: Optional[Any] = None, - proxy_logging_obj: Optional[Any] = None, -) -> Optional[Dict[str, Any]]: + org_admin_org_ids: list[str] | None = None, + user_api_key_cache: UserApiKeyCache | None = None, + proxy_logging_obj: ProxyLogging | None = None, +) -> dict[str, object] | None: """ Build where conditions for team list query. Returns None when the query is guaranteed to yield no results (e.g. user has no team memberships), allowing the caller to skip the DB round-trip. """ - where_conditions: Dict[str, Any] = {} + where_conditions: dict[str, object] = {} if team_id: where_conditions["team_id"] = team_id @@ -4052,8 +4052,8 @@ async def _build_team_list_where_conditions( async def _batch_resolve_access_group_resources( - all_access_group_ids: List[str], -) -> Dict[str, Dict[str, List[str]]]: + all_access_group_ids: list[str], +) -> dict[str, dict[str, list[str]]]: """ Batch-fetch access groups in a single DB query and return a per-group resource map. @@ -4067,11 +4067,11 @@ async def _batch_resolve_access_group_resources( return {} unique_ids = list(set(all_access_group_ids)) - rows = await AccessGroupRepository(_prisma_client).table.find_many( + rows: list[LiteLLM_AccessGroupTable] = await AccessGroupRepository(_prisma_client).table.find_many( where={"access_group_id": {"in": unique_ids}}, ) - result: Dict[str, Dict[str, List[str]]] = {} + result: dict[str, dict[str, list[str]]] = {} for row in rows: result[row.access_group_id] = { "models": list(row.access_model_names or []), @@ -4084,10 +4084,10 @@ async def _batch_resolve_access_group_resources( def _convert_teams_to_response_models( teams: list, use_deleted_table: bool, - keys_count_by_team: Optional[Dict[str, int]] = None, -) -> List[Union[TeamListItem, LiteLLM_TeamTable, LiteLLM_DeletedTeamTable]]: + keys_count_by_team: dict[str, int] | None = None, +) -> list[TeamListItem | LiteLLM_TeamTable | LiteLLM_DeletedTeamTable]: """Convert raw Prisma team rows to response models.""" - team_list: List[Union[TeamListItem, LiteLLM_TeamTable, LiteLLM_DeletedTeamTable]] = [] + team_list: list[TeamListItem | LiteLLM_TeamTable | LiteLLM_DeletedTeamTable] = [] counts = keys_count_by_team or {} for team in teams: try: @@ -4114,10 +4114,15 @@ def _convert_teams_to_response_models( return team_list +class _TeamIdGroupCount(TypedDict): + team_id: str + _count: dict[str, int] + + async def _get_keys_count_by_team( - prisma_client: Any, + prisma_client: PrismaClient, teams: list, -) -> Dict[str, int]: +) -> dict[str, int]: """Aggregate virtual-key counts per team for the given page of teams. Runs a single GROUP BY against LiteLLM_VerificationToken. The IN clause is @@ -4128,22 +4133,22 @@ async def _get_keys_count_by_team( if not page_team_ids: return {} - grouped = await VerificationTokenRepository(prisma_client).table.group_by( + grouped: list[_TeamIdGroupCount] = await VerificationTokenRepository(prisma_client).table.group_by( by=["team_id"], where={"team_id": {"in": page_team_ids}}, count={"team_id": True}, ) - return {row["team_id"]: row.get("_count", {}).get("team_id", 0) for row in grouped if row.get("team_id")} + return {row["team_id"]: row["_count"].get("team_id", 0) for row in grouped if row["team_id"]} async def _enforce_list_team_v2_access( user_api_key_dict: UserAPIKeyAuth, - user_id: Optional[str], - organization_id: Optional[str], - prisma_client: Any, - user_api_key_cache: Any, - proxy_logging_obj: Any, -) -> Tuple[Optional[str], Optional[List[str]]]: + user_id: str | None, + organization_id: str | None, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> tuple[str | None, list[str] | None]: """Enforce access control for list_team_v2. - Proxy admins and admin viewers can query any teams. @@ -4153,7 +4158,7 @@ async def _enforce_list_team_v2_access( Returns the (possibly overridden) user_id and org_admin_org_ids. """ is_proxy_admin = _user_has_admin_view(user_api_key_dict) - org_admin_org_ids: Optional[List[str]] = None + org_admin_org_ids: list[str] | None = None if is_proxy_admin: return user_id, org_admin_org_ids @@ -4216,21 +4221,17 @@ async def _enforce_list_team_v2_access( @management_endpoint_wrapper async def list_team_v2( http_request: Request, - user_id: Optional[str] = fastapi.Query( - default=None, description="Only return teams which this 'user_id' belongs to" - ), - organization_id: Optional[str] = fastapi.Query( + user_id: str | None = fastapi.Query(default=None, description="Only return teams which this 'user_id' belongs to"), + organization_id: str | None = fastapi.Query( default=None, description="Only return teams which this 'organization_id' belongs to", ), - team_id: Optional[str] = fastapi.Query( - default=None, description="Only return teams which this 'team_id' belongs to" - ), - team_alias: Optional[str] = fastapi.Query( + team_id: str | None = fastapi.Query(default=None, description="Only return teams which this 'team_id' belongs to"), + team_alias: str | None = fastapi.Query( default=None, description="Only return teams which this 'team_alias' belongs to. Supports partial matching.", ), - search: Optional[str] = fastapi.Query( + search: str | None = fastapi.Query( default=None, description="Combined search: matches teams whose 'team_id' matches the value OR whose 'team_alias' contains it (case-insensitive).", ), @@ -4242,12 +4243,12 @@ async def list_team_v2( ] = "exact", page: int = fastapi.Query(default=1, description="Page number for pagination", ge=1), page_size: int = fastapi.Query(default=10, description="Number of teams per page", ge=1, le=100), - sort_by: Optional[str] = fastapi.Query( + sort_by: str | None = fastapi.Query( default=None, description="Column to sort by (e.g. 'team_id', 'team_alias', 'created_at')", ), sort_order: str = fastapi.Query(default="asc", description="Sort order ('asc' or 'desc')"), - status: Optional[str] = fastapi.Query(default=None, description="Filter by status (e.g. 'deleted')"), + status: str | None = fastapi.Query(default=None, description="Filter by status (e.g. 'deleted')"), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -4341,30 +4342,32 @@ async def list_team_v2( # Get teams with pagination if use_deleted_table: - teams = await DeletedTeamRepository(prisma_client).table.find_many( + teams: list[LiteLLM_DeletedTeamTable] | list[LiteLLM_TeamTable] = await DeletedTeamRepository( + prisma_client + ).table.find_many( where=where_conditions, skip=skip, take=page_size, order=order_by if order_by else {"created_at": "desc"}, # Default sort ) # Get total count for pagination - total_count = await DeletedTeamRepository(prisma_client).table.count(where=where_conditions) + total_count: int = await DeletedTeamRepository(prisma_client).table.count(where=where_conditions) else: - teams = await TeamRepository(prisma_client).table.find_many( + teams = await TeamRepository(prisma_client).find_many( where=where_conditions, skip=skip, take=page_size, order=order_by if order_by else {"created_at": "desc"}, # Default sort ) # Get total count for pagination - total_count = await TeamRepository(prisma_client).table.count(where=where_conditions) + total_count = await TeamRepository(prisma_client).count(where=where_conditions) # Calculate total pages total_pages = -(-total_count // page_size) # Ceiling division # Aggregate virtual-key counts per team for the current page. The deleted # table does not carry keys_count, so it is skipped. - keys_count_by_team: Dict[str, int] = {} + keys_count_by_team: dict[str, int] = {} if not use_deleted_table: keys_count_by_team = await _get_keys_count_by_team(prisma_client, teams) @@ -4397,13 +4400,24 @@ async def list_team_v2( } +class _RawTeamRecord(Protocol): + """Shape of an unconverted `litellm_teamtable` row as returned by a raw + `.table.find_many()` call: `members_with_roles` is still the raw JSON-decoded + list of dicts, not the parsed `List[Member]` that `LiteLLM_TeamTable` exposes.""" + + team_id: str + members_with_roles: list[Mapping[str, object]] | None + + def model_dump(self) -> dict[str, object]: ... + + async def _authorize_and_filter_teams( user_api_key_dict: UserAPIKeyAuth, - user_id: Optional[str], - prisma_client: Any, - user_api_key_cache: Any, - proxy_logging_obj: Any, -) -> list: + user_id: str | None, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> list[_RawTeamRecord]: """ Authorize the /team/list request and return filtered teams. @@ -4413,7 +4427,7 @@ async def _authorize_and_filter_teams( - Others: 401. """ is_proxy_admin = _user_has_admin_view(user_api_key_dict) - allowed_org_ids: Optional[List[str]] = None + allowed_org_ids: list[str] | None = None if not is_proxy_admin: is_own_query = ( @@ -4450,7 +4464,7 @@ async def _authorize_and_filter_teams( if allowed_org_ids is not None: # Org admin: query DB for teams in their orgs - org_teams = await TeamRepository(prisma_client).table.find_many( + org_teams: list[_RawTeamRecord] = await TeamRepository(prisma_client).table.find_many( where={"organization_id": {"in": allowed_org_ids}}, include={"litellm_model_table": True}, ) @@ -4464,7 +4478,9 @@ async def _authorize_and_filter_teams( ] elif user_id: # Regular user: fetch all and filter by membership (Prisma can't filter JSON arrays) - response = await TeamRepository(prisma_client).table.find_many(include={"litellm_model_table": True}) + response: list[_RawTeamRecord] = await TeamRepository(prisma_client).table.find_many( + include={"litellm_model_table": True} + ) return [ team for team in response @@ -4472,17 +4488,18 @@ async def _authorize_and_filter_teams( ] else: # Proxy admin: all teams - return list(await TeamRepository(prisma_client).table.find_many(include={"litellm_model_table": True})) + all_teams: list[_RawTeamRecord] = await TeamRepository(prisma_client).table.find_many( + include={"litellm_model_table": True} + ) + return all_teams @router.get("/team/list", tags=["team management"], dependencies=[Depends(user_api_key_auth)]) @management_endpoint_wrapper async def list_team( http_request: Request, - user_id: Optional[str] = fastapi.Query( - default=None, description="Only return teams which this 'user_id' belongs to" - ), - organization_id: Optional[str] = None, + user_id: str | None = fastapi.Query(default=None, description="Only return teams which this 'user_id' belongs to"), + organization_id: str | None = None, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -4518,22 +4535,24 @@ async def list_team( _team_ids = [team.team_id for team in filtered_response] returned_tm = await get_all_team_memberships(prisma_client, _team_ids, user_id=user_id) - returned_responses: List[TeamListResponseObject] = [] + returned_responses: list[TeamListResponseObject] = [] for team in filtered_response: - _team_memberships: List[LiteLLM_TeamMembership] = [] + _team_memberships: list[LiteLLM_TeamMembership] = [] for tm in returned_tm: if tm.team_id == team.team_id: _team_memberships.append(tm) # add all keys that belong to the team - keys = await VerificationTokenRepository(prisma_client).table.find_many(where={"team_id": team.team_id}) + keys = await VerificationTokenRepository(prisma_client).find_many(where={"team_id": team.team_id}) try: returned_responses.append( - TeamListResponseObject( - **team.model_dump(), - team_memberships=_team_memberships, - keys=keys, + TeamListResponseObject.model_validate( + { + **team.model_dump(), + "team_memberships": _team_memberships, + "keys": keys, + } ) ) except Exception as e: @@ -4558,7 +4577,7 @@ async def get_paginated_teams( prisma_client: PrismaClient, page_size: int = 10, page: int = 1, -) -> Tuple[List[LiteLLM_TeamTable], int]: +) -> tuple[list[LiteLLM_TeamTable], int]: """ Get paginated list of teams from team table @@ -4574,10 +4593,10 @@ async def get_paginated_teams( # Calculate skip for pagination skip = (page - 1) * page_size # Get total count - total_count = await TeamRepository(prisma_client).table.count() + total_count = await TeamRepository(prisma_client).count() # Get paginated teams - teams = await TeamRepository(prisma_client).table.find_many( + teams = await TeamRepository(prisma_client).find_many( skip=skip, take=page_size, order={"team_alias": "asc"}, # Sort by team_alias @@ -4594,12 +4613,12 @@ async def get_paginated_teams( dependencies=[Depends(user_api_key_auth)], include_in_schema=False, responses={ - 200: {"model": List[LiteLLM_TeamTable]}, + 200: {"model": list[LiteLLM_TeamTable]}, }, ) async def ui_view_teams( - team_id: Optional[str] = fastapi.Query(default=None, description="Team ID in the request parameters"), - team_alias: Optional[str] = fastapi.Query(default=None, description="Team alias in the request parameters"), + team_id: str | None = fastapi.Query(default=None, description="Team ID in the request parameters"), + team_alias: str | None = fastapi.Query(default=None, description="Team alias in the request parameters"), page: int = fastapi.Query(default=1, description="Page number for pagination", ge=1), page_size: int = fastapi.Query(default=50, description="Number of items per page", ge=1, le=100), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), @@ -4627,7 +4646,7 @@ async def ui_view_teams( skip = (page - 1) * page_size # Build where conditions based on provided parameters - where_conditions = {} + where_conditions: dict[str, object] = {} if team_id: where_conditions["team_id"] = { @@ -4642,7 +4661,7 @@ async def ui_view_teams( } # Query users with pagination and filters - teams = await TeamRepository(prisma_client).table.find_many( + teams = await TeamRepository(prisma_client).find_many( where=where_conditions, skip=skip, take=page_size, @@ -4658,7 +4677,7 @@ async def ui_view_teams( raise HTTPException(status_code=500, detail=f"Error searching teams: {str(e)}") -def add_new_models_to_team(team_obj: LiteLLM_TeamTable, new_models: List[str]) -> List[str]: +def add_new_models_to_team(team_obj: LiteLLM_TeamTable, new_models: list[str]) -> list[str]: """ Add new models to a team's allowed model list. """ @@ -4710,16 +4729,14 @@ async def team_model_add( raise HTTPException(status_code=500, detail={"error": "No db connected"}) # Get existing team - team_row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": data.team_id}) + team_obj = await TeamRepository(prisma_client).find_by_id(data.team_id) - if team_row is None: + if team_obj is None: raise HTTPException( status_code=404, detail={"error": f"Team not found, passed team_id={data.team_id}"}, ) - team_obj = LiteLLM_TeamTable.model_validate(team_row.model_dump()) - # Authorization check - only proxy admin, team admin, or org admin can add models if ( user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value @@ -4756,7 +4773,7 @@ async def team_model_add( # the writer and lets Prisma bump updated_at. # `include` mirrors the relations the auth path consumes off the cached # team object so that `_refresh_cached_team` doesn't null them out. - updated_team = await TeamRepository(prisma_client).table.update( + updated_team: LiteLLM_TeamTable = await TeamRepository(prisma_client).table.update( where={"team_id": data.team_id}, data={"updated_at": datetime.now(timezone.utc)}, include={"object_permission": True}, # type: ignore @@ -4810,16 +4827,14 @@ async def team_model_delete( raise HTTPException(status_code=500, detail={"error": "No db connected"}) # Get existing team - team_row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": data.team_id}) + team_obj = await TeamRepository(prisma_client).find_by_id(data.team_id) - if team_row is None: + if team_obj is None: raise HTTPException( status_code=404, detail={"error": f"Team not found, passed team_id={data.team_id}"}, ) - team_obj = LiteLLM_TeamTable.model_validate(team_row.model_dump()) - # Authorization check - only proxy admin, team admin, or org admin can remove models if ( user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value @@ -4838,7 +4853,7 @@ async def team_model_delete( updated_models = [m for m in current_models if m not in data.models] # Update team. See team_model_add for the rationale on `include`. - updated_team = await TeamRepository(prisma_client).table.update( + updated_team: LiteLLM_TeamTable = await TeamRepository(prisma_client).table.update( where={"team_id": data.team_id}, data={"models": updated_models}, include={"object_permission": True}, # type: ignore @@ -4973,7 +4988,7 @@ async def update_team_member_permissions( }, ) # Update the team member permissions - updated_team = await TeamRepository(prisma_client).table.update( + updated_team: LiteLLM_TeamTable = await TeamRepository(prisma_client).table.update( where={"team_id": data.team_id}, data={"team_member_permissions": data.team_member_permissions}, ) @@ -5065,9 +5080,11 @@ async def _compute_and_batch_updates(prisma_client, teams, permissions_to_add: s return len(updates) -async def _append_permissions_to_specific_teams(prisma_client, team_ids: List[str], permissions_to_add: set) -> int: +async def _append_permissions_to_specific_teams( + prisma_client: PrismaClient, team_ids: list[str], permissions_to_add: set +) -> int: """Fetch specific teams by ID and append permissions.""" - teams = await TeamRepository(prisma_client).table.find_many( + teams = await TeamRepository(prisma_client).find_many( where={"team_id": {"in": team_ids}}, ) @@ -5082,14 +5099,14 @@ async def _append_permissions_to_specific_teams(prisma_client, team_ids: List[st return await _compute_and_batch_updates(prisma_client, teams, permissions_to_add) -async def _append_permissions_to_all_teams(prisma_client, permissions_to_add: set) -> int: +async def _append_permissions_to_all_teams(prisma_client: PrismaClient, permissions_to_add: set) -> int: """Paginated read + batched write across all teams.""" teams_updated = 0 - cursor = None + cursor: str | None = None BATCH_SIZE = 500 while True: - find_args: dict = { + find_args: dict[str, object] = { "take": BATCH_SIZE, "order": {"team_id": "asc"}, } @@ -5097,7 +5114,7 @@ async def _append_permissions_to_all_teams(prisma_client, permissions_to_add: se find_args["cursor"] = {"team_id": cursor} find_args["skip"] = 1 - teams = await TeamRepository(prisma_client).table.find_many(**find_args) + teams: list[LiteLLM_TeamTable] = await TeamRepository(prisma_client).table.find_many(**find_args) if not teams: break @@ -5118,14 +5135,14 @@ async def _append_permissions_to_all_teams(prisma_client, permissions_to_add: se tags=["team management"], ) async def get_team_daily_activity( - team_ids: Optional[str] = None, - start_date: Optional[str] = None, - end_date: Optional[str] = None, - model: Optional[str] = None, - api_key: Optional[str] = None, + team_ids: str | None = None, + start_date: str | None = None, + end_date: str | None = None, + model: str | None = None, + api_key: str | None = None, page: int = 1, page_size: int = 10, - exclude_team_ids: Optional[str] = None, + exclude_team_ids: str | None = None, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -5157,7 +5174,7 @@ async def get_team_daily_activity( # Convert comma-separated tags string to list if provided team_ids_list = team_ids.split(",") if team_ids else None - exclude_team_ids_list: Optional[List[str]] = None + exclude_team_ids_list: list[str] | None = None if exclude_team_ids: exclude_team_ids_list = exclude_team_ids.split(",") if exclude_team_ids else None @@ -5194,11 +5211,11 @@ async def get_team_daily_activity( ) ## Fetch team aliases and check team admin status - where_condition = {} + where_condition: dict[str, object] = {} if team_ids_list: where_condition["team_id"] = {"in": list(team_ids_list)} - team_aliases = await TeamRepository(prisma_client).table.find_many(where=where_condition) - team_alias_metadata = {t.team_id: {"team_alias": t.team_alias} for t in team_aliases} + team_aliases = await TeamRepository(prisma_client).find_many(where=where_condition) + team_alias_metadata: dict[str, dict[str, object]] = {t.team_id: {"team_alias": t.team_alias} for t in team_aliases} # Check if user is team admin or has /team/daily/activity permission # If not, filter by user's API keys. @@ -5210,11 +5227,10 @@ async def get_team_daily_activity( # only has admin/permission for a strict subset, fall back to # filtering the entire response by their own API keys (they can re- # request the admin-only teams separately to get the wider view). - user_api_keys: Optional[List[str]] = None + user_api_keys: list[str] | None = None if not _user_has_admin_view(user_api_key_dict) and team_ids_list and team_aliases: has_full_team_view = True - for team_alias in team_aliases: - team_obj = LiteLLM_TeamTable.model_validate(team_alias.model_dump()) + for team_obj in team_aliases: is_admin = _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team_obj) has_perm = _team_member_has_permission( user_api_key_dict=user_api_key_dict, @@ -5228,7 +5244,7 @@ async def get_team_daily_activity( # If user does not have full team view, filter by their API keys if not has_full_team_view: # Get all API keys for this user - user_keys = await VerificationTokenRepository(prisma_client).table.find_many( + user_keys = await VerificationTokenRepository(prisma_client).find_many( where={"user_id": user_api_key_dict.user_id} ) user_api_keys = [key.token for key in user_keys if key.token] @@ -5237,7 +5253,7 @@ async def get_team_daily_activity( user_api_keys = [""] # Use empty string to ensure no matches # If api_key parameter is provided, use it; otherwise use user_api_keys if set - final_api_key_filter: Optional[Union[str, List[str]]] = api_key + final_api_key_filter: str | list[str] | None = api_key if final_api_key_filter is None and user_api_keys is not None: final_api_key_filter = user_api_keys diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9023f5fb23f..544a5d3a103 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -15,19 +15,14 @@ import threading import time import traceback import warnings -from collections.abc import Mapping +from collections.abc import AsyncGenerator, Callable, Mapping, Sequence from datetime import datetime, timedelta, timezone from typing import ( TYPE_CHECKING, Any, - AsyncGenerator, - Callable, - Dict, - List, Literal, Optional, Set, - Tuple, TypedDict, Union, cast, @@ -131,15 +126,20 @@ from litellm.utils import ( if TYPE_CHECKING: from aiohttp import ClientSession from opentelemetry.trace import Span as _Span + from prisma.models import LiteLLM_AccessGroupTable as PrismaAccessGroupTable + from prisma.models import LiteLLM_InvitationLink as PrismaInvitationLink + from prisma.models import LiteLLM_PromptTable as PrismaPromptTable + from prisma.models import LiteLLM_ProxyModelTable as PrismaProxyModelTable + from prisma.models import LiteLLM_TeamTable as PrismaTeamTable + from prisma.models import LiteLLM_UserTable as PrismaUserTable from litellm.integrations.opentelemetry import OpenTelemetry Span = Union[_Span, Any] else: Span = Any - OpenTelemetry = Any -REALTIME_REQUEST_SCOPE_TEMPLATE: Dict[str, Any] = { +REALTIME_REQUEST_SCOPE_TEMPLATE: dict[str, Any] = { "type": "http", "method": "POST", "path": "/v1/realtime", @@ -308,11 +308,6 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( encrypt_value_helper, ) from litellm.proxy.common_utils.html_forms.ui_login import build_ui_login_form -from litellm.proxy.config_resolvers import resolve_fields -from litellm.proxy.config_resolvers.alerting import ( - EMAIL_DESCRIPTORS, - SLACK_DESCRIPTORS, -) from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, _safe_get_request_headers, @@ -338,6 +333,11 @@ from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, get_management_object_ttl, ) +from litellm.proxy.config_resolvers import resolve_fields +from litellm.proxy.config_resolvers.alerting import ( + EMAIL_DESCRIPTORS, + SLACK_DESCRIPTORS, +) from litellm.proxy.container_endpoints.endpoints import router as container_router from litellm.proxy.credential_endpoints.endpoints import router as credential_router from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import SpendLogCleanup @@ -387,6 +387,8 @@ from litellm.proxy.management_endpoints.common_utils import ( ) from litellm.proxy.management_endpoints.coordination_redis_endpoints import ( get_persisted_coordination_redis_settings, +) +from litellm.proxy.management_endpoints.coordination_redis_endpoints import ( router as coordination_redis_settings_router, ) from litellm.proxy.management_endpoints.cost_tracking_settings import ( @@ -395,16 +397,6 @@ from litellm.proxy.management_endpoints.cost_tracking_settings import ( from litellm.proxy.management_endpoints.customer_endpoints import ( router as customer_router, ) -from litellm.proxy.management_endpoints.management_v1 import ( - router as management_v1_router, -) -from litellm.proxy.management_endpoints.management_v1.common import ( - MANAGEMENT_V1_PREFIX, - PROBLEM_TYPE_BASE, - ManagementProblem, - problem_response, -) -from litellm.types.proxy.management_endpoints.management_v1 import ProblemDetail from litellm.proxy.management_endpoints.fallback_management_endpoints import ( router as fallback_management_router, ) @@ -422,6 +414,15 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( from litellm.proxy.management_endpoints.key_management_endpoints import ( router as key_management_router, ) +from litellm.proxy.management_endpoints.management_v1 import ( + router as management_v1_router, +) +from litellm.proxy.management_endpoints.management_v1.common import ( + MANAGEMENT_V1_PREFIX, + PROBLEM_TYPE_BASE, + ManagementProblem, + problem_response, +) from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( router as model_access_group_management_router, ) @@ -475,6 +476,7 @@ from litellm.proxy.plugin_routes import ( from litellm.proxy.plugin_routes import ( router as plugin_router, ) +from litellm.types.proxy.management_endpoints.management_v1 import ProblemDetail try: from litellm.proxy.enterprise_billing.billing_metrics import ( @@ -484,8 +486,8 @@ try: shutdown_billing_metrics_recorder as _shutdown_billing_metrics_recorder, ) - build_billing_metrics_recorder: Optional[Callable[..., Optional[BillingRecorder]]] = _build_billing_metrics_recorder - shutdown_billing_metrics_recorder: Optional[Callable[[], None]] = _shutdown_billing_metrics_recorder + build_billing_metrics_recorder: Callable[..., BillingRecorder | None] | None = _build_billing_metrics_recorder + shutdown_billing_metrics_recorder: Callable[[], None] | None = _shutdown_billing_metrics_recorder except ImportError: build_billing_metrics_recorder = None shutdown_billing_metrics_recorder = None @@ -649,7 +651,7 @@ from fastapi.responses import ( RedirectResponse, StreamingResponse, ) -from fastapi.routing import APIRouter +from fastapi.routing import APIRoute, APIRouter from fastapi.security import OAuth2PasswordBearer from fastapi.security.api_key import APIKeyHeader from fastapi.staticfiles import StaticFiles @@ -675,7 +677,7 @@ try: from litellm_enterprise.proxy.proxy_server import EnterpriseProxyConfig enterprise_router = _enterprise_router - enterprise_proxy_config: Optional[EnterpriseProxyConfig] = EnterpriseProxyConfig() + enterprise_proxy_config: EnterpriseProxyConfig | None = EnterpriseProxyConfig() except ImportError: enterprise_proxy_config = None ################### @@ -684,7 +686,7 @@ server_root_path = get_server_root_path() _license_check = LicenseCheck() premium_user: bool = _license_check.is_premium() premium_user_data: Optional["EnterpriseLicenseData"] = _license_check.airgapped_license_data -global_max_parallel_request_retries_env: Optional[str] = os.getenv("LITELLM_GLOBAL_MAX_PARALLEL_REQUEST_RETRIES") +global_max_parallel_request_retries_env: str | None = os.getenv("LITELLM_GLOBAL_MAX_PARALLEL_REQUEST_RETRIES") proxy_state = ProxyState() SENSITIVE_DATA_MASKER = SensitiveDataMasker() @@ -729,7 +731,7 @@ if global_max_parallel_request_retries_env is None: else: global_max_parallel_request_retries = int(global_max_parallel_request_retries_env) -global_max_parallel_request_retry_timeout_env: Optional[str] = os.getenv( +global_max_parallel_request_retry_timeout_env: str | None = os.getenv( "LITELLM_GLOBAL_MAX_PARALLEL_REQUEST_RETRY_TIMEOUT" ) if global_max_parallel_request_retry_timeout_env is None: @@ -840,21 +842,16 @@ async def _initialize_shared_aiohttp_session(): _build_aiohttp_keepalive_socket_factory, ) - connector_kwargs: Dict[str, Any] = { - "keepalive_timeout": AIOHTTP_KEEPALIVE_TIMEOUT, - "ttl_dns_cache": AIOHTTP_TTL_DNS_CACHE, - } - if AIOHTTP_NEEDS_CLEANUP_CLOSED: - connector_kwargs["enable_cleanup_closed"] = True - if AIOHTTP_CONNECTOR_LIMIT > 0: - connector_kwargs["limit"] = AIOHTTP_CONNECTOR_LIMIT - if AIOHTTP_CONNECTOR_LIMIT_PER_HOST > 0: - connector_kwargs["limit_per_host"] = AIOHTTP_CONNECTOR_LIMIT_PER_HOST socket_factory = _build_aiohttp_keepalive_socket_factory() - if socket_factory is not None: - connector_kwargs["socket_factory"] = socket_factory - connector = TCPConnector(**connector_kwargs) + connector = TCPConnector( + keepalive_timeout=AIOHTTP_KEEPALIVE_TIMEOUT, + ttl_dns_cache=AIOHTTP_TTL_DNS_CACHE, + enable_cleanup_closed=AIOHTTP_NEEDS_CLEANUP_CLOSED, + limit=AIOHTTP_CONNECTOR_LIMIT if AIOHTTP_CONNECTOR_LIMIT > 0 else 100, + limit_per_host=(AIOHTTP_CONNECTOR_LIMIT_PER_HOST if AIOHTTP_CONNECTOR_LIMIT_PER_HOST > 0 else 0), + socket_factory=socket_factory, + ) session = ClientSession(connector=connector) verbose_proxy_logger.info( @@ -923,8 +920,8 @@ async def proxy_startup_event(app: FastAPI): ## CHECK MASTER KEY IN ENVIRONMENT ## master_key = get_secret_str("LITELLM_MASTER_KEY") ### LOAD CONFIG ### - worker_config: Optional[Union[str, dict]] = get_secret("WORKER_CONFIG") # type: ignore - env_config_yaml: Optional[str] = get_secret_str("CONFIG_FILE_PATH") + worker_config: str | dict | None = get_secret("WORKER_CONFIG") # type: ignore + env_config_yaml: str | None = get_secret_str("CONFIG_FILE_PATH") verbose_proxy_logger.debug("worker_config: %s", _redact_worker_config_for_logging(worker_config)) # check if it's a valid file path if env_config_yaml is not None: @@ -961,7 +958,7 @@ async def proxy_startup_event(app: FastAPI): # check if DATABASE_URL in environment - load from there if prisma_client is None: - _db_url: Optional[str] = get_secret("DATABASE_URL", None) # type: ignore + _db_url: str | None = get_secret("DATABASE_URL", None) # type: ignore prisma_client = await ProxyStartupEvent._setup_prisma_client( database_url=_db_url, proxy_logging_obj=proxy_logging_obj, @@ -1156,7 +1153,7 @@ async def proxy_startup_event(app: FastAPI): await proxy_shutdown_event() # type: ignore[reportGeneralTypeIssues] -def _generate_stable_operation_id(route: Any) -> str: +def _generate_stable_operation_id(route: APIRoute) -> str: operation_id = re.sub(r"\W", "_", f"{route.name}{route.path_format}") route_methods = sorted(route.methods or []) if len(route_methods) == 1: @@ -1193,11 +1190,11 @@ def _strip_operation_id_method_suffix(operation_id: str) -> str: def ensure_unique_openapi_operation_ids( - openapi_schema: Dict[str, Any], - reserved_operation_ids: Optional[Set[str]] = None, -) -> Dict[str, Any]: + openapi_schema: dict[str, Any], + reserved_operation_ids: Set[str] | None = None, +) -> dict[str, Any]: operation_entries = [] - operation_id_counts: Dict[str, int] = {} + operation_id_counts: dict[str, int] = {} for path_item in openapi_schema.get("paths", {}).values(): if not isinstance(path_item, dict): continue @@ -1414,7 +1411,7 @@ async def openai_exception_handler(request: Request, exc: ProxyException): ) -def _close_dangling_otel_server_span(request: Request, status_code: int, exc: Optional[Exception] = None) -> None: +def _close_dangling_otel_server_span(request: Request, status_code: int, exc: Exception | None = None) -> None: parent_otel_span = getattr(request.state, "parent_otel_span", None) if parent_otel_span is None: return @@ -1501,8 +1498,8 @@ router = APIRouter() def _get_cors_config( - cors_origins_env: Optional[str] = None, - cors_credentials_env: Optional[str] = None, + cors_origins_env: str | None = None, + cors_credentials_env: str | None = None, ): """ Compute CORS allowed origins and credentials flag from environment variables. @@ -1972,7 +1969,7 @@ def mount_swagger_ui(): mount_swagger_ui() docs_url = _get_docs_url() -root_redirect_url: Optional[str] = os.getenv("ROOT_REDIRECT_URL") +root_redirect_url: str | None = os.getenv("ROOT_REDIRECT_URL") if docs_url != "/" and root_redirect_url is not None: @app.get("/", include_in_schema=False) @@ -1980,8 +1977,6 @@ if docs_url != "/" and root_redirect_url is not None: return RedirectResponse(url=root_redirect_url) # type: ignore[arg-type] -from typing import Dict - user_api_base = None user_model = None user_debug = False @@ -1991,20 +1986,20 @@ user_temperature = None user_telemetry = True user_config = None user_headers = None -user_config_file_path: Optional[str] = None +user_config_file_path: str | None = None local_logging = True # writes logs to a local api_log.json file for debugging experimental = False #### GLOBAL VARIABLES #### -llm_router: Optional[Router] = None -llm_model_list: Optional[list] = None +llm_router: Router | None = None +llm_model_list: list | None = None general_settings: dict = {} -config_passthrough_endpoints: Optional[List[Dict[str, Any]]] = None +config_passthrough_endpoints: list[dict[str, Any]] | None = None log_file = "api_log.json" worker_config = None -master_key: Optional[str] = None -config_agents: Optional[List[AgentConfig]] = None +master_key: str | None = None +config_agents: list[AgentConfig] | None = None otel_logging = False -prisma_client: Optional[PrismaClient] = None +prisma_client: PrismaClient | None = None shared_aiohttp_session: Optional["ClientSession"] = None # Global shared session for connection reuse user_api_key_cache: UserApiKeyCache = UserApiKeyCache( default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value @@ -2013,9 +2008,9 @@ spend_counter_cache = DualCache(default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_ cli_sso_session_cache = DualCache(default_in_memory_ttl=CLI_SSO_SESSION_TTL_SECONDS) model_max_budget_limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=user_api_key_cache) litellm.logging_callback_manager.add_litellm_callback(model_max_budget_limiter) -redis_usage_cache: Optional[RedisCache] = None # redis cache used for tracking spend, tpm/rpm limits -polling_via_cache_enabled: Union[Literal["all"], List[str], bool] = False -native_background_mode: List[str] = [] # Models that should use native provider background mode instead of polling +redis_usage_cache: RedisCache | None = None # redis cache used for tracking spend, tpm/rpm limits +polling_via_cache_enabled: Literal["all"] | list[str] | bool = False +native_background_mode: list[str] = [] # Models that should use native provider background mode instead of polling polling_cache_ttl: int = 3600 # Default 1 hour TTL for polling cache user_custom_auth = None user_custom_key_generate = None @@ -2032,13 +2027,13 @@ use_queue = False health_check_interval = None health_check_concurrency = None health_check_details = None -health_check_results: Dict[str, Union[int, List[Dict[str, Any]]]] = {} +health_check_results: dict[str, int | list[dict[str, Any]]] = {} background_health_check_loop_active = False background_health_check_cycle_seq = 0 -queue: List = [] +queue: list = [] litellm_proxy_budget_name = LITELLM_PROXY_BUDGET_NAME litellm_proxy_admin_name = LITELLM_PROXY_ADMIN_NAME -ui_access_mode: Union[Literal["admin", "all"], Dict] = "all" +ui_access_mode: Literal["admin", "all"] | dict = "all" proxy_budget_rescheduler_min_time = PROXY_BUDGET_RESCHEDULER_MIN_TIME proxy_budget_rescheduler_max_time = PROXY_BUDGET_RESCHEDULER_MAX_TIME proxy_batch_polling_interval = PROXY_BATCH_POLLING_INTERVAL @@ -2047,9 +2042,9 @@ proxy_config_reload_interval_seconds = PROXY_CONFIG_RELOAD_INTERVAL_SECONDS litellm_master_key_hash = None disable_spend_logs = False jwt_handler = JWTHandler() -prompt_injection_detection_obj: Optional[_OPTIONAL_PromptInjectionDetection] = None +prompt_injection_detection_obj: _OPTIONAL_PromptInjectionDetection | None = None store_model_in_db: bool = False -open_telemetry_logger: Optional[OpenTelemetry] = None +open_telemetry_logger: Optional["OpenTelemetry"] = None ### INITIALIZE GLOBAL LOGGING OBJECT ### proxy_logging_obj: ProxyLogging = ProxyLogging(user_api_key_cache=user_api_key_cache, premium_user=premium_user) ### REDIS QUEUE ### @@ -2066,7 +2061,7 @@ last_anthropic_beta_headers_reload = None ### DB WRITER ### -db_writer_client: Optional[AsyncHTTPHandler] = None +db_writer_client: AsyncHTTPHandler | None = None ### logger ### @@ -2084,7 +2079,7 @@ def _resolve_typed_dict_type(typ): return None -def _resolve_pydantic_type(typ) -> List: +def _resolve_pydantic_type(typ) -> list: """Resolve the actual TypedDict class from a potentially wrapped type.""" origin = get_origin(typ) typs = [] @@ -2370,14 +2365,14 @@ async def _read_spend_counter_estimate(counter_key: str, fallback_spend: float) async def increment_spend_counters( - token: Optional[str], - team_id: Optional[str], - user_id: Optional[str], - response_cost: Optional[float], - org_id: Optional[str] = None, - budget_reservation: Optional[dict] = None, - end_user_id: Optional[str] = None, - tags: Optional[List[str]] = None, + token: str | None, + team_id: str | None, + user_id: str | None, + response_cost: float | None, + org_id: str | None = None, + budget_reservation: dict | None = None, + end_user_id: str | None = None, + tags: list[str] | None = None, ): """ Atomically increment spend counters for budget enforcement. @@ -2531,8 +2526,8 @@ async def increment_spend_counters( async def _reconcile_budget_reservation_for_counter_update( - budget_reservation: Optional[dict], - response_cost: Optional[float], + budget_reservation: dict | None, + response_cost: float | None, ) -> Set[str]: if budget_reservation is None: return set() @@ -2566,8 +2561,8 @@ async def _reconcile_budget_reservation_for_counter_update( async def _increment_end_user_and_tag_spend_counters( - end_user_id: Optional[str], - tags: Optional[List[str]], + end_user_id: str | None, + tags: list[str] | None, response_cost: float, reserved_counter_keys: Set[str], ) -> None: @@ -2596,7 +2591,7 @@ async def _increment_end_user_and_tag_spend_counters( async def _increment_org_spend_counter( - org_id: Optional[str], + org_id: str | None, response_cost: float, reserved_counter_keys: Set[str], ) -> None: @@ -2613,7 +2608,7 @@ async def _increment_org_spend_counter( async def _init_and_increment_unreserved_spend_counter( counter_key: str, - source_cache_key: Union[str, List[str]], + source_cache_key: str | list[str], increment: float, reserved_counter_keys: Set[str], ) -> None: @@ -2629,7 +2624,7 @@ async def _init_and_increment_unreserved_spend_counter( async def _init_and_increment_spend_counter( counter_key: str, - source_cache_key: Union[str, List[str]], + source_cache_key: str | list[str], increment: float, ): """ @@ -2659,7 +2654,7 @@ async def _init_and_increment_window_spend_counter( counter_key: str, entity_type: str, entity_id: str, - window_start: Optional[datetime], + window_start: datetime | None, increment: float, ): if window_start is None: @@ -2682,7 +2677,7 @@ async def _init_and_increment_window_spend_counter( async def _ensure_spend_counter_initialized( counter_key: str, - source_cache_key: Union[str, List[str]], + source_cache_key: str | list[str], ): is_warm = await _is_spend_counter_cache_warm(counter_key=counter_key) if is_warm is False: @@ -2701,7 +2696,7 @@ async def _ensure_spend_counter_initialized( async def _get_source_cache_base_spend( - source_cache_key: Union[str, List[str]], + source_cache_key: str | list[str], ) -> float: source_cache_keys = [source_cache_key] if isinstance(source_cache_key, str) else source_cache_key for cache_key in source_cache_keys: @@ -2802,13 +2797,13 @@ async def _invalidate_spend_counter(counter_key: str): async def update_cache( - token: Optional[str], - user_id: Optional[str], - end_user_id: Optional[str], - team_id: Optional[str], - response_cost: Optional[float], - parent_otel_span: Optional[Span], # type: ignore - tags: Optional[List[str]] = None, + token: str | None, + user_id: str | None, + end_user_id: str | None, + team_id: str | None, + response_cost: float | None, + parent_otel_span: Span | None, + tags: list[str] | None = None, ): """ Use this to update the cache with new user spend. @@ -2816,7 +2811,7 @@ async def update_cache( Put any alerting logic in here. """ - values_to_update_in_cache: List[Tuple[Any, Any]] = [] + values_to_update_in_cache: list[tuple[Any, Any]] = [] ### UPDATE KEY SPEND ### async def _update_key_cache(token: str, response_cost: float): @@ -2976,7 +2971,7 @@ async def update_cache( if cached_team is None: # do nothing if team not in api key cache return - existing_spend_obj: Optional[LiteLLM_TeamTableCachedObj] = CacheCodec.deserialize( + existing_spend_obj: LiteLLM_TeamTableCachedObj | None = CacheCodec.deserialize( cached_team, LiteLLM_TeamTableCachedObj ) if existing_spend_obj is None: @@ -3107,7 +3102,7 @@ def run_ollama_serve(): """) -def _get_process_rss_mb() -> Optional[float]: +def _get_process_rss_mb() -> float | None: """ Get process RSS memory in MB. On Linux, ru_maxrss is in KB. On macOS, ru_maxrss is in bytes. @@ -3137,13 +3132,13 @@ def _is_unexpected_keyword_argument_type_error(exc: BaseException) -> bool: async def _run_direct_health_check_with_instrumentation( model_list: list, - details: Optional[bool], - max_concurrency: Optional[int], + details: bool | None, + max_concurrency: int | None, instrumentation_context: dict, ): """Call ``perform_health_check``, retrying with fewer kwargs on unexpected-kw TypeErrors.""" _hc_filter = health_check_filter_kwargs_from_general_settings(general_settings) - last_type_error: Optional[TypeError] = None + last_type_error: TypeError | None = None for extra_kwargs in ( { "instrumentation_context": instrumentation_context, @@ -3215,7 +3210,7 @@ def _get_endpoint_exception_status(endpoint: dict, exceptions: dict) -> int: def _write_health_state_to_router_cache( healthy_endpoints: list, unhealthy_endpoints: list, - exceptions_by_model_id: Optional[dict] = None, + exceptions_by_model_id: dict | None = None, ) -> None: """ Write deployment health states to the router's health state cache @@ -3501,7 +3496,7 @@ class StreamingCallbackError(Exception): # active), so the runtime gate cannot distinguish a YAML-sourced value # from a DB-sourced value. Scrubbing at the merge boundary closes that # gap without tracking source on every config dict entry. -_DB_OVERLAY_REMOTE_MODULE_STR_FIELDS: Dict[str, Tuple[str, ...]] = { +_DB_OVERLAY_REMOTE_MODULE_STR_FIELDS: dict[str, tuple[str, ...]] = { "litellm_settings": ("post_call_rules",), "general_settings": ( "custom_auth", @@ -3511,7 +3506,7 @@ _DB_OVERLAY_REMOTE_MODULE_STR_FIELDS: Dict[str, Tuple[str, ...]] = { "custom_ui_sso_sign_in_handler", ), } -_DB_OVERLAY_REMOTE_MODULE_LIST_FIELDS: Dict[str, Tuple[str, ...]] = { +_DB_OVERLAY_REMOTE_MODULE_LIST_FIELDS: dict[str, tuple[str, ...]] = { "litellm_settings": ( "callbacks", "success_callback", @@ -3525,7 +3520,7 @@ def _is_remote_module_url(value: Any) -> bool: return isinstance(value, str) and (value.startswith("s3://") or value.startswith("gcs://")) -def _scrub_guardrail_inner(inner: Dict[str, Any]) -> None: +def _scrub_guardrail_inner(inner: dict[str, Any]) -> None: """Strip remote-URL entries from a guardrail's ``callbacks`` list and ``guardrail`` (v2 module-path) field. Mutates in place.""" cbs = inner.get("callbacks") @@ -3645,7 +3640,7 @@ def _scrub_db_overlay_remote_module_loads(section: str, db_value: Any) -> Any: return sanitized -def _normalize_user_url_validation(value: object) -> Optional[bool]: +def _normalize_user_url_validation(value: object) -> bool | None: if value is None: return None if isinstance(value, str): @@ -3835,10 +3830,10 @@ class ProxyConfig: """ def __init__(self) -> None: - self.config: Dict[str, Any] = {} - self._last_semantic_filter_config: Optional[Dict[str, Any]] = None - self._last_hashicorp_vault_config: Optional[Dict[str, Any]] = None - self.worker_registry: List["WorkerRegistryEntry"] = [] + self.config: dict[str, Any] = {} + self._last_semantic_filter_config: dict[str, Any] | None = None + self._last_hashicorp_vault_config: dict[str, Any] | None = None + self.worker_registry: list[WorkerRegistryEntry] = [] def is_yaml(self, config_file_path: str) -> bool: if not os.path.isfile(config_file_path): @@ -3857,7 +3852,7 @@ class ProxyConfig: except Exception as e: raise Exception(f"Error loading yaml file {file_path}: {str(e)}") - async def _get_config_from_file(self, config_file_path: Optional[str] = None) -> dict: + async def _get_config_from_file(self, config_file_path: str | None = None) -> dict: """ Given a config file path, load the config from the file. Args: @@ -4040,7 +4035,7 @@ class ProxyConfig: config[key] = get_secret(value) return config - def _get_team_config(self, team_id: str, all_teams_config: List[Dict]) -> Dict: + def _get_team_config(self, team_id: str, all_teams_config: list[dict]) -> dict: team_config: dict = {} for team in all_teams_config: if "team_id" not in team: @@ -4163,7 +4158,7 @@ class ProxyConfig: llm_router.cache_responses = True verbose_proxy_logger.debug("Set router.cache_responses=True after initializing cache") - async def get_config(self, config_file_path: Optional[str] = None) -> dict: + async def get_config(self, config_file_path: str | None = None) -> dict: """ Load config file Supports reading from: @@ -4235,7 +4230,7 @@ class ProxyConfig: ) return {} - def load_credential_list(self, config: dict) -> List[CredentialItem]: + def load_credential_list(self, config: dict) -> list[CredentialItem]: """ Load the credential list from the database """ @@ -4245,7 +4240,7 @@ class ProxyConfig: credential_list = [CredentialItem(**cred) for cred in credential_list_dict] return credential_list - def parse_search_tools(self, config: dict) -> Optional[List[SearchToolTypedDict]]: + def parse_search_tools(self, config: dict) -> list[SearchToolTypedDict] | None: """ Parse and validate search tools from config. Loads environment variables and casts to SearchToolTypedDict. @@ -4266,7 +4261,7 @@ class ProxyConfig: if not search_tools_raw: return None - search_tools_parsed: List[SearchToolTypedDict] = [] + search_tools_parsed: list[SearchToolTypedDict] = [] print( # noqa: T201 "\033[32mLiteLLM: Proxy initialized with Search Tools:\033[0m" @@ -4336,7 +4331,7 @@ class ProxyConfig: # ``` ######################################################### if isinstance(value, str) and value.startswith("os.environ/"): - resolved_secret_string: Optional[str] = get_secret_str(secret_name=value) + resolved_secret_string: str | None = get_secret_str(secret_name=value) if resolved_secret_string is not None: os.environ[key] = resolved_secret_string else: @@ -4355,7 +4350,7 @@ class ProxyConfig: premium_user = _license_check.is_premium() return - async def load_config(self, router: Optional[litellm.Router], config_file_path: str): + async def load_config(self, router: litellm.Router | None, config_file_path: str): """ Load config values into proxy global state """ @@ -4972,7 +4967,7 @@ class ProxyConfig: run_ollama_serve() ## ASSISTANT SETTINGS - assistants_config: Optional[AssistantsTypedDict] = None + assistants_config: AssistantsTypedDict | None = None assistant_settings = config.get("assistant_settings", None) if assistant_settings: for k, v in assistant_settings["litellm_params"].items(): @@ -4983,7 +4978,7 @@ class ProxyConfig: assistants_config = AssistantsTypedDict(**assistant_settings) # type: ignore ## SEARCH TOOLS SETTINGS - search_tools: Optional[List[SearchToolTypedDict]] = self.parse_search_tools(config) + search_tools: list[SearchToolTypedDict] | None = self.parse_search_tools(config) ## SANDBOX TOOLS SETTINGS from litellm.sandbox.sandbox_tools import register_sandbox_tools @@ -5045,7 +5040,7 @@ class ProxyConfig: router._update_redis_cache(cache=redis_usage_cache) # Guardrail settings - guardrails_v2: Optional[List[Dict]] = None + guardrails_v2: list[dict] | None = None if config is not None: guardrails_v2 = config.get("guardrails", None) @@ -5064,7 +5059,7 @@ class ProxyConfig: ) ## Prompt settings - prompts: Optional[List[Dict]] = None + prompts: list[dict] | None = None if config is not None: prompts = config.get("prompts", None) if prompts: @@ -5081,7 +5076,7 @@ class ProxyConfig: return router, router.get_model_list(), general_settings - async def _init_non_llm_configs(self, config: dict, config_file_path: Optional[str] = None): + async def _init_non_llm_configs(self, config: dict, config_file_path: str | None = None): """ Initialize non-LLM configs eg. MCP tools, vector stores, etc. """ @@ -5135,7 +5130,7 @@ class ProxyConfig: async def _init_policy_engine( self, - config: Optional[dict], + config: dict | None, prisma_client: Optional["PrismaClient"], llm_router: Optional["Router"], ): @@ -5215,8 +5210,8 @@ class ProxyConfig: def initialize_secret_manager( self, - key_management_system: Optional[str], - config_file_path: Optional[str] = None, + key_management_system: str | None, + config_file_path: str | None = None, ): """ Initialize the relevant secret manager if `key_management_system` is provided @@ -5278,7 +5273,7 @@ class ProxyConfig: Return model info w/ id """ - _id: Optional[str] = getattr(model, "model_id", None) + _id: str | None = getattr(model, "model_id", None) if _id is not None: model.model_info["id"] = _id model.model_info["db_model"] = True @@ -5451,7 +5446,7 @@ class ProxyConfig: async def _update_llm_router( self, - new_models: Optional[Json], + new_models: Json | None, proxy_logging_obj: ProxyLogging, ) -> frozenset[str] | None: global llm_router, llm_model_list, master_key, general_settings @@ -5535,7 +5530,7 @@ class ProxyConfig: def _add_callback_from_db_to_in_memory_litellm_callbacks( self, callback: str, - event_types: List[Literal["success", "failure"]], + event_types: list[Literal["success", "failure"]], existing_callbacks: list, ) -> None: """ @@ -5590,7 +5585,7 @@ class ProxyConfig: existing_callbacks=litellm.callbacks, ) - def _encrypt_env_variables(self, environment_variables: dict, new_encryption_key: Optional[str] = None) -> dict: + def _encrypt_env_variables(self, environment_variables: dict, new_encryption_key: str | None = None) -> dict: """ Encrypts a dictionary of environment variables and returns them. """ @@ -5631,9 +5626,7 @@ class ProxyConfig: decrypted_variables[k] = decrypted_value return decrypted_variables - def _encrypt_env_variables_for_db( - self, environment_variables: dict, new_encryption_key: Optional[str] = None - ) -> dict: + def _encrypt_env_variables_for_db(self, environment_variables: dict, new_encryption_key: str | None = None) -> dict: """ Idempotently encrypt environment variables for a DB write. @@ -5655,7 +5648,7 @@ class ProxyConfig: ) @staticmethod - def _parse_router_settings_value(value: Any) -> Optional[dict]: + def _parse_router_settings_value(value: Any) -> dict | None: """ Parse a router_settings value that may be a dict or a JSON/YAML string. @@ -5664,7 +5657,7 @@ class ProxyConfig: if value is None: return None - parsed: Optional[dict] = None + parsed: dict | None = None if isinstance(value, dict): parsed = value elif isinstance(value, str): @@ -5685,9 +5678,9 @@ class ProxyConfig: async def _get_hierarchical_router_settings( self, user_api_key_dict: Optional["UserAPIKeyAuth"], - prisma_client: Optional[PrismaClient], + prisma_client: PrismaClient | None, proxy_logging_obj: Optional["ProxyLogging"] = None, - ) -> Optional[dict]: + ) -> dict | None: """ Get router_settings in priority order: Key > Team @@ -5731,8 +5724,8 @@ class ProxyConfig: async def _add_router_settings_from_db_config( self, config_data: dict, - llm_router: Optional[Router], - prisma_client: Optional[PrismaClient], + llm_router: Router | None, + prisma_client: PrismaClient | None, ) -> None: """ Adds router settings from DB config to litellm proxy @@ -5901,7 +5894,7 @@ class ProxyConfig: except ValueError: verbose_proxy_logger.error("Invalid maximum_spend_logs_retention_interval value") - async def _update_general_settings(self, db_general_settings: Optional[Json]): + async def _update_general_settings(self, db_general_settings: Json | None): """ Pull from DB, read general settings value """ @@ -6072,7 +6065,7 @@ class ProxyConfig: self, prisma_client: PrismaClient, config: dict, - store_model_in_db: Optional[bool], + store_model_in_db: bool | None, ): if store_model_in_db is not True: verbose_proxy_logger.info("'store_model_in_db' is not True, skipping db updates") @@ -6110,7 +6103,7 @@ class ProxyConfig: return config - def _should_load_db_object(self, object_type: Union[str, SupportedDBObjectType]) -> bool: + def _should_load_db_object(self, object_type: str | SupportedDBObjectType) -> bool: """ Check if an object type should be loaded from the database based on general_settings.supported_db_objects. @@ -6142,7 +6135,7 @@ class ProxyConfig: # Check if the object type is in the list (supports both str and enum values) return any(str(obj) == object_type_str for obj in supported_db_objects) - async def _get_models_from_db(self, prisma_client: PrismaClient) -> Optional[list]: + async def _get_models_from_db(self, prisma_client: PrismaClient) -> list | None: """ Fetch all model deployments from the DB. @@ -6152,7 +6145,7 @@ class ProxyConfig: as "all models deleted" and must not evict existing router deployments. """ try: - new_models = await ModelRepository(prisma_client).table.find_many() + new_models: list[PrismaProxyModelTable] = list(await ModelRepository(prisma_client).table.find_many()) return new_models except Exception as e: verbose_proxy_logger.exception( @@ -6637,7 +6630,7 @@ class ProxyConfig: from litellm.types.prompts.init_prompts import PromptSpec try: - prompts_in_db = await PromptRepository(prisma_client).table.find_many() + prompts_in_db: Sequence[PrismaPromptTable] = await PromptRepository(prisma_client).table.find_many() for prompt in prompts_in_db: # Convert DB object to dict and create versioned prompt_id prompt_spec = self._get_prompt_spec_for_db_prompt(db_prompt=prompt) @@ -6655,7 +6648,7 @@ class ProxyConfig: ) try: - guardrails_in_db: List[Guardrail] = await GuardrailRegistry.get_all_guardrails_from_db( + guardrails_in_db: list[Guardrail] = await GuardrailRegistry.get_all_guardrails_from_db( prisma_client=prisma_client ) verbose_proxy_logger.debug("guardrails from the DB %s", str(guardrails_in_db)) @@ -6909,7 +6902,7 @@ class ProxyConfig: await initialize_pass_through_endpoints_in_db() - def decrypt_credentials(self, credential: Union[dict, BaseModel]) -> CredentialItem: + def decrypt_credentials(self, credential: dict | BaseModel) -> CredentialItem: if isinstance(credential, dict): credential_object = CredentialItem(**credential) elif isinstance(credential, BaseModel): @@ -6922,7 +6915,7 @@ class ProxyConfig: credential_object.credential_values = decrypted_credential_values return credential_object - async def delete_credentials(self, db_credentials: List[CredentialItem]): + async def delete_credentials(self, db_credentials: list[CredentialItem]): """ Create all-up list of db credentials + local credentials Compare to the litellm.credential_list @@ -7300,7 +7293,7 @@ def _restamp_streaming_chunk_model( def _fast_serialize_simple_model_response_stream( chunk: ModelResponseStream, -) -> Optional[bytes]: +) -> bytes | None: """ Serialize the common OpenAI text streaming chunk without the full Pydantic serializer. Fall back for richer chunks so tool calls, logprobs, usage, and @@ -7375,7 +7368,7 @@ def _fast_serialize_simple_model_response_stream( return orjson.dumps(payload) -def _serialize_streaming_chunk(chunk: BaseModel) -> Union[str, bytes]: +def _serialize_streaming_chunk(chunk: BaseModel) -> str | bytes: if isinstance(chunk, ModelResponseStream): serialized_chunk = _fast_serialize_simple_model_response_stream(chunk) if serialized_chunk is not None: @@ -7409,7 +7402,7 @@ async def _apply_streaming_chunk_hooks( user_api_key_dict: UserAPIKeyAuth, request_data: dict, str_so_far: str, -) -> Tuple[Any, str]: +) -> tuple[Any, str]: chunk = await proxy_logging_obj.async_post_call_streaming_hook( user_api_key_dict=user_api_key_dict, response=chunk, @@ -7424,7 +7417,7 @@ async def _apply_streaming_chunk_hooks( return chunk, str_so_far -def _format_streaming_sse_chunk(chunk: Union[str, bytes]) -> Union[str, bytes]: +def _format_streaming_sse_chunk(chunk: str | bytes) -> str | bytes: if isinstance(chunk, bytes): return b"data: " + chunk + b"\n\n" return f"data: {chunk}\n\n" @@ -7456,7 +7449,7 @@ async def async_data_generator( stream_completed = False client_disconnected = False try: - error_message: Optional[str] = None + error_message: str | None = None requested_model_from_client = _get_client_requested_model_for_streaming(request_data=request_data) ( fallback_was_attempted, @@ -7717,9 +7710,9 @@ class ProxyStartupEvent: @classmethod def _initialize_startup_logging( cls, - llm_router: Optional[Router], + llm_router: Router | None, proxy_logging_obj: ProxyLogging, - redis_usage_cache: Optional[RedisCache], + redis_usage_cache: RedisCache | None, ): """Initialize logging and alerting on startup""" ## COST TRACKING ## @@ -7761,7 +7754,7 @@ class ProxyStartupEvent: @staticmethod def _validate_redis_transaction_buffer_config( general_settings: dict, - redis_usage_cache: Optional[RedisCache], + redis_usage_cache: RedisCache | None, ): """ Validates that when use_redis_transaction_buffer is enabled, @@ -7769,9 +7762,7 @@ class ProxyStartupEvent: """ from litellm.secret_managers.main import str_to_bool - _use_redis_transaction_buffer: Optional[Union[bool, str]] = general_settings.get( - "use_redis_transaction_buffer", False - ) + _use_redis_transaction_buffer: bool | str | None = general_settings.get("use_redis_transaction_buffer", False) if isinstance(_use_redis_transaction_buffer, str): _use_redis_transaction_buffer = str_to_bool(_use_redis_transaction_buffer) @@ -7793,7 +7784,7 @@ class ProxyStartupEvent: @staticmethod async def _init_coordination_redis_from_db( litellm_settings: Mapping[str, object], - llm_router: Optional[Router], + llm_router: Router | None, ) -> RedisCache | None: """ Applies a coordination_redis block saved to the database, which the admin @@ -7856,8 +7847,8 @@ class ProxyStartupEvent: @classmethod async def _initialize_semantic_tool_filter( cls, - llm_router: Optional[Router], - litellm_settings: Dict[str, Any], + llm_router: Router | None, + litellm_settings: dict[str, Any], ): """Initialize MCP semantic tool filter if configured""" from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook @@ -7889,7 +7880,7 @@ class ProxyStartupEvent: def _initialize_jwt_auth( cls, general_settings: dict, - prisma_client: Optional[PrismaClient], + prisma_client: PrismaClient | None, user_api_key_cache: UserApiKeyCache, ): """Initialize JWT auth on startup""" @@ -8432,7 +8423,7 @@ class ProxyStartupEvent: LITELLM_KEY_ROTATION_ENABLED, ) - key_rotation_enabled: Optional[bool] = str_to_bool(LITELLM_KEY_ROTATION_ENABLED) + key_rotation_enabled: bool | None = str_to_bool(LITELLM_KEY_ROTATION_ENABLED) verbose_proxy_logger.debug(f"key_rotation_enabled: {key_rotation_enabled}") if key_rotation_enabled is True: @@ -8485,7 +8476,7 @@ class ProxyStartupEvent: LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_INTERVAL_SECONDS, ) - expired_ui_session_key_cleanup_enabled: Optional[bool] = str_to_bool( + expired_ui_session_key_cleanup_enabled: bool | None = str_to_bool( LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED ) verbose_proxy_logger.debug(f"expired_ui_session_key_cleanup_enabled: {expired_ui_session_key_cleanup_enabled}") @@ -8585,16 +8576,16 @@ class ProxyStartupEvent: @classmethod async def _setup_prisma_client( cls, - database_url: Optional[str], + database_url: str | None, proxy_logging_obj: ProxyLogging, user_api_key_cache: UserApiKeyCache, - ) -> Optional[PrismaClient]: + ) -> PrismaClient | None: """ - Sets up prisma client - Adds necessary views to proxy """ try: - prisma_client: Optional[PrismaClient] = None + prisma_client: PrismaClient | None = None if database_url is not None: try: prisma_client = PrismaClient(database_url=database_url, proxy_logging_obj=proxy_logging_obj) @@ -8752,14 +8743,14 @@ class ProxyStartupEvent: ) # if project requires model list async def model_list( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - return_wildcard_routes: Optional[bool] = False, - team_id: Optional[str] = None, - include_model_access_groups: Optional[bool] = False, - only_model_access_groups: Optional[bool] = False, - include_metadata: Optional[bool] = False, - fallback_type: Optional[str] = None, - scope: Optional[str] = None, - healthy_only: Optional[bool] = False, + return_wildcard_routes: bool | None = False, + team_id: str | None = None, + include_model_access_groups: bool | None = False, + only_model_access_groups: bool | None = False, + include_metadata: bool | None = False, + fallback_type: str | None = None, + scope: str | None = None, + healthy_only: bool | None = False, ): """ Use `/model/info` - to get detailed model information, example - pricing, mode, etc. @@ -8936,8 +8927,8 @@ async def model_list( async def model_info( model_id: str, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - team_id: Optional[str] = None, - healthy_only: Optional[bool] = False, + team_id: str | None = None, + healthy_only: bool | None = False, ): """ Retrieve information about a specific model accessible to your API key. @@ -9021,7 +9012,7 @@ async def model_info( ) -def _blocked_response_usage(original_response: Optional[Any]) -> "litellm.Usage": +def _blocked_response_usage(original_response: Any | None) -> "litellm.Usage": """ Token usage for a synthetic guardrail-blocked response. @@ -9061,7 +9052,7 @@ def _blocked_response_usage(original_response: Optional[Any]) -> "litellm.Usage" async def chat_completion( request: Request, fastapi_response: Response, - model: Optional[str] = None, + model: str | None = None, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -9136,6 +9127,7 @@ async def chat_completion( _data = e.request_data # Capture logging_obj before post_call_failure_hook pops it from _data. _logging_obj = _data.get("litellm_logging_obj") + assert isinstance(_logging_obj, LiteLLMLoggingObj) await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, @@ -9182,11 +9174,13 @@ async def chat_completion( if data.get("stream", None) is not None and data["stream"] is True: _iterator = litellm.utils.ModelResponseIterator(model_response=_chat_response, convert_to_delta=True) + _rejected_logging_obj = _data.get("litellm_logging_obj") + assert isinstance(_rejected_logging_obj, LiteLLMLoggingObj) _streaming_response = litellm.CustomStreamWrapper( completion_stream=_iterator, model=data.get("model", ""), custom_llm_provider="cached_response", - logging_obj=_data.get("litellm_logging_obj", None), + logging_obj=_rejected_logging_obj, ) selected_data_generator = select_data_generator( response=_streaming_response, @@ -9226,7 +9220,7 @@ async def chat_completion( async def completion( request: Request, fastapi_response: Response, - model: Optional[str] = None, + model: str | None = None, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -9406,7 +9400,7 @@ async def completion( async def embeddings( request: Request, fastapi_response: Response, - model: Optional[str] = None, + model: str | None = None, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -9538,7 +9532,7 @@ async def moderations( ``` """ global proxy_logging_obj - data: Dict = {} + data: dict = {} try: # Use orjson to parse JSON data, orjson speeds up requests significantly body = await request.body() @@ -9660,7 +9654,7 @@ async def audio_speech( https://platform.openai.com/docs/api-reference/audio/createSpeech """ global proxy_logging_obj - data: Dict = {} + data: dict = {} try: # Use orjson to parse JSON data, orjson speeds up requests significantly body = await request.body() @@ -9782,7 +9776,7 @@ async def audio_transcriptions( https://platform.openai.com/docs/api-reference/audio/createTranscription?lang=curl """ global proxy_logging_obj - data: Dict = {} + data: dict = {} try: # Use orjson to parse JSON data, orjson speeds up requests significantly form_data = await get_form_data(request) @@ -9928,15 +9922,15 @@ async def audio_transcriptions( @app.websocket("/vertex_ai/live") async def vertex_ai_live_passthrough_endpoint( websocket: WebSocket, - model: Optional[str] = fastapi.Query( + model: str | None = fastapi.Query( None, description="Optional model name, used to determine Vertex region for global models.", ), - vertex_project: Optional[str] = fastapi.Query( + vertex_project: str | None = fastapi.Query( None, description="Override the Vertex AI project id used for the upstream connection.", ), - vertex_location: Optional[str] = fastapi.Query( + vertex_location: str | None = fastapi.Query( None, description="Override the Vertex AI region (for example, 'us-central1').", ), @@ -9964,12 +9958,12 @@ async def vertex_ai_live_passthrough_endpoint( @lru_cache(maxsize=_REALTIME_BODY_CACHE_SIZE) -def _realtime_query_params_template(model: Optional[str], intent: Optional[str]) -> Tuple[Tuple[str, str], ...]: +def _realtime_query_params_template(model: str | None, intent: str | None) -> tuple[tuple[str, str], ...]: """ Build a hashable representation of the realtime query params so we can cache the repetitive model/intent combinations. """ - params: List[Tuple[str, str]] = [] + params: list[tuple[str, str]] = [] if model is not None: params.append(("model", model)) if intent is not None: @@ -9982,9 +9976,9 @@ def _realtime_query_params_template(model: Optional[str], intent: Optional[str]) @app.websocket("/realtime") async def realtime_websocket_endpoint( websocket: WebSocket, - model: Optional[str] = fastapi.Query(None, description="The model to use for the websocket connection."), - intent: Optional[str] = fastapi.Query(None, description="The intent of the websocket connection."), - guardrails: Optional[str] = fastapi.Query( + model: str | None = fastapi.Query(None, description="The model to use for the websocket connection."), + intent: str | None = fastapi.Query(None, description="The intent of the websocket connection."), + guardrails: str | None = fastapi.Query( None, description="Comma-separated list of guardrail names to apply to this request.", ), @@ -10020,7 +10014,7 @@ async def realtime_websocket_endpoint( # Only use explicit parameters, not all query params query_params = cast(RealtimeQueryParams, dict(_realtime_query_params_template(model, intent))) - data: Dict[str, Any] = { + data: dict[str, Any] = { "model": route_model, "websocket": websocket, "query_params": query_params, # Only explicit params @@ -10136,7 +10130,7 @@ async def get_assistants( API Reference docs - https://platform.openai.com/docs/api-reference/assistants/listAssistants """ global proxy_logging_obj - data: Dict = {} + data: dict = {} try: # Use orjson to parse JSON data, orjson speeds up requests significantly await request.body() @@ -10319,7 +10313,7 @@ async def delete_assistant( API Reference docs - https://platform.openai.com/docs/api-reference/assistants/createAssistant """ global proxy_logging_obj - data: Dict = {} + data: dict = {} try: # Use orjson to parse JSON data, orjson speeds up requests significantly @@ -10409,7 +10403,7 @@ async def create_threads( API Reference - https://platform.openai.com/docs/api-reference/threads/createThread """ global proxy_logging_obj - data: Dict = {} + data: dict = {} try: # Use orjson to parse JSON data, orjson speeds up requests significantly await request.body() @@ -10499,7 +10493,7 @@ async def get_thread( API Reference - https://platform.openai.com/docs/api-reference/threads/getThread """ global proxy_logging_obj - data: Dict = {} + data: dict = {} try: # Include original request and headers in the data data = await add_litellm_data_to_request( @@ -10586,7 +10580,7 @@ async def add_messages( API Reference - https://platform.openai.com/docs/api-reference/messages/createMessage """ global proxy_logging_obj - data: Dict = {} + data: dict = {} try: # Use orjson to parse JSON data, orjson speeds up requests significantly body = await request.body() @@ -10677,7 +10671,7 @@ async def get_messages( API Reference - https://platform.openai.com/docs/api-reference/messages/listMessages """ global proxy_logging_obj - data: Dict = {} + data: dict = {} try: # Include original request and headers in the data data = await add_litellm_data_to_request( @@ -10764,7 +10758,7 @@ async def run_thread( API Reference: https://platform.openai.com/docs/api-reference/runs/createRun """ global proxy_logging_obj - data: Dict = {} + data: dict = {} try: body = await request.body() data = orjson.loads(body) @@ -10868,7 +10862,7 @@ from litellm.repositories.user_repository import UserRepository def _get_provider_token_counter( deployment: dict, model_to_use: str -) -> Tuple[Optional[BaseTokenCounter], Optional[str], Optional[str]]: +) -> tuple[BaseTokenCounter | None, str | None, str | None]: """ Auto-route to the correct provider's token counter based on model/deployment. Uses the existing get_provider_model_info infrastructure with switch-case pattern. @@ -10879,8 +10873,8 @@ def _get_provider_token_counter( from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider full_model = deployment.get("litellm_params", {}).get("model", "") - model: Optional[str] = None - custom_llm_provider: Optional[str] = None + model: str | None = None + custom_llm_provider: str | None = None try: # Use existing LiteLLM logic to determine provider @@ -10923,14 +10917,14 @@ def _get_provider_token_counter( async def _try_provider_token_count( provider_counter: "BaseTokenCounter", - custom_llm_provider: Optional[str], + custom_llm_provider: str | None, model_to_use: str, - messages: Optional[list], - contents: Optional[list], - deployment: Optional[Dict[str, Any]], + messages: list | None, + contents: list | None, + deployment: dict[str, Any] | None, request_model: str, - tools: Optional[list] = None, - system: Optional[str] = None, + tools: list | None = None, + system: str | None = None, ) -> Optional["TokenCountResponse"]: """Attempt provider-specific token counting. Returns result on success, None to fall through to local counting.""" if not provider_counter.should_use_token_counting_api(custom_llm_provider=custom_llm_provider): @@ -11001,9 +10995,9 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) if prompt is None and messages is None and contents is None: raise HTTPException(status_code=400, detail="prompt or messages or contents must be provided") - deployment: Optional[Dict[str, Any]] = None + deployment: dict[str, Any] | None = None litellm_model_name = None - model_info: Optional[ModelMapInfo] = None + model_info: ModelMapInfo | None = None if llm_router is not None: # get 1 deployment corresponding to the model try: @@ -11029,8 +11023,8 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) ) # use litellm model name, if it's not avalable then fallback to request.model # Try provider-specific token counting first - only for non-direct requests (from provider endpoints) - provider_counter: Optional[BaseTokenCounter] = None - custom_llm_provider: Optional[str] = None + provider_counter: BaseTokenCounter | None = None + custom_llm_provider: str | None = None if call_endpoint is True and deployment is not None: # Auto-route to the correct provider based on model provider_counter, _model, custom_llm_provider = _get_provider_token_counter(deployment, model_to_use) @@ -11062,10 +11056,10 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) ) # Default LiteLLM token counting - custom_tokenizer: Optional[CustomHuggingfaceTokenizer] = None + custom_tokenizer: CustomHuggingfaceTokenizer | None = None if model_info is not None: custom_tokenizer = cast( - Optional[CustomHuggingfaceTokenizer], + CustomHuggingfaceTokenizer | None, model_info.get("custom_tokenizer", None), ) _tokenizer_used = litellm.utils._select_tokenizer(model=model_to_use, custom_tokenizer=custom_tokenizer) @@ -11136,10 +11130,10 @@ async def transform_request(request: TransformRequestBody): async def _check_if_model_is_user_added( - models: List[Dict], + models: list[dict], user_api_key_dict: UserAPIKeyAuth, - prisma_client: Optional[PrismaClient], -) -> List[Dict]: + prisma_client: PrismaClient | None, +) -> list[dict]: """ Check if model is in db @@ -11157,36 +11151,38 @@ async def _check_if_model_is_user_added( id = model.get("model_info", {}).get("id", None) if id is None: continue - db_model = await ModelRepository(prisma_client).table.find_unique(where={"model_id": id}) + db_model: PrismaProxyModelTable | None = await ModelRepository(prisma_client).table.find_unique( + where={"model_id": id} + ) if db_model is not None: if db_model.created_by == user_api_key_dict.user_id: filtered_models.append(model) return filtered_models -def _check_if_model_is_team_model(models: List[DeploymentTypedDict], user_row: LiteLLM_UserTable) -> List[Dict]: +def _check_if_model_is_team_model(models: list[DeploymentTypedDict], user_row: LiteLLM_UserTable) -> list[dict]: """ Check if model is a team model Check if user is a member of the team that the model belongs to """ - user_team_models: List[Dict] = [] + user_team_models: list[dict] = [] for model in models: model_team_id = model.get("model_info", {}).get("team_id", None) if model_team_id is not None: if model_team_id in user_row.teams: - user_team_models.append(cast(Dict, model)) + user_team_models.append(cast(dict, model)) return user_team_models async def non_admin_all_models( - all_models: List[Dict], + all_models: list[dict], llm_router: Router, user_api_key_dict: UserAPIKeyAuth, - prisma_client: Optional[PrismaClient], + prisma_client: PrismaClient | None, ): """ Check if model is in db @@ -11228,13 +11224,13 @@ async def non_admin_all_models( def _add_team_models_to_all_models( - team_db_objects_typed: List[LiteLLM_TeamTable], + team_db_objects_typed: list[LiteLLM_TeamTable], llm_router: Router, -) -> Dict[str, Set[str]]: +) -> dict[str, Set[str]]: """ Add team models to all models """ - team_models: Dict[str, Set[str]] = {} + team_models: dict[str, Set[str]] = {} for team_object in team_db_objects_typed: if ( @@ -11269,11 +11265,11 @@ def _add_team_models_to_all_models( async def _add_access_group_models_to_team_models( - team_db_objects_typed: List[LiteLLM_TeamTable], + team_db_objects_typed: list[LiteLLM_TeamTable], llm_router: Router, prisma_client: PrismaClient, - team_models: Dict[str, Set[str]], -) -> Dict[str, Set[str]]: + team_models: dict[str, Set[str]], +) -> dict[str, Set[str]]: """ Resolve models reachable via team access groups and merge them into team_models. @@ -11284,7 +11280,7 @@ async def _add_access_group_models_to_team_models( (not directly in team.models) are included in the UI model listing. """ # First pass: identify eligible teams and collect all distinct access group IDs - eligible_teams: List[LiteLLM_TeamTable] = [] + eligible_teams: list[LiteLLM_TeamTable] = [] all_access_group_ids: Set[str] = set() for team_object in team_db_objects_typed: @@ -11303,10 +11299,10 @@ async def _add_access_group_models_to_team_models( return team_models # Single batch fetch for all access groups - access_group_rows = await AccessGroupRepository(prisma_client).table.find_many( + access_group_rows: Sequence[PrismaAccessGroupTable] = await AccessGroupRepository(prisma_client).table.find_many( where={"access_group_id": {"in": list(all_access_group_ids)}} ) - ag_model_map: Dict[str, List[str]] = { + ag_model_map: dict[str, list[str]] = { row.access_group_id: row.access_model_names or [] for row in access_group_rows } @@ -11328,10 +11324,10 @@ async def _add_access_group_models_to_team_models( async def get_all_team_models( - user_teams: Union[List[str], Literal["*"]], + user_teams: list[str] | Literal["*"], prisma_client: PrismaClient, llm_router: Router, -) -> Dict[str, List[str]]: +) -> dict[str, list[str]]: """ Get all models across all teams user is in. @@ -11340,10 +11336,10 @@ async def get_all_team_models( 3. Return {"model_id": ["team_id1", "team_id2"]} """ - team_db_objects_typed: List[LiteLLM_TeamTable] = [] + team_db_objects_typed: list[LiteLLM_TeamTable] = [] if user_teams == "*": - team_db_objects = await TeamRepository(prisma_client).table.find_many() + team_db_objects: Sequence[PrismaTeamTable] = await TeamRepository(prisma_client).table.find_many() team_db_objects_typed = [ LiteLLM_TeamTable.model_validate(team_db_object.model_dump()) for team_db_object in team_db_objects ] @@ -11368,7 +11364,7 @@ async def get_all_team_models( ) # convert set to list - returned_team_models: Dict[str, List[str]] = {} + returned_team_models: dict[str, list[str]] = {} for model_id, team_ids in team_models.items(): returned_team_models[model_id] = list(team_ids) @@ -11378,7 +11374,7 @@ async def get_all_team_models( def get_direct_access_models( user_db_object: LiteLLM_UserTable, llm_router: Router, -) -> List[str]: +) -> list[str]: """ Get all models that user has direct access to. @@ -11396,7 +11392,7 @@ def get_direct_access_models( ] -def _filter_models_to_user_accessible(all_models: List[Dict]) -> List[Dict]: +def _filter_models_to_user_accessible(all_models: list[dict]) -> list[dict]: """Keep only deployments the caller can use via direct access or team membership.""" return [ _model @@ -11410,19 +11406,19 @@ async def _populate_team_access_on_models( user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient, llm_router: Router, - all_models: List[Dict], -) -> List[Dict]: + all_models: list[dict], +) -> list[dict]: """ Populate `model_info.access_via_team_ids` and `model_info.direct_access` without filtering the model list. """ - user_teams: Optional[Union[List[str], Literal["*"]]] = None - direct_access_models: List[str] = [] + user_teams: list[str] | Literal["*"] | None = None + direct_access_models: list[str] = [] if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: user_teams = "*" direct_access_models = llm_router.get_model_ids(exclude_team_models=True) # has access to all models elif user_api_key_dict.user_id is not None: - user_db_object = await UserRepository(prisma_client).table.find_unique( + user_db_object: PrismaUserTable | None = await UserRepository(prisma_client).table.find_unique( where={"user_id": user_api_key_dict.user_id} ) if user_db_object is not None: @@ -11465,8 +11461,8 @@ async def get_all_team_and_direct_access_models( user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient, llm_router: Router, - all_models: List[Dict], -) -> List[Dict]: + all_models: list[dict], +) -> list[dict]: """ Get all models across all teams user is in. """ @@ -11480,8 +11476,8 @@ async def get_all_team_and_direct_access_models( def _enrich_model_info_with_litellm_data( - model: Dict[str, Any], debug: bool = False, llm_router: Optional[Router] = None -) -> Dict[str, Any]: + model: dict[str, Any], debug: bool = False, llm_router: Router | None = None +) -> dict[str, Any]: """ Enrich a model dictionary with litellm model info (pricing, context window, etc.) and remove sensitive information. @@ -11542,9 +11538,9 @@ def _enrich_model_info_with_litellm_data( async def _get_caller_byok_team_scope( - user_api_key_dict: Optional[UserAPIKeyAuth], - prisma_client: Optional[Any], -) -> Optional[Set[str]]: + user_api_key_dict: UserAPIKeyAuth | None, + prisma_client: Any | None, +) -> Set[str] | None: """ Return the team IDs whose BYOK rows the caller is allowed to see via `/v2/model/info` search results. @@ -11567,7 +11563,9 @@ async def _get_caller_byok_team_scope( if user_id is None: return key_team_scope try: - user_row = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id}) + user_row: PrismaUserTable | None = await UserRepository(prisma_client).table.find_unique( + where={"user_id": user_id} + ) except Exception: verbose_proxy_logger.exception( "Failed to look up caller teams while scoping BYOK search; defaulting to key team scope only." @@ -11578,7 +11576,7 @@ async def _get_caller_byok_team_scope( return key_team_scope | set(user_row.teams or []) -def _byok_row_outside_caller_teams(model_info_dict: Dict[str, Any], allowed_team_ids: Optional[Set[str]]) -> bool: +def _byok_row_outside_caller_teams(model_info_dict: dict[str, Any], allowed_team_ids: Set[str] | None) -> bool: """Whether a team BYOK row belongs to a team the caller is not a member of. `team_id` is only set on team BYOK rows; non-team rows fall through @@ -11607,9 +11605,9 @@ async def _fetch_db_models_for_search( router_models_count: int, page: int, size: int, - sort_by: Optional[str], - is_byok_outside_caller_teams: Callable[[Dict[str, Any]], bool], -) -> Tuple[List[Dict[str, Any]], int]: + sort_by: str | None, + is_byok_outside_caller_teams: Callable[[dict[str, Any]], bool], +) -> tuple[list[dict[str, Any]], int]: """ Run the bounded DB query that backs `/v2/model/info?search=`. Returns `(decrypted_models, total_count)` where `total_count` is the cheap @@ -11625,7 +11623,7 @@ async def _fetch_db_models_for_search( filter for `team_public_model_name` instead and keep the DB cost bounded by `search`. """ - db_where_condition: Dict[str, Any] = {"model_name": {"contains": search_lower, "mode": "insensitive"}} + db_where_condition: dict[str, Any] = {"model_name": {"contains": search_lower, "mode": "insensitive"}} if db_model_ids_in_router: db_where_condition["model_id"] = {"not": {"in": list(db_model_ids_in_router)}} @@ -11639,7 +11637,7 @@ async def _fetch_db_models_for_search( db_models_total_count = await ModelRepository(prisma_client).table.count(where=db_where_condition) - db_models_raw: list = [] + db_models_raw: Sequence[PrismaProxyModelTable] = [] if take_limit > 0: db_models_raw = await ModelRepository(prisma_client).table.find_many( where=db_where_condition, @@ -11654,7 +11652,7 @@ async def _fetch_db_models_for_search( if not is_byok_outside_caller_teams(m.model_info if isinstance(m.model_info, dict) else {}) ] - decrypted: List[Dict[str, Any]] = [] + decrypted: list[dict[str, Any]] = [] for db_model in matching_db_rows: decrypted_models = proxy_config.decrypt_model_list_from_db([db_model]) if decrypted_models: @@ -11664,15 +11662,15 @@ async def _fetch_db_models_for_search( async def _apply_search_filter_to_models( - all_models: List[Dict[str, Any]], + all_models: list[dict[str, Any]], search: str, - prisma_client: Optional[Any], + prisma_client: Any | None, proxy_config: Any, - user_api_key_dict: Optional[UserAPIKeyAuth] = None, + user_api_key_dict: UserAPIKeyAuth | None = None, page: int = 1, size: int = 50, - sort_by: Optional[str] = None, -) -> Tuple[List[Dict[str, Any]], Optional[int]]: + sort_by: str | None = None, +) -> tuple[list[dict[str, Any]], int | None]: """ Apply search filter to models, querying database for additional matching models. @@ -11706,10 +11704,10 @@ async def _apply_search_filter_to_models( prisma_client=prisma_client, ) - def _is_byok_outside_caller_teams(model_info_dict: Dict[str, Any]) -> bool: + def _is_byok_outside_caller_teams(model_info_dict: dict[str, Any]) -> bool: return _byok_row_outside_caller_teams(model_info_dict, allowed_team_ids) - def _model_matches_search(m: Dict[str, Any]) -> bool: + def _model_matches_search(m: dict[str, Any]) -> bool: # Team BYOK models persist an internal `model_name` # (e.g. `model_name_{team_id}_{uuid}`) and expose the user-facing # name via `model_info.team_public_model_name`. Match both so the @@ -11748,7 +11746,7 @@ async def _apply_search_filter_to_models( router_models_count = config_models_count + db_models_in_router_count # Query database for additional models with search term - db_models: List[Dict[str, Any]] = [] + db_models: list[dict[str, Any]] = [] if prisma_client is not None: try: db_models, db_models_total_count = await _fetch_db_models_for_search( @@ -11772,7 +11770,7 @@ async def _apply_search_filter_to_models( return filtered_router_models + db_models, search_total_count -def _normalize_datetime_for_sorting(dt: Any) -> Optional[datetime]: +def _normalize_datetime_for_sorting(dt: Any) -> datetime | None: """ Normalize a datetime value to a timezone-aware UTC datetime for sorting. @@ -11815,10 +11813,10 @@ def _normalize_datetime_for_sorting(dt: Any) -> Optional[datetime]: def _sort_models( - all_models: List[Dict[str, Any]], - sort_by: Optional[str], + all_models: list[dict[str, Any]], + sort_by: str | None, sort_order: str = "asc", -) -> List[Dict[str, Any]]: +) -> list[dict[str, Any]]: """ Sort models by the specified field and order. @@ -11841,7 +11839,7 @@ def _sort_models( reverse = sort_order.lower() == "desc" - def get_sort_key(model: Dict[str, Any]) -> Any: + def get_sort_key(model: dict[str, Any]) -> Any: model_info = model.get("model_info", {}) if sort_by == "model_name": @@ -11920,12 +11918,12 @@ def _is_auto_router_model(model: Mapping[str, object]) -> bool: def _paginate_models_response( - all_models: List[Dict[str, Any]], + all_models: list[dict[str, Any]], page: int, size: int, - total_count: Optional[int], - search: Optional[str], -) -> Dict[str, Any]: + total_count: int | None, + search: str | None, +) -> dict[str, Any]: """ Paginate models and return response dictionary. @@ -11959,9 +11957,9 @@ def _paginate_models_response( } -def _team_models_resolve_to_names(team_models: List[str], access_groups: Dict[str, Any]) -> List[str]: +def _team_models_resolve_to_names(team_models: list[str], access_groups: dict[str, Any]) -> list[str]: """Expand team model entries (including access group names) to concrete model names.""" - resolved: List[str] = [] + resolved: list[str] = [] for name in team_models: if name in access_groups: resolved.extend(access_groups[name]) @@ -11970,10 +11968,12 @@ def _team_models_resolve_to_names(team_models: List[str], access_groups: Dict[st return resolved -async def _load_team_object_for_model_filter(team_id: str, prisma_client: PrismaClient) -> Optional[LiteLLM_TeamTable]: +async def _load_team_object_for_model_filter(team_id: str, prisma_client: PrismaClient) -> LiteLLM_TeamTable | None: """Load team row from DB; returns None if missing or on error.""" try: - team_db_object = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id}) + team_db_object: PrismaTeamTable | None = await TeamRepository(prisma_client).table.find_unique( + where={"team_id": team_id} + ) if team_db_object is None: verbose_proxy_logger.warning(f"Team {team_id} not found in database") return None @@ -12022,7 +12022,7 @@ async def _gather_team_accessible_model_ids( try: if team_object.models and SpecialModelNames.all_proxy_models.value not in team_object.models: _resolved_names = _team_models_resolve_to_names(team_object.models, access_groups) - db_models = await ModelRepository(prisma_client).table.find_many( + db_models: Sequence[PrismaProxyModelTable] = await ModelRepository(prisma_client).table.find_many( where={"model_name": {"in": _resolved_names}} ) for db_model in db_models: @@ -12059,7 +12059,9 @@ async def _authorize_team_id_query( detail={"error": "Not authorized to view this team's models"}, ) try: - user_row = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id}) + user_row: PrismaUserTable | None = await UserRepository(prisma_client).table.find_unique( + where={"user_id": user_id} + ) except Exception: verbose_proxy_logger.exception("Failed to look up caller teams while authorizing teamId filter") raise HTTPException( @@ -12075,12 +12077,12 @@ async def _authorize_team_id_query( async def _filter_models_by_team_id( - all_models: List[Dict[str, Any]], + all_models: list[dict[str, Any]], team_id: str, prisma_client: PrismaClient, llm_router: Router, - user_api_key_dict: Optional[UserAPIKeyAuth] = None, -) -> List[Dict[str, Any]]: + user_api_key_dict: UserAPIKeyAuth | None = None, +) -> list[dict[str, Any]]: """ Filter models by team ID. Returns models where: - team_id matches the model's BYOK team_id, OR @@ -12143,11 +12145,11 @@ async def _filter_models_by_team_id( async def _find_model_by_id( model_id: str, - search: Optional[str], + search: str | None, llm_router, prisma_client, proxy_config, -) -> tuple[list, Optional[int]]: +) -> tuple[list, int | None]: """Find a model by its ID and optionally filter by search term.""" found_model = None @@ -12160,7 +12162,9 @@ async def _find_model_by_id( # If not found in config, search in database if found_model is None: try: - db_model = await ModelRepository(prisma_client).table.find_unique(where={"model_id": model_id}) + db_model: PrismaProxyModelTable | None = await ModelRepository(prisma_client).table.find_unique( + where={"model_id": model_id} + ) if db_model: # Convert database model to router format decrypted_models = proxy_config.decrypt_model_list_from_db([db_model]) @@ -12180,7 +12184,7 @@ async def _find_model_by_id( # Set all_models to the found model or empty list all_models = [found_model] if found_model is not None else [] - search_total_count: Optional[int] = len(all_models) + search_total_count: int | None = len(all_models) return all_models, search_total_count @@ -12191,25 +12195,25 @@ async def _find_model_by_id( ) async def model_info_v2( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - model: Optional[str] = fastapi.Query(None, description="Specify the model name (optional)"), - user_models_only: Optional[bool] = fastapi.Query(False, description="Only return models added by this user"), - include_team_models: Optional[bool] = fastapi.Query( + model: str | None = fastapi.Query(None, description="Specify the model name (optional)"), + user_models_only: bool | None = fastapi.Query(False, description="Only return models added by this user"), + include_team_models: bool | None = fastapi.Query( False, description="Return all models across all teams user is in." ), - debug: Optional[bool] = False, + debug: bool | None = False, page: int = Query(1, description="Page number", ge=1), size: int = Query(50, description="Page size", ge=1), - search: Optional[str] = fastapi.Query(None, description="Search model names (case-insensitive partial match)"), - modelId: Optional[str] = fastapi.Query(None, description="Search for a specific model by its unique ID"), - teamId: Optional[str] = fastapi.Query( + search: str | None = fastapi.Query(None, description="Search model names (case-insensitive partial match)"), + modelId: str | None = fastapi.Query(None, description="Search for a specific model by its unique ID"), + teamId: str | None = fastapi.Query( None, description="Filter models by team ID. Returns models with direct_access=True or teamId in access_via_team_ids", ), - sortBy: Optional[str] = fastapi.Query( + sortBy: str | None = fastapi.Query( None, description="Field to sort by. Options: model_name, created_at, updated_at, costs, status", ), - sortOrder: Optional[str] = fastapi.Query( + sortOrder: str | None = fastapi.Query( "asc", description="Sort order. Options: asc, desc", ), @@ -12418,9 +12422,9 @@ async def model_info_v2( ) async def model_streaming_metrics( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - _selected_model_group: Optional[str] = None, - startTime: Optional[datetime] = None, - endTime: Optional[datetime] = None, + _selected_model_group: str | None = None, + startTime: datetime | None = None, + endTime: datetime | None = None, ): global prisma_client, llm_router if prisma_client is None: @@ -12523,7 +12527,7 @@ async def model_streaming_metrics( """ # convert daily entries to list of dicts - response: List[dict] = [] + response: list[dict] = [] # sort daily entries by date _daily_entries = dict(sorted(_daily_entries.items(), key=lambda item: item[0])) @@ -12548,11 +12552,11 @@ async def model_streaming_metrics( ) async def model_metrics( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - _selected_model_group: Optional[str] = "gpt-4-32k", - startTime: Optional[datetime] = None, - endTime: Optional[datetime] = None, - api_key: Optional[str] = None, - customer: Optional[str] = None, + _selected_model_group: str | None = "gpt-4-32k", + startTime: datetime | None = None, + endTime: datetime | None = None, + api_key: str | None = None, + customer: str | None = None, ): global prisma_client, llm_router if prisma_client is None: @@ -12638,7 +12642,7 @@ async def model_metrics( """ # convert daily entries to list of dicts - response: List[dict] = [] + response: list[dict] = [] # sort daily entries by date _daily_entries = dict(sorted(_daily_entries.items(), key=lambda item: item[0])) @@ -12663,11 +12667,11 @@ async def model_metrics( ) async def model_metrics_slow_responses( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - _selected_model_group: Optional[str] = "gpt-4-32k", - startTime: Optional[datetime] = None, - endTime: Optional[datetime] = None, - api_key: Optional[str] = None, - customer: Optional[str] = None, + _selected_model_group: str | None = "gpt-4-32k", + startTime: datetime | None = None, + endTime: datetime | None = None, + api_key: str | None = None, + customer: str | None = None, ): global prisma_client, llm_router, proxy_logging_obj if prisma_client is None: @@ -12752,11 +12756,11 @@ ORDER BY ) async def model_metrics_exceptions( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - _selected_model_group: Optional[str] = None, - startTime: Optional[datetime] = None, - endTime: Optional[datetime] = None, - api_key: Optional[str] = None, - customer: Optional[str] = None, + _selected_model_group: str | None = None, + startTime: datetime | None = None, + endTime: datetime | None = None, + api_key: str | None = None, + customer: str | None = None, ): global prisma_client, llm_router if prisma_client is None: @@ -12798,7 +12802,7 @@ async def model_metrics_exceptions( LIMIT 200; """ db_response = await prisma_client.db.query_raw(sql_query, startTime, endTime, _selected_model_group, api_key) - response: List[dict] = [] + response: list[dict] = [] exception_types = set() """ @@ -12829,7 +12833,7 @@ async def model_metrics_exceptions( return {"data": response, "exception_types": list(exception_types)} -def _deployment_matches_allowed_model_names(model: Dict[str, Any], allowed_model_names: Set[str]) -> bool: +def _deployment_matches_allowed_model_names(model: dict[str, Any], allowed_model_names: Set[str]) -> bool: """Match a router deployment against allowed public model names. Team-scoped rows store an internal routing key in ``model_name``; callers @@ -12848,7 +12852,7 @@ def _deployment_matches_allowed_model_names(model: Dict[str, Any], allowed_model def _get_v1_model_info_allowed_model_names( user_api_key_dict: UserAPIKeyAuth, llm_router: Router, -) -> Optional[Set[str]]: +) -> Set[str] | None: """Return key/team allowlisted public model names, or None if unrestricted.""" model_access_groups = llm_router.get_model_access_groups() proxy_model_list = llm_router.get_model_names() @@ -12878,9 +12882,9 @@ def _get_v1_model_info_allowed_model_names( def _filter_v1_model_info_deployments( - all_models: List[dict], - allowed_model_names: Optional[Set[str]], -) -> List[dict]: + all_models: list[dict], + allowed_model_names: Set[str] | None, +) -> list[dict]: if allowed_model_names is None: return all_models return [model for model in all_models if _deployment_matches_allowed_model_names(model, allowed_model_names)] @@ -12964,12 +12968,12 @@ def _get_proxy_model_info(model: dict) -> dict: ) async def model_info_v1( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - litellm_model_id: Optional[str] = None, - include_team_models: Optional[bool] = fastapi.Query( + litellm_model_id: str | None = None, + include_team_models: bool | None = fastapi.Query( False, description="When true, filter to deployments the caller can use via direct access or team membership.", ), - teamId: Optional[str] = fastapi.Query( + teamId: str | None = fastapi.Query( None, description="Filter models by team ID. Returns models with direct_access=True or teamId in access_via_team_ids", ), @@ -13022,7 +13026,7 @@ async def model_info_v1( if user_model is not None: # user is trying to get specific model from litellm router try: - model_info: Dict = cast(Dict, litellm.get_model_info(model=user_model)) + model_info: dict = cast(dict, litellm.get_model_info(model=user_model)) except Exception: model_info = {} _deployment_info = Deployment( @@ -13070,7 +13074,7 @@ async def model_info_v1( detail={"error": f"Model id = {litellm_model_id} not found on litellm proxy"}, ) _deployment_info_dict = _get_proxy_model_info(model=deployment_info.model_dump(exclude_none=True)) - single_model_list: List[dict] = [_deployment_info_dict] + single_model_list: list[dict] = [_deployment_info_dict] if prisma_client is not None: single_model_list = await _populate_team_access_on_models( user_api_key_dict=user_api_key_dict, @@ -13094,7 +13098,7 @@ async def model_info_v1( # expanded model names from get_complete_model_list(). Team-scoped rows # use internal routing keys (model_name_{team_id}_{uuid}) and were omitted # when v1 resolved models only via public model_name strings. - all_models: List[dict] = copy.deepcopy(llm_router.model_list) + all_models: list[dict] = copy.deepcopy(llm_router.model_list) alias_models = copy.deepcopy(llm_router.get_model_list_from_model_alias()) all_models.extend(alias_models) @@ -13153,9 +13157,9 @@ async def model_info_v1( def _get_model_group_info( - llm_router: Router, all_models_str: List[str], model_group: Optional[str] -) -> List[ModelGroupInfoProxy]: - model_groups: List[ModelGroupInfoProxy] = [] + llm_router: Router, all_models_str: list[str], model_group: str | None +) -> list[ModelGroupInfoProxy]: + model_groups: list[ModelGroupInfoProxy] = [] unique_models = [] for model in all_models_str: @@ -13193,7 +13197,7 @@ def _get_model_group_info( ) async def model_group_info( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - model_group: Optional[str] = None, + model_group: str | None = None, ): """ Get information about all the deployments on litellm proxy, including config.yaml descriptions (except api key and api base) @@ -13367,7 +13371,7 @@ async def model_group_info( return_wildcard_routes=False, user_api_key_cache=user_api_key_cache, ) - model_groups: List[ModelGroupInfoProxy] = _get_model_group_info( + model_groups: list[ModelGroupInfoProxy] = _get_model_group_info( llm_router=llm_router, all_models_str=all_models_str, model_group=model_group ) @@ -13462,7 +13466,7 @@ async def alerting_settings( if db_general_settings is not None and db_general_settings.param_value is not None: db_general_settings_dict = dict(db_general_settings.param_value) alerting_args_dict: dict = db_general_settings_dict.get("alerting_args", {}) # type: ignore - alerting_values: Optional[list] = db_general_settings_dict.get("alerting") # type: ignore + alerting_values: list | None = db_general_settings_dict.get("alerting") # type: ignore else: alerting_args_dict = {} alerting_values = None @@ -13503,7 +13507,7 @@ async def alerting_settings( for field_name, field_info in SlackAlertingArgs.model_fields.items(): if field_name in allowed_args: - _stored_in_db: Optional[bool] = None + _stored_in_db: bool | None = None if field_name in alerting_args_dict: _stored_in_db = True else: @@ -13532,7 +13536,7 @@ async def alerting_settings( async def async_queue_request( request: Request, fastapi_response: Response, - model: Optional[str] = None, + model: str | None = None, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): global general_settings, user_debug, proxy_logging_obj @@ -13975,7 +13979,9 @@ async def onboarding(invite_link: str, request: Request): detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - invite_obj = await InvitationLinkRepository(prisma_client).table.find_unique(where={"id": invite_link}) + invite_obj: PrismaInvitationLink | None = await InvitationLinkRepository(prisma_client).table.find_unique( + where={"id": invite_link} + ) if invite_obj is None: raise HTTPException(status_code=401, detail={"error": "Invitation link does not exist in db."}) #### CHECK IF EXPIRED @@ -14167,7 +14173,9 @@ async def claim_onboarding_link(data: InvitationClaim, request: Request): detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - invite_obj = await InvitationLinkRepository(prisma_client).table.find_unique(where={"id": data.invitation_link}) + invite_obj: PrismaInvitationLink | None = await InvitationLinkRepository(prisma_client).table.find_unique( + where={"id": data.invitation_link} + ) if invite_obj is None: raise HTTPException(status_code=401, detail={"error": "Invitation link does not exist in db."}) #### CHECK IF EXPIRED @@ -14502,7 +14510,9 @@ async def invitation_info(invitation_id: str, user_api_key_dict: UserAPIKeyAuth }, ) - response = await InvitationLinkRepository(prisma_client).table.find_unique(where={"id": invitation_id}) + response: PrismaInvitationLink | None = await InvitationLinkRepository(prisma_client).table.find_unique( + where={"id": invitation_id} + ) if response is None: raise HTTPException( @@ -14550,7 +14560,7 @@ async def invitation_update( ) current_time = litellm.utils.get_utc_datetime() - response = await InvitationLinkRepository(prisma_client).table.update( + response: PrismaInvitationLink | None = await InvitationLinkRepository(prisma_client).table.update( where={"id": data.invitation_id}, data={ "id": data.invitation_id, @@ -14621,7 +14631,9 @@ async def invitation_delete( # Org admins can only delete invitations they created if is_other_admin and not is_proxy_admin: - invitation = await InvitationLinkRepository(prisma_client).table.find_unique(where={"id": data.invitation_id}) + invitation: PrismaInvitationLink | None = await InvitationLinkRepository(prisma_client).table.find_unique( + where={"id": data.invitation_id} + ) if invitation is None: raise HTTPException( status_code=400, @@ -14633,7 +14645,9 @@ async def invitation_delete( detail={"error": "Organization admins can only delete invitations they created."}, ) - response = await InvitationLinkRepository(prisma_client).table.delete(where={"id": data.invitation_id}) + response: PrismaInvitationLink | None = await InvitationLinkRepository(prisma_client).table.delete( + where={"id": data.invitation_id} + ) if response is None: raise HTTPException( @@ -14974,7 +14988,7 @@ def _redact_secret_values_in_obj(value: JsonValue, depth: int = 0) -> JsonValue: return value -def _redact_config_param_value_for_logging(param_name: Optional[str], param_value: JsonValue) -> JsonValue: +def _redact_config_param_value_for_logging(param_name: str | None, param_value: JsonValue) -> JsonValue: if param_name == "environment_variables" and isinstance(param_value, dict): return {key: "REDACTED" for key in param_value} if isinstance(param_value, (dict, list)): @@ -14992,7 +15006,7 @@ def _redact_general_setting_value(field_name: str, value: JsonValue, is_full_adm return value -def _dump_redacted_config(value: Optional[JsonValue], *, redact_all_values: bool = False) -> Optional[str]: +def _dump_redacted_config(value: JsonValue | None, *, redact_all_values: bool = False) -> str | None: # `default=str` matches the sibling audit-log serializers in # team_endpoints.py and the LiteLLM_AuditLogs validator, so a YAML-loaded # value with a non-JSON-native leaf (datetime, custom object) cannot turn @@ -15007,8 +15021,8 @@ def _dump_redacted_config(value: Optional[JsonValue], *, redact_all_values: bool async def create_config_audit_log( param_name: str, action: AUDIT_ACTIONS, - before_value: Optional[JsonValue], - after_value: Optional[JsonValue], + before_value: JsonValue | None, + after_value: JsonValue | None, user_api_key_dict: UserAPIKeyAuth, table_name: LitellmTableNames = LitellmTableNames.CONFIG_TABLE_NAME, ) -> None: @@ -15044,7 +15058,7 @@ _EXTRA_SECRET_CALLBACK_ENV_VARS = frozenset( ) -def _redact_callback_env_vars(env_vars: dict[str, Optional[str]]) -> dict[str, Optional[str]]: +def _redact_callback_env_vars(env_vars: dict[str, str | None]) -> dict[str, str | None]: """Return a copy of ``env_vars`` with values for keys classified as sensitive by ``is_sensitive_callback_key`` replaced with ``"REDACTED"``. ``None`` values pass through unchanged. @@ -15277,7 +15291,7 @@ async def _reset_general_settings_ui_litellm_field(field_name: str, user_api_key async def get_config_list( config_type: Literal["general_settings"], user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -) -> List[ConfigList]: +) -> list[ConfigList]: """ List the available fields + current values for a given type of setting (currently just 'general_settings'user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),) """ @@ -15434,7 +15448,7 @@ async def get_config_list( for litellm_field_name, spec in _GENERAL_SETTINGS_UI_LITELLM_FIELDS.items(): current_value: GeneralSettingsUILiteLLMValue = getattr(litellm, litellm_field_name, None) default_value = _general_settings_ui_litellm_default(spec) - stored_in_db_litellm: Optional[bool] + stored_in_db_litellm: bool | None if litellm_field_name in db_litellm_settings: stored_in_db_litellm = True elif current_value != default_value: @@ -16675,7 +16689,7 @@ async def _mcp_forward_as_path(path_segment: str, request: Request): return await _stream_mcp_asgi_response(handle_streamable_http_mcp, scope, request.receive) -async def _resolve_mcp_csv_tokens(csv_segment: str, client_ip: Optional[str]) -> List[str]: +async def _resolve_mcp_csv_tokens(csv_segment: str, client_ip: str | None) -> list[str]: """Validate a comma-separated ``/{name1,name2,...}/mcp`` segment. For each token, check (in order) whether it is a registered MCP server @@ -16703,7 +16717,7 @@ async def _resolve_mcp_csv_tokens(csv_segment: str, client_ip: Optional[str]) -> ) seen: set = set() - deduped: List[str] = [] + deduped: list[str] = [] for raw in csv_segment.split(","): token = raw.strip() if not token or token in seen: @@ -16713,7 +16727,7 @@ async def _resolve_mcp_csv_tokens(csv_segment: str, client_ip: Optional[str]) -> if len(deduped) >= DEFAULT_MCP_NAMESPACE_CSV_MAX_TOKENS: break - resolved: List[str] = [] + resolved: list[str] = [] for token in deduped: if global_mcp_server_manager.get_mcp_server_by_name(token, client_ip=client_ip): resolved.append(token) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 2ca251a3211..6dcf2dfd177 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -10,6 +10,7 @@ import sys import threading import time import traceback +from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass, field from datetime import date, datetime, timedelta, timezone from email.mime.multipart import MIMEMultipart @@ -17,17 +18,12 @@ from email.mime.text import MIMEText from typing import ( TYPE_CHECKING, Any, - AsyncGenerator, - Awaitable, - Callable, ClassVar, - Dict, - List, Literal, - Mapping, Optional, - Sequence, - Tuple, + Protocol, + TypedDict, + TypeVar, Union, cast, overload, @@ -98,9 +94,6 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, ModifyResponseException, ) -from litellm.proxy.hooks.sensitive_data_routing import ( - _PROXY_SensitiveDataRoutingHandler, -) from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.prometheus import PrometheusLogger from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting @@ -147,6 +140,9 @@ from litellm.proxy.hooks.parallel_request_limiter import ( from litellm.proxy.hooks.parallel_request_limiter_v3 import ( _PROXY_MaxParallelRequestsHandler_v3, ) +from litellm.proxy.hooks.sensitive_data_routing import ( + _PROXY_SensitiveDataRoutingHandler, +) from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor from litellm.repositories.budget_repository import BudgetRepository @@ -164,6 +160,7 @@ from litellm.repositories.verification_token_repository import ( ) from litellm.secret_managers.main import str_to_bool from litellm.types.integrations.slack_alerting import DEFAULT_ALERT_TYPES +from litellm.types.llms.base import HiddenParams from litellm.types.mcp import ( MCPDuringCallResponseObject, MCPPreCallRequestObject, @@ -176,6 +173,7 @@ if TYPE_CHECKING: from mcp.types import CallToolResult from opentelemetry.trace import Span as _Span from prisma.client import TransactionManager + from prisma.types import HttpConfig from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy.db.spend_log_tool_index import ToolUsageTransaction @@ -189,6 +187,8 @@ unified_guardrail = UnifiedLLMGuardrails() NON_OPENAI_STREAM_GUARDRAIL_TRANSLATION_CALL_TYPES: "frozenset[CallTypes]" = frozenset({CallTypes.anthropic_messages}) +_GuardrailHookResultT = TypeVar("_GuardrailHookResultT") + def print_verbose(print_statement): """ @@ -239,10 +239,10 @@ class InternalUsageCache: async def async_get_cache( self, key, - litellm_parent_otel_span: Union[Span, None], + litellm_parent_otel_span: Span | None, local_only: bool = False, **kwargs, - ) -> Any: + ) -> Any: # any-ok: DualCache stores arbitrary value types; return type mirrors its own untyped signature return await self.dual_cache.async_get_cache( key=key, local_only=local_only, @@ -254,7 +254,7 @@ class InternalUsageCache: self, key, value, - litellm_parent_otel_span: Union[Span, None], + litellm_parent_otel_span: Span | None, local_only: bool = False, **kwargs, ) -> None: @@ -268,8 +268,8 @@ class InternalUsageCache: async def async_batch_set_cache( self, - cache_list: List, - litellm_parent_otel_span: Union[Span, None], + cache_list: list, + litellm_parent_otel_span: Span | None, local_only: bool = False, **kwargs, ) -> None: @@ -283,7 +283,7 @@ class InternalUsageCache: 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, ): return await self.dual_cache.async_batch_get_cache( @@ -296,7 +296,7 @@ class InternalUsageCache: self, key, value: float, - litellm_parent_otel_span: Union[Span, None], + litellm_parent_otel_span: Span | None, local_only: bool = False, **kwargs, ): @@ -327,7 +327,7 @@ class InternalUsageCache: key, local_only: bool = False, **kwargs, - ) -> Any: + ) -> Any: # any-ok: DualCache stores arbitrary value types; return type mirrors its own untyped signature return self.dual_cache.get_cache( key=key, local_only=local_only, @@ -338,7 +338,7 @@ class InternalUsageCache: ### LOGGING ### # Cache for inspect.signature checks — avoids repeated introspection per request -_CALLBACK_ACCEPTS_CALL_INFO: Dict[int, bool] = {} +_CALLBACK_ACCEPTS_CALL_INFO: dict[int, bool] = {} def _accepts_litellm_call_info(cb: CustomLogger) -> bool: @@ -349,7 +349,7 @@ def _accepts_litellm_call_info(cb: CustomLogger) -> bool: return _CALLBACK_ACCEPTS_CALL_INFO[key] -def _enrich_http_exception_with_guardrail_context(exc: BaseException, callback: Any) -> None: +def _enrich_http_exception_with_guardrail_context(exc: BaseException, callback: CustomLogger) -> None: """ If `exc` is an HTTPException with a dict `detail`, mutate it in place to add `guardrail_name` and `guardrail_mode` taken from the callback instance. @@ -371,6 +371,13 @@ def _enrich_http_exception_with_guardrail_context(exc: BaseException, callback: detail.setdefault("guardrail_mode", event_hook) +class _McpPreCallHookResult(TypedDict): + should_proceed: bool + modified_arguments: dict[str, object] + error_message: str | None + hidden_params: HiddenParams + + def _exception_changes_request_flow(exc: BaseException) -> bool: """ True for guardrail exceptions the proxy turns into an alternate request flow @@ -398,11 +405,11 @@ class _CallbackCapabilities: # Tuple[(resolved_callback, "override" | "apply_guardrail"), ...] # Ordered the same as ``litellm.callbacks``; used to build the streaming # iterator chain without re-scanning per request. - iterator_overrides: Tuple[Tuple[Any, str], ...] = field(default_factory=tuple) + iterator_overrides: tuple[tuple[CustomLogger, str], ...] = field(default_factory=tuple) # Resolved CustomLogger callbacks in original order. Pre-resolving once # avoids the per-request ``get_custom_logger_compatible_class`` walk for # every string entry in ``litellm.callbacks``. - resolved_callbacks: Tuple[Any, ...] = field(default_factory=tuple) + resolved_callbacks: tuple[CustomLogger, ...] = field(default_factory=tuple) class ProxyLogging: @@ -428,16 +435,16 @@ class ProxyLogging: self.max_parallel_request_limiter = _PROXY_MaxParallelRequestsHandler(self.internal_usage_cache) self.max_budget_limiter = _PROXY_MaxBudgetLimiter() self.cache_control_check = _PROXY_CacheControlCheck() - self.alerting: Optional[List] = None + self.alerting: list | None = None self.alerting_threshold: float = 300 # default to 5 min. threshold - self.alert_types: List[AlertType] = DEFAULT_ALERT_TYPES - self.alert_to_webhook_url: Optional[dict] = None + self.alert_types: list[AlertType] = DEFAULT_ALERT_TYPES + self.alert_to_webhook_url: dict | None = None self.slack_alerting_instance: SlackAlerting = SlackAlerting( alerting_threshold=self.alerting_threshold, alerting=self.alerting, internal_usage_cache=self.internal_usage_cache.dual_cache, ) - self.email_logging_instance: Optional[Any] = None + self.email_logging_instance: Any | None = None if BaseEmailLogger is not None: email_logger_class = _get_email_logger_class() if email_logger_class is not None: @@ -448,7 +455,7 @@ class ProxyLogging: self.premium_user = premium_user self.service_logging_obj = ServiceLogging() self.db_spend_update_writer = DBSpendUpdateWriter() - self.proxy_hook_mapping: Dict[str, CustomLogger] = {} + self.proxy_hook_mapping: dict[str, CustomLogger] = {} # Guard flags to prevent duplicate background tasks self.daily_report_started: bool = False @@ -456,8 +463,8 @@ class ProxyLogging: def startup_event( self, - llm_router: Optional[Router], - redis_usage_cache: Optional[RedisCache], + llm_router: Router | None, + redis_usage_cache: RedisCache | None, ): """Initialize logging and alerting on proxy startup""" ## UPDATE SLACK ALERTING ## @@ -494,13 +501,13 @@ class ProxyLogging: def update_values( self, - alerting: Optional[List] = None, - alerting_threshold: Optional[float] = None, - redis_cache: Optional[RedisCache] = None, - alert_types: Optional[List[AlertType]] = None, - alerting_args: Optional[dict] = None, - alert_to_webhook_url: Optional[dict] = None, - alert_type_config: Optional[dict] = None, + alerting: list | None = None, + alerting_threshold: float | None = None, + redis_cache: RedisCache | None = None, + alert_types: list[AlertType] | None = None, + alerting_args: dict | None = None, + alert_to_webhook_url: dict | None = None, + alert_type_config: dict | None = None, ): updated_slack_alerting: bool = False if alerting is not None: @@ -546,7 +553,7 @@ class ProxyLogging: self.db_spend_update_writer.redis_update_buffer.redis_cache = redis_cache self.db_spend_update_writer.pod_lock_manager.redis_cache = redis_cache - def _add_proxy_hooks(self, llm_router: Optional[Router] = None): + def _add_proxy_hooks(self, llm_router: Router | None = None): """ Add proxy hooks to litellm.callbacks """ @@ -555,7 +562,7 @@ class ProxyLogging: for hook in PROXY_HOOKS: proxy_hook = get_proxy_hook(hook) expected_args = inspect.getfullargspec(proxy_hook).args - passed_in_args: Dict[str, Any] = {} + passed_in_args: dict[str, Any] = {} if "internal_usage_cache" in expected_args: passed_in_args["internal_usage_cache"] = self.internal_usage_cache if "prisma_client" in expected_args: @@ -565,20 +572,20 @@ class ProxyLogging: self.proxy_hook_mapping[hook] = proxy_hook_obj - def get_proxy_hook(self, hook: str) -> Optional[CustomLogger]: + def get_proxy_hook(self, hook: str) -> CustomLogger | None: """ Get a proxy hook from the proxy_hook_mapping """ return self.proxy_hook_mapping.get(hook) - def _init_litellm_callbacks(self, llm_router: Optional[Router] = None): + def _init_litellm_callbacks(self, llm_router: Router | None = None): self._add_proxy_hooks(llm_router) litellm.logging_callback_manager.add_litellm_callback(self.service_logging_obj) # type: ignore # Track string callbacks and their initialized instances so we can # replace them in-place, preventing duplicates (string + instance) in # litellm.callbacks which caused double-counting of metrics. - string_callbacks_to_replace: Dict[int, CustomLogger] = {} + string_callbacks_to_replace: dict[int, CustomLogger] = {} for idx, callback in enumerate(litellm.callbacks): if isinstance(callback, str): @@ -680,7 +687,7 @@ class ProxyLogging: return synthetic_data - def _convert_llm_result_to_mcp_response(self, llm_result, request_obj) -> Optional[Any]: + def _convert_llm_result_to_mcp_response(self, llm_result, request_obj) -> MCPPreCallResponseObject | None: """ Convert LLM guardrail result back to MCP response format. """ @@ -740,7 +747,7 @@ class ProxyLogging: return None - def _extract_modified_arguments_from_content(self, masked_content: str, request_obj) -> Optional[dict]: + def _extract_modified_arguments_from_content(self, masked_content: str, request_obj) -> dict | None: """ Extract modified/masked arguments from the guardrail response content. """ @@ -777,7 +784,7 @@ class ProxyLogging: verbose_proxy_logger.error(f"Error extracting modified arguments: {e}") return None - def _parse_arguments_manually(self, args_text: str, original_args: dict) -> Optional[dict]: + def _parse_arguments_manually(self, args_text: str, original_args: dict) -> dict | None: """ Try to manually parse arguments when JSON parsing fails. This is a fallback for cases where the guardrail modifies the format. @@ -806,7 +813,7 @@ class ProxyLogging: verbose_proxy_logger.error(f"Error in manual argument parsing: {e}") return None - def _convert_llm_result_to_mcp_during_response(self, llm_result, request_obj) -> Optional[Any]: + def _convert_llm_result_to_mcp_during_response(self, llm_result, request_obj) -> MCPDuringCallResponseObject | None: """ Convert LLM guardrail result back to MCP during call response format. """ @@ -843,7 +850,7 @@ class ProxyLogging: return None - def get_combined_callback_list(self, dynamic_success_callbacks: Optional[List], global_callbacks: List) -> List: + def get_combined_callback_list(self, dynamic_success_callbacks: list | None, global_callbacks: list) -> list: if dynamic_success_callbacks is None: return list(global_callbacks) return list(dict.fromkeys(dynamic_success_callbacks + global_callbacks)) @@ -852,7 +859,7 @@ class ProxyLogging: self, response: MCPPreCallResponseObject, original_request: MCPPreCallRequestObject, - ) -> Dict[str, Any]: + ) -> _McpPreCallHookResult: """ Parse the response from the pre_mcp_tool_call_hook @@ -860,7 +867,7 @@ class ProxyLogging: 2. Apply any argument modifications 3. Handle validation errors """ - result = { + result: _McpPreCallHookResult = { "should_proceed": response.should_proceed, "modified_arguments": response.modified_arguments or original_request.arguments, "error_message": response.error_message, @@ -885,7 +892,7 @@ class ProxyLogging: hidden_params=HiddenParams(), ) - def _convert_mcp_hook_response_to_kwargs(self, response_data: Optional[dict], original_kwargs: dict) -> dict: + def _convert_mcp_hook_response_to_kwargs(self, response_data: dict | None, original_kwargs: dict) -> dict: """ Helper function to convert pre_call_hook response back to kwargs for MCP usage. @@ -953,10 +960,10 @@ class ProxyLogging: callback: "CustomGuardrail", hook_type: str, data: dict, - user_api_key_dict: Optional[UserAPIKeyAuth], + user_api_key_dict: UserAPIKeyAuth | None, call_type: CallTypesLiteral, - response: Optional[Any] = None, - ) -> Any: + response: Any | None = None, # any-ok: CustomLogger hook methods declare response/return as Any + ) -> Any: # any-ok: mirrors whichever dispatched CustomLogger hook, itself declared -> Any """ Execute a single guardrail's hook. @@ -1008,10 +1015,10 @@ class ProxyLogging: guardrail_name: str, hook_type: str, data: dict, - user_api_key_dict: Optional[UserAPIKeyAuth], + user_api_key_dict: UserAPIKeyAuth | None, call_type: CallTypesLiteral, - response: Optional[Any] = None, - ) -> Any: + response: Any | None = None, # any-ok: forwarded verbatim to _execute_guardrail_hook + ) -> Any: # any-ok: forwards _execute_guardrail_hook's return, itself declared -> Any """ Execute a guardrail using the router's load balancing. @@ -1034,7 +1041,7 @@ class ProxyLogging: # Select guardrail using router's load balancing selected_guardrail = llm_router.get_available_guardrail(guardrail_name=guardrail_name) - callback = selected_guardrail.get("callback") + callback: CustomGuardrail | None = selected_guardrail.get("callback") if callback is None: raise ValueError(f"No callback found for guardrail: {guardrail_name}") @@ -1051,10 +1058,10 @@ class ProxyLogging: self, callback: CustomGuardrail, data: dict, - user_api_key_dict: Optional[UserAPIKeyAuth], + user_api_key_dict: UserAPIKeyAuth | None, call_type: CallTypesLiteral, event_type: GuardrailEventHooks, - ) -> Optional[dict]: + ) -> dict | None: """ Process a guardrail callback during pre-call hook. @@ -1145,9 +1152,9 @@ class ProxyLogging: async def _process_prompt_template( self, data: dict, - litellm_logging_obj: Any, - prompt_id: Any, - prompt_version: Any, + litellm_logging_obj: Any, # any-ok: some kwargs in the call below don't match the real signature + prompt_id: Any, # any-ok: sourced from an untyped `data.get(...)` lookup + prompt_version: Any, # any-ok: sourced from an untyped `data.get(...)` lookup call_type: CallTypesLiteral, ) -> None: """Process prompt template if applicable.""" @@ -1169,7 +1176,7 @@ class ProxyLogging: custom_logger = IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id(lookup_prompt_id) prompt_spec = IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id(lookup_prompt_id) - litellm_prompt_id: Optional[str] = None + litellm_prompt_id: str | None = None if prompt_spec is not None: litellm_prompt_id = prompt_spec.litellm_params.prompt_id data.pop("prompt_id", None) @@ -1349,9 +1356,9 @@ class ProxyLogging: async def pre_call_hook( self, user_api_key_dict: UserAPIKeyAuth, - data: Optional[dict], + data: dict | None, call_type: CallTypesLiteral, - ) -> Optional[dict]: + ) -> dict | None: """ Allows users to modify/reject the incoming request to the proxy, without having to deal with parsing Request body. @@ -1418,7 +1425,7 @@ class ProxyLogging: and not (cb.guardrail_name and cb.guardrail_name in pipeline_managed) ) - deferred_route_exc: Optional[SensitiveDataRouteException] = None + deferred_route_exc: SensitiveDataRouteException | None = None for _callback in caps.resolved_callbacks: start_time = time.time() try: @@ -1442,9 +1449,7 @@ class ProxyLogging: data = result elif ( - _callback is not None - and isinstance(_callback, CustomLogger) - and "async_pre_call_hook" in vars(_callback.__class__) + "async_pre_call_hook" in vars(_callback.__class__) and _callback.__class__.async_pre_call_hook != CustomLogger.async_pre_call_hook ): if call_type == "call_mcp_tool" and user_api_key_dict is None: @@ -1550,9 +1555,9 @@ class ProxyLogging: async def _handle_sensitive_data_route_exception( self, exc: SensitiveDataRouteException, - data: Optional[dict], - user_api_key_dict: Optional[UserAPIKeyAuth], - ) -> Optional[dict]: + data: dict | None, + user_api_key_dict: UserAPIKeyAuth | None, + ) -> dict | None: """ Handle SensitiveDataRouteException by rerouting the current request to the target model and, when sticky_session_routing is enabled, persisting @@ -1603,7 +1608,7 @@ class ProxyLogging: guardrail_name: str, latency_seconds: float, status: str, - error_type: Optional[str], + error_type: str | None, hook_type: str, ) -> None: for prom_callback in litellm.callbacks: @@ -1618,17 +1623,19 @@ class ProxyLogging: break @staticmethod - async def _run_guardrail_with_metrics(callback: Any, coro: Awaitable[Any], hook_type: str) -> Any: + async def _run_guardrail_with_metrics( + callback: CustomGuardrail, coro: Awaitable[_GuardrailHookResultT], hook_type: str + ) -> _GuardrailHookResultT: """ Await `coro`, recording its latency and status to the `litellm_guardrail_latency_seconds` metric under `hook_type`, and enriching any raised HTTPException with the originating callback's `guardrail_name`/`guardrail_mode` before re-raising. """ - guardrail_name = getattr(callback, "guardrail_name", None) or type(callback).__name__ + guardrail_name = callback.guardrail_name or type(callback).__name__ start_time = time.perf_counter() status = "success" - error_type: Optional[str] = None + error_type: str | None = None try: return await coro except SensitiveDataRouteException: @@ -1650,8 +1657,8 @@ class ProxyLogging: @staticmethod async def _wrap_streaming_iterator_with_enrichment( - callback: Any, gen: AsyncGenerator[Any, None] - ) -> AsyncGenerator[Any, None]: + callback: CustomLogger, gen: AsyncGenerator[_GuardrailHookResultT, None] + ) -> AsyncGenerator[_GuardrailHookResultT, None]: """ Yield from `gen`; if iteration raises an HTTPException with dict detail, enrich the detail with the originating callback's `guardrail_name` and @@ -1670,7 +1677,7 @@ class ProxyLogging: # Cache for callback-capability detection. Keyed on a signature of # litellm.callbacks (length + each item's id) so we recompute when the # callback list mutates (add/remove) without iterating every request. - _callback_capabilities_cache: ClassVar[Dict[Tuple[int, Tuple[int, ...]], "_CallbackCapabilities"]] = {} + _callback_capabilities_cache: ClassVar[dict[tuple[int, tuple[int, ...]], "_CallbackCapabilities"]] = {} @staticmethod def _callback_capabilities() -> "_CallbackCapabilities": @@ -1695,12 +1702,13 @@ class ProxyLogging: has_streaming_chunk_override = False has_guardrail = False has_pre_call_override = False - iterator_overrides: List[Tuple[Any, str]] = [] # (callback, kind) - resolved_callbacks: List[Any] = [] + iterator_overrides: list[tuple[CustomLogger, str]] = [] # (callback, kind) + resolved_callbacks: list[CustomLogger] = [] for callback in callbacks: + resolved: Callable[..., object] | CustomLogger | None if isinstance(callback, str): - resolved: Any = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class( + resolved = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class( cast(_custom_logger_compatible_callbacks_literal, callback) ) else: @@ -1816,7 +1824,7 @@ class ProxyLogging: async def during_call_hook( self, data: dict, - user_api_key_dict: Optional[UserAPIKeyAuth], + user_api_key_dict: UserAPIKeyAuth | None, call_type: CallTypesLiteral, ): """ @@ -1958,7 +1966,7 @@ class ProxyLogging: message: str, level: Literal["Low", "Medium", "High"], alert_type: AlertType, - request_data: Optional[dict] = None, + request_data: dict | None = None, ): """ Alerting based on thresholds: - https://github.com/BerriAI/litellm/issues/1298 @@ -2062,10 +2070,10 @@ class ProxyLogging: request_data: dict, original_exception: Exception, user_api_key_dict: UserAPIKeyAuth, - error_type: Optional[ProxyErrorTypes] = None, - route: Optional[str] = None, - traceback_str: Optional[str] = None, - ) -> Optional[HTTPException]: + error_type: ProxyErrorTypes | None = None, + route: str | None = None, + traceback_str: str | None = None, + ) -> HTTPException | None: """ Allows users to raise custom exceptions/log when a call fails, without having to deal with parsing Request body. Callbacks can return or raise HTTPException to transform error responses sent to clients. @@ -2149,11 +2157,11 @@ class ProxyLogging: request_data.pop("litellm_logging_obj", None) # Track the first HTTPException returned or raised by any callback - transformed_exception: Optional[HTTPException] = None + transformed_exception: HTTPException | None = None for callback in litellm.callbacks: try: - _callback: Optional[CustomLogger] = None + _callback: CustomLogger | None = None if isinstance(callback, str): _callback = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class( cast(_custom_logger_compatible_callbacks_literal, callback) @@ -2188,8 +2196,8 @@ class ProxyLogging: def _is_proxy_only_llm_api_error( self, original_exception: Exception, - error_type: Optional[ProxyErrorTypes] = None, - route: Optional[str] = None, + error_type: ProxyErrorTypes | None = None, + route: str | None = None, ) -> bool: """ Return True if the error is a Proxy Only LLM API Error @@ -2222,15 +2230,15 @@ class ProxyLogging: self, request_data: dict, user_api_key_dict: UserAPIKeyAuth, - route: Optional[str] = None, - original_exception: Optional[Exception] = None, + route: str | None = None, + original_exception: Exception | None = None, ): """ Handle logging for proxy only errors by calling `litellm_logging_obj.async_failure_handler` Is triggered when self._is_proxy_only_error() returns True """ - litellm_logging_obj: Optional[Logging] = request_data.get("litellm_logging_obj", None) + litellm_logging_obj: Logging | None = request_data.get("litellm_logging_obj", None) if litellm_logging_obj is None: from litellm._uuid import uuid @@ -2268,8 +2276,8 @@ class ProxyLogging: litellm_params=_litellm_params, ) - input: Union[list, str, dict] = "" - normalized_call_type: Optional[str] = None + input: list | str | dict = "" + normalized_call_type: str | None = None if "messages" in request_data and isinstance(request_data["messages"], list): input = request_data["messages"] litellm_logging_obj.model_call_details["messages"] = input @@ -2338,11 +2346,11 @@ class ProxyLogging: from litellm.proxy.proxy_server import llm_router from litellm.types.guardrails import GuardrailEventHooks - guardrail_callbacks: List[CustomGuardrail] = [] - other_callbacks: List[CustomLogger] = [] + guardrail_callbacks: list[CustomGuardrail] = [] + other_callbacks: list[CustomLogger] = [] try: for callback in litellm.callbacks: - _callback: Optional[CustomLogger] = None + _callback: CustomLogger | None = None if isinstance(callback, str): _callback = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class( cast(_custom_logger_compatible_callbacks_literal, callback) @@ -2380,7 +2388,7 @@ class ProxyLogging: ): continue - guardrail_response: Optional[Any] = None + guardrail_response: Any | None = None # any-ok: CustomLogger.async_post_call_success_hook -> Any if "apply_guardrail" in type(callback).__dict__: data["guardrail_to_apply"] = callback @@ -2546,8 +2554,8 @@ class ProxyLogging: data: dict, user_api_key_dict: UserAPIKeyAuth, response: Any, - request_headers: Optional[Dict[str, str]] = None, - ) -> Dict[str, str]: + request_headers: dict[str, str] | None = None, + ) -> dict[str, str]: """ Calls async_post_call_response_headers_hook on all CustomLogger callbacks. Merges all returned header dicts (later callbacks override earlier ones). @@ -2555,7 +2563,7 @@ class ProxyLogging: Returns: Dict[str, str]: Merged headers from all callbacks. """ - merged_headers: Dict[str, str] = {} + merged_headers: dict[str, str] = {} # Outer call sites in common_request_processing.py already gate this # call with ``has_post_call_response_headers_callbacks()``. The # cached detection makes the redundant interior guard cheap, but the @@ -2569,7 +2577,7 @@ class ProxyLogging: litellm_call_info = self._build_litellm_call_info(data=data, response=response) for callback in litellm.callbacks: - _callback: Optional[CustomLogger] = None + _callback: CustomLogger | None = None if isinstance(callback, str): _callback = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class( cast(_custom_logger_compatible_callbacks_literal, callback) @@ -2601,7 +2609,7 @@ class ProxyLogging: return merged_headers @staticmethod - def _build_litellm_call_info(data: dict, response: Any) -> Dict[str, Any]: + def _build_litellm_call_info(data: dict, response: Any) -> dict[str, Any]: """ Build a normalized dict of routing metadata from response._hidden_params and data, abstracting away the metadata vs litellm_metadata split. @@ -2630,9 +2638,9 @@ class ProxyLogging: async def async_post_call_streaming_hook( self, data: dict, - response: Union[ModelResponse, EmbeddingResponse, ImageResponse, ModelResponseStream], + response: ModelResponse | EmbeddingResponse | ImageResponse | ModelResponseStream, user_api_key_dict: UserAPIKeyAuth, - str_so_far: Optional[str] = None, + str_so_far: str | None = None, ): """ Allow user to modify outgoing streaming data -> per chunk @@ -2652,7 +2660,7 @@ class ProxyLogging: from litellm.proxy.proxy_server import llm_router - response_str: Optional[str] = None + response_str: str | None = None if isinstance(response, (ModelResponse, ModelResponseStream)): response_str = litellm.get_response_string(response_obj=response) elif isinstance(response, dict) and self.is_a2a_streaming_response(response): @@ -2662,12 +2670,12 @@ class ProxyLogging: if response_str is not None: # Cache model-level guardrails check per-request to avoid repeated # dict lookups + llm_router.get_deployment() per callback per chunk. - _cached_guardrail_data: Optional[dict] = None + _cached_guardrail_data: dict | None = None _guardrail_data_computed = False for callback in litellm.callbacks: try: - _callback: Optional[CustomLogger] = None + _callback: CustomLogger | None = None if isinstance(callback, CustomGuardrail): # Main - V2 Guardrails implementation from litellm.types.guardrails import GuardrailEventHooks @@ -2769,7 +2777,10 @@ class ProxyLogging: ), ) else: - # kind == "apply_guardrail": route through unified_guardrail + # kind == "apply_guardrail": route through unified_guardrail. Only + # callbacks with an "apply_guardrail" method reach this branch, and + # that method is only ever defined on CustomGuardrail subclasses. + assert isinstance(resolved_callback, CustomGuardrail) current_response = self._wrap_streaming_iterator_with_enrichment( resolved_callback, unified_guardrail.async_post_call_streaming_iterator_hook( @@ -2835,7 +2846,7 @@ class ProxyLogging: return await limiter.async_release_max_parallel_requests_on_disconnect(user_api_key_dict) - def _init_response_taking_too_long_task(self, data: Optional[dict] = None): + def _init_response_taking_too_long_task(self, data: dict | None = None): """ Initialize the response taking too long task if user is using slack alerting @@ -2878,9 +2889,9 @@ _DEPRECATED_KEY_CACHE_TTL_SECONDS = 60 async def _lookup_deprecated_key( - db: Any, + db: PrismaWrapper | RoutingPrismaWrapper, hashed_token: str, -) -> Optional[str]: +) -> str | None: """ Check if a token exists in the deprecated keys table and is still within its grace period. @@ -2937,20 +2948,27 @@ class _ConfigRow: __slots__ = ("param_name", "param_value") - def __init__(self, param_name: str, param_value: Any) -> None: + def __init__(self, param_name: str, param_value: Any) -> None: # any-ok: config values are arbitrary JSON self.param_name = param_name self.param_value = param_value +class _ConfigRowLike(Protocol): + """Structural shape shared by Prisma litellm_config rows and ``_ConfigRow``.""" + + param_name: str + param_value: object + + def _config_cache_key(param_name: str) -> str: return f"litellm_config:param:{param_name}" -def _pack_config_row(row: Any) -> Dict[str, Any]: +def _pack_config_row(row: _ConfigRowLike) -> dict[str, object]: return {"param_name": row.param_name, "param_value": row.param_value} -def _unpack_config_row(cached: Any) -> Optional[_ConfigRow]: +def _unpack_config_row(cached: Any) -> _ConfigRow | None: # any-ok: DualCache value is untyped if cached is None or cached == _CONFIG_CACHE_MISS: return None if isinstance(cached, dict): @@ -2958,7 +2976,9 @@ def _unpack_config_row(cached: Any) -> Optional[_ConfigRow]: return None -async def get_config_param(prisma_client: Any, param_name: str) -> Optional[Any]: +async def get_config_param( + prisma_client: "PrismaClient", param_name: str +) -> Any | None: # any-ok: row is either a Prisma model, a _ConfigRow shim, or None """Cached read of a LiteLLM_Config row; returns row, _ConfigRow shim, or None.""" cache_key = _config_cache_key(param_name) cached = await litellm_config_cache.async_get_cache(cache_key) @@ -2966,7 +2986,7 @@ async def get_config_param(prisma_client: Any, param_name: str) -> Optional[Any] return _unpack_config_row(cached) row = await prisma_client.get_generic_data(key="param_name", value=param_name, table_name="config") - cache_value: Any = _pack_config_row(row) if row is not None else _CONFIG_CACHE_MISS + cache_value = _pack_config_row(row) if row is not None else _CONFIG_CACHE_MISS await litellm_config_cache.async_set_cache(cache_key, cache_value, ttl=LITELLM_CONFIG_CACHE_TTL_SECONDS) return row @@ -2976,7 +2996,7 @@ async def invalidate_config_param(param_name: str) -> None: await litellm_config_cache.async_delete_cache(_config_cache_key(param_name)) -async def prefetch_config_params(prisma_client: Any, param_names: List[str]) -> None: +async def prefetch_config_params(prisma_client: "PrismaClient", param_names: list[str]) -> None: """Batch-load LiteLLM_Config rows into the cache with one find_many.""" if not param_names: return @@ -2993,27 +3013,27 @@ async def prefetch_config_params(prisma_client: Any, param_names: List[str]) -> by_name = {row.param_name: row for row in rows} for name in param_names: row = by_name.get(name) - cache_value: Any = _pack_config_row(row) if row is not None else _CONFIG_CACHE_MISS + cache_value = _pack_config_row(row) if row is not None else _CONFIG_CACHE_MISS await litellm_config_cache.async_set_cache( _config_cache_key(name), cache_value, ttl=LITELLM_CONFIG_CACHE_TTL_SECONDS ) class PrismaClient: - spend_log_transactions: List = [] + spend_log_transactions: list = [] _spend_log_transactions_lock = asyncio.Lock() - tool_usage_transactions: List["ToolUsageTransaction"] = [] + tool_usage_transactions: list["ToolUsageTransaction"] = [] _tool_usage_transactions_lock = asyncio.Lock() def __init__( self, database_url: str, proxy_logging_obj: ProxyLogging, - http_client: Optional[Any] = None, + http_client: Optional["HttpConfig"] = None, ): ## init logging object self.proxy_logging_obj = proxy_logging_obj - self.iam_token_db_auth: Optional[bool] = str_to_bool(os.getenv("IAM_TOKEN_DB_AUTH")) + self.iam_token_db_auth: bool | None = str_to_bool(os.getenv("IAM_TOKEN_DB_AUTH")) verbose_proxy_logger.debug("Creating Prisma Client..") try: from prisma import Prisma # type: ignore @@ -3046,7 +3066,7 @@ class PrismaClient: # reader endpoint and writes stay on the writer. Falls back to the # writer-only wrapper when the env var is unset, preserving existing # single-DB deployments. - self.db: Union[PrismaWrapper, RoutingPrismaWrapper] + self.db: PrismaWrapper | RoutingPrismaWrapper if read_replica_url: try: # If IAM auth is enabled, the reader refreshes its own token on @@ -3074,7 +3094,8 @@ class PrismaClient: ) read_replica_url = reader_iam_endpoint.build_url(reader_token) os.environ["DATABASE_URL_READ_REPLICA"] = read_replica_url - reader_kwargs: Dict[str, Any] = {"datasource": {"url": read_replica_url}} + # any-ok: Prisma's generated DatasourceOverride TypedDict rejects a matching dict literal here + reader_kwargs: dict[str, Any] = {"datasource": {"url": read_replica_url}} if http_client is not None: reader_prisma = Prisma(http=http_client, **reader_kwargs) else: @@ -3110,7 +3131,7 @@ class PrismaClient: else: self.db = writer_wrapper # Client to connect to Prisma db self._db_reconnect_lock = asyncio.Lock() - self._db_health_watchdog_task: Optional[asyncio.Task] = None + self._db_health_watchdog_task: asyncio.Task | None = None self._db_last_reconnect_attempt_ts: float = 0.0 self._db_reconnect_cooldown_seconds: int = max(1, int(os.getenv("PRISMA_RECONNECT_COOLDOWN_SECONDS", "15"))) self._db_health_watchdog_interval_seconds: int = max( @@ -3139,7 +3160,7 @@ class PrismaClient: self._engine_pid: int = 0 self._watching_engine: bool = False self._engine_confirmed_dead: bool = False - self._engine_wait_thread: Optional[threading.Thread] = None + self._engine_wait_thread: threading.Thread | None = None verbose_proxy_logger.debug("Success - Created Prisma Client") @property @@ -3157,7 +3178,7 @@ class PrismaClient: """ return cast("TransactionManager", self.db.tx()) # cast-ok: wrappers delegate tx via __getattr__ (untyped) - def get_request_status(self, payload: Union[dict, SpendLogsPayload]) -> Literal["success", "failure"]: + def get_request_status(self, payload: dict | SpendLogsPayload) -> Literal["success", "failure"]: """ Determine if a request was successful or failed based on payload metadata. @@ -3169,9 +3190,9 @@ class PrismaClient: """ try: # Get metadata and convert to dict if it's a JSON string - payload_metadata: Union[Dict, SpendLogsMetadata, str] = payload.get("metadata", {}) + payload_metadata: dict | SpendLogsMetadata | str = payload.get("metadata", {}) if isinstance(payload_metadata, str): - payload_metadata_json: Union[Dict, SpendLogsMetadata] = cast(Dict, json.loads(payload_metadata)) + payload_metadata_json: dict | SpendLogsMetadata = cast(dict, json.loads(payload_metadata)) else: payload_metadata_json = payload_metadata @@ -3302,7 +3323,7 @@ class PrismaClient: async def get_generic_data( self, key: str, - value: Any, + value: object, table_name: Literal["users", "keys", "config", "spend"], ): """ @@ -3359,7 +3380,7 @@ class PrismaClient: raise e - async def _query_first_with_cached_plan_fallback(self, sql_query: str, *args) -> Optional[dict]: + async def _query_first_with_cached_plan_fallback(self, sql_query: str, *args) -> dict | None: """ Execute a query, recovering once from PostgreSQL's "cached plan must not change result type" error. @@ -3415,38 +3436,29 @@ class PrismaClient: @log_db_metrics async def get_data( self, - token: Optional[Union[str, list]] = None, - user_id: Optional[str] = None, - user_id_list: Optional[list] = None, - team_id: Optional[str] = None, - team_id_list: Optional[list] = None, - key_val: Optional[dict] = None, - table_name: Optional[ - Literal[ - "user", - "key", - "config", - "spend", - "enduser", - "budget", - "team", - "user_notification", - "combined_view", - ] - ] = None, + token: str | list | None = None, + user_id: str | None = None, + user_id_list: list | None = None, + team_id: str | None = None, + team_id_list: list | None = None, + key_val: dict | None = None, + table_name: Literal[ + "user", "key", "config", "spend", "enduser", "budget", "team", "user_notification", "combined_view" + ] + | None = None, query_type: Literal["find_unique", "find_all"] = "find_unique", - expires: Optional[datetime] = None, - reset_at: Optional[datetime] = None, - offset: Optional[int] = None, # pagination, what row number to start from - limit: Optional[int] = None, # pagination, number of rows to getch when find_all==True - parent_otel_span: Optional[Span] = None, - proxy_logging_obj: Optional[ProxyLogging] = None, - budget_id_list: Optional[List[str]] = None, + expires: datetime | None = None, + reset_at: datetime | None = None, + offset: int | None = None, # pagination, what row number to start from + limit: int | None = None, # pagination, number of rows to getch when find_all==True + parent_otel_span: Span | None = None, + proxy_logging_obj: ProxyLogging | None = None, + budget_id_list: list[str] | None = None, check_deprecated: bool = True, ): args_passed_in = locals() start_time = time.time() - hashed_token: Optional[str] = None + hashed_token: str | None = None try: response: Any = None if (token is not None and table_name is None) or (table_name is not None and table_name == "key"): @@ -3777,7 +3789,7 @@ class PrismaClient: if response["team_blocked"] is None: response["team_blocked"] = False - team_member: Optional[Member] = None + team_member: Member | None = None if response["team_members_with_roles"] is not None and response["user_id"] is not None: ## find the team member corresponding to user id """ @@ -3991,15 +4003,15 @@ class PrismaClient: ) async def update_data( self, - token: Optional[str] = None, + token: str | None = None, data: dict = {}, - data_list: Optional[List] = None, - user_id: Optional[str] = None, - team_id: Optional[str] = None, + data_list: list | None = None, + user_id: str | None = None, + team_id: str | None = None, query_type: Literal["update", "update_many"] = "update", - table_name: Optional[Literal["user", "key", "config", "spend", "team", "enduser", "budget"]] = None, - update_key_values: Optional[dict] = None, - update_key_values_custom_query: Optional[dict] = None, + table_name: Literal["user", "key", "config", "spend", "team", "enduser", "budget"] | None = None, + update_key_values: dict | None = None, + update_key_values_custom_query: dict | None = None, ): """ Update existing data @@ -4240,10 +4252,10 @@ class PrismaClient: ) async def delete_data( self, - tokens: Optional[List] = None, - team_id_list: Optional[List] = None, - table_name: Optional[Literal["user", "key", "config", "spend", "team"]] = None, - user_id: Optional[str] = None, + tokens: list | None = None, + team_id_list: list | None = None, + table_name: Literal["user", "key", "config", "spend", "team"] | None = None, + user_id: str | None = None, ): """ Allow user to delete a key(s) @@ -4252,7 +4264,7 @@ class PrismaClient: """ start_time = time.time() try: - if tokens is not None and isinstance(tokens, List): + if tokens is not None and isinstance(tokens, list): hashed_tokens = [] for token in tokens: if isinstance(token, str) and token.startswith("sk-"): @@ -4271,11 +4283,11 @@ class PrismaClient: ) verbose_proxy_logger.debug("deleted_tokens: %s", deleted_tokens) return {"deleted_keys": deleted_tokens} - elif table_name == "team" and team_id_list is not None and isinstance(team_id_list, List): + elif table_name == "team" and team_id_list is not None and isinstance(team_id_list, list): # admin only endpoint -> `/team/delete` await TeamRepository(self).table.delete_many(where={"team_id": {"in": team_id_list}}) return {"deleted_teams": team_id_list} - elif table_name == "key" and team_id_list is not None and isinstance(team_id_list, List): + elif table_name == "key" and team_id_list is not None and isinstance(team_id_list, list): # admin only endpoint -> `/team/delete` await VerificationTokenRepository(self).table.delete_many(where={"team_id": {"in": team_id_list}}) except Exception as e: @@ -4719,7 +4731,7 @@ class PrismaClient: self._cleanup_engine_watcher() asyncio.create_task(self._start_engine_watcher()) - async def _run_reconnect_cycle(self, timeout_seconds: Optional[float] = None) -> None: + async def _run_reconnect_cycle(self, timeout_seconds: float | None = None) -> None: """ Run a reconnect cycle with a single overall timeout budget. @@ -4829,7 +4841,7 @@ class PrismaClient: self, force: bool, reason: str, - timeout_seconds: Optional[float], + timeout_seconds: float | None, ) -> bool: now = time.time() if force is False and now - self._db_last_reconnect_attempt_ts < self._db_reconnect_cooldown_seconds: @@ -4877,8 +4889,8 @@ class PrismaClient: self, reason: str, force: bool = False, - timeout_seconds: Optional[float] = None, - lock_timeout_seconds: Optional[float] = None, + timeout_seconds: float | None = None, + lock_timeout_seconds: float | None = None, ) -> bool: """ Attempt to reconnect the Prisma client in a singleflight manner. @@ -5097,7 +5109,7 @@ class PrismaClient: ) # Health Check Database Methods - def _validate_response_time(self, response_time_ms: Optional[float]) -> Optional[float]: + def _validate_response_time(self, response_time_ms: float | None) -> float | None: """Validate and clean response time value""" if response_time_ms is None: return None @@ -5108,7 +5120,7 @@ class PrismaClient: verbose_proxy_logger.warning(f"Invalid response_time_ms value: {response_time_ms}") return None - def _clean_details(self, details: Optional[dict]) -> Optional[dict]: + def _clean_details(self, details: dict | None) -> dict | None: """Clean and validate details JSON""" if not isinstance(details, dict): return None @@ -5124,11 +5136,11 @@ class PrismaClient: status: str, healthy_count: int = 0, unhealthy_count: int = 0, - error_message: Optional[str] = None, - response_time_ms: Optional[float] = None, - details: Optional[dict] = None, - checked_by: Optional[str] = None, - model_id: Optional[str] = None, + error_message: str | None = None, + response_time_ms: float | None = None, + details: dict | None = None, + checked_by: str | None = None, + model_id: str | None = None, ): """Save health check result to database""" try: @@ -5161,10 +5173,10 @@ class PrismaClient: async def get_health_check_history( self, - model_name: Optional[str] = None, + model_name: str | None = None, limit: int = 100, offset: int = 0, - status_filter: Optional[str] = None, + status_filter: str | None = None, ): """ Get health check history with optional filtering @@ -5244,9 +5256,9 @@ def _create_smtp_connection(smtp_host: str, smtp_port: int) -> smtplib.SMTP: async def send_email( - receiver_email: Optional[str] = None, - subject: Optional[str] = None, - html: Optional[str] = None, + receiver_email: str | None = None, + subject: str | None = None, + html: str | None = None, ): """ smtp_host, @@ -5397,7 +5409,7 @@ class ProxyUpdateSpend: n_retry_times: int, prisma_client: PrismaClient, proxy_logging_obj: ProxyLogging, - end_user_list_transactions: Dict[str, float], + end_user_list_transactions: dict[str, float], ): for i in range(n_retry_times + 1): start_time = time.time() @@ -5435,9 +5447,9 @@ class ProxyUpdateSpend: async def update_spend_logs( n_retry_times: int, prisma_client: PrismaClient, - db_writer_client: Optional[AsyncHTTPHandler], + db_writer_client: AsyncHTTPHandler | None, proxy_logging_obj: ProxyLogging, - logs_to_process: Optional[List[Dict[str, Any]]] = None, + logs_to_process: list[dict[str, Any]] | None = None, ): BATCH_SIZE = 1000 # Preferred size of each batch to write to the database MAX_LOGS_PER_INTERVAL = 10000 # Maximum number of logs to flush in a single interval @@ -5530,7 +5542,7 @@ class ProxyUpdateSpend: async def update_spend( prisma_client: PrismaClient, - db_writer_client: Optional[AsyncHTTPHandler], + db_writer_client: AsyncHTTPHandler | None, proxy_logging_obj: ProxyLogging, ): """ @@ -5623,7 +5635,7 @@ async def update_daily_tag_spend( async def update_spend_logs_job( prisma_client: PrismaClient, - db_writer_client: Optional[AsyncHTTPHandler], + db_writer_client: AsyncHTTPHandler | None, proxy_logging_obj: ProxyLogging, ): """ @@ -5696,7 +5708,7 @@ async def update_spend_logs_job( async def _monitor_spend_logs_queue( prisma_client: PrismaClient, - db_writer_client: Optional[AsyncHTTPHandler], + db_writer_client: AsyncHTTPHandler | None, proxy_logging_obj: ProxyLogging, ): """ @@ -5853,7 +5865,7 @@ def _get_month_end_date(today: date) -> date: return date(today.year, today.month + 1, 1) - timedelta(days=1) -def _is_projected_spend_over_limit(current_spend: float, soft_budget_limit: Optional[float]): +def _is_projected_spend_over_limit(current_spend: float, soft_budget_limit: float | None): if soft_budget_limit is None: # If there's no limit, we can't exceed it. return False @@ -5880,7 +5892,7 @@ def _is_projected_spend_over_limit(current_spend: float, soft_budget_limit: Opti return False -def _get_projected_spend_over_limit(current_spend: float, soft_budget_limit: Optional[float]) -> Optional[tuple]: +def _get_projected_spend_over_limit(current_spend: float, soft_budget_limit: float | None) -> tuple | None: if soft_budget_limit is None: return None @@ -5932,7 +5944,7 @@ def _to_ns(dt): def _check_and_merge_model_level_guardrails( data: dict, - llm_router: Optional[Router], + llm_router: Router | None, trust_client_model_info: bool = True, ) -> dict: """ @@ -5966,7 +5978,7 @@ def _check_and_merge_model_level_guardrails( # Medium on #29654). team_id = metadata.get("user_api_key_team_id") or litellm_metadata.get("user_api_key_team_id") - model_level_guardrails: Optional[list] = None + model_level_guardrails: list | None = None if model_id is not None: deployment = llm_router.get_deployment(model_id=model_id) if deployment is None: @@ -6015,7 +6027,10 @@ def _check_and_merge_model_level_guardrails( return _merge_guardrails_with_existing(data, model_level_guardrails) -def _merge_guardrails_with_existing(data: dict, model_level_guardrails: Any) -> dict: +def _merge_guardrails_with_existing( + data: dict, + model_level_guardrails: Any, # any-ok: sourced from an untyped litellm_params.get("guardrails") lookup +) -> dict: """ Merge model-level guardrails with any existing guardrails in the request data. @@ -6063,7 +6078,7 @@ def get_error_message_str(e: Exception) -> str: return error_message -def _get_redoc_url() -> Optional[str]: +def _get_redoc_url() -> str | None: """ Get the Redoc URL from the environment variables. @@ -6080,7 +6095,7 @@ def _get_redoc_url() -> Optional[str]: return "/redoc" -def _get_docs_url() -> Optional[str]: +def _get_docs_url() -> str | None: """ Get the docs (Swagger UI) URL from the environment variables. @@ -6097,7 +6112,7 @@ def _get_docs_url() -> Optional[str]: return "/" -def _get_openapi_url() -> Optional[str]: +def _get_openapi_url() -> str | None: """ Get the OpenAPI JSON URL from the environment variables. @@ -6140,7 +6155,7 @@ def handle_exception_on_proxy(e: Exception) -> ProxyException: ) -def _premium_user_check(feature: Optional[str] = None): +def _premium_user_check(feature: str | None = None): """ Raises an HTTPException if the user is not a premium user """ @@ -6160,7 +6175,7 @@ def _premium_user_check(feature: Optional[str] = None): ) -def is_known_model(model: Optional[str], llm_router: Optional[Router]) -> bool: +def is_known_model(model: str | None, llm_router: Router | None) -> bool: """ Returns True if the model is in the llm_router model names """ @@ -6210,7 +6225,7 @@ def join_paths(base_path: str, route: str) -> str: return final_path -def get_custom_url(request_base_url: str, route: Optional[str] = None) -> str: +def get_custom_url(request_base_url: str, route: str | None = None) -> str: # Use environment variable value, otherwise use URL from request server_base_url = get_proxy_base_url() if server_base_url is not None: @@ -6230,7 +6245,7 @@ def get_custom_url(request_base_url: str, route: Optional[str] = None) -> str: return join_paths(base_url, server_root_path) -def get_proxy_base_url() -> Optional[str]: +def get_proxy_base_url() -> str | None: """ Get the proxy base url from the environment variables. """ @@ -6247,7 +6262,7 @@ def get_server_root_path() -> str: return os.getenv("SERVER_ROOT_PATH", "") -def normalize_route_for_root_path(route: str) -> Optional[str]: +def normalize_route_for_root_path(route: str) -> str | None: """Strip SERVER_ROOT_PATH prefix. Returns de-prefixed route, or None if route is not under root path.""" root_path = get_server_root_path() if root_path and root_path != "/": @@ -6287,7 +6302,7 @@ def is_valid_api_key(key: str) -> bool: return False -def construct_database_url_from_env_vars() -> Optional[str]: +def construct_database_url_from_env_vars() -> str | None: """ Construct a DATABASE_URL from individual environment variables. Returns: @@ -6328,15 +6343,15 @@ async def get_available_models_for_user( user_api_key_dict: "UserAPIKeyAuth", llm_router: Optional["Router"], general_settings: dict, - user_model: Optional[str], + user_model: str | None, prisma_client: Optional["PrismaClient"] = None, proxy_logging_obj: Optional["ProxyLogging"] = None, - team_id: Optional[str] = None, + team_id: str | None = None, include_model_access_groups: bool = False, only_model_access_groups: bool = False, return_wildcard_routes: bool = False, user_api_key_cache: Optional["UserApiKeyCache"] = None, -) -> List[str]: +) -> list[str]: """ Get the list of models available to a user based on their API key and team permissions. @@ -6380,7 +6395,7 @@ async def get_available_models_for_user( ) # Get team models - team_models: List[str] = user_api_key_dict.team_models + team_models: list[str] = user_api_key_dict.team_models # If specific team_id is provided, validate and get team models if team_id and prisma_client and proxy_logging_obj and user_api_key_cache: @@ -6425,7 +6440,7 @@ def create_model_info_response( model_id: str, provider: str, include_metadata: bool = False, - fallback_type: Optional[str] = None, + fallback_type: str | None = None, llm_router: Optional["Router"] = None, get_model_info: Callable[[str], ModelInfo] = litellm.get_model_info, ) -> ModelInfoResponse: @@ -6498,7 +6513,7 @@ def create_model_info_response( def validate_model_access( model_id: str, - available_models: List[str], + available_models: list[str], ) -> None: """ Validate that a model is accessible to the user. @@ -6531,7 +6546,7 @@ def validate_model_access( ) -_PRESERVED_NONE_FIELDS: List[tuple[str, str]] = [ +_PRESERVED_NONE_FIELDS: list[tuple[str, str]] = [ ("message", "content"), # null when tool_calls present (issue #6677) ("message", "role"), # always required by OpenAI spec ("delta", "content"), # null in streaming chunks @@ -6540,9 +6555,9 @@ _PRESERVED_NONE_FIELDS: List[tuple[str, str]] = [ def model_dump_with_preserved_fields( obj: Any, - preserve_fields: Optional[List[str]] = None, + preserve_fields: list[str] | None = None, exclude_unset: bool = True, -) -> Dict[str, Any]: +) -> dict[str, Any]: """ Serialize a Pydantic model to a dictionary while preserving specific fields even if they are None. diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index 5c3e0cf0902..c731227c752 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -3,9 +3,6 @@ import logging from typing import ( Any, - List, - Optional, - Union, cast, ) @@ -18,10 +15,10 @@ from litellm.utils import CustomStreamWrapper def _add_mcp_metadata_to_response( - response: Union[ModelResponse, CustomStreamWrapper], - openai_tools: Optional[List], - tool_calls: Optional[List] = None, - tool_results: Optional[List] = None, + response: ModelResponse | CustomStreamWrapper, + openai_tools: list | None, + tool_calls: list | None = None, + tool_results: list | None = None, ) -> None: """ Add MCP metadata to response's provider_specific_fields. @@ -80,10 +77,10 @@ def _add_mcp_metadata_to_response( async def acompletion_with_mcp( model: str, - messages: List, - tools: Optional[List] = None, + messages: list, + tools: list | None = None, **kwargs: Any, -) -> Union[ModelResponse, CustomStreamWrapper]: +) -> ModelResponse | CustomStreamWrapper: """ Async completion with MCP integration. @@ -229,10 +226,10 @@ async def acompletion_with_mcp( self.openai_tools = openai_tools self.base_call_args = base_call_args self.request_tags = request_tags - self.collected_chunks: List[ModelResponseStream] = [] - self.tool_calls: Optional[List] = None - self.tool_results: Optional[List] = None - self.complete_response: Optional[ModelResponse] = None + self.collected_chunks: list[ModelResponseStream] = [] + self.tool_calls: list | None = None + self.tool_results: list | None = None + self.complete_response: ModelResponse | None = None self.stream_exhausted = False self.tool_execution_done = False self.follow_up_stream = None @@ -503,12 +500,12 @@ async def acompletion_with_mcp( # Create a wrapper class that delegates to our custom iterator # We'll use a simple approach: just replace the __aiter__ method class MCPStreamWrapper(CustomStreamWrapper): - def __init__(self, original_wrapper, custom_iterator): + def __init__(self, original_wrapper: CustomStreamWrapper, custom_iterator): # Initialize with the same parameters as original wrapper super().__init__( completion_stream=None, model=getattr(original_wrapper, "model", "unknown"), - logging_obj=getattr(original_wrapper, "logging_obj", None), + logging_obj=original_wrapper.logging_obj, custom_llm_provider=getattr(original_wrapper, "custom_llm_provider", None), stream_options=getattr(original_wrapper, "stream_options", None), make_call=getattr(original_wrapper, "make_call", None), diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index dab666ff0d9..799a60598e9 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -5,10 +5,11 @@ import json import time import traceback import uuid +from collections.abc import Mapping from datetime import datetime from functools import lru_cache from types import MappingProxyType -from typing import Any, Dict, List, Literal, Mapping, Optional +from typing import Any, Literal import httpx from openai._streaming import SSEDecoder @@ -29,7 +30,11 @@ from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( from litellm.litellm_core_utils.thread_pool_executor import executor from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.responses.utils import ResponseAPILoggingUtils, ResponsesAPIRequestUtils -from litellm.types.llms.openai import ResponsesAPIStreamEvents +from litellm.types.llms.openai import ( + PART_UNION_TYPES, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, +) from litellm.types.utils import CallTypes from litellm.utils import async_post_call_success_deployment_hook @@ -41,7 +46,7 @@ def _get_openai_response_types(): return openai_types -def _log_background_task_failure(task: "asyncio.Task[Any]", *, task_name: str) -> None: +def _log_background_task_failure(task: asyncio.Task[None], *, task_name: str) -> None: if task.cancelled(): return exception = task.exception() @@ -78,7 +83,7 @@ _ERROR_CODE_HTTP_STATUS: Mapping[str, int] = MappingProxyType( ) -def _error_event_fields(error_obj: object) -> tuple[str, Optional[str], Optional[str]]: +def _error_event_fields(error_obj: object) -> tuple[str, str | None, str | None]: if isinstance(error_obj, dict): raw_message = error_obj.get("message") raw_type = error_obj.get("type") @@ -97,7 +102,7 @@ def _error_event_fields(error_obj: object) -> tuple[str, Optional[str], Optional return message, error_type, code -def _status_code_for_error_fields(error_type: Optional[str], error_code: Optional[str]) -> int: +def _status_code_for_error_fields(error_type: str | None, error_code: str | None) -> int: fields = tuple(field for field in (error_code, error_type) if field is not None) if any(field.startswith("rate_limit") or field == "insufficient_quota" for field in fields): return 429 @@ -118,34 +123,34 @@ class BaseResponsesAPIStreamingIterator: self, response: httpx.Response, model: str, - responses_api_provider_config: Optional[BaseResponsesAPIConfig], + responses_api_provider_config: BaseResponsesAPIConfig | None, logging_obj: LiteLLMLoggingObj, - litellm_metadata: Optional[Dict[str, Any]] = None, - custom_llm_provider: Optional[str] = None, - request_data: Optional[Dict[str, Any]] = None, - call_type: Optional[str] = None, + litellm_metadata: dict[str, Any] | None = None, + custom_llm_provider: str | None = None, + request_data: dict[str, Any] | None = None, + call_type: str | None = None, ): self.response = response self.model = model self.logging_obj = logging_obj self.finished = False self.responses_api_provider_config = responses_api_provider_config - self.completed_response: Optional[Any] = None + self.completed_response: Any | None = None self.start_time = getattr(logging_obj, "start_time", datetime.now()) self._failure_handled = False # Track if failure handler has been called self._yielded_first_chunk = False self._generated_content = "" self._completed_response_cached = False self._completed_response_logged = False - self._completed_response_cache_hit: Optional[bool] = None + self._completed_response_cache_hit: bool | None = None self._persist_completed_response_before_logging = True self._stream_created_time: float = time.time() # track request context for hooks self.litellm_metadata = litellm_metadata self.custom_llm_provider = custom_llm_provider - self.request_data: Dict[str, Any] = request_data or {} - self.call_type: Optional[str] = call_type + self.request_data: dict[str, Any] = request_data or {} + self.call_type: str | None = call_type # set hidden params for response headers (e.g., x-litellm-model-id) # This matches the stream wrapper in litellm/litellm_core_utils/streaming_handler.py @@ -153,7 +158,7 @@ class BaseResponsesAPIStreamingIterator: model=model or "", optional_params=self.logging_obj.model_call_details.get("litellm_params", {}), ) - _model_info: Dict = litellm_metadata.get("model_info", {}) if litellm_metadata else {} + _model_info: dict = litellm_metadata.get("model_info", {}) if litellm_metadata else {} self._hidden_params = { "model_id": _model_info.get("id", None), "api_base": _api_base, @@ -175,7 +180,7 @@ class BaseResponsesAPIStreamingIterator: llm_provider=self.custom_llm_provider or "", ) - def _process_chunk(self, chunk) -> Optional[Any]: + def _process_chunk(self, chunk: str) -> Any | None: """Process a single chunk of data from the stream""" if not chunk: return None @@ -298,14 +303,12 @@ class BaseResponsesAPIStreamingIterator: self.completed_response = openai_responses_api_chunk # Add cost to usage object if include_cost_in_streaming_usage is True if litellm.include_cost_in_streaming_usage and self.logging_obj is not None: - response_obj: Optional[Any] = getattr(openai_responses_api_chunk, "response", None) + response_obj: Any | None = getattr(openai_responses_api_chunk, "response", None) if response_obj: - usage_obj: Optional[Any] = getattr(response_obj, "usage", None) + usage_obj: Any | None = getattr(response_obj, "usage", None) if usage_obj is not None: try: - cost: Optional[float] = self.logging_obj._response_cost_calculator( - result=response_obj - ) + cost: float | None = self.logging_obj._response_cost_calculator(result=response_obj) if cost is not None: setattr(usage_obj, "cost", cost) except Exception: @@ -403,10 +406,10 @@ class BaseResponsesAPIStreamingIterator: ) self._handle_failure(exception) - def _record_failed_response_usage(self, response_obj: Optional[Any]) -> None: + def _record_failed_response_usage(self, response_obj: ResponsesAPIResponse | None) -> None: if response_obj is None or self.logging_obj is None: return - usage_obj = getattr(response_obj, "usage", None) + usage_obj = response_obj.usage if usage_obj is None: return try: @@ -453,7 +456,7 @@ class BaseResponsesAPIStreamingIterator: is_pre_first_chunk=not self._yielded_first_chunk, ) - def _get_completed_response_object(self) -> Optional[Any]: + def _get_completed_response_object(self) -> ResponsesAPIResponse | None: openai_types = _get_openai_response_types() completed_response = self.completed_response if isinstance(completed_response, openai_types.ResponsesAPIResponse): @@ -535,7 +538,7 @@ class BaseResponsesAPIStreamingIterator: """ try: # Align with chat pipeline: use logging_obj model_call_details + call_type - typed_call_type: Optional[CallTypes] = None + typed_call_type: CallTypes | None = None if self.call_type is not None: try: typed_call_type = CallTypes(self.call_type) @@ -579,7 +582,7 @@ class BaseResponsesAPIStreamingIterator: if self.completed_response is None: return - request_payload: Dict[str, Any] = {} + request_payload: dict[str, Any] = {} if isinstance(self.request_data, dict): request_payload.update(self.request_data) try: @@ -609,7 +612,7 @@ class BaseResponsesAPIStreamingIterator: pass try: - typed_call_type: Optional[CallTypes] = None + typed_call_type: CallTypes | None = None if self.call_type is not None: try: typed_call_type = CallTypes(self.call_type) @@ -689,10 +692,10 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): model: str, responses_api_provider_config: BaseResponsesAPIConfig, logging_obj: LiteLLMLoggingObj, - litellm_metadata: Optional[Dict[str, Any]] = None, - custom_llm_provider: Optional[str] = None, - request_data: Optional[Dict[str, Any]] = None, - call_type: Optional[str] = None, + litellm_metadata: dict[str, Any] | None = None, + custom_llm_provider: str | None = None, + request_data: dict[str, Any] | None = None, + call_type: str | None = None, ): super().__init__( response, @@ -771,10 +774,10 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): model: str, responses_api_provider_config: BaseResponsesAPIConfig, logging_obj: LiteLLMLoggingObj, - litellm_metadata: Optional[Dict[str, Any]] = None, - custom_llm_provider: Optional[str] = None, - request_data: Optional[Dict[str, Any]] = None, - call_type: Optional[str] = None, + litellm_metadata: dict[str, Any] | None = None, + custom_llm_provider: str | None = None, + request_data: dict[str, Any] | None = None, + call_type: str | None = None, ): super().__init__( response, @@ -858,10 +861,10 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): model: str, responses_api_provider_config: BaseResponsesAPIConfig, logging_obj: LiteLLMLoggingObj, - litellm_metadata: Optional[Dict[str, Any]] = None, - custom_llm_provider: Optional[str] = None, - request_data: Optional[Dict[str, Any]] = None, - call_type: Optional[str] = None, + litellm_metadata: dict[str, Any] | None = None, + custom_llm_provider: str | None = None, + request_data: dict[str, Any] | None = None, + call_type: str | None = None, ): transformed = responses_api_provider_config.transform_response_api_response( model=model, @@ -882,7 +885,7 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): def _set_events_from_response( self, - transformed: Any, + transformed: ResponsesAPIResponse, logging_obj: LiteLLMLoggingObj, ) -> None: self._events = _build_synthetic_response_events( @@ -925,10 +928,10 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): def __init__( self, - response: Any, + response: ResponsesAPIResponse, logging_obj: LiteLLMLoggingObj, - request_data: Optional[Dict[str, Any]] = None, - call_type: Optional[str] = None, + request_data: dict[str, Any] | None = None, + call_type: str | None = None, ): BaseResponsesAPIStreamingIterator.__init__( self, @@ -943,13 +946,13 @@ class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): ) self._completed_response_cache_hit = True self._persist_completed_response_before_logging = False - self._events: List[Any] = [] + self._events: list[Any] = [] self._idx = 0 self._set_events_from_response(transformed=response, logging_obj=logging_obj) def _set_events_from_response( self, - transformed: Any, + transformed: ResponsesAPIResponse, logging_obj: LiteLLMLoggingObj, ) -> None: self._events = _build_synthetic_response_events( @@ -989,7 +992,7 @@ class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): return evt -def _dump_response_object(obj: Any) -> Dict[str, Any]: +def _dump_response_object(obj: Any) -> dict[str, Any]: if hasattr(obj, "model_dump"): return obj.model_dump() if isinstance(obj, dict): @@ -1002,7 +1005,7 @@ def _build_response_status_event( "response.created", "response.in_progress", ], - transformed: Any, + transformed: ResponsesAPIResponse, ) -> Any: openai_types = _get_openai_response_types() in_progress_response = transformed.model_copy( @@ -1019,11 +1022,11 @@ def _build_content_part_done_event( item_id: str, output_index: int, content_index: int, - part_payload: Dict[str, Any], -) -> Optional[Any]: + part_payload: dict[str, Any], +) -> Any | None: openai_types = _get_openai_response_types() part_type = part_payload.get("type") - part: Any + part: PART_UNION_TYPES if part_type == "output_text": annotations = [ openai_types.BaseLiteLLMOpenAIResponseObject(**annotation) @@ -1059,11 +1062,11 @@ def _build_content_part_done_event( def _add_text_like_part_events( *, - events: List[Any], + events: list[Any], item_id: str, output_index: int, content_index: int, - part_payload: Dict[str, Any], + part_payload: dict[str, Any], chunk_size: int, ) -> None: openai_types = _get_openai_response_types() @@ -1125,22 +1128,22 @@ def _add_text_like_part_events( def _build_synthetic_response_events( *, - transformed: Any, + transformed: ResponsesAPIResponse, logging_obj: LiteLLMLoggingObj, chunk_size: int, -) -> List[Any]: +) -> list[Any]: openai_types = _get_openai_response_types() if litellm.include_cost_in_streaming_usage and logging_obj is not None: - usage_obj: Optional[Any] = getattr(transformed, "usage", None) + usage_obj = transformed.usage if usage_obj is not None: try: - cost: Optional[float] = logging_obj._response_cost_calculator(result=transformed) + cost: float | None = logging_obj._response_cost_calculator(result=transformed) if cost is not None: - setattr(usage_obj, "cost", cost) + usage_obj.cost = cost except Exception: pass - events: List[Any] = [ + events: list[Any] = [ _build_response_status_event(openai_types.ResponsesAPIStreamEvents.RESPONSE_CREATED, transformed), _build_response_status_event(openai_types.ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS, transformed), ] @@ -1297,26 +1300,26 @@ class ResponsesWebSocketStreaming: websocket: Any, backend_ws: Any, logging_obj: LiteLLMLoggingObj, - user_api_key_dict: Optional[Any] = None, - request_data: Optional[Dict] = None, - first_message: Optional[str] = None, - guardrail_callbacks: Optional[List[Any]] = None, - output_guardrail_callbacks: Optional[List[Any]] = None, - authorized_model: Optional[str] = None, + user_api_key_dict: Any | None = None, + request_data: dict | None = None, + first_message: str | None = None, + guardrail_callbacks: list[Any] | None = None, + output_guardrail_callbacks: list[Any] | None = None, + authorized_model: str | None = None, ): self.websocket = websocket self.backend_ws = backend_ws self.logging_obj = logging_obj self.user_api_key_dict = user_api_key_dict - self.request_data: Dict = request_data or {} - self.messages: list[Dict] = [] - self.input_messages: list[Dict[str, str]] = [] + self.request_data: dict = request_data or {} + self.messages: list[dict] = [] + self.input_messages: list[dict[str, str]] = [] self.first_message = first_message - self.guardrail_callbacks: List[Any] = guardrail_callbacks or [] - self.output_guardrail_callbacks: List[Any] = output_guardrail_callbacks or [] + self.guardrail_callbacks: list[Any] = guardrail_callbacks or [] + self.output_guardrail_callbacks: list[Any] = output_guardrail_callbacks or [] # Model name authorized at connection time; enforced on every # response.create frame to prevent deployment-substitution attacks. - self.authorized_model: Optional[str] = authorized_model + self.authorized_model: str | None = authorized_model def _should_store_event(self, event_obj: dict) -> bool: return event_obj.get("type") in RESPONSES_WS_LOGGED_EVENT_TYPES @@ -1592,7 +1595,7 @@ class ResponsesWebSocketStreaming: if not self.guardrail_callbacks: return response_str - pii_tokens: Dict[str, str] = (self.request_data.get("metadata") or {}).get("pii_tokens", {}) + pii_tokens: dict[str, str] = (self.request_data.get("metadata") or {}).get("pii_tokens", {}) if not pii_tokens: return response_str @@ -1797,22 +1800,22 @@ class ManagedResponsesWebSocketHandler: self, websocket: Any, model: str, - logging_obj: "LiteLLMLoggingObj", - user_api_key_dict: Optional[Any] = None, - litellm_metadata: Optional[Dict[str, Any]] = None, - api_key: Optional[str] = None, - api_base: Optional[str] = None, - timeout: Optional[float] = None, - custom_llm_provider: Optional[str] = None, - first_message: Optional[str] = None, + logging_obj: LiteLLMLoggingObj, + user_api_key_dict: Any | None = None, + litellm_metadata: dict[str, Any] | None = None, + api_key: str | None = None, + api_base: str | None = None, + timeout: float | None = None, + custom_llm_provider: str | None = None, + first_message: str | None = None, **kwargs: Any, ) -> None: self.websocket = websocket self.model = model self.logging_obj = logging_obj self.user_api_key_dict = user_api_key_dict - self.litellm_metadata: Dict[str, Any] = litellm_metadata or {} - self.model_group: Optional[str] = self.litellm_metadata.get("model_group") or self.litellm_metadata.get( + self.litellm_metadata: dict[str, Any] = litellm_metadata or {} + self.model_group: str | None = self.litellm_metadata.get("model_group") or self.litellm_metadata.get( "deployment_model_name" ) self.api_key = api_key @@ -1822,19 +1825,19 @@ class ManagedResponsesWebSocketHandler: self._connection_provider = self._resolve_provider(model) or custom_llm_provider self.first_message = first_message # Carry through safe pass-through kwargs (e.g. extra_headers) - self.extra_kwargs: Dict[str, Any] = {k: v for k, v in kwargs.items() if k not in _MANAGED_WS_SKIP_KWARGS} + self.extra_kwargs: dict[str, Any] = {k: v for k, v in kwargs.items() if k not in _MANAGED_WS_SKIP_KWARGS} # In-memory session history: response_id → full accumulated message list. # Keyed by the DECODED (pre-encoding) response ID from response.completed. # This avoids the async DB-write race condition where spend logs haven't # been committed yet when the next response.create arrives. - self._session_history: Dict[str, List[Dict[str, Any]]] = {} + self._session_history: dict[str, list[dict[str, Any]]] = {} # ------------------------------------------------------------------ # Internal helpers # ------------------------------------------------------------------ @staticmethod - def _serialize_chunk(chunk: Any) -> Optional[str]: + def _serialize_chunk(chunk: Any) -> str | None: """Serialize a streaming chunk to a JSON string for WebSocket transmission.""" try: if hasattr(chunk, "model_dump_json"): @@ -1856,7 +1859,7 @@ class ManagedResponsesWebSocketHandler: except Exception: pass - def _get_history_messages(self, previous_response_id: str) -> List[Dict[str, Any]]: + def _get_history_messages(self, previous_response_id: str) -> list[dict[str, Any]]: """ Return accumulated message history for *previous_response_id*. @@ -1867,7 +1870,7 @@ class ManagedResponsesWebSocketHandler: raw_id = decoded.get("response_id", previous_response_id) return list(self._session_history.get(raw_id, [])) - def _store_history(self, response_id: str, messages: List[Dict[str, Any]]) -> None: + def _store_history(self, response_id: str, messages: list[dict[str, Any]]) -> None: """ Store the complete accumulated message history for *response_id*. @@ -1877,13 +1880,13 @@ class ManagedResponsesWebSocketHandler: self._session_history[response_id] = messages @staticmethod - def _extract_response_id(completed_event: Dict[str, Any]) -> Optional[str]: + def _extract_response_id(completed_event: dict[str, Any]) -> str | None: """ Pull the raw (decoded) response ID out of a ``response.completed`` event. Returns *None* if the event doesn't contain a usable ID. """ resp_obj = completed_event.get("response", {}) - encoded_id: Optional[str] = resp_obj.get("id") if isinstance(resp_obj, dict) else None + encoded_id: str | None = resp_obj.get("id") if isinstance(resp_obj, dict) else None if not encoded_id: return None decoded = ResponsesAPIRequestUtils._decode_responses_api_response_id(encoded_id) @@ -1891,8 +1894,8 @@ class ManagedResponsesWebSocketHandler: @staticmethod def _extract_output_messages( - completed_event: Dict[str, Any], - ) -> List[Dict[str, Any]]: + completed_event: dict[str, Any], + ) -> list[dict[str, Any]]: """ Convert the output items in a ``response.completed`` event into Responses API message dicts suitable for the next turn's ``input``. @@ -1900,7 +1903,7 @@ class ManagedResponsesWebSocketHandler: resp_obj = completed_event.get("response", {}) if not isinstance(resp_obj, dict): return [] - messages: List[Dict[str, Any]] = [] + messages: list[dict[str, Any]] = [] for item in resp_obj.get("output", []) or []: if not isinstance(item, dict): continue @@ -1927,7 +1930,7 @@ class ManagedResponsesWebSocketHandler: return messages @staticmethod - def _input_to_messages(input_val: Any) -> List[Dict[str, Any]]: + def _input_to_messages(input_val: Any) -> list[dict[str, Any]]: """ Normalise the ``input`` field of a ``response.create`` event to a list of Responses API message dicts. @@ -1948,7 +1951,7 @@ class ManagedResponsesWebSocketHandler: # _process_response_create sub-methods # ------------------------------------------------------------------ - async def _parse_message(self, raw_message: str) -> Optional[Dict[str, Any]]: + async def _parse_message(self, raw_message: str) -> dict[str, Any] | None: """Parse raw WS text; return the message dict or None (JSON error / ignored type).""" try: msg_obj = json.loads(raw_message) @@ -1961,14 +1964,14 @@ class ManagedResponsesWebSocketHandler: return msg_obj @staticmethod - def _is_warmup_frame(msg_obj: Dict[str, Any]) -> bool: + def _is_warmup_frame(msg_obj: dict[str, Any]) -> bool: """Return True for a response.create whose generate flag is false.""" nested = msg_obj.get("response") source = nested if isinstance(nested, dict) and nested else msg_obj return source.get("generate") is False @staticmethod - def _is_warmup_response_id(response_id: Optional[str]) -> bool: + def _is_warmup_response_id(response_id: str | None) -> bool: """Return True for synthetic warmup IDs that only exist on this connection.""" if not response_id: return False @@ -1977,13 +1980,13 @@ class ManagedResponsesWebSocketHandler: return str(raw_id).startswith(_WARMUP_RESPONSE_ID_PREFIX) @staticmethod - def _warmup_source_params(msg_obj: Dict[str, Any]) -> Dict[str, Any]: + def _warmup_source_params(msg_obj: dict[str, Any]) -> dict[str, Any]: nested = msg_obj.get("response") if isinstance(nested, dict) and nested: return nested return {k: v for k, v in msg_obj.items() if k != "type"} - def _build_warmup_response(self, msg_obj: Dict[str, Any]) -> Dict[str, Any]: + def _build_warmup_response(self, msg_obj: dict[str, Any]) -> dict[str, Any]: """Build a minimal completed Responses API object for a warmup ack.""" source = self._warmup_source_params(msg_obj) wire_model = source.get("model") or self.model_group or self.model @@ -2001,7 +2004,7 @@ class ManagedResponsesWebSocketHandler: }, } - async def _send_warmup_ack(self, msg_obj: Dict[str, Any]) -> None: + async def _send_warmup_ack(self, msg_obj: dict[str, Any]) -> None: """ Acknowledge a generate=false prewarm without calling the provider. @@ -2024,14 +2027,14 @@ class ManagedResponsesWebSocketHandler: await self.websocket.send_text(serialized) @staticmethod - def _build_base_call_kwargs(msg_obj: Dict[str, Any]) -> Dict[str, Any]: + def _build_base_call_kwargs(msg_obj: dict[str, Any]) -> dict[str, Any]: """ Extract Responses API params from the event, handling both wire formats: Nested: {"type": "response.create", "response": {"input": [...], ...}} Flat: {"type": "response.create", "input": [...], "model": "...", ...} """ nested = msg_obj.get("response") - response_params: Dict[str, Any] = ( + response_params: dict[str, Any] = ( nested if isinstance(nested, dict) and nested else {k: v for k, v in msg_obj.items() if k != "type"} ) return { @@ -2042,10 +2045,10 @@ class ManagedResponsesWebSocketHandler: def _apply_history( self, - call_kwargs: Dict[str, Any], - previous_response_id: Optional[str], - current_messages: List[Dict[str, Any]], - prior_history: List[Dict[str, Any]], + call_kwargs: dict[str, Any], + previous_response_id: str | None, + current_messages: list[dict[str, Any]], + prior_history: list[dict[str, Any]], ) -> None: """Prepend in-memory turn history, or fall back to DB-based reconstruction.""" if not previous_response_id: @@ -2074,7 +2077,7 @@ class ManagedResponsesWebSocketHandler: call_kwargs["previous_response_id"] = previous_response_id @staticmethod - def _resolve_provider(model: Optional[str]) -> Optional[str]: + def _resolve_provider(model: str | None) -> str | None: """Resolve the LLM provider for a model string, or None if unresolvable.""" if not model: return None @@ -2086,7 +2089,7 @@ class ManagedResponsesWebSocketHandler: except Exception: return None - def _same_provider(self, model: Optional[str]) -> bool: + def _same_provider(self, model: str | None) -> bool: """Return True if model uses the same LLM provider as the connection model.""" if model is None or model == self.model: return True @@ -2095,7 +2098,7 @@ class ManagedResponsesWebSocketHandler: return False return event_provider == self._connection_provider - def _inject_credentials(self, call_kwargs: Dict[str, Any], model: Optional[str] = None) -> None: + def _inject_credentials(self, call_kwargs: dict[str, Any], model: str | None = None) -> None: """Inject connection-level credentials and metadata into call_kwargs.""" if self.api_key is not None: call_kwargs["api_key"] = self.api_key @@ -2114,7 +2117,7 @@ class ManagedResponsesWebSocketHandler: call_kwargs["litellm_metadata"] = dict(self.litellm_metadata) @staticmethod - def _update_proxy_request(call_kwargs: Dict[str, Any], model: str) -> None: + def _update_proxy_request(call_kwargs: dict[str, Any], model: str) -> None: """Update proxy_server_request body so spend logs record the full request.""" proxy_server_request = (call_kwargs.get("litellm_metadata") or {}).get("proxy_server_request") or {} if not isinstance(proxy_server_request, dict): @@ -2133,7 +2136,7 @@ class ManagedResponsesWebSocketHandler: call_kwargs.setdefault("litellm_params", {}) call_kwargs["litellm_params"]["proxy_server_request"] = proxy_server_request - async def _stream_and_forward(self, model: str, call_kwargs: Dict[str, Any]) -> Optional[Dict[str, Any]]: + async def _stream_and_forward(self, model: str, call_kwargs: dict[str, Any]) -> dict[str, Any] | None: """ Stream ``litellm.aresponses`` and forward every chunk over the WebSocket. @@ -2141,7 +2144,7 @@ class ManagedResponsesWebSocketHandler: directly (before serialization) to avoid a redundant JSON round-trip on every chunk. Returns the completed event dict, or ``None``. """ - completed_event: Optional[Dict[str, Any]] = None + completed_event: dict[str, Any] | None = None stream_response = await litellm.aresponses(model=model, **call_kwargs) async for chunk in stream_response: # type: ignore[union-attr] if chunk is None: @@ -2165,9 +2168,9 @@ class ManagedResponsesWebSocketHandler: def _save_turn_history( self, - completed_event: Optional[Dict[str, Any]], - prior_history: List[Dict[str, Any]], - current_messages: List[Dict[str, Any]], + completed_event: dict[str, Any] | None, + prior_history: list[dict[str, Any]], + current_messages: list[dict[str, Any]], ) -> None: """Store this turn in in-memory history for future previous_response_id lookups.""" if completed_event is None: @@ -2236,7 +2239,7 @@ class ManagedResponsesWebSocketHandler: else: model = requested_model - previous_response_id: Optional[str] = call_kwargs.pop("previous_response_id", None) + previous_response_id: str | None = call_kwargs.pop("previous_response_id", None) current_messages = self._input_to_messages(call_kwargs.get("input")) # Fetch history once; reused in both _apply_history and _save_turn_history diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 6fb3ed748b6..265b8a9bf1a 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,6 +1,6 @@ { "ANN001": { - "limit": 3118 + "limit": 3100 }, "ANN002": { "limit": 69 @@ -9,10 +9,10 @@ "limit": 831 }, "ANN201": { - "limit": 2138 + "limit": 2135 }, "ANN202": { - "limit": 944 + "limit": 941 }, "ANN204": { "limit": 724 @@ -24,7 +24,7 @@ "limit": 130 }, "ANN401": { - "limit": 2009 + "limit": 1808 }, "ASYNC230": { "limit": 14 @@ -42,7 +42,7 @@ "limit": 84 }, "B010": { - "limit": 194 + "limit": 191 }, "B018": { "limit": 5 @@ -60,7 +60,7 @@ "limit": 4 }, "BLE001": { - "limit": 2899 + "limit": 2893 }, "C401": { "limit": 11 @@ -123,7 +123,7 @@ "limit": 52 }, "I001": { - "limit": 270 + "limit": 265 }, "LOG015": { "limit": 8 @@ -135,7 +135,7 @@ "limit": 30 }, "PERF401": { - "limit": 142 + "limit": 139 }, "PERF402": { "limit": 9 @@ -222,7 +222,7 @@ "limit": 38 }, "RET504": { - "limit": 716 + "limit": 698 }, "RUF010": { "limit": 874 @@ -306,7 +306,7 @@ "limit": 9 }, "TID251": { - "limit": 2652 + "limit": 2627 }, "TRY002": { "limit": 547 @@ -324,10 +324,10 @@ "limit": 879 }, "UP006": { - "limit": 12135 + "limit": 11225 }, "UP007": { - "limit": 2526 + "limit": 2129 }, "UP008": { "limit": 5 @@ -354,15 +354,15 @@ "limit": 4 }, "UP035": { - "limit": 2232 + "limit": 2191 }, "UP036": { "limit": 4 }, "UP037": { - "limit": 105 + "limit": 102 }, "UP045": { - "limit": 17805 + "limit": 16130 } } diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 05499e83c42..56bdc17af57 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 23250 + "limit": 23244 }, "LIT002": { - "limit": 27277 + "limit": 27176 }, "LIT003": { "limit": 292 @@ -24,6 +24,6 @@ "limit": 1004 }, "LIT009": { - "limit": 2473 + "limit": 2431 } }