diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 75d4d13eb71..685fb4af1d3 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,6 +1,6 @@ { "reportAny": { - "limit": 37484 + "limit": 32503 }, "reportArgumentType": { "limit": 2704 @@ -24,7 +24,7 @@ "limit": 42 }, "reportExplicitAny": { - "limit": 10389 + "limit": 9875 }, "reportFunctionMemberAccess": { "limit": 11 @@ -54,10 +54,10 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5900 + "limit": 5846 }, "reportMissingTypeArgument": { - "limit": 15903 + "limit": 15857 }, "reportMissingTypeStubs": { "limit": 41 @@ -99,19 +99,19 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 45894 + "limit": 45607 }, "reportUnknownLambdaType": { "limit": 113 }, "reportUnknownMemberType": { - "limit": 40539 + "limit": 40441 }, "reportUnknownParameterType": { - "limit": 20403 + "limit": 20307 }, "reportUnknownVariableType": { - "limit": 32141 + "limit": 32003 }, "reportUnnecessaryCast": { "limit": 177 @@ -138,7 +138,7 @@ "limit": 206 }, "reportUnusedImport": { - "limit": 1005 + "limit": 1004 }, "reportUnusedVariable": { "limit": 1297 diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 60dbf7c644a..43522c2008d 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -14,8 +14,11 @@ Dict, Iterator, List, + Mapping, NoReturn, Optional, + Protocol, + Sequence, Union, cast, ) @@ -23,6 +26,7 @@ import anyio import httpx from pydantic import BaseModel +from typing_extensions import Required, TypedDict import litellm from litellm import verbose_logger @@ -60,10 +64,71 @@ _SYNC_ITER_EXHAUSTED = object() -_GCHUNK_FIELDS: frozenset = frozenset(GChunk.__annotations__) +_GCHUNK_FIELDS: frozenset[str] = frozenset(GChunk.__annotations__) -def _next_sync_or_exhausted(it: Any) -> Any: +class _LoggedLiteLLMParams(TypedDict, total=False): + merge_reasoning_content_in_choices: bool | None + + +class _DumpedDeltaKwargs(TypedDict, total=False): + role: str | None + finish_reason: str | None + tool_calls: Sequence[Mapping[str, object]] | None + + +class _ModelResponseStreamInitKwargs(TypedDict, total=False): + model: str | None + + +class _PredibaseStreamDetails(TypedDict): + finish_reason: str + + +class _PredibaseStreamData(TypedDict, total=False): + token: Mapping[str, str] + details: Required[_PredibaseStreamDetails] + generated_text: str + error: str + + +class _Ai21StreamData(TypedDict): + completions: Sequence[Mapping[str, Mapping[str, str]]] + + +class _MaritalkStreamData(TypedDict): + answer: str + + +class _NlpCloudStreamData(TypedDict): + generated_text: str + + +class _AlephAlphaStreamData(TypedDict): + completions: Sequence[Mapping[str, str]] + + +class _AzureStreamChoice(TypedDict): + delta: Mapping[str, str] | None + finish_reason: str | None + + +class _AzureStreamData(TypedDict): + choices: Sequence[_AzureStreamChoice] + + +class _BasetenStreamData(TypedDict, total=False): + token: Mapping[str, str] + model_output: Mapping[str, Union[Sequence[str], str]] | str + completion: object + + +class _TextCompletionStreamChoice(Protocol): + text: str + finish_reason: str | None + + +def _next_sync_or_exhausted(it: Iterator[object]) -> object: """ Call next(it) from a thread and return _SYNC_ITER_EXHAUSTED on StopIteration. @@ -77,7 +142,7 @@ def _next_sync_or_exhausted(it: Any) -> Any: return _SYNC_ITER_EXHAUSTED -def is_async_iterable(obj: Any) -> bool: +def is_async_iterable(obj: object) -> bool: """ Check if an object is an async iterable (can be used with 'async for'). @@ -90,7 +155,7 @@ def is_async_iterable(obj: Any) -> bool: return isinstance(obj, collections.abc.AsyncIterable) -def print_verbose(print_statement): +def print_verbose(print_statement: object): try: if litellm.set_verbose: print(print_statement) # noqa: T201 @@ -116,7 +181,7 @@ def __init__( self, completion_stream, model, - logging_obj: Any, + logging_obj: LiteLLMLoggingObject, custom_llm_provider: Optional[str] = None, stream_options=None, make_call: Optional[Callable] = None, @@ -131,9 +196,8 @@ def __init__( self.sent_last_chunk = False self._stream_created_time: float = time.time() - litellm_params: GenericLiteLLMParams = GenericLiteLLMParams( - **self.logging_obj.model_call_details.get("litellm_params", {}) - ) + _logged_litellm_params: _LoggedLiteLLMParams = self.logging_obj.model_call_details.get("litellm_params", {}) + litellm_params: GenericLiteLLMParams = GenericLiteLLMParams(**_logged_litellm_params) self.merge_reasoning_content_in_choices: bool = litellm_params.merge_reasoning_content_in_choices or False self.sent_first_thinking_block = False self.sent_last_thinking_block = False @@ -195,7 +259,7 @@ def __init__( # 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, object] = { **self._hidden_params, "response_cost": None, } @@ -257,7 +321,7 @@ def check_is_function_call(self, logging_obj) -> bool: return False - def process_chunk(self, chunk: str): + def process_chunk(self, chunk: str) -> str: """ NLP Cloud streaming returns the entire response, for each chunk. Process this, to only return the delta. """ @@ -358,7 +422,7 @@ def handle_predibase_chunk(self, chunk): finish_reason = "" print_verbose(f"chunk: {chunk}") if chunk.startswith("data:"): - data_json = json.loads(chunk[5:]) + data_json: _PredibaseStreamData = json.loads(chunk[5:]) print_verbose(f"data json: {data_json}") if "token" in data_json and "text" in data_json["token"]: text = data_json["token"]["text"] @@ -388,7 +452,7 @@ def handle_predibase_chunk(self, chunk): def handle_ai21_chunk(self, chunk): # fake streaming chunk = chunk.decode("utf-8") - data_json = json.loads(chunk) + data_json: _Ai21StreamData = json.loads(chunk) try: text = data_json["completions"][0]["data"]["text"] is_finished = True @@ -403,7 +467,7 @@ def handle_ai21_chunk(self, chunk): # fake streaming def handle_maritalk_chunk(self, chunk): # fake streaming chunk = chunk.decode("utf-8") - data_json = json.loads(chunk) + data_json: _MaritalkStreamData = json.loads(chunk) try: text = data_json["answer"] is_finished = True @@ -424,7 +488,7 @@ def handle_nlp_cloud_chunk(self, chunk): if self.model and "dolphin" in self.model: chunk = self.process_chunk(chunk=chunk) else: - data_json = json.loads(chunk) + data_json: _NlpCloudStreamData = json.loads(chunk) chunk = data_json["generated_text"] text = chunk if "[DONE]" in text: @@ -441,7 +505,7 @@ def handle_nlp_cloud_chunk(self, chunk): def handle_aleph_alpha_chunk(self, chunk): chunk = chunk.decode("utf-8") - data_json = json.loads(chunk) + data_json: _AlephAlphaStreamData = json.loads(chunk) try: text = data_json["completions"][0]["completion"] is_finished = True @@ -454,7 +518,7 @@ def handle_aleph_alpha_chunk(self, chunk): except Exception: raise ValueError(f"Unable to parse response. Original response: {chunk}") - def handle_azure_chunk(self, chunk): + def handle_azure_chunk(self, chunk: str): is_finished = False finish_reason = "" text = "" @@ -469,7 +533,7 @@ def handle_azure_chunk(self, chunk): "finish_reason": finish_reason, } elif chunk.startswith("data:"): - data_json = json.loads(chunk[5:]) # chunk.startswith("data:"): + data_json: _AzureStreamData = json.loads(chunk[5:]) # chunk.startswith("data:"): try: if len(data_json["choices"]) > 0: delta = data_json["choices"][0]["delta"] @@ -558,7 +622,7 @@ def handle_azure_text_completion_chunk(self, chunk): text = "" is_finished = False finish_reason = None - choices = getattr(chunk, "choices", []) + choices: Sequence[_TextCompletionStreamChoice] = getattr(chunk, "choices", []) if len(choices) > 0: text = choices[0].text if choices[0].finish_reason is not None: @@ -579,7 +643,7 @@ def handle_openai_text_completion_chunk(self, chunk): is_finished = False finish_reason = None usage = None - choices = getattr(chunk, "choices", []) + choices: Sequence[_TextCompletionStreamChoice] = getattr(chunk, "choices", []) if len(choices) > 0: text = choices[0].text if choices[0].finish_reason is not None: @@ -601,7 +665,7 @@ def handle_baseten_chunk(self, chunk): chunk = chunk.decode("utf-8") if len(chunk) > 0: if chunk.startswith("data:"): - data_json = json.loads(chunk[5:]) + data_json: _BasetenStreamData = json.loads(chunk[5:]) if "token" in data_json and "text" in data_json["token"]: return data_json["token"]["text"] else: @@ -665,19 +729,23 @@ def handle_triton_stream(self, chunk): 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[str, object] | 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, object] = {"model": _model} else: chunk.pop("model", None) args = {"model": _model} if chunk: args.update({k: v for k, v in chunk.items() if k != "stream"}) - model_response = ModelResponseStream(**args) + model_response = ModelResponseStream( + **cast( # cast-ok: heterogeneous chunk payload consumed via ModelResponseStream(**kwargs) + _ModelResponseStreamInitKwargs, args + ) + ) if self.response_id is not None: model_response.id = self.response_id if self.system_fingerprint is not None: @@ -750,7 +818,11 @@ def copy_model_response_level_provider_specific_fields( """ Copy provider_specific_fields from original_chunk to model_response. """ - provider_specific_fields = getattr(original_chunk, "provider_specific_fields", None) + provider_specific_fields: dict[str, object] | None = ( + getattr( # mutable-ok: assigned straight to ModelResponseStream.provider_specific_fields, declared Optional[Dict[str, Any]] + original_chunk, "provider_specific_fields", None + ) + ) if provider_specific_fields is not None: model_response.provider_specific_fields = provider_specific_fields for k, v in provider_specific_fields.items(): @@ -814,7 +886,9 @@ def strip_role_from_delta(self, model_response: ModelResponseStream) -> ModelRes model_response.choices[0].delta["role"] = "assistant" self.sent_first_chunk = True elif self.sent_first_chunk is True and hasattr(model_response.choices[0].delta, "role"): - _initial_delta = model_response.choices[0].delta.model_dump() + _initial_delta = cast( # cast-ok: Delta.model_dump payload re-enters Delta(**kwargs) which takes extras + _DumpedDeltaKwargs, model_response.choices[0].delta.model_dump() + ) _initial_delta.pop("role", None) model_response.choices[0].delta = Delta(**_initial_delta) @@ -904,14 +978,17 @@ def return_processed_chunk_logic( # noqa: C901 if hold is False: ## check if openai/azure chunk - original_chunk = response_obj.get("original_chunk", None) + original_chunk: ModelResponseStream | None = response_obj.get("original_chunk", None) if original_chunk: if len(original_chunk.choices) > 0: choices = [] for choice in original_chunk.choices: try: if isinstance(choice, BaseModel): - choice_json = choice.model_dump() # type: ignore + choice_json = cast( # cast-ok: choice dump feeds StreamingChoices(**kwargs) which takes extras + _DumpedDeltaKwargs, + choice.model_dump(), + ) choice_json.pop( "finish_reason", None ) # for mistral etc. which return a value in their last chunk (not-openai compatible). @@ -946,7 +1023,11 @@ def return_processed_chunk_logic( # noqa: C901 self.sent_first_chunk = True 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) + model_response.choices[0].delta = Delta( + **cast( # cast-ok: completion payload feeds Delta(**kwargs) which takes extras + _DumpedDeltaKwargs, completion_obj + ) + ) _index: Optional[int] = completion_obj.get("index") if _index is not None: model_response.choices[0].index = _index @@ -1111,7 +1192,9 @@ def _dispatch_provider_chunk( for key, value in anthropic_response_obj["provider_specific_fields"].items(): setattr(model_response, key, value) - response_obj = cast(dict[str, Any], anthropic_response_obj) + response_obj = cast( # cast-ok: GenericStreamingChunk narrows to the plain response mapping + dict[str, object], anthropic_response_obj + ) elif self.model == "replicate" or self.custom_llm_provider == "replicate": response_obj = self.handle_replicate_chunk(chunk) completion_obj["content"] = response_obj["text"] @@ -1295,15 +1378,19 @@ def _dispatch_provider_chunk( if response_obj["is_finished"]: self.received_finish_reason = response_obj["finish_reason"] elif self.custom_llm_provider == "cached_response": - chunk = cast(ModelResponseStream, chunk) - chunk_finish_reason = chunk.choices[0].finish_reason + cached_chunk = cast( # cast-ok: cached_response streams replay ModelResponseStream chunks + ModelResponseStream, chunk + ) + chunk_finish_reason = cached_chunk.choices[0].finish_reason response_obj = { - "text": chunk.choices[0].delta.content, + "text": cached_chunk.choices[0].delta.content, "is_finished": chunk_finish_reason is not None, "finish_reason": chunk_finish_reason, - "original_chunk": chunk, + "original_chunk": cached_chunk, "tool_calls": ( - chunk.choices[0].delta.tool_calls if hasattr(chunk.choices[0].delta, "tool_calls") else None + cached_chunk.choices[0].delta.tool_calls + if hasattr(cached_chunk.choices[0].delta, "tool_calls") + else None ), } @@ -1311,11 +1398,11 @@ def _dispatch_provider_chunk( if response_obj["tool_calls"] is not None: completion_obj["tool_calls"] = response_obj["tool_calls"] print_verbose(f"completion obj content: {completion_obj['content']}") - if hasattr(chunk, "id"): - model_response.id = chunk.id - self.response_id = chunk.id - if hasattr(chunk, "system_fingerprint"): - self.system_fingerprint = chunk.system_fingerprint + if hasattr(cached_chunk, "id"): + model_response.id = cached_chunk.id + self.response_id = cached_chunk.id + if hasattr(cached_chunk, "system_fingerprint"): + self.system_fingerprint = cached_chunk.system_fingerprint if response_obj["is_finished"]: self.received_finish_reason = response_obj["finish_reason"] else: # openai / azure chat model @@ -1392,7 +1479,9 @@ def chunk_creator(self, chunk: Any): # type: ignore model_response.model = self.model ## FUNCTION CALL PARSING - original_chunk = response_obj.get("original_chunk") if response_obj is not None else None + original_chunk: ModelResponseStream | None = ( + response_obj.get("original_chunk") if response_obj is not None else None + ) if ( original_chunk is not None ): # function / tool calling branch - only set for openai/azure compatible endpoints @@ -1430,7 +1519,9 @@ def chunk_creator(self, chunk: Any): # type: ignore is None ): t.function.arguments = "" - _json_delta = delta.model_dump() + _json_delta = cast( # cast-ok: delta dump re-enters Delta(**kwargs) which takes extras + _DumpedDeltaKwargs, delta.model_dump() + ) if "role" not in _json_delta or _json_delta["role"] is None: _json_delta["role"] = "assistant" # mistral's api returns role as None if "tool_calls" in _json_delta and isinstance(_json_delta["tool_calls"], list): @@ -1458,7 +1549,11 @@ def chunk_creator(self, chunk: Any): # type: ignore if original_chunk.choices[0].delta is None else dict(original_chunk.choices[0].delta) ) - model_response.choices[0].delta = Delta(**delta) + model_response.choices[0].delta = Delta( + **cast( # cast-ok: raw delta mapping feeds Delta(**kwargs) which takes extras + _DumpedDeltaKwargs, delta + ) + ) except Exception: model_response.choices[0].delta = Delta() else: @@ -1497,7 +1592,7 @@ def chunk_creator(self, chunk: Any): # type: ignore original_exception=e, ) - def set_logging_event_loop(self, loop): + def set_logging_event_loop(self, loop: asyncio.AbstractEventLoop): """ import litellm, asyncio @@ -1672,7 +1767,9 @@ def run_success_logging_and_cache_storage(self, processed_chunk, cache_hit: bool else: asyncio.run(self.logging_obj.async_success_handler(processed_chunk, None, None, cache_hit)) ## SYNC LOGGING — only for sync SDK entrypoints; async proxy paths export via async_success_handler - litellm_params = self.logging_obj.model_call_details.get("litellm_params", {}) + litellm_params: dict[str, object] = self.logging_obj.model_call_details.get( + "litellm_params", {} + ) # mutable-ok: passed to Logging._is_sync_litellm_request, whose parameter is typed dict if self.logging_obj._is_sync_litellm_request(litellm_params): self.logging_obj.success_handler(processed_chunk, None, None, cache_hit) @@ -1700,7 +1797,9 @@ def _propagate_usage_cost_to_hidden_params( usage.cost, copy it into _hidden_params so litellm's cost calculator uses it instead of a token-based estimate. """ - _usage = getattr(response, "usage", None) + _usage = cast( # cast-ok: assembled ModelResponse carries an Optional[Usage] attribute + Usage | None, getattr(response, "usage", None) + ) if _usage is not None and hasattr(_usage, "cost") and _usage.cost is not None: if "additional_headers" not in response._hidden_params: response._hidden_params["additional_headers"] = {} @@ -2254,7 +2353,7 @@ def _strip_sse_data_from_chunk(chunk: Optional[str]) -> Optional[str]: return chunk -def calculate_total_usage(chunks: List[ModelResponse]) -> Usage: +def calculate_total_usage(chunks: Sequence[ModelResponse]) -> Usage: """Assume most recent usage chunk has total usage uptil then.""" prompt_tokens: int = 0 completion_tokens: int = 0 @@ -2287,7 +2386,7 @@ def calculate_total_usage(chunks: List[ModelResponse]) -> Usage: return returned_usage_chunk -def generic_chunk_has_all_required_fields(chunk: dict) -> bool: +def generic_chunk_has_all_required_fields(chunk: Mapping[str, object]) -> bool: """ Checks if the provided chunk dictionary contains all required fields for GenericStreamingChunk. diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index ec1301e5923..54ac71f5fdb 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -15,6 +15,7 @@ Literal, Optional, Tuple, + TypedDict, Union, cast, get_type_hints, @@ -22,6 +23,7 @@ from urllib.parse import parse_qs, urlencode, urlparse, urlunparse import httpx # type: ignore +from typing_extensions import Unpack from openai.types.file_deleted import FileDeleted import litellm @@ -161,8 +163,14 @@ def _rust_responses_websocket_enabled( if TYPE_CHECKING: from aiohttp import ClientSession + from fastapi import WebSocket + from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.llms.base_llm.realtime.http_transformation import ( + BaseRealtimeHTTPConfig, + ) + from litellm.proxy._types import UserAPIKeyAuth from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( AnthropicMessagesStreamingResponse, ) @@ -183,17 +191,41 @@ def _rust_responses_websocket_enabled( LiteLLMLoggingObj = Any +class _RawBodyKwargs(TypedDict, total=False): + data: bytes + json: dict[str, object] + + +class _DeleteRequestKwargs(TypedDict, total=False): + url: str + headers: dict[str, str] + timeout: float | httpx.Timeout | None + json: dict[str, object] + + +class _MediaUploadKwargs(TypedDict, total=False): + headers: dict[str, str] + content: Iterator[bytes] | AsyncIterator[bytes] + timeout: float | httpx.Timeout + + +class _ResponsesWebSocketLiteLLMParams(TypedDict, total=False): + api_version: str | None + organization: str | None + max_retries: int | None + + def _google_genai_streaming_hidden_params( *, api_base: str, 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 +242,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 +253,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 +265,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 @@ -359,7 +391,7 @@ async def async_completion( messages: list, optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: object, api_key: Optional[str] = None, client: Optional[AsyncHTTPHandler] = None, json_mode: bool = False, @@ -425,7 +457,7 @@ def completion( api_base: Optional[str], custom_llm_provider: str, model_response: ModelResponse, - encoding, + encoding: object, logging_obj: LiteLLMLoggingObj, optional_params: dict, timeout: Union[float, httpx.Timeout], @@ -434,7 +466,7 @@ def completion( stream: Optional[bool] = False, fake_stream: bool = False, api_key: Optional[str] = None, - headers: Optional[Dict[str, Any]] = None, + headers: dict[str, str] | None = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, provider_config: Optional[BaseConfig] = None, shared_session: Optional["ClientSession"] = None, @@ -644,14 +676,14 @@ def make_sync_call( original_data: dict, model: str, messages: list, - logging_obj, + logging_obj: LiteLLMLoggingObj, optional_params: dict, litellm_params: dict, timeout: Union[float, httpx.Timeout], fake_stream: bool = False, client: Optional[HTTPHandler] = None, json_mode: bool = False, - ) -> Tuple[Any, dict]: + ) -> tuple[object, dict]: if client is None or not isinstance(client, HTTPHandler): sync_httpx_client = _get_httpx_client( { @@ -691,7 +723,7 @@ def make_sync_call( json_mode=json_mode, ) - completion_stream: Any = MockResponseIterator(model_response=model_response, json_mode=json_mode) + completion_stream: object = MockResponseIterator(model_response=model_response, json_mode=json_mode) else: completion_stream = provider_config.get_model_response_iterator( streaming_response=response.iter_lines(), @@ -783,7 +815,7 @@ async def make_async_call_stream_helper( client: Optional[AsyncHTTPHandler] = None, json_mode: Optional[bool] = None, signed_json_body: Optional[bytes] = None, - ) -> Tuple[Any, httpx.Headers]: + ) -> tuple[object, httpx.Headers]: """ Helper function for making an async call with stream. @@ -827,7 +859,7 @@ async def make_async_call_stream_helper( json_mode=json_mode, ) - completion_stream: Any = MockResponseIterator(model_response=model_response, json_mode=json_mode) + completion_stream: object = 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 @@ -877,7 +909,7 @@ def embedding( api_key: Optional[str] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, aembedding: Optional[bool] = False, - headers: Optional[Dict[str, Any]] = None, + headers: dict[str, str] | None = None, ) -> EmbeddingResponse: provider_config = ProviderConfigManager.get_provider_embedding_config( model=model, provider=litellm.LlmProviders(custom_llm_provider) @@ -1051,7 +1083,7 @@ def rerank( timeout: Optional[Union[float, httpx.Timeout]], model_response: RerankResponse, _is_async: bool = False, - headers: Optional[Dict[str, Any]] = None, + headers: dict[str, str] | None = None, api_key: Optional[str] = None, api_base: Optional[str] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, @@ -1177,7 +1209,7 @@ def _prepare_audio_transcription_request( logging_obj: LiteLLMLoggingObj, api_key: Optional[str], api_base: Optional[str], - headers: Optional[Dict[str, Any]], + headers: dict[str, str] | None, provider_config: BaseAudioTranscriptionConfig, ) -> Tuple[dict, str, Union[dict, bytes, None], Optional[dict]]: """ @@ -1266,10 +1298,10 @@ def audio_transcriptions( custom_llm_provider: str, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, atranscription: bool = False, - headers: Optional[Dict[str, Any]] = None, + headers: dict[str, str] | None = None, provider_config: Optional[BaseAudioTranscriptionConfig] = 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}") @@ -1351,7 +1383,7 @@ async def async_audio_transcriptions( api_base: Optional[str], custom_llm_provider: str, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - headers: Optional[Dict[str, Any]] = None, + headers: dict[str, str] | None = None, provider_config: Optional[BaseAudioTranscriptionConfig] = None, shared_session: Optional["ClientSession"] = None, ) -> TranscriptionResponse: @@ -1417,10 +1449,10 @@ def _prepare_ocr_request( logging_obj: LiteLLMLoggingObj, api_key: Optional[str], api_base: Optional[str], - headers: Optional[Dict[str, Any]], + headers: dict[str, str] | None, provider_config: BaseOCRConfig, litellm_params: dict, - ) -> Tuple[Dict[str, Any], str, Dict[str, Any], None]: + ) -> tuple[dict[str, str], str, dict[str, object], None]: """ Shared logic for preparing OCR requests. Returns: (headers, complete_url, data, files) @@ -1483,10 +1515,10 @@ async def _async_prepare_ocr_request( logging_obj: LiteLLMLoggingObj, api_key: Optional[str], api_base: Optional[str], - headers: Optional[Dict[str, Any]], + headers: dict[str, str] | None, provider_config: BaseOCRConfig, litellm_params: dict, - ) -> Tuple[Dict[str, Any], str, Dict[str, Any], None]: + ) -> tuple[dict[str, str], str, dict[str, object], None]: """ Async version of _prepare_ocr_request for providers that need async transforms. Returns: (headers, complete_url, data, files) @@ -1567,10 +1599,10 @@ def ocr( custom_llm_provider: str, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, aocr: bool = False, - headers: Optional[Dict[str, Any]] = None, + headers: dict[str, str] | None = None, provider_config: Optional[BaseOCRConfig] = None, litellm_params: Optional[dict] = None, - ) -> Union[OCRResponse, Coroutine[Any, Any, OCRResponse]]: + ) -> OCRResponse | Coroutine[None, None, OCRResponse]: """ Sync OCR handler. """ @@ -1641,7 +1673,7 @@ async def async_ocr( api_base: Optional[str], custom_llm_provider: str, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - headers: Optional[Dict[str, Any]] = None, + headers: dict[str, str] | None = None, provider_config: Optional[BaseOCRConfig] = None, litellm_params: Optional[dict] = None, ) -> OCRResponse: @@ -1703,9 +1735,9 @@ def search( custom_llm_provider: str, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, asearch: bool = False, - headers: Optional[Dict[str, Any]] = None, + headers: dict[str, str] | None = None, provider_config: Optional[BaseSearchConfig] = None, - ) -> Union[SearchResponse, Coroutine[Any, Any, SearchResponse]]: + ) -> SearchResponse | Coroutine[None, None, SearchResponse]: """ Sync Search handler. """ @@ -1798,7 +1830,7 @@ async def async_search( api_base: Optional[str], custom_llm_provider: str, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - headers: Optional[Dict[str, Any]] = None, + headers: dict[str, str] | None = None, provider_config: Optional[BaseSearchConfig] = None, ) -> SearchResponse: """ @@ -1976,7 +2008,7 @@ async def async_anthropic_messages_handler( litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, client: Optional[AsyncHTTPHandler] = None, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, api_key: Optional[str] = None, api_base: Optional[str] = None, stream: Optional[bool] = False, @@ -2342,7 +2374,7 @@ def anthropic_messages_handler( kwargs: Optional[Dict[str, Any]] = None, ) -> Union[ AnthropicMessagesResponse, - Coroutine[Any, Any, Union[AnthropicMessagesResponse, AsyncIterator]], + Coroutine[None, None, AnthropicMessagesResponse | AsyncIterator], ]: """ LLM HTTP Handler for Anthropic Messages @@ -2449,18 +2481,18 @@ def response_api_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = 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]], + Coroutine[None, None, ResponsesAPIResponse | BaseResponsesAPIStreamingIterator], ]: """ Handles responses API requests. @@ -2543,7 +2575,7 @@ def response_api_handler( # 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, object] = {"input": input} try: request_context.update(response_api_optional_request_params) except Exception: @@ -2573,7 +2605,9 @@ def response_api_handler( 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: _RawBodyKwargs = ( + _RawBodyKwargs(data=signed_body) if signed_body is not None else _RawBodyKwargs(json=data) + ) ## LOGGING logging_obj.pre_call( @@ -2663,12 +2697,12 @@ async def async_response_api_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = 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]: """ @@ -2720,7 +2754,7 @@ async def async_response_api_handler( # 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, object] = {"input": input} try: request_context.update(response_api_optional_request_params) except Exception: @@ -2747,7 +2781,9 @@ async def async_response_api_handler( 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: _RawBodyKwargs = ( + _RawBodyKwargs(data=signed_body) if signed_body is not None else _RawBodyKwargs(json=data) + ) ## LOGGING logging_obj.pre_call( @@ -2847,8 +2883,8 @@ async def async_delete_response_api_handler( 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, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, _is_async: bool = False, @@ -2902,7 +2938,7 @@ async def async_delete_response_api_handler( }, ) - delete_kwargs: Dict[str, Any] = { + delete_kwargs: _DeleteRequestKwargs = { "url": url, "headers": headers, "timeout": timeout, @@ -2931,13 +2967,13 @@ def delete_response_api_handler( 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, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = 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 +3028,7 @@ def delete_response_api_handler( }, ) - delete_kwargs: Dict[str, Any] = { + delete_kwargs: _DeleteRequestKwargs = { "url": url, "headers": headers, "timeout": timeout, @@ -3021,13 +3057,13 @@ def get_responses( 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, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = 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 @@ -3102,8 +3138,8 @@ async def async_get_responses( 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, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, shared_session: Optional["ClientSession"] = None, @@ -3183,12 +3219,12 @@ def list_responses_input_items( include: Optional[List[str]] = None, limit: int = 20, order: Literal["asc", "desc"] = "desc", - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = 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, @@ -3269,7 +3305,7 @@ async def async_list_responses_input_items( include: Optional[List[str]] = None, limit: int = 20, order: Literal["asc", "desc"] = "desc", - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, shared_session: Optional["ClientSession"] = None, @@ -3357,7 +3393,7 @@ def _extract_upload_url_from_response( else: # Response body style (e.g., Manus, S3 presigned URLs) try: - response_data = response.json() + response_data: dict[str, str] = response.json() upload_url = response_data.get(upload_url_key) return upload_url, response_data if upload_url else None except Exception: @@ -3375,7 +3411,7 @@ def create_file( _is_async: bool = False, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, - ) -> Union[OpenAIFileObject, Coroutine[Any, Any, OpenAIFileObject]]: + ) -> OpenAIFileObject | Coroutine[None, None, OpenAIFileObject]: """ Creates a file using Gemini's two-step upload process """ @@ -3731,7 +3767,7 @@ def _upload_media( timeout: Optional[Union[float, httpx.Timeout]], ) -> httpx.Response: headers = {**base_headers, "Content-Type": content_type} - kwargs: Dict[str, Any] = { + kwargs: _MediaUploadKwargs = { "headers": headers, "content": self._iter_in_blocks(body_stream.iter_bytes(), self._MEDIA_UPLOAD_BLOCK_SIZE), } @@ -3768,7 +3804,7 @@ async def _abody() -> AsyncIterator[bytes]: break yield cast(bytes, block) - kwargs: Dict[str, Any] = {"headers": headers, "content": _abody()} + kwargs: _MediaUploadKwargs = {"headers": headers, "content": _abody()} if timeout is not None: kwargs["timeout"] = timeout resp = await client.client.post(url, **kwargs) @@ -3789,7 +3825,7 @@ def create_batch( client: Optional[Union["HTTPHandler", "AsyncHTTPHandler"]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, model: Optional[str] = None, - ) -> Union["LiteLLMBatch", Coroutine[Any, Any, "LiteLLMBatch"]]: + ) -> Union["LiteLLMBatch", Coroutine[None, None, "LiteLLMBatch"]]: """ Creates a batch using provider-specific batch creation process """ @@ -3901,7 +3937,7 @@ def retrieve_batch( client: Optional[Union["HTTPHandler", "AsyncHTTPHandler"]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, model: Optional[str] = None, - ) -> Union["LiteLLMBatch", Coroutine[Any, Any, "LiteLLMBatch"]]: + ) -> Union["LiteLLMBatch", Coroutine[None, None, "LiteLLMBatch"]]: """ Retrieve a batch using provider-specific configuration. """ @@ -4138,13 +4174,13 @@ def cancel_response_api_handler( 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, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = 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. @@ -4218,8 +4254,8 @@ async def async_cancel_response_api_handler( 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, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, _is_async: bool = False, @@ -4294,13 +4330,13 @@ def compact_response_api_handler( 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, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = 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 +4393,9 @@ def compact_response_api_handler( 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: _RawBodyKwargs = ( + _RawBodyKwargs(data=signed_body) if signed_body is not None else _RawBodyKwargs(json=data) + ) ## LOGGING logging_obj.pre_call( @@ -4393,8 +4431,8 @@ async def async_compact_response_api_handler( 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, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, _is_async: bool = False, @@ -4448,7 +4486,9 @@ async def async_compact_response_api_handler( 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: _RawBodyKwargs = ( + _RawBodyKwargs(data=signed_body) if signed_body is not None else _RawBodyKwargs(json=data) + ) ## LOGGING logging_obj.pre_call( @@ -4485,7 +4525,7 @@ def retrieve_file( _is_async: bool = False, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, - ) -> Union[OpenAIFileObject, Coroutine[Any, Any, OpenAIFileObject]]: + ) -> OpenAIFileObject | Coroutine[None, None, OpenAIFileObject]: """ Retrieve file metadata by ID """ @@ -4609,7 +4649,7 @@ def delete_file( _is_async: bool = False, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, - ) -> Union["FileDeleted", Coroutine[Any, Any, "FileDeleted"]]: + ) -> Union["FileDeleted", Coroutine[None, None, "FileDeleted"]]: """ Delete a file by ID """ @@ -4733,7 +4773,7 @@ def list_files( _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]]]: + ) -> list[OpenAIFileObject] | Coroutine[None, None, list[OpenAIFileObject]]: """ List all files """ @@ -4857,7 +4897,7 @@ def retrieve_file_content( _is_async: bool = False, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, - ) -> Union["HttpxBinaryResponseContent", Coroutine[Any, Any, "HttpxBinaryResponseContent"]]: + ) -> Union["HttpxBinaryResponseContent", Coroutine[None, None, "HttpxBinaryResponseContent"]]: """ Retrieve file content by ID """ @@ -4988,9 +5028,9 @@ async def async_retrieve_file_content( def _prepare_fake_stream_request( self, stream: bool, - data: dict, + data: dict[str, object], fake_stream: bool, - ) -> Tuple[bool, dict]: + ) -> tuple[bool, dict[str, object]]: """ Handles preparing a request when `fake_stream` is True. """ @@ -5081,7 +5121,7 @@ async def _execute_anthropic_agentic_plan( fingerprints: List[str], fingerprint: str, stream: bool = False, - callback: Optional[Any] = None, + callback: Optional["CustomLogger"] = None, ) -> Any: from litellm.anthropic_interface import messages as anthropic_messages @@ -5160,7 +5200,7 @@ async def _execute_responses_agentic_plan( max_loops: int, fingerprints: list[str], fingerprint: str, - callback: Any | None = None, + callback: Optional["CustomLogger"] = None, ) -> Any: patch = plan.request_patch or AgenticLoopRequestPatch() if patch.messages is None: @@ -5714,6 +5754,7 @@ def _handle_error( BaseSkillsAPIConfig, "BasePassthroughConfig", "BaseContainerConfig", + "BaseRealtimeHTTPConfig", BaseEvalsAPIConfig, ], ): @@ -5769,7 +5810,7 @@ async def _open_realtime_backend_ws( websockets_module: Any, url: str, headers: dict, - ssl_context: Any, + ssl_context: bool | str | ssl.SSLContext, *, open_timeout: float = 8.0, max_attempts: int = 3, @@ -5819,16 +5860,16 @@ async def _open_realtime_backend_ws( async def async_realtime( self, model: str, - websocket: Any, + websocket: "WebSocket", logging_obj: LiteLLMLoggingObj, provider_config: BaseRealtimeConfig, headers: dict, api_base: Optional[str] = None, api_key: Optional[str] = None, - client: Optional[Any] = None, + client: object | None = None, timeout: Optional[float] = None, - user_api_key_dict: Optional[Any] = None, - litellm_metadata: Optional[Dict[str, Any]] = None, + user_api_key_dict: Optional["UserAPIKeyAuth"] = None, + litellm_metadata: dict[str, object] | None = None, query_params: Optional[RealtimeQueryParams] = None, ): import websockets @@ -5850,7 +5891,7 @@ async def async_realtime( 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, object] = {} if litellm_metadata: _request_data["litellm_metadata"] = litellm_metadata realtime_streaming = RealTimeStreaming( @@ -5922,12 +5963,12 @@ async def async_realtime_client_secret_handler( self, api_base: str, api_key: str, - request_data: Dict[str, Any], + request_data: dict[str, object], logging_obj: LiteLLMLoggingObj, timeout: Union[float, httpx.Timeout], - provider_config: Optional[Any] = None, + provider_config: Optional["BaseRealtimeHTTPConfig"] = None, model: Optional[str] = None, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, api_version: Optional[str] = None, ) -> httpx.Response: @@ -5955,12 +5996,12 @@ async def async_realtime_transcription_session_handler( self, api_base: str, api_key: str, - request_data: Dict[str, Any], + request_data: dict[str, object], logging_obj: LiteLLMLoggingObj, timeout: Union[float, httpx.Timeout], - provider_config: Optional[Any] = None, + provider_config: Optional["BaseRealtimeHTTPConfig"] = None, model: Optional[str] = None, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, api_version: Optional[str] = None, ) -> httpx.Response: @@ -5984,12 +6025,12 @@ async def _async_realtime_session_post( endpoint: Literal["client_secrets", "transcription_sessions"], api_base: str, api_key: str, - request_data: Dict[str, Any], + request_data: dict[str, object], logging_obj: LiteLLMLoggingObj, timeout: Union[float, httpx.Timeout], - provider_config: Optional[Any] = None, + provider_config: Optional["BaseRealtimeHTTPConfig"] = None, model: Optional[str] = None, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, api_version: Optional[str] = None, ) -> httpx.Response: @@ -6014,7 +6055,7 @@ async def _async_realtime_session_post( ) 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, str] = provider_config.validate_environment( headers={}, model=model or "", api_key=api_key ) else: @@ -6059,10 +6100,10 @@ async def async_realtime_calls_handler( sdp_body: bytes, logging_obj: LiteLLMLoggingObj, timeout: Union[float, httpx.Timeout], - provider_config: Optional[Any] = None, + provider_config: Optional["BaseRealtimeHTTPConfig"] = None, model: Optional[str] = None, - session_config: Optional[Dict[str, Any]] = None, - extra_headers: Optional[Dict[str, Any]] = None, + session_config: dict[str, object] | None = None, + extra_headers: dict[str, str] | None = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, api_version: Optional[str] = None, ) -> httpx.Response: @@ -6085,7 +6126,7 @@ async def async_realtime_calls_handler( 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, str] = provider_config.get_realtime_calls_headers(ephemeral_key=openai_ephemeral_key) else: url = f"{api_base.rstrip('/')}/v1/realtime/calls" headers = { @@ -6137,17 +6178,17 @@ async def async_realtime_calls_handler( async def async_responses_websocket( self, model: str, - websocket: Any, + websocket: "WebSocket", 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, + user_api_key_dict: Optional["UserAPIKeyAuth"] = None, + litellm_metadata: dict[str, object] | None = None, custom_llm_provider: Optional[str] = None, first_message: Optional[str] = None, - **kwargs: Any, + **kwargs: Unpack[_ResponsesWebSocketLiteLLMParams], ): """ Handles Responses API WebSocket mode. @@ -6251,7 +6292,7 @@ async def _backend_connection(): yield backend async with _backend_connection() as backend_ws: - _request_data: Dict[str, Any] = {} + _request_data: dict[str, object] = {} if litellm_metadata: _request_data["litellm_metadata"] = litellm_metadata @@ -6322,15 +6363,15 @@ def image_edit_handler( 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, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, _is_async: bool = False, fake_stream: bool = False, - litellm_metadata: Optional[Dict[str, Any]] = None, + litellm_metadata: dict[str, object] | None = None, ) -> Union[ ImageResponse, - Coroutine[Any, Any, ImageResponse], + Coroutine[None, None, ImageResponse], ]: """ @@ -6442,11 +6483,11 @@ async def async_image_edit_handler( 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, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = 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. @@ -6540,16 +6581,16 @@ def image_generation_handler( 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, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, _is_async: bool = False, fake_stream: bool = False, - litellm_metadata: Optional[Dict[str, Any]] = None, + litellm_metadata: dict[str, object] | None = None, api_key: Optional[str] = None, ) -> Union[ ImageResponse, - Coroutine[Any, Any, ImageResponse], + Coroutine[None, None, ImageResponse], ]: """ Handles image generation requests. @@ -6667,11 +6708,11 @@ async def async_image_generation_handler( 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, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, fake_stream: bool = False, - litellm_metadata: Optional[Dict[str, Any]] = None, + litellm_metadata: dict[str, object] | None = None, api_key: Optional[str] = None, ) -> ImageResponse: """ @@ -6775,16 +6816,16 @@ def video_generation_handler( 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, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, _is_async: bool = False, fake_stream: bool = False, - litellm_metadata: Optional[Dict[str, Any]] = None, + litellm_metadata: dict[str, object] | None = None, api_key: Optional[str] = None, ) -> Union[ VideoObject, - Coroutine[Any, Any, VideoObject], + Coroutine[None, None, VideoObject], ]: """ Handles video generation requests. @@ -6899,11 +6940,11 @@ async def async_video_generation_handler( 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, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, fake_stream: bool = False, - litellm_metadata: Optional[Dict[str, Any]] = None, + litellm_metadata: dict[str, object] | None = None, api_key: Optional[str] = None, ) -> VideoObject: """ @@ -6999,12 +7040,12 @@ def video_content_handler( litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, timeout: Union[float, httpx.Timeout], - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, api_key: Optional[str] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, _is_async: bool = False, variant: Optional[str] = None, - ) -> Union[bytes, Coroutine[Any, Any, bytes]]: + ) -> bytes | Coroutine[None, None, bytes]: """ Handle video content download requests. """ @@ -7089,7 +7130,7 @@ async def async_video_content_handler( litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, timeout: Union[float, httpx.Timeout], - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, api_key: Optional[str] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, variant: Optional[str] = None, @@ -7165,13 +7206,13 @@ def video_remix_handler( prompt: str, video_remix_provider_config: BaseVideoConfig, custom_llm_provider: str, - litellm_params, - logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + litellm_params: GenericLiteLLMParams, + logging_obj: LiteLLMLoggingObj, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, timeout: Optional[float] = None, _is_async: bool = False, - client=None, + client: HTTPHandler | AsyncHTTPHandler | None = None, api_key: Optional[str] = None, ): """ @@ -7264,12 +7305,12 @@ async def async_video_remix_handler( prompt: str, video_remix_provider_config: BaseVideoConfig, custom_llm_provider: str, - litellm_params, - logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + litellm_params: GenericLiteLLMParams, + logging_obj: LiteLLMLoggingObj, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, timeout: Optional[float] = None, - client=None, + client: HTTPHandler | AsyncHTTPHandler | None = None, api_key: Optional[str] = None, ): """ @@ -7344,15 +7385,15 @@ async def async_video_remix_handler( def video_create_character_handler( self, name: str, - video: Any, + video: object, video_provider_config: BaseVideoConfig, custom_llm_provider: str, - litellm_params, - logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, + litellm_params: GenericLiteLLMParams, + logging_obj: LiteLLMLoggingObj, + extra_headers: dict[str, str] | None = None, timeout: Optional[float] = None, _is_async: bool = False, - client=None, + client: HTTPHandler | AsyncHTTPHandler | None = None, api_key: Optional[str] = None, ): if _is_async: @@ -7428,14 +7469,14 @@ def video_create_character_handler( async def async_video_create_character_handler( self, name: str, - video: Any, + video: object, video_provider_config: BaseVideoConfig, custom_llm_provider: str, - litellm_params, - logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, + litellm_params: GenericLiteLLMParams, + logging_obj: LiteLLMLoggingObj, + extra_headers: dict[str, str] | None = None, timeout: Optional[float] = None, - client=None, + client: HTTPHandler | AsyncHTTPHandler | None = None, api_key: Optional[str] = None, ): if client is None or not isinstance(client, AsyncHTTPHandler): @@ -7502,12 +7543,12 @@ def video_get_character_handler( character_id: str, video_provider_config: BaseVideoConfig, custom_llm_provider: str, - litellm_params, - logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, + litellm_params: GenericLiteLLMParams, + logging_obj: LiteLLMLoggingObj, + extra_headers: dict[str, str] | None = None, timeout: Optional[float] = None, _is_async: bool = False, - client=None, + client: HTTPHandler | AsyncHTTPHandler | None = None, api_key: Optional[str] = None, ): if _is_async: @@ -7571,11 +7612,11 @@ async def async_video_get_character_handler( character_id: str, video_provider_config: BaseVideoConfig, custom_llm_provider: str, - litellm_params, - logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, + litellm_params: GenericLiteLLMParams, + logging_obj: LiteLLMLoggingObj, + extra_headers: dict[str, str] | None = None, timeout: Optional[float] = None, - client=None, + client: HTTPHandler | AsyncHTTPHandler | None = None, api_key: Optional[str] = None, ): if client is None or not isinstance(client, AsyncHTTPHandler): @@ -7630,13 +7671,13 @@ def video_edit_handler( video_id: str, video_provider_config: BaseVideoConfig, custom_llm_provider: str, - litellm_params, - logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + litellm_params: GenericLiteLLMParams, + logging_obj: LiteLLMLoggingObj, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, timeout: Optional[float] = None, _is_async: bool = False, - client=None, + client: HTTPHandler | AsyncHTTPHandler | None = None, api_key: Optional[str] = None, ): if _is_async: @@ -7739,12 +7780,12 @@ async def async_video_edit_handler( video_id: str, video_provider_config: BaseVideoConfig, custom_llm_provider: str, - litellm_params, - logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + litellm_params: GenericLiteLLMParams, + logging_obj: LiteLLMLoggingObj, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, timeout: Optional[float] = None, - client=None, + client: HTTPHandler | AsyncHTTPHandler | None = None, api_key: Optional[str] = None, ): if client is None or not isinstance(client, AsyncHTTPHandler): @@ -7836,13 +7877,13 @@ def video_extension_handler( seconds: str, video_provider_config: BaseVideoConfig, custom_llm_provider: str, - litellm_params, - logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + litellm_params: GenericLiteLLMParams, + logging_obj: LiteLLMLoggingObj, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, timeout: Optional[float] = None, _is_async: bool = False, - client=None, + client: HTTPHandler | AsyncHTTPHandler | None = None, api_key: Optional[str] = None, ): if _is_async: @@ -7925,12 +7966,12 @@ async def async_video_extension_handler( seconds: str, video_provider_config: BaseVideoConfig, custom_llm_provider: str, - litellm_params, - logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + litellm_params: GenericLiteLLMParams, + logging_obj: LiteLLMLoggingObj, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, timeout: Optional[float] = None, - client=None, + client: HTTPHandler | AsyncHTTPHandler | None = None, api_key: Optional[str] = None, ): if client is None or not isinstance(client, AsyncHTTPHandler): @@ -8000,13 +8041,13 @@ def video_list_handler( order: Optional[str], 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, + litellm_params: GenericLiteLLMParams, + logging_obj: LiteLLMLoggingObj, + extra_headers: dict[str, str] | None = None, + extra_query: dict[str, object] | None = None, timeout: Optional[float] = None, _is_async: bool = False, - client=None, + client: HTTPHandler | AsyncHTTPHandler | None = None, api_key: Optional[str] = None, ): """ @@ -8054,12 +8095,12 @@ async def async_video_list_handler( order: Optional[str], 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, + litellm_params: GenericLiteLLMParams, + logging_obj: LiteLLMLoggingObj, + extra_headers: dict[str, str] | None = None, + extra_query: dict[str, object] | None = None, timeout: Optional[float] = None, - client=None, + client: HTTPHandler | AsyncHTTPHandler | None = None, api_key: Optional[str] = None, ): """ @@ -8135,11 +8176,11 @@ async def async_video_delete_handler( video_id: str, video_delete_provider_config: BaseVideoConfig, custom_llm_provider: str, - litellm_params, - logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, + litellm_params: GenericLiteLLMParams, + logging_obj: LiteLLMLoggingObj, + extra_headers: dict[str, str] | None = None, timeout: Optional[float] = None, - client=None, + client: HTTPHandler | AsyncHTTPHandler | None = None, api_key: Optional[str] = None, ): """ @@ -8211,13 +8252,13 @@ def video_status_handler( video_id: str, video_status_provider_config: BaseVideoConfig, custom_llm_provider: str, - litellm_params, - logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + litellm_params: GenericLiteLLMParams, + logging_obj: LiteLLMLoggingObj, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, timeout: Optional[float] = None, _is_async: bool = False, - client=None, + client: HTTPHandler | AsyncHTTPHandler | None = None, api_key: Optional[str] = None, ): """ @@ -8316,12 +8357,12 @@ async def async_video_status_handler( video_id: str, video_status_provider_config: BaseVideoConfig, custom_llm_provider: str, - litellm_params, - logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + litellm_params: GenericLiteLLMParams, + logging_obj: LiteLLMLoggingObj, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, timeout: Optional[float] = None, - client=None, + client: HTTPHandler | AsyncHTTPHandler | None = None, api_key: Optional[str] = None, ): """ @@ -8408,11 +8449,11 @@ def container_create_handler( container_provider_config: "BaseContainerConfig", litellm_params: GenericLiteLLMParams, logging_obj: "LiteLLMLoggingObj", - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Union[float, httpx.Timeout] = 600, _is_async: bool = False, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - ) -> Union["ContainerObject", Coroutine[Any, Any, "ContainerObject"]]: + ) -> 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( @@ -8495,7 +8536,7 @@ async def async_container_create_handler( container_provider_config: "BaseContainerConfig", litellm_params: GenericLiteLLMParams, logging_obj: "LiteLLMLoggingObj", - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Union[float, httpx.Timeout] = 600, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, ) -> "ContainerObject": @@ -8572,12 +8613,12 @@ def container_list_handler( 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, + extra_headers: dict[str, str] | None = None, + extra_query: dict[str, object] | None = None, timeout: Union[float, httpx.Timeout] = 600, _is_async: bool = False, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - ) -> Union["ContainerListResponse", Coroutine[Any, Any, "ContainerListResponse"]]: + ) -> 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( @@ -8662,8 +8703,8 @@ async def async_container_list_handler( 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, + extra_headers: dict[str, str] | None = None, + extra_query: dict[str, object] | None = None, timeout: Union[float, httpx.Timeout] = 600, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, ) -> "ContainerListResponse": @@ -8737,12 +8778,12 @@ def container_retrieve_handler( container_provider_config: "BaseContainerConfig", litellm_params: GenericLiteLLMParams, logging_obj: "LiteLLMLoggingObj", - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, + extra_query: dict[str, object] | None = None, timeout: Union[float, httpx.Timeout] = 600, _is_async: bool = False, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - ) -> Union["ContainerObject", Coroutine[Any, Any, "ContainerObject"]]: + ) -> 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( @@ -8825,8 +8866,8 @@ async def async_container_retrieve_handler( container_provider_config: "BaseContainerConfig", litellm_params: GenericLiteLLMParams, logging_obj: "LiteLLMLoggingObj", - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, + extra_query: dict[str, object] | None = None, timeout: Union[float, httpx.Timeout] = 600, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, ) -> "ContainerObject": @@ -8902,12 +8943,12 @@ def container_delete_handler( container_provider_config: "BaseContainerConfig", litellm_params: GenericLiteLLMParams, logging_obj: "LiteLLMLoggingObj", - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, + extra_query: dict[str, object] | None = None, timeout: Union[float, httpx.Timeout] = 600, _is_async: bool = False, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - ) -> Union["DeleteContainerResult", Coroutine[Any, Any, "DeleteContainerResult"]]: + ) -> 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( @@ -8990,8 +9031,8 @@ async def async_container_delete_handler( container_provider_config: "BaseContainerConfig", litellm_params: GenericLiteLLMParams, logging_obj: "LiteLLMLoggingObj", - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, + extra_query: dict[str, object] | None = None, timeout: Union[float, httpx.Timeout] = 600, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, ) -> "DeleteContainerResult": @@ -9070,12 +9111,12 @@ def container_file_list_handler( 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, + extra_headers: dict[str, str] | None = None, + extra_query: dict[str, object] | None = None, timeout: Union[float, httpx.Timeout] = 600, _is_async: bool = False, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - ) -> Union["ContainerFileListResponse", Coroutine[Any, Any, "ContainerFileListResponse"]]: + ) -> Union["ContainerFileListResponse", Coroutine[None, None, "ContainerFileListResponse"]]: if _is_async: return self.async_container_file_list_handler( container_id=container_id, @@ -9162,8 +9203,8 @@ async def async_container_file_list_handler( 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, + extra_headers: dict[str, str] | None = None, + extra_query: dict[str, object] | None = None, timeout: Union[float, httpx.Timeout] = 600, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, ) -> "ContainerFileListResponse": @@ -9239,11 +9280,11 @@ def container_file_content_handler( container_provider_config: "BaseContainerConfig", litellm_params: GenericLiteLLMParams, logging_obj: "LiteLLMLoggingObj", - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Union[float, httpx.Timeout] = 600, _is_async: bool = False, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - ) -> Union[bytes, Coroutine[Any, Any, bytes]]: + ) -> bytes | Coroutine[None, None, bytes]: if _is_async: return self.async_container_file_content_handler( container_id=container_id, @@ -9325,7 +9366,7 @@ async def async_container_file_content_handler( container_provider_config: "BaseContainerConfig", litellm_params: GenericLiteLLMParams, logging_obj: "LiteLLMLoggingObj", - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Union[float, httpx.Timeout] = 600, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, ) -> bytes: @@ -9404,8 +9445,8 @@ async def async_vector_store_search_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, _is_async: bool = False, @@ -9457,7 +9498,7 @@ async def async_vector_store_search_handler( litellm_params=dict(litellm_params), extra_body=extra_body, ) - all_optional_params: Dict[str, Any] = dict(litellm_params) + all_optional_params: dict[str, object] = 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, @@ -9502,12 +9543,12 @@ def vector_store_search_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = 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, @@ -9553,7 +9594,7 @@ def vector_store_search_handler( extra_body=extra_body, ) - all_optional_params: Dict[str, Any] = dict(litellm_params) + all_optional_params: dict[str, object] = dict(litellm_params) all_optional_params.update(vector_store_search_optional_params or {}) headers, signed_json_body = vector_store_provider_config.sign_request( @@ -9596,8 +9637,8 @@ async def async_vector_store_create_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, _is_async: bool = False, @@ -9656,12 +9697,12 @@ def vector_store_create_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = 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, @@ -9726,8 +9767,8 @@ async def async_vector_store_retrieve_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, ) -> VectorStoreCreateResponse: @@ -9779,12 +9820,12 @@ def vector_store_retrieve_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = 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, @@ -9846,8 +9887,8 @@ async def async_vector_store_list_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, ): @@ -9873,7 +9914,7 @@ async def async_vector_store_list_handler( url = api_base - params: Dict[str, Any] = {} + params: dict[str, object] = {} if after is not None: params["after"] = after if before is not None: @@ -9910,8 +9951,8 @@ def vector_store_list_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, _is_async: bool = False, @@ -9951,7 +9992,7 @@ def vector_store_list_handler( url = api_base - params: Dict[str, Any] = {} + params: dict[str, object] = {} if after is not None: params["after"] = after if before is not None: @@ -9986,8 +10027,8 @@ async def async_vector_store_update_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, ) -> VectorStoreCreateResponse: @@ -10052,12 +10093,12 @@ def vector_store_update_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = 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, @@ -10129,8 +10170,8 @@ async def async_vector_store_delete_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, ): @@ -10180,8 +10221,8 @@ def vector_store_delete_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, _is_async: bool = False, @@ -10247,8 +10288,8 @@ async def async_vector_store_file_create_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, ) -> VectorStoreFileObject: @@ -10312,12 +10353,12 @@ def vector_store_file_create_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = 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, @@ -10389,8 +10430,8 @@ async def async_vector_store_file_list_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, + extra_query: dict[str, object] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, ) -> VectorStoreFileListResponse: @@ -10453,12 +10494,12 @@ def vector_store_file_list_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, + extra_query: dict[str, object] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = 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, @@ -10529,7 +10570,7 @@ async def async_vector_store_file_retrieve_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, ) -> VectorStoreFileObject: @@ -10588,11 +10629,11 @@ def vector_store_file_retrieve_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = 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, @@ -10658,7 +10699,7 @@ async def async_vector_store_file_content_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, ) -> VectorStoreFileContentResponse: @@ -10719,13 +10760,13 @@ def vector_store_file_content_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, _is_async: bool = False, ) -> Union[ VectorStoreFileContentResponse, - Coroutine[Any, Any, VectorStoreFileContentResponse], + Coroutine[None, None, VectorStoreFileContentResponse], ]: if _is_async: return self.async_vector_store_file_content_handler( @@ -10795,8 +10836,8 @@ async def async_vector_store_file_update_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, ) -> VectorStoreFileObject: @@ -10861,12 +10902,12 @@ def vector_store_file_update_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = 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, @@ -10939,7 +10980,7 @@ async def async_vector_store_file_delete_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, ) -> VectorStoreFileDeleteResponse: @@ -10998,13 +11039,13 @@ def vector_store_file_delete_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, _is_async: bool = False, ) -> Union[ VectorStoreFileDeleteResponse, - Coroutine[Any, Any, VectorStoreFileDeleteResponse], + Coroutine[None, None, VectorStoreFileDeleteResponse], ]: if _is_async: return self.async_vector_store_file_delete_handler( @@ -11068,21 +11109,21 @@ def vector_store_file_delete_handler( def generate_content_handler( self, model: str, - contents: Any, + contents: object, generate_content_provider_config: BaseGoogleGenAIGenerateContentConfig, generate_content_config_dict: Dict, - tools: Any, + tools: object, custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, _is_async: bool = False, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = 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: object | None = None, ) -> Any: """ Handles Google GenAI generate content requests. @@ -11200,20 +11241,20 @@ def generate_content_handler( async def async_generate_content_handler( self, model: str, - contents: Any, + contents: object, generate_content_provider_config: BaseGoogleGenAIGenerateContentConfig, generate_content_config_dict: Dict, - tools: Any, + tools: object, custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, object] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[AsyncHTTPHandler] = 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: object | None = None, ) -> Any: """ Async version of the generate content handler. @@ -11326,12 +11367,12 @@ def text_to_speech_handler( litellm_params: Dict, logging_obj: LiteLLMLoggingObj, timeout: Union[float, httpx.Timeout], - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, _is_async: bool = False, ) -> Union[ "HttpxBinaryResponseContent", - Coroutine[Any, Any, "HttpxBinaryResponseContent"], + Coroutine[None, None, "HttpxBinaryResponseContent"], ]: """ Handles text-to-speech requests. @@ -11441,7 +11482,7 @@ async def async_text_to_speech_handler( litellm_params: Dict, logging_obj: LiteLLMLoggingObj, timeout: Union[float, httpx.Timeout], - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, ) -> "HttpxBinaryResponseContent": """ @@ -11573,12 +11614,12 @@ def create_skill_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = 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( @@ -11639,7 +11680,7 @@ async def async_create_skill_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, shared_session: Optional["ClientSession"] = None, @@ -11695,12 +11736,12 @@ def list_skills_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = 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( @@ -11754,7 +11795,7 @@ async def async_list_skills_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, shared_session: Optional["ClientSession"] = None, @@ -11800,12 +11841,12 @@ def get_skill_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = 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( @@ -11856,7 +11897,7 @@ async def async_get_skill_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, shared_session: Optional["ClientSession"] = None, @@ -11901,12 +11942,12 @@ def delete_skill_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = 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( @@ -11957,7 +11998,7 @@ async def async_delete_skill_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, shared_session: Optional["ClientSession"] = None, @@ -12007,12 +12048,12 @@ def create_eval_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = 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( @@ -12066,7 +12107,7 @@ async def async_create_eval_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, shared_session: Optional["ClientSession"] = None, @@ -12113,12 +12154,12 @@ def list_evals_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = 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( @@ -12172,7 +12213,7 @@ async def async_list_evals_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, shared_session: Optional["ClientSession"] = None, @@ -12218,12 +12259,12 @@ def get_eval_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = 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( @@ -12274,7 +12315,7 @@ async def async_get_eval_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, shared_session: Optional["ClientSession"] = None, @@ -12320,12 +12361,12 @@ def update_eval_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = 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( @@ -12379,7 +12420,7 @@ async def async_update_eval_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, shared_session: Optional["ClientSession"] = None, @@ -12425,12 +12466,12 @@ def delete_eval_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = 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( @@ -12481,7 +12522,7 @@ async def async_delete_eval_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, shared_session: Optional["ClientSession"] = None, @@ -12526,12 +12567,12 @@ def cancel_eval_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = 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( @@ -12582,7 +12623,7 @@ async def async_cancel_eval_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, shared_session: Optional["ClientSession"] = None, @@ -12632,12 +12673,12 @@ def create_run_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = 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( @@ -12691,7 +12732,7 @@ async def async_create_run_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, shared_session: Optional["ClientSession"] = None, @@ -12738,12 +12779,12 @@ def list_runs_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = 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( @@ -12797,7 +12838,7 @@ async def async_list_runs_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, shared_session: Optional["ClientSession"] = None, @@ -12843,12 +12884,12 @@ def get_run_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = 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( @@ -12899,7 +12940,7 @@ async def async_get_run_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, shared_session: Optional["ClientSession"] = None, @@ -12944,12 +12985,12 @@ def cancel_run_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = 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( @@ -13000,7 +13041,7 @@ async def async_cancel_run_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, shared_session: Optional["ClientSession"] = None, @@ -13045,12 +13086,12 @@ def delete_run_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = 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( @@ -13101,7 +13142,7 @@ async def async_delete_run_handler( custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, str] | None = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, shared_session: Optional["ClientSession"] = None, diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index 6b191144a11..98248bfe4e2 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -61,6 +61,116 @@ drop_params_from_unprocessable_entity_error, ) +if TYPE_CHECKING: + from openai.types.batch import Errors as BatchErrors + from openai.types.batch_request_counts import BatchRequestCounts + from openai.types.beta.threads.message import Attachment as MessageAttachment + from openai.types.beta.threads.message import ( + IncompleteDetails as MessageIncompleteDetails, + ) + from typing_extensions import NotRequired, TypedDict + + from litellm.types.utils import Usage as LiteLLMUsage + + class _FileObjectDump(TypedDict): + id: str + bytes: int + created_at: int + filename: str + object: Literal["file"] + purpose: OpenAIFilesPurpose + status: Literal["uploaded", "processed", "error"] + expires_at: int | None + status_details: str | None + + class _BatchDump(TypedDict): + id: str + completion_window: str + created_at: int + endpoint: str + input_file_id: str + object: Literal["batch"] + status: Literal[ + "validating", + "failed", + "in_progress", + "finalizing", + "completed", + "expired", + "cancelling", + "cancelled", + ] + cancelled_at: int | None + cancelling_at: int | None + completed_at: int | None + error_file_id: str | None + errors: BatchErrors | None + expired_at: int | None + expires_at: int | None + failed_at: int | None + finalizing_at: int | None + in_progress_at: int | None + metadata: Optional[ + Dict[str, str] + ] # mutable-ok: unpacked into LiteLLMBatch, whose metadata field is declared Dict[str, str] + model: str | None + output_file_id: str | None + request_counts: BatchRequestCounts | None + usage: LiteLLMUsage | None + + class _MessageDump(TypedDict): + id: str + assistant_id: str | None + attachments: Optional[ + list[MessageAttachment] + ] # mutable-ok: unpacked into OpenAIMessage, whose attachments field is declared List[Attachment] + completed_at: int | None + content: list[ + MessageContent + ] # mutable-ok: unpacked into OpenAIMessage, whose content field is declared List[MessageContent] + created_at: int + incomplete_at: int | None + incomplete_details: MessageIncompleteDetails | None + metadata: Optional[ + Dict[str, str] + ] # mutable-ok: unpacked into OpenAIMessage, whose metadata field is declared Dict[str, str] + object: Literal["thread.message"] + role: Literal["user", "assistant"] + run_id: str | None + status: Literal["in_progress", "incomplete", "completed"] + thread_id: str + + class _ThreadDump(TypedDict): + id: str + created_at: int + metadata: object | None + object: Literal["thread"] + + class _RunThreadStreamData(TypedDict): + thread_id: str + assistant_id: str + additional_instructions: str | None + instructions: str | None + metadata: Optional[ + Dict[str, str] + ] # mutable-ok: unpacked into runs.stream(), whose metadata param is declared Dict[str, str] + model: str | None + tools: Iterable[AssistantToolParam] | None + event_handler: NotRequired[AssistantEventHandler] + + class _AsyncRunThreadStreamData(TypedDict): + thread_id: str + assistant_id: str + additional_instructions: str | None + instructions: str | None + metadata: Optional[ + Dict[str, str] + ] # mutable-ok: unpacked into runs.stream(), whose metadata param is declared Dict[str, str] + model: str | None + tools: Iterable[AssistantToolParam] | None + event_handler: NotRequired[AsyncAssistantEventHandler] + + openaiOSeriesConfig = OpenAIOSeriesConfig() openAIGPT5Config = OpenAIGPT5Config() @@ -73,7 +183,9 @@ class MistralEmbeddingConfig: def __init__( self, ) -> None: - locals_ = locals().copy() + locals_: Dict[str, object] = ( + locals().copy() + ) # mutable-ok: invariant dict pins locals()'s values to object; Mapping re-widens them to Any for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) @@ -165,7 +277,9 @@ def __init__( top_p: Optional[int] = None, response_format: Optional[dict] = None, ) -> None: - locals_ = locals().copy() + locals_: Dict[str, object] = ( + locals().copy() + ) # mutable-ok: invariant dict pins locals()'s values to object; Mapping re-widens them to Any for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) @@ -275,16 +389,19 @@ def transform_response( messages: List[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: object, api_key: Optional[str] = None, json_mode: Optional[bool] = None, ) -> ModelResponse: logging_obj.post_call(original_response=raw_response.text) logging_obj.model_call_details["response_headers"] = raw_response.headers + raw_response_json = cast( # cast-ok: chat completion responses parse to a JSON object + "Dict[str, object]", raw_response.json() + ) final_response_obj = cast( ModelResponse, convert_to_model_response_object( - response_object=raw_response.json(), + response_object=raw_response_json, model_response_object=model_response, hidden_params={"headers": raw_response.headers}, _response_headers=dict(raw_response.headers), @@ -313,7 +430,7 @@ def get_model_response_iterator( streaming_response: Union[Iterator[str], AsyncIterator[str], ModelResponse], sync_stream: bool, json_mode: Optional[bool] = False, - ) -> Any: + ) -> "OpenAIChatCompletionResponseIterator": return OpenAIChatCompletionResponseIterator( streaming_response=streaming_response, sync_stream=sync_stream, @@ -488,14 +605,14 @@ def make_sync_openai_chat_completion_request( async def _call_agentic_completion_hooks_openai( self, - response: Any, + response: object, model: str, messages: List[Dict], optional_params: Dict, logging_obj: LiteLLMLoggingObj, stream: bool, litellm_params: Dict, - ) -> Optional[Any]: + ) -> ModelResponse | None: """ Call agentic completion hooks for all custom loggers (OpenAI Chat Completions API). @@ -546,15 +663,18 @@ async def _call_agentic_completion_hooks_openai( kwargs_with_provider["custom_llm_provider"] = custom_llm_provider # For OpenAI Chat Completions, use the chat completion agentic loop method - agentic_response = await callback.async_run_chat_completion_agentic_loop( - tools=tool_calls, - model=model, - messages=messages, - response=response, - optional_params=optional_params, - logging_obj=logging_obj, - stream=stream, - kwargs=kwargs_with_provider, + agentic_response = cast( # cast-ok: chat completion agentic loop yields a ModelResponse + "ModelResponse | None", + await callback.async_run_chat_completion_agentic_loop( + tools=tool_calls, + model=model, + messages=messages, + response=response, + optional_params=optional_params, + logging_obj=logging_obj, + stream=stream, + kwargs=kwargs_with_provider, + ), ) # First hook that runs agentic loop wins return agentic_response @@ -816,7 +936,7 @@ def completion( # type: ignore status_code = getattr(e, "status_code", 500) error_headers = getattr(e, "headers", None) error_text = getattr(e, "text", str(e)) - error_response = getattr(e, "response", None) + error_response: httpx.Response | None = getattr(e, "response", None) error_body = getattr(e, "body", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) @@ -935,7 +1055,7 @@ async def acompletion( raise e # e.message except Exception as e: - exception_response = getattr(e, "response", None) + exception_response: httpx.Response | None = getattr(e, "response", None) status_code = getattr(e, "status_code", 500) exception_body = getattr(e, "body", None) error_headers = getattr(e, "headers", None) @@ -1092,7 +1212,7 @@ async def async_streaming( error_headers = getattr(e, "headers", None) status_code = getattr(e, "status_code", 500) - error_response = getattr(e, "response", None) + error_response: httpx.Response | None = getattr(e, "response", None) exception_body = getattr(e, "body", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) @@ -1247,7 +1367,7 @@ async def aembedding( status_code = getattr(e, "status_code", 500) error_headers = getattr(e, "headers", None) error_text = getattr(e, "text", str(e)) - error_response = getattr(e, "response", None) + error_response: httpx.Response | None = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers) @@ -1333,7 +1453,7 @@ def embedding( # type: ignore status_code = getattr(e, "status_code", 500) error_headers = getattr(e, "headers", None) error_text = getattr(e, "text", str(e)) - error_response = getattr(e, "response", None) + error_response: httpx.Response | None = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers) @@ -1626,7 +1746,10 @@ async def acreate_file( openai_client: AsyncOpenAI, ) -> OpenAIFileObject: response = await openai_client.files.create(**create_file_data) # type: ignore[arg-type] - return OpenAIFileObject(**response.model_dump()) + response_dict = cast( # cast-ok: FileObject.model_dump returns its typed fields + "_FileObjectDump", response.model_dump() + ) + return OpenAIFileObject(**response_dict) def create_file( self, @@ -1638,7 +1761,7 @@ def create_file( max_retries: Optional[int], organization: Optional[str], client: Optional[Union[OpenAI, AsyncOpenAI]] = None, - ) -> Union[OpenAIFileObject, Coroutine[Any, Any, OpenAIFileObject]]: + ) -> OpenAIFileObject | Coroutine[None, None, OpenAIFileObject]: openai_client: Optional[Union[OpenAI, AsyncOpenAI]] = self.get_openai_client( api_key=api_key, api_base=api_base, @@ -1662,7 +1785,10 @@ def create_file( create_file_data=create_file_data, openai_client=openai_client ) response = cast(OpenAI, openai_client).files.create(**create_file_data) # type: ignore[arg-type] - return OpenAIFileObject(**response.model_dump()) + response_dict = cast( # cast-ok: FileObject.model_dump returns its typed fields + "_FileObjectDump", response.model_dump() + ) + return OpenAIFileObject(**response_dict) async def afile_content( self, @@ -1682,7 +1808,7 @@ def file_content( max_retries: Optional[int], organization: Optional[str], client: Optional[Union[OpenAI, AsyncOpenAI]] = None, - ) -> Union[HttpxBinaryResponseContent, Coroutine[Any, Any, HttpxBinaryResponseContent]]: + ) -> HttpxBinaryResponseContent | Coroutine[None, None, HttpxBinaryResponseContent]: openai_client: Optional[Union[OpenAI, AsyncOpenAI]] = self.get_openai_client( api_key=api_key, api_base=api_base, @@ -1986,7 +2112,8 @@ async def acreate_batch( openai_client: AsyncOpenAI, ) -> LiteLLMBatch: response = await openai_client.batches.create(**create_batch_data) # type: ignore[arg-type] - return LiteLLMBatch(**response.model_dump()) + batch_dict = cast("_BatchDump", response.model_dump()) # cast-ok: Batch.model_dump returns its typed fields + return LiteLLMBatch(**batch_dict) def create_batch( self, @@ -1998,7 +2125,7 @@ def create_batch( max_retries: Optional[int], organization: Optional[str], client: Optional[Union[OpenAI, AsyncOpenAI]] = None, - ) -> Union[LiteLLMBatch, Coroutine[Any, Any, LiteLLMBatch]]: + ) -> LiteLLMBatch | Coroutine[None, None, LiteLLMBatch]: openai_client: Optional[Union[OpenAI, AsyncOpenAI]] = self.get_openai_client( api_key=api_key, api_base=api_base, @@ -2023,7 +2150,8 @@ def create_batch( ) response = cast(OpenAI, openai_client).batches.create(**create_batch_data) # type: ignore[arg-type] - return LiteLLMBatch(**response.model_dump()) + batch_dict = cast("_BatchDump", response.model_dump()) # cast-ok: Batch.model_dump returns its typed fields + return LiteLLMBatch(**batch_dict) async def aretrieve_batch( self, @@ -2032,7 +2160,8 @@ async def aretrieve_batch( ) -> LiteLLMBatch: verbose_logger.debug("retrieving batch, args= %s", retrieve_batch_data) response = await openai_client.batches.retrieve(**retrieve_batch_data) # type: ignore[arg-type] - return LiteLLMBatch(**response.model_dump()) + batch_dict = cast("_BatchDump", response.model_dump()) # cast-ok: Batch.model_dump returns its typed fields + return LiteLLMBatch(**batch_dict) def retrieve_batch( self, @@ -2068,7 +2197,8 @@ def retrieve_batch( retrieve_batch_data=retrieve_batch_data, openai_client=openai_client ) response = cast(OpenAI, openai_client).batches.retrieve(**retrieve_batch_data) # type: ignore[arg-type] - return LiteLLMBatch(**response.model_dump()) + batch_dict = cast("_BatchDump", response.model_dump()) # cast-ok: Batch.model_dump returns its typed fields + return LiteLLMBatch(**batch_dict) async def acancel_batch( self, @@ -2077,7 +2207,8 @@ async def acancel_batch( ) -> LiteLLMBatch: verbose_logger.debug("async cancelling batch, args= %s", cancel_batch_data) response = await openai_client.batches.cancel(**cancel_batch_data) - return LiteLLMBatch(**response.model_dump()) + batch_dict = cast("_BatchDump", response.model_dump()) # cast-ok: Batch.model_dump returns its typed fields + return LiteLLMBatch(**batch_dict) def cancel_batch( self, @@ -2117,7 +2248,8 @@ def cancel_batch( if not isinstance(openai_client, OpenAI): raise ValueError("OpenAI client is not an instance of OpenAI. Make sure you passed a sync OpenAI client.") response = openai_client.batches.cancel(**cancel_batch_data) - return LiteLLMBatch(**response.model_dump()) + batch_dict = cast("_BatchDump", response.model_dump()) # cast-ok: Batch.model_dump returns its typed fields + return LiteLLMBatch(**batch_dict) async def alist_batches( self, @@ -2477,9 +2609,11 @@ async def a_add_message( response_obj: Optional[OpenAIMessage] = None if getattr(thread_message, "status", None) is None: thread_message.status = "completed" - response_obj = OpenAIMessage(**thread_message.dict()) + message_dict = cast("_MessageDump", thread_message.dict()) # cast-ok: Message.dict returns its typed fields + response_obj = OpenAIMessage(**message_dict) else: - response_obj = OpenAIMessage(**thread_message.dict()) + message_dict = cast("_MessageDump", thread_message.dict()) # cast-ok: Message.dict returns its typed fields + response_obj = OpenAIMessage(**message_dict) return response_obj # fmt: off @@ -2556,9 +2690,11 @@ def add_message( response_obj: Optional[OpenAIMessage] = None if getattr(thread_message, "status", None) is None: thread_message.status = "completed" - response_obj = OpenAIMessage(**thread_message.dict()) + message_dict = cast("_MessageDump", thread_message.dict()) # cast-ok: Message.dict returns its typed fields + response_obj = OpenAIMessage(**message_dict) else: - response_obj = OpenAIMessage(**thread_message.dict()) + message_dict = cast("_MessageDump", thread_message.dict()) # cast-ok: Message.dict returns its typed fields + response_obj = OpenAIMessage(**message_dict) return response_obj async def async_get_messages( @@ -2680,7 +2816,10 @@ async def async_create_thread( message_thread = await openai_client.beta.threads.create(**data) # type: ignore - return Thread(**message_thread.dict()) + thread_dict = cast( # cast-ok: thread dump carries litellm Thread's typed fields + "_ThreadDump", message_thread.dict() + ) + return Thread(**thread_dict) # fmt: off @@ -2766,7 +2905,10 @@ def create_thread( message_thread = openai_client.beta.threads.create(**data) # type: ignore - return Thread(**message_thread.dict()) + thread_dict = cast( # cast-ok: thread dump carries litellm Thread's typed fields + "_ThreadDump", message_thread.dict() + ) + return Thread(**thread_dict) async def async_get_thread( self, @@ -2789,7 +2931,8 @@ async def async_get_thread( response = await openai_client.beta.threads.retrieve(thread_id=thread_id) - return Thread(**response.dict()) + thread_dict = cast("_ThreadDump", response.dict()) # cast-ok: thread dump carries litellm Thread's typed fields + return Thread(**thread_dict) # fmt: off @@ -2855,7 +2998,8 @@ def get_thread( response = openai_client.beta.threads.retrieve(thread_id=thread_id) - return Thread(**response.dict()) + thread_dict = cast("_ThreadDump", response.dict()) # cast-ok: thread dump carries litellm Thread's typed fields + return Thread(**thread_dict) def delete_thread(self): pass @@ -2912,7 +3056,7 @@ def async_run_thread_stream( tools: Optional[Iterable[AssistantToolParam]], event_handler: Optional[AssistantEventHandler], ) -> AsyncAssistantStreamManager[AsyncAssistantEventHandler]: - data: Dict[str, Any] = { + data: _AsyncRunThreadStreamData = { "thread_id": thread_id, "assistant_id": assistant_id, "additional_instructions": additional_instructions, @@ -2922,7 +3066,9 @@ def async_run_thread_stream( "tools": tools, } if event_handler is not None: - data["event_handler"] = event_handler + data["event_handler"] = cast( # cast-ok: async runs stream requires an async event handler + "AsyncAssistantEventHandler", event_handler + ) return client.beta.threads.runs.stream(**data) # type: ignore def run_thread_stream( @@ -2937,7 +3083,7 @@ def run_thread_stream( tools: Optional[Iterable[AssistantToolParam]], event_handler: Optional[AssistantEventHandler], ) -> AssistantStreamManager[AssistantEventHandler]: - data: Dict[str, Any] = { + data: _RunThreadStreamData = { "thread_id": thread_id, "assistant_id": assistant_id, "additional_instructions": additional_instructions, diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 9fe970f7fa9..652852fe693 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -3,7 +3,24 @@ import hashlib import json from datetime import datetime, timedelta, timezone -from typing import TYPE_CHECKING, Any, Awaitable, Callable, Dict, Iterable, List, Optional, Set, Union, cast +from typing import ( + TYPE_CHECKING, + Any, + Awaitable, + Callable, + Dict, + Iterable, + List, + Mapping, + Optional, + Protocol, + Sequence, + Set, + TypedDict, + TypeVar, + Union, + cast, +) from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid @@ -16,7 +33,6 @@ from litellm.proxy._types import ( LiteLLM_MCPServerTable, LiteLLM_ObjectPermissionTable, - LiteLLM_TeamTable, MCPApprovalStatus, MCPEnvVarScope, MCPSubmissionsSummary, @@ -45,9 +61,19 @@ from litellm.types.mcp import MCPCredentials if TYPE_CHECKING: + from prisma import Prisma + from prisma.models import LiteLLM_MCPServerOAuthClient as PrismaMCPServerOAuthClientRow + from prisma.models import LiteLLM_MCPServerTable as PrismaMCPServerTableRow + from prisma.models import LiteLLM_MCPUserCredentials as PrismaMCPUserCredentialsRow + from prisma.models import LiteLLM_MCPUserEnvVars as PrismaMCPUserEnvVarsRow + from prisma.models import LiteLLM_ObjectPermissionTable as PrismaObjectPermissionTableRow + from prisma.models import LiteLLM_TeamTable as PrismaTeamTableRow + from prisma.models import LiteLLM_VerificationToken as PrismaVerificationTokenRow + + from litellm.models.mcp_server import MCPEnvVar from litellm.types.mcp_server.mcp_server_manager import MCPServer -_AUTH_FLOW_SCOPED_FIELDS: frozenset = frozenset( +_AUTH_FLOW_SCOPED_FIELDS: frozenset[str] = frozenset( { "issuer", "authorization_url", @@ -76,7 +102,7 @@ def _blank_to_none(value: Optional[str]) -> Optional[str]: # the current code has never written — a cleared column can then never be # silently resurrected by a stale blob copy. These keys are stored plaintext # (endpoints/identifiers, not secrets), so values lift as-is. -_TOKEN_EXCHANGE_COLUMN_FIELDS: frozenset = frozenset( +_TOKEN_EXCHANGE_COLUMN_FIELDS: frozenset[str] = frozenset( { "token_exchange_endpoint", "audience", @@ -89,10 +115,122 @@ def _blank_to_none(value: Optional[str]) -> Optional[str]: # OAuth app (client_id/client_secret) plus the same authorize relay, and neither mints anything the # gateway keeps. So a switch WITHIN this class must preserve the stored app, unlike a cross-class # switch (e.g. an oauth2 row whose client may be DCR-minted and is not reusable elsewhere). -_CLIENT_FORWARDED_AUTH_TYPES: frozenset = frozenset({"true_passthrough", "oauth_delegate"}) +_CLIENT_FORWARDED_AUTH_TYPES: frozenset[str] = frozenset({"true_passthrough", "oauth_delegate"}) # Minted token material that must never survive a client rotation on a persisted row. -_MINTED_TOKEN_CREDENTIAL_FIELDS: frozenset = frozenset({"access_token", "refresh_token", "expires_in"}) +_MINTED_TOKEN_CREDENTIAL_FIELDS: frozenset[str] = frozenset({"access_token", "refresh_token", "expires_in"}) + + +_RowT_co = TypeVar("_RowT_co", covariant=True) + + +class _PrismaTable(Protocol[_RowT_co]): + async def find_unique( + self, + where: Mapping[str, object], + include: Mapping[str, bool] | None = None, + ) -> _RowT_co | None: ... + + async def find_many( + self, + where: Mapping[str, object] | None = None, + include: Mapping[str, bool] | None = None, + order: Mapping[str, str] | None = None, + take: int | None = None, + ) -> Sequence[_RowT_co]: ... + + async def create(self, data: Mapping[str, object]) -> _RowT_co: ... + + async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> _RowT_co: ... + + async def upsert(self, where: Mapping[str, object], data: Mapping[str, Mapping[str, object]]) -> _RowT_co: ... + + async def delete(self, where: Mapping[str, object]) -> _RowT_co | None: ... + + async def delete_many(self, where: Mapping[str, object]) -> int: ... + + +def _mcp_server_table(prisma_client: PrismaClient) -> "_PrismaTable[PrismaMCPServerTableRow]": + return cast( # cast-ok: repository ``table`` is typed Any; narrow it to the prisma table surface used here + "_PrismaTable[PrismaMCPServerTableRow]", + MCPServerRepository(prisma_client).table, + ) + + +def _verification_token_table(prisma_client: PrismaClient) -> "_PrismaTable[PrismaVerificationTokenRow]": + return cast( # cast-ok: repository ``table`` is typed Any; narrow it to the prisma table surface used here + "_PrismaTable[PrismaVerificationTokenRow]", + VerificationTokenRepository(prisma_client).table, + ) + + +def _team_table(prisma_client: PrismaClient) -> "_PrismaTable[PrismaTeamTableRow]": + return cast( # cast-ok: repository ``table`` is typed Any; narrow it to the prisma table surface used here + "_PrismaTable[PrismaTeamTableRow]", + TeamRepository(prisma_client).table, + ) + + +def _object_permission_table(prisma_client: PrismaClient) -> "_PrismaTable[PrismaObjectPermissionTableRow]": + return cast( # cast-ok: repository ``table`` is typed Any; narrow it to the prisma table surface used here + "_PrismaTable[PrismaObjectPermissionTableRow]", + ObjectPermissionRepository(prisma_client).table, + ) + + +def _mcp_user_credentials_table(prisma_client: PrismaClient) -> "_PrismaTable[PrismaMCPUserCredentialsRow]": + return cast( # cast-ok: repository ``table`` is typed Any; narrow it to the prisma table surface used here + "_PrismaTable[PrismaMCPUserCredentialsRow]", + MCPUserCredentialsRepository(prisma_client).table, + ) + + +def _mcp_oauth_client_table(prisma_client: PrismaClient) -> "_PrismaTable[PrismaMCPServerOAuthClientRow]": + return cast( # cast-ok: repository ``table`` is typed Any; narrow it to the prisma table surface used here + "_PrismaTable[PrismaMCPServerOAuthClientRow]", + MCPServerOAuthClientRepository(prisma_client).table, + ) + + +def _prisma_db(prisma_client: PrismaClient) -> "Prisma": + return cast( # cast-ok: PrismaClient.db is an untyped proxy around the generated prisma client + "Prisma", prisma_client.db + ) + + +def _mcp_user_env_vars_table(prisma_client: PrismaClient) -> "_PrismaTable[PrismaMCPUserEnvVarsRow]": + return cast( # cast-ok: the generated table client methods reference types pyright cannot load; narrow to the surface used here + "_PrismaTable[PrismaMCPUserEnvVarsRow]", + _prisma_db(prisma_client).litellm_mcpuserenvvars, + ) + + +def _credentials_blob(value: object) -> str | Mapping[str, object] | None: + return cast( # cast-ok: prisma types Json columns opaquely; at runtime the column holds the deserialized blob + "str | Mapping[str, object] | None", value + ) + + +def _env_vars_blob(value: object) -> Sequence[Mapping[str, str]] | None: + return cast( # cast-ok: prisma types Json columns opaquely; at runtime the column holds the deserialized blob + "Sequence[Mapping[str, str]] | None", value + ) + + +def _loads_json_object(value: str) -> object: + return cast( # cast-ok: json.loads returns Any; the parsed payload is handled as an opaque object + object, json.loads(value) + ) + + +def _parse_credentials_blob_dict( + blob: str | Mapping[str, object], +) -> dict[str, object]: # mutable-ok: callers pop lifted token-exchange keys off the returned copy + if isinstance(blob, str): + return cast( # cast-ok: the credentials column always stores a serialized JSON object + dict[str, object], json.loads(blob) + ) + return dict(blob) def _credential_auth_class(auth_type: Optional[str]) -> Optional[str]: @@ -104,7 +242,7 @@ def _credential_auth_class(auth_type: Optional[str]) -> Optional[str]: return auth_type -def _drop_stale_minted_on_client_rotation(merged: Dict[str, Any], new_creds: Dict[str, Any]) -> Dict[str, Any]: +def _drop_stale_minted_on_client_rotation(merged: dict[str, object], new_creds: dict[str, object]) -> dict[str, object]: """When the update rotates the client, drop stale minted token keys it did not itself set, so an old app's access/refresh token never rides forward under the new client. A no-op when no client key changed.""" if "client_id" not in new_creds and "client_secret" not in new_creds: @@ -114,13 +252,13 @@ def _drop_stale_minted_on_client_rotation(merged: Dict[str, Any], new_creds: Dic } -def _is_global_env_var_scope(scope: Any) -> bool: +def _is_global_env_var_scope(scope: object) -> bool: """``scope="user"`` entries are placeholders the user fills in; everything else (including a missing scope) is an admin-supplied global value.""" return scope != MCPEnvVarScope.user and scope != "user" -def _encrypt_global_env_var_values(env_vars: Iterable[Dict[str, Any]]) -> None: +def _encrypt_global_env_var_values(env_vars: Iterable[dict[str, str]]) -> None: """Encrypt ``scope="global"`` env var values in place before persisting. Global values hold admin-supplied secrets (API keys, passwords) that get @@ -136,7 +274,11 @@ def _encrypt_global_env_var_values(env_vars: Iterable[Dict[str, Any]]) -> None: entry["value"] = encrypt_value_helper(value) -def decrypt_global_env_var_values(env_vars: Optional[Iterable[Any]]) -> None: +def decrypt_global_env_var_values( + env_vars: Optional[ + Iterable[Union[dict[str, str], "MCPEnvVar"]] + ], # mutable-ok: each entry's value is decrypted in place +) -> None: """Decrypt ``scope="global"`` env var values in place after reading the DB. Accepts ``MCPEnvVar`` models (``LiteLLM_MCPServerTable``) or plain dicts @@ -175,7 +317,7 @@ def decrypt_global_env_var_values(env_vars: Optional[Iterable[Any]]) -> None: entry.value = decrypted -def _decrypt_env_vars_on_returned_row(row: Any) -> None: +def _decrypt_env_vars_on_returned_row(row: object) -> None: """Decrypt ``scope="global"`` env var values on a row returned by Prisma create/update. Prisma may hand back ``env_vars`` either as a parsed list (the common case for @@ -187,16 +329,21 @@ def _decrypt_env_vars_on_returned_row(row: Any) -> None: write the decrypted list back onto the row so downstream consumers see plain values. """ - env_vars = getattr(row, "env_vars", None) + env_vars = cast( # cast-ok: prisma returns the Json column as a parsed list or a raw JSON string + "str | list[dict[str, str]] | None", getattr(row, "env_vars", None) + ) if env_vars is None: return if isinstance(env_vars, str): try: - env_vars = json.loads(env_vars) + parsed = _loads_json_object(env_vars) except (json.JSONDecodeError, TypeError): return - if not isinstance(env_vars, list): + if not isinstance(parsed, list): return + env_vars = cast( # cast-ok: the parsed payload is the stored env var entry list + "list[dict[str, str]]", parsed + ) try: setattr(row, "env_vars", env_vars) except (AttributeError, TypeError): @@ -205,8 +352,8 @@ def _decrypt_env_vars_on_returned_row(row: Any) -> None: def _reencrypt_global_env_var_values( - env_vars: Optional[Iterable[Any]], new_encryption_key: str -) -> Optional[List[Dict[str, Any]]]: + env_vars: Iterable[Mapping[str, str]] | None, new_encryption_key: str +) -> Sequence[Mapping[str, str]] | None: """Re-encrypt ``scope="global"`` env var values for master-key rotation. Each global value is decrypted with the current salt key and re-encrypted @@ -219,7 +366,9 @@ def _reencrypt_global_env_var_values( return None if isinstance(env_vars, str): try: - env_vars = json.loads(env_vars) + env_vars = cast( # cast-ok: the env var column stores a serialized list of entries + "list[dict[str, str]]", json.loads(env_vars) + ) except (json.JSONDecodeError, TypeError): return None if not env_vars: @@ -438,12 +587,12 @@ async def get_all_mcp_servers( Pass approval_status=None to return all servers regardless of approval state. """ try: - where: Dict[str, Any] = {} + where: dict[str, str] = {} if approval_status is not None: where["approval_status"] = approval_status - mcp_servers = await MCPServerRepository(prisma_client).table.find_many(where=where if where else {}) + mcp_servers = await _mcp_server_table(prisma_client).find_many(where=where if where else {}) - tables = [LiteLLM_MCPServerTable(**mcp_server.model_dump()) for mcp_server in mcp_servers] + tables = [LiteLLM_MCPServerTable.model_validate(mcp_server.model_dump()) for mcp_server in mcp_servers] for table in tables: decrypt_global_env_var_values(table.env_vars) return tables @@ -458,14 +607,14 @@ async def get_mcp_server(prisma_client: PrismaClient, server_id: str) -> Optiona """ Returns the matching mcp server from the db iff exists """ - mcp_server: Optional[LiteLLM_MCPServerTable] = await MCPServerRepository(prisma_client).table.find_unique( + mcp_server = await _mcp_server_table(prisma_client).find_unique( where={ "server_id": server_id, } ) if mcp_server is None: return None - table = LiteLLM_MCPServerTable(**mcp_server.model_dump()) + table = LiteLLM_MCPServerTable.model_validate(mcp_server.model_dump()) decrypt_global_env_var_values(table.env_vars) return table @@ -474,14 +623,14 @@ async def get_mcp_servers(prisma_client: PrismaClient, server_ids: Iterable[str] """ Returns the matching mcp servers from the db with the server_ids """ - _mcp_servers: List[LiteLLM_MCPServerTable] = await MCPServerRepository(prisma_client).table.find_many( + _mcp_servers = await _mcp_server_table(prisma_client).find_many( where={ - "server_id": {"in": server_ids}, + "server_id": {"in": list(server_ids)}, } ) final_mcp_servers: List[LiteLLM_MCPServerTable] = [] for _mcp_server in _mcp_servers: - table = LiteLLM_MCPServerTable(**_mcp_server.model_dump()) + table = LiteLLM_MCPServerTable.model_validate(_mcp_server.model_dump()) decrypt_global_env_var_values(table.env_vars) final_mcp_servers.append(table) @@ -492,7 +641,7 @@ async def get_mcp_servers_by_verificationtoken(prisma_client: PrismaClient, toke """ Returns the mcp servers from the db for the verification token """ - verification_token_record: LiteLLM_TeamTable = await VerificationTokenRepository(prisma_client).table.find_unique( + verification_token_record = await _verification_token_table(prisma_client).find_unique( where={ "token": token, }, @@ -511,7 +660,7 @@ async def get_mcp_servers_by_team(prisma_client: PrismaClient, team_id: str) -> """ Returns the mcp servers from the db for the team id """ - team_record: LiteLLM_TeamTable = await TeamRepository(prisma_client).table.find_unique( + team_record = await _team_table(prisma_client).find_unique( where={ "team_id": team_id, }, @@ -561,7 +710,7 @@ async def get_objectpermissions_for_mcp_server( """ Get all the object permissions records and the associated team and verficiationtoken records that have access to the mcp server """ - object_permission_records = await ObjectPermissionRepository(prisma_client).table.find_many( + object_permission_records = await _object_permission_table(prisma_client).find_many( where={ "mcp_servers": {"has": mcp_server_id}, }, @@ -571,21 +720,23 @@ async def get_objectpermissions_for_mcp_server( }, ) - return object_permission_records + return cast( # cast-ok: callers consume the prisma rows under the table-model annotation; runtime is unchanged + list[LiteLLM_ObjectPermissionTable], object_permission_records + ) -async def get_virtualkeys_for_mcp_server(prisma_client: PrismaClient, server_id: str) -> List: +async def get_virtualkeys_for_mcp_server( + prisma_client: PrismaClient, server_id: str +) -> Sequence["PrismaVerificationTokenRow"]: """ Get all the virtual keys that have access to the mcp server """ - virtual_keys = await VerificationTokenRepository(prisma_client).table.find_many( + virtual_keys = await _verification_token_table(prisma_client).find_many( where={ "mcp_servers": {"has": server_id}, }, ) - if virtual_keys is None: - return [] return virtual_keys @@ -626,7 +777,7 @@ async def delete_mcp_server( Returns the deleted mcp server record if it exists, otherwise None """ - deleted_server = await MCPServerRepository(prisma_client).table.delete( + deleted_server = await _mcp_server_table(prisma_client).delete( where={ "server_id": server_id, }, @@ -634,9 +785,7 @@ async def delete_mcp_server( if deleted_server is not None: credential_user_ids: List[str] = [] try: - credential_rows = await prisma_client.db.litellm_mcpusercredentials.find_many( - where={"server_id": server_id} - ) + credential_rows = await _mcp_user_credentials_table(prisma_client).find_many(where={"server_id": server_id}) credential_user_ids = [row.user_id for row in credential_rows] except Exception as e: # noqa: BLE001 - enumeration is best-effort; cached tokens expire by TTL verbose_proxy_logger.warning( @@ -645,9 +794,9 @@ async def delete_mcp_server( e, ) for model, label in ( - (prisma_client.db.litellm_mcpusercredentials, "credential"), - (prisma_client.db.litellm_mcpuserenvvars, "env var"), - (prisma_client.db.litellm_mcpserveroauthclient, "OAuth client"), + (_mcp_user_credentials_table(prisma_client), "credential"), + (_mcp_user_env_vars_table(prisma_client), "env var"), + (_mcp_oauth_client_table(prisma_client), "OAuth client"), ): try: await model.delete_many(where={"server_id": server_id}) @@ -668,7 +817,9 @@ async def delete_mcp_server( invalidate_token_cache = global_mcp_server_manager.invalidate_user_oauth_token_cache for user_id in credential_user_ids: await invalidate_token_cache(user_id, server_id) - return deleted_server + return cast( # cast-ok: callers consume the prisma row under the table-model annotation; runtime is unchanged + LiteLLM_MCPServerTable | None, deleted_server + ) async def create_mcp_server( @@ -687,12 +838,12 @@ async def create_mcp_server( data_dict["created_by"] = touched_by data_dict["updated_by"] = touched_by - new_mcp_server = await MCPServerRepository(prisma_client).table.create( - data=data_dict # type: ignore - ) + new_mcp_server = await _mcp_server_table(prisma_client).create(data=data_dict) _decrypt_env_vars_on_returned_row(new_mcp_server) - return new_mcp_server + return cast( # cast-ok: callers consume the prisma row under the table-model annotation; runtime is unchanged + LiteLLM_MCPServerTable, new_mcp_server + ) async def update_mcp_server( @@ -704,8 +855,6 @@ async def update_mcp_server( """ Update a new mcp server record in the db """ - import json - from litellm.litellm_core_utils.safe_json_dumps import safe_dumps # Use helper to prepare data with proper JSON serialization. @@ -723,7 +872,8 @@ async def update_mcp_server( url_provided = "url" in data_dict and data_dict["url"] is not None issuer_provided = "issuer" in data_dict if data.auth_type or has_credentials or explicit_te_write or url_provided or issuer_provided: - existing = await MCPServerRepository(prisma_client).table.find_unique(where={"server_id": data.server_id}) + existing = await _mcp_server_table(prisma_client).find_unique(where={"server_id": data.server_id}) + existing_credentials = _credentials_blob(existing.credentials) if existing is not None else None auth_type_changed = bool( data.auth_type @@ -762,10 +912,8 @@ async def update_mcp_server( # place, and the next credentials update's migrate-on-write would silently # repopulate the column the admin just cleared. (When credentials ARE in the # update, the merge below performs the same migration.) - if explicit_te_write and "credentials" not in data_dict and existing is not None and existing.credentials: - existing_creds = ( - json.loads(existing.credentials) if isinstance(existing.credentials, str) else dict(existing.credentials) - ) + if explicit_te_write and "credentials" not in data_dict and existing is not None and existing_credentials: + existing_creds = _parse_credentials_blob_dict(existing_credentials) if _TOKEN_EXCHANGE_COLUMN_FIELDS & existing_creds.keys(): for te_field in _TOKEN_EXCHANGE_COLUMN_FIELDS: legacy_value = existing_creds.pop(te_field, None) @@ -777,23 +925,15 @@ async def update_mcp_server( # Without this, a partial credential update (e.g. changing only region) # would wipe encrypted secrets that the UI cannot display back. if "credentials" in data_dict and data_dict["credentials"] is not None: - if existing and existing.credentials: + if existing and existing_credentials: # Only merge when the credential CLASS is unchanged. A cross-class switch # (e.g. oauth2 → api_key, or oauth2 → true_passthrough) replaces credentials # entirely to avoid stale secrets from the previous class lingering; a switch # within the client-forwarded class (true_passthrough ↔ oauth_delegate) keeps # the same declared app and so must merge, not replace. if not auth_type_changed: - existing_creds = ( - json.loads(existing.credentials) - if isinstance(existing.credentials, str) - else dict(existing.credentials) - ) - new_creds = ( - json.loads(data_dict["credentials"]) - if isinstance(data_dict["credentials"], str) - else dict(data_dict["credentials"]) - ) + existing_creds = _parse_credentials_blob_dict(existing_credentials) + new_creds = _parse_credentials_blob_dict(data_dict["credentials"]) # New values override existing; existing keys not in update are preserved. A client # rotation additionally drops the previous app's stale minted token keys. merged = _drop_stale_minted_on_client_rotation({**existing_creds, **new_creds}, new_creds) @@ -823,13 +963,15 @@ async def update_mcp_server( data_dict["credentials"] = Json(None) - updated_mcp_server = await MCPServerRepository(prisma_client).table.update( + updated_mcp_server = await _mcp_server_table(prisma_client).update( where={"server_id": data.server_id}, - data=data_dict, # type: ignore + data=data_dict, ) _decrypt_env_vars_on_returned_row(updated_mcp_server) - return updated_mcp_server + return cast( # cast-ok: callers consume the prisma row under the table-model annotation; runtime is unchanged + LiteLLM_MCPServerTable, updated_mcp_server + ) async def get_mcp_server_oauth_client_credentials(prisma_client: PrismaClient, server_id: str) -> object | None: @@ -838,7 +980,7 @@ async def get_mcp_server_oauth_client_credentials(prisma_client: PrismaClient, s LiteLLM_MCPServerTable row, so their dynamically registered client lives here keyed by server_id. The returned value is the raw credentials blob for ``_get_persisted_dcr_credentials`` to parse.""" - row = await MCPServerOAuthClientRepository(prisma_client).table.find_unique(where={"server_id": server_id}) + row = await _mcp_oauth_client_table(prisma_client).find_unique(where={"server_id": server_id}) if row is None: return None return row.credentials @@ -856,7 +998,7 @@ async def upsert_mcp_server_oauth_client_credentials( encrypted = encrypt_credentials(credentials=dict(credentials), encryption_key=_get_salt_key()) blob = safe_dumps(encrypted) - await MCPServerOAuthClientRepository(prisma_client).table.upsert( + await _mcp_oauth_client_table(prisma_client).upsert( where={"server_id": server_id}, data={ "create": {"server_id": server_id, "credentials": blob}, @@ -883,17 +1025,17 @@ def _reencrypt_mcp_credentials_blob(credentials: object, new_master_key: str) -> async def rotate_mcp_server_credentials_master_key(prisma_client: PrismaClient, touched_by: str, new_master_key: str): from litellm.litellm_core_utils.safe_json_dumps import safe_dumps # noqa: PLC0415 # avoids circular import - mcp_servers = await MCPServerRepository(prisma_client).table.find_many() + mcp_servers = await _mcp_server_table(prisma_client).find_many() updated = 0 for mcp_server in mcp_servers: - update_data: Dict[str, Any] = {} + update_data: dict[str, str] = {} rotated_credentials = _reencrypt_mcp_credentials_blob(mcp_server.credentials, new_master_key) if rotated_credentials is not None: update_data["credentials"] = rotated_credentials - rotated_env_vars = _reencrypt_global_env_var_values(mcp_server.env_vars, new_master_key) + rotated_env_vars = _reencrypt_global_env_var_values(_env_vars_blob(mcp_server.env_vars), new_master_key) if rotated_env_vars is not None: update_data["env_vars"] = safe_dumps(rotated_env_vars) @@ -901,19 +1043,19 @@ async def rotate_mcp_server_credentials_master_key(prisma_client: PrismaClient, continue update_data["updated_by"] = touched_by - await MCPServerRepository(prisma_client).table.update( + await _mcp_server_table(prisma_client).update( where={"server_id": mcp_server.server_id}, data=update_data, ) updated += 1 - oauth_clients = await MCPServerOAuthClientRepository(prisma_client).table.find_many() + oauth_clients = await _mcp_oauth_client_table(prisma_client).find_many() oauth_updated = 0 for oauth_client in oauth_clients: rotated_credentials = _reencrypt_mcp_credentials_blob(oauth_client.credentials, new_master_key) if rotated_credentials is None: continue - await MCPServerOAuthClientRepository(prisma_client).table.update( + await _mcp_oauth_client_table(prisma_client).update( where={"server_id": oauth_client.server_id}, data={"credentials": rotated_credentials}, ) @@ -959,7 +1101,7 @@ def _decode_oauth_payload(stored: str) -> Optional[Dict[str, Any]]: if decoded is None: return None try: - parsed = json.loads(decoded) + parsed = _loads_json_object(decoded) except (ValueError, TypeError): return None if isinstance(parsed, dict) and parsed.get("type") == "oauth2": @@ -975,7 +1117,7 @@ async def rotate_mcp_user_credentials_master_key(prisma_client: PrismaClient, ne under the new master key. Rows that are unreadable under both paths are logged and skipped so one corrupt row does not abort the rotation. """ - rows = await MCPUserCredentialsRepository(prisma_client).table.find_many() + rows = await _mcp_user_credentials_table(prisma_client).find_many() rotated = 0 skipped = 0 for row in rows: @@ -990,7 +1132,7 @@ async def rotate_mcp_user_credentials_master_key(prisma_client: PrismaClient, ne skipped += 1 continue re_encrypted = encrypt_value_helper(plaintext, new_encryption_key=new_master_key) - await MCPUserCredentialsRepository(prisma_client).table.update( + await _mcp_user_credentials_table(prisma_client).update( where={ "user_id_server_id": { "user_id": row.user_id, @@ -1015,7 +1157,7 @@ async def rotate_mcp_user_env_vars_master_key(prisma_client: PrismaClient, new_m skipped so one corrupt row does not abort the rotation nor overwrite values that may still be recoverable. """ - rows = await prisma_client.db.litellm_mcpuserenvvars.find_many() + rows = await _mcp_user_env_vars_table(prisma_client).find_many() rotated = 0 skipped = 0 for row in rows: @@ -1034,7 +1176,7 @@ async def rotate_mcp_user_env_vars_master_key(prisma_client: PrismaClient, new_m skipped += 1 continue re_encrypted = encrypt_value_helper(plaintext, new_encryption_key=new_master_key) - await prisma_client.db.litellm_mcpuserenvvars.update( + await _mcp_user_env_vars_table(prisma_client).update( where={ "user_id_server_id": { "user_id": row.user_id, @@ -1060,7 +1202,7 @@ async def store_user_credential( """Store a user credential for a BYOK MCP server.""" encoded = encrypt_value_helper(credential) - await MCPUserCredentialsRepository(prisma_client).table.upsert( + await _mcp_user_credentials_table(prisma_client).upsert( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}}, data={ "create": { @@ -1080,7 +1222,7 @@ async def get_user_credential( ) -> Optional[str]: """Return credential for a user+server pair, or None.""" - row = await MCPUserCredentialsRepository(prisma_client).table.find_unique( + row = await _mcp_user_credentials_table(prisma_client).find_unique( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} ) if row is None: @@ -1094,7 +1236,7 @@ async def has_user_credential( server_id: str, ) -> bool: """Return True if the user has a stored credential for this server.""" - row = await MCPUserCredentialsRepository(prisma_client).table.find_unique( + row = await _mcp_user_credentials_table(prisma_client).find_unique( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} ) return row is not None @@ -1106,7 +1248,7 @@ async def delete_user_credential( server_id: str, ) -> None: """Delete the user's stored credential for a BYOK MCP server.""" - await MCPUserCredentialsRepository(prisma_client).table.delete( + await _mcp_user_credentials_table(prisma_client).delete( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} ) @@ -1135,7 +1277,7 @@ async def store_user_oauth_credential( if expires_in is not None: expires_at = (datetime.now(timezone.utc) + timedelta(seconds=expires_in)).isoformat() - payload: Dict[str, Any] = { + payload: dict[str, object] = { "type": "oauth2", "access_token": access_token, "connected_at": datetime.now(timezone.utc).isoformat(), @@ -1151,7 +1293,7 @@ async def store_user_oauth_credential( # Skip the guard when the caller knows the row is already an OAuth2 credential # (e.g. during token refresh), saving an extra DB round-trip. if not skip_byok_guard: - existing = await MCPUserCredentialsRepository(prisma_client).table.find_unique( + existing = await _mcp_user_credentials_table(prisma_client).find_unique( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} ) if existing is not None and _decode_oauth_payload(existing.credential_b64) is None: @@ -1166,7 +1308,7 @@ async def store_user_oauth_credential( ) encoded = encrypt_value_helper(json.dumps(payload)) - await MCPUserCredentialsRepository(prisma_client).table.upsert( + await _mcp_user_credentials_table(prisma_client).upsert( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}}, data={ "create": { @@ -1207,7 +1349,7 @@ async def get_user_oauth_credential( ) -> Optional[Dict[str, Any]]: """Return the decoded OAuth2 payload dict for a user+server pair, or None.""" - row = await MCPUserCredentialsRepository(prisma_client).table.find_unique( + row = await _mcp_user_credentials_table(prisma_client).find_unique( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} ) if row is None: @@ -1221,7 +1363,7 @@ async def list_user_oauth_credentials( ) -> List[Dict[str, Any]]: """Return all OAuth2 credential payloads for a user, tagged with server_id.""" - rows = await MCPUserCredentialsRepository(prisma_client).table.find_many(where={"user_id": user_id}) + rows = await _mcp_user_credentials_table(prisma_client).find_many(where={"user_id": user_id}) results: List[Dict[str, Any]] = [] for row in rows: payload = _decode_oauth_payload(row.credential_b64) @@ -1260,7 +1402,7 @@ def mcp_oauth_token_identity(server: object) -> tuple[object, ...]: creds = getattr(server, "credentials", None) if isinstance(creds, str): try: - parsed: object = json.loads(creds) + parsed: object = _loads_json_object(creds) except ValueError: parsed = None else: @@ -1302,12 +1444,12 @@ async def purge_user_oauth_credentials_for_server( invalidate_token_cache is injectable for tests; it defaults to the manager's shared invalidate_user_oauth_token_cache, the single invalidation point for per-user tokens.""" - repo = MCPUserCredentialsRepository(prisma_client) - rows = await repo.table.find_many(where={"server_id": server_id}) + repo = _mcp_user_credentials_table(prisma_client) + rows = await repo.find_many(where={"server_id": server_id}) oauth_rows = [row for row in rows if _decode_oauth_payload(row.credential_b64) is not None] if not oauth_rows: return 0 - deleted_count = await repo.table.delete_many( + deleted_count = await repo.delete_many( where={"server_id": server_id, "user_id": {"in": [row.user_id for row in oauth_rows]}} ) if invalidate_token_cache is None: @@ -1330,10 +1472,17 @@ async def purge_user_oauth_credentials_for_server( return deleted_count +class _OAuthTokenResponse(TypedDict, total=False): + access_token: str + refresh_token: str + expires_in: int + scope: str + + async def refresh_user_oauth_token( prisma_client: PrismaClient, user_id: str, - server: Any, + server: "MCPServer", cred: Dict[str, Any], ) -> Optional[Dict[str, Any]]: """Attempt to refresh a per-user OAuth2 token using its stored refresh_token. @@ -1384,7 +1533,9 @@ async def refresh_user_oauth_token( data=token_data, ) response.raise_for_status() - body: Dict[str, Any] = response.json() + body = cast( # cast-ok: the token endpoint returns the RFC 6749 token-response JSON object + "_OAuthTokenResponse", response.json() + ) except Exception as exc: verbose_proxy_logger.warning( "refresh_user_oauth_token: refresh request failed for user=%s server=%s: %s", @@ -1439,7 +1590,7 @@ async def refresh_user_oauth_token( async def resolve_valid_user_oauth_token( user_id: str, - server: Any, + server: "MCPServer", cred: Optional[Dict[str, Any]], prisma_client: Optional[PrismaClient] = None, ) -> Optional[Dict[str, Any]]: @@ -1568,7 +1719,7 @@ async def get_active_submitted_mcp_server_ids_for_user( if not user_id: return [] - rows = await MCPServerRepository(prisma_client).table.find_many( + rows = await _mcp_server_table(prisma_client).find_many( where={ "submitted_by": user_id, "approval_status": MCPApprovalStatus.active, @@ -1584,7 +1735,7 @@ async def approve_mcp_server( ) -> LiteLLM_MCPServerTable: """Set approval_status=active and record reviewed_at.""" now = datetime.now(timezone.utc) - updated = await MCPServerRepository(prisma_client).table.update( + updated = await _mcp_server_table(prisma_client).update( where={"server_id": server_id}, data={ "approval_status": MCPApprovalStatus.active, @@ -1592,7 +1743,7 @@ async def approve_mcp_server( "updated_by": touched_by, }, ) - table = LiteLLM_MCPServerTable(**updated.model_dump()) + table = LiteLLM_MCPServerTable.model_validate(updated.model_dump()) decrypt_global_env_var_values(table.env_vars) return table @@ -1605,18 +1756,18 @@ async def reject_mcp_server( ) -> LiteLLM_MCPServerTable: """Set approval_status=rejected, record reviewed_at and review_notes.""" now = datetime.now(timezone.utc) - data: Dict[str, Any] = { + data: dict[str, object] = { "approval_status": MCPApprovalStatus.rejected, "reviewed_at": now, "updated_by": touched_by, } if review_notes is not None: data["review_notes"] = review_notes - updated = await MCPServerRepository(prisma_client).table.update( + updated = await _mcp_server_table(prisma_client).update( where={"server_id": server_id}, data=data, ) - table = LiteLLM_MCPServerTable(**updated.model_dump()) + table = LiteLLM_MCPServerTable.model_validate(updated.model_dump()) decrypt_global_env_var_values(table.env_vars) return table @@ -1629,12 +1780,12 @@ async def get_mcp_submissions( along with a summary count breakdown by approval_status. Mirrors get_guardrail_submissions() from guardrail_endpoints.py. """ - rows = await MCPServerRepository(prisma_client).table.find_many( + rows = await _mcp_server_table(prisma_client).find_many( where={"submitted_at": {"not": None}}, order={"submitted_at": "desc"}, take=500, # safety cap; paginate if needed in a future iteration ) - items = [LiteLLM_MCPServerTable(**r.model_dump()) for r in rows] + items = [LiteLLM_MCPServerTable.model_validate(r.model_dump()) for r in rows] for item in items: decrypt_global_env_var_values(item.env_vars) @@ -1671,7 +1822,7 @@ def _decode_user_env_vars(stored: str) -> Dict[str, str]: ) return {} try: - parsed = json.loads(decrypted) + parsed = _loads_json_object(decrypted) except (ValueError, TypeError): return {} if not isinstance(parsed, dict): @@ -1685,7 +1836,7 @@ async def get_user_env_vars( server_id: str, ) -> Dict[str, str]: """Return the calling user's env var dict for ``server_id`` (empty if none).""" - row = await prisma_client.db.litellm_mcpuserenvvars.find_unique( + row = await _mcp_user_env_vars_table(prisma_client).find_unique( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} ) if row is None: @@ -1705,7 +1856,7 @@ async def get_user_env_vars_bulk( ids = list(server_ids) if not ids: return {} - rows = await prisma_client.db.litellm_mcpuserenvvars.find_many(where={"user_id": user_id, "server_id": {"in": ids}}) + rows = await _mcp_user_env_vars_table(prisma_client).find_many(where={"user_id": user_id, "server_id": {"in": ids}}) return {row.server_id: _decode_user_env_vars(row.values_b64) for row in rows} @@ -1730,15 +1881,18 @@ async def merge_user_env_vars( "big", signed=True, ) - async with prisma_client.db.tx() as tx: + async with _prisma_db(prisma_client).tx() as tx: await tx.execute_raw("SELECT pg_advisory_xact_lock($1::bigint)", lock_key) - row = await tx.litellm_mcpuserenvvars.find_unique( + tx_env_vars_table = cast( # cast-ok: the generated table client methods reference types pyright cannot load; narrow to the surface used here + "_PrismaTable[PrismaMCPUserEnvVarsRow]", tx.litellm_mcpuserenvvars + ) + row = await tx_env_vars_table.find_unique( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} ) existing = _decode_user_env_vars(row.values_b64) if row is not None else {} merged = {k: v for k, v in {**existing, **updates}.items() if k in allowed} encoded = encrypt_value_helper(json.dumps(merged)) - await tx.litellm_mcpuserenvvars.upsert( + await tx_env_vars_table.upsert( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}}, data={ "create": { @@ -1762,4 +1916,4 @@ async def delete_user_env_vars( Uses ``delete_many`` so a missing row is a no-op; real DB errors still propagate to the caller instead of being silently swallowed. """ - await prisma_client.db.litellm_mcpuserenvvars.delete_many(where={"user_id": user_id, "server_id": server_id}) + await _mcp_user_env_vars_table(prisma_client).delete_many(where={"user_id": user_id, "server_id": server_id}) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index ce82ca74267..2719418ac82 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -13,7 +13,24 @@ import math import re import time -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Type, Union, cast +from datetime import datetime +from typing import ( + TYPE_CHECKING, + Any, + Dict, + Iterable, + List, + Literal, + Mapping, + Optional, + Protocol, + Sequence, + Type, + TypedDict, + TypeVar, + Union, + cast, +) from fastapi import HTTPException, Request, status from pydantic import BaseModel @@ -119,6 +136,393 @@ all_routes = LiteLLMRoutes.openai_routes.value + LiteLLMRoutes.management_routes.value +class _BudgetRowFields(TypedDict): + budget_id: str | None + soft_budget: float | None + max_budget: float | None + max_parallel_requests: int | None + tpm_limit: int | None + rpm_limit: int | None + model_max_budget: Optional[ + dict[str, object] + ] # mutable-ok: splatted into LiteLLM_BudgetTable, which declares this field as a dict + budget_duration: str | None + allowed_models: Optional[ + list[str] + ] # mutable-ok: splatted into LiteLLM_BudgetTable, which declares this field as a list + + +class _EndUserRowFields(TypedDict): + user_id: str + blocked: bool + alias: str | None + spend: float + allowed_model_region: Literal["eu", "us"] | None + default_model: str | None + budget_id: str | None + + +class _TagRowFields(TypedDict): + tag_name: str + description: str | None + models: list[str] # mutable-ok: splatted into LiteLLM_TagTable, which declares this field as a list + model_info: Optional[ + dict[str, object] + ] # mutable-ok: splatted into LiteLLM_TagTable, which declares this field as a dict + spend: float + budget_id: str | None + created_at: datetime | None + created_by: str | None + updated_at: datetime | None + + +class _TeamMembershipRowFields(TypedDict): + user_id: str + team_id: str + budget_id: str | None + spend: float | None + + +class _OrgMembershipRowFields(TypedDict): + user_id: str + organization_id: str + user_role: str | None + spend: float + budget_id: str | None + created_at: datetime + updated_at: datetime + user_email: str | None + + +class _UserRowFields(TypedDict): + user_id: str + user_alias: str | None + team_id: str | None + sso_user_id: str | None + organization_id: str | None + object_permission_id: str | None + teams: list[str] # mutable-ok: splatted into LiteLLM_UserTable, which declares this field as a list + user_role: str | None + max_budget: float | None + spend: float + user_email: str | None + models: list[str] # mutable-ok: splatted into LiteLLM_UserTable, which declares this field as a list + metadata: Optional[ + dict[str, object] + ] # mutable-ok: splatted into LiteLLM_UserTable, which declares this field as a dict + max_parallel_requests: int | None + tpm_limit: int | None + rpm_limit: int | None + budget_duration: str | None + budget_reset_at: datetime | None + allowed_cache_controls: list[ + str + ] # mutable-ok: splatted into LiteLLM_UserTable, which declares this field as a list + policies: list[str] # mutable-ok: splatted into LiteLLM_UserTable, which declares this field as a list + model_spend: Optional[ + dict[str, object] + ] # mutable-ok: splatted into LiteLLM_UserTable, which declares this field as a dict + model_max_budget: Optional[ + dict[str, object] + ] # mutable-ok: splatted into LiteLLM_UserTable, which declares this field as a dict + created_at: datetime | None + updated_at: datetime | None + + +class _TeamRowFields(TypedDict): + team_alias: str | None + team_id: str + organization_id: str | None + admins: list[str] # mutable-ok: splatted into LiteLLM_TeamTable, which declares this field as a list + members: list[str] # mutable-ok: splatted into LiteLLM_TeamTable, which declares this field as a list + team_member_permissions: Optional[ + list[str] + ] # mutable-ok: splatted into LiteLLM_TeamTable, which declares this field as a list + metadata: Optional[ + dict[str, object] + ] # mutable-ok: splatted into LiteLLM_TeamTable, which declares this field as a dict + tpm_limit: int | None + rpm_limit: int | None + max_budget: float | None + soft_budget: float | None + budget_duration: str | None + models: list[str] # mutable-ok: splatted into LiteLLM_TeamTable, which declares this field as a list + blocked: bool + router_settings: Optional[ + dict[str, object] + ] # mutable-ok: splatted into LiteLLM_TeamTable, which declares this field as a dict + access_group_ids: Optional[ + list[str] + ] # mutable-ok: splatted into LiteLLM_TeamTable, which declares this field as a list + default_team_member_models: Optional[ + list[str] + ] # mutable-ok: splatted into LiteLLM_TeamTable, which declares this field as a list + spend: float | None + max_parallel_requests: int | None + budget_reset_at: datetime | None + model_id: int | None + model_spend: Optional[ + dict[str, object] + ] # mutable-ok: splatted into LiteLLM_TeamTable, which declares this field as a dict + model_max_budget: Optional[ + dict[str, object] + ] # mutable-ok: splatted into LiteLLM_TeamTable, which declares this field as a dict + policies: list[str] | None # mutable-ok: splatted into LiteLLM_TeamTable, which declares this field as a list + allow_team_guardrail_config: bool | None + object_permission_id: str | None + updated_at: datetime | None + created_at: datetime | None + + +class _AccessGroupRowFields(TypedDict): + access_group_id: str + access_group_name: str + description: str | None + access_model_names: list[ + str + ] # mutable-ok: splatted into LiteLLM_AccessGroupTable, which declares this field as a list + access_mcp_server_ids: list[ + str + ] # mutable-ok: splatted into LiteLLM_AccessGroupTable, which declares this field as a list + access_agent_ids: list[ + str + ] # mutable-ok: splatted into LiteLLM_AccessGroupTable, which declares this field as a list + assigned_team_ids: list[ + str + ] # mutable-ok: splatted into LiteLLM_AccessGroupTable, which declares this field as a list + assigned_key_ids: list[ + str + ] # mutable-ok: splatted into LiteLLM_AccessGroupTable, which declares this field as a list + created_at: datetime | None + created_by: str | None + updated_at: datetime | None + updated_by: str | None + + +class _OrgRowFields(TypedDict): + organization_id: str | None + organization_alias: str | None + budget_id: str + spend: float + metadata: Optional[ + dict[str, object] + ] # mutable-ok: splatted into LiteLLM_OrganizationTable, which declares this field as a dict + models: list[str] # mutable-ok: splatted into LiteLLM_OrganizationTable, which declares this field as a list + model_spend: Optional[ + dict[str, object] + ] # mutable-ok: splatted into LiteLLM_OrganizationTable, which declares this field as a dict + created_by: str + updated_by: str + object_permission_id: str | None + + +class _ObjectPermissionRowFields(TypedDict): + object_permission_id: str + mcp_servers: Optional[ + list[str] + ] # mutable-ok: splatted into LiteLLM_ObjectPermissionTable, which declares this field as a list + mcp_access_groups: Optional[ + list[str] + ] # mutable-ok: splatted into LiteLLM_ObjectPermissionTable, which declares this field as a list + mcp_tool_permissions: Optional[ + dict[str, list[str]] + ] # mutable-ok: splatted into LiteLLM_ObjectPermissionTable, which declares this field as a dict + vector_stores: Optional[ + list[str] + ] # mutable-ok: splatted into LiteLLM_ObjectPermissionTable, which declares this field as a list + agents: Optional[ + list[str] + ] # mutable-ok: splatted into LiteLLM_ObjectPermissionTable, which declares this field as a list + agent_access_groups: Optional[ + list[str] + ] # mutable-ok: splatted into LiteLLM_ObjectPermissionTable, which declares this field as a list + models: Optional[ + list[str] + ] # mutable-ok: splatted into LiteLLM_ObjectPermissionTable, which declares this field as a list + mcp_toolsets: Optional[ + list[str] + ] # mutable-ok: splatted into LiteLLM_ObjectPermissionTable, which declares this field as a list + blocked_tools: Optional[ + list[str] + ] # mutable-ok: splatted into LiteLLM_ObjectPermissionTable, which declares this field as a list + search_tools: Optional[ + list[str] + ] # mutable-ok: splatted into LiteLLM_ObjectPermissionTable, which declares this field as a list + mcp_tool_search_enabled: bool | None + + +class _VectorStoreRowFields(TypedDict): + vector_store_id: str + custom_llm_provider: str + vector_store_name: str | None + vector_store_description: str | None + vector_store_metadata: Optional[ + dict[str, object] + ] # mutable-ok: splatted into LiteLLM_ManagedVectorStoresTable, which declares this field as a dict + created_at: datetime | None + updated_at: datetime | None + litellm_credential_name: str | None + litellm_params: Optional[ + dict[str, object] + ] # mutable-ok: splatted into LiteLLM_ManagedVectorStoresTable, which declares this field as a dict + team_id: str | None + user_id: str | None + + +class _ProjectRowFields(TypedDict): + project_id: str + project_alias: str | None + description: str | None + team_id: str | None + budget_id: str | None + metadata: Optional[ + dict[str, object] + ] # mutable-ok: splatted into LiteLLM_ProjectTableCachedObj, which declares this field as a dict + models: list[str] # mutable-ok: splatted into LiteLLM_ProjectTableCachedObj, which declares this field as a list + spend: float + model_spend: Optional[ + dict[str, object] + ] # mutable-ok: splatted into LiteLLM_ProjectTableCachedObj, which declares this field as a dict + model_rpm_limit: Optional[ + dict[str, object] + ] # mutable-ok: splatted into LiteLLM_ProjectTableCachedObj, which declares this field as a dict + model_tpm_limit: Optional[ + dict[str, object] + ] # mutable-ok: splatted into LiteLLM_ProjectTableCachedObj, which declares this field as a dict + blocked: bool + object_permission_id: str | None + created_by: str | None + updated_by: str | None + created_at: datetime | None + updated_at: datetime | None + + +class _UserAPIKeyAuthFields(TypedDict, total=False): + token: str | None + key_name: str | None + key_alias: str | None + spend: float + max_budget: float | None + user_id: str | None + team_id: str | None + models: list[str] # mutable-ok: splatted into UserAPIKeyAuth, which declares this field as a list + budget_duration: str | None + metadata: dict[str, object] # mutable-ok: splatted into UserAPIKeyAuth, which declares this field as a dict + + +class _OrgQueryKwargs(TypedDict, total=False): + where: Mapping[str, object] + include: Mapping[str, object] + + +class _BudgetRow(Protocol): + def dict(self) -> _BudgetRowFields: ... + + +class _EndUserRow(Protocol): + def dict(self) -> _EndUserRowFields: ... + + +class _TagRow(Protocol): + tag_name: str + + def dict(self) -> _TagRowFields: ... + + +class _TeamMembershipRow(Protocol): + def dict(self) -> _TeamMembershipRowFields: ... + + +class _OrgMembershipRow(Protocol): + def model_dump(self) -> _OrgMembershipRowFields: ... + + +class _UserRow(Protocol): + user_id: str + + @property + def organization_memberships(self) -> Sequence[_OrgMembershipRow | None] | None: ... + + @organization_memberships.setter + def organization_memberships(self, value: Sequence[LiteLLM_OrganizationMembershipTable] | None) -> None: ... + + def keys(self) -> Iterable[str]: ... + + def __getitem__(self, key: str) -> object: ... + + +class _TeamRow(Protocol): + def dict(self) -> _TeamRowFields: ... + + def model_dump(self) -> _TeamRowFields: ... + + +class _AccessGroupRow(Protocol): + def dict(self) -> _AccessGroupRowFields: ... + + +class _JWTKeyMappingRow(Protocol): + token: str + + +class _OrgRow(Protocol): + def model_dump(self) -> _OrgRowFields: ... + + +class _ObjectPermissionRow(Protocol): + def dict(self) -> _ObjectPermissionRowFields: ... + + +class _VectorStoreRow(Protocol): + def model_dump(self) -> _VectorStoreRowFields: ... + + def dict(self) -> _VectorStoreRowFields: ... + + def keys(self) -> Iterable[str]: ... + + def __getitem__(self, key: str) -> object: ... + + +class _ProjectRow(Protocol): + def model_dump(self) -> _ProjectRowFields: ... + + +_PrismaRowT_co = TypeVar("_PrismaRowT_co", covariant=True) + + +class _PrismaTable(Protocol[_PrismaRowT_co]): + async def find_unique( + self, + where: Mapping[str, object] | None = None, + include: Mapping[str, object] | None = None, + ) -> _PrismaRowT_co | None: ... + + async def find_first( + self, + where: Mapping[str, object] | None = None, + include: Mapping[str, object] | None = None, + ) -> _PrismaRowT_co | None: ... + + async def find_many( + self, + where: Mapping[str, object] | None = None, + include: Mapping[str, object] | None = None, + take: int | None = None, + ) -> Sequence[_PrismaRowT_co]: ... + + async def create( + self, + data: Mapping[str, object], + include: Mapping[str, object] | None = None, + ) -> _PrismaRowT_co: ... + + async def update( + self, + where: Mapping[str, object], + data: Mapping[str, object], + ) -> _PrismaRowT_co | None: ... + + def _log_budget_lookup_failure(entity: str, error: Exception) -> None: """ Log a warning when budget lookup fails; cache will not be populated. @@ -377,7 +781,7 @@ def _global_proxy_budget_check(global_proxy_spend: Optional[float], skip_budget_ ) -def _guardrail_modification_check(request_body: dict, team_object: Optional[LiteLLM_TeamTable]) -> None: +def _guardrail_modification_check(request_body: Mapping[str, object], team_object: LiteLLM_TeamTable | None) -> None: """ Reject user-supplied metadata flags that would modify guardrail behavior unless the team has explicit permission. Checked keys include the plural @@ -392,7 +796,7 @@ def _guardrail_modification_check(request_body: dict, team_object: Optional[Lite """ from litellm.proxy.guardrails.guardrail_helpers import can_modify_guardrails - def _coerce_to_dict(container: Any) -> Optional[dict]: + def _coerce_to_dict(container: object) -> dict | None: """Accept dict or JSON-string (from multipart/form-data or extra_body). Without this, an attacker can smuggle guardrail keys past the check by @@ -404,11 +808,11 @@ def _coerce_to_dict(container: Any) -> Optional[dict]: if isinstance(container, dict): return container if isinstance(container, str): - parsed = safe_json_loads(container) + parsed = cast("object", safe_json_loads(container)) # cast-ok: safe_json_loads is untyped return parsed if isinstance(parsed, dict) else None return None - def _user_requested_modification(container: Any) -> bool: + def _user_requested_modification(container: object) -> bool: coerced = _coerce_to_dict(container) if coerced is None: return False @@ -716,7 +1120,10 @@ async def _user_max_budget_check() -> None: _enforce_user_param_check(general_settings, request, request_body, route) _global_proxy_budget_check(global_proxy_spend, skip_budget_checks, route) - _guardrail_modification_check(request_body, team_object) + _guardrail_modification_check( + cast("dict[str, object]", request_body), # cast-ok: request body dicts are string-keyed JSON payloads + team_object, + ) # 10 [OPTIONAL] Organization RBAC checks organization_role_based_access_check(user_object=user_object, route=route, request_body=request_body) @@ -943,9 +1350,9 @@ async def get_default_end_user_budget( # Fetch from database try: - budget_record = await BudgetRepository(prisma_client).table.find_unique( - where={"budget_id": litellm.max_end_user_budget_id} - ) + budget_record = await cast( # cast-ok: repository .table returns an untyped prisma client + "_PrismaTable[_BudgetRow]", BudgetRepository(prisma_client).table + ).find_unique(where={"budget_id": litellm.max_end_user_budget_id}) if budget_record is None: verbose_proxy_logger.warning( @@ -995,14 +1402,19 @@ async def get_team_member_default_budget( cache_key = f"team_member_default_budget:{budget_id}" - cached_budget = await user_api_key_cache.async_get_cache(key=cache_key) + cached_budget = cast( # cast-ok: untyped cache returns the budget model or its serialized dict + "LiteLLM_BudgetTable | _BudgetRowFields | None", + await user_api_key_cache.async_get_cache(key=cache_key), + ) if isinstance(cached_budget, LiteLLM_BudgetTable): return cached_budget if isinstance(cached_budget, dict): return LiteLLM_BudgetTable(**cached_budget) try: - budget_record = await BudgetRepository(prisma_client).table.find_unique(where={"budget_id": budget_id}) + budget_record = await cast( # cast-ok: repository .table returns an untyped prisma client + "_PrismaTable[_BudgetRow]", BudgetRepository(prisma_client).table + ).find_unique(where={"budget_id": budget_id}) if budget_record is None: verbose_proxy_logger.warning(f"Team-default member budget not found in database: {budget_id}") @@ -1159,7 +1571,9 @@ async def get_end_user_object( # Fetch from database try: - response = await EndUserRepository(prisma_client).table.find_unique( + response = await cast( # cast-ok: repository .table returns an untyped prisma client + "_PrismaTable[_EndUserRow]", EndUserRepository(prisma_client).table + ).find_unique( where={"user_id": end_user_id}, include={"litellm_budget_table": True, "object_permission": True}, ) @@ -1231,7 +1645,9 @@ async def resolve_and_validate_end_user_id( return raw_end_user_id cache_key = f"end_user_validation:{raw_end_user_id}" - cached = await user_api_key_cache.async_get_cache(key=cache_key) + cached = cast( # cast-ok: untyped cache returns the validation marker string + "object", await user_api_key_cache.async_get_cache(key=cache_key) + ) if cached == "valid": return raw_end_user_id if cached == "invalid": @@ -1333,8 +1749,10 @@ async def get_tag_objects_batch( if not tag_names: return {} - tag_objects = {} - uncached_tags = [] + tag_objects: dict[ + str, LiteLLM_TagTable + ] = {} # mutable-ok: filled per tag_name below and returned as the batch result + uncached_tags: list[str] = [] # mutable-ok: appended to in the cache-miss loop below # Try to get all tags from cache first for tag_name in tag_names: @@ -1351,7 +1769,9 @@ async def get_tag_objects_batch( # Batch fetch uncached tags from DB in one query if uncached_tags: try: - db_tags = await TagRepository(prisma_client).table.find_many( + db_tags = await cast( # cast-ok: repository .table returns an untyped prisma client + "_PrismaTable[_TagRow]", TagRepository(prisma_client).table + ).find_many( where={"tag_name": {"in": uncached_tags}}, include={"litellm_budget_table": True}, ) @@ -1445,7 +1865,9 @@ async def get_team_membership( # else, check db try: - response = await TeamMembershipRepository(prisma_client).table.find_unique( + response = await cast( # cast-ok: repository .table returns an untyped prisma client + "_PrismaTable[_TeamMembershipRow]", TeamMembershipRepository(prisma_client).table + ).find_unique( where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}, include={"litellm_budget_table": True}, ) @@ -1514,7 +1936,7 @@ def _should_check_db(key: str, last_db_access_time: LimitedSizeOrderedDict, db_c return False -def _update_last_db_access_time(key: str, value: Optional[Any], last_db_access_time: LimitedSizeOrderedDict): +def _update_last_db_access_time(key: str, value: object | None, last_db_access_time: LimitedSizeOrderedDict): last_db_access_time[key] = (value, time.time()) @@ -1535,7 +1957,9 @@ def _get_role_based_permissions( for role_based_permission in role_based_permissions: if role_based_permission.role == rbac_role: - return getattr(role_based_permission, key) + return cast( # cast-ok: key is "models" or "routes", both Optional[List[str]] attributes + "list[str] | None", getattr(role_based_permission, key) + ) return None @@ -1576,7 +2000,7 @@ async def _get_fuzzy_user_object( prisma_client: PrismaClient, sso_user_id: Optional[str] = None, user_email: Optional[str] = None, -) -> Optional[LiteLLM_UserTable]: +) -> Optional["_UserRow"]: """ Checks if sso user is in db. @@ -1590,7 +2014,9 @@ async def _get_fuzzy_user_object( response = None if sso_user_id is not None: - response = await UserRepository(prisma_client).table.find_unique( + response = await cast( # cast-ok: repository .table returns an untyped prisma client + "_PrismaTable[_UserRow]", UserRepository(prisma_client).table + ).find_unique( where={"sso_user_id": sso_user_id}, include={"organization_memberships": True}, ) @@ -1598,14 +2024,18 @@ async def _get_fuzzy_user_object( if response is None and user_email is not None: # Use case-insensitive query to handle emails with different casing # This matches the pattern used in _check_duplicate_user_email - response = await UserRepository(prisma_client).table.find_first( + response = await cast( # cast-ok: repository .table returns an untyped prisma client + "_PrismaTable[_UserRow]", UserRepository(prisma_client).table + ).find_first( where={"user_email": {"equals": user_email, "mode": "insensitive"}}, include={"organization_memberships": True}, ) if response is not None and sso_user_id is not None: # update sso_user_id asyncio.create_task( # background task to update user with sso id - UserRepository(prisma_client).table.update( + cast( # cast-ok: repository .table returns an untyped prisma client + "_PrismaTable[_UserRow]", UserRepository(prisma_client).table + ).update( where={"user_id": response.user_id}, data={"sso_user_id": sso_user_id}, ) @@ -1655,9 +2085,9 @@ async def get_user_object( ) if should_check_db: - response = await UserRepository(prisma_client).table.find_unique( - where={"user_id": user_id}, include={"organization_memberships": True} - ) + response = await cast( # cast-ok: repository .table returns an untyped prisma client + "_PrismaTable[_UserRow]", UserRepository(prisma_client).table + ).find_unique(where={"user_id": user_id}, include={"organization_memberships": True}) if response is None: response = await _get_fuzzy_user_object( @@ -1671,7 +2101,7 @@ async def get_user_object( if response is None: if user_id_upsert: - new_user_params: Dict[str, Any] = { + new_user_params: dict[str, object] = { "user_id": user_id, } if user_email is not None: @@ -1683,10 +2113,14 @@ async def get_user_object( and new_user_params.get("budget_reset_at") is None ): new_user_params["budget_reset_at"] = get_budget_reset_time( - budget_duration=new_user_params["budget_duration"] + budget_duration=cast( # cast-ok: guarded non-None above; budget durations are strings + "str", new_user_params["budget_duration"] + ) ) - response = await UserRepository(prisma_client).table.create( + response = await cast( # cast-ok: repository .table returns an untyped prisma client + "_PrismaTable[_UserRow]", UserRepository(prisma_client).table + ).create( data=new_user_params, include={"organization_memberships": True}, ) @@ -1708,7 +2142,7 @@ async def get_user_object( ] response.organization_memberships = _dumped_memberships - _response = LiteLLM_UserTable(**dict(response)) + _response = LiteLLM_UserTable(**cast("_UserRowFields", dict(response))) # cast-ok: prisma row field mapping response_dict = _response.model_dump() # save the user object to cache @@ -1736,7 +2170,7 @@ async def get_user_object( async def _cache_management_object( key: str, - value: Union[BaseModel, Dict[str, Any]], + value: BaseModel | dict[str, object], user_api_key_cache: UserApiKeyCache, proxy_logging_obj: Optional[ProxyLogging], *, @@ -1830,7 +2264,9 @@ async def _delete_cache_key_object( @log_db_metrics async def _get_team_db_check(team_id: str, prisma_client: PrismaClient, team_id_upsert: Optional[bool] = None): - response = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id}) + response = await cast( # cast-ok: repository .table returns an untyped prisma client + "_PrismaTable[_TeamRow]", TeamRepository(prisma_client).table + ).find_unique(where={"team_id": team_id}) if response is None and team_id_upsert: from litellm.proxy.management_endpoints.team_endpoints import new_team @@ -1845,12 +2281,17 @@ async def _get_team_db_check(team_id: str, prisma_client: PrismaClient, team_id_ http_request=mock_request, user_api_key_dict=system_admin_user, ) - response = LiteLLM_TeamTable(**created_team_dict) + response = cast( # cast-ok: the validated team model exposes the same row surface + "_TeamRow", + LiteLLM_TeamTable(**cast("_TeamRowFields", created_team_dict)), # cast-ok: new_team returns an untyped dict + ) return response async def _get_team_object_from_db(team_id: str, prisma_client: PrismaClient): - return await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id}) + return await cast( # cast-ok: repository .table returns an untyped prisma client + "_PrismaTable[_TeamRow]", TeamRepository(prisma_client).table + ).find_unique(where={"team_id": team_id}) async def _get_team_object_from_user_api_key_cache( @@ -2058,9 +2499,9 @@ async def get_access_object( # Not in cache - fetch from DB try: - response = await AccessGroupRepository(prisma_client).table.find_unique( - where={"access_group_id": access_group_id} - ) + response = await cast( # cast-ok: repository .table returns an untyped prisma client + "_PrismaTable[_AccessGroupRow]", AccessGroupRepository(prisma_client).table + ).find_unique(where={"access_group_id": access_group_id}) if response is None: raise HTTPException( @@ -2134,7 +2575,9 @@ async def get_team_object_by_alias( # Query database by team_alias try: - teams = await TeamRepository(prisma_client).table.find_many(where={"team_alias": team_alias}) + teams = await cast( # cast-ok: repository .table returns an untyped prisma client + "_PrismaTable[_TeamRow]", TeamRepository(prisma_client).table + ).find_many(where={"team_alias": team_alias}) if not teams: raise HTTPException( @@ -2236,7 +2679,9 @@ async def get_org_object_by_alias( # Query database by organization_alias try: - orgs = await OrganizationRepository(prisma_client).table.find_many(where={"organization_alias": org_alias}) + orgs = await cast( # cast-ok: repository .table returns an untyped prisma client + "_PrismaTable[_OrgRow]", OrganizationRepository(prisma_client).table + ).find_many(where={"organization_alias": org_alias}) if not orgs: raise HTTPException( @@ -2396,7 +2841,9 @@ def get_key_object_from_ui_hash_key( if decrypted_token is None: return None try: - return UserAPIKeyAuth(**json.loads(decrypted_token)) + return UserAPIKeyAuth( + **cast("_UserAPIKeyAuthFields", json.loads(decrypted_token)) # cast-ok: decrypted UI session payload + ) except Exception as e: raise Exception(f"Invalid hash key. Hash key={hashed_token}. Decrypted token={decrypted_token}. Error: {e}") @@ -2453,7 +2900,9 @@ async def get_jwt_key_mapping_object( Returns the hashed token (str) if a matching active mapping is found, else None. """ - mapping = await JWTKeyMappingRepository(prisma_client).table.find_first( + mapping = await cast( # cast-ok: repository .table returns an untyped prisma client + "_PrismaTable[_JWTKeyMappingRow]", JWTKeyMappingRepository(prisma_client).table + ).find_first( where={ "jwt_claim_name": jwt_claim_name, "jwt_claim_value": jwt_claim_value, @@ -2515,7 +2964,9 @@ async def get_key_object( code=status.HTTP_401_UNAUTHORIZED, ) - _response = UserAPIKeyAuth(**_valid_token.model_dump(exclude_none=True)) + _response = UserAPIKeyAuth( + **cast("_UserAPIKeyAuthFields", _valid_token.model_dump(exclude_none=True)) # cast-ok: combined view dump + ) # Load object_permission if object_permission_id exists but object_permission is not loaded if _response.object_permission_id and not _response.object_permission: @@ -2581,9 +3032,9 @@ async def get_object_permission( # else, check db try: - response = await ObjectPermissionRepository(prisma_client).table.find_unique( - where={"object_permission_id": object_permission_id} - ) + response = await cast( # cast-ok: repository .table returns an untyped prisma client + "_PrismaTable[_ObjectPermissionRow]", ObjectPermissionRepository(prisma_client).table + ).find_unique(where={"object_permission_id": object_permission_id}) if response is None: return None @@ -2637,7 +3088,9 @@ async def get_managed_vector_store_rows_by_uuids( if not cache_misses: return result - rows = await ManagedVectorStoresRepository(prisma_client).table.find_many( + rows = await cast( # cast-ok: repository .table returns an untyped prisma client + "_PrismaTable[_VectorStoreRow]", ManagedVectorStoresRepository(prisma_client).table + ).find_many( where={"vector_store_id": {"in": cache_misses}}, take=len(cache_misses), ) @@ -2648,7 +3101,9 @@ async def get_managed_vector_store_rows_by_uuids( row_dict = dict(row) if hasattr(row, "__dict__") else {} if not row_dict: continue - cached_obj = LiteLLM_ManagedVectorStoresTable(**row_dict) + cached_obj = LiteLLM_ManagedVectorStoresTable( + **cast("_VectorStoreRowFields", row_dict) # cast-ok: prisma row field mapping + ) key = "managed_vector_store_id:{}".format(cached_obj.vector_store_id) await user_api_key_cache.async_set_cache( key=key, @@ -2711,11 +3166,13 @@ async def get_org_object( return deserialized_org # else, check db try: - query_kwargs: Dict[str, Any] = {"where": {"organization_id": org_id}} + query_kwargs: _OrgQueryKwargs = {"where": {"organization_id": org_id}} if include_budget_table: query_kwargs["include"] = {"litellm_budget_table": True} - response = await OrganizationRepository(prisma_client).table.find_unique(**query_kwargs) + response = await cast( # cast-ok: repository .table returns an untyped prisma client + "_PrismaTable[_OrgRow]", OrganizationRepository(prisma_client).table + ).find_unique(**query_kwargs) except Exception: # An operational failure (DB down, timeout, cache fault) is NOT the same fact as a confirmed # missing row, and relabelling it as "doesn't exist" made every caller unable to tell them @@ -3670,17 +4127,18 @@ async def _virtual_key_soft_budget_check( ) -def _parse_email_list(raw: Any) -> List[str]: +def _parse_email_list(raw: object) -> list[str]: """Parse emails from a list or comma-separated string.""" if isinstance(raw, list): - return [e.strip() for e in raw if isinstance(e, str) and e.strip()] + entries = cast("list[object]", raw) # cast-ok: narrowed to list; elements validated below + return [e.strip() for e in entries if isinstance(e, str) and e.strip()] elif isinstance(raw, str): return [e.strip() for e in raw.split(",") if e.strip()] return [] def _normalize_alert_emails( - cfg: Optional[Dict[str, Any]], + cfg: Mapping[str, object] | None, ) -> Dict[str, List[str]]: """Coerce user-supplied threshold→recipients mapping to Dict[str, List[str]]. @@ -3693,8 +4151,8 @@ def _normalize_alert_emails( def _merge_budget_alert_email_configs( - global_cfg: Optional[Dict[str, Any]], - per_key_cfg: Optional[Dict[str, Any]], + global_cfg: Mapping[str, object] | None, + per_key_cfg: Mapping[str, object] | None, ) -> Optional[Dict[str, List[str]]]: """ Per-threshold additive merge: each threshold's recipient list is the union @@ -4197,7 +4655,9 @@ async def get_project_object( return deserialized_project # Fetch from DB - project_row = await ProjectRepository(prisma_client).table.find_unique( + project_row = await cast( # cast-ok: repository .table returns an untyped prisma client + "_PrismaTable[_ProjectRow]", ProjectRepository(prisma_client).table + ).find_unique( where={"project_id": project_id}, include={"litellm_budget_table": True}, ) @@ -4498,7 +4958,9 @@ async def vector_store_access_check( ######################################################### # Check if the key can access the vector store if valid_token is not None and valid_token.object_permission_id is not None: - key_object_permission = await ObjectPermissionRepository(prisma_client).table.find_unique( + key_object_permission = await cast( # cast-ok: repository .table returns an untyped prisma client + "_PrismaTable[LiteLLM_ObjectPermissionTable]", ObjectPermissionRepository(prisma_client).table + ).find_unique( where={"object_permission_id": valid_token.object_permission_id}, ) if key_object_permission is not None: @@ -4510,7 +4972,9 @@ async def vector_store_access_check( # Check if the team can access the vector store if team_object is not None and team_object.object_permission_id is not None: - team_object_permission = await ObjectPermissionRepository(prisma_client).table.find_unique( + team_object_permission = await cast( # cast-ok: repository .table returns an untyped prisma client + "_PrismaTable[LiteLLM_ObjectPermissionTable]", ObjectPermissionRepository(prisma_client).table + ).find_unique( where={"object_permission_id": team_object.object_permission_id}, ) if team_object_permission is not None: diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index a5ecf4e7f93..ee285bad5df 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -1,9 +1,22 @@ import asyncio from datetime import datetime from types import SimpleNamespace -from typing import Any, Awaitable, Callable, Dict, List, Optional, Set, Tuple, Union +from typing import ( + Awaitable, + Callable, + Dict, + List, + Mapping, + Optional, + Protocol, + Set, + Tuple, + Union, + cast, +) from fastapi import HTTPException, status +from typing_extensions import TypedDict from litellm._logging import verbose_proxy_logger from litellm.proxy._types import CommonProxyErrors @@ -16,6 +29,7 @@ BreakdownMetrics, DailySpendData, DailySpendMetadata, + GroupedData, KeyMetadata, KeyMetricWithMetadata, MetricWithMetadata, @@ -33,8 +47,152 @@ "litellm_dailytagspend": "LiteLLM_DailyTagSpend", } +_WhereFilter = dict[str, str | list[str] | dict[str, list[str]]] +_WhereConditions = dict[str, str | _WhereFilter] + + +class _SpendMetricsRecord(Protocol): + @property + def spend(self) -> float | None: ... + + @property + def prompt_tokens(self) -> int | None: ... + + @property + def completion_tokens(self) -> int | None: ... + + @property + def cache_read_input_tokens(self) -> int | None: ... + + @property + def cache_creation_input_tokens(self) -> int | None: ... + + @property + def compression_saved_tokens(self) -> int | None: ... + + @property + def compression_savings_spend(self) -> float | None: ... + + @property + def prompt_caching_savings_spend(self) -> float | None: ... + + @property + def api_requests(self) -> int | None: ... + + @property + def successful_requests(self) -> int | None: ... + + @property + def failed_requests(self) -> int | None: ... + + +class _DailySpendRecord(_SpendMetricsRecord, Protocol): + @property + def spend(self) -> float: ... + + @property + def date(self) -> str: ... + + @property + def api_key(self) -> str: ... + + @property + def model(self) -> str | None: ... + + @property + def model_group(self) -> str | None: ... + + @property + def custom_llm_provider(self) -> str | None: ... + + @property + def mcp_namespaced_tool_name(self) -> str | None: ... + + @property + def endpoint(self) -> str | None: ... + + +class _GroupingSetsRow(_SpendMetricsRecord, Protocol): + @property + def group_level(self) -> int: ... + + @property + def date(self) -> str: ... + + @property + def api_key(self) -> str | None: ... + + @property + def model(self) -> str | None: ... + + @property + def model_group(self) -> str | None: ... -def update_metrics(existing_metrics: SpendMetrics, record: Any) -> SpendMetrics: + @property + def custom_llm_provider(self) -> str | None: ... + + @property + def mcp_namespaced_tool_name(self) -> str | None: ... + + @property + def endpoint(self) -> str | None: ... + + +class _TokenMetadataRecord(Protocol): + @property + def token(self) -> str: ... + + @property + def key_alias(self) -> str | None: ... + + @property + def team_id(self) -> str | None: ... + + +class _TokenTable(Protocol): + def find_many( + self, + where: Mapping[str, object], + order: Mapping[str, str] | None = None, + ) -> Awaitable[list[_TokenMetadataRecord]]: ... + + +class _TokenTableRepository(Protocol): + @property + def table(self) -> _TokenTable: ... + + +class _RawQueryDb(Protocol): + def query_raw(self, query: str, *params: str) -> Awaitable[list[dict[str, object]] | None]: ... + + +def _token_table(repository: _TokenTableRepository) -> _TokenTable: + return repository.table + + +class _DailySpendTable(Protocol): + def count(self, where: Mapping[str, object]) -> Awaitable[int]: ... + + def find_many( + self, + where: Mapping[str, object], + order: list[Mapping[str, str]], + skip: int, + take: int, + ) -> Awaitable[list[_DailySpendRecord]]: ... + + +class _KeyMetadataInfo(TypedDict, total=False): + key_alias: str | None + team_id: str | None + + +class _AggregatedSpendData(TypedDict): + results: list[DailySpendData] + totals: SpendMetrics + + +def update_metrics(existing_metrics: SpendMetrics, record: _SpendMetricsRecord) -> SpendMetrics: """Update metrics with new record data. Rollup rows can carry None for numeric fields when SUM() spans zero rows @@ -66,19 +224,19 @@ def _is_user_agent_tag(tag: Optional[str]) -> bool: return normalized_tag.startswith("user-agent:") or normalized_tag.startswith("user agent:") -def compute_tag_metadata_totals(records: List[Any]) -> SpendMetrics: +def compute_tag_metadata_totals(records: list[_DailySpendRecord]) -> SpendMetrics: """ Deduplicate spend metrics for tags using request_id, ignoring User-Agent prefixed tags. Each unique request_id contributes at most one record (the tag with max spend) to metadata. """ - deduped_records: Dict[str, Any] = {} + deduped_records: dict[str, _DailySpendRecord] = {} for record in records: - request_id = getattr(record, "request_id", None) + request_id: str | None = getattr(record, "request_id", None) if not request_id: continue - tag_value = getattr(record, "tag", None) + tag_value: str | None = getattr(record, "tag", None) if _is_user_agent_tag(tag_value): continue @@ -94,12 +252,12 @@ def compute_tag_metadata_totals(records: List[Any]) -> SpendMetrics: def update_breakdown_metrics( breakdown: BreakdownMetrics, - record: Any, - model_metadata: Dict[str, Dict[str, Any]], - provider_metadata: Dict[str, Dict[str, Any]], - api_key_metadata: Dict[str, Dict[str, Any]], + record: _DailySpendRecord, + model_metadata: dict[str, dict[str, object]], + provider_metadata: dict[str, dict[str, object]], + api_key_metadata: dict[str, _KeyMetadataInfo], entity_id_field: Optional[str] = None, - entity_metadata_field: Optional[Dict[str, dict]] = None, + entity_metadata_field: dict[str, dict[str, object]] | None = None, ) -> BreakdownMetrics: """Updates breakdown metrics for a single record using the existing update_metrics function""" @@ -241,7 +399,7 @@ def update_breakdown_metrics( # Update entity-specific metrics if entity_id_field is provided if entity_id_field: - entity_value = getattr(record, entity_id_field, None) + entity_value: str | None = getattr(record, entity_id_field, None) entity_value = entity_value if entity_value else "Unassigned" # allow for null entity_id_field if entity_value not in breakdown.entities: breakdown.entities[entity_value] = MetricWithMetadata( @@ -270,22 +428,24 @@ def update_breakdown_metrics( async def get_api_key_metadata( prisma_client: PrismaClient, api_keys: Set[str], -) -> Dict[str, Dict[str, Any]]: +) -> dict[str, _KeyMetadataInfo]: """Get api key metadata, falling back to deleted keys table for keys not found in active table. This ensures that key_alias and team_id are preserved in historical activity logs even after a key is deleted or regenerated. """ - key_records = await VerificationTokenRepository(prisma_client).table.find_many( + key_records = await _token_table(VerificationTokenRepository(prisma_client)).find_many( where={"token": {"in": list(api_keys)}} ) - result = {k.token: {"key_alias": k.key_alias, "team_id": k.team_id} for k in key_records} + result: dict[str, _KeyMetadataInfo] = { + k.token: {"key_alias": k.key_alias, "team_id": k.team_id} for k in key_records + } # For any keys not found in the active table, check the deleted keys table missing_keys = api_keys - set(result.keys()) if missing_keys: try: - deleted_key_records = await DeletedVerificationTokenRepository(prisma_client).table.find_many( + deleted_key_records = await _token_table(DeletedVerificationTokenRepository(prisma_client)).find_many( where={"token": {"in": list(missing_keys)}}, order={"deleted_at": "desc"}, ) @@ -342,12 +502,12 @@ def _build_where_conditions( api_key: Optional[Union[str, List[str]]], exclude_entity_ids: Optional[List[str]] = None, timezone_offset_minutes: Optional[int] = None, -) -> Dict[str, Any]: +) -> _WhereConditions: """Build prisma where clause for daily activity queries.""" # Adjust dates for timezone if provided adjusted_start, adjusted_end = _adjust_dates_for_timezone(start_date, end_date, timezone_offset_minutes) - where_conditions: Dict[str, Any] = { + where_conditions: _WhereConditions = { "date": { "gte": adjusted_start, "lte": adjusted_end, @@ -369,7 +529,7 @@ def _build_where_conditions( where_conditions[entity_id_field] = {"equals": entity_id} if exclude_entity_ids: - current = where_conditions.get(entity_id_field, {}) + current: str | _WhereFilter = where_conditions.get(entity_id_field, {}) if isinstance(current, str): current = {"equals": current} current["not"] = {"in": exclude_entity_ids} @@ -389,7 +549,7 @@ def _build_aggregated_sql_query( api_key: Optional[str], exclude_entity_ids: Optional[List[str]] = None, timezone_offset_minutes: Optional[int] = None, -) -> Tuple[str, List[Any]]: +) -> tuple[str, list[str]]: """Build a parameterized SQL GROUP BY query for aggregated daily activity. Groups by (date, api_key, model, model_group, custom_llm_provider, @@ -407,7 +567,7 @@ def _build_aggregated_sql_query( adjusted_start, adjusted_end = _adjust_dates_for_timezone(start_date, end_date, timezone_offset_minutes) sql_conditions: List[str] = [] - sql_params: List[Any] = [] + sql_params: list[str] = [] p = 1 # parameter index (1-based for PostgreSQL $N placeholders) # Date range (always present) @@ -506,17 +666,17 @@ def _build_aggregated_sql_query( def _aggregate_spend_records_sync( *, - records: List[Any], - api_key_metadata: Dict[str, Dict[str, Any]], + records: list[_DailySpendRecord], + api_key_metadata: dict[str, _KeyMetadataInfo], entity_id_field: Optional[str], - entity_metadata_field: Optional[Dict[str, dict]], -) -> Dict[str, Any]: - model_metadata: Dict[str, Dict[str, Any]] = {} - provider_metadata: Dict[str, Dict[str, Any]] = {} + entity_metadata_field: dict[str, dict[str, object]] | None, +) -> _AggregatedSpendData: + model_metadata: dict[str, dict[str, object]] = {} + provider_metadata: dict[str, dict[str, object]] = {} results: List[DailySpendData] = [] total_metrics = SpendMetrics() - grouped_data: Dict[str, Dict[str, Any]] = {} + grouped_data: dict[str, GroupedData] = {} for record in records: date_str = record.date @@ -557,10 +717,10 @@ def _aggregate_spend_records_sync( async def _aggregate_spend_records( *, prisma_client: PrismaClient, - records: List[Any], + records: list[_DailySpendRecord], entity_id_field: Optional[str], - entity_metadata_field: Optional[Dict[str, dict]], -) -> Dict[str, Any]: + entity_metadata_field: dict[str, dict[str, object]] | None, +) -> _AggregatedSpendData: """Aggregate rows into DailySpendData list and total metrics. The per-row loop is offloaded to a worker thread via asyncio.to_thread so @@ -568,7 +728,7 @@ async def _aggregate_spend_records( """ api_keys: Set[str] = {record.api_key for record in records if record.api_key} - api_key_metadata: Dict[str, Dict[str, Any]] = {} + api_key_metadata: dict[str, _KeyMetadataInfo] = {} if api_keys: api_key_metadata = await get_api_key_metadata(prisma_client, api_keys) @@ -603,7 +763,7 @@ async def _aggregate_spend_records( _GROUP_DATE_ENDPOINT_API_KEY = 30 # 0b0011110 -def _record_to_spend_metrics(record: Any) -> SpendMetrics: +def _record_to_spend_metrics(record: _SpendMetricsRecord) -> SpendMetrics: """Build a SpendMetrics directly from one already-aggregated rollup row. SUM() over zero rows is SQL NULL, so rollup rows (notably the grand-total @@ -627,16 +787,16 @@ def _record_to_spend_metrics(record: Any) -> SpendMetrics: ) -def _key_metadata(api_key_metadata: Dict[str, Dict[str, Any]], api_key: str) -> KeyMetadata: +def _key_metadata(api_key_metadata: dict[str, _KeyMetadataInfo], api_key: str) -> KeyMetadata: meta = api_key_metadata.get(api_key, {}) return KeyMetadata(key_alias=meta.get("key_alias"), team_id=meta.get("team_id")) def _aggregate_grouping_sets_records_sync( *, - records: List[Any], - api_key_metadata: Dict[str, Dict[str, Any]], -) -> Dict[str, Any]: + records: list[_GroupingSetsRow], + api_key_metadata: dict[str, _KeyMetadataInfo], +) -> _AggregatedSpendData: """Build the response from rollup rows produced by the GROUPING SETS query. Each row carries a `group_level` bitmask (from Postgres GROUPING()) that @@ -645,10 +805,10 @@ def _aggregate_grouping_sets_records_sync( summing in Python and no nested update_metrics calls. """ total_metrics = SpendMetrics() - grouped_data: Dict[str, Dict[str, Any]] = {} + grouped_data: dict[str, GroupedData] = {} - def ensure_date(date_str: str) -> Dict[str, Any]: - bucket = grouped_data.get(date_str) + def ensure_date(date_str: str) -> GroupedData: + bucket: GroupedData | None = grouped_data.get(date_str) if bucket is None: bucket = {"metrics": SpendMetrics(), "breakdown": BreakdownMetrics()} grouped_data[date_str] = bucket @@ -753,12 +913,12 @@ def assign_api_key_breakdown( async def _aggregate_grouping_sets_records( *, prisma_client: PrismaClient, - records: List[Any], -) -> Dict[str, Any]: + records: list[_GroupingSetsRow], +) -> _AggregatedSpendData: """Async wrapper: fetch api_key_metadata, then dispatch on a worker thread.""" api_keys: Set[str] = {r.api_key for r in records if r.api_key} - api_key_metadata: Dict[str, Dict[str, Any]] = {} + api_key_metadata: dict[str, _KeyMetadataInfo] = {} if api_keys: api_key_metadata = await get_api_key_metadata(prisma_client, api_keys) @@ -774,7 +934,7 @@ async def get_daily_activity( table_name: str, entity_id_field: str, entity_id: Optional[Union[str, List[str]]], - entity_metadata_field: Optional[Dict[str, dict]], + entity_metadata_field: dict[str, dict[str, object]] | None, start_date: Optional[str], end_date: Optional[str], model: Optional[str], @@ -782,9 +942,11 @@ async def get_daily_activity( page: int, page_size: int, exclude_entity_ids: Optional[List[str]] = None, - metadata_metrics_func: Optional[Callable[[List[Any]], SpendMetrics]] = None, + metadata_metrics_func: Callable[[list[_DailySpendRecord]], SpendMetrics] | None = None, timezone_offset_minutes: Optional[int] = None, - resolve_entity_metadata: Optional[Callable[[list[Any]], Awaitable[dict[str, dict]]]] = None, + resolve_entity_metadata: Optional[ + Callable[[list[_DailySpendRecord]], Awaitable[dict[str, dict[str, object]]]] + ] = None, ) -> SpendAnalyticsPaginatedResponse: """Common function to get daily activity for any entity type. @@ -818,8 +980,12 @@ async def get_daily_activity( timezone_offset_minutes=timezone_offset_minutes, ) + table = cast( # cast-ok: prisma table accessor is dynamic + _DailySpendTable, getattr(prisma_client.db, table_name) + ) + # Get total count for pagination - total_count = await getattr(prisma_client.db, table_name).count(where=where_conditions) + total_count = await table.count(where=where_conditions) # Fetch paginated results. # ``date`` alone is not a unique sort key -- a busy tenant has many @@ -831,7 +997,7 @@ async def get_daily_activity( # total. Adding ``id`` (the row's UUID primary key, present on both # LiteLLM_DailyUserSpend and LiteLLM_DailyTeamSpend) as a tiebreaker # gives every page a stable cursor (#30164). - daily_spend_data = await getattr(prisma_client.db, table_name).find_many( + daily_spend_data = await table.find_many( where=where_conditions, order=[ {"date": "desc"}, @@ -893,7 +1059,7 @@ async def get_daily_activity_aggregated( table_name: str, entity_id_field: str, entity_id: Optional[Union[str, List[str]]], - entity_metadata_field: Optional[Dict[str, dict]], + entity_metadata_field: dict[str, dict[str, object]] | None, start_date: Optional[str], end_date: Optional[str], model: Optional[str], @@ -935,11 +1101,17 @@ async def get_daily_activity_aggregated( ) # Execute GROUPING SETS query — returns one row per rollup level. - rows = await prisma_client.db.query_raw(sql_query, *sql_params) + db = cast(_RawQueryDb, prisma_client.db) # cast-ok: query_raw is resolved dynamically on the prisma wrapper + rows = await db.query_raw(sql_query, *sql_params) if rows is None: rows = [] - records = [SimpleNamespace(**row) for row in rows] + records = [ + cast( # cast-ok: raw rollup rows expose the selected SQL columns as attributes + _GroupingSetsRow, SimpleNamespace(**row) + ) + for row in rows + ] # The grouping-sets dispatcher places each row directly in its bucket # using the row's GROUPING() bitmask. No Python-side summing needed. diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index f741783134e..9a5bcf24ae2 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -16,7 +16,7 @@ import json import traceback from datetime import datetime, timezone -from typing import Any, Dict, List, Optional, Union, cast +from typing import Any, Dict, List, Mapping, Optional, Protocol, Sequence, TypeVar, Union, cast import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Request, status @@ -72,10 +72,100 @@ ) if TYPE_CHECKING: + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.proxy_server import PrismaClient + from litellm.proxy.utils import ProxyLogging router = APIRouter() +_PrismaRowT_co = TypeVar("_PrismaRowT_co", covariant=True) + + +class _TeamTableRow(Protocol): + team_id: str + members_with_roles: str + + def model_dump(self) -> Mapping[str, Any]: ... + + +class _OrgMembershipRow(Protocol): + user_id: str + organization_id: str | None + + +class _PrismaTableClient(Protocol[_PrismaRowT_co]): + async def find_first(self, *, where: Mapping[str, object]) -> _PrismaRowT_co | None: ... + + async def find_unique(self, *, where: Mapping[str, object]) -> _PrismaRowT_co | None: ... + + async def find_many( + self, + *, + where: Mapping[str, object] = ..., + skip: int = ..., + take: int = ..., + order: Mapping[str, str] = ..., + ) -> Sequence[_PrismaRowT_co]: ... + + async def count(self, *, where: Mapping[str, object] = ...) -> int: ... + + async def create(self, *, data: Mapping[str, object]) -> _PrismaRowT_co: ... + + async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> _PrismaRowT_co | None: ... + + async def update_many(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> int: ... + + async def delete_many(self, *, where: Mapping[str, object]) -> int: ... + + +def _user_table(prisma_client: Optional["PrismaClient"]) -> _PrismaTableClient[LiteLLM_UserTable]: + return cast( # cast-ok: repository .table is typed Any at the prisma seam + _PrismaTableClient[LiteLLM_UserTable], + UserRepository(prisma_client).table, + ) + + +def _team_table(prisma_client: Optional["PrismaClient"]) -> _PrismaTableClient[_TeamTableRow]: + return cast( # cast-ok: repository .table is typed Any at the prisma seam + _PrismaTableClient[_TeamTableRow], + TeamRepository(prisma_client).table, + ) + + +def _verification_token_table(prisma_client: Optional["PrismaClient"]) -> _PrismaTableClient[object]: + return cast( # cast-ok: repository .table is typed Any at the prisma seam + _PrismaTableClient[object], + VerificationTokenRepository(prisma_client).table, + ) + + +def _invitation_link_table(prisma_client: Optional["PrismaClient"]) -> _PrismaTableClient[object]: + return cast( # cast-ok: repository .table is typed Any at the prisma seam + _PrismaTableClient[object], + InvitationLinkRepository(prisma_client).table, + ) + + +def _team_membership_table(prisma_client: Optional["PrismaClient"]) -> _PrismaTableClient[object]: + return cast( # cast-ok: repository .table is typed Any at the prisma seam + _PrismaTableClient[object], + TeamMembershipRepository(prisma_client).table, + ) + + +def _org_membership_table(prisma_client: Optional["PrismaClient"]) -> _PrismaTableClient[_OrgMembershipRow]: + return cast( # cast-ok: repository .table is typed Any at the prisma seam + _PrismaTableClient[_OrgMembershipRow], + OrganizationMembershipRepository(prisma_client).table, + ) + + +def _organization_table(prisma_client: Optional["PrismaClient"]) -> _PrismaTableClient[object]: + return cast( # cast-ok: repository .table is typed Any at the prisma seam + _PrismaTableClient[object], + OrganizationRepository(prisma_client).table, + ) + def _hash_password_in_dict(data: dict) -> None: """Hash password field in-place if present.""" @@ -128,7 +218,7 @@ def _update_internal_new_user_params(data_json: dict, data: NewUserRequest) -> d async def _check_duplicate_user_field( field_name: str, field_value: Optional[str], - prisma_client: Any, + prisma_client: Optional["PrismaClient"], *, case_insensitive: bool = False, label: Optional[str] = None, @@ -156,7 +246,7 @@ async def _check_duplicate_user_field( if case_insensitive: where_clause[field_name]["mode"] = "insensitive" - existing_user = await UserRepository(prisma_client).table.find_first(where=where_clause) + existing_user = await _user_table(prisma_client).find_first(where=where_clause) if existing_user is not None: existing_value = getattr(existing_user, field_name, value) @@ -167,7 +257,7 @@ async def _check_duplicate_user_field( ) -async def _check_duplicate_user_email(user_email: Optional[str], prisma_client: Any) -> None: +async def _check_duplicate_user_email(user_email: str | None, prisma_client: Optional["PrismaClient"]) -> None: """ Helper function to check if a user email already exists in the database. """ @@ -180,7 +270,7 @@ async def _check_duplicate_user_email(user_email: Optional[str], prisma_client: ) -async def _check_duplicate_user_id(user_id: Optional[str], prisma_client: Any) -> None: +async def _check_duplicate_user_id(user_id: str | None, prisma_client: Optional["PrismaClient"]) -> None: """ Helper function to check if a user id already exists in the database. """ @@ -641,15 +731,15 @@ def _enforce_user_info_access(user_id: Optional[str], user_api_key_dict: UserAPI async def _get_user_info_teams( - prisma_client: Any, + prisma_client: "PrismaClient", user_id: Optional[str], - user_info: Optional[Any], + user_info: object | None, user_api_key_dict: UserAPIKeyAuth, -) -> tuple[list[Any], Optional[list[Any]]]: +) -> tuple[list[TeamListResponseObject], list[TeamListResponseObject] | None]: """Fetch and merge teams from membership + user.teams field.""" from litellm.proxy.management_endpoints.team_endpoints import list_team - team_list: list[Any] = [] + team_list: list[TeamListResponseObject] = [] team_id_list: list[str] = [] teams_1 = await list_team( @@ -664,7 +754,7 @@ async def _get_user_info_teams( team_list = teams_1 team_id_list = [team.team_id for team in teams_1] - teams_2: Optional[list[Any]] = None + teams_2: list[TeamListResponseObject] | None = None target_team_ids = getattr(user_info, "teams", None) if target_team_ids and isinstance(target_team_ids, list): @@ -711,10 +801,12 @@ def _redact_scim_enterprise_metadata( def _build_user_info_response( user_id: Optional[str], - user_info: Optional[Any], + user_info: Optional[ + BaseModel | dict[str, Any] + ], # mutable-ok: password and metadata are stripped from this dict in place below keys: Optional[List[LiteLLM_VerificationToken]], - team_list: list[Any], - teams_1: Optional[list[Any]], + team_list: list[TeamListResponseObject], + teams_1: list[TeamListResponseObject] | None, ) -> UserInfoResponse: """Create UserInfoResponse while filtering sensitive fields.""" if user_info is None and keys is not None: @@ -839,8 +931,8 @@ async def _check_user_info_v2_access( return None # Helper: fetch the target user row (reused across branches) - async def _fetch_target_user(): - return await UserRepository(prisma_client).table.find_unique(where={"user_id": target_user_id}) + async def _fetch_target_user() -> LiteLLM_UserTable | None: + return await _user_table(prisma_client).find_unique(where={"user_id": target_user_id}) # Rule 1: Proxy admins — fetch and return the target row directly if _user_has_admin_view(user_api_key_dict): @@ -853,9 +945,7 @@ async def _fetch_target_user(): # Rule 3: Team admins can look up users in their teams if user_api_key_dict.user_id is not None: # Get caller's teams - caller_user = await UserRepository(prisma_client).table.find_unique( - where={"user_id": user_api_key_dict.user_id} - ) + caller_user = await _user_table(prisma_client).find_unique(where={"user_id": user_api_key_dict.user_id}) if caller_user is not None and caller_user.teams: # Fetch the target user ONCE, before the loop target_user = await _fetch_target_user() @@ -863,9 +953,9 @@ async def _fetch_target_user(): return None # Get all teams the caller belongs to - teams = await TeamRepository(prisma_client).table.find_many(where={"team_id": {"in": caller_user.teams}}) + teams = await _team_table(prisma_client).find_many(where={"team_id": {"in": caller_user.teams}}) for team in teams: - team_obj = LiteLLM_TeamTable(**team.model_dump()) + team_obj = LiteLLM_TeamTable.model_validate(team.model_dump()) if _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team_obj): # Check if target user is in this team if team.team_id in (target_user.teams or []): @@ -989,7 +1079,10 @@ async def _get_user_info_for_proxy_admin(user_api_key_dict: UserAPIKeyAuth): "Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" ) - results = await prisma_client.db.query_raw(sql_query) + results = cast( # cast-ok: prisma query_raw is untyped at the wrapper seam + list[dict[str, Any]], + await prisma_client.db.query_raw(sql_query), + ) verbose_proxy_logger.debug("results_keys: %s", results) @@ -1121,7 +1214,7 @@ def _update_internal_user_params( async def _schedule_user_update_audit_log( - response: Dict[str, Any], + response: Mapping[str, object], existing_user_row: Optional[BaseModel], litellm_changed_by: Optional[str], user_api_key_dict: UserAPIKeyAuth, @@ -1132,9 +1225,9 @@ async def _schedule_user_update_audit_log( if prisma_client is None: return try: - updated_user_row = await UserRepository(prisma_client).table.find_first(where={"user_id": response["user_id"]}) + updated_user_row = await _user_table(prisma_client).find_first(where={"user_id": response["user_id"]}) if updated_user_row: - user_row_typed = LiteLLM_UserTable(**updated_user_row.model_dump(exclude_none=True)) + user_row_typed = LiteLLM_UserTable.model_validate(updated_user_row.model_dump(exclude_none=True)) asyncio.create_task( UserManagementEventHooks.create_internal_user_audit_log( user_id=user_row_typed.user_id, @@ -1160,7 +1253,7 @@ def _check_user_update_authz( raise HTTPException(status_code=403, detail="Only proxy admins can modify user roles.") if existing_user_row is not None: - typed_row = LiteLLM_UserTable(**existing_user_row.model_dump(exclude_none=True)) + typed_row = LiteLLM_UserTable.model_validate(existing_user_row.model_dump(exclude_none=True)) if not can_user_call_user_update(user_api_key_dict=user_api_key_dict, user_info=typed_row): raise HTTPException( status_code=403, @@ -1179,7 +1272,7 @@ def _check_user_update_authz( async def _invalidate_user_spend_counter_if_changed( - non_default_values: dict[str, Any], + non_default_values: Mapping[str, object], ) -> None: """Invalidate the cross-pod spend counter after a direct ``spend`` change. @@ -1225,18 +1318,14 @@ async def _update_single_user_helper( existing_user_row: Optional[BaseModel] = None if user_request.user_id: - existing_user_row = await UserRepository(prisma_client).table.find_first( - where={"user_id": user_request.user_id} - ) + existing_user_row = await _user_table(prisma_client).find_first(where={"user_id": user_request.user_id}) elif user_request.user_email: - existing_user_row = await UserRepository(prisma_client).table.find_first( - where={"user_email": user_request.user_email} - ) + existing_user_row = await _user_table(prisma_client).find_first(where={"user_email": user_request.user_email}) _check_user_update_authz(user_request, user_api_key_dict, existing_user_row) if existing_user_row is not None: - existing_user_row = LiteLLM_UserTable(**existing_user_row.model_dump(exclude_none=True)) + existing_user_row = LiteLLM_UserTable.model_validate(existing_user_row.model_dump(exclude_none=True)) # Prevent budget self-escalation (GHSA-wvg4-6222-3q4r): non-admin callers # must not be able to raise their own budget/spend fields. @@ -1585,7 +1674,7 @@ async def bulk_user_update( detail="Only proxy admins can update all users at once.", ) # Optimized path for updating all users directly in database - all_users_in_db = await UserRepository(prisma_client).table.find_many(order={"created_at": "desc"}) + all_users_in_db = await _user_table(prisma_client).find_many(order={"created_at": "desc"}) if not all_users_in_db: raise HTTPException( @@ -1617,7 +1706,7 @@ async def bulk_user_update( try: # Perform bulk database update - await UserRepository(prisma_client).table.update_many( + await _user_table(prisma_client).update_many( where={}, data=non_default_values, # Update all users ) @@ -1700,9 +1789,9 @@ async def bulk_user_update( async def get_user_key_counts( - prisma_client, + prisma_client: Optional["PrismaClient"], user_ids: Optional[List[str]] = None, -): +) -> Mapping[str, int]: """ Helper function to get the count of keys for each user using Prisma's count method. @@ -1718,11 +1807,11 @@ async def get_user_key_counts( if not user_ids or len(user_ids) == 0: return {} - result = {} + result: dict[str, int] = {} # mutable-ok: one key assigned per user_id inside the loop below # Get count for each user_id individually for user_id in user_ids: - count = await VerificationTokenRepository(prisma_client).table.count( + count = await _verification_token_table(prisma_client).count( where={ "user_id": user_id, "OR": [ @@ -1771,9 +1860,9 @@ def _validate_sort_params(sort_by: Optional[str], sort_order: str) -> Optional[D async def _authorize_user_list_request( user_api_key_dict: UserAPIKeyAuth, organization_ids: Optional[str], - prisma_client: Any, - user_api_key_cache: Any, - proxy_logging_obj: Any, + prisma_client: "PrismaClient", + user_api_key_cache: "UserApiKeyCache", + proxy_logging_obj: "ProxyLogging", ) -> Optional[str]: """ Authorize the /user/list request and return the (possibly scoped) organization_ids string. @@ -1911,7 +2000,7 @@ async def get_users( skip = (page - 1) * page_size # Build where conditions based on provided parameters - where_conditions: Dict[str, Any] = {} + where_conditions: dict[str, object] = {} if role: where_conditions["user_role"] = role @@ -1959,7 +2048,7 @@ async def get_users( _validate_sort_params(sort_by, sort_order) if sort_by is not None and isinstance(sort_by, str) else None ) - users = await UserRepository(prisma_client).table.find_many( + users: Sequence[LiteLLM_UserTable] | None = await _user_table(prisma_client).find_many( where=where_conditions, skip=skip, take=page_size, @@ -1967,9 +2056,10 @@ async def get_users( ) # Get total count of user rows - total_count = await UserRepository(prisma_client).table.count(where=where_conditions) + total_count = await _user_table(prisma_client).count(where=where_conditions) # Get key count for each user + user_key_counts: Mapping[str, int] if users is not None: user_key_counts = await get_user_key_counts(prisma_client, [user.user_id for user in users]) else: @@ -1986,7 +2076,11 @@ async def get_users( for user in users: user_dump = user.model_dump() user_dump["metadata"] = _redact_scim_enterprise_metadata(user_dump.get("metadata")) - user_list.append(LiteLLM_UserTableWithKeyCount(**user_dump, key_count=user_key_counts.get(user.user_id, 0))) + user_list.append( + LiteLLM_UserTableWithKeyCount.model_validate( + {**user_dump, "key_count": user_key_counts.get(user.user_id, 0)} + ) + ) else: user_list = [] @@ -2056,10 +2150,10 @@ async def delete_user( # loop an org-admin of org-A could delete users in org-B by supplying # {"user_ids": [victim_in_org_B], "organization_id": "org-A"}. caller_is_proxy_admin = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value - caller_admin_org_ids: set = set() + caller_admin_org_ids: set[str] = set() if not caller_is_proxy_admin: caller_memberships = ( - await OrganizationMembershipRepository(prisma_client).table.find_many( + await _org_membership_table(prisma_client).find_many( where={ "user_id": user_api_key_dict.user_id, "user_role": LitellmUserRoles.ORG_ADMIN.value, @@ -2077,9 +2171,9 @@ async def delete_user( # Batch-fetch target memberships once before the per-user loop. Avoids # an N+1 DB call when delete_user is called with a large user_ids list. - target_org_ids_by_user: Dict[str, set] = {} + target_org_ids_by_user: dict[str, set[str]] = {} if not caller_is_proxy_admin: - all_target_memberships = await OrganizationMembershipRepository(prisma_client).table.find_many( + all_target_memberships = await _org_membership_table(prisma_client).find_many( where={"user_id": {"in": data.user_ids}} ) for m in all_target_memberships: @@ -2089,7 +2183,7 @@ async def delete_user( # check that all teams passed exist for user_id in data.user_ids: - user_row = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id}) + user_row = await _user_table(prisma_client).find_unique(where={"user_id": user_id}) if user_row is None: raise HTTPException( @@ -2141,11 +2235,11 @@ async def delete_user( ) ## CLEANUP MEMBERS_WITH_ROLES - fetch_all_teams = await TeamRepository(prisma_client).table.find_many(where={"team_id": {"in": user_row.teams}}) - teams_to_update = [] + fetch_all_teams = await _team_table(prisma_client).find_many(where={"team_id": {"in": user_row.teams}}) + teams_to_update: list[_TeamTableRow] = [] # mutable-ok: appended to inside the membership-cleanup loop below for team in fetch_all_teams: is_member_in_team, new_team_members = _cleanup_members_with_roles( - existing_team_row=LiteLLM_TeamTable(**team.model_dump()), + existing_team_row=LiteLLM_TeamTable.model_validate(team.model_dump()), data=TeamMemberDeleteRequest( team_id=team.team_id, user_id=user_row.user_id, @@ -2160,17 +2254,17 @@ async def delete_user( ## update teams for team in teams_to_update: - await TeamRepository(prisma_client).table.update( + await _team_table(prisma_client).update( where={"team_id": team.team_id}, data={"members_with_roles": team.members_with_roles}, ) # End of Audit logging ## DELETE ASSOCIATED KEYS - await VerificationTokenRepository(prisma_client).table.delete_many(where={"user_id": {"in": data.user_ids}}) + await _verification_token_table(prisma_client).delete_many(where={"user_id": {"in": data.user_ids}}) ## DELETE ASSOCIATED INVITATION LINKS - await InvitationLinkRepository(prisma_client).table.delete_many( + await _invitation_link_table(prisma_client).delete_many( where={ "OR": [ {"user_id": {"in": data.user_ids}}, @@ -2181,13 +2275,13 @@ async def delete_user( ) ## DELETE ASSOCIATED ORGANIZATION MEMBERSHIPS - await OrganizationMembershipRepository(prisma_client).table.delete_many(where={"user_id": {"in": data.user_ids}}) + await _org_membership_table(prisma_client).delete_many(where={"user_id": {"in": data.user_ids}}) ## DELETE ASSOCIATED TEAM MEMBERSHIPS - await TeamMembershipRepository(prisma_client).table.delete_many(where={"user_id": {"in": data.user_ids}}) + await _team_membership_table(prisma_client).delete_many(where={"user_id": {"in": data.user_ids}}) ## DELETE USERS - deleted_users = await UserRepository(prisma_client).table.delete_many(where={"user_id": {"in": data.user_ids}}) + deleted_users = await _user_table(prisma_client).delete_many(where={"user_id": {"in": data.user_ids}}) return deleted_users @@ -2215,14 +2309,14 @@ async def add_internal_user_to_organization( try: # Check if organization_id exists - organization_row = await OrganizationRepository(prisma_client).table.find_unique( + organization_row = await _organization_table(prisma_client).find_unique( where={"organization_id": organization_id} ) if organization_row is None: raise Exception(f"Organization not found, passed organization_id={organization_id}") # Create a new organization membership entry - new_membership = await OrganizationMembershipRepository(prisma_client).table.create( + new_membership = await _org_membership_table(prisma_client).create( data={ "user_id": user_id, "organization_id": organization_id, @@ -2239,9 +2333,9 @@ async def add_internal_user_to_organization( async def _resolve_org_filter_for_user_search( user_api_key_dict: UserAPIKeyAuth, team_id: Optional[str], - prisma_client: Any, - user_api_key_cache: Any, - proxy_logging_obj: Any, + prisma_client: "PrismaClient", + user_api_key_cache: "UserApiKeyCache", + proxy_logging_obj: "ProxyLogging", ) -> Optional[List[str]]: """ Return a list of org IDs to filter by, or ``None`` for no filter. @@ -2305,9 +2399,9 @@ async def _resolve_org_filter_for_user_search( async def _resolve_team_org_filter( user_api_key_dict: UserAPIKeyAuth, team_id: str, - prisma_client: Any, - user_api_key_cache: Any, - proxy_logging_obj: Any, + prisma_client: "PrismaClient", + user_api_key_cache: "UserApiKeyCache", + proxy_logging_obj: "ProxyLogging", ) -> List[str]: """Look up the team and return its org as a filter list, or raise 403.""" from litellm.proxy.management_endpoints.common_utils import _is_user_team_admin @@ -2397,7 +2491,7 @@ async def ui_view_users( skip = (page - 1) * page_size # Build where conditions based on provided parameters - where_conditions: Dict[str, Any] = {} + where_conditions: dict[str, object] = {} if user_id: where_conditions["user_id"] = { @@ -2416,7 +2510,7 @@ async def ui_view_users( where_conditions["organization_memberships"] = {"some": {"organization_id": {"in": org_filter_ids}}} # Query users with pagination and filters - users: Optional[List[BaseModel]] = await UserRepository(prisma_client).table.find_many( + users: Sequence[LiteLLM_UserTable] | None = await _user_table(prisma_client).find_many( where=where_conditions, skip=skip, take=page_size, @@ -2426,7 +2520,7 @@ async def ui_view_users( if not users: return [] - return [LiteLLM_UserTableFiltered(**user.model_dump()) for user in users] + return [LiteLLM_UserTableFiltered.model_validate(user.model_dump()) for user in users] except HTTPException: raise @@ -2438,13 +2532,15 @@ async def ui_view_users( # Using shared metric helper implementations from common_daily_activity -async def _resolve_user_email_metadata(prisma_client: "PrismaClient", records: list[Any]) -> dict[str, dict]: +async def _resolve_user_email_metadata( + prisma_client: Optional["PrismaClient"], records: Sequence[Any] +) -> dict[str, dict]: """Map each user_id on the page to its email/alias so the Usage dashboard can label the 'Spend Per User' chart with the email instead of the raw UUID.""" user_ids = {record.user_id for record in records if getattr(record, "user_id", None)} if not user_ids: return {} - users = await UserRepository(prisma_client).table.find_many(where={"user_id": {"in": list(user_ids)}}) + users = await _user_table(prisma_client).find_many(where={"user_id": {"in": list(user_ids)}}) return {user.user_id: {"user_email": user.user_email, "user_alias": user.user_alias} for user in users} diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index ac6a2a4a7db..96e2579e6b0 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -18,11 +18,12 @@ import re import secrets import traceback -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from datetime import datetime, timedelta, timezone from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, cast import fastapi +from typing_extensions import Never import yaml from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, status @@ -127,6 +128,7 @@ from litellm.types.router import Deployment from litellm.types.utils import ( BudgetConfig, + CredentialItem, PersonalUIKeyGenerationConfig, TeamUIKeyGenerationConfig, ) @@ -490,7 +492,7 @@ def handle_key_type(data: GenerateKeyRequest, data_json: dict) -> dict: def _validate_caller_can_change_key_ownership( data: Optional[BaseModel], - existing_key_row: Any, + existing_key_row: LiteLLM_VerificationToken, user_api_key_dict: UserAPIKeyAuth, ) -> None: """ @@ -889,7 +891,7 @@ 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: LiteLLM_BudgetTable = await BudgetRepository(prisma_client).table.create( data={ **new_budget, # type: ignore "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, @@ -1095,7 +1097,7 @@ async def _common_key_generation_helper( def _check_key_model_specific_limits( - keys: List[LiteLLM_VerificationToken], + keys: Sequence[LiteLLM_VerificationToken], data: Union[GenerateKeyRequest, UpdateKeyRequest], entity_rpm_limit: Optional[int], entity_tpm_limit: Optional[int], @@ -1166,7 +1168,7 @@ def _check_key_model_specific_limits( def _check_key_rpm_tpm_limits( - keys: List[LiteLLM_VerificationToken], + keys: Sequence[LiteLLM_VerificationToken], data: Union[GenerateKeyRequest, UpdateKeyRequest], entity_rpm_limit: Optional[int], entity_tpm_limit: Optional[int], @@ -1204,7 +1206,7 @@ def _check_key_rpm_tpm_limits( def check_team_key_model_specific_limits( - keys: List[LiteLLM_VerificationToken], + keys: Sequence[LiteLLM_VerificationToken], team_table: LiteLLM_TeamTableCachedObj, data: Union[GenerateKeyRequest, UpdateKeyRequest], ) -> None: @@ -1229,7 +1231,7 @@ def check_team_key_model_specific_limits( def check_team_key_rpm_tpm_limits( - keys: List[LiteLLM_VerificationToken], + keys: Sequence[LiteLLM_VerificationToken], team_table: LiteLLM_TeamTableCachedObj, data: Union[GenerateKeyRequest, UpdateKeyRequest], ) -> None: @@ -1261,7 +1263,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( + keys: Sequence[LiteLLM_VerificationToken] = await VerificationTokenRepository(prisma_client).table.find_many( where={"team_id": team_table.team_id}, ) # Exclude the key being updated to avoid double-counting its limits. @@ -1331,7 +1333,7 @@ async def _check_project_key_limits( def check_org_key_model_specific_limits( - keys: List[LiteLLM_VerificationToken], + keys: Sequence[LiteLLM_VerificationToken], org_table: LiteLLM_OrganizationTable, data: Union[GenerateKeyRequest, UpdateKeyRequest], ) -> None: @@ -1364,7 +1366,7 @@ def check_org_key_model_specific_limits( def check_org_key_rpm_tpm_limits( - keys: List[LiteLLM_VerificationToken], + keys: Sequence[LiteLLM_VerificationToken], org_table: LiteLLM_OrganizationTable, data: Union[GenerateKeyRequest, UpdateKeyRequest], ) -> None: @@ -1405,7 +1407,7 @@ 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: LiteLLM_UserTable | None = await UserRepository(prisma_client).table.find_unique( where={"user_id": user_api_key_dict.user_id}, include={"organization_memberships": True}, ) @@ -1443,7 +1445,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: Sequence[LiteLLM_VerificationToken] = await VerificationTokenRepository(prisma_client).table.find_many( where={"organization_id": org_table.organization_id}, ) # Exclude the key being updated to avoid double-counting its limits. @@ -1587,7 +1589,7 @@ 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 + result: Mapping[str, object] = await user_custom_key_generate(data) else: raise ValueError("user_custom_key_generate must be a coroutine") decision = result.get("decision", True) @@ -1786,7 +1788,7 @@ 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 + result: Mapping[str, object] = await user_custom_key_generate(data) else: raise ValueError("user_custom_key_generate must be a coroutine") decision = result.get("decision", True) @@ -1855,7 +1857,7 @@ def prepare_metadata_fields(data: BaseModel, non_default_values: dict, existing_ ) casted_metadata[reserved_field] = existing_value - data_json = data.model_dump(exclude_unset=True, exclude_none=True) + data_json: dict[str, object] = data.model_dump(exclude_unset=True, exclude_none=True) try: for k, v in data_json.items(): @@ -2046,7 +2048,9 @@ async def _get_and_validate_existing_key( hashed_token = _hash_token_if_needed(token=token) - existing_key_row = await VerificationTokenRepository(prisma_client).table.find_unique(where={"token": hashed_token}) + existing_key_row: LiteLLM_VerificationToken | None = await VerificationTokenRepository( + prisma_client + ).table.find_unique(where={"token": hashed_token}) if existing_key_row is None: raise ProxyException( @@ -2065,11 +2069,11 @@ async def _process_single_key_update( litellm_changed_by: Optional[str], prisma_client: Optional[PrismaClient], user_api_key_cache: UserApiKeyCache, - proxy_logging_obj: Any, + proxy_logging_obj: ProxyLogging, llm_router: Optional[Router], user_custom_key_update: Optional[Callable] = None, existing_key_row: Optional[LiteLLM_VerificationToken] = None, -) -> Dict[str, Any]: +) -> dict[str, object]: """ Process a single key update with all validations and checks. @@ -2218,9 +2222,9 @@ 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]: """Validate MCP servers in object_permission against the effective team.""" @@ -2255,12 +2259,12 @@ async def _validate_mcp_servers_for_key_update( async def _validate_update_key_data( 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, + user_api_key_cache: UserApiKeyCache, ) -> None: """Validate permissions and constraints for key update.""" # Reject NaN/±inf spend before it can reach the DB / spend counter. @@ -2614,7 +2618,7 @@ 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) + result: Mapping[str, object] = await user_custom_key_update(data) else: raise ValueError("user_custom_key_update must be a coroutine") decision = result.get("decision", True) @@ -2892,7 +2896,7 @@ 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() @@ -2962,7 +2966,9 @@ 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: Sequence[LiteLLM_VerificationToken] = await VerificationTokenRepository( + prisma_client + ).table.find_many( where={ "team_id": data.team_id, "AND": [ @@ -2980,7 +2986,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 = [cast(str, row.token) for row in existing_keys] # cast-ok: DB rows always carry a token else: if data.key_ids is None or len(data.key_ids) == 0: raise HTTPException( @@ -3050,7 +3056,7 @@ async def bulk_update_team_keys( update_key_request = UpdateKeyRequest( key=token, team_id=data.team_id, - **update_field_dict, + **cast(dict[str, Never], update_field_dict), # cast-ok: pydantic validates the dumped fields ) updated_key_info = await _process_single_key_update( update_key_request=update_key_request, @@ -3367,7 +3373,9 @@ 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: Sequence[LiteLLM_VerificationToken] = await VerificationTokenRepository( + prisma_client + ).table.find_many( where={"key_alias": {"in": data.key_aliases}}, include={"litellm_budget_table": True}, ) @@ -4039,7 +4047,7 @@ def _transform_verification_tokens_to_deleted_records( keys: List[LiteLLM_VerificationToken], user_api_key_dict: UserAPIKeyAuth, litellm_changed_by: Optional[str] = None, -) -> List[Dict[str, Any]]: +) -> list[dict[str, object]]: """Transform verification tokens into deleted token records ready for persistence.""" if not keys: return [] @@ -4049,13 +4057,15 @@ def _transform_verification_tokens_to_deleted_records( for key in keys: key_payload = key.model_dump() deleted_record = LiteLLM_DeletedVerificationToken( - **key_payload, + **cast(dict[str, Never], key_payload), # cast-ok: pydantic revalidates the dumped fields 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() + record: dict[str, object] = ( + deleted_record.model_dump() + ) # mutable-ok: org_id/relation keys are popped and JSON fields rewritten in place below # Map org_id to organization_id (model uses org_id, but schema expects organization_id) org_id_value = record.pop("org_id", None) @@ -4090,7 +4100,7 @@ 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.""" @@ -4124,9 +4134,9 @@ async def delete_key_aliases( 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( - where={"key_alias": {"in": key_aliases}} - ) + _keys_being_deleted: Sequence[LiteLLM_VerificationToken] = await VerificationTokenRepository( + prisma_client + ).table.find_many(where={"key_alias": {"in": key_aliases}}) tokens = [key.token for key in _keys_being_deleted] return await delete_verification_tokens( @@ -4161,7 +4171,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: list[LiteLLM_ProxyModelTable] | None = await ModelRepository(prisma_client).table.find_many() except Exception: models = None # 2. process model table @@ -4171,14 +4181,16 @@ async def _rotate_master_key( new_models = [] for model in decrypted_models: new_model = await _add_model_to_db( - model_params=Deployment(**model), + model_params=Deployment(**cast(dict[str, Never], model)), # cast-ok: pydantic revalidates user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, new_encryption_key=new_master_key, should_create_model_in_db=False, ) if new_model: - _dumped = new_model.model_dump(exclude_none=True) + _dumped: dict[str, object] = new_model.model_dump( + exclude_none=True + ) # mutable-ok: litellm_params/model_info are rewritten in place as prisma.Json _dumped["litellm_params"] = prisma.Json(_dumped["litellm_params"]) # type: ignore[attr-defined] _dumped["model_info"] = prisma.Json(_dumped["model_info"]) # type: ignore[attr-defined] new_models.append(_dumped) @@ -4191,7 +4203,7 @@ async def _rotate_master_key( ) # 3. process config table try: - config = await ConfigRepository(prisma_client).table.find_many() + config: Sequence[LiteLLM_Config] | None = await ConfigRepository(prisma_client).table.find_many() except Exception: config = None @@ -4256,7 +4268,7 @@ async def _rotate_master_key( # 5. process credentials table try: - credentials = await CredentialsRepository(prisma_client).table.find_many() + credentials: Sequence[CredentialItem] | None = await CredentialsRepository(prisma_client).table.find_many() except Exception: credentials = None if credentials: @@ -4270,7 +4282,9 @@ async def _rotate_master_key( updated_patch=decrypted_cred, new_encryption_key=new_master_key, ) - _cred_data = encrypted_cred.model_dump(exclude_none=True) + _cred_data: dict[str, object] = encrypted_cred.model_dump( + exclude_none=True + ) # mutable-ok: credential_values/credential_info are rewritten in place as prisma.Json if "credential_values" in _cred_data: _cred_data["credential_values"] = prisma.Json( # type: ignore[attr-defined] _cred_data["credential_values"] @@ -4520,7 +4534,7 @@ 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: LiteLLM_VerificationToken | None = await VerificationTokenRepository(prisma_client).table.update( where={"token": hashed_api_key}, data=update_data, # type: ignore ) @@ -4535,7 +4549,7 @@ async def _execute_virtual_key_regeneration( proxy_logging_obj=proxy_logging_obj, ) - response = GenerateKeyResponse(**updated_token_dict) + response = GenerateKeyResponse(**cast(dict[str, Never], updated_token_dict)) # cast-ok: pydantic revalidates asyncio.create_task( KeyManagementEventHooks.async_key_rotated_hook( data=data, @@ -4721,7 +4735,9 @@ async def regenerate_key_fn( else: hashed_api_key = hash_token(key) - _key_in_db = await VerificationTokenRepository(prisma_client).table.find_unique( + _key_in_db: LiteLLM_VerificationToken | None = await VerificationTokenRepository( + prisma_client + ).table.find_unique( where={"token": hashed_api_key}, ) if _key_in_db is None: @@ -4853,7 +4869,7 @@ async def _check_proxy_or_team_admin_for_key( ) -def _validate_reset_spend_value(reset_to: Any, key_in_db: LiteLLM_VerificationToken) -> float: +def _validate_reset_spend_value(reset_to: object, key_in_db: LiteLLM_VerificationToken) -> float: if not isinstance(reset_to, (int, float)): raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, @@ -4925,7 +4941,9 @@ 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: LiteLLM_VerificationToken | None = await VerificationTokenRepository( + prisma_client + ).table.find_unique( where={"token": hashed_api_key}, include={"litellm_budget_table": True}, ) @@ -4945,7 +4963,7 @@ async def reset_key_spend_fn( user_api_key_cache=user_api_key_cache, ) - updated_key = await VerificationTokenRepository(prisma_client).table.update( + updated_key: LiteLLM_VerificationToken | None = await VerificationTokenRepository(prisma_client).table.update( where={"token": hashed_api_key}, data={"spend": reset_to}, ) @@ -5029,7 +5047,9 @@ async def validate_key_list_check( code=status.HTTP_403_FORBIDDEN, ) - complete_user_info = LiteLLM_UserTable(**complete_user_info_db_obj.model_dump()) + complete_user_info = LiteLLM_UserTable( + **cast(dict[str, Never], complete_user_info_db_obj.model_dump()) # cast-ok: pydantic revalidates + ) # internal user can only see their own keys if user_id: @@ -5063,7 +5083,7 @@ async def validate_key_list_check( if key_hash: try: - key_info = await VerificationTokenRepository(prisma_client).table.find_unique( + key_info: LiteLLM_VerificationToken = await VerificationTokenRepository(prisma_client).table.find_unique( where={"token": key_hash}, ) except Exception: @@ -5102,7 +5122,10 @@ async def _fetch_user_team_objects( if teams is None: return [] - return [LiteLLM_TeamTable(**team.model_dump()) for team in teams] + return [ + LiteLLM_TeamTable(**cast(dict[str, Never], team.model_dump())) # cast-ok: pydantic revalidates + for team in teams + ] def _get_admin_team_ids_from_objects( @@ -5370,8 +5393,8 @@ async def list_keys( async def _apply_non_admin_alias_scope( user_api_key_dict: UserAPIKeyAuth, - prisma_client: Any, - query_params: List[Any], + prisma_client: PrismaClient, + query_params: list[object], where_parts: List[str], ) -> None: """Append SQL scope conditions so non-admin users only see aliases for @@ -5384,7 +5407,9 @@ async def _apply_non_admin_alias_scope( # Look up the user's teams from the user table 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: LiteLLM_UserTable | None = await UserRepository(prisma_client).table.find_unique( + where={"user_id": user_api_key_dict.user_id} + ) if user_row is not None: user_teams = getattr(user_row, "teams", []) or [] @@ -5442,7 +5467,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[object] = [UI_SESSION_TOKEN_TEAM_ID] where_parts = [ "key_alias IS NOT NULL", "key_alias != ''", @@ -5469,7 +5494,7 @@ async def key_aliases( where_sql = " AND ".join(where_parts) count_sql = f'SELECT COUNT(*) AS count FROM "LiteLLM_VerificationToken" WHERE {where_sql}' - count_rows = await prisma_client.db.query_raw(count_sql, *query_params) + count_rows: Sequence[Mapping[str, int]] = await prisma_client.db.query_raw(count_sql, *query_params) total_count = int(count_rows[0]["count"]) if count_rows else 0 aliases_params = query_params + [size, (page - 1) * size] @@ -5482,7 +5507,7 @@ async def key_aliases( f" ORDER BY key_alias ASC" f" LIMIT ${limit_idx} OFFSET ${offset_idx}" ) - alias_rows = await prisma_client.db.query_raw(aliases_sql, *aliases_params) + alias_rows: Sequence[Mapping[str, str]] = 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")] total_pages = -(-total_count // size) if total_count > 0 else 0 @@ -5550,7 +5575,7 @@ 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}}]} @@ -5764,7 +5789,9 @@ async def _list_key_helper( # Fetch keys with pagination if use_deleted_table: - keys = await DeletedVerificationTokenRepository(prisma_client).table.find_many( + keys: Sequence[LiteLLM_VerificationToken] = await DeletedVerificationTokenRepository( + prisma_client + ).table.find_many( where=where, # type: ignore skip=skip, # type: ignore take=size, # type: ignore @@ -5797,7 +5824,7 @@ async def _list_key_helper( # Get total count of keys if use_deleted_table: - total_count = await DeletedVerificationTokenRepository(prisma_client).table.count( + total_count: int = await DeletedVerificationTokenRepository(prisma_client).table.count( where=where # type: ignore ) else: @@ -5811,13 +5838,15 @@ async def _list_key_helper( total_pages = -(-total_count // size) # Ceiling division # Fetch user information if expand includes "user" - user_map = {} + user_map: Mapping[str, LiteLLM_UserTable] = {} if expand and "user" in expand: user_ids = [key.user_id for key in keys if key.user_id] 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: Sequence[LiteLLM_UserTable] = await UserRepository(prisma_client).table.find_many( + where={"user_id": {"in": all_ids}} + ) user_map = {user.user_id: user for user in users} # Prepare response @@ -5851,9 +5880,12 @@ async def _list_key_helper( if return_full_object is True or (expand and "user" in expand): if use_deleted_table: # Use deleted key type to preserve deleted_at, deleted_by, etc. - key_list.append(LiteLLM_DeletedVerificationToken(**key_dict)) + key_list.append( + LiteLLM_DeletedVerificationToken(**cast(dict[str, Never], key_dict)) # cast-ok: revalidated + ) else: - key_list.append(UserAPIKeyAuth(**key_dict)) # Return full key object + # Return full key object + key_list.append(UserAPIKeyAuth(**cast(dict[str, Never], key_dict))) # cast-ok: revalidated else: _token = key_dict.get("token") key_list.append(cast(str, _token)) # Return only the token @@ -5880,8 +5912,8 @@ 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, + hashed_token: str | None, + prisma_client: PrismaClient, user_api_key_cache: UserApiKeyCache, route: str, ) -> None: @@ -5900,7 +5932,9 @@ 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: LiteLLM_VerificationToken | None = await VerificationTokenRepository( + prisma_client + ).table.find_unique(where={"token": hashed_token}) if target_key_row is None: raise HTTPException( status_code=404, @@ -5996,7 +6030,9 @@ 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: LiteLLM_VerificationToken | None = await VerificationTokenRepository( + prisma_client + ).table.find_unique(where={"token": hashed_token}) if existing_record is None: raise ProxyException( message="Key not found.", @@ -6026,7 +6062,7 @@ async def block_key( ) ) - record = await VerificationTokenRepository(prisma_client).table.update( + record: LiteLLM_VerificationToken | None = await VerificationTokenRepository(prisma_client).table.update( where={"token": hashed_token}, data={"blocked": True}, # type: ignore ) @@ -6107,7 +6143,9 @@ 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: LiteLLM_VerificationToken | None = await VerificationTokenRepository( + prisma_client + ).table.find_unique(where={"token": hashed_token}) if existing_record is None: raise ProxyException( message="Key not found.", @@ -6137,7 +6175,7 @@ async def unblock_key( ) ) - record = await VerificationTokenRepository(prisma_client).table.update( + record: LiteLLM_VerificationToken | None = await VerificationTokenRepository(prisma_client).table.update( where={"token": hashed_token}, data={"blocked": False}, # type: ignore ) @@ -6392,7 +6430,7 @@ def _validate_key_alias_format(key_alias: Optional[str]) -> None: async def _enforce_unique_key_alias( key_alias: Optional[str], - prisma_client: Any, + prisma_client: PrismaClient | None, existing_key_token: Optional[str] = None, ) -> None: """ @@ -6408,7 +6446,7 @@ 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} diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 5a289d22f99..bb0867705c0 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -13,7 +13,16 @@ #### ORGANIZATION MANAGEMENT #### -from typing import Annotated, Any, Dict, List, Mapping, Optional, Tuple +from typing import ( + Annotated, + List, + Mapping, + Optional, + Protocol, + Sequence, + Tuple, + TypeVar, +) import fastapi from fastapi import APIRouter, Depends, HTTPException, Request, status @@ -108,6 +117,91 @@ async def _verify_org_access( _BUDGET_SETTABLE_FIELDS = frozenset(LiteLLM_BudgetTable.model_fields.keys()) - {"budget_id"} _ORG_COLUMN_FIELDS = frozenset({"organization_alias", "models"}) +_PrismaModelT_co = TypeVar("_PrismaModelT_co", bound=BaseModel, covariant=True) + + +class _TableClient(Protocol[_PrismaModelT_co]): + async def find_unique( + self, + where: Mapping[str, object], + include: Mapping[str, object] | None = None, + ) -> _PrismaModelT_co | None: ... + + async def find_many( + self, + where: Mapping[str, object] | None = None, + include: Mapping[str, object] | None = None, + ) -> Sequence[_PrismaModelT_co]: ... + + async def create( + self, + data: Mapping[str, object], + include: Mapping[str, object] | None = None, + ) -> _PrismaModelT_co: ... + + async def update( + self, + where: Mapping[str, object], + data: Mapping[str, object], + include: Mapping[str, object] | None = None, + ) -> _PrismaModelT_co | None: ... + + async def upsert( + self, + where: Mapping[str, object], + data: Mapping[str, object], + include: Mapping[str, object] | None = None, + ) -> _PrismaModelT_co: ... + + async def delete( + self, + where: Mapping[str, object], + include: Mapping[str, object] | None = None, + ) -> _PrismaModelT_co | None: ... + + async def delete_many(self, where: Mapping[str, object] | None = None) -> int: ... + + +class _OrganizationTx(Protocol): + @property + def litellm_objectpermissiontable(self) -> "_TableClient[LiteLLM_ObjectPermissionTable]": ... + + @property + def litellm_budgettable(self) -> "_TableClient[LiteLLM_BudgetTableFull]": ... + + @property + def litellm_organizationtable(self) -> "_TableClient[LiteLLM_OrganizationTableWithMembers]": ... + + +def _organization_table(prisma_client: PrismaClient) -> "_TableClient[LiteLLM_OrganizationTableWithMembers]": + return OrganizationRepository(prisma_client).table + + +def _organization_membership_table( + prisma_client: PrismaClient, +) -> "_TableClient[LiteLLM_OrganizationMembershipTable]": + return OrganizationMembershipRepository(prisma_client).table + + +def _user_table(prisma_client: PrismaClient) -> "_TableClient[LiteLLM_UserTable]": + return UserRepository(prisma_client).table + + +def _budget_table(prisma_client: PrismaClient) -> "_TableClient[LiteLLM_BudgetTableFull]": + return BudgetRepository(prisma_client).table + + +def _object_permission_table(prisma_client: PrismaClient) -> "_TableClient[LiteLLM_ObjectPermissionTable]": + return ObjectPermissionRepository(prisma_client).table + + +def _team_table(prisma_client: PrismaClient) -> "_TableClient[LiteLLM_TeamTable]": + return TeamRepository(prisma_client).table + + +def _verification_token_table(prisma_client: PrismaClient) -> "_TableClient[LiteLLM_VerificationToken]": + return VerificationTokenRepository(prisma_client).table + def build_budget_write_data(budget_updates: Mapping[str, object], updated_by: str) -> Mapping[str, object]: """ @@ -263,10 +357,9 @@ async def new_organization( if user_api_key_dict.user_id is not None: try: - user_object = await UserRepository(prisma_client).table.find_unique( - where={"user_id": user_api_key_dict.user_id} - ) - user_object_correct_type = LiteLLM_UserTable(**user_object.model_dump()) + user_object = await _user_table(prisma_client).find_unique(where={"user_id": user_api_key_dict.user_id}) + if user_object is not None: + user_object_correct_type = LiteLLM_UserTable.model_validate(user_object.model_dump()) except Exception: pass @@ -279,19 +372,21 @@ async def new_organization( budget_params = LiteLLM_BudgetTable.model_fields.keys() # Only include Budget Params when creating an entry in litellm_budgettable - _json_data = data.json(exclude_none=True) + _json_data = _STR_OBJECT_DICT_ADAPTER.validate_python(data.json(exclude_none=True)) _budget_data = {k: v for k, v in _json_data.items() if k in budget_params} - budget_row = LiteLLM_BudgetTable(**_budget_data) + budget_row = LiteLLM_BudgetTable.model_validate(_budget_data) - new_budget = prisma_client.jsonify_object(budget_row.json(exclude_none=True)) + new_budget = _STR_OBJECT_DICT_ADAPTER.validate_python( + prisma_client.jsonify_object(budget_row.json(exclude_none=True)) + ) - _budget = await BudgetRepository(prisma_client).table.create( + _budget = await _budget_table(prisma_client).create( data={ - **new_budget, # type: ignore + **new_budget, "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, } - ) # type: ignore + ) data.budget_id = _budget.budget_id @@ -318,11 +413,13 @@ async def new_organization( for m in data.models: await can_user_call_model(m, llm_router=llm_router, user_object=user_object_correct_type) - organization_row = LiteLLM_OrganizationTable( - **data.json(exclude_none=True), - object_permission_id=object_permission_id, - 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, + organization_row = LiteLLM_OrganizationTable.model_validate( + { + **data.json(exclude_none=True), + "object_permission_id": object_permission_id, + "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, + } ) for field in LiteLLM_ManagementEndpoint_MetadataFields: @@ -333,11 +430,13 @@ async def new_organization( value=getattr(data, field), ) - new_organization_row = prisma_client.jsonify_object(organization_row.json(exclude_none=True)) + new_organization_row = _STR_OBJECT_DICT_ADAPTER.validate_python( + prisma_client.jsonify_object(organization_row.json(exclude_none=True)) + ) verbose_proxy_logger.info(f"new_organization_row: {json.dumps(new_organization_row, indent=2)}") - response = await OrganizationRepository(prisma_client).table.create( + response = await _organization_table(prisma_client).create( data={ - **new_organization_row, # type: ignore + **new_organization_row, }, include={"litellm_budget_table": True}, ) @@ -382,7 +481,7 @@ async def get_organization_daily_activity( # Restrict non-proxy-admins to only organizations where they are org_admin if not _user_has_admin_view(user_api_key_dict): - memberships = await OrganizationMembershipRepository(prisma_client).table.find_many( + memberships = await _organization_membership_table(prisma_client).find_many( where={"user_id": user_api_key_dict.user_id} ) admin_org_ids = [m.organization_id for m in memberships if m.user_role == LitellmUserRoles.ORG_ADMIN.value] @@ -399,11 +498,19 @@ async def get_organization_daily_activity( ) # Fetch organization aliases for metadata - where_condition = {} + where_condition: dict[ + str, object + ] = {} # mutable-ok: the organization_id filter is inserted below only when org ids were requested if org_ids_list: where_condition["organization_id"] = {"in": list(org_ids_list)} - org_aliases = await OrganizationRepository(prisma_client).table.find_many(where=where_condition) - org_alias_metadata = {o.organization_id: {"organization_alias": o.organization_alias} for o in org_aliases} + org_aliases = await _organization_table(prisma_client).find_many(where=where_condition) + org_alias_metadata: dict[ + str, dict[str, object] + ] = { # mutable-ok: get_daily_activity's entity_metadata_field parameter is typed Optional[Dict[str, Dict[str, object]]] + o.organization_id: {"organization_alias": o.organization_alias} + for o in org_aliases + if o.organization_id is not None + } # Query daily activity for organizations return await get_daily_activity( @@ -436,7 +543,7 @@ async def _set_object_permission( return None if data.object_permission is not None: - created_object_permission = await ObjectPermissionRepository(prisma_client).table.create( + created_object_permission = await _object_permission_table(prisma_client).create( data=data.object_permission.model_dump(exclude_none=True), ) del data.object_permission @@ -474,7 +581,7 @@ async def update_organization( ) # Transform UI payload to expected format - raw_data = await request.json() + raw_data = _STR_OBJECT_DICT_ADAPTER.validate_python(await request.json()) raw_data_with_flat_budget_fields = handle_nested_budget_structure_in_organization_update_request(raw_data) # Create validated data model @@ -510,22 +617,24 @@ async def update_organization( prisma_client=prisma_client, ) - existing_organization_row = await OrganizationRepository(prisma_client).table.find_unique( + existing_organization_row = await _organization_table(prisma_client).find_unique( where={"organization_id": data.organization_id}, ) if existing_organization_row is None: raise ValueError(f"Organization not found for organization_id={data.organization_id}") - updated_organization_row_json = data.model_dump(exclude_none=True) + updated_organization_row_json = _STR_OBJECT_DICT_ADAPTER.validate_python(data.model_dump(exclude_none=True)) # Merge metadata from existing organization with updated metadata if updated_organization_row_json.get("metadata") is not None: - existing_metadata = existing_organization_row.metadata or {} - updated_metadata = updated_organization_row_json.get("metadata", {}) + existing_metadata = _STR_OBJECT_DICT_ADAPTER.validate_python(existing_organization_row.metadata or {}) + updated_metadata = _STR_OBJECT_DICT_ADAPTER.validate_python(updated_organization_row_json.get("metadata", {})) merged_metadata = _update_dictionary(existing_dict=existing_metadata.copy(), new_dict=updated_metadata) updated_organization_row_json["metadata"] = merged_metadata - updated_organization_row = prisma_client.jsonify_object(updated_organization_row_json) + updated_organization_row = _STR_OBJECT_DICT_ADAPTER.validate_python( + prisma_client.jsonify_object(updated_organization_row_json) + ) if data.object_permission is not None: updated_organization_row = await handle_update_object_permission( data_json=updated_organization_row, @@ -534,12 +643,16 @@ async def update_organization( # Handle budget updates if budget fields are provided budget_fields = { - k: v for k, v in data.model_dump().items() if k in LiteLLM_BudgetTable.model_fields.keys() and v is not None + k: v + for k, v in _STR_OBJECT_DICT_ADAPTER.validate_python(data.model_dump()).items() + if k in LiteLLM_BudgetTable.model_fields.keys() and v is not None } if budget_fields and existing_organization_row.budget_id: await update_budget( - budget_obj=BudgetNewRequest(budget_id=existing_organization_row.budget_id, **budget_fields), + budget_obj=BudgetNewRequest.model_validate( + {"budget_id": existing_organization_row.budget_id, **budget_fields} + ), user_api_key_dict=user_api_key_dict, ) @@ -547,7 +660,7 @@ async def update_organization( for field in LiteLLM_BudgetTable.model_fields.keys(): updated_organization_row.pop(field, None) - response = await OrganizationRepository(prisma_client).table.update( + response = await _organization_table(prisma_client).update( where={"organization_id": data.organization_id}, data=updated_organization_row, include={"members": True, "teams": True, "litellm_budget_table": True}, @@ -557,9 +670,9 @@ async def update_organization( async def handle_update_object_permission( - data_json: dict, + data_json: dict[str, object], existing_organization_row: LiteLLM_OrganizationTable, -) -> dict: +) -> dict[str, object]: """ Handle the update of object permission for an organization. @@ -665,7 +778,7 @@ async def update_organization_v2( prisma_client=prisma_client, ) - existing_organization_row = await OrganizationRepository(prisma_client).table.find_unique( + existing_organization_row = await _organization_table(prisma_client).find_unique( where={"organization_id": organization_id}, ) if existing_organization_row is None: @@ -698,14 +811,17 @@ async def update_organization_v2( else ({"object_permission_id": None} if object_permission_cleared else {}) ) - organization_write_data = prisma_client.jsonify_object( - { - **org_column_updates, - **object_permission_write, - "updated_by": user_api_key_dict.user_id, - } + organization_write_data = _STR_OBJECT_DICT_ADAPTER.validate_python( + prisma_client.jsonify_object( + { + **org_column_updates, + **object_permission_write, + "updated_by": user_api_key_dict.user_id, + } + ) ) + tx: _OrganizationTx async with prisma_client.db.tx() as tx: if object_permission_upsert is not None: await tx.litellm_objectpermissiontable.upsert( @@ -718,8 +834,10 @@ async def update_organization_v2( if budget_updates: await tx.litellm_budgettable.update( where={"budget_id": existing_organization_row.budget_id}, - data=prisma_client.jsonify_object( - dict(build_budget_write_data(budget_updates, user_api_key_dict.user_id)) + data=_STR_OBJECT_DICT_ADAPTER.validate_python( + prisma_client.jsonify_object( + dict(build_budget_write_data(budget_updates, user_api_key_dict.user_id)) + ) ), ) response = await tx.litellm_organizationtable.update( @@ -762,18 +880,18 @@ async def delete_organization( detail={"error": "Only proxy admins can delete organizations"}, ) - deleted_orgs = [] + deleted_orgs: list[ + LiteLLM_OrganizationTableWithMembers + ] = [] # mutable-ok: each deleted organization is appended by the loop below for organization_id in data.organization_ids: # delete all teams in the organization - await TeamRepository(prisma_client).table.delete_many(where={"organization_id": organization_id}) + await _team_table(prisma_client).delete_many(where={"organization_id": organization_id}) # delete all members in the organization - await OrganizationMembershipRepository(prisma_client).table.delete_many( - where={"organization_id": organization_id} - ) + await _organization_membership_table(prisma_client).delete_many(where={"organization_id": organization_id}) # delete all keys in the organization - await VerificationTokenRepository(prisma_client).table.delete_many(where={"organization_id": organization_id}) + await _verification_token_table(prisma_client).delete_many(where={"organization_id": organization_id}) # delete the organization - deleted_org = await OrganizationRepository(prisma_client).table.delete( + deleted_org = await _organization_table(prisma_client).delete( where={"organization_id": organization_id}, include={"members": True, "teams": True, "litellm_budget_table": True}, ) @@ -836,7 +954,9 @@ async def list_organization( ) # Build where conditions based on provided filters - where_conditions: Dict[str, Any] = {} + where_conditions: dict[ + str, object + ] = {} # mutable-ok: org_id / org_alias filter keys are inserted below as the query params arrive if org_id: where_conditions["organization_id"] = org_id @@ -848,14 +968,15 @@ async def list_organization( } # if proxy admin or admin viewer - get all orgs (with optional filters) + response: Sequence[LiteLLM_OrganizationTableWithMembers] if _user_has_admin_view(user_api_key_dict): - response = await OrganizationRepository(prisma_client).table.find_many( + response = await _organization_table(prisma_client).find_many( where=where_conditions if where_conditions else None, include={"litellm_budget_table": True, "members": True, "teams": True}, ) # if internal user - get orgs they are a member of (with optional filters) else: - org_memberships = await OrganizationMembershipRepository(prisma_client).table.find_many( + org_memberships = await _organization_membership_table(prisma_client).find_many( where={"user_id": user_api_key_dict.user_id} ) membership_org_ids = [membership.organization_id for membership in org_memberships] @@ -869,7 +990,7 @@ async def list_organization( response = [] else: where_conditions["organization_id"] = org_id - response = await OrganizationRepository(prisma_client).table.find_many( + response = await _organization_table(prisma_client).find_many( where=where_conditions, include={ "litellm_budget_table": True, @@ -880,7 +1001,7 @@ async def list_organization( else: # Filter by membership and any additional filters where_conditions["organization_id"] = {"in": membership_org_ids} - response = await OrganizationRepository(prisma_client).table.find_many( + response = await _organization_table(prisma_client).find_many( where=where_conditions, include={ "litellm_budget_table": True, @@ -920,9 +1041,7 @@ async def info_organization( prisma_client=prisma_client, ) - response: Optional[LiteLLM_OrganizationTableWithMembers] = await OrganizationRepository( - prisma_client - ).table.find_unique( + response: LiteLLM_OrganizationTableWithMembers | None = await _organization_table(prisma_client).find_unique( where={"organization_id": organization_id}, include={ "litellm_budget_table": True, @@ -939,7 +1058,7 @@ async def info_organization( if response is None: raise HTTPException(status_code=404, detail={"error": "Organization not found"}) - response_pydantic_obj = LiteLLM_OrganizationTableWithMembers(**response.model_dump()) + response_pydantic_obj = LiteLLM_OrganizationTableWithMembers.model_validate(response.model_dump()) return response_pydantic_obj @@ -975,7 +1094,7 @@ async def deprecated_info_organization( prisma_client=prisma_client, ) - response = await OrganizationRepository(prisma_client).table.find_many( + response = await _organization_table(prisma_client).find_many( where={"organization_id": {"in": data.organizations}}, include={"litellm_budget_table": True}, ) @@ -1052,7 +1171,7 @@ async def organization_member_add( ) # Check if organization exists - existing_organization_row = await OrganizationRepository(prisma_client).table.find_unique( + existing_organization_row = await _organization_table(prisma_client).find_unique( where={"organization_id": data.organization_id} ) if existing_organization_row is None: @@ -1125,7 +1244,7 @@ async def find_member_if_email(user_email: str, prisma_client: PrismaClient) -> "error": f"Unique user not found for user_email={user_email}. Potential duplicate OR non-existent user_email in LiteLLM_UserTable. Use 'user_id' instead." }, ) - existing_user_email_row_pydantic = LiteLLM_UserTable(**existing_user_email_row.model_dump()) + existing_user_email_row_pydantic = LiteLLM_UserTable.model_validate(existing_user_email_row.model_dump()) return existing_user_email_row_pydantic @@ -1163,7 +1282,7 @@ async def organization_member_update( ) # Check if organization exists - existing_organization_row = await OrganizationRepository(prisma_client).table.find_unique( + existing_organization_row = await _organization_table(prisma_client).find_unique( where={"organization_id": data.organization_id} ) if existing_organization_row is None: @@ -1180,7 +1299,7 @@ async def organization_member_update( data.user_id = existing_user_email_row.user_id try: - existing_organization_membership = await OrganizationMembershipRepository(prisma_client).table.find_unique( + existing_organization_membership = await _organization_membership_table(prisma_client).find_unique( where={ "user_id_organization_id": { "user_id": data.user_id, @@ -1205,7 +1324,7 @@ async def organization_member_update( # org-scoped operations. An org-admin of any org could otherwise # alter a PROXY_ADMIN user's per-org role, which has downstream # effects on admin UI filtering and scope derivation. - target_user_row = await UserRepository(prisma_client).table.find_unique(where={"user_id": data.user_id}) + target_user_row = await _user_table(prisma_client).find_unique(where={"user_id": data.user_id}) if target_user_row is not None and getattr(target_user_row, "user_role", None) in ( LitellmUserRoles.PROXY_ADMIN.value, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, @@ -1222,7 +1341,7 @@ async def organization_member_update( # Update member role if data.role is not None: - await OrganizationMembershipRepository(prisma_client).table.update( + await _organization_membership_table(prisma_client).update( where={ "user_id_organization_id": { "user_id": data.user_id, @@ -1245,7 +1364,7 @@ async def organization_member_update( ) # update organization membership with new budget_id - await OrganizationMembershipRepository(prisma_client).table.update( + await _organization_membership_table(prisma_client).update( where={ "user_id_organization_id": { "user_id": data.user_id, @@ -1254,9 +1373,7 @@ async def organization_member_update( }, data={"budget_id": budget_id}, ) - final_organization_membership: Optional[BaseModel] = await OrganizationMembershipRepository( - prisma_client - ).table.find_unique( + final_organization_membership = await _organization_membership_table(prisma_client).find_unique( where={ "user_id_organization_id": { "user_id": data.user_id, @@ -1272,8 +1389,8 @@ async def organization_member_update( detail={"error": f"Member not found in organization={data.organization_id} for user_id={data.user_id}"}, ) - final_organization_membership_pydantic = LiteLLM_OrganizationMembershipTable( - **final_organization_membership.model_dump(exclude_none=True) + final_organization_membership_pydantic = LiteLLM_OrganizationMembershipTable.model_validate( + final_organization_membership.model_dump(exclude_none=True) ) return final_organization_membership_pydantic except Exception as e: @@ -1315,7 +1432,7 @@ async def organization_member_delete( existing_user_email_row = await find_member_if_email(data.user_email, prisma_client) data.user_id = existing_user_email_row.user_id - member_to_delete = await OrganizationMembershipRepository(prisma_client).table.delete( + member_to_delete = await _organization_membership_table(prisma_client).delete( where={ "user_id_organization_id": { "user_id": data.user_id, @@ -1374,16 +1491,16 @@ async def add_member_to_organization( _returned_user = await prisma_client.insert_data(data=new_user_defaults, table_name="user") # type: ignore if _returned_user is not None: - user_object = LiteLLM_UserTable(**_returned_user.model_dump()) + user_object = LiteLLM_UserTable.model_validate(_returned_user.model_dump()) elif existing_user_email_row is not None and len(existing_user_email_row) > 1: raise HTTPException( status_code=400, detail={"error": "Multiple users with this email found in db. Please use 'user_id' instead."}, ) elif existing_user_email_row is not None: - user_object = LiteLLM_UserTable(**existing_user_email_row.model_dump()) + user_object = LiteLLM_UserTable.model_validate(existing_user_email_row.model_dump()) elif existing_user_id_row is not None: - user_object = LiteLLM_UserTable(**existing_user_id_row.model_dump()) + user_object = LiteLLM_UserTable.model_validate(existing_user_id_row.model_dump()) else: raise HTTPException( status_code=404, @@ -1396,14 +1513,16 @@ async def add_member_to_organization( ) # Add user to organization - _organization_membership = await OrganizationMembershipRepository(prisma_client).table.create( + _organization_membership = await _organization_membership_table(prisma_client).create( data={ "organization_id": organization_id, "user_id": user_object.user_id, "user_role": member.role, } ) - organization_membership = LiteLLM_OrganizationMembershipTable(**_organization_membership.model_dump()) + organization_membership = LiteLLM_OrganizationMembershipTable.model_validate( + _organization_membership.model_dump() + ) return user_object, organization_membership except Exception as e: diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 59b0cbc4ae7..457087bc2a3 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -14,11 +14,25 @@ import math import traceback from datetime import datetime, timezone -from typing import Annotated, Any, Dict, List, Mapping, Optional, Tuple, Union, cast +from typing import ( + AbstractSet, + Annotated, + Dict, + List, + Mapping, + Optional, + Protocol, + Sequence, + Tuple, + TypeVar, + Union, + cast, +) import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Request, status from pydantic import BaseModel +from typing_extensions import Required, TypedDict import litellm from litellm._logging import verbose_proxy_logger @@ -28,13 +42,20 @@ from litellm.proxy._types import ( UI_TEAM_ID, BlockTeamRequest, + BudgetLimitEntry, CommonProxyErrors, DeleteTeamRequest, + LiteLLM_AccessGroupTable, LiteLLM_AuditLogs, + LiteLLM_BudgetTable, + LiteLLM_BudgetTableFull, LiteLLM_DeletedTeamTable, LiteLLM_ManagementEndpoint_MetadataFields, LiteLLM_ManagementEndpoint_MetadataFields_Premium, LiteLLM_ModelTable, + LiteLLM_ObjectPermissionBase, + LiteLLM_ObjectPermissionTable, + LiteLLM_OrganizationMembershipTable, LiteLLM_OrganizationTable, LiteLLM_OrganizationTableWithMembers, LiteLLM_TeamMembership, @@ -78,6 +99,7 @@ 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 +128,7 @@ 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 ( @@ -141,7 +163,7 @@ router = APIRouter() -def _sanitize_for_log(value: Any) -> str: +def _sanitize_for_log(value: object) -> str: """Strip CR/LF from user-controlled values to prevent log injection.""" try: text = str(value) @@ -150,10 +172,375 @@ def _sanitize_for_log(value: Any) -> str: return text.replace("\r", "").replace("\n", "") +class _TeamRowDump(TypedDict, total=False): + team_id: Required[str] + team_alias: str | None + organization_id: str | None + admins: list[str] # mutable-ok: splatted into LiteLLM_TeamTable(**dump), whose pydantic field is declared list + members: list[str] # mutable-ok: splatted into LiteLLM_TeamTable(**dump), whose pydantic field is declared list + members_with_roles: list[ + Member + ] # mutable-ok: splatted into LiteLLM_TeamTable(**dump), whose pydantic field is declared list + team_member_permissions: Optional[ + list[str] + ] # mutable-ok: splatted into LiteLLM_TeamTable(**dump), whose pydantic field is declared list + metadata: Optional[ + dict[str, object] + ] # mutable-ok: splatted into LiteLLM_TeamTable(**dump), whose pydantic field is declared dict + tpm_limit: int | None + rpm_limit: int | None + max_budget: float | None + soft_budget: float | None + budget_duration: str | None + budget_limits: Optional[ + list[BudgetLimitEntry] + ] # mutable-ok: splatted into LiteLLM_TeamTable(**dump), whose pydantic field is declared list + models: list[str] # mutable-ok: splatted into LiteLLM_TeamTable(**dump), whose pydantic field is declared list + blocked: bool + router_settings: Optional[ + dict[str, object] + ] # mutable-ok: splatted into LiteLLM_TeamTable(**dump), whose pydantic field is declared dict + access_group_ids: Optional[ + list[str] + ] # mutable-ok: splatted into LiteLLM_TeamTable(**dump), whose pydantic field is declared list + default_team_member_models: Optional[ + list[str] + ] # mutable-ok: splatted into LiteLLM_TeamTable(**dump), whose pydantic field is declared list + spend: float | None + max_parallel_requests: int | None + budget_reset_at: datetime | None + model_id: int | None + model_spend: Optional[ + dict[str, float] + ] # mutable-ok: splatted into LiteLLM_TeamTable(**dump), whose pydantic field is declared dict + model_max_budget: Optional[ + dict[str, float] + ] # mutable-ok: splatted into LiteLLM_TeamTable(**dump), whose pydantic field is declared dict + policies: Optional[ + list[str] + ] # mutable-ok: splatted into LiteLLM_TeamTable(**dump), whose pydantic field is declared list + allow_team_guardrail_config: bool | None + litellm_model_table: LiteLLM_ModelTable | None + object_permission: LiteLLM_ObjectPermissionTable | None + object_permission_id: str | None + updated_at: datetime | None + created_at: datetime | None + + +class _TeamMembershipRowDump(TypedDict, total=False): + user_id: Required[str] + team_id: Required[str] + budget_id: str | None + spend: float | None + litellm_budget_table: LiteLLM_BudgetTableFull | LiteLLM_BudgetTable | None + + +class _UserRowDump(TypedDict, total=False): + user_id: Required[str] + user_alias: str | None + team_id: str | None + sso_user_id: str | None + organization_id: str | None + object_permission_id: str | None + password: str | None + teams: list[str] # mutable-ok: splatted into LiteLLM_UserTable(**dump), whose pydantic field is declared list + user_role: str | None + max_budget: float | None + spend: float + user_email: str | None + models: list[str] # mutable-ok: splatted into LiteLLM_UserTable(**dump), whose pydantic field is declared list + metadata: Optional[ + dict[str, object] + ] # mutable-ok: splatted into LiteLLM_UserTable(**dump), whose pydantic field is declared dict + max_parallel_requests: int | None + tpm_limit: int | None + rpm_limit: int | None + budget_duration: str | None + budget_reset_at: datetime | None + allowed_cache_controls: list[ + str + ] # mutable-ok: splatted into LiteLLM_UserTable(**dump), whose pydantic field is declared list + policies: list[str] # mutable-ok: splatted into LiteLLM_UserTable(**dump), whose pydantic field is declared list + model_spend: Optional[ + dict[str, float] + ] # mutable-ok: splatted into LiteLLM_UserTable(**dump), whose pydantic field is declared dict + model_max_budget: Optional[ + dict[str, float] + ] # mutable-ok: splatted into LiteLLM_UserTable(**dump), whose pydantic field is declared dict + created_at: datetime | None + updated_at: datetime | None + object_permission: LiteLLM_ObjectPermissionTable | None + + +class _OrgRowDump(TypedDict, total=False): + organization_id: str | None + organization_alias: str | None + budget_id: Required[str] + spend: float + metadata: Optional[ + dict[str, object] + ] # mutable-ok: splatted into LiteLLM_OrganizationTableWithMembers(**dump), whose pydantic field is declared dict + models: list[ + str + ] # mutable-ok: splatted into LiteLLM_OrganizationTableWithMembers(**dump), whose pydantic field is declared list + model_spend: Optional[ + dict[str, float] + ] # mutable-ok: splatted into LiteLLM_OrganizationTableWithMembers(**dump), whose pydantic field is declared dict + created_by: Required[str] + updated_by: Required[str] + users: Optional[ + list[LiteLLM_UserTable] + ] # mutable-ok: splatted into LiteLLM_OrganizationTableWithMembers(**dump), whose pydantic field is declared list + litellm_budget_table: LiteLLM_BudgetTable | None + object_permission: LiteLLM_ObjectPermissionTable | None + object_permission_id: str | None + members: list[ + LiteLLM_OrganizationMembershipTable + ] # mutable-ok: splatted into LiteLLM_OrganizationTableWithMembers(**dump), whose pydantic field is declared list + teams: list[ + LiteLLM_TeamTable + ] # mutable-ok: splatted into LiteLLM_OrganizationTableWithMembers(**dump), whose pydantic field is declared list + created_at: Required[datetime] + updated_at: Required[datetime] + + +class _TeamPatchDump(TypedDict, total=False): + team_alias: str | None + organization_id: str | None + metadata: Optional[ + dict[str, object] + ] # mutable-ok: splatted into UpdateTeamRequest(**patch), whose pydantic field is declared dict + tpm_limit: int | None + rpm_limit: int | None + max_budget: float | None + soft_budget: float | None + models: Optional[ + list[str] + ] # mutable-ok: splatted into UpdateTeamRequest(**patch), whose pydantic field is declared list + blocked: bool | None + budget_duration: str | None + tags: Optional[ + list[str] + ] # mutable-ok: splatted into UpdateTeamRequest(**patch), whose pydantic field is declared list + model_aliases: Optional[ + dict[str, str] + ] # mutable-ok: splatted into UpdateTeamRequest(**patch), whose pydantic field is declared dict + guardrails: Optional[ + list[str] + ] # mutable-ok: splatted into UpdateTeamRequest(**patch), whose pydantic field is declared list + policies: Optional[ + list[str] + ] # mutable-ok: splatted into UpdateTeamRequest(**patch), whose pydantic field is declared list + object_permission: LiteLLM_ObjectPermissionBase | None + disable_global_guardrails: bool | None + team_member_budget: float | None + team_member_budget_duration: str | None + team_member_rpm_limit: int | None + team_member_tpm_limit: int | None + team_member_key_duration: str | None + allowed_passthrough_routes: Optional[ + list[str] + ] # mutable-ok: splatted into UpdateTeamRequest(**patch), whose pydantic field is declared list + secret_manager_settings: Optional[ + dict[str, object] + ] # mutable-ok: splatted into UpdateTeamRequest(**patch), whose pydantic field is declared dict + prompts: Optional[ + list[str] + ] # mutable-ok: splatted into UpdateTeamRequest(**patch), whose pydantic field is declared list + model_rpm_limit: Optional[ + dict[str, int] + ] # mutable-ok: splatted into UpdateTeamRequest(**patch), whose pydantic field is declared dict + model_tpm_limit: Optional[ + dict[str, int] + ] # mutable-ok: splatted into UpdateTeamRequest(**patch), whose pydantic field is declared dict + mcp_rpm_limit: Optional[ + dict[str, int] + ] # mutable-ok: splatted into UpdateTeamRequest(**patch), whose pydantic field is declared dict + router_settings: Optional[ + dict[str, object] + ] # mutable-ok: splatted into UpdateTeamRequest(**patch), whose pydantic field is declared dict + access_group_ids: Optional[ + list[str] + ] # mutable-ok: splatted into UpdateTeamRequest(**patch), whose pydantic field is declared list + budget_limits: Optional[ + list[BudgetLimitEntry] + ] # mutable-ok: splatted into UpdateTeamRequest(**patch), whose pydantic field is declared list + default_team_member_models: Optional[ + list[str] + ] # mutable-ok: splatted into UpdateTeamRequest(**patch), whose pydantic field is declared list + + +_RowT = TypeVar("_RowT") + + +class _PrismaTableProtocol(Protocol[_RowT]): + async def find_unique( + self, + where: Mapping[str, object], + include: Mapping[str, object] | None = None, + ) -> _RowT | None: ... + + async def find_first( + self, + where: Mapping[str, object] | None = None, + order: Mapping[str, str] | None = None, + ) -> _RowT | None: ... + + async def find_many( + self, + where: Mapping[str, object] | None = None, + include: Mapping[str, object] | None = None, + order: Mapping[str, str] | Sequence[Mapping[str, str]] | None = None, + take: int | None = None, + skip: int | None = None, + cursor: Mapping[str, object] | None = None, + ) -> list[_RowT]: ... # mutable-ok: prisma returns a real list that callers hand to List-typed helpers + + async def create( + self, + data: Mapping[str, object], + include: Mapping[str, object] | None = None, + ) -> _RowT: ... + + async def create_many( + self, + data: Sequence[Mapping[str, object]], + skip_duplicates: bool | None = None, + ) -> int: ... + + async def update( + self, + where: Mapping[str, object], + data: Mapping[str, object], + include: Mapping[str, object] | None = None, + ) -> _RowT | None: ... + + async def update_many( + self, + where: Mapping[str, object], + data: Mapping[str, object], + ) -> int: ... + + async def upsert( + self, + where: Mapping[str, object], + data: Mapping[str, object], + include: Mapping[str, object] | None = None, + ) -> _RowT: ... + + async def delete_many( + self, + where: Mapping[str, object] | None = None, + ) -> int: ... + + async def count( + self, + where: Mapping[str, object] | None = None, + ) -> int: ... + + +def _team_table(prisma_client: PrismaClient) -> _PrismaTableProtocol[LiteLLM_TeamTable]: + return cast( # cast-ok: prisma table is exposed untyped by the repository seam + _PrismaTableProtocol[LiteLLM_TeamTable], TeamRepository(prisma_client).table + ) + + +def _team_membership_table(prisma_client: PrismaClient) -> _PrismaTableProtocol[LiteLLM_TeamMembership]: + return cast( # cast-ok: prisma table is exposed untyped by the repository seam + _PrismaTableProtocol[LiteLLM_TeamMembership], TeamMembershipRepository(prisma_client).table + ) + + +def _user_table(prisma_client: PrismaClient) -> _PrismaTableProtocol[LiteLLM_UserTable]: + return cast( # cast-ok: prisma table is exposed untyped by the repository seam + _PrismaTableProtocol[LiteLLM_UserTable], UserRepository(prisma_client).table + ) + + +def _verification_token_table(prisma_client: PrismaClient) -> _PrismaTableProtocol[LiteLLM_VerificationToken]: + return cast( # cast-ok: prisma table is exposed untyped by the repository seam + _PrismaTableProtocol[LiteLLM_VerificationToken], VerificationTokenRepository(prisma_client).table + ) + + +def _model_table(prisma_client: PrismaClient) -> _PrismaTableProtocol[LiteLLM_ModelTable]: + return cast( # cast-ok: prisma table is exposed untyped by the repository seam + _PrismaTableProtocol[LiteLLM_ModelTable], ModelTableRepository(prisma_client).table + ) + + +def _organization_table(prisma_client: PrismaClient) -> _PrismaTableProtocol[LiteLLM_OrganizationTableWithMembers]: + return cast( # cast-ok: prisma table is exposed untyped by the repository seam + _PrismaTableProtocol[LiteLLM_OrganizationTableWithMembers], OrganizationRepository(prisma_client).table + ) + + +def _organization_membership_table( + prisma_client: PrismaClient, +) -> _PrismaTableProtocol[LiteLLM_OrganizationMembershipTable]: + return cast( # cast-ok: prisma table is exposed untyped by the repository seam + _PrismaTableProtocol[LiteLLM_OrganizationMembershipTable], OrganizationMembershipRepository(prisma_client).table + ) + + +def _deleted_team_table(prisma_client: PrismaClient) -> _PrismaTableProtocol[LiteLLM_DeletedTeamTable]: + return cast( # cast-ok: prisma table is exposed untyped by the repository seam + _PrismaTableProtocol[LiteLLM_DeletedTeamTable], DeletedTeamRepository(prisma_client).table + ) + + +def _access_group_table(prisma_client: PrismaClient) -> _PrismaTableProtocol[LiteLLM_AccessGroupTable]: + return cast( # cast-ok: prisma table is exposed untyped by the repository seam + _PrismaTableProtocol[LiteLLM_AccessGroupTable], AccessGroupRepository(prisma_client).table + ) + + +def _budget_table(prisma_client: PrismaClient) -> _PrismaTableProtocol[LiteLLM_BudgetTableFull]: + return cast( # cast-ok: prisma table is exposed untyped by the repository seam + _PrismaTableProtocol[LiteLLM_BudgetTableFull], BudgetRepository(prisma_client).table + ) + + +class _TeamPageArgs(TypedDict, total=False): + take: int + order: Mapping[str, str] + cursor: Mapping[str, str] + skip: int + + +class _TokenTeamGroupByRow(TypedDict): + team_id: str + _count: Mapping[str, int] + + +def _team_row_dump(team_row: BaseModel) -> _TeamRowDump: + return cast( # cast-ok: team row dump mirrors LiteLLM_TeamTable constructor fields + _TeamRowDump, team_row.model_dump() + ) + + +def _team_membership_row_dump(membership_row: BaseModel) -> _TeamMembershipRowDump: + return cast( # cast-ok: membership row dump mirrors LiteLLM_TeamMembership constructor fields + _TeamMembershipRowDump, membership_row.model_dump() + ) + + +def _user_row_dump(user_row: BaseModel) -> _UserRowDump: + return cast( # cast-ok: user row dump mirrors LiteLLM_UserTable constructor fields + _UserRowDump, user_row.model_dump() + ) + + +def _org_row_dump(org_row: BaseModel) -> _OrgRowDump: + return cast( # cast-ok: org row dump mirrors LiteLLM_OrganizationTableWithMembers constructor fields + _OrgRowDump, org_row.model_dump() + ) + + 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, ) -> None: """ Refresh the in-memory cached team object after a DB write. @@ -171,7 +558,7 @@ async def _refresh_cached_team( """ await _cache_team_object( team_id=team_row.team_id, - team_table=LiteLLM_TeamTableCachedObj(**team_row.model_dump()), + team_table=LiteLLM_TeamTableCachedObj(**_team_row_dump(team_row)), user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) @@ -274,9 +661,12 @@ async def create_team_member_budget_table( if team_member_budget_duration is not None: budget_request.budget_duration = team_member_budget_duration - team_member_budget_table = await new_budget( - budget_obj=budget_request, - user_api_key_dict=user_api_key_dict, + team_member_budget_table = cast( # cast-ok: new_budget returns the created prisma budget row + LiteLLM_BudgetTableFull, + await new_budget( + budget_obj=budget_request, + user_api_key_dict=user_api_key_dict, + ), ) # Add team_member_budget_id as metadata field to team table @@ -322,9 +712,12 @@ async def upsert_team_member_budget_table( if team_member_budget_duration is not None: budget_request.budget_duration = team_member_budget_duration - budget_row = await update_budget( - budget_obj=budget_request, - user_api_key_dict=user_api_key_dict, + budget_row = cast( # cast-ok: update_budget returns the updated prisma budget row + LiteLLM_BudgetTableFull, + await update_budget( + budget_obj=budget_request, + user_api_key_dict=user_api_key_dict, + ), ) verbose_proxy_logger.info( f"Updated team member budget table: {budget_row.budget_id}, with team_member_budget={team_member_budget}, team_member_rpm_limit={team_member_rpm_limit}, team_member_tpm_limit={team_member_tpm_limit}" @@ -395,7 +788,9 @@ async def clear_team_member_budget_fields( @staticmethod async def backfill_team_member_budget_entries( team_id: str, - members_with_roles: List[Union[Member, dict]], + members_with_roles: Sequence[ + Member | dict[str, object] + ], # mutable-ok: the raw-member alternative is narrowed with isinstance(m, dict) below team_member_budget_id: str, prisma_client: PrismaClient, ) -> None: @@ -414,7 +809,7 @@ async def backfill_team_member_budget_entries( 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 = await _team_membership_table(prisma_client).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. @@ -433,7 +828,7 @@ async def backfill_team_member_budget_entries( ) if missing: - await TeamMembershipRepository(prisma_client).table.create_many( + await _team_membership_table(prisma_client).create_many( data=missing, skip_duplicates=True, # safety net against concurrent races ) @@ -447,7 +842,7 @@ async def backfill_team_member_budget_entries( # 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 = await _team_membership_table(prisma_client).update_many( where={"team_id": team_id, "budget_id": None}, data={"budget_id": team_member_budget_id}, ) @@ -460,7 +855,7 @@ async def backfill_team_member_budget_entries( ) -def _get_default_team_param(field: str) -> Any: +def _get_default_team_param(field: str) -> object | None: """ Returns a default value for the given field from litellm.default_team_params config. Returns None if no default is configured. @@ -503,14 +898,14 @@ 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 = await _team_membership_table(prisma_client).find_many( where=where_obj, include={"litellm_budget_table": True}, ) returned_tm: List[LiteLLM_TeamMembership] = [] for tm in team_memberships: - returned_tm.append(LiteLLM_TeamMembership(**tm.model_dump())) + returned_tm.append(LiteLLM_TeamMembership(**_team_membership_row_dump(tm))) return returned_tm @@ -765,14 +1160,14 @@ 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( + teams = await _team_table(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(**team.model_dump())) + team_objs.append(LiteLLM_TeamTable(**_team_row_dump(team))) check_org_team_model_specific_limits( teams=team_objs, @@ -790,7 +1185,7 @@ async def _check_user_team_limits( data: Union[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. @@ -1051,7 +1446,7 @@ async def new_team( ) # Check if license is over limit - total_teams = await TeamRepository(prisma_client).table.count() + total_teams = await _team_table(prisma_client).count() if total_teams and _license_check.is_team_count_over_limit(team_count=total_teams): raise HTTPException( status_code=403, @@ -1153,7 +1548,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 = await _model_table(prisma_client).create( {**litellm_modeltable.json(exclude_none=True)} # type: ignore ) # type: ignore @@ -1254,7 +1649,7 @@ async def new_team( complete_team_data_dict = prisma_client.jsonify_team_object(db_data=complete_team_data_dict) - team_row: LiteLLM_TeamTable = await TeamRepository(prisma_client).table.create( + team_row: LiteLLM_TeamTable = await _team_table(prisma_client).create( data=complete_team_data_dict, include={"litellm_model_table": True}, # type: ignore ) @@ -1357,11 +1752,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,11 +1769,11 @@ async def _update_model_table( updated_by=user_api_key_dict.user_id or litellm_proxy_admin_name, ) if model_id is None: - model_dict = await ModelTableRepository(prisma_client).table.create( + model_dict = await _model_table(prisma_client).create( data={**litellm_modeltable.json(exclude_none=True)} # type: ignore ) else: - model_dict = await ModelTableRepository(prisma_client).table.upsert( + model_dict = await _model_table(prisma_client).upsert( where={"id": model_id}, data={ "update": {**litellm_modeltable.json(exclude_none=True)}, # type: ignore @@ -1394,7 +1789,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. @@ -1432,11 +1827,11 @@ async def _auto_add_team_members_to_organization( async def fetch_and_validate_organization( organization_id: str, - existing_team_row: Any, + existing_team_row: LiteLLM_TeamTable, llm_router: Optional[Router], - prisma_client: Any, + prisma_client: PrismaClient, user_api_key_dict: Optional[UserAPIKeyAuth] = None, -) -> Any: +) -> LiteLLM_OrganizationTableWithMembers: """ Fetch and validate an organization for team update operations. @@ -1455,7 +1850,7 @@ 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 = await _organization_table(prisma_client).find_unique( where={"organization_id": organization_id}, include={"litellm_budget_table": True, "members": True, "teams": True}, ) @@ -1467,9 +1862,9 @@ 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(**organization_row.model_dump()) + organization = LiteLLM_OrganizationTableWithMembers(**_org_row_dump(organization_row)) validate_team_org_change( - team=LiteLLM_TeamTable(**existing_team_row.model_dump()), + team=LiteLLM_TeamTable(**_team_row_dump(existing_team_row)), organization=organization, llm_router=llm_router, is_proxy_admin=is_proxy_admin, @@ -1477,7 +1872,7 @@ async def fetch_and_validate_organization( if is_proxy_admin: await _auto_add_team_members_to_organization( - team=LiteLLM_TeamTable(**existing_team_row.model_dump()), + team=LiteLLM_TeamTable(**_team_row_dump(existing_team_row)), organization=organization, prisma_client=prisma_client, ) @@ -1704,7 +2099,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 _team_table(prisma_client).find_unique(where={"team_id": data.team_id}) if existing_team_row is None: raise HTTPException( @@ -1714,7 +2109,7 @@ async def update_team( # Verify caller has access to manage this team await _verify_team_access( - team_obj=LiteLLM_TeamTable(**existing_team_row.model_dump()), + team_obj=LiteLLM_TeamTable(**_team_row_dump(existing_team_row)), user_api_key_dict=user_api_key_dict, ) @@ -1757,7 +2152,7 @@ async def update_team( ): # Is the caller org_admin of the destination org? caller_memberships = ( - await OrganizationMembershipRepository(prisma_client).table.find_many( + await _organization_membership_table(prisma_client).find_many( where={ "user_id": user_api_key_dict.user_id, "organization_id": data.organization_id, @@ -1908,7 +2303,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 _team_table(prisma_client).update( where={"team_id": data.team_id}, data=updated_kv, # `object_permission` is included so `_refresh_cached_team` @@ -2004,7 +2399,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 _team_table(prisma_client).find_unique(where={"team_id": team_id}) if existing_team_row is None: raise HTTPException( status_code=404, @@ -2013,7 +2408,12 @@ async def patch_team( existing_metadata = existing_team_row.metadata if isinstance(existing_team_row.metadata, dict) else {} patch_fields["metadata"] = apply_json_merge_patch(existing_metadata, patch_fields["metadata"]) - update_request = UpdateTeamRequest(team_id=team_id, **patch_fields) + update_request = UpdateTeamRequest( + team_id=team_id, + **cast( + _TeamPatchDump, patch_fields + ), # cast-ok: exclude_unset dump of PatchTeamRequest mirrors UpdateTeamRequest fields + ) result = await update_team( data=update_request, @@ -2486,7 +2886,7 @@ 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 = await _user_table(prisma_client).find_first( where={"user_email": {"equals": member.user_email, "mode": "insensitive"}} ) @@ -2515,7 +2915,7 @@ async def _validate_and_populate_member_user_info( # 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 _user_table(prisma_client).find_unique(where={"user_id": member.user_id}) if user_by_id is None: # User doesn't exist yet - allow it to pass with user_email as None @@ -2591,7 +2991,7 @@ async def team_member_add( detail={"error": f"Team not found for team_id={getattr(data, 'team_id', None)}"}, ) - complete_team_data = LiteLLM_TeamTable(**existing_team_row.model_dump()) + complete_team_data = LiteLLM_TeamTable(**_team_row_dump(existing_team_row)) team_member_add_duplication_check( data=data, @@ -2637,7 +3037,7 @@ async def team_member_add( _emit_team_members_metric(complete_team_data) return TeamAddMemberResponse( - **updated_team.model_dump(), + **_team_row_dump(updated_team), updated_users=updated_users, updated_team_memberships=updated_team_memberships, ) @@ -2704,14 +3104,14 @@ 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 _team_table(prisma_client).find_unique(where={"team_id": data.team_id}) 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(**_existing_team_row.model_dump()) + existing_team_row = LiteLLM_TeamTable(**_team_row_dump(_existing_team_row)) ## CHECK IF USER IS PROXY ADMIN OR TEAM ADMIN OR ORG ADMIN @@ -2742,7 +3142,7 @@ async def team_member_delete( _db_new_team_members: List[dict] = [m.model_dump() for m in new_team_members] - _ = await TeamRepository(prisma_client).table.update( + _ = await _team_table(prisma_client).update( where={ "team_id": data.team_id, }, @@ -2768,7 +3168,7 @@ async def team_member_delete( if data.team_id in existing_user.teams: team_list = existing_user.teams team_list.remove(data.team_id) - await UserRepository(prisma_client).table.update( + await _user_table(prisma_client).update( where={ "user_id": existing_user.user_id, }, @@ -2785,9 +3185,7 @@ async def team_member_delete( user_ids_to_delete.add(existing_user.user_id) for _uid in user_ids_to_delete: - await TeamMembershipRepository(prisma_client).table.delete_many( - where={"team_id": data.team_id, "user_id": _uid} - ) + await _team_membership_table(prisma_client).delete_many(where={"team_id": data.team_id, "user_id": _uid}) ## DELETE KEYS CREATED BY USER FOR THIS TEAM if user_ids_to_delete: @@ -2796,9 +3194,9 @@ async def team_member_delete( ) # Fetch keys before deletion to persist them - keys_to_delete: List[LiteLLM_VerificationToken] = await VerificationTokenRepository( + keys_to_delete: list[LiteLLM_VerificationToken] = await _verification_token_table( prisma_client - ).table.find_many( + ).find_many( # mutable-ok: passed to _persist_deleted_verification_tokens(keys: List[LiteLLM_VerificationToken]) where={ "user_id": {"in": list(user_ids_to_delete)}, "team_id": data.team_id, @@ -2813,7 +3211,7 @@ async def team_member_delete( litellm_changed_by=None, ) - await VerificationTokenRepository(prisma_client).table.delete_many( + await _verification_token_table(prisma_client).delete_many( where={ "user_id": {"in": list(user_ids_to_delete)}, "team_id": data.team_id, @@ -2832,7 +3230,9 @@ async def team_member_delete( } -def _build_member_budget_patch(data: TeamMemberUpdateRequest) -> Dict[str, Any]: +def _build_member_budget_patch( + data: TeamMemberUpdateRequest, +) -> dict[str, object]: # mutable-ok: handed to _upsert_budget_and_membership(budget_patch: Dict[str, Any]) """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.""" @@ -2908,14 +3308,14 @@ 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 _team_table(prisma_client).find_unique(where={"team_id": data.team_id}) 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(**_existing_team_row.model_dump()) + existing_team_row = LiteLLM_TeamTable(**_team_row_dump(_existing_team_row)) ## CHECK IF USER IS PROXY ADMIN OR TEAM ADMIN OR ORG ADMIN @@ -2975,7 +3375,7 @@ async def team_member_update( ### upsert new budget budget_patch = _build_member_budget_patch(data) - async with prisma_client.db.tx() as tx: + async with prisma_client.tx() as tx: await _upsert_budget_and_membership( tx=tx, team_id=data.team_id, @@ -3004,7 +3404,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] - await TeamRepository(prisma_client).table.update( + await _team_table(prisma_client).update( where={"team_id": data.team_id}, data={"members_with_roles": json.dumps(_db_team_members)}, # type: ignore ) @@ -3135,7 +3535,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 _user_table(prisma_client).find_many(order={"created_at": "desc"}) data.members = [ Member( user_id=user.user_id, @@ -3251,9 +3651,7 @@ async def delete_team( 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} - ) + team_row_base: BaseModel | None = await _team_table(prisma_client).find_unique(where={"team_id": team_id}) if team_row_base is None: raise Exception except Exception: @@ -3261,7 +3659,7 @@ async def delete_team( status_code=404, detail={"error": f"Team not found, passed team_id={team_id}"}, ) - team_row_pydantic = LiteLLM_TeamTable(**team_row_base.model_dump()) + team_row_pydantic = LiteLLM_TeamTable(**_team_row_dump(team_row_base)) # Verify caller has access to manage this team await _verify_team_access( @@ -3320,7 +3718,9 @@ async def delete_team( _persist_deleted_verification_tokens, ) - keys_to_delete: List[LiteLLM_VerificationToken] = await VerificationTokenRepository(prisma_client).table.find_many( + keys_to_delete: list[LiteLLM_VerificationToken] = await _verification_token_table( + prisma_client + ).find_many( # mutable-ok: passed to _persist_deleted_verification_tokens(keys: List[LiteLLM_VerificationToken]) where={"team_id": {"in": data.team_ids}} ) @@ -3376,7 +3776,7 @@ def _transform_teams_to_deleted_records( teams: List[LiteLLM_TeamTable], user_api_key_dict: UserAPIKeyAuth, litellm_changed_by: Optional[str] = None, -) -> List[Dict[str, Any]]: +) -> Sequence[Mapping[str, object]]: """Transform teams into deleted team records ready for persistence.""" if not teams: return [] @@ -3384,7 +3784,7 @@ def _transform_teams_to_deleted_records( deleted_at = datetime.now(timezone.utc) records = [] for team in teams: - team_payload = team.model_dump() + team_payload = _team_row_dump(team) deleted_record = LiteLLM_DeletedTeamTable( **team_payload, deleted_at=deleted_at, @@ -3419,13 +3819,13 @@ def _transform_teams_to_deleted_records( async def _save_deleted_team_records( - records: List[Dict[str, Any]], + records: Sequence[Mapping[str, object]], prisma_client: PrismaClient, ) -> None: """Save deleted team records to the database.""" if not records: return - await DeletedTeamRepository(prisma_client).table.create_many(data=records) + await _deleted_team_table(prisma_client).create_many(data=records) async def _persist_deleted_team_records( @@ -3501,9 +3901,7 @@ async def _add_team_member_budget_table( team_info_response_object: TeamInfoResponseObjectTeamTable, ) -> TeamInfoResponseObjectTeamTable: try: - team_budget = await BudgetRepository(prisma_client).table.find_unique( - where={"budget_id": team_member_budget_id} - ) + team_budget = await _budget_table(prisma_client).find_unique(where={"budget_id": team_member_budget_id}) team_info_response_object.team_member_budget_table = team_budget except Exception: verbose_proxy_logger.info( @@ -3513,7 +3911,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: @@ -3567,7 +3965,7 @@ async def team_info( ) try: - team_info: Optional[BaseModel] = await TeamRepository(prisma_client).table.find_unique( + team_info: BaseModel | None = await _team_table(prisma_client).find_unique( where={"team_id": team_id}, include={"litellm_model_table": True, "object_permission": True}, ) @@ -3580,7 +3978,7 @@ async def team_info( ) await validate_membership( user_api_key_dict=user_api_key_dict, - team_table=LiteLLM_TeamTable(**team_info.model_dump()), + team_table=LiteLLM_TeamTable(**_team_row_dump(team_info)), ) ## GET ALL KEYS ## @@ -3617,7 +4015,7 @@ async def team_info( if isinstance(team_info, dict): _team_info = TeamInfoResponseObjectTeamTable(**team_info) elif isinstance(team_info, BaseModel): - _team_info = TeamInfoResponseObjectTeamTable(**team_info.model_dump()) + _team_info = TeamInfoResponseObjectTeamTable(**_team_row_dump(team_info)) else: _team_info = TeamInfoResponseObjectTeamTable() @@ -3814,7 +4212,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 _team_table(prisma_client).find_unique(where={"team_id": data.team_id}) if existing_team is None: raise HTTPException( status_code=404, @@ -3823,11 +4221,11 @@ async def block_team( # Verify caller has access to manage this team await _verify_team_access( - team_obj=LiteLLM_TeamTable(**existing_team.model_dump()), + team_obj=LiteLLM_TeamTable(**_team_row_dump(existing_team)), user_api_key_dict=user_api_key_dict, ) - record = await TeamRepository(prisma_client).table.update( + record = await _team_table(prisma_client).update( where={"team_id": data.team_id}, data={"blocked": True}, # type: ignore ) @@ -3863,7 +4261,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 _team_table(prisma_client).find_unique(where={"team_id": data.team_id}) if existing_team is None: raise HTTPException( status_code=404, @@ -3872,11 +4270,11 @@ async def unblock_team( # Verify caller has access to manage this team await _verify_team_access( - team_obj=LiteLLM_TeamTable(**existing_team.model_dump()), + team_obj=LiteLLM_TeamTable(**_team_row_dump(existing_team)), user_api_key_dict=user_api_key_dict, ) - record = await TeamRepository(prisma_client).table.update( + record = await _team_table(prisma_client).update( where={"team_id": data.team_id}, data={"blocked": False}, # type: ignore ) @@ -3910,28 +4308,28 @@ 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 = await _user_table(prisma_client).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(**user_info.model_dump()) + user_info_correct_type = LiteLLM_UserTable(**_user_row_dump(user_info)) available_teams = [team for team in available_teams if team not in user_info_correct_type.teams] - available_teams_db = await TeamRepository(prisma_client).table.find_many(where={"team_id": {"in": available_teams}}) + available_teams_db = await _team_table(prisma_client).find_many(where={"team_id": {"in": available_teams}}) - available_teams_correct_type = [LiteLLM_TeamTable(**team.model_dump()) for team in available_teams_db] + available_teams_correct_type = [LiteLLM_TeamTable(**_team_row_dump(team)) for team in available_teams_db] return available_teams_correct_type async def _get_org_admin_org_ids( user_id: str, - prisma_client: Any, - user_api_key_cache: Any, - proxy_logging_obj: Any, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, ) -> Optional[List[str]]: """ Return the list of organization IDs where the user is an org admin. @@ -3970,16 +4368,18 @@ async def _build_team_list_where_conditions( use_deleted_table: bool, search: Optional[str] = None, 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]]: + user_api_key_cache: UserApiKeyCache | None = None, + proxy_logging_obj: ProxyLogging | None = None, +) -> Mapping[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 + ] = {} # mutable-ok: filter keys are assigned in place as each query parameter is applied below if team_id: where_conditions["team_id"] = team_id @@ -4007,7 +4407,9 @@ async def _build_team_list_where_conditions( user_object_correct_type = await get_user_object( user_id=user_id, prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, # type: ignore[arg-type] + user_api_key_cache=cast( # cast-ok: list endpoints always pass the proxy cache + UserApiKeyCache, user_api_key_cache + ), user_id_upsert=False, proxy_logging_obj=proxy_logging_obj, ) @@ -4061,7 +4463,7 @@ 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 = await _access_group_table(_prisma_client).find_many( where={"access_group_id": {"in": unique_ids}}, ) @@ -4085,9 +4487,11 @@ def _convert_teams_to_response_models( counts = keys_count_by_team or {} for team in teams: try: - team_dict = team.model_dump() + team_dict = _team_row_dump(team) except Exception: - team_dict = team.dict() + team_dict = cast( # cast-ok: pydantic v1 fallback dump mirrors LiteLLM_TeamTable fields + _TeamRowDump, team.dict() + ) if use_deleted_table: team_list.append(LiteLLM_DeletedTeamTable(**team_dict)) @@ -4109,7 +4513,7 @@ def _convert_teams_to_response_models( async def _get_keys_count_by_team( - prisma_client: Any, + prisma_client: PrismaClient, teams: list, ) -> Dict[str, int]: """Aggregate virtual-key counts per team for the given page of teams. @@ -4122,10 +4526,13 @@ async def _get_keys_count_by_team( if not page_team_ids: return {} - grouped = await VerificationTokenRepository(prisma_client).table.group_by( - by=["team_id"], - where={"team_id": {"in": page_team_ids}}, - count={"team_id": True}, + grouped = cast( # cast-ok: prisma group_by returns plain aggregation dicts + list[_TokenTeamGroupByRow], + 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")} @@ -4134,9 +4541,9 @@ 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, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, ) -> Tuple[Optional[str], Optional[List[str]]]: """Enforce access control for list_team_v2. @@ -4328,23 +4735,23 @@ async def list_team_v2( # Get teams with pagination if use_deleted_table: - teams = await DeletedTeamRepository(prisma_client).table.find_many( + teams = await _deleted_team_table(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 DeletedTeamRepository(prisma_client).table.count(where=where_conditions) + total_count = await _deleted_team_table(prisma_client).count(where=where_conditions) else: - teams = await TeamRepository(prisma_client).table.find_many( + teams = await _team_table(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 _team_table(prisma_client).count(where=where_conditions) # Calculate total pages total_pages = -(-total_count // page_size) # Ceiling division @@ -4387,9 +4794,9 @@ async def list_team_v2( 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, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, ) -> list: """ Authorize the /team/list request and return filtered teams. @@ -4437,7 +4844,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 = await _team_table(prisma_client).find_many( where={"organization_id": {"in": allowed_org_ids}}, include={"litellm_model_table": True}, ) @@ -4447,19 +4854,31 @@ async def _authorize_and_filter_teams( return [ team for team in org_teams - if team.members_with_roles and any(m.get("user_id") == user_id for m in team.members_with_roles) + if team.members_with_roles + and any( + m.get("user_id") == user_id + for m in cast( # cast-ok: members_with_roles is raw JSON dicts on prisma rows + Sequence[Mapping[str, object]], team.members_with_roles + ) + ) ] 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 = await _team_table(prisma_client).find_many(include={"litellm_model_table": True}) return [ team for team in response - if team.members_with_roles and any(m.get("user_id") == user_id for m in team.members_with_roles) + if team.members_with_roles + and any( + m.get("user_id") == user_id + for m in cast( # cast-ok: members_with_roles is raw JSON dicts on prisma rows + Sequence[Mapping[str, object]], team.members_with_roles + ) + ) ] else: # Proxy admin: all teams - return list(await TeamRepository(prisma_client).table.find_many(include={"litellm_model_table": True})) + return list(await _team_table(prisma_client).find_many(include={"litellm_model_table": True})) @router.get("/team/list", tags=["team management"], dependencies=[Depends(user_api_key_auth)]) @@ -4513,7 +4932,7 @@ async def list_team( _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 _verification_token_table(prisma_client).find_many(where={"team_id": team.team_id}) try: returned_responses.append( @@ -4561,10 +4980,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 _team_table(prisma_client).count() # Get paginated teams - teams = await TeamRepository(prisma_client).table.find_many( + teams = await _team_table(prisma_client).find_many( skip=skip, take=page_size, order={"team_alias": "asc"}, # Sort by team_alias @@ -4629,7 +5048,7 @@ async def ui_view_teams( } # Query users with pagination and filters - teams = await TeamRepository(prisma_client).table.find_many( + teams = await _team_table(prisma_client).find_many( where=where_conditions, skip=skip, take=page_size, @@ -4697,7 +5116,7 @@ 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_row = await _team_table(prisma_client).find_unique(where={"team_id": data.team_id}) if team_row is None: raise HTTPException( @@ -4705,7 +5124,7 @@ async def team_model_add( detail={"error": f"Team not found, passed team_id={data.team_id}"}, ) - team_obj = LiteLLM_TeamTable(**team_row.model_dump()) + team_obj = LiteLLM_TeamTable(**_team_row_dump(team_row)) # Authorization check - only proxy admin, team admin, or org admin can add models if ( @@ -4743,10 +5162,13 @@ 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( - where={"team_id": data.team_id}, - data={"updated_at": datetime.now(timezone.utc)}, - include={"object_permission": True}, # type: ignore + updated_team = cast( # cast-ok: update on an existing, just-validated team returns the row + LiteLLM_TeamTable, + await _team_table(prisma_client).update( + where={"team_id": data.team_id}, + data={"updated_at": datetime.now(timezone.utc)}, + include={"object_permission": True}, # type: ignore + ), ) await _refresh_cached_team( @@ -4797,7 +5219,7 @@ 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_row = await _team_table(prisma_client).find_unique(where={"team_id": data.team_id}) if team_row is None: raise HTTPException( @@ -4805,7 +5227,7 @@ async def team_model_delete( detail={"error": f"Team not found, passed team_id={data.team_id}"}, ) - team_obj = LiteLLM_TeamTable(**team_row.model_dump()) + team_obj = LiteLLM_TeamTable(**_team_row_dump(team_row)) # Authorization check - only proxy admin, team admin, or org admin can remove models if ( @@ -4825,10 +5247,13 @@ 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( - where={"team_id": data.team_id}, - data={"models": updated_models}, - include={"object_permission": True}, # type: ignore + updated_team = cast( # cast-ok: update on an existing, just-validated team returns the row + LiteLLM_TeamTable, + await _team_table(prisma_client).update( + where={"team_id": data.team_id}, + data={"models": updated_models}, + include={"object_permission": True}, # type: ignore + ), ) await _refresh_cached_team( @@ -4873,7 +5298,7 @@ async def team_member_permissions( check_db_only=True, ) - complete_team_data = LiteLLM_TeamTable(**existing_team_row.model_dump()) + complete_team_data = LiteLLM_TeamTable(**_team_row_dump(existing_team_row)) # Admin Viewer follows the read-parity rule: see team permissions like # a Proxy Admin would. Team / org admins keep their existing scope. @@ -4940,7 +5365,7 @@ async def update_team_member_permissions( check_db_only=True, ) - complete_team_data = LiteLLM_TeamTable(**existing_team_row.model_dump()) + complete_team_data = LiteLLM_TeamTable(**_team_row_dump(existing_team_row)) # Available-team self-join must NOT grant write access to team-wide # permission policies; only proxy/team/org admins can update them. @@ -4960,9 +5385,12 @@ async def update_team_member_permissions( }, ) # Update the team member permissions - updated_team = await TeamRepository(prisma_client).table.update( - where={"team_id": data.team_id}, - data={"team_member_permissions": data.team_member_permissions}, + updated_team = cast( # cast-ok: update on an existing, just-validated team returns the row + LiteLLM_TeamTable, + await _team_table(prisma_client).update( + where={"team_id": data.team_id}, + data={"team_member_permissions": data.team_member_permissions}, + ), ) return updated_team @@ -5030,7 +5458,7 @@ async def bulk_update_team_member_permissions( } -async def _compute_and_batch_updates(prisma_client, teams, permissions_to_add: set) -> int: +async def _compute_and_batch_updates(prisma_client, teams, permissions_to_add: AbstractSet[str]) -> int: """Compute merged permissions and batch-write updates. Returns count of teams updated.""" updates = [] for team in teams: @@ -5052,9 +5480,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, team_ids: Sequence[str], permissions_to_add: AbstractSet[str] +) -> int: """Fetch specific teams by ID and append permissions.""" - teams = await TeamRepository(prisma_client).table.find_many( + teams = await _team_table(prisma_client).find_many( where={"team_id": {"in": team_ids}}, ) @@ -5069,14 +5499,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: AbstractSet[str]) -> int: """Paginated read + batched write across all teams.""" teams_updated = 0 cursor = None BATCH_SIZE = 500 while True: - find_args: dict = { + find_args: _TeamPageArgs = { "take": BATCH_SIZE, "order": {"team_id": "asc"}, } @@ -5084,7 +5514,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 = await _team_table(prisma_client).find_many(**find_args) if not teams: break @@ -5184,8 +5614,12 @@ async def get_team_daily_activity( where_condition = {} 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 _team_table(prisma_client).find_many(where=where_condition) + team_alias_metadata: dict[ + str, dict[str, object] + ] = { # mutable-ok: get_daily_activity takes entity_metadata_field: 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. @@ -5201,7 +5635,7 @@ async def get_team_daily_activity( 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(**team_alias.model_dump()) + team_obj = LiteLLM_TeamTable(**_team_row_dump(team_alias)) 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, @@ -5215,7 +5649,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 _verification_token_table(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] diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index a20b557e38b..34edaa359ee 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -15,7 +15,7 @@ import time import traceback import warnings -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from datetime import datetime, timedelta, timezone from typing import ( TYPE_CHECKING, @@ -635,7 +635,7 @@ def generate_feedback_box(): 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 @@ -820,13 +820,23 @@ async def proxy_shutdown_event(): async def _initialize_shared_aiohttp_session(): """Initialize shared aiohttp session for connection reuse with connection limits.""" try: + import socket + from aiohttp import ClientSession, TCPConnector from litellm.llms.custom_httpx.http_handler import ( _build_aiohttp_keepalive_socket_factory, ) - connector_kwargs: Dict[str, Any] = { + class ConnectorKwargs(TypedDict, total=False): + keepalive_timeout: float + ttl_dns_cache: int + enable_cleanup_closed: bool + limit: int + limit_per_host: int + socket_factory: Callable[[tuple[object, ...]], socket.socket] + + connector_kwargs: ConnectorKwargs = { "keepalive_timeout": AIOHTTP_KEEPALIVE_TIMEOUT, "ttl_dns_cache": AIOHTTP_TTL_DNS_CACHE, } @@ -958,7 +968,9 @@ async def proxy_startup_event(app: FastAPI): async def _run_pw_migration(): try: - result = await migrate_passwords_to_scrypt_async(prisma_client) + result = await migrate_passwords_to_scrypt_async( + cast(PrismaClient, prisma_client) # cast-ok: guarded by the enclosing prisma_client None check + ) verbose_proxy_logger.info(f"Password migration: {result}") except Exception as e: verbose_proxy_logger.warning(f"Password migration skipped: {e}") @@ -1141,7 +1153,7 @@ async def _run_pw_migration(): 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: @@ -2782,7 +2794,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[str, object]] = [] ### UPDATE KEY SPEND ### async def _update_key_cache(token: str, response_cost: float): @@ -3513,7 +3525,7 @@ def _scrub_guardrail_inner(inner: Dict[str, Any]) -> None: inner["guardrail"] = None -def _scrub_db_overlay_remote_module_loads(section: str, db_value: Any) -> Any: +def _scrub_db_overlay_remote_module_loads(section: str, db_value: object) -> object: """Strip ``s3://`` / ``gcs://`` entries from the DB-overlay value for fields whose contents reach ``get_instance_fn``. The same scheme is allowed from a YAML config (the documented operator flow) but a @@ -3915,7 +3927,9 @@ async def save_config(self, new_config: dict, include_env_vars: bool = False): # if using - db for config - models are in ModelTable # Make a copy to avoid mutating the original config - config_to_save = new_config.copy() + config_to_save: dict[str, Any] = ( + new_config.copy() + ) # mutable-ok: keys are popped and re-encrypted in place before the DB write # environment_variables are persisted to the DB only when a caller # explicitly opts in. Most callers reach save_config after @@ -5390,8 +5404,12 @@ def _add_deployment(self, db_models: list) -> int: added_models += 1 return added_models - def decrypt_model_list_from_db(self, new_models: list) -> list: - _model_list: list = [] + def decrypt_model_list_from_db( + self, new_models: Sequence[Any] + ) -> list[ + dict[str, Any] + ]: # mutable-ok: result is handed to litellm.Router(model_list=...), which requires a list of dicts + _model_list: list[dict[str, Any]] = [] # mutable-ok: appended to once per decrypted deployment below for m in new_models: _litellm_params = m.litellm_params if isinstance(_litellm_params, BaseModel): @@ -5400,7 +5418,9 @@ def decrypt_model_list_from_db(self, new_models: list) -> list: # decrypt values for k, v in _litellm_params.items(): _litellm_params[k] = self._resolve_db_litellm_param(key=k, value=v) - _litellm_params = LiteLLM_Params(**_litellm_params) + _litellm_params = cast( # cast-ok: db params dict has untyped values, validated by pydantic at runtime + Callable[..., LiteLLM_Params], LiteLLM_Params + )(**_litellm_params) else: verbose_proxy_logger.error( f"Invalid model added to proxy db. Invalid litellm params. litellm_params={_litellm_params}" @@ -5448,7 +5468,7 @@ async def _update_llm_router( ) return - models_list: list = new_models if isinstance(new_models, list) else [] + models_list: list[Any] = new_models if isinstance(new_models, list) else [] if llm_router is None and master_key is not None: verbose_proxy_logger.debug(f"len new_models: {len(models_list)}") @@ -5620,7 +5640,7 @@ def _encrypt_env_variables_for_db( ) @staticmethod - def _parse_router_settings_value(value: Any) -> Optional[dict]: + def _parse_router_settings_value(value: object) -> dict | None: """ Parse a router_settings value that may be a dict or a JSON/YAML string. @@ -6251,7 +6271,9 @@ async def _init_semantic_filter_settings_in_db(self, prisma_client: PrismaClient if isinstance(litellm_settings, str): litellm_settings = json.loads(litellm_settings) - mcp_semantic_filter_config = litellm_settings.get("mcp_semantic_tool_filter", None) + mcp_semantic_filter_config = cast( # cast-ok: litellm_settings config rows store a JSON object + dict[str, Any], litellm_settings + ).get("mcp_semantic_tool_filter", None) if mcp_semantic_filter_config is None: return @@ -6386,7 +6408,9 @@ async def _check_and_reload_model_cost_map(self, prisma_client: PrismaClient): if config_record is None or config_record.param_value is None: return # No configuration found, skip reload - config = config_record.param_value + config = cast( # cast-ok: reload config rows store a JSON object + dict[str, Any], config_record.param_value + ) interval_hours = config.get("interval_hours") force_reload = config.get("force_reload", False) @@ -6486,7 +6510,9 @@ async def _check_and_reload_anthropic_beta_headers(self, prisma_client: PrismaCl if config_record is None or config_record.param_value is None: return # No configuration found, skip reload - config = config_record.param_value + config = cast( # cast-ok: reload config rows store a JSON object + dict[str, Any], config_record.param_value + ) interval_hours = config.get("interval_hours") force_reload = config.get("force_reload", False) @@ -7719,7 +7745,11 @@ async def _initialize_semantic_tool_filter( """Initialize MCP semantic tool filter if configured""" from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook - mcp_semantic_filter_config = litellm_settings.get("mcp_semantic_tool_filter", None) + mcp_semantic_filter_config: dict[str, Any] | None = ( + litellm_settings.get( # mutable-ok: forwarded to SemanticToolFilterHook.initialize_from_config, which takes Dict + "mcp_semantic_tool_filter", None + ) + ) # Only proceed if the feature is configured and enabled if not mcp_semantic_filter_config or not mcp_semantic_filter_config.get("enabled", False): @@ -8857,7 +8887,7 @@ async def model_info( ) -def _blocked_response_usage(original_response: Optional[Any]) -> "litellm.Usage": +def _blocked_response_usage(original_response: object | None) -> "litellm.Usage": """ Token usage for a synthetic guardrail-blocked response. @@ -8971,7 +9001,9 @@ async def chat_completion( # Guardrail flagged content in passthrough mode - return 200 with violation message _data = e.request_data # Capture logging_obj before post_call_failure_hook pops it from _data. - _logging_obj = _data.get("litellm_logging_obj") + _logging_obj = cast( # cast-ok: chat request data always carries the litellm logging object + LiteLLMLoggingObj, _data.get("litellm_logging_obj") + ) await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, @@ -9022,7 +9054,9 @@ async def chat_completion( completion_stream=_iterator, model=data.get("model", ""), custom_llm_provider="cached_response", - logging_obj=_data.get("litellm_logging_obj", None), + logging_obj=cast( # cast-ok: chat request data always carries the litellm logging object + LiteLLMLoggingObj, _data.get("litellm_logging_obj", None) + ), ) selected_data_generator = select_data_generator( response=_streaming_response, @@ -11178,13 +11212,20 @@ async def get_all_team_models( team_db_objects_typed: List[LiteLLM_TeamTable] = [] + team_table_from_row = cast( # cast-ok: prisma row dump has untyped values, validated by pydantic at runtime + Callable[..., LiteLLM_TeamTable], LiteLLM_TeamTable + ) if user_teams == "*": team_db_objects = await TeamRepository(prisma_client).table.find_many() - team_db_objects_typed = [LiteLLM_TeamTable(**team_db_object.model_dump()) for team_db_object in team_db_objects] + team_db_objects_typed = [ + team_table_from_row(**team_db_object.model_dump()) for team_db_object in team_db_objects + ] else: team_db_objects = await TeamRepository(prisma_client).table.find_many(where={"team_id": {"in": user_teams}}) - team_db_objects_typed = [LiteLLM_TeamTable(**team_db_object.model_dump()) for team_db_object in team_db_objects] + team_db_objects_typed = [ + team_table_from_row(**team_db_object.model_dump()) for team_db_object in team_db_objects + ] team_models = _add_team_models_to_all_models( team_db_objects_typed=team_db_objects_typed, @@ -11258,7 +11299,9 @@ async def _populate_team_access_on_models( where={"user_id": user_api_key_dict.user_id} ) if user_db_object is not None: - user_object = LiteLLM_UserTable(**user_db_object.model_dump()) + user_object = cast( # cast-ok: prisma row dump has untyped values, validated by pydantic at runtime + Callable[..., LiteLLM_UserTable], LiteLLM_UserTable + )(**user_db_object.model_dump()) user_teams = user_object.teams or [] direct_access_models = get_direct_access_models( user_db_object=user_object, @@ -11327,7 +11370,9 @@ def _enrich_model_info_with_litellm_data( Enriched model dictionary with sensitive info removed """ # provided model_info in config.yaml - model_info = model.get("model_info", {}) + model_info: dict[str, Any] = model.get( + "model_info", {} + ) # mutable-ok: litellm model-cost keys are written into it before it is stored back on the model if debug is True: _openai_client = "None" if llm_router is not None: @@ -11344,7 +11389,7 @@ def _enrich_model_info_with_litellm_data( # 2nd pass on the model, try seeing if we can find model in litellm model_cost map if litellm_model_info == {}: # use litellm_param model_name to get model_info - litellm_params = model.get("litellm_params", {}) + litellm_params: dict[str, Any] = model.get("litellm_params", {}) litellm_model = litellm_params.get("model", None) try: litellm_model_info = litellm.get_model_info(model=litellm_model) @@ -11375,7 +11420,7 @@ def _enrich_model_info_with_litellm_data( async def _get_caller_byok_team_scope( user_api_key_dict: Optional[UserAPIKeyAuth], - prisma_client: Optional[Any], + prisma_client: PrismaClient | None, ) -> Optional[Set[str]]: """ Return the team IDs whose BYOK rows the caller is allowed to see via @@ -11432,8 +11477,8 @@ def _byok_row_outside_caller_teams(model_info_dict: Dict[str, Any], allowed_team async def _fetch_db_models_for_search( - prisma_client: Any, - proxy_config: Any, + prisma_client: PrismaClient, + proxy_config: ProxyConfig, search_lower: str, db_model_ids_in_router: Set[str], router_models_count: int, @@ -11498,8 +11543,8 @@ async def _fetch_db_models_for_search( async def _apply_search_filter_to_models( all_models: List[Dict[str, Any]], search: str, - prisma_client: Optional[Any], - proxy_config: Any, + prisma_client: PrismaClient | None, + proxy_config: ProxyConfig, user_api_key_dict: Optional[UserAPIKeyAuth] = None, page: int = 1, size: int = 50, @@ -11566,7 +11611,7 @@ def _model_matches_search(m: Dict[str, Any]) -> bool: db_model_ids_in_router = set() for m in filtered_router_models: - model_info = m.get("model_info", {}) + model_info: Mapping[str, Any] = m.get("model_info", {}) is_db_model = model_info.get("db_model", False) model_id = model_info.get("id") @@ -11604,7 +11649,7 @@ def _model_matches_search(m: Dict[str, Any]) -> bool: return filtered_router_models + db_models, search_total_count -def _normalize_datetime_for_sorting(dt: Any) -> Optional[datetime]: +def _normalize_datetime_for_sorting(dt: object) -> datetime | None: """ Normalize a datetime value to a timezone-aware UTC datetime for sorting. @@ -11674,7 +11719,7 @@ def _sort_models( reverse = sort_order.lower() == "desc" def get_sort_key(model: Dict[str, Any]) -> Any: - model_info = model.get("model_info", {}) + model_info: Mapping[str, Any] = model.get("model_info", {}) if sort_by == "model_name": # Team BYOK models persist an internal `model_name` (e.g. @@ -11775,7 +11820,7 @@ 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: Mapping[str, Sequence[str]]) -> list[str]: """Expand team model entries (including access group names) to concrete model names.""" resolved: List[str] = [] for name in team_models: @@ -11793,7 +11838,9 @@ async def _load_team_object_for_model_filter(team_id: str, prisma_client: Prisma if team_db_object is None: verbose_proxy_logger.warning(f"Team {team_id} not found in database") return None - return LiteLLM_TeamTable(**team_db_object.model_dump()) + return cast( # cast-ok: prisma row dump has untyped values, validated by pydantic at runtime + Callable[..., LiteLLM_TeamTable], LiteLLM_TeamTable + )(**team_db_object.model_dump()) except Exception as e: verbose_proxy_logger.exception(f"Error fetching team {team_id}: {str(e)}") return None @@ -11935,7 +11982,7 @@ async def _filter_models_by_team_id( # every public model the admin can call. filtered_models = [] for _model in all_models: - model_info = _model.get("model_info", {}) + model_info: dict[str, Any] = _model.get("model_info", {}) model_id = model_info.get("id", None) # BYOK rows owned by this team are always accessible to it, even if @@ -12286,7 +12333,9 @@ async def model_streaming_metrics( """ _all_api_bases = set() - db_response = await prisma_client.db.query_raw(sql_query, _selected_model_group, startTime, endTime) + db_response: Sequence[Mapping[str, Any]] | None = await prisma_client.db.query_raw( + sql_query, _selected_model_group, startTime, endTime + ) _daily_entries: dict = {} # {"Jun 23": {"model1": 0.002, "model2": 0.003}} if db_response is not None: for model_data in db_response: @@ -12408,7 +12457,7 @@ async def model_metrics( avg_latency_per_token DESC; """ _all_api_bases = set() - db_response = await prisma_client.db.query_raw( + db_response: Sequence[Mapping[str, Any]] | None = await prisma_client.db.query_raw( sql_query, _selected_model_group, startTime, endTime, api_key, customer ) _daily_entries: dict = {} # {"Jun 23": {"model1": 0.002, "model2": 0.003}} @@ -12526,7 +12575,9 @@ async def model_metrics_slow_responses( slow_count DESC; """ - db_response = await prisma_client.db.query_raw( + db_response: Optional[ + list[dict[str, Any]] + ] = await prisma_client.db.query_raw( # mutable-ok: each row's api_base is rewritten in place below sql_query, alerting_threshold, _selected_model_group, @@ -12971,7 +13022,11 @@ def _get_model_group_info( _model_group_info = llm_router.get_model_group_info(model_group=model) if _model_group_info is not None: - model_groups.append(ModelGroupInfoProxy(**_model_group_info.model_dump())) + model_groups.append( + cast( # cast-ok: model group info dump has untyped values, validated by pydantic at runtime + Callable[..., ModelGroupInfoProxy], ModelGroupInfoProxy + )(**_model_group_info.model_dump()) + ) else: model_group_info = ModelGroupInfoProxy( model_group=model, @@ -13553,7 +13608,7 @@ async def login_v2(request: Request): from litellm.proxy.utils import get_custom_url try: - body = await request.json() + body: Mapping[str, Any] = await request.json() username = str(body.get("username")) password = str(body.get("password")) @@ -13626,7 +13681,7 @@ async def login_v3(request: Request): code=status.HTTP_404_NOT_FOUND, ) - body = await request.json() + body: Mapping[str, Any] = await request.json() username = str(body.get("username")) password = str(body.get("password")) @@ -13697,7 +13752,7 @@ async def login_v3_exchange(request: Request): code=status.HTTP_404_NOT_FOUND, ) - body = await request.json() + body: Mapping[str, Any] = await request.json() code = body.get("code") if not code: raise ProxyException( @@ -14690,7 +14745,9 @@ async def update_config_general_settings( ) try: - ConfigGeneralSettings(**{data.field_name: data.field_value}) + cast( # cast-ok: field name/value are dynamic user input, validated by pydantic at runtime + Callable[..., ConfigGeneralSettings], ConfigGeneralSettings + )(**{data.field_name: data.field_value}) except Exception: raise HTTPException( status_code=400, @@ -15003,7 +15060,7 @@ def _general_settings_ui_litellm_default( return False if spec["type"] == "Boolean" else None -def _validate_general_settings_ui_litellm_value(field_name: str, value: Any) -> GeneralSettingsUILiteLLMValue: +def _validate_general_settings_ui_litellm_value(field_name: str, value: object) -> GeneralSettingsUILiteLLMValue: spec = _GENERAL_SETTINGS_UI_LITELLM_FIELDS[field_name] field_type = spec["type"] if value is None or value == "": @@ -15043,7 +15100,7 @@ def _validate_general_settings_ui_litellm_value(field_name: str, value: Any) -> async def _persist_general_settings_ui_litellm_field( - field_name: str, value: Any, user_api_key_dict: UserAPIKeyAuth + field_name: str, value: object, user_api_key_dict: UserAPIKeyAuth ) -> dict: validated = _validate_general_settings_ui_litellm_value(field_name, value) config = await proxy_config.get_config() diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 171e17ce650..97f9d92a91c 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -26,8 +26,11 @@ Literal, Mapping, Optional, + Protocol, Sequence, Tuple, + TypedDict, + TypeVar, Union, cast, overload, @@ -173,6 +176,21 @@ if TYPE_CHECKING: from opentelemetry.trace import Span as _Span from prisma.client import TransactionManager + from prisma.types import DatasourceOverride + from prisma.models import ( + LiteLLM_BudgetTable, + LiteLLM_Config, + LiteLLM_DeprecatedVerificationToken, + LiteLLM_EndUserTable, + LiteLLM_HealthCheckTable, + LiteLLM_SpendLogs, + LiteLLM_TeamTable, + LiteLLM_UserNotifications, + LiteLLM_UserTable, + LiteLLM_VerificationToken, + ) + + from pydantic import BaseModel from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -180,6 +198,105 @@ else: Span = Any +_T = TypeVar("_T") +_PrismaModelT_co = TypeVar("_PrismaModelT_co", covariant=True) + + +class _PrismaTable(Protocol[_PrismaModelT_co]): + """Typed view of the Prisma table delegates the repositories expose as ``Any``.""" + + def find_unique( + self, + where: Mapping[str, object], + include: Mapping[str, object] | None = None, + ) -> Awaitable[_PrismaModelT_co | None]: ... + + def find_first( + self, + where: Mapping[str, object] | None = None, + order: Mapping[str, str] | Sequence[Mapping[str, str]] | None = None, + include: Mapping[str, object] | None = None, + ) -> Awaitable[_PrismaModelT_co | None]: ... + + def find_many( + self, + take: int | None = None, + skip: int | None = None, + where: Mapping[str, object] | None = None, + include: Mapping[str, object] | None = None, + order: Mapping[str, str] | Sequence[Mapping[str, str]] | None = None, + distinct: Sequence[str] | None = None, + ) -> Awaitable[Sequence[_PrismaModelT_co]]: ... + + def create( + self, + data: Mapping[str, object], + include: Mapping[str, object] | None = None, + ) -> Awaitable[_PrismaModelT_co]: ... + + def create_many( + self, + data: Sequence[Mapping[str, object]], + skip_duplicates: bool | None = None, + ) -> Awaitable[int]: ... + + def update( + self, + data: Mapping[str, object], + where: Mapping[str, object], + include: Mapping[str, object] | None = None, + ) -> Awaitable[_PrismaModelT_co | None]: ... + + def upsert( + self, + where: Mapping[str, object], + data: Mapping[str, object], + include: Mapping[str, object] | None = None, + ) -> Awaitable[_PrismaModelT_co]: ... + + def delete_many( + self, + where: Mapping[str, object] | None = None, + ) -> Awaitable[int]: ... + + +class _PrismaBatchTable(Protocol): + """Typed view of a table on a Prisma batch writer (sync queueing methods).""" + + def update(self, data: Mapping[str, object], where: Mapping[str, object]) -> None: ... + + def upsert(self, where: Mapping[str, object], data: Mapping[str, object]) -> None: ... + + +class _PrismaBatcher(Protocol): + """Typed view of the untyped ``db.batch_()`` handle.""" + + litellm_usertable: _PrismaBatchTable + litellm_endusertable: _PrismaBatchTable + litellm_budgettable: _PrismaBatchTable + litellm_teamtable: _PrismaBatchTable + litellm_verificationtoken: _PrismaBatchTable + + def commit(self) -> Awaitable[None]: ... + + +class _PrismaRawClient(Protocol): + """Typed view of the raw-SQL surface the prisma wrappers delegate via ``__getattr__``.""" + + def query_raw(self, query: str, *args: object) -> Awaitable[Sequence[Mapping[str, Any]]]: ... + + def query_first(self, query: str, *args: object) -> Awaitable[dict[str, Any] | None]: ... + + def execute_raw(self, query: str, *args: object) -> Awaitable[int]: ... + + def batch_(self) -> _PrismaBatcher: ... + + +class _LegacyModelDump(Protocol): + """Pydantic v1 dump surface, reached only when ``model_dump`` is unavailable.""" + + def dict(self) -> dict[str, Any]: ... + unified_guardrail = UnifiedLLMGuardrails() @@ -343,7 +460,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. @@ -392,11 +509,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: @@ -670,7 +787,7 @@ def _convert_mcp_to_llm_format(self, request_obj, kwargs: dict) -> dict: 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. """ @@ -796,7 +913,7 @@ def _parse_arguments_manually(self, args_text: str, original_args: dict) -> Opti 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. """ @@ -1136,8 +1253,8 @@ async def _process_prompt_template( self, data: dict, litellm_logging_obj: Any, - prompt_id: Any, - prompt_version: Any, + prompt_id: str, + prompt_version: int | None, call_type: CallTypesLiteral, ) -> None: """Process prompt template if applicable.""" @@ -1358,8 +1475,10 @@ async def pre_call_hook( return None litellm_logging_obj = cast(Optional["LiteLLMLoggingObj"], data.get("litellm_logging_obj", None)) - prompt_id = data.get("prompt_id", None) - prompt_version = data.get("prompt_version", None) + prompt_id = cast(str | None, data.get("prompt_id", None)) # cast-ok: request payload is an untyped dict + prompt_version = cast( # cast-ok: request payload is an untyped dict + int | None, data.get("prompt_version", None) + ) ## PROMPT TEMPLATE CHECK ## @@ -1432,8 +1551,7 @@ async def pre_call_hook( data = result elif ( - _callback is not None - and isinstance(_callback, CustomLogger) + isinstance(_callback, CustomLogger) and "async_pre_call_hook" in vars(_callback.__class__) and _callback.__class__.async_pre_call_hook != CustomLogger.async_pre_call_hook ): @@ -1608,7 +1726,7 @@ def _emit_guardrail_metrics( break @staticmethod - async def _run_guardrail_with_metrics(callback: Any, coro: Awaitable[Any], hook_type: str) -> Any: + async def _run_guardrail_with_metrics(callback: CustomLogger, coro: Awaitable[_T], hook_type: str) -> _T: """ Await `coro`, recording its latency and status to the `litellm_guardrail_latency_seconds` metric under `hook_type`, and @@ -1640,8 +1758,8 @@ async def _run_guardrail_with_metrics(callback: Any, coro: Awaitable[Any], hook_ @staticmethod async def _wrap_streaming_iterator_with_enrichment( - callback: Any, gen: AsyncGenerator[Any, None] - ) -> AsyncGenerator[Any, None]: + callback: CustomLogger, gen: AsyncGenerator[_T, None] + ) -> AsyncGenerator[_T, None]: """ Yield from `gen`; if iteration raises an HTTPException with dict detail, enrich the detail with the originating callback's `guardrail_name` and @@ -1685,12 +1803,12 @@ def _callback_capabilities() -> "_CallbackCapabilities": 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: 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: @@ -2468,7 +2586,7 @@ async def post_call_response_headers_hook( self, data: dict, user_api_key_dict: UserAPIKeyAuth, - response: Any, + response: object, request_headers: Optional[Dict[str, str]] = None, ) -> Dict[str, str]: """ @@ -2524,7 +2642,7 @@ async def post_call_response_headers_hook( return merged_headers @staticmethod - def _build_litellm_call_info(data: dict, response: Any) -> Dict[str, Any]: + def _build_litellm_call_info(data: dict, response: object) -> dict[str, Any]: """ Build a normalized dict of routing metadata from response._hidden_params and data, abstracting away the metadata vs litellm_metadata split. @@ -2636,7 +2754,7 @@ async def async_post_call_streaming_iterator_hook( response, user_api_key_dict: UserAPIKeyAuth, request_data: dict, - ): + ) -> AsyncGenerator[Any, None]: """ Allow user to modify outgoing streaming data -> Given a whole response iterator. This hook is best used when you need to modify multiple chunks of the response at once. @@ -2810,7 +2928,11 @@ async def _lookup_deprecated_key( _deprecated_key_cache.pop(hashed_token, None) try: - deprecated_row = await db.litellm_deprecatedverificationtoken.find_first( + deprecated_table = cast( # cast-ok: db attributes are delegated via untyped __getattr__; typed table protocol + "_PrismaTable[LiteLLM_DeprecatedVerificationToken]", + db.litellm_deprecatedverificationtoken, + ) + deprecated_row = await deprecated_table.find_first( where={ "token": hashed_token, "revoke_at": {"gt": now}, @@ -2848,20 +2970,25 @@ 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: object) -> None: self.param_name = param_name self.param_value = param_value +class _PackedConfigRow(TypedDict): + 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: Union["LiteLLM_Config", _ConfigRow]) -> _PackedConfigRow: return {"param_name": row.param_name, "param_value": row.param_value} -def _unpack_config_row(cached: Any) -> Optional[_ConfigRow]: +def _unpack_config_row(cached: _PackedConfigRow | str | None) -> _ConfigRow | None: if cached is None or cached == _CONFIG_CACHE_MISS: return None if isinstance(cached, dict): @@ -2869,15 +2996,25 @@ 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 +) -> Union["LiteLLM_Config", _ConfigRow] | 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) + cached = ( + cast( # cast-ok: DualCache values are untyped; this cache only ever stores packed rows or the miss sentinel + _PackedConfigRow | str | None, + await litellm_config_cache.async_get_cache(cache_key), + ) + ) if cached is not None: 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 + row = cast( # cast-ok: get_generic_data is untyped; the "config" table returns LiteLLM_Config rows + Optional["LiteLLM_Config"], + await prisma_client.get_generic_data(key="param_name", value=param_name, table_name="config"), + ) + cache_value: _PackedConfigRow | str = _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 @@ -2887,14 +3024,16 @@ 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 try: - rows = await ConfigRepository(prisma_client).table.find_many( - where={"param_name": {"in": param_names}} # type: ignore + config_table = cast( # cast-ok: repository .table is declared Any; typed table protocol pins the real shape + "_PrismaTable[LiteLLM_Config]", + ConfigRepository(prisma_client).table, ) + rows = await config_table.find_many(where={"param_name": {"in": param_names}}) except Exception as e: verbose_proxy_logger.debug( "prefetch_config_params failed, falling through to per-param queries: %s", @@ -2904,7 +3043,7 @@ 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: _PackedConfigRow | str = _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 ) @@ -2983,11 +3122,13 @@ def __init__( ) 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}} + reader_datasource = cast( # cast-ok: prisma's DatasourceOverride is not inferred from a dict literal + "DatasourceOverride", {"url": read_replica_url} + ) if http_client is not None: - reader_prisma = Prisma(http=http_client, **reader_kwargs) + reader_prisma = Prisma(http=http_client, datasource=reader_datasource) else: - reader_prisma = Prisma(**reader_kwargs) + reader_prisma = Prisma(datasource=reader_datasource) reader_wrapper = PrismaWrapper( original_prisma=reader_prisma, iam_token_db_auth=iam_flag, @@ -3066,6 +3207,71 @@ def tx(self) -> "TransactionManager": """ return cast("TransactionManager", self.db.tx()) # cast-ok: wrappers delegate tx via __getattr__ (untyped) + @property + def raw_db(self) -> "_PrismaRawClient": + """Raw-SQL surface of the underlying Prisma client. + + Callers go through this instead of reaching into ``self.db`` so the + wrapper's untyped ``__getattr__`` delegation stays behind one seam. + """ + return cast("_PrismaRawClient", self.db) # cast-ok: wrappers delegate raw SQL via __getattr__ (untyped) + + @property + def user_table(self) -> "_PrismaTable[LiteLLM_UserTable]": + return cast( # cast-ok: repository .table is declared Any; typed table protocol pins the real shape + "_PrismaTable[LiteLLM_UserTable]", + UserRepository(self).table, + ) + + @property + def token_table(self) -> "_PrismaTable[LiteLLM_VerificationToken]": + return cast( # cast-ok: repository .table is declared Any; typed table protocol pins the real shape + "_PrismaTable[LiteLLM_VerificationToken]", + VerificationTokenRepository(self).table, + ) + + @property + def config_table(self) -> "_PrismaTable[LiteLLM_Config]": + return cast( # cast-ok: repository .table is declared Any; typed table protocol pins the real shape + "_PrismaTable[LiteLLM_Config]", + ConfigRepository(self).table, + ) + + @property + def spend_log_table(self) -> "_PrismaTable[LiteLLM_SpendLogs]": + return cast( # cast-ok: repository .table is declared Any; typed table protocol pins the real shape + "_PrismaTable[LiteLLM_SpendLogs]", + SpendLogsRepository(self).table, + ) + + @property + def budget_table(self) -> "_PrismaTable[LiteLLM_BudgetTable]": + return cast( # cast-ok: repository .table is declared Any; typed table protocol pins the real shape + "_PrismaTable[LiteLLM_BudgetTable]", + BudgetRepository(self).table, + ) + + @property + def end_user_table(self) -> "_PrismaTable[LiteLLM_EndUserTable]": + return cast( # cast-ok: repository .table is declared Any; typed table protocol pins the real shape + "_PrismaTable[LiteLLM_EndUserTable]", + EndUserRepository(self).table, + ) + + @property + def team_table(self) -> "_PrismaTable[LiteLLM_TeamTable]": + return cast( # cast-ok: repository .table is declared Any; typed table protocol pins the real shape + "_PrismaTable[LiteLLM_TeamTable]", + TeamRepository(self).table, + ) + + @property + def user_notification_table(self) -> "_PrismaTable[LiteLLM_UserNotifications]": + return cast( # cast-ok: repository .table is declared Any; typed table protocol pins the real shape + "_PrismaTable[LiteLLM_UserNotifications]", + UserNotificationsRepository(self).table, + ) + def get_request_status(self, payload: Union[dict, SpendLogsPayload]) -> Literal["success", "failure"]: """ Determine if a request was successful or failed based on payload metadata. @@ -3097,7 +3303,7 @@ def hash_token(self, token: str): return hashed_token - def jsonify_object(self, data: dict) -> dict: + def jsonify_object(self, data: dict[str, Any]) -> dict[str, Any]: db_data = copy.deepcopy(data) for k, v in db_data.items(): @@ -3144,7 +3350,7 @@ async def check_view_exists(self): required_view = "LiteLLM_VerificationTokenView" expected_views_str = ", ".join(f"'{view}'" for view in expected_views) pg_schema = os.getenv("DATABASE_SCHEMA", "public") - ret = await self.db.query_raw(f""" + ret = await self.raw_db.query_raw(f""" WITH existing_views AS ( SELECT viewname FROM pg_views @@ -3165,7 +3371,7 @@ async def check_view_exists(self): ## check if required view exists ## if ret[0]["view_names"] and required_view not in ret[0]["view_names"]: await self.health_check() # make sure we can connect to db - await self.db.execute_raw(""" + await self.raw_db.execute_raw(""" CREATE VIEW "LiteLLM_VerificationTokenView" AS SELECT v.*, @@ -3211,7 +3417,7 @@ async def check_view_exists(self): async def get_generic_data( self, key: str, - value: Any, + value: object, table_name: Literal["users", "keys", "config", "spend"], ): """ @@ -3227,21 +3433,22 @@ async def get_generic_data( async def _do_query(): if table_name == "users": - return await UserRepository(self).table.find_first( + return await self.user_table.find_first( where={key: value} # type: ignore ) elif table_name == "keys": - return await VerificationTokenRepository(self).table.find_first( # type: ignore + return await self.token_table.find_first( where={key: value} # type: ignore ) elif table_name == "config": - return await ConfigRepository(self).table.find_first( # type: ignore + return await self.config_table.find_first( where={key: value} # type: ignore ) elif table_name == "spend": - return await self.db.l.find_first( # type: ignore - where={key: value} # type: ignore - ) + return await cast( # cast-ok: delegated via untyped __getattr__; typed table protocol + "_PrismaTable[LiteLLM_SpendLogs]", + self.db.l, + ).find_first(where={key: value}) return None try: @@ -3268,7 +3475,7 @@ async def _do_query(): 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: object) -> dict[str, Any] | None: """ Execute a query, recovering once from PostgreSQL's "cached plan must not change result type" error. @@ -3301,7 +3508,7 @@ async def _query_first_with_cached_plan_fallback(self, sql_query: str, *args) -> attempt reconnects once the cooldown elapses. """ try: - return await self.db.query_first(sql_query, *args) + return await self.raw_db.query_first(sql_query, *args) except Exception as e: if "cached plan must not change result type" not in str(e): raise @@ -3312,7 +3519,7 @@ async def _query_first_with_cached_plan_fallback(self, sql_query: str, *args) -> "changes are applied." ) await self.attempt_db_reconnect(reason="postgres_cached_plan_error") - return await self.db.query_first(sql_query, *args) + return await self.raw_db.query_first(sql_query, *args) @backoff.on_exception( backoff.expo, @@ -3326,10 +3533,10 @@ async def get_data( self, token: Optional[Union[str, list]] = None, user_id: Optional[str] = None, - user_id_list: Optional[list] = None, + user_id_list: Sequence[str] | None = None, team_id: Optional[str] = None, - team_id_list: Optional[list] = None, - key_val: Optional[dict] = None, + team_id_list: Sequence[str] | None = None, + key_val: Mapping[str, object] | None = None, table_name: Optional[ Literal[ "user", @@ -3370,7 +3577,7 @@ async def get_data( status_code=400, detail={"error": f"No token passed in. Token={token}"}, ) - response = await VerificationTokenRepository(self).table.find_unique( + response = await self.token_table.find_unique( where={"token": hashed_token}, # type: ignore include={"litellm_budget_table": True}, ) @@ -3385,7 +3592,7 @@ async def get_data( detail=f"Authentication Error: invalid user key - user key does not exist in db. User Key={token}", ) elif query_type == "find_all" and user_id is not None: - response = await VerificationTokenRepository(self).table.find_many( + response = await self.token_table.find_many( where={"user_id": user_id}, include={"litellm_budget_table": True}, ) @@ -3394,7 +3601,7 @@ async def get_data( if isinstance(r.expires, datetime): r.expires = r.expires.isoformat() elif query_type == "find_all" and team_id is not None: - response = await VerificationTokenRepository(self).table.find_many( + response = await self.token_table.find_many( take=limit, where={"team_id": team_id}, include={"litellm_budget_table": True}, @@ -3404,7 +3611,7 @@ async def get_data( if isinstance(r.expires, datetime): r.expires = r.expires.isoformat() elif query_type == "find_all" and expires is not None and reset_at is not None: - response = await VerificationTokenRepository(self).table.find_many( + response = await self.token_table.find_many( where={ # type: ignore "OR": [ {"expires": None}, @@ -3418,7 +3625,7 @@ async def get_data( if isinstance(r.expires, datetime): r.expires = r.expires.isoformat() elif query_type == "find_all": - where_filter: dict = {} + where_filter: dict[str, dict[str, object]] = {} if token is not None: where_filter["token"] = {} if isinstance(token, str): @@ -3434,7 +3641,7 @@ async def get_data( else: hashed_tokens.append(t) where_filter["token"]["in"] = hashed_tokens - response = await VerificationTokenRepository(self).table.find_many( + response = await self.token_table.find_many( order={"spend": "desc"}, where=where_filter, # type: ignore include={"litellm_budget_table": True}, @@ -3452,17 +3659,15 @@ async def get_data( if key_val is None: key_val = {"user_id": user_id} - response = await UserRepository(self).table.find_unique( # type: ignore - where=key_val, # type: ignore + response = await self.user_table.find_unique( + where=key_val, include={"organization_memberships": True}, ) elif query_type == "find_all" and key_val is not None: - response = await UserRepository(self).table.find_many( - where=key_val # type: ignore - ) # type: ignore + response = await self.user_table.find_many(where=key_val) elif query_type == "find_all" and reset_at is not None: - response = await UserRepository(self).table.find_many( + response = await self.user_table.find_many( where={ # type: ignore # A user seeded from default_internal_user_params # (or created via /user/new without an explicit @@ -3485,15 +3690,15 @@ async def get_data( } ) elif query_type == "find_all" and user_id_list is not None: - response = await UserRepository(self).table.find_many(where={"user_id": {"in": user_id_list}}) + response = await self.user_table.find_many(where={"user_id": {"in": user_id_list}}) elif query_type == "find_all": if expires is not None: - response = await UserRepository(self).table.find_many( # type: ignore + response = await self.user_table.find_many( order={"spend": "desc"}, - where={ # type: ignore + where={ "OR": [ - {"expires": None}, # type: ignore - {"expires": {"gt": expires}}, # type: ignore + {"expires": None}, + {"expires": {"gt": expires}}, ], }, ) @@ -3512,32 +3717,32 @@ async def get_data( LIMIT $1 OFFSET $2 """ - response = await self.db.query_raw(sql_query, limit, offset) + response = await self.raw_db.query_raw(sql_query, limit, offset) return response elif table_name == "spend": verbose_proxy_logger.debug("PrismaClient: get_data: table_name == 'spend'") if key_val is not None: if query_type == "find_unique": - response = await SpendLogsRepository(self).table.find_unique( # type: ignore - where={ # type: ignore - key_val["key"]: key_val["value"], # type: ignore - } + response = await self.spend_log_table.find_unique( + where=cast( # cast-ok: spend lookups filter on a caller-chosen column + Mapping[str, object], {key_val["key"]: key_val["value"]} + ) ) elif query_type == "find_all": - response = await SpendLogsRepository(self).table.find_many( # type: ignore - where={ - key_val["key"]: key_val["value"], # type: ignore - } + response = await self.spend_log_table.find_many( + where=cast( # cast-ok: spend lookups filter on a caller-chosen column + Mapping[str, object], {key_val["key"]: key_val["value"]} + ) ) return response else: - response = await SpendLogsRepository(self).table.find_many( # type: ignore + response = await self.spend_log_table.find_many( order={"startTime": "desc"}, ) return response elif table_name == "budget" and reset_at is not None: if query_type == "find_all": - response = await BudgetRepository(self).table.find_many( + response = await self.budget_table.find_many( where={ # type: ignore "OR": [ { @@ -3554,18 +3759,16 @@ async def get_data( elif table_name == "enduser" and budget_id_list is not None: if query_type == "find_all": - response = await EndUserRepository(self).table.find_many( - where={"budget_id": {"in": budget_id_list}} - ) + response = await self.end_user_table.find_many(where={"budget_id": {"in": budget_id_list}}) return response elif table_name == "team": if query_type == "find_unique": - response = await TeamRepository(self).table.find_unique( + response = await self.team_table.find_unique( where={"team_id": team_id}, # type: ignore include={"litellm_model_table": True}, # type: ignore ) elif query_type == "find_all" and reset_at is not None: - response = await TeamRepository(self).table.find_many( + response = await self.team_table.find_many( where={ # type: ignore # Same NULL budget_reset_at gap as the user query # above: a team with a budget_duration but no @@ -3582,24 +3785,24 @@ async def get_data( } ) elif query_type == "find_all" and user_id is not None: - response = await TeamRepository(self).table.find_many( + response = await self.team_table.find_many( where={ "members": {"has": user_id}, }, include={"litellm_budget_table": True}, ) elif query_type == "find_all" and team_id_list is not None: - response = await TeamRepository(self).table.find_many(where={"team_id": {"in": team_id_list}}) + response = await self.team_table.find_many(where={"team_id": {"in": team_id_list}}) elif query_type == "find_all" and team_id_list is None: - response = await TeamRepository(self).table.find_many(take=MAX_TEAM_LIST_LIMIT) + response = await self.team_table.find_many(take=MAX_TEAM_LIST_LIMIT) return response elif table_name == "user_notification": if query_type == "find_unique": - response = await UserNotificationsRepository(self).table.find_unique( # type: ignore + response = await self.user_notification_table.find_unique( where={"user_id": user_id} # type: ignore ) elif query_type == "find_all": - response = await UserNotificationsRepository(self).table.find_many() # type: ignore + response = await self.user_notification_table.find_many() return response elif table_name == "combined_view": # check if plain text or hash @@ -3735,7 +3938,7 @@ async def get_data( ) raise e - def jsonify_team_object(self, db_data: dict): + def jsonify_team_object(self, db_data: dict[str, Any]) -> dict[str, Any]: db_data = self.jsonify_object(data=db_data) if db_data.get("members_with_roles", None) is not None and isinstance(db_data["members_with_roles"], list): db_data["members_with_roles"] = json.dumps(db_data["members_with_roles"]) @@ -3753,7 +3956,7 @@ def jsonify_team_object(self, db_data: dict): ) async def insert_data( self, - data: dict, + data: dict[str, Any], table_name: Literal["user", "key", "config", "spend", "team", "user_notification"], ): """ @@ -3775,7 +3978,7 @@ async def insert_data( if db_data.get("budget_limits") is None: db_data.pop("budget_limits", None) print_verbose("PrismaClient: Before upsert into litellm_verificationtoken") - new_verification_token = await VerificationTokenRepository(self).table.upsert( # type: ignore + new_verification_token = await self.token_table.upsert( where={ "token": hashed_token, }, @@ -3790,7 +3993,7 @@ async def insert_data( elif table_name == "user": db_data = self.jsonify_object(data=data) try: - new_user_row = await UserRepository(self).table.upsert( + new_user_row = await self.user_table.upsert( where={"user_id": data["user_id"]}, data={ "create": {**db_data}, # type: ignore @@ -3813,7 +4016,7 @@ async def insert_data( return new_user_row elif table_name == "team": db_data = self.jsonify_team_object(db_data=data) - new_team_row = await TeamRepository(self).table.upsert( + new_team_row = await self.team_table.upsert( where={"team_id": data["team_id"]}, data={ "create": {**db_data}, # type: ignore @@ -3835,7 +4038,7 @@ async def insert_data( for k, v in data.items(): updated_data = v updated_data = json.dumps(updated_data) - updated_table_row = ConfigRepository(self).table.upsert( + updated_table_row = self.config_table.upsert( where={"param_name": k}, # type: ignore data={ "create": {"param_name": k, "param_value": updated_data}, # type: ignore @@ -3851,7 +4054,7 @@ async def insert_data( verbose_proxy_logger.info("Data Inserted into Config Table") elif table_name == "spend": db_data = self.jsonify_object(data=data) - new_spend_row = await SpendLogsRepository(self).table.upsert( + new_spend_row = await self.spend_log_table.upsert( where={"request_id": data["request_id"]}, data={ "create": {**db_data}, # type: ignore @@ -3862,7 +4065,7 @@ async def insert_data( return new_spend_row elif table_name == "user_notification": db_data = self.jsonify_object(data=data) - new_user_notification_row = await UserNotificationsRepository(self).table.upsert( # type: ignore + new_user_notification_row = await self.user_notification_table.upsert( where={"request_id": data["request_id"]}, data={ "create": {**db_data}, # type: ignore @@ -3901,14 +4104,14 @@ async def insert_data( async def update_data( self, token: Optional[str] = None, - data: dict = {}, + data: dict[str, Any] = {}, data_list: Optional[List] = None, user_id: Optional[str] = None, team_id: Optional[str] = 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, + update_key_values: dict[str, Any] | None = None, + update_key_values_custom_query: dict[str, Any] | None = None, ): """ Update existing data @@ -3924,17 +4127,17 @@ async def update_data( # check if plain text or hash token = _hash_token_if_needed(token=token) db_data["token"] = token - response = await VerificationTokenRepository(self).table.update( + response = await self.token_table.update( where={"token": token}, # type: ignore data={**db_data}, # type: ignore ) verbose_proxy_logger.debug("\033[91m" + f"DB Token Table update succeeded {response}" + "\033[0m") - _data: dict = {} + _data: dict[str, Any] = {} if response is not None: try: _data = response.model_dump() # type: ignore except Exception: - _data = response.dict() + _data = cast("_LegacyModelDump", response).dict() # cast-ok: pydantic v1 dump fallback return {"token": token, "data": _data} elif user_id is not None or (table_name is not None and table_name == "user") and query_type == "update": """ @@ -3947,7 +4150,7 @@ async def update_data( update_key_values = update_key_values_custom_query else: update_key_values = db_data - update_user_row = await UserRepository(self).table.upsert( + update_user_row = await self.user_table.upsert( where={"user_id": user_id}, # type: ignore data={ "create": {**db_data}, # type: ignore @@ -3976,7 +4179,7 @@ async def update_data( update_key_values["members_with_roles"], list ): update_key_values["members_with_roles"] = json.dumps(update_key_values["members_with_roles"]) - update_team_row = await TeamRepository(self).table.upsert( + update_team_row = await self.team_table.upsert( where={"team_id": team_id}, # type: ignore data={ "create": {**db_data}, # type: ignore @@ -3999,7 +4202,7 @@ async def update_data( """ Batch write update queries """ - batcher = self.db.batch_() + batcher = self.raw_db.batch_() for idx, t in enumerate(data_list): # check if plain text or hash if t.token.startswith("sk-"): # type: ignore @@ -4024,7 +4227,7 @@ async def update_data( """ Batch write update queries """ - batcher = self.db.batch_() + batcher = self.raw_db.batch_() for idx, user in enumerate(data_list): try: data_json = self.jsonify_object(data=user.model_dump(exclude_none=True)) @@ -4051,7 +4254,7 @@ async def update_data( """ Batch write update queries """ - batcher = self.db.batch_() + batcher = self.raw_db.batch_() for enduser in data_list: try: data_json = self.jsonify_object(data=enduser.model_dump(exclude_none=True)) @@ -4078,7 +4281,7 @@ async def update_data( """ Batch write update queries """ - batcher = self.db.batch_() + batcher = self.raw_db.batch_() for budget in data_list: try: data_json = self.jsonify_object(data=budget.model_dump(exclude_none=True)) @@ -4103,7 +4306,7 @@ async def update_data( and isinstance(data_list, list) ): # Batch write update queries - batcher = self.db.batch_() + batcher = self.raw_db.batch_() for idx, team in enumerate(data_list): try: data_json = self.jsonify_team_object(db_data=team.model_dump(exclude_none=True)) @@ -4150,7 +4353,7 @@ async def update_data( async def delete_data( self, tokens: Optional[List] = None, - team_id_list: Optional[List] = None, + team_id_list: Sequence[str] | None = None, table_name: Optional[Literal["user", "key", "config", "spend", "team"]] = None, user_id: Optional[str] = None, ): @@ -4169,24 +4372,24 @@ async def delete_data( else: hashed_token = token hashed_tokens.append(hashed_token) - filter_query: dict = {} + filter_query: dict[str, object] = {} if user_id is not None: filter_query = {"AND": [{"token": {"in": hashed_tokens}}, {"user_id": user_id}]} else: filter_query = {"token": {"in": hashed_tokens}} - deleted_tokens = await VerificationTokenRepository(self).table.delete_many( + deleted_tokens = await self.token_table.delete_many( where=filter_query # type: ignore ) 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): # admin only endpoint -> `/team/delete` - await TeamRepository(self).table.delete_many(where={"team_id": {"in": team_id_list}}) + await self.team_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): # admin only endpoint -> `/team/delete` - await VerificationTokenRepository(self).table.delete_many(where={"team_id": {"in": team_id_list}}) + await self.token_table.delete_many(where={"team_id": {"in": team_id_list}}) except Exception as e: import traceback @@ -4901,7 +5104,7 @@ async def _db_health_watchdog_loop(self) -> None: try: await asyncio.sleep(self._db_health_watchdog_interval_seconds) await asyncio.wait_for( - self.db.query_raw("SELECT 1"), + self.raw_db.query_raw("SELECT 1"), timeout=self._db_health_watchdog_probe_timeout_seconds, ) if isinstance(self.db, RoutingPrismaWrapper) and self.db.writer_unavailable: @@ -4937,7 +5140,7 @@ async def health_check(self): # Execute the raw query # The asterisk before `user_id_list` unpacks the list into separate arguments - response = await self.db.query_raw(sql_query) + response = await self.raw_db.query_raw(sql_query) return response except Exception as e: import traceback @@ -4975,7 +5178,7 @@ async def _fetch_row_count() -> int: FROM pg_class WHERE oid = '"LiteLLM_SpendLogs"'::regclass; """ - result = await self.db.query_raw(query=sql_query) + result = await self.raw_db.query_raw(query=sql_query) return result[0]["reltuples"] try: @@ -5062,7 +5265,13 @@ async def save_health_check_result( health_check_data.update({k: v for k, v in optional_fields.items() if v is not None}) verbose_proxy_logger.debug(f"Saving health check data: {health_check_data}") - return await HealthCheckRepository(self).table.create(data=health_check_data) + health_check_table = ( + cast( # cast-ok: repository .table is declared Any; typed table protocol pins the real shape + "_PrismaTable[LiteLLM_HealthCheckTable]", + HealthCheckRepository(self).table, + ) + ) + return await health_check_table.create(data=health_check_data) except Exception as e: verbose_proxy_logger.error(f"Error saving health check result for model {model_name}: {e}") @@ -5074,18 +5283,24 @@ async def get_health_check_history( limit: int = 100, offset: int = 0, status_filter: Optional[str] = None, - ): + ) -> Sequence["LiteLLM_HealthCheckTable"]: """ Get health check history with optional filtering """ try: - where_clause = {} + where_clause: dict[str, str] = {} if model_name: where_clause["model_name"] = model_name if status_filter: where_clause["status"] = status_filter - results = await HealthCheckRepository(self).table.find_many( + health_check_table = ( + cast( # cast-ok: repository .table is declared Any; typed table protocol pins the real shape + "_PrismaTable[LiteLLM_HealthCheckTable]", + HealthCheckRepository(self).table, + ) + ) + results = await health_check_table.find_many( where=where_clause, order={"checked_at": "desc"}, take=limit, @@ -5096,7 +5311,7 @@ async def get_health_check_history( verbose_proxy_logger.error(f"Error getting health check history: {e}") return [] - async def get_all_latest_health_checks(self): + async def get_all_latest_health_checks(self) -> Sequence["LiteLLM_HealthCheckTable"]: """ Get the latest health check for each model. @@ -5104,7 +5319,13 @@ async def get_all_latest_health_checks(self): (via Prisma ``distinct`` + ``order``) so we never load the full history into memory. """ try: - return await HealthCheckRepository(self).table.find_many( + health_check_table = ( + cast( # cast-ok: repository .table is declared Any; typed table protocol pins the real shape + "_PrismaTable[LiteLLM_HealthCheckTable]", + HealthCheckRepository(self).table, + ) + ) + return await health_check_table.find_many( distinct=["model_id", "model_name"], order=[ {"model_id": "asc"}, @@ -5261,13 +5482,17 @@ def verify_password(password: str, stored: str) -> bool: return False -async def migrate_passwords_to_scrypt_async(prisma_client) -> str: +async def migrate_passwords_to_scrypt_async(prisma_client: "PrismaClient") -> str: """ Migrate plaintext passwords in the DB to scrypt. SHA256 passwords are left alone (they migrate on next login via the SHA256 fallback). Skips quickly if no plaintext passwords exist. """ - all_with_pw = await UserRepository(prisma_client).table.find_many( + user_table = cast( # cast-ok: repository .table is declared Any; typed table protocol pins the real shape + "_PrismaTable[LiteLLM_UserTable]", + UserRepository(prisma_client).table, + ) + all_with_pw = await user_table.find_many( where={"password": {"not": None}}, ) @@ -5281,7 +5506,9 @@ def _is_sha256_hex(s: str) -> bool: return "No plaintext passwords found" for user in plaintext_users: - await UserRepository(prisma_client).table.update( + if user.password is None: + continue + await user_table.update( where={"user_id": user.user_id}, data={"password": hash_password(user.password)}, ) @@ -6382,7 +6609,7 @@ def validate_model_access( def model_dump_with_preserved_fields( - obj: Any, + obj: "BaseModel", preserve_fields: Optional[List[str]] = None, exclude_unset: bool = True, ) -> Dict[str, Any]: @@ -6403,11 +6630,17 @@ def model_dump_with_preserved_fields( """ result = obj.model_dump(exclude_none=True, exclude_unset=exclude_unset) - choices = result.get("choices") + choices = cast( # cast-ok: model_dump of a response model yields a JSON dict per choice + list[dict[str, dict[str, object]]] | None, + result.get("choices"), + ) if not choices: return result - obj_choices = obj.choices + obj_choices = cast( # cast-ok: models whose dump contains "choices" expose the matching attribute + Sequence[object], + getattr(obj, "choices"), + ) for choice_obj, choice_dict in zip(obj_choices, choices): for sub_object, field_name in _PRESERVED_NONE_FIELDS: sub_dict = choice_dict.get(sub_object)