From add2a1ce3f6c5c7f2462091d97b70cb7c889bbc0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 1 Aug 2026 03:57:48 -0700 Subject: [PATCH] chore(typing): clear 2.4k basedpyright errors across 15 Any hotspot files Replace Any-typed seams with real types in the files carrying the highest reportAny/reportExplicitAny density. The dominant source was the repository layer: BaseRepository.table is declared Any, so every repository read poisoned its rows and every downstream call. Typed pass-through accessors under a _PrismaTableActions Protocol pay that crossing once per table, and TypedDicts and Protocols replace the remaining Any-typed request, row, and tool payloads across the team, key, SCIM, spend, MCP, guardrail, video, and websearch surfaces No casts, no type: ignore, no noqa, no new Any annotations, no behavior changes. Whole-tree basedpyright: reportAny 20,840 -> 19,397, reportExplicitAny 7,253 -> 6,518, all rules 151,424 -> 149,066, with no rule increased in any file. Budgets ratcheted: basedpyright -2,358, ruff-strict -300, type-discipline -68 --- basedpyright-code-budget.json | 18 +- .../websearch_interception/handler.py | 107 ++-- .../context_management/editors/compact.py | 134 +++-- litellm/llms/custom_httpx/llm_http_handler.py | 558 +++++++++--------- litellm/llms/openai/videos/transformation.py | 61 +- .../mcp_server/mcp_server_manager.py | 188 ++++-- .../mcp_server/sampling_handler.py | 209 ++++--- .../proxy/guardrails/guardrail_endpoints.py | 96 ++- .../internal_user_endpoints.py | 159 +++-- .../key_management_endpoints.py | 174 ++++-- .../model_management_endpoints.py | 17 +- .../management_endpoints/scim/scim_v2.py | 255 +++++--- .../management_endpoints/team_endpoints.py | 297 +++++++--- .../spend_management_endpoints.py | 368 +++++++++--- .../mcp/litellm_proxy_mcp_handler.py | 72 ++- litellm/videos/main.py | 204 +++---- ruff-strict-budget.json | 20 +- type-discipline-budget.json | 8 +- 18 files changed, 1903 insertions(+), 1042 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 43df27ea2e2..f6dd90077b1 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,6 +1,6 @@ { "reportAny": { - "limit": 31256 + "limit": 29813 }, "reportArgumentType": { "limit": 2645 @@ -24,7 +24,7 @@ "limit": 42 }, "reportExplicitAny": { - "limit": 10208 + "limit": 9473 }, "reportFunctionMemberAccess": { "limit": 11 @@ -54,10 +54,10 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5869 + "limit": 5855 }, "reportMissingTypeArgument": { - "limit": 15861 + "limit": 15852 }, "reportMissingTypeStubs": { "limit": 41 @@ -99,19 +99,19 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 45357 + "limit": 45324 }, "reportUnknownLambdaType": { "limit": 113 }, "reportUnknownMemberType": { - "limit": 40477 + "limit": 40452 }, "reportUnknownParameterType": { - "limit": 20338 + "limit": 20309 }, "reportUnknownVariableType": { - "limit": 32047 + "limit": 31978 }, "reportUnnecessaryCast": { "limit": 177 @@ -123,7 +123,7 @@ "limit": 7 }, "reportUnnecessaryIsInstance": { - "limit": 1205 + "limit": 1204 }, "reportUntypedBaseClass": { "limit": 165 diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 21d990e8e60..531caf273f1 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -9,7 +9,8 @@ server-side using litellm router's search tools. import asyncio import math import uuid -from typing import Any, Dict, List, Optional, Tuple, Union, cast +from collections.abc import AsyncIterator, Mapping +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast import litellm from litellm._logging import verbose_logger @@ -29,19 +30,31 @@ from litellm.integrations.websearch_interception.transformation import ( WebSearchTransformation, ) from litellm.llms.base_llm.search.transformation import SearchResponse -from litellm.types.integrations.websearch_interception import ( - WebSearchInterceptionConfig, -) from litellm.types.integrations.custom_logger import ( CHAT_COMPLETION_AGENTIC_SURFACE, RESPONSES_AGENTIC_SURFACE, AgenticLoopPlan, AgenticLoopRequestPatch, ) +from litellm.types.integrations.websearch_interception import ( + WebSearchInterceptionConfig, +) from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import CallTypes, LlmProviders from litellm.utils import ProviderConfigManager +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.llms.base_llm.anthropic_messages.transformation import ( + BaseAnthropicMessagesConfig, + ) + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.llms.anthropic_messages.anthropic_response import ( + AnthropicMessagesResponse, + ) + from litellm.types.utils import ModelResponse + from litellm.utils import CustomStreamWrapper + # Key used to flag, on per-request kwargs, that the originating client sent # an Anthropic-native ``web_search_*`` tool — meaning the final response # should include ``web_search_tool_result`` content blocks so the client @@ -94,8 +107,8 @@ class WebSearchInterceptionLogger(CustomLogger): messages: List[Dict], tools: Optional[List[Dict]], custom_llm_provider: Optional[str], - kwargs: Optional[dict[str, Any]] = None, - ) -> Optional[Dict[str, Any]]: + kwargs: Mapping[str, object] | None = None, + ) -> dict[str, object] | None: """ Short-circuit web-search-only requests by executing the search directly. @@ -188,7 +201,7 @@ class WebSearchInterceptionLogger(CustomLogger): verbose_logger.error(f"WebSearchInterception: Short-circuit search failed: {e}") search_result_text, structured = f"Search failed: {e}", None - content: List[Dict[str, Any]] = [] + content: list[dict[str, object]] = [] if native_tool is not None: tool_use_id = f"srvtoolu_{uuid.uuid4().hex}" tool_name = native_tool.get("name") or "web_search" @@ -210,7 +223,7 @@ class WebSearchInterceptionLogger(CustomLogger): # github_copilot, etc.) see the same payload they always have. content.append({"type": "text", "text": search_result_text}) - response: Dict[str, Any] = { + response: dict[str, object] = { "id": f"msg_{str(uuid.uuid4())}", "type": "message", "role": "assistant", @@ -228,7 +241,9 @@ class WebSearchInterceptionLogger(CustomLogger): ) return response - async def async_pre_call_deployment_hook(self, kwargs: Dict[str, Any], call_type: Optional[Any]) -> Optional[dict]: + async def async_pre_call_deployment_hook( + self, kwargs: Dict[str, Any], call_type: Optional[CallTypes] + ) -> Optional[dict]: """ Pre-call hook to convert native Anthropic web_search tools to regular tools. @@ -297,7 +312,7 @@ class WebSearchInterceptionLogger(CustomLogger): return kwargs - def _convert_responses_tools(self, kwargs: dict[str, Any], tools: list[dict[str, Any]]) -> dict | None: + def _convert_responses_tools(self, kwargs: Mapping[str, object], tools: list[dict[str, object]]) -> dict | None: """Convert Responses API web search tools to the LiteLLM standard function tool.""" if not any(is_web_search_tool_responses(tool) for tool in tools): return None @@ -370,7 +385,7 @@ class WebSearchInterceptionLogger(CustomLogger): return tool.get("name") @classmethod - def _sync_forced_tool_choice(cls, tool_choice: Any, converted_tools: list[dict[str, Any]]) -> Any: + def _sync_forced_tool_choice(cls, tool_choice: Any, converted_tools: list[dict[str, object]]) -> object: """Repoint a forced ``tool_choice`` at ``litellm_web_search`` when it names a web-search tool that was just converted away. @@ -468,7 +483,7 @@ class WebSearchInterceptionLogger(CustomLogger): async def async_should_run_agentic_loop( self, - response: Any, + response: object, model: str, messages: List[Dict], tools: Optional[List[Dict]], @@ -578,7 +593,7 @@ class WebSearchInterceptionLogger(CustomLogger): async def async_should_run_chat_completion_agentic_loop( self, - response: Any, + response: object, model: str, messages: List[Dict], tools: Optional[List[Dict]], @@ -636,7 +651,7 @@ class WebSearchInterceptionLogger(CustomLogger): async def async_should_run_responses_agentic_loop( self, - response: Any, + response: object, model: str, messages: list[dict], tools: list[dict] | None, @@ -687,13 +702,13 @@ class WebSearchInterceptionLogger(CustomLogger): tools: Dict, model: str, messages: List[Dict], - response: Any, - anthropic_messages_provider_config: Any, + response: object, + anthropic_messages_provider_config: "BaseAnthropicMessagesConfig | None", anthropic_messages_optional_request_params: Dict, - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj | None", stream: bool, kwargs: Dict, - ) -> Any: + ) -> "AnthropicMessagesResponse | AsyncIterator[object]": """ Execute agentic loop with WebSearch execution for Anthropic Messages API. @@ -721,10 +736,10 @@ class WebSearchInterceptionLogger(CustomLogger): tools: Dict, model: str, messages: List[Dict], - response: Any, - anthropic_messages_provider_config: Any, + response: object, + anthropic_messages_provider_config: "BaseAnthropicMessagesConfig | None", anthropic_messages_optional_request_params: Dict, - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj | None", stream: bool, kwargs: Dict, ) -> AgenticLoopPlan: @@ -764,7 +779,7 @@ class WebSearchInterceptionLogger(CustomLogger): kwargs=kwargs, ) - metadata: Dict[str, Any] = { + metadata: dict[str, object] = { "tool_type": "websearch", "response_format": "anthropic", } @@ -787,10 +802,10 @@ class WebSearchInterceptionLogger(CustomLogger): async def async_post_agentic_loop_response_hook( self, - response: Any, + response: object, plan: AgenticLoopPlan, kwargs: Dict, - ) -> Any: + ) -> object: """ Inject Anthropic-native ``web_search_tool_result`` blocks into the final response when the originating client used a native @@ -810,9 +825,9 @@ class WebSearchInterceptionLogger(CustomLogger): def _build_native_result_blocks( tool_calls: List[Dict], structured_results: List[Optional[SearchResponse]], - ) -> List[Dict[str, Any]]: + ) -> list[dict[str, object]]: """Build one ``web_search_tool_result`` block per tool_call.""" - blocks: List[Dict[str, Any]] = [] + blocks: list[dict[str, object]] = [] for i, tool_call in enumerate(tool_calls): tool_use_id = tool_call.get("id") or "" structured = structured_results[i] if i < len(structured_results) else None @@ -825,7 +840,7 @@ class WebSearchInterceptionLogger(CustomLogger): return blocks @staticmethod - def _inject_native_blocks(response: Any, native_blocks: List[Dict[str, Any]]) -> Any: + def _inject_native_blocks(response: Any, native_blocks: list[dict[str, object]]) -> Any: """Prepend native blocks to response content, dict or object form.""" if not native_blocks: return response @@ -849,12 +864,12 @@ class WebSearchInterceptionLogger(CustomLogger): tools: Dict, model: str, messages: List[Dict], - response: Any, + response: object, optional_params: Dict, - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj | None", stream: bool, kwargs: Dict, - ) -> Any: + ) -> "ModelResponse | CustomStreamWrapper": """ Execute agentic loop with WebSearch execution for Chat Completions API. @@ -884,9 +899,9 @@ class WebSearchInterceptionLogger(CustomLogger): tools: Dict, model: str, messages: List[Dict], - response: Any, + response: object, optional_params: Dict, - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj | None", stream: bool, kwargs: Dict, ) -> AgenticLoopPlan: @@ -911,9 +926,9 @@ class WebSearchInterceptionLogger(CustomLogger): tools: dict, model: str, messages: list[dict], - response: Any, + response: object, optional_params: dict, - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj | None", stream: bool, kwargs: dict, ) -> AgenticLoopPlan: @@ -1023,7 +1038,7 @@ class WebSearchInterceptionLogger(CustomLogger): return [] @staticmethod - def _extract_search_text(result: Any) -> str: + def _extract_search_text(result: object) -> str: if isinstance(result, Exception): verbose_logger.error(f"WebSearchInterception: Responses search failed with error: {str(result)}") return f"Search failed: {str(result)}" @@ -1091,10 +1106,10 @@ class WebSearchInterceptionLogger(CustomLogger): tool_calls: List[Dict], thinking_blocks: List[Dict], anthropic_messages_optional_request_params: Dict, - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj | None", stream: bool, kwargs: Dict, - ) -> Any: + ) -> "AnthropicMessagesResponse | AsyncIterator[object]": """Legacy path: execute search + build patch + run follow-up call.""" request_patch, structured_results = await self._build_anthropic_request_patch( model=model, @@ -1118,7 +1133,7 @@ class WebSearchInterceptionLogger(CustomLogger): if max_tokens is None: max_tokens = cast(int, kwargs.get("max_tokens", 1024)) - response = await anthropic_messages.acreate( + response: AnthropicMessagesResponse | AsyncIterator[object] = await anthropic_messages.acreate( max_tokens=max_tokens, messages=request_patch.messages, model=request_patch.model or model, @@ -1145,7 +1160,7 @@ class WebSearchInterceptionLogger(CustomLogger): tool_calls: List[Dict], thinking_blocks: List[Dict], anthropic_messages_optional_request_params: Dict, - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj | None", kwargs: Dict, ) -> Tuple[AgenticLoopRequestPatch, List[Optional[SearchResponse]]]: """ @@ -1238,7 +1253,7 @@ class WebSearchInterceptionLogger(CustomLogger): return patch, structured_results async def _execute_search( - self, query: str, kwargs: Optional[dict[str, Any]] = None + self, query: str, kwargs: Mapping[str, object] | None = None ) -> Tuple[str, Optional[SearchResponse]]: """ Execute a single web search using router's search tools. @@ -1300,8 +1315,8 @@ class WebSearchInterceptionLogger(CustomLogger): async def _authorize_search_tool( self, - search_tool: dict[str, Any], - kwargs: Optional[dict[str, Any]], + search_tool: Mapping[str, object], + kwargs: Mapping[str, object] | None, ) -> None: search_tool_name = search_tool.get("search_tool_name") if not isinstance(search_tool_name, str) or not search_tool_name: @@ -1343,7 +1358,7 @@ class WebSearchInterceptionLogger(CustomLogger): ) @staticmethod - def _get_user_api_key_auth_from_kwargs(kwargs: Optional[dict[str, Any]]) -> Any: + def _get_user_api_key_auth_from_kwargs(kwargs: Mapping[str, object] | None) -> "UserAPIKeyAuth | None": if not kwargs: return None @@ -1363,7 +1378,7 @@ class WebSearchInterceptionLogger(CustomLogger): return None - def _select_search_tool_from_router(self, llm_router: Any) -> Optional[dict[str, Any]]: + def _select_search_tool_from_router(self, llm_router: object) -> Optional[dict[str, Any]]: if llm_router is None or not hasattr(llm_router, "search_tools"): return None search_tools = list(getattr(llm_router, "search_tools") or []) @@ -1405,11 +1420,11 @@ class WebSearchInterceptionLogger(CustomLogger): messages: List[Dict], tool_calls: List[Dict], optional_params: Dict, - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj | None", stream: bool, kwargs: Dict, response_format: str = "openai", - ) -> Any: + ) -> "ModelResponse | CustomStreamWrapper": """Legacy path: execute search + build patch + run follow-up call.""" request_patch = await self._build_chat_completion_request_patch( model=model, diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py index f18a9f41939..c87014ffc1e 100644 --- a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py +++ b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py @@ -13,7 +13,8 @@ Mirrors Anthropic's native ``compact_20260112`` for non-Anthropic providers: """ import re -from typing import Any, Dict, List, Literal, Optional, Tuple, Union, cast +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union, cast import litellm from litellm._logging import verbose_logger @@ -23,6 +24,18 @@ from litellm.types.llms.anthropic import ( UsageIteration, ) +if TYPE_CHECKING: + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + from litellm.proxy._types import UserAPIKeyAuth + from litellm.router import Router + from litellm.types.llms.anthropic import ( + AllAnthropicToolsValues, + AnthopicMessagesAssistantMessageParam, + AnthropicMessagesUserMessageParam, + ) + from litellm.types.llms.openai import ChatCompletionToolParam + from litellm.types.utils import ModelResponse + from ..constants import ( COMPACT_DEFAULT_INSTRUCTIONS, COMPACT_DEFAULT_TRIGGER_TOKENS, @@ -98,9 +111,9 @@ def _read_summary_max_tokens_setting() -> int: async def _check_summary_model_access( - user_api_key_auth: Any, + user_api_key_auth: Optional["UserAPIKeyAuth"], summary_model: str, - llm_router: Any, + llm_router: Optional["Router"], ) -> bool: """Return True when every model-allowlist scope on the parent request is satisfied for ``summary_model``. @@ -294,7 +307,7 @@ async def _check_summary_model_access( async def _check_summary_model_budget( - user_api_key_auth: Any, + user_api_key_auth: Optional["UserAPIKeyAuth"], summary_model: str, ) -> bool: """Return True when the caller is within their per-model budget for @@ -357,7 +370,7 @@ async def _check_summary_model_budget( async def _check_summary_model_rate_limit( - user_api_key_auth: Any, + user_api_key_auth: Optional["UserAPIKeyAuth"], summary_model: str, ) -> bool: """Return True when the caller is within their configured RPM/TPM limits @@ -433,7 +446,7 @@ async def _check_summary_model_rate_limit( def _find_latest_compaction_index( - messages: List[Dict[str, Any]], + messages: List[Dict[str, object]], ) -> Tuple[Optional[int], Optional[int]]: """Return (message_index, block_index) of the most recent compaction block. @@ -453,7 +466,7 @@ def _find_latest_compaction_index( def _slice_around_compaction_block( messages: List[Dict[str, Any]], -) -> Tuple[List[Dict[str, Any]], Optional[Dict[str, Any]]]: +) -> Tuple[List[Dict[str, object]], Optional[Dict[str, object]]]: """Apply Anthropic's "drop everything before the compaction block" rule. Returns ``(sliced_messages_with_compaction_block, compaction_block_dict)`` @@ -468,27 +481,26 @@ def _slice_around_compaction_block( original_msg = messages[msg_idx] original_content = original_msg["content"] - compaction_block = cast(Dict[str, Any], original_content[blk_idx]) + compaction_block = cast(Dict[str, object], original_content[blk_idx]) # Per Anthropic's contract everything before the compaction block is # dropped, including earlier blocks within the same assistant message. sliced_content = list(original_content[blk_idx:]) - sliced_first_msg = {**original_msg, "content": sliced_content} - sliced_messages: List[Dict[str, Any]] = [sliced_first_msg] + sliced_messages: List[Dict[str, object]] = [{**original_msg, "content": sliced_content}] sliced_messages.extend(messages[msg_idx + 1 :]) return sliced_messages, compaction_block def _strip_compaction_blocks( - messages: List[Dict[str, Any]], -) -> List[Dict[str, Any]]: + messages: List[Dict[str, object]], +) -> List[Dict[str, object]]: """Drop any ``compaction`` content blocks from messages. Used to build the downstream-bound message list — the adapter has no concept of a compaction block, so it must not see one. """ - cleaned: List[Dict[str, Any]] = [] + cleaned: List[Dict[str, object]] = [] for msg in messages: content = msg.get("content") if not isinstance(content, list): @@ -503,9 +515,9 @@ def _strip_compaction_blocks( def _augment_system_with_summary( - system: Optional[Union[str, List[Dict[str, Any]]]], + system: Optional[Union[str, List[Dict[str, object]]]], summary_text: str, -) -> Union[str, List[Dict[str, Any]]]: +) -> Union[str, List[Dict[str, object]]]: """Prepend a "Previous conversation summary: ..." block to ``system``.""" prefix = f"{COMPACT_SUMMARY_SYSTEM_PREFIX}{summary_text}\n\n" if system is None: @@ -522,7 +534,7 @@ def _augment_system_with_summary( return [{"type": "text", "text": prefix.rstrip()}, *system] -def _resolve_trigger_tokens(edit_spec: Dict[str, Any]) -> Tuple[int, List[str]]: +def _resolve_trigger_tokens(edit_spec: Dict[str, object]) -> Tuple[int, List[str]]: """Validate and resolve ``trigger.value``. Raises ``AnthropicContextManagementError`` if the explicitly-supplied value @@ -556,7 +568,7 @@ def _resolve_trigger_tokens(edit_spec: Dict[str, Any]) -> Tuple[int, List[str]]: return value, warnings -def _build_summary_prompt(edit_spec: Dict[str, Any], tools: Optional[List[Dict[str, Any]]]) -> str: +def _build_summary_prompt(edit_spec: Dict[str, object], tools: Optional[List[Dict[str, object]]]) -> str: custom = edit_spec.get("instructions") if isinstance(custom, str) and custom.strip(): return custom @@ -567,8 +579,8 @@ def _build_summary_prompt(edit_spec: Dict[str, Any], tools: Optional[List[Dict[s def _propagate_metadata( - parent_litellm_metadata: Optional[Dict[str, Any]], -) -> Dict[str, Any]: + parent_litellm_metadata: Optional[Mapping[str, object]], +) -> Dict[str, object]: """Extract the parent request's auth/spend-attribution fields for the summary subcall. The proxy attaches ``user_api_key``, ``user_api_key_team_id`` etc. to @@ -579,7 +591,7 @@ def _propagate_metadata( """ if not parent_litellm_metadata: return {} - propagated: Dict[str, Any] = {} + propagated: Dict[str, object] = {} for key in _PROPAGATED_METADATA_KEYS: if key in parent_litellm_metadata: propagated[key] = parent_litellm_metadata[key] @@ -588,10 +600,10 @@ def _propagate_metadata( def _count_effective_tokens( model: str, - effective_messages: List[Dict[str, Any]], - compaction_block: Optional[Dict[str, Any]], - tools: Optional[List[Dict[str, Any]]], - system: Optional[Union[str, List[Dict[str, Any]]]] = None, + effective_messages: List[Dict[str, object]], + compaction_block: Optional[CompactionBlock], + tools: Optional[List[Dict[str, object]]], + system: Optional[Union[str, List[Dict[str, object]]]] = None, ) -> int: """Token-count the conversation as it will appear downstream. @@ -609,25 +621,32 @@ def _count_effective_tokens( messages_without_compaction = _strip_compaction_blocks(effective_messages) adapter = LiteLLMAnthropicMessagesAdapter() try: - openai_shape = adapter.translate_anthropic_messages_to_openai(messages=cast(Any, messages_without_compaction)) + openai_shape = adapter.translate_anthropic_messages_to_openai( + messages=cast( + "List[Union[AnthropicMessagesUserMessageParam, AnthopicMessagesAssistantMessageParam]]", + messages_without_compaction, + ) + ) except Exception as e: verbose_logger.debug( "compact_20260112: anthropic→openai translation failed during token " "count, falling back to raw messages: %s", e, ) - openai_shape = cast(Any, messages_without_compaction) + openai_shape = messages_without_compaction # Translate Anthropic-shaped tools (``input_schema``) to OpenAI-shaped # tools (``{"type": "function", "function": {...}}``) so ``token_counter`` # gets a consistent format regardless of which counting path it uses. # An inaccurate tool token count here could cause the polyfill to skip # needed compaction or trigger unnecessary summarization. - openai_tools: Optional[List[Dict[str, Any]]] = None + openai_tools: Optional[List[Dict[str, object]]] = None if tools: try: - translated_tools, _ = adapter.translate_anthropic_tools_to_openai(tools=cast(Any, tools)) - openai_tools = cast(List[Dict[str, Any]], translated_tools) + translated_tools, _ = adapter.translate_anthropic_tools_to_openai( + tools=cast("List[AllAnthropicToolsValues]", tools) + ) + openai_tools = cast(List[Dict[str, object]], translated_tools) except Exception as e: verbose_logger.debug( "compact_20260112: anthropic→openai tools translation failed " @@ -638,8 +657,8 @@ def _count_effective_tokens( total = litellm.token_counter( model=model, - messages=cast(Any, openai_shape), - tools=cast(Any, openai_tools), + messages=cast(List[Dict[str, object]], openai_shape), + tools=cast("Optional[List[ChatCompletionToolParam]]", openai_tools), ) if compaction_block is not None: content = compaction_block.get("content") or "" @@ -652,7 +671,7 @@ def _count_effective_tokens( def _system_to_text( - system: Optional[Union[str, List[Dict[str, Any]]]], + system: Optional[Union[str, List[Dict[str, object]]]], ) -> str: """Flatten an Anthropic-style ``system`` value into a single string for token counting. Returns ``""`` when ``system`` carries no text.""" @@ -670,8 +689,8 @@ def _system_to_text( def _select_last_user_question( - messages: List[Dict[str, Any]], -) -> List[Dict[str, Any]]: + messages: List[Dict[str, object]], +) -> List[Dict[str, object]]: """Pick the most recent ``user`` turn that is a real question. Returns a one-element message list with any ``tool_result`` blocks @@ -735,10 +754,10 @@ def _system_to_openai_message( def _build_summary_messages( - effective_messages: List[Dict[str, Any]], + effective_messages: List[Dict[str, object]], prompt: str, - system: Optional[Union[str, List[Dict[str, Any]]]] = None, -) -> List[Dict[str, Any]]: + system: Optional[Union[str, List[Dict[str, object]]]] = None, +) -> List[Dict[str, object]]: """Build the OpenAI-shape message list for the summary call. The caller's ``system`` prompt is prepended (the default summarization @@ -753,7 +772,10 @@ def _build_summary_messages( stripped = _strip_compaction_blocks(effective_messages) try: openai_messages = LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai( - messages=cast(Any, stripped) + messages=cast( + "List[Union[AnthropicMessagesUserMessageParam, AnthopicMessagesAssistantMessageParam]]", + stripped, + ) ) except Exception as e: verbose_logger.warning( @@ -761,9 +783,9 @@ def _build_summary_messages( "building summary call; falling back to raw shape: %s", e, ) - openai_messages = cast(Any, stripped) + openai_messages = stripped - summary_messages: List[Dict[str, Any]] = [] + summary_messages: List[Dict[str, object]] = [] system_message = _system_to_openai_message(system) if system_message is not None: summary_messages.append(system_message) @@ -783,7 +805,7 @@ def _build_summary_messages( return summary_messages -def _is_user_message(msg: Any) -> bool: +def _is_user_message(msg: object) -> bool: return isinstance(msg, dict) and msg.get("role") == "user" @@ -805,12 +827,12 @@ def _append_text_to_content(content: Any, extra_text: str) -> Any: async def _call_summary_model( *, summary_model: str, - summary_messages: List[Dict[str, Any]], - metadata: Dict[str, Any], + summary_messages: List[Dict[str, object]], + metadata: Mapping[str, object], llm_router: Any, allowed_model_region: Optional[str] = None, max_tokens: int = COMPACT_SUMMARY_MAX_TOKENS, -) -> Any: +) -> Union["ModelResponse", "CustomStreamWrapper"]: """Invoke the configured summary model. Prefers ``llm_router.acompletion`` so the model alias resolves against the @@ -877,7 +899,7 @@ def _extract_response_text(response: Any) -> Optional[str]: return None -def _extract_usage(response: Any) -> Tuple[int, int]: +def _extract_usage(response: object) -> Tuple[int, int]: usage = getattr(response, "usage", None) if usage is None: return 0, 0 @@ -889,8 +911,8 @@ def _extract_usage(response: Any) -> Tuple[int, int]: def apply_client_compaction_block_history( *, - messages: List[Dict[str, Any]], - system: Optional[Union[str, List[Dict[str, Any]]]], + messages: List[Dict[str, object]], + system: Optional[Union[str, List[Dict[str, object]]]], ) -> Optional[PolyfillResult]: """Honor client-sent compaction blocks without a ``compact_20260112`` edit. @@ -911,7 +933,7 @@ def apply_client_compaction_block_history( ) prior_summary_text = prior_compaction_block.get("content") or "" - augmented_system: Union[str, List[Dict[str, Any]], None] = system + augmented_system: Union[str, List[Dict[str, object]], None] = system if isinstance(prior_summary_text, str) and prior_summary_text: augmented_system = _augment_system_with_summary(system, prior_summary_text) verbose_logger.info( @@ -936,13 +958,13 @@ def apply_client_compaction_block_history( async def apply_compact_20260112( *, model: str, - messages: List[Dict[str, Any]], - tools: Optional[List[Dict[str, Any]]], - system: Optional[Union[str, List[Dict[str, Any]]]], - edit_spec: Dict[str, Any], - litellm_metadata: Optional[Dict[str, Any]] = None, - llm_router: Any = None, - user_api_key_auth: Any = None, + messages: List[Dict[str, object]], + tools: Optional[List[Dict[str, object]]], + system: Optional[Union[str, List[Dict[str, object]]]], + edit_spec: Dict[str, object], + litellm_metadata: Optional[Mapping[str, object]] = None, + llm_router: Optional["Router"] = None, + user_api_key_auth: Optional["UserAPIKeyAuth"] = None, ) -> PolyfillResult: """Apply ``compact_20260112``; return a ``PolyfillResult``. @@ -971,7 +993,7 @@ async def apply_compact_20260112( # non-Anthropic backends (which would reject them). effective_messages, prior_compaction_block = _slice_around_compaction_block(messages) prior_summary_text = prior_compaction_block.get("content") if prior_compaction_block else None - augmented_system: Union[str, List[Dict[str, Any]], None] = system + augmented_system: Union[str, List[Dict[str, object]], None] = system if isinstance(prior_summary_text, str) and prior_summary_text: augmented_system = _augment_system_with_summary(system, prior_summary_text) verbose_logger.info( diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index d6acbaae434..b64d8e3dc7e 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -13,8 +13,10 @@ from typing import ( Iterator, List, Literal, + Mapping, Optional, Tuple, + TypeVar, Union, cast, get_type_hints, @@ -162,7 +164,11 @@ from .http_handler import get_shared_realtime_ssl_context if TYPE_CHECKING: from aiohttp import ClientSession + from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import ( + FakeAnthropicMessagesStreamIterator, + ) from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( AnthropicMessagesStreamingResponse, ) @@ -182,6 +188,8 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_ResponseT = TypeVar("_ResponseT") + def _google_genai_streaming_hidden_params( *, @@ -189,11 +197,11 @@ def _google_genai_streaming_hidden_params( litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, response_headers: httpx.Headers, -) -> Dict[str, Any]: +) -> Dict[str, object]: """Pre-stream metadata for proxy response headers (mirrors CustomStreamWrapper._hidden_params).""" from litellm.litellm_core_utils.core_helpers import process_response_headers - _model_info: Dict[str, Any] = dict(getattr(litellm_params, "model_info", None) or {}) + _model_info: Mapping[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 +218,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 +229,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 +241,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 +367,7 @@ class BaseLLMHTTPHandler: 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 +433,7 @@ class BaseLLMHTTPHandler: 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], @@ -474,7 +482,7 @@ class BaseLLMHTTPHandler: litellm_params=litellm_params, ) - data = provider_config.transform_request( + data: Dict[str, object] = provider_config.transform_request( model=model, messages=messages, optional_params=optional_params, @@ -651,7 +659,7 @@ class BaseLLMHTTPHandler: 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 +699,7 @@ class BaseLLMHTTPHandler: 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 +791,7 @@ class BaseLLMHTTPHandler: 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 +835,7 @@ class BaseLLMHTTPHandler: 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 @@ -846,10 +854,10 @@ class BaseLLMHTTPHandler: def _add_stream_param_to_request_body( self, - data: dict, + data: Dict[str, object], provider_config: BaseConfig, fake_stream: bool, - ) -> dict: + ) -> Dict[str, object]: """ Some providers like Bedrock invoke do not support the stream parameter in the request body, we only pass `stream` in the request body the provider supports it. """ @@ -1051,7 +1059,7 @@ class BaseLLMHTTPHandler: timeout: Optional[Union[float, httpx.Timeout]], model_response: RerankResponse, _is_async: bool = False, - headers: Optional[Dict[str, Any]] = None, + headers: Optional[Dict[str, object]] = None, api_key: Optional[str] = None, api_base: Optional[str] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, @@ -1177,7 +1185,7 @@ class BaseLLMHTTPHandler: logging_obj: LiteLLMLoggingObj, api_key: Optional[str], api_base: Optional[str], - headers: Optional[Dict[str, Any]], + headers: Optional[Dict[str, object]], provider_config: BaseAudioTranscriptionConfig, ) -> Tuple[dict, str, Union[dict, bytes, None], Optional[dict]]: """ @@ -1266,10 +1274,10 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, atranscription: bool = False, - headers: Optional[Dict[str, Any]] = None, + headers: Optional[Dict[str, object]] = None, provider_config: Optional[BaseAudioTranscriptionConfig] = None, shared_session: Optional["ClientSession"] = None, - ) -> Union[TranscriptionResponse, Coroutine[Any, Any, TranscriptionResponse]]: + ) -> Union[TranscriptionResponse, Coroutine[object, object, TranscriptionResponse]]: if provider_config is None: raise ValueError(f"No provider config found for model: {model} and provider: {custom_llm_provider}") @@ -1351,7 +1359,7 @@ class BaseLLMHTTPHandler: api_base: Optional[str], custom_llm_provider: str, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - headers: Optional[Dict[str, Any]] = None, + headers: Optional[Dict[str, object]] = None, provider_config: Optional[BaseAudioTranscriptionConfig] = None, shared_session: Optional["ClientSession"] = None, ) -> TranscriptionResponse: @@ -1417,7 +1425,7 @@ class BaseLLMHTTPHandler: logging_obj: LiteLLMLoggingObj, api_key: Optional[str], api_base: Optional[str], - headers: Optional[Dict[str, Any]], + headers: Optional[Dict[str, object]], provider_config: BaseOCRConfig, litellm_params: dict, ) -> Tuple[Dict[str, Any], str, Dict[str, Any], None]: @@ -1483,7 +1491,7 @@ class BaseLLMHTTPHandler: logging_obj: LiteLLMLoggingObj, api_key: Optional[str], api_base: Optional[str], - headers: Optional[Dict[str, Any]], + headers: Optional[Dict[str, object]], provider_config: BaseOCRConfig, litellm_params: dict, ) -> Tuple[Dict[str, Any], str, Dict[str, Any], None]: @@ -1567,10 +1575,10 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, aocr: bool = False, - headers: Optional[Dict[str, Any]] = None, + headers: Optional[Dict[str, object]] = None, provider_config: Optional[BaseOCRConfig] = None, litellm_params: Optional[dict] = None, - ) -> Union[OCRResponse, Coroutine[Any, Any, OCRResponse]]: + ) -> Union[OCRResponse, Coroutine[object, object, OCRResponse]]: """ Sync OCR handler. """ @@ -1641,7 +1649,7 @@ class BaseLLMHTTPHandler: api_base: Optional[str], custom_llm_provider: str, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - headers: Optional[Dict[str, Any]] = None, + headers: Optional[Dict[str, object]] = None, provider_config: Optional[BaseOCRConfig] = None, litellm_params: Optional[dict] = None, ) -> OCRResponse: @@ -1703,9 +1711,9 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, asearch: bool = False, - headers: Optional[Dict[str, Any]] = None, + headers: Optional[Dict[str, object]] = None, provider_config: Optional[BaseSearchConfig] = None, - ) -> Union[SearchResponse, Coroutine[Any, Any, SearchResponse]]: + ) -> Union[SearchResponse, Coroutine[object, object, SearchResponse]]: """ Sync Search handler. """ @@ -1798,7 +1806,7 @@ class BaseLLMHTTPHandler: api_base: Optional[str], custom_llm_provider: str, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - headers: Optional[Dict[str, Any]] = None, + headers: Optional[Dict[str, object]] = None, provider_config: Optional[BaseSearchConfig] = None, ) -> SearchResponse: """ @@ -1976,7 +1984,7 @@ class BaseLLMHTTPHandler: litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, client: Optional[AsyncHTTPHandler] = None, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, api_key: Optional[str] = None, api_base: Optional[str] = None, stream: Optional[bool] = False, @@ -2339,10 +2347,10 @@ class BaseLLMHTTPHandler: api_key: Optional[str] = None, api_base: Optional[str] = None, stream: Optional[bool] = False, - kwargs: Optional[Dict[str, Any]] = None, + kwargs: Optional[Dict[str, object]] = None, ) -> Union[ AnthropicMessagesResponse, - Coroutine[Any, Any, Union[AnthropicMessagesResponse, AsyncIterator]], + Coroutine[object, object, Union[AnthropicMessagesResponse, AsyncIterator]], ]: """ LLM HTTP Handler for Anthropic Messages @@ -2449,18 +2457,18 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_body: Optional[Mapping[str, object]] = 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: Optional[Dict[str, object]] = None, shared_session: Optional["ClientSession"] = None, ) -> Union[ ResponsesAPIResponse, BaseResponsesAPIStreamingIterator, - Coroutine[Any, Any, Union[ResponsesAPIResponse, BaseResponsesAPIStreamingIterator]], + Coroutine[object, object, Union[ResponsesAPIResponse, BaseResponsesAPIStreamingIterator]], ]: """ Handles responses API requests. @@ -2543,7 +2551,7 @@ class BaseLLMHTTPHandler: # Preserve the OpenAI-style request context (not sent to the provider) for streaming # hooks/metadata; the streaming iterator now consumes this to run deployment hooks # with the same info as chat, including litellm_params. - request_context: Dict[str, Any] = {"input": input} + request_context: Dict[str, object] = {"input": input} try: request_context.update(response_api_optional_request_params) except Exception: @@ -2663,12 +2671,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_body: Optional[Mapping[str, object]] = 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: Optional[Dict[str, object]] = None, shared_session: Optional["ClientSession"] = None, ) -> Union[ResponsesAPIResponse, BaseResponsesAPIStreamingIterator]: """ @@ -2720,7 +2728,7 @@ class BaseLLMHTTPHandler: # Preserve the OpenAI-style request context (not sent to the provider) for streaming # hooks/metadata; the streaming iterator now consumes this to run deployment hooks # with the same info as chat, including litellm_params. - request_context: Dict[str, Any] = {"input": input} + request_context: Dict[str, object] = {"input": input} try: request_context.update(response_api_optional_request_params) except Exception: @@ -2847,8 +2855,8 @@ class BaseLLMHTTPHandler: 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: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, _is_async: bool = False, @@ -2931,13 +2939,13 @@ class BaseLLMHTTPHandler: 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: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = 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]]: + ) -> Union[DeleteResponseResult, Coroutine[object, object, DeleteResponseResult]]: """ Async version of the responses API handler. Uses async HTTP client to make requests. @@ -3021,13 +3029,13 @@ class BaseLLMHTTPHandler: 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: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = 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]]: + ) -> Union[ResponsesAPIResponse, Coroutine[object, object, ResponsesAPIResponse]]: """ Get a response by ID Uses GET /v1/responses/{response_id} endpoint in the responses API @@ -3102,8 +3110,8 @@ class BaseLLMHTTPHandler: 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: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, shared_session: Optional["ClientSession"] = None, @@ -3183,12 +3191,12 @@ class BaseLLMHTTPHandler: include: Optional[List[str]] = None, limit: int = 20, order: Literal["asc", "desc"] = "desc", - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = 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]]: + ) -> Union[Dict, Coroutine[object, object, Dict]]: if _is_async: return self.async_list_responses_input_items( response_id=response_id, @@ -3269,7 +3277,7 @@ class BaseLLMHTTPHandler: include: Optional[List[str]] = None, limit: int = 20, order: Literal["asc", "desc"] = "desc", - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, shared_session: Optional["ClientSession"] = None, @@ -3375,7 +3383,7 @@ class BaseLLMHTTPHandler: _is_async: bool = False, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, - ) -> Union[OpenAIFileObject, Coroutine[Any, Any, OpenAIFileObject]]: + ) -> Union[OpenAIFileObject, Coroutine[object, object, OpenAIFileObject]]: """ Creates a file using Gemini's two-step upload process """ @@ -3789,7 +3797,7 @@ class BaseLLMHTTPHandler: 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[object, object, "LiteLLMBatch"]]: """ Creates a batch using provider-specific batch creation process """ @@ -3901,7 +3909,7 @@ class BaseLLMHTTPHandler: 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[object, object, "LiteLLMBatch"]]: """ Retrieve a batch using provider-specific configuration. """ @@ -4138,13 +4146,13 @@ class BaseLLMHTTPHandler: 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: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = 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]]: + ) -> Union[ResponsesAPIResponse, Coroutine[object, object, ResponsesAPIResponse]]: """ Async version of the responses API handler. Uses async HTTP client to make requests. @@ -4218,8 +4226,8 @@ class BaseLLMHTTPHandler: 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: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, _is_async: bool = False, @@ -4294,13 +4302,13 @@ class BaseLLMHTTPHandler: 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: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = 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]]: + ) -> Union[ResponsesAPIResponse, Coroutine[object, object, ResponsesAPIResponse]]: """ Handler for the compact responses API. """ @@ -4393,8 +4401,8 @@ class BaseLLMHTTPHandler: 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: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, _is_async: bool = False, @@ -4485,7 +4493,7 @@ class BaseLLMHTTPHandler: _is_async: bool = False, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, - ) -> Union[OpenAIFileObject, Coroutine[Any, Any, OpenAIFileObject]]: + ) -> Union[OpenAIFileObject, Coroutine[object, object, OpenAIFileObject]]: """ Retrieve file metadata by ID """ @@ -4609,7 +4617,7 @@ class BaseLLMHTTPHandler: _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[object, object, "FileDeleted"]]: """ Delete a file by ID """ @@ -4733,7 +4741,7 @@ class BaseLLMHTTPHandler: _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]]]: + ) -> Union[List[OpenAIFileObject], Coroutine[object, object, List[OpenAIFileObject]]]: """ List all files """ @@ -4857,7 +4865,7 @@ class BaseLLMHTTPHandler: _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[object, object, "HttpxBinaryResponseContent"]]: """ Retrieve file content by ID """ @@ -5008,7 +5016,7 @@ class BaseLLMHTTPHandler: return depth, max(max_loops, 1), fingerprints @staticmethod - def _has_agentic_completion_hook(logging_obj: Any) -> bool: + def _has_agentic_completion_hook(logging_obj: LiteLLMLoggingObj) -> bool: """ True if any registered callback actually overrides ``async_should_run_agentic_loop`` (the gate every agentic hook goes @@ -5039,7 +5047,7 @@ class BaseLLMHTTPHandler: @staticmethod def _check_agentic_loop_safety( - tool_calls: Any, + tool_calls: object, fingerprints: List[str], depth: int, max_loops: int, @@ -5062,7 +5070,7 @@ class BaseLLMHTTPHandler: return fingerprint @staticmethod - def _fingerprint_agentic_tools(tools: Dict) -> str: + def _fingerprint_agentic_tools(tools: object) -> str: try: return json.dumps(tools, sort_keys=True, default=str) except Exception: @@ -5081,8 +5089,8 @@ class BaseLLMHTTPHandler: fingerprints: List[str], fingerprint: str, stream: bool = False, - callback: Optional[Any] = None, - ) -> Any: + callback: Optional["CustomLogger"] = None, + ) -> Union[AnthropicMessagesResponse, AsyncIterator[object]]: from litellm.anthropic_interface import messages as anthropic_messages patch = plan.request_patch or AgenticLoopRequestPatch() @@ -5091,7 +5099,7 @@ class BaseLLMHTTPHandler: full_model_name = model if logging_obj is not None: - agentic_params = logging_obj.model_call_details.get("agentic_loop_params", {}) + agentic_params: Mapping[str, object] = logging_obj.model_call_details.get("agentic_loop_params", {}) full_model_name = cast(str, agentic_params.get("model", model)) optional_params = dict(anthropic_messages_optional_request_params) @@ -5121,7 +5129,7 @@ class BaseLLMHTTPHandler: kwargs_for_followup["max_agentic_loops"] = max_loops kwargs_for_followup["_agentic_loop_fingerprints"] = fingerprints + [fingerprint] - response = await anthropic_messages.acreate( + response: Union[AnthropicMessagesResponse, AsyncIterator[object]] = await anthropic_messages.acreate( **{ "max_tokens": max_tokens, "messages": patch.messages, @@ -5160,8 +5168,8 @@ class BaseLLMHTTPHandler: max_loops: int, fingerprints: list[str], fingerprint: str, - callback: Any | None = None, - ) -> Any: + callback: Optional["CustomLogger"] = None, + ) -> Union[ResponsesAPIResponse, BaseResponsesAPIStreamingIterator]: patch = plan.request_patch or AgenticLoopRequestPatch() if patch.messages is None: raise ValueError("Agentic loop plan missing patched responses input") @@ -5192,7 +5200,7 @@ class BaseLLMHTTPHandler: kwargs_for_followup["_agentic_loop_fingerprints"] = fingerprints + [fingerprint] try: - response = await litellm.aresponses( + response: Union[ResponsesAPIResponse, BaseResponsesAPIStreamingIterator] = await litellm.aresponses( model=patch.model or model, input=patch.messages, **optional_params, @@ -5227,7 +5235,7 @@ class BaseLLMHTTPHandler: @staticmethod async def _run_agentic_loop_cleanup( - callback: Any, + callback: "CustomLogger", plan: AgenticLoopPlan, kwargs: dict, logging_obj: "LiteLLMLoggingObj", @@ -5248,10 +5256,10 @@ class BaseLLMHTTPHandler: self, result: Any, model: str, - responses_api_provider_config: Any, + responses_api_provider_config: BaseResponsesAPIConfig, logging_obj: "LiteLLMLoggingObj", custom_llm_provider: str, - ) -> Any: + ) -> MockResponsesAPIStreamingIterator: """ Wrap a completed responses result as a synthetic stream. @@ -5330,10 +5338,10 @@ class BaseLLMHTTPHandler: def _maybe_wrap_in_fake_stream( self, - response: Any, + response: _ResponseT, logging_obj: Optional["LiteLLMLoggingObj"], api_surface: str, - ) -> Any: + ) -> Union[_ResponseT, "FakeAnthropicMessagesStreamIterator"]: """ If the original request was streaming but converted to non-streaming for WebSearch interception, wrap the dict response in a FakeAnthropicMessagesStreamIterator. @@ -5402,7 +5410,7 @@ class BaseLLMHTTPHandler: continue should_run: bool = False - tool_calls: Any = None + tool_calls: object = None try: # First: Check if agentic loop should run. Wrap in try/except # to shield from buggy user callbacks — a callback crash should @@ -5449,7 +5457,7 @@ class BaseLLMHTTPHandler: callback.__class__.async_build_agentic_loop_plan is not CustomLogger.async_build_agentic_loop_plan ) if not build_plan_overridden: - agentic_result = await callback.async_run_agentic_loop( + agentic_result: object = await callback.async_run_agentic_loop( tools=tool_calls, model=model, messages=messages, @@ -5571,7 +5579,7 @@ class BaseLLMHTTPHandler: continue should_run: bool = False - tool_calls: Any = None + tool_calls: object = None try: ( should_run, @@ -5828,7 +5836,7 @@ class BaseLLMHTTPHandler: client: Optional[Any] = None, timeout: Optional[float] = None, user_api_key_dict: Optional[Any] = None, - litellm_metadata: Optional[Dict[str, Any]] = None, + litellm_metadata: Optional[Dict[str, object]] = None, query_params: Optional[RealtimeQueryParams] = None, ): import websockets @@ -5927,7 +5935,7 @@ class BaseLLMHTTPHandler: timeout: Union[float, httpx.Timeout], provider_config: Optional[Any] = None, model: Optional[str] = None, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, api_version: Optional[str] = None, ) -> httpx.Response: @@ -5960,7 +5968,7 @@ class BaseLLMHTTPHandler: timeout: Union[float, httpx.Timeout], provider_config: Optional[Any] = None, model: Optional[str] = None, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, api_version: Optional[str] = None, ) -> httpx.Response: @@ -5989,7 +5997,7 @@ class BaseLLMHTTPHandler: timeout: Union[float, httpx.Timeout], provider_config: Optional[Any] = None, model: Optional[str] = None, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, api_version: Optional[str] = None, ) -> httpx.Response: @@ -6061,8 +6069,8 @@ class BaseLLMHTTPHandler: timeout: Union[float, httpx.Timeout], provider_config: Optional[Any] = None, model: Optional[str] = None, - session_config: Optional[Dict[str, Any]] = None, - extra_headers: Optional[Dict[str, Any]] = None, + session_config: Optional[Dict[str, object]] = None, + extra_headers: Optional[Dict[str, object]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, api_version: Optional[str] = None, ) -> httpx.Response: @@ -6144,7 +6152,7 @@ class BaseLLMHTTPHandler: api_key: Optional[str] = None, timeout: Optional[float] = None, user_api_key_dict: Optional[Any] = None, - litellm_metadata: Optional[Dict[str, Any]] = None, + litellm_metadata: Optional[Dict[str, object]] = None, custom_llm_provider: Optional[str] = None, first_message: Optional[str] = None, **kwargs: Any, @@ -6324,15 +6332,15 @@ class BaseLLMHTTPHandler: 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: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, _is_async: bool = False, fake_stream: bool = False, - litellm_metadata: Optional[Dict[str, Any]] = None, + litellm_metadata: Optional[Dict[str, object]] = None, ) -> Union[ ImageResponse, - Coroutine[Any, Any, ImageResponse], + Coroutine[object, object, ImageResponse], ]: """ @@ -6444,11 +6452,11 @@ class BaseLLMHTTPHandler: 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: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, fake_stream: bool = False, - litellm_metadata: Optional[Dict[str, Any]] = None, + litellm_metadata: Optional[Dict[str, object]] = None, ) -> ImageResponse: """ Async version of the image edit handler. @@ -6542,16 +6550,16 @@ class BaseLLMHTTPHandler: 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: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, _is_async: bool = False, fake_stream: bool = False, - litellm_metadata: Optional[Dict[str, Any]] = None, + litellm_metadata: Optional[Dict[str, object]] = None, api_key: Optional[str] = None, ) -> Union[ ImageResponse, - Coroutine[Any, Any, ImageResponse], + Coroutine[object, object, ImageResponse], ]: """ Handles image generation requests. @@ -6669,11 +6677,11 @@ class BaseLLMHTTPHandler: 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: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, fake_stream: bool = False, - litellm_metadata: Optional[Dict[str, Any]] = None, + litellm_metadata: Optional[Dict[str, object]] = None, api_key: Optional[str] = None, ) -> ImageResponse: """ @@ -6777,16 +6785,16 @@ class BaseLLMHTTPHandler: 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: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, _is_async: bool = False, fake_stream: bool = False, - litellm_metadata: Optional[Dict[str, Any]] = None, + litellm_metadata: Optional[Dict[str, object]] = None, api_key: Optional[str] = None, ) -> Union[ VideoObject, - Coroutine[Any, Any, VideoObject], + Coroutine[object, object, VideoObject], ]: """ Handles video generation requests. @@ -6901,11 +6909,11 @@ class BaseLLMHTTPHandler: 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: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, fake_stream: bool = False, - litellm_metadata: Optional[Dict[str, Any]] = None, + litellm_metadata: Optional[Dict[str, object]] = None, api_key: Optional[str] = None, ) -> VideoObject: """ @@ -7001,12 +7009,12 @@ class BaseLLMHTTPHandler: litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, timeout: Union[float, httpx.Timeout], - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = 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]]: + ) -> Union[bytes, Coroutine[object, object, bytes]]: """ Handle video content download requests. """ @@ -7091,7 +7099,7 @@ class BaseLLMHTTPHandler: litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, timeout: Union[float, httpx.Timeout], - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, api_key: Optional[str] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, variant: Optional[str] = None, @@ -7169,8 +7177,8 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params, logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, timeout: Optional[float] = None, _is_async: bool = False, client=None, @@ -7268,8 +7276,8 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params, logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, timeout: Optional[float] = None, client=None, api_key: Optional[str] = None, @@ -7351,7 +7359,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params, logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, timeout: Optional[float] = None, _is_async: bool = False, client=None, @@ -7435,7 +7443,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params, logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, timeout: Optional[float] = None, client=None, api_key: Optional[str] = None, @@ -7506,7 +7514,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params, logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, timeout: Optional[float] = None, _is_async: bool = False, client=None, @@ -7575,7 +7583,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params, logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, timeout: Optional[float] = None, client=None, api_key: Optional[str] = None, @@ -7634,8 +7642,8 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params, logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, timeout: Optional[float] = None, _is_async: bool = False, client=None, @@ -7743,8 +7751,8 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params, logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, timeout: Optional[float] = None, client=None, api_key: Optional[str] = None, @@ -7840,8 +7848,8 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params, logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, timeout: Optional[float] = None, _is_async: bool = False, client=None, @@ -7929,8 +7937,8 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params, logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, timeout: Optional[float] = None, client=None, api_key: Optional[str] = None, @@ -8004,8 +8012,8 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params, logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_query: Optional[Dict[str, object]] = None, timeout: Optional[float] = None, _is_async: bool = False, client=None, @@ -8058,8 +8066,8 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params, logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_query: Optional[Dict[str, object]] = None, timeout: Optional[float] = None, client=None, api_key: Optional[str] = None, @@ -8139,7 +8147,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params, logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, timeout: Optional[float] = None, client=None, api_key: Optional[str] = None, @@ -8215,8 +8223,8 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params, logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, timeout: Optional[float] = None, _is_async: bool = False, client=None, @@ -8320,8 +8328,8 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params, logging_obj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, timeout: Optional[float] = None, client=None, api_key: Optional[str] = None, @@ -8410,11 +8418,11 @@ class BaseLLMHTTPHandler: container_provider_config: "BaseContainerConfig", litellm_params: GenericLiteLLMParams, logging_obj: "LiteLLMLoggingObj", - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = 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[object, object, "ContainerObject"]]: if _is_async: # Return the async coroutine if called with _is_async=True return self.async_container_create_handler( @@ -8497,7 +8505,7 @@ class BaseLLMHTTPHandler: container_provider_config: "BaseContainerConfig", litellm_params: GenericLiteLLMParams, logging_obj: "LiteLLMLoggingObj", - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, timeout: Union[float, httpx.Timeout] = 600, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, ) -> "ContainerObject": @@ -8574,12 +8582,12 @@ class BaseLLMHTTPHandler: 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: Optional[Dict[str, object]] = None, + extra_query: Optional[Dict[str, object]] = 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[object, object, "ContainerListResponse"]]: if _is_async: # Return the async coroutine if called with _is_async=True return self.async_container_list_handler( @@ -8664,8 +8672,8 @@ class BaseLLMHTTPHandler: 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: Optional[Dict[str, object]] = None, + extra_query: Optional[Dict[str, object]] = None, timeout: Union[float, httpx.Timeout] = 600, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, ) -> "ContainerListResponse": @@ -8739,12 +8747,12 @@ class BaseLLMHTTPHandler: container_provider_config: "BaseContainerConfig", litellm_params: GenericLiteLLMParams, logging_obj: "LiteLLMLoggingObj", - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_query: Optional[Dict[str, object]] = 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[object, object, "ContainerObject"]]: if _is_async: # Return the async coroutine if called with _is_async=True return self.async_container_retrieve_handler( @@ -8827,8 +8835,8 @@ class BaseLLMHTTPHandler: container_provider_config: "BaseContainerConfig", litellm_params: GenericLiteLLMParams, logging_obj: "LiteLLMLoggingObj", - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_query: Optional[Dict[str, object]] = None, timeout: Union[float, httpx.Timeout] = 600, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, ) -> "ContainerObject": @@ -8904,12 +8912,12 @@ class BaseLLMHTTPHandler: container_provider_config: "BaseContainerConfig", litellm_params: GenericLiteLLMParams, logging_obj: "LiteLLMLoggingObj", - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_query: Optional[Dict[str, object]] = 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[object, object, "DeleteContainerResult"]]: if _is_async: # Return the async coroutine if called with _is_async=True return self.async_container_delete_handler( @@ -8992,8 +9000,8 @@ class BaseLLMHTTPHandler: container_provider_config: "BaseContainerConfig", litellm_params: GenericLiteLLMParams, logging_obj: "LiteLLMLoggingObj", - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_query: Optional[Dict[str, object]] = None, timeout: Union[float, httpx.Timeout] = 600, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, ) -> "DeleteContainerResult": @@ -9072,12 +9080,12 @@ class BaseLLMHTTPHandler: 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: Optional[Dict[str, object]] = None, + extra_query: Optional[Dict[str, object]] = 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[object, object, "ContainerFileListResponse"]]: if _is_async: return self.async_container_file_list_handler( container_id=container_id, @@ -9164,8 +9172,8 @@ class BaseLLMHTTPHandler: 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: Optional[Dict[str, object]] = None, + extra_query: Optional[Dict[str, object]] = None, timeout: Union[float, httpx.Timeout] = 600, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, ) -> "ContainerFileListResponse": @@ -9241,11 +9249,11 @@ class BaseLLMHTTPHandler: container_provider_config: "BaseContainerConfig", litellm_params: GenericLiteLLMParams, logging_obj: "LiteLLMLoggingObj", - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, timeout: Union[float, httpx.Timeout] = 600, _is_async: bool = False, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - ) -> Union[bytes, Coroutine[Any, Any, bytes]]: + ) -> Union[bytes, Coroutine[object, object, bytes]]: if _is_async: return self.async_container_file_content_handler( container_id=container_id, @@ -9327,7 +9335,7 @@ class BaseLLMHTTPHandler: container_provider_config: "BaseContainerConfig", litellm_params: GenericLiteLLMParams, logging_obj: "LiteLLMLoggingObj", - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, timeout: Union[float, httpx.Timeout] = 600, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, ) -> bytes: @@ -9406,8 +9414,8 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, _is_async: bool = False, @@ -9504,12 +9512,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, _is_async: bool = False, - ) -> Union[VectorStoreSearchResponse, Coroutine[Any, Any, VectorStoreSearchResponse]]: + ) -> Union[VectorStoreSearchResponse, Coroutine[object, object, VectorStoreSearchResponse]]: if _is_async: return self.async_vector_store_search_handler( vector_store_id=vector_store_id, @@ -9598,8 +9606,8 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, _is_async: bool = False, @@ -9658,12 +9666,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, _is_async: bool = False, - ) -> Union[VectorStoreCreateResponse, Coroutine[Any, Any, VectorStoreCreateResponse]]: + ) -> Union[VectorStoreCreateResponse, Coroutine[object, object, VectorStoreCreateResponse]]: if _is_async: return self.async_vector_store_create_handler( vector_store_create_optional_params=vector_store_create_optional_params, @@ -9728,8 +9736,8 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, ) -> VectorStoreCreateResponse: @@ -9781,12 +9789,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, _is_async: bool = False, - ) -> Union[VectorStoreCreateResponse, Coroutine[Any, Any, VectorStoreCreateResponse]]: + ) -> Union[VectorStoreCreateResponse, Coroutine[object, object, VectorStoreCreateResponse]]: if _is_async: return self.async_vector_store_retrieve_handler( vector_store_id=vector_store_id, @@ -9848,8 +9856,8 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, ): @@ -9912,8 +9920,8 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, _is_async: bool = False, @@ -9988,8 +9996,8 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, ) -> VectorStoreCreateResponse: @@ -10054,12 +10062,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, _is_async: bool = False, - ) -> Union[VectorStoreCreateResponse, Coroutine[Any, Any, VectorStoreCreateResponse]]: + ) -> Union[VectorStoreCreateResponse, Coroutine[object, object, VectorStoreCreateResponse]]: if _is_async: return self.async_vector_store_update_handler( vector_store_id=vector_store_id, @@ -10131,8 +10139,8 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, ): @@ -10182,8 +10190,8 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, _is_async: bool = False, @@ -10249,8 +10257,8 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, str]] = None, + extra_body: Optional[Dict[str, object]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, ) -> VectorStoreFileObject: @@ -10314,12 +10322,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, str]] = None, + extra_body: Optional[Dict[str, object]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, _is_async: bool = False, - ) -> Union[VectorStoreFileObject, Coroutine[Any, Any, VectorStoreFileObject]]: + ) -> Union[VectorStoreFileObject, Coroutine[object, object, VectorStoreFileObject]]: if _is_async: return self.async_vector_store_file_create_handler( vector_store_id=vector_store_id, @@ -10391,8 +10399,8 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, str]] = None, + extra_query: Optional[Dict[str, object]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, ) -> VectorStoreFileListResponse: @@ -10455,12 +10463,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, str]] = None, + extra_query: Optional[Dict[str, object]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, _is_async: bool = False, - ) -> Union[VectorStoreFileListResponse, Coroutine[Any, Any, VectorStoreFileListResponse]]: + ) -> Union[VectorStoreFileListResponse, Coroutine[object, object, VectorStoreFileListResponse]]: if _is_async: return self.async_vector_store_file_list_handler( vector_store_id=vector_store_id, @@ -10531,7 +10539,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, str]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, ) -> VectorStoreFileObject: @@ -10590,11 +10598,11 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, str]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, _is_async: bool = False, - ) -> Union[VectorStoreFileObject, Coroutine[Any, Any, VectorStoreFileObject]]: + ) -> Union[VectorStoreFileObject, Coroutine[object, object, VectorStoreFileObject]]: if _is_async: return self.async_vector_store_file_retrieve_handler( vector_store_id=vector_store_id, @@ -10660,7 +10668,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, str]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, ) -> VectorStoreFileContentResponse: @@ -10721,13 +10729,13 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, str]] = 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[object, object, VectorStoreFileContentResponse], ]: if _is_async: return self.async_vector_store_file_content_handler( @@ -10797,8 +10805,8 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, str]] = None, + extra_body: Optional[Dict[str, object]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, ) -> VectorStoreFileObject: @@ -10863,12 +10871,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, str]] = None, + extra_body: Optional[Dict[str, object]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, _is_async: bool = False, - ) -> Union[VectorStoreFileObject, Coroutine[Any, Any, VectorStoreFileObject]]: + ) -> Union[VectorStoreFileObject, Coroutine[object, object, VectorStoreFileObject]]: if _is_async: return self.async_vector_store_file_update_handler( vector_store_id=vector_store_id, @@ -10941,7 +10949,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, str]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, ) -> VectorStoreFileDeleteResponse: @@ -11000,13 +11008,13 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, str]] = 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[object, object, VectorStoreFileDeleteResponse], ]: if _is_async: return self.async_vector_store_file_delete_handler( @@ -11077,13 +11085,13 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = 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, + litellm_metadata: Optional[Dict[str, object]] = None, system_instruction: Optional[Any] = None, ) -> Any: """ @@ -11209,12 +11217,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[AsyncHTTPHandler] = None, stream: bool = False, - litellm_metadata: Optional[Dict[str, Any]] = None, + litellm_metadata: Optional[Dict[str, object]] = None, system_instruction: Optional[Any] = None, ) -> Any: """ @@ -11328,12 +11336,12 @@ class BaseLLMHTTPHandler: litellm_params: Dict, logging_obj: LiteLLMLoggingObj, timeout: Union[float, httpx.Timeout], - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, _is_async: bool = False, ) -> Union[ "HttpxBinaryResponseContent", - Coroutine[Any, Any, "HttpxBinaryResponseContent"], + Coroutine[object, object, "HttpxBinaryResponseContent"], ]: """ Handles text-to-speech requests. @@ -11443,7 +11451,7 @@ class BaseLLMHTTPHandler: litellm_params: Dict, logging_obj: LiteLLMLoggingObj, timeout: Union[float, httpx.Timeout], - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, ) -> "HttpxBinaryResponseContent": """ @@ -11575,12 +11583,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = 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[object, object, "Skill"]]: """Create a skill""" if _is_async: return self.async_create_skill_handler( @@ -11641,7 +11649,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, shared_session: Optional["ClientSession"] = None, @@ -11697,12 +11705,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = 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[object, object, "ListSkillsResponse"]]: """List skills""" if _is_async: return self.async_list_skills_handler( @@ -11756,7 +11764,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, shared_session: Optional["ClientSession"] = None, @@ -11802,12 +11810,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = 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[object, object, "Skill"]]: """Get a skill""" if _is_async: return self.async_get_skill_handler( @@ -11858,7 +11866,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, shared_session: Optional["ClientSession"] = None, @@ -11903,12 +11911,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = 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[object, object, "DeleteSkillResponse"]]: """Delete a skill""" if _is_async: return self.async_delete_skill_handler( @@ -11959,7 +11967,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, shared_session: Optional["ClientSession"] = None, @@ -12009,12 +12017,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = 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[object, object, "Eval"]]: """Create an eval""" if _is_async: return self.async_create_eval_handler( @@ -12068,7 +12076,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, shared_session: Optional["ClientSession"] = None, @@ -12115,12 +12123,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = 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[object, object, "ListEvalsResponse"]]: """List evals""" if _is_async: return self.async_list_evals_handler( @@ -12174,7 +12182,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, shared_session: Optional["ClientSession"] = None, @@ -12220,12 +12228,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = 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[object, object, "Eval"]]: """Get an eval""" if _is_async: return self.async_get_eval_handler( @@ -12276,7 +12284,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, shared_session: Optional["ClientSession"] = None, @@ -12322,12 +12330,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = 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[object, object, "Eval"]]: """Update an eval""" if _is_async: return self.async_update_eval_handler( @@ -12381,7 +12389,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, shared_session: Optional["ClientSession"] = None, @@ -12427,12 +12435,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = 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[object, object, "DeleteEvalResponse"]]: """Delete an eval""" if _is_async: return self.async_delete_eval_handler( @@ -12483,7 +12491,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, shared_session: Optional["ClientSession"] = None, @@ -12528,12 +12536,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = 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[object, object, "CancelEvalResponse"]]: """Cancel an eval""" if _is_async: return self.async_cancel_eval_handler( @@ -12584,7 +12592,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, shared_session: Optional["ClientSession"] = None, @@ -12634,12 +12642,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = 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[object, object, "Run"]]: """Create a run""" if _is_async: return self.async_create_run_handler( @@ -12693,7 +12701,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, shared_session: Optional["ClientSession"] = None, @@ -12740,12 +12748,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = 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[object, object, "ListRunsResponse"]]: """List runs""" if _is_async: return self.async_list_runs_handler( @@ -12799,7 +12807,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, shared_session: Optional["ClientSession"] = None, @@ -12845,12 +12853,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = 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[object, object, "Run"]]: """Get a run""" if _is_async: return self.async_get_run_handler( @@ -12901,7 +12909,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, shared_session: Optional["ClientSession"] = None, @@ -12946,12 +12954,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = 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[object, object, "CancelRunResponse"]]: """Cancel a run""" if _is_async: return self.async_cancel_run_handler( @@ -13002,7 +13010,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, shared_session: Optional["ClientSession"] = None, @@ -13047,12 +13055,12 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = 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[object, object, "RunDeleteResponse"]]: """Delete a run""" if _is_async: return self.async_delete_run_handler( @@ -13103,7 +13111,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = 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/videos/transformation.py b/litellm/llms/openai/videos/transformation.py index 684601367b6..855bc410cea 100644 --- a/litellm/llms/openai/videos/transformation.py +++ b/litellm/llms/openai/videos/transformation.py @@ -4,7 +4,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast from urllib.parse import quote import httpx -from httpx._types import RequestFiles +from httpx._types import FileContent, FileTypes, RequestFiles import litellm from litellm.litellm_core_utils.url_utils import encode_url_path_segment @@ -128,7 +128,7 @@ class OpenAIVideoConfig(BaseVideoConfig): # Handle input_reference parameter if provided _input_reference = video_create_optional_request_params.get("input_reference") data_without_files = {k: v for k, v in request_dict.items() if k not in ["input_reference"]} - files_list: List[Tuple[str, Any]] = [] + files_list: List[Tuple[str, FileTypes]] = [] # Handle input_reference parameter if _input_reference is not None: @@ -177,9 +177,7 @@ class OpenAIVideoConfig(BaseVideoConfig): request_data: Optional[Dict] = None, ) -> VideoObject: """Transform the OpenAI video creation response.""" - response_data = raw_response.json() - - video_obj = VideoObject(**response_data) # type: ignore[arg-type] + video_obj = VideoObject.model_validate(raw_response.json()) if custom_llm_provider and video_obj.id: video_obj.id = encode_video_id_with_provider(video_obj.id, custom_llm_provider, model) @@ -223,7 +221,7 @@ class OpenAIVideoConfig(BaseVideoConfig): url = f"{url}?variant={quote(variant, safe='')}" # No additional data needed for GET content request - data: Dict[str, Any] = {} + data: Dict[str, object] = {} return url, data @@ -274,10 +272,8 @@ class OpenAIVideoConfig(BaseVideoConfig): """ Transform the OpenAI video remix response. """ - response_data = raw_response.json() - # Transform the response data - video_obj = VideoObject(**response_data) # type: ignore[arg-type] + video_obj = VideoObject.model_validate(raw_response.json()) if custom_llm_provider and video_obj.id: video_obj.id = encode_video_id_with_provider(video_obj.id, custom_llm_provider, None) @@ -392,7 +388,7 @@ class OpenAIVideoConfig(BaseVideoConfig): url = f"{api_base.rstrip('/')}/{encoded_video_id}" # No data needed for DELETE request - data: Dict[str, Any] = {} + data: Dict[str, object] = {} return url, data @@ -404,10 +400,8 @@ class OpenAIVideoConfig(BaseVideoConfig): """ Transform the OpenAI video delete response. """ - response_data = raw_response.json() - # Transform the response data - video_obj = VideoObject(**response_data) # type: ignore[arg-type] # type: ignore[arg-type] + video_obj = VideoObject.model_validate(raw_response.json()) return video_obj @@ -429,7 +423,7 @@ class OpenAIVideoConfig(BaseVideoConfig): url = f"{api_base.rstrip('/')}/{encoded_video_id}" # No additional data needed for GET request - data: Dict[str, Any] = {} + data: Dict[str, object] = {} return url, data @@ -442,9 +436,8 @@ class OpenAIVideoConfig(BaseVideoConfig): """ Transform the OpenAI video retrieve response. """ - response_data = raw_response.json() # Transform the response data - video_obj = VideoObject(**response_data) # type: ignore[arg-type] + video_obj = VideoObject.model_validate(raw_response.json()) if custom_llm_provider and video_obj.id: video_obj.id = encode_video_id_with_provider(video_obj.id, custom_llm_provider, None) @@ -465,22 +458,22 @@ class OpenAIVideoConfig(BaseVideoConfig): def transform_video_create_character_request( self, name: str, - video: Any, + video: FileContent, api_base: str, litellm_params: GenericLiteLLMParams, headers: dict, ) -> Tuple[str, list]: url = f"{api_base.rstrip('/')}/characters" - files_list: List[Tuple[str, Any]] = [("name", (None, name))] + files_list: List[Tuple[str, FileTypes]] = [("name", (None, name))] self._add_video_to_files(files_list, video, "video") return url, files_list def transform_video_create_character_response( self, raw_response: httpx.Response, - logging_obj: Any, + logging_obj: LiteLLMLoggingObj, ) -> CharacterObject: - return CharacterObject(**raw_response.json()) + return CharacterObject.model_validate(raw_response.json()) def transform_video_get_character_request( self, @@ -497,9 +490,9 @@ class OpenAIVideoConfig(BaseVideoConfig): def transform_video_get_character_response( self, raw_response: httpx.Response, - logging_obj: Any, + logging_obj: LiteLLMLoggingObj, ) -> CharacterObject: - return CharacterObject(**raw_response.json()) + return CharacterObject.model_validate(raw_response.json()) def transform_video_edit_request( self, @@ -508,12 +501,12 @@ class OpenAIVideoConfig(BaseVideoConfig): api_base: str, litellm_params: GenericLiteLLMParams, headers: dict, - extra_body: Optional[Dict[str, Any]] = None, - prefetched_source_data: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, object]] = None, + prefetched_source_data: Optional[Dict[str, object]] = None, ) -> Tuple[str, Dict]: original_video_id = extract_original_video_id(video_id) url = f"{api_base.rstrip('/')}/edits" - data: Dict[str, Any] = {"prompt": prompt, "video": {"id": original_video_id}} + data: Dict[str, object] = {"prompt": prompt, "video": {"id": original_video_id}} if extra_body: data.update(extra_body) return url, data @@ -521,11 +514,11 @@ class OpenAIVideoConfig(BaseVideoConfig): def transform_video_edit_response( self, raw_response: httpx.Response, - logging_obj: Any, + logging_obj: LiteLLMLoggingObj, custom_llm_provider: Optional[str] = None, request_data: Optional[Dict] = None, ) -> VideoObject: - video_obj = VideoObject(**raw_response.json()) + video_obj = VideoObject.model_validate(raw_response.json()) if custom_llm_provider and video_obj.id: video_obj.id = encode_video_id_with_provider(video_obj.id, custom_llm_provider, None) return video_obj @@ -538,11 +531,11 @@ class OpenAIVideoConfig(BaseVideoConfig): api_base: str, litellm_params: GenericLiteLLMParams, headers: dict, - extra_body: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, object]] = None, ) -> Tuple[str, Dict]: original_video_id = extract_original_video_id(video_id) url = f"{api_base.rstrip('/')}/extensions" - data: Dict[str, Any] = { + data: Dict[str, object] = { "prompt": prompt, "seconds": seconds, "video": {"id": original_video_id}, @@ -554,10 +547,10 @@ class OpenAIVideoConfig(BaseVideoConfig): def transform_video_extension_response( self, raw_response: httpx.Response, - logging_obj: Any, + logging_obj: LiteLLMLoggingObj, custom_llm_provider: Optional[str] = None, ) -> VideoObject: - video_obj = VideoObject(**raw_response.json()) + video_obj = VideoObject.model_validate(raw_response.json()) if custom_llm_provider and video_obj.id: video_obj.id = encode_video_id_with_provider(video_obj.id, custom_llm_provider, None) return video_obj @@ -578,8 +571,8 @@ class OpenAIVideoConfig(BaseVideoConfig): def _add_video_to_files( self, - files_list: List[Tuple[str, Any]], - video: Any, + files_list: List[Tuple[str, FileTypes]], + video: FileContent, field_name: str, ) -> None: """ @@ -592,7 +585,7 @@ class OpenAIVideoConfig(BaseVideoConfig): content_type = self._get_video_content_type(video=video, filename=filename) files_list.append((field_name, (filename, video, content_type))) - def _get_video_content_type(self, video: Any, filename: str) -> str: + def _get_video_content_type(self, video: FileContent, filename: str) -> str: guessed_content_type, _ = mimetypes.guess_type(filename) if guessed_content_type and guessed_content_type.startswith("video/"): return guessed_content_type diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index e3f7352e8ca..0a6d72881be 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -15,8 +15,19 @@ import re import time from collections.abc import Sequence from contextlib import asynccontextmanager -from typing import Any, AsyncIterator, Callable, Literal, Optional, Union, cast -from urllib.parse import urlparse +from typing import ( + TYPE_CHECKING, + Any, + AsyncIterator, + Callable, + Literal, + Optional, + TypeAlias, + TypedDict, + Union, + cast, +) +from urllib.parse import ParseResult, urlparse import anyio import httpx @@ -32,7 +43,7 @@ from mcp.types import ( ResourceTemplate, ) from mcp.types import Tool as MCPTool -from pydantic import AnyUrl +from pydantic import AnyUrl, BaseModel import litellm from litellm._logging import verbose_logger @@ -139,10 +150,15 @@ from litellm.proxy._types import ( from litellm.proxy.auth.ip_address_utils import IPAddressUtils from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper from litellm.proxy.common_utils.user_api_key_cache import get_management_object_ttl -from litellm.proxy.utils import ProxyLogging, get_server_root_path +from litellm.proxy.utils import PrismaClient, ProxyLogging, get_server_root_path from litellm.repositories.table_repositories import MCPServerRepository from litellm.types.llms.custom_http import httpxSpecialProvider -from litellm.types.mcp import DEFAULT_SUBJECT_TOKEN_TYPE, MCPAuth, MCPStdioConfig +from litellm.types.mcp import ( + DEFAULT_SUBJECT_TOKEN_TYPE, + MCPAuth, + MCPStdioConfig, + MCPTokenEndpointAuthMethod, +) from litellm.types.mcp_server.mcp_server_manager import ( MCPInfo, MCPOAuthMetadata, @@ -150,6 +166,14 @@ from litellm.types.mcp_server.mcp_server_manager import ( ) from litellm.types.utils import CallTypes +if TYPE_CHECKING: + from mcp.client.session import ClientSession + from mcp.shared.context import RequestContext + from mcp.types import CreateMessageRequestParams + + from litellm.caching.caching import InMemoryCache + from litellm.types.mcp_server.mcp_toolset import MCPToolset + try: from mcp.shared.tool_name_validation import ( SEP_986_URL, @@ -209,6 +233,95 @@ _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES: tuple[MCPAuth, ...] = ( _OAUTH_DISCOVERY_RETRY_BASE_SECONDS = 30.0 _OAUTH_DISCOVERY_RETRY_MAX_SECONDS = 900.0 +_StringList: TypeAlias = list[str] +_StringMap: TypeAlias = dict[str, str] +_ToolParamMap: TypeAlias = dict[str, list[str]] +_EnvVarList: TypeAlias = list[dict[str, object]] +_InMemoryCacheDict: TypeAlias = dict[str, object] +_ToolArguments: TypeAlias = dict[str, object] + + +class MCPServerConfig(TypedDict, total=False): + """Shape of a single ``mcp_servers`` entry in config.yaml, as consumed by + :meth:`MCPServerManager.load_servers_from_config`. Every key is optional: YAML supplies + whatever the admin wrote, and each read applies its own default.""" + + alias: str + description: str + mcp_info: MCPInfo + url: str + spec_path: str + transport: MCPTransportType + auth_type: MCPAuthType + authentication_token: str + auth_value: str + instructions: str + command: str + args: _StringList + env: _StringMap + client_id: str + client_secret: str + oauth2_flow: str + issuer: str + authorization_url: str + token_url: str + registration_url: str + token_endpoint_auth_method: MCPTokenEndpointAuthMethod + scopes: str | Sequence[str] + dcr_bridge: object + extra_headers: _StringList + allowed_tools: _StringList + disallowed_tools: _StringList + allowed_params: _ToolParamMap + access_groups: _StringList + static_headers: _StringMap + env_vars: _EnvVarList + allow_all_keys: bool + available_on_public_internet: bool + delegate_auth_to_upstream: bool + oauth_passthrough: bool + allow_sampling: bool + allow_elicitation: bool + aws_access_key_id: str + aws_secret_access_key: str + aws_session_token: str + aws_region_name: str + aws_service_name: str + aws_role_name: str + aws_session_name: str + token_exchange_endpoint: str + token_exchange_profile: str + audience: str + subject_token_type: str + upstream_resource: str + id_jag_resource_token_endpoint: str + id_jag_resource: str + client_private_key: str + client_private_key_id: str + client_assertion_signing_alg: str + timeout: float + max_concurrent_requests: int + + +class _ProtectedResourceMetadataPayload(TypedDict, total=False): + """The RFC 9728 protected-resource metadata document fields this gateway reads.""" + + authorization_servers: Sequence[object] + scopes_supported: Sequence[str] + scopes: Sequence[str] + + +class _AuthorizationServerMetadataPayload(TypedDict, total=False): + """The RFC 8414 / OpenID Discovery authorization-server metadata fields this gateway reads.""" + + issuer: str + authorization_endpoint: str + token_endpoint: str + registration_endpoint: str + scopes_supported: Sequence[str] + grant_types_supported: Sequence[str] + token_endpoint_auth_methods_supported: Sequence[str] + def _blank_to_none(value: str | None) -> str | None: """Collapse an absent, empty, or whitespace-only string to ``None``. @@ -968,7 +1081,7 @@ def _warn_internal_delegate_pkce_if_applicable(server: MCPServer, *, source: str ) -def _deserialize_json_dict(data: Any) -> Optional[dict[str, str]]: +def _deserialize_json_dict(data: str | _StringMap | None) -> Optional[dict[str, str]]: """ Deserialize optional JSON mappings stored in the database. @@ -1057,7 +1170,7 @@ def _normalize_mcp_server_cost_info(mcp_info: MCPInfo) -> None: mcp_info["mcp_server_cost_info"] = normalized -def _create_sampling_callback(user_api_key_auth: Optional[Any] = None): +def _create_sampling_callback(user_api_key_auth: Optional[UserAPIKeyAuth] = None): """ Create a sampling callback for MCP ClientSession. Returns a callable that handles sampling/createMessage requests from @@ -1066,7 +1179,10 @@ def _create_sampling_callback(user_api_key_auth: Optional[Any] = None): if not MCP_SAMPLING_AVAILABLE: return None - async def _sampling_callback(context, params): + async def _sampling_callback( + context: "RequestContext[ClientSession, object]", + params: "CreateMessageRequestParams", + ): import litellm from litellm.proxy._experimental.mcp_server.sampling_handler import ( handle_sampling_create_message, @@ -1309,8 +1425,9 @@ class MCPServerManager: if state is None: return True failures, attempted_at = state + backoff_multiplier: int = 2 ** max(failures - 1, 0) delay = min( - _OAUTH_DISCOVERY_RETRY_BASE_SECONDS * (2 ** max(failures - 1, 0)), + _OAUTH_DISCOVERY_RETRY_BASE_SECONDS * backoff_multiplier, _OAUTH_DISCOVERY_RETRY_MAX_SECONDS, ) return (time.monotonic() - attempted_at) >= delay @@ -1324,7 +1441,7 @@ class MCPServerManager: self._oauth_discovery_retry_state[server.server_id] = (failures + 1, time.monotonic()) def _remember_upstream_initialize_instructions(self, server: MCPServer, client: MCPClient) -> None: - raw = getattr(client, "_last_initialize_instructions", None) + raw: str | None = getattr(client, "_last_initialize_instructions", None) if raw and str(raw).strip(): self._upstream_initialize_instructions_by_server_id[server.server_id] = str(raw).strip() @@ -1430,9 +1547,10 @@ class MCPServerManager: # Track which aliases have been used to ensure only first occurrence is used used_aliases = set() - for server_name, server_config in mcp_servers_config.items(): + for server_name, raw_server_config in mcp_servers_config.items(): + server_config: MCPServerConfig = raw_server_config validate_mcp_server_name(server_name) - _mcp_info: dict[str, Any] = server_config.get("mcp_info", None) or {} + _mcp_info: MCPInfo = server_config.get("mcp_info", None) or {} # Preserve all custom fields from config while setting defaults for core fields mcp_info: MCPInfo = _mcp_info.copy() # Set default values for core fields if not present @@ -1895,7 +2013,7 @@ class MCPServerManager: mcp_server: LiteLLM_MCPServerTable, *, env_vars_are_encrypted: bool, - ) -> Optional[list[dict[str, Any]]]: + ) -> Optional[_EnvVarList]: env_vars_list = _deserialize_json_list(getattr(mcp_server, "env_vars", None)) if env_vars_are_encrypted: from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 @@ -2279,7 +2397,7 @@ class MCPServerManager: async def _get_active_submitted_mcp_server_ids_for_user( self, user_api_key_auth: UserAPIKeyAuth | None ) -> list[str]: - submitter_user_id = getattr(user_api_key_auth, "user_id", None) if user_api_key_auth else None + submitter_user_id: str | None = getattr(user_api_key_auth, "user_id", None) if user_api_key_auth else None if not submitter_user_id: return [] @@ -2551,10 +2669,10 @@ class MCPServerManager: try: from litellm.proxy.proxy_server import user_api_key_cache - in_mem = getattr(user_api_key_cache, "in_memory_cache", None) + in_mem: InMemoryCache | None = getattr(user_api_key_cache, "in_memory_cache", None) if in_mem is None: return - cache_dict = getattr(in_mem, "cache_dict", {}) + cache_dict: _InMemoryCacheDict = getattr(in_mem, "cache_dict", {}) if toolset_id is None: keys_to_remove = [k for k in cache_dict if k.startswith("toolset_")] else: @@ -2574,9 +2692,9 @@ class MCPServerManager: async def get_toolset_by_name_cached( self, - prisma_client: Any, + prisma_client: PrismaClient, toolset_name: str, - ) -> Optional[Any]: + ) -> "Optional[MCPToolset]": """Return a toolset by name, cached in ``user_api_key_cache`` (Redis-backed ``DualCache`` in production) to avoid a DB hit on every routed request. @@ -2803,7 +2921,7 @@ class MCPServerManager: and report ``unknown`` instead of a misleading ``unhealthy``. """ static_headers = server.static_headers - env_vars = getattr(server, "env_vars", None) + env_vars: _EnvVarList | None = getattr(server, "env_vars", None) if not static_headers or not env_vars: return False _global_values, user_specs = parse_admin_env_vars(env_vars) @@ -2929,7 +3047,7 @@ class MCPServerManager: """ if user_api_key_auth is None: return {} - user_id = getattr(user_api_key_auth, "user_id", None) + user_id: str | None = getattr(user_api_key_auth, "user_id", None) if not user_id: return {} @@ -2976,7 +3094,7 @@ class MCPServerManager: match await provider.resolve_credentials(to_subject(user_api_key_auth, subject_token), spec): case Ok(auth): # NoOpAuth has no header_name and so never conflicts. - header_name = getattr(auth, "header_name", None) + header_name: str | None = getattr(auth, "header_name", None) conflicts = bool( header_name and extra_headers and any(key.lower() == header_name.lower() for key in extra_headers) ) @@ -3546,7 +3664,7 @@ class MCPServerManager: self, server: MCPServer, prompt_name: str, - arguments: Optional[dict[str, Any]] = None, + arguments: Optional[dict[str, str]] = None, mcp_auth_header: Optional[Union[str, dict[str, str]]] = None, extra_headers: Optional[dict[str, str]] = None, raw_headers: Optional[dict[str, str]] = None, @@ -3606,7 +3724,7 @@ class MCPServerManager: and base_port == target_port ) - async def _fetch_oauth_discovery_url(self, url: str, server_url: str) -> Any: + async def _fetch_oauth_discovery_url(self, url: str, server_url: str) -> httpx.Response: client = get_async_httpx_client( llm_provider=httpxSpecialProvider.MCP, params={"timeout": MCP_METADATA_TIMEOUT}, @@ -3807,7 +3925,7 @@ class MCPServerManager: try: response = await self._fetch_oauth_discovery_url(resource_metadata_url, server_url) response.raise_for_status() - data = response.json() + data: _ProtectedResourceMetadataPayload = response.json() except SSRFError as exc: verbose_logger.warning( "MCP OAuth discovery: refusing to fetch resource metadata from %s " @@ -3932,7 +4050,7 @@ class MCPServerManager: try: response = await self._fetch_oauth_discovery_url(url, server_url) response.raise_for_status() - data = response.json() + data: _AuthorizationServerMetadataPayload = response.json() except SSRFError as exc: verbose_logger.warning( "MCP OAuth discovery: refusing to fetch authorization-server " @@ -3993,7 +4111,7 @@ class MCPServerManager: @staticmethod def _build_azure_authorization_server_metadata( - parsed_issuer_url: Any, + parsed_issuer_url: ParseResult, ) -> Optional[MCPOAuthMetadata]: path_parts = [part for part in (parsed_issuer_url.path or "").split("/") if part] if parsed_issuer_url.netloc not in _AZURE_ENTRA_HOSTS or len(path_parts) != 2 or path_parts[1] != "v2.0": @@ -4054,7 +4172,7 @@ class MCPServerManager: "aws_session_name": credentials_dict.get("aws_session_name"), } - def _extract_scopes(self, scopes_value: Any) -> Optional[list[str]]: + def _extract_scopes(self, scopes_value: str | Sequence[object] | None) -> Optional[list[str]]: if isinstance(scopes_value, str): scopes = [s.strip() for s in scopes_value.split() if s.strip()] return scopes or None @@ -4292,7 +4410,7 @@ class MCPServerManager: return match_known_tool_name(tool_name, server, server.allowed_tools or ()) is not None return match_known_tool_name(tool_name, server, server.disallowed_tools or ()) is None - def validate_allowed_params(self, tool_name: str, arguments: dict[str, Any], server: MCPServer) -> None: + def validate_allowed_params(self, tool_name: str, arguments: _ToolArguments, server: MCPServer) -> None: """ Filter arguments to only include allowed parameters for the given tool. @@ -4373,7 +4491,7 @@ class MCPServerManager: self, server: MCPServer, tool_name: str, - arguments: dict[str, Any], + arguments: _ToolArguments, ) -> CallToolResult: """ Call an OpenAPI tool handler directly. @@ -4537,7 +4655,7 @@ class MCPServerManager: def _create_during_hook_task( self, name: str, - arguments: dict[str, Any], + arguments: _ToolArguments, server_name_from_prefix: Optional[str], user_api_key_auth: Optional[UserAPIKeyAuth], proxy_logging_obj: ProxyLogging, @@ -4636,7 +4754,7 @@ class MCPServerManager: self, mcp_server: MCPServer, original_tool_name: str, - arguments: dict[str, Any], + arguments: _ToolArguments, tasks: list, mcp_auth_header: Optional[str], mcp_server_auth_headers: Optional[dict[str, dict[str, str]]], @@ -4990,7 +5108,7 @@ class MCPServerManager: # shadow the resolver, double-resolving and hiding the per-server challenge. return oauth2_headers - user_id = getattr(user_api_key_auth, "user_id", None) + user_id: str | None = getattr(user_api_key_auth, "user_id", None) if not user_id: return oauth2_headers @@ -5091,7 +5209,7 @@ class MCPServerManager: self, server_name: str, name: str, - arguments: dict[str, Any], + arguments: _ToolArguments, user_api_key_auth: Optional[UserAPIKeyAuth] = None, mcp_auth_header: Optional[str] = None, mcp_server_auth_headers: Optional[dict[str, dict[str, str]]] = None, @@ -5330,7 +5448,7 @@ class MCPServerManager: # Pending/rejected servers are excluded at the DB level so we never load them. from litellm.proxy._experimental.mcp_server.db import LiteLLM_MCPServerTable - raw_rows = await MCPServerRepository(prisma_client).table.find_many( + raw_rows: Sequence[BaseModel] = await MCPServerRepository(prisma_client).table.find_many( where={ "OR": [ {"approval_status": None}, @@ -5836,7 +5954,7 @@ class MCPServerManager: @staticmethod def _env_vars_to_models( - env_vars: Optional[list[dict[str, Any]]], + env_vars: Optional[_EnvVarList], ) -> Optional[list[MCPEnvVar]]: if env_vars is None: return None diff --git a/litellm/proxy/_experimental/mcp_server/sampling_handler.py b/litellm/proxy/_experimental/mcp_server/sampling_handler.py index 65630f74e90..5bad530f37b 100644 --- a/litellm/proxy/_experimental/mcp_server/sampling_handler.py +++ b/litellm/proxy/_experimental/mcp_server/sampling_handler.py @@ -10,16 +10,23 @@ MCP Spec Reference: https://modelcontextprotocol.io/specification/2025-11-25/client/sampling """ -from typing import Any, Dict, List, Optional, Union import typing +from collections.abc import Mapping, Sequence +from typing import Any, Dict, List, NamedTuple, Optional, Protocol, Union if typing.TYPE_CHECKING: + from fastapi import Request + from mcp.client.session import ClientSession + from mcp.shared.context import RequestContext + from mcp.types import ContentBlock, SamplingMessageContentBlock + + from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.utils import ProxyLogging -from litellm._logging import verbose_logger - from fastapi import HTTPException +from litellm._logging import verbose_logger + # Guard imports that require the mcp package try: from mcp.types import ( @@ -65,7 +72,7 @@ def _resolve_model_from_preferences( import litellm # Build list of available model names from proxy Router or litellm.model_list - available_model_names: list = [] + available_model_names: list[str] = [] try: from litellm.proxy.proxy_server import llm_router @@ -83,7 +90,7 @@ def _resolve_model_from_preferences( available_model_names.append(entry) if model_preferences and model_preferences.hints: for hint in model_preferences.hints: - hint_name = getattr(hint, "name", None) + hint_name: str | None = getattr(hint, "name", None) if not hint_name: continue # Try direct match first @@ -133,7 +140,7 @@ def _resolve_model_from_preferences( ) return available_model_names[0] # Last resort - use LiteLLM default or raise error - default_sampling_model = getattr(litellm, "default_mcp_sampling_model", None) + default_sampling_model: str | None = getattr(litellm, "default_mcp_sampling_model", None) if default_sampling_model: verbose_logger.debug( "MCP sampling model resolution: using litellm.default_mcp_sampling_model='%s'", @@ -153,6 +160,13 @@ def _has_priorities(model_preferences: "ModelPreferences") -> bool: ) +class _ScoredModel(NamedTuple): + name: str + cost: float + max_output: float + output_tps: float + + def _select_model_by_priority( model_names: List[str], model_preferences: "ModelPreferences", @@ -183,12 +197,12 @@ def _select_model_by_priority( """ import litellm as _litellm - cost_weight = getattr(model_preferences, "costPriority", None) or 0.0 - speed_weight = getattr(model_preferences, "speedPriority", None) or 0.0 - intel_weight = getattr(model_preferences, "intelligencePriority", None) or 0.0 + cost_weight: float = getattr(model_preferences, "costPriority", None) or 0.0 + speed_weight: float = getattr(model_preferences, "speedPriority", None) or 0.0 + intel_weight: float = getattr(model_preferences, "intelligencePriority", None) or 0.0 # Gather raw metrics for each model - scored: List[Dict[str, Any]] = [] + scored: list[_ScoredModel] = [] for name in model_names: try: info = _litellm.get_model_info(name) @@ -200,12 +214,12 @@ def _select_model_by_priority( max_output = info.get("max_output_tokens") or info.get("max_tokens") or 0 output_tps = info.get("output_tokens_per_second") or 0.0 scored.append( - { - "name": name, - "cost": total_cost, - "max_output": max_output, - "output_tps": output_tps, - } + _ScoredModel( + name=name, + cost=total_cost, + max_output=max_output, + output_tps=output_tps, + ) ) if not scored: @@ -222,9 +236,9 @@ def _select_model_by_priority( normed = [1.0 - n for n in normed] return normed - costs = [s["cost"] for s in scored] - max_outputs = [float(s["max_output"]) for s in scored] - output_tps_values = [s["output_tps"] for s in scored] + costs = [s.cost for s in scored] + max_outputs = [float(s.max_output) for s in scored] + output_tps_values = [s.output_tps for s in scored] # costPriority: lower cost → higher score (invert) cost_scores = _normalise(costs, invert=True) @@ -243,7 +257,7 @@ def _select_model_by_priority( score = cost_weight * cost_scores[i] + speed_weight * speed_scores[i] + intel_weight * intel_scores[i] verbose_logger.debug( "MCP priority scoring: model=%s cost_score=%.3f speed_score=%.3f intel_score=%.3f → weighted=%.3f", - entry["name"], + entry.name, cost_scores[i], speed_scores[i], intel_scores[i], @@ -251,14 +265,14 @@ def _select_model_by_priority( ) if score > best_score: best_score = score - best_name = entry["name"] + best_name = entry.name return best_name def _convert_mcp_content_to_openai( - content: Any, -) -> Union[str, Dict[str, Any], List[Dict[str, Any]]]: + content: "SamplingMessageContentBlock | Sequence[SamplingMessageContentBlock]", +) -> "str | dict[str, object] | list[dict[str, object]]": """ Convert MCP SamplingMessage content to OpenAI message content format. Handles: @@ -283,7 +297,7 @@ def _convert_mcp_content_to_openai( def _convert_single_content( content: Any, -) -> Union[Dict[str, Any], List[Dict[str, Any]]]: +) -> "dict[str, object] | list[dict[str, object]]": """Convert a single MCP content item to OpenAI format. For text/image/audio content, returns a single content-part dict. @@ -339,7 +353,7 @@ def _convert_single_content( # Marked so the message-level converter can emit it as a # separate ``{"role": "tool", ...}`` message. tool_use_id = getattr(content, "toolUseId", "") - nested_content = getattr(content, "content", []) + nested_content: Sequence[ContentBlock] = getattr(content, "content", []) if isinstance(nested_content, list): text_parts = [getattr(c, "text", str(c)) for c in nested_content if getattr(c, "type", None) == "text"] result_text = "\n".join(text_parts) if text_parts else "" @@ -358,7 +372,7 @@ def _convert_single_content( def _convert_mcp_messages_to_openai( messages: List["SamplingMessage"], system_prompt: Optional[str] = None, -) -> List[Dict[str, Any]]: +) -> "Sequence[Mapping[str, object]]": """ Convert MCP SamplingMessage list to OpenAI messages format. MCP messages use: @@ -369,7 +383,7 @@ def _convert_mcp_messages_to_openai( - role: "system" | "user" | "assistant" | "tool" - content: str | list[content_part] """ - openai_messages: List[Dict[str, Any]] = [] + openai_messages: list[Mapping[str, object]] = [] # Add system prompt if provided if system_prompt: openai_messages.append({"role": "system", "content": system_prompt}) @@ -380,7 +394,7 @@ def _convert_mcp_messages_to_openai( if role == "assistant" and _has_tool_use(content): tool_calls = _extract_tool_calls(content) if tool_calls: - openai_msg: Dict[str, Any] = { + openai_msg: dict[str, object] = { "role": "assistant", "tool_calls": tool_calls, } @@ -400,7 +414,7 @@ def _convert_mcp_messages_to_openai( # tool_use / tool_result that slipped past the fast-path checks # above (e.g. unexpected role, single non-list content). converted = _convert_mcp_content_to_openai(content) - converted_parts = ( + converted_parts: Sequence[Mapping[str, object]] = ( converted if isinstance(converted, list) else ([converted] if isinstance(converted, dict) else []) ) @@ -422,7 +436,7 @@ def _convert_mcp_messages_to_openai( # Emit assistant message with tool_calls if any were found if tool_call_markers: - openai_msg_tc: Dict[str, Any] = { + openai_msg_tc: dict[str, object] = { "role": "assistant", "tool_calls": tool_call_markers, } @@ -442,21 +456,25 @@ def _convert_mcp_messages_to_openai( return openai_messages -def _has_tool_use(content: Any) -> bool: +def _has_tool_use(content: "SamplingMessageContentBlock | Sequence[SamplingMessageContentBlock]") -> bool: """Check if content contains ToolUseContent.""" if isinstance(content, list): return any(getattr(c, "type", None) == "tool_use" for c in content) - return getattr(content, "type", None) == "tool_use" + content_type: str | None = getattr(content, "type", None) + return content_type == "tool_use" -def _has_tool_result(content: Any) -> bool: +def _has_tool_result(content: "SamplingMessageContentBlock | Sequence[SamplingMessageContentBlock]") -> bool: """Check if content contains ToolResultContent.""" if isinstance(content, list): return any(getattr(c, "type", None) == "tool_result" for c in content) - return getattr(content, "type", None) == "tool_result" + content_type: str | None = getattr(content, "type", None) + return content_type == "tool_result" -def _extract_tool_calls(content: Any) -> List[Dict[str, Any]]: +def _extract_tool_calls( + content: "SamplingMessageContentBlock | Sequence[SamplingMessageContentBlock]", +) -> "Sequence[Mapping[str, object]]": """Extract OpenAI-format tool_calls from MCP ToolUseContent.""" import json @@ -477,7 +495,9 @@ def _extract_tool_calls(content: Any) -> List[Dict[str, Any]]: return tool_calls -def _extract_text_parts(content: Any) -> Optional[str]: +def _extract_text_parts( + content: "SamplingMessageContentBlock | Sequence[SamplingMessageContentBlock]", +) -> Optional[str]: """Extract text parts from mixed content.""" items = content if isinstance(content, list) else [content] texts = [] @@ -487,7 +507,9 @@ def _extract_text_parts(content: Any) -> Optional[str]: return "\n".join(texts) if texts else None -def _extract_tool_results(content: Any) -> List[Dict[str, Any]]: +def _extract_tool_results( + content: "SamplingMessageContentBlock | Sequence[SamplingMessageContentBlock]", +) -> "Sequence[Mapping[str, object]]": """Extract OpenAI-format tool messages from MCP ToolResultContent.""" items = content if isinstance(content, list) else [content] results = [] @@ -495,7 +517,7 @@ def _extract_tool_results(content: Any) -> List[Dict[str, Any]]: if getattr(item, "type", None) == "tool_result": tool_use_id = getattr(item, "toolUseId", "") # Extract text from nested content - nested_content = getattr(item, "content", []) + nested_content: Sequence[ContentBlock] = getattr(item, "content", []) if isinstance(nested_content, list): text_parts = [getattr(c, "text", str(c)) for c in nested_content if getattr(c, "type", None) == "text"] result_text = "\n".join(text_parts) if text_parts else "" @@ -513,7 +535,7 @@ def _extract_tool_results(content: Any) -> List[Dict[str, Any]]: def _convert_mcp_tools_to_openai( tools: Optional[List["Tool"]], -) -> Optional[List[Dict[str, Any]]]: +) -> "Sequence[Mapping[str, object]] | None": """ Convert MCP Tool definitions to OpenAI function calling format. MCP Tool: {name, description, inputSchema} @@ -541,7 +563,7 @@ def _convert_mcp_tools_to_openai( def _convert_mcp_tool_choice_to_openai( tool_choice: Optional["ToolChoice"], -) -> Optional[Union[str, Dict[str, Any]]]: +) -> "str | None": """ Convert MCP ToolChoice to OpenAI tool_choice format. MCP: {mode: "auto"} | {mode: "required"} | {mode: "none"} @@ -559,8 +581,32 @@ def _convert_mcp_tool_choice_to_openai( return "auto" +class _SamplingResponseMessage(Protocol): + @property + def content(self) -> str | None: ... + + @property + def tool_calls(self) -> Sequence[object] | None: ... + + +class _SamplingResponseChoice(Protocol): + @property + def message(self) -> _SamplingResponseMessage: ... + + @property + def finish_reason(self) -> str | None: ... + + +class _SamplingCompletionResponse(Protocol): + @property + def choices(self) -> Sequence[_SamplingResponseChoice]: ... + + @property + def model(self) -> str | None: ... + + def _convert_openai_response_to_mcp_result( - response: Any, + response: _SamplingCompletionResponse, model_name: str, ) -> Union["CreateMessageResult", "CreateMessageResultWithTools", "ErrorData"]: """ @@ -593,12 +639,12 @@ def _convert_openai_response_to_mcp_result( stop_reason = "maxTokens" else: stop_reason = "endTurn" - actual_model = getattr(response, "model", model_name) or model_name + actual_model: str = getattr(response, "model", model_name) or model_name # Check if response has tool calls tool_calls = getattr(message, "tool_calls", None) if tool_calls: # Build ToolUseContent items - content_parts: "List[Any]" = [] + content_parts: list[SamplingMessageContentBlock] = [] # Include text content if present if message.content: content_parts.append(TextContent(type="text", text=message.content)) @@ -636,7 +682,7 @@ def _convert_openai_response_to_mcp_result( ) -async def _check_model_access(model: str, user_api_key_auth: Any) -> Optional["ErrorData"]: +async def _check_model_access(model: str, user_api_key_auth: "UserAPIKeyAuth | None") -> Optional["ErrorData"]: """Enforce model-permission checks for MCP sampling requests. Runs the same authorization checks as ``/chat/completions``: @@ -678,14 +724,14 @@ async def _check_model_access(model: str, user_api_key_auth: Any) -> Optional["E try: import litellm from litellm.proxy.auth.auth_checks import ( + _check_team_member_model_access, can_key_call_model, + can_project_access_model, can_team_access_model, can_user_call_model, - can_project_access_model, - _check_team_member_model_access, + get_project_object, get_team_object, get_user_object, - get_project_object, ) try: @@ -700,16 +746,20 @@ async def _check_model_access(model: str, user_api_key_auth: Any) -> Optional["E llm_router=_llm_router, ) - _team_id = getattr(user_api_key_auth, "team_id", None) - _user_id = getattr(user_api_key_auth, "user_id", None) - _project_id = getattr(user_api_key_auth, "project_id", None) + _team_id: str | None = getattr(user_api_key_auth, "team_id", None) + _user_id: str | None = getattr(user_api_key_auth, "user_id", None) + _project_id: str | None = getattr(user_api_key_auth, "project_id", None) try: from litellm.proxy.proxy_server import ( prisma_client as _prisma_client, - user_api_key_cache as _user_api_key_cache, + ) + from litellm.proxy.proxy_server import ( proxy_logging_obj as _proxy_logging_obj, ) + from litellm.proxy.proxy_server import ( + user_api_key_cache as _user_api_key_cache, + ) except ImportError: _prisma_client = None _user_api_key_cache = None # type: ignore[assignment] @@ -799,7 +849,7 @@ async def _check_model_access(model: str, user_api_key_auth: Any) -> Optional["E async def _run_budget_checks( model: str, - user_api_key_auth: Any, + user_api_key_auth: "UserAPIKeyAuth", raw_headers: Optional[Dict[str, str]] = None, client_ip: Optional[str] = None, ) -> Optional["ErrorData"]: @@ -811,25 +861,33 @@ async def _run_budget_checks( Returns None if all checks pass, or an ErrorData describing the denial. """ try: - from litellm.proxy.auth.auth_checks import common_checks - from litellm.proxy.proxy_server import ( - general_settings, - llm_router as _llm_router, - prisma_client as _prisma_client, - proxy_logging_obj as _proxy_logging_obj, - user_api_key_cache as _user_api_key_cache, - ) + import litellm from litellm.proxy.auth.auth_checks import ( + common_checks, get_team_object, get_user_object, ) - import litellm + from litellm.proxy.proxy_server import ( + general_settings, + ) + from litellm.proxy.proxy_server import ( + llm_router as _llm_router, + ) + from litellm.proxy.proxy_server import ( + prisma_client as _prisma_client, + ) + from litellm.proxy.proxy_server import ( + proxy_logging_obj as _proxy_logging_obj, + ) + from litellm.proxy.proxy_server import ( + user_api_key_cache as _user_api_key_cache, + ) except ImportError as import_err: verbose_logger.warning("MCP sampling: budget check imports unavailable: %s", import_err) return None # Can't enforce budgets without the modules - _team_id = getattr(user_api_key_auth, "team_id", None) - _user_id = getattr(user_api_key_auth, "user_id", None) + _team_id: str | None = getattr(user_api_key_auth, "team_id", None) + _user_id: str | None = getattr(user_api_key_auth, "user_id", None) team_obj = None if _team_id and _prisma_client and _user_api_key_cache: @@ -889,7 +947,7 @@ async def _run_budget_checks( # common_checks runs. _tag_max_budget_check inside common_checks only # inspects request_body; without this pre-merge, header-supplied tags # bypass per-tag budget enforcement (mirroring the regular auth path). - request_body: Dict[str, Any] = {"model": model} + request_body: dict[str, object] = {"model": model} try: from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup @@ -935,7 +993,7 @@ async def _run_budget_checks( def _build_sampling_request( raw_headers: Optional[Dict[str, str]] = None, client_ip: Optional[str] = None, -) -> Any: +) -> "Request": """Build a synthetic FastAPI Request for sampling sub-calls. Converts the original MCP connection's HTTP headers into ASGI @@ -961,7 +1019,7 @@ def _build_sampling_request( from fastapi import Request # --- Build ASGI headers --- - _scope_headers: list = [(b"content-type", b"application/json")] + _scope_headers: list[tuple[bytes, bytes]] = [(b"content-type", b"application/json")] # Hop-by-hop headers that must NOT be forwarded into the # synthetic request (they describe the original HTTP framing, # not the logical request). @@ -1001,8 +1059,8 @@ def _build_sampling_request( try: import litellm.proxy.proxy_server as proxy_server - _proxy_host = getattr(proxy_server, "server_host", None) - _proxy_port = getattr(proxy_server, "server_port", None) + _proxy_host: str | None = getattr(proxy_server, "server_host", None) + _proxy_port: str | int | None = getattr(proxy_server, "server_port", None) if _proxy_host: _server_host = str(_proxy_host) @@ -1016,7 +1074,7 @@ def _build_sampling_request( if client_ip: _client_tuple = (client_ip, 0) - scope: Dict[str, Any] = { + scope: dict[str, object] = { "type": "http", "method": "POST", "path": "/mcp/sampling/createMessage", @@ -1035,7 +1093,7 @@ def _build_sampling_request( async def _build_completion_kwargs( params: "CreateMessageRequestParams", model: str, - user_api_key_auth: Any, + user_api_key_auth: "UserAPIKeyAuth", raw_headers: Optional[Dict[str, str]], client_ip: Optional[str], ) -> Dict[str, Any]: @@ -1078,7 +1136,7 @@ async def _build_completion_kwargs( async def _run_guardrails_and_call_llm( completion_kwargs: Dict[str, Any], - user_api_key_auth: Any, + user_api_key_auth: "UserAPIKeyAuth", ) -> Any: try: from litellm.proxy.proxy_server import proxy_logging_obj as _plo @@ -1111,10 +1169,10 @@ async def _run_guardrails_and_call_llm( async def handle_sampling_create_message( - context: Any, + context: "RequestContext[ClientSession, object]", params: "CreateMessageRequestParams", default_model: Optional[str] = None, - user_api_key_auth: Optional[Any] = None, + user_api_key_auth: "UserAPIKeyAuth | None" = None, raw_headers: Optional[Dict[str, str]] = None, client_ip: Optional[str] = None, ) -> Union["CreateMessageResult", "CreateMessageResultWithTools", "ErrorData"]: @@ -1184,7 +1242,7 @@ async def handle_sampling_create_message( client_ip=client_ip, ) - openai_messages = completion_kwargs["messages"] + openai_messages: Sequence[Mapping[str, object]] = completion_kwargs["messages"] openai_tools = completion_kwargs.get("tools") verbose_logger.debug( "MCP sampling: calling litellm.acompletion with model=%s, num_messages=%d, has_tools=%s", @@ -1193,7 +1251,7 @@ async def handle_sampling_create_message( bool(openai_tools), ) - response = await _run_guardrails_and_call_llm( + response: _SamplingCompletionResponse = await _run_guardrails_and_call_llm( completion_kwargs=completion_kwargs, user_api_key_auth=user_api_key_auth, ) @@ -1214,7 +1272,6 @@ async def handle_sampling_create_message( RateLimitError, ServiceUnavailableError, ) - from litellm.proxy._types import ProxyException if isinstance( diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index da7a72e1cff..a29899a3965 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -6,8 +6,20 @@ import concurrent.futures import inspect import json import os +from collections.abc import Mapping, Sequence from datetime import datetime, timezone -from typing import Any, Dict, List, Literal, Optional, Type, TypeVar, Union, cast +from typing import ( + TYPE_CHECKING, + Any, + Dict, + List, + Literal, + Optional, + Type, + TypeVar, + Union, + cast, +) from urllib.parse import urlparse from fastapi import APIRouter, Depends, HTTPException, Request @@ -49,12 +61,44 @@ from litellm.types.guardrails import ( ToolPermissionGuardrailConfigModel, ) +if TYPE_CHECKING: + from types import CodeType + + from prisma.actions import LiteLLM_GuardrailsTableActions + from prisma.models import LiteLLM_GuardrailsTable + + from litellm.proxy.utils import PrismaClient + #### GUARDRAILS ENDPOINTS #### router = APIRouter() GUARDRAIL_REGISTRY = GuardrailRegistry() +def _guardrails_table(prisma_client: "PrismaClient") -> "LiteLLM_GuardrailsTableActions[LiteLLM_GuardrailsTable]": + table: LiteLLM_GuardrailsTableActions[LiteLLM_GuardrailsTable] = GuardrailsRepository(prisma_client).table + return table + + +async def _create_guardrail_row(prisma_client: "PrismaClient", data: Mapping[str, object]) -> "LiteLLM_GuardrailsTable": + row: LiteLLM_GuardrailsTable = await GuardrailsRepository(prisma_client).table.create(data=data) + return row + + +async def _delete_guardrail_row(prisma_client: "PrismaClient", where: Mapping[str, object]) -> None: + await GuardrailsRepository(prisma_client).table.delete(where=where) + + +async def _find_team_guardrail_rows( + prisma_client: "PrismaClient", where: Mapping[str, object] +) -> "Sequence[LiteLLM_GuardrailsTable]": + rows: Sequence[LiteLLM_GuardrailsTable] = await GuardrailsRepository(prisma_client).table.find_many( + where=where, + order={"created_at": "desc"}, + ) + return rows + + def _get_guardrails_list_response( guardrails_config: List[Dict], ) -> ListGuardrailsResponse: @@ -363,7 +407,7 @@ async def create_guardrail( # Configuration error — roll back the DB write so the guardrail isn't orphaned if prisma_client is not None: try: - await GuardrailsRepository(prisma_client).table.delete(where={"guardrail_id": guardrail_id}) + await _delete_guardrail_row(prisma_client, where={"guardrail_id": guardrail_id}) except Exception as rollback_err: verbose_proxy_logger.warning(f"Rollback failed for guardrail '{guardrail_id}': {rollback_err}") raise HTTPException( @@ -571,7 +615,7 @@ class RegisterGuardrailRequest(BaseModel): guardrail_name: str litellm_params: Dict[str, Any] # guardrail, mode, api_base required; api_key, headers, etc. optional - guardrail_info: Optional[Dict[str, Any]] = None + guardrail_info: Optional[Dict[str, object]] = None team_id: Optional[str] = None def get_litellm_params_dict(self) -> Dict[str, Any]: @@ -600,8 +644,8 @@ class GuardrailSubmissionItem(BaseModel): team_guardrail: bool = ( False # True when submitted via team (team_id set); use to distinguish team vs regular guardrails ) - litellm_params: Optional[Dict[str, Any]] = None - guardrail_info: Optional[Dict[str, Any]] = None + litellm_params: Optional[Dict[str, object]] = None + guardrail_info: Optional[Dict[str, object]] = None submitted_by_user_id: Optional[str] = None submitted_by_email: Optional[str] = None submitted_at: Optional[datetime] = None @@ -685,9 +729,7 @@ async def register_guardrail( ) try: - existing = await GuardrailsRepository(prisma_client).table.find_unique( - where={"guardrail_name": request.guardrail_name} - ) + existing = await _guardrails_table(prisma_client).find_unique(where={"guardrail_name": request.guardrail_name}) if existing is not None: raise HTTPException( status_code=400, @@ -708,7 +750,8 @@ async def register_guardrail( guardrail_info_str = safe_dumps(guardrail_info) try: - created = await GuardrailsRepository(prisma_client).table.create( + created = await _create_guardrail_row( + prisma_client, data={ "guardrail_name": request.guardrail_name, "litellm_params": litellm_params_str, @@ -718,7 +761,7 @@ async def register_guardrail( "submitted_at": now, "created_at": now, "updated_at": now, - } + }, ) return RegisterGuardrailResponse( guardrail_id=created.guardrail_id, @@ -731,7 +774,7 @@ async def register_guardrail( raise HTTPException(status_code=500, detail=str(e)) -def _parse_json_field(value: Any) -> Optional[Dict[str, Any]]: +def _parse_json_field(value: object) -> Optional[Dict[str, Any]]: if value is None: return None if isinstance(value, dict): @@ -768,7 +811,7 @@ async def _get_user_team_ids(user_api_key_dict: UserAPIKeyAuth) -> List[str]: return [t for t in user_obj.teams if t] -def _row_to_submission_item(row: Any) -> GuardrailSubmissionItem: +def _row_to_submission_item(row: "LiteLLM_GuardrailsTable") -> GuardrailSubmissionItem: from litellm.litellm_core_utils.litellm_logging import _get_masked_values guardrail_info = _parse_json_field(row.guardrail_info) or {} @@ -835,7 +878,7 @@ async def list_guardrail_submissions( ) try: - where_clause: Dict[str, Any] = {"team_id": {"not": None}} + where_clause: Dict[str, object] = {"team_id": {"not": None}} if visible_team_ids is not None: if not visible_team_ids: # Non-admin with no team memberships: nothing visible. @@ -846,10 +889,7 @@ async def list_guardrail_submissions( where_clause["team_id"] = {"in": visible_team_ids} # Single query: fetch team guardrails visible to the caller - all_team_rows = await GuardrailsRepository(prisma_client).table.find_many( - where=where_clause, - order={"created_at": "desc"}, - ) + all_team_rows = await _find_team_guardrail_rows(prisma_client, where_clause) # Derive summary counts from the full result set total = len(all_team_rows) @@ -909,7 +949,7 @@ async def get_guardrail_submission( is_admin = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN try: - row = await GuardrailsRepository(prisma_client).table.find_unique(where={"guardrail_id": guardrail_id}) + row = await _guardrails_table(prisma_client).find_unique(where={"guardrail_id": guardrail_id}) if row is None: raise HTTPException(status_code=404, detail="Guardrail submission not found") if not is_admin: @@ -946,7 +986,7 @@ async def approve_guardrail_submission( raise HTTPException(status_code=500, detail="Prisma client not initialized") try: - row = await GuardrailsRepository(prisma_client).table.find_unique(where={"guardrail_id": guardrail_id}) + row = await _guardrails_table(prisma_client).find_unique(where={"guardrail_id": guardrail_id}) if row is None: raise HTTPException(status_code=404, detail="Guardrail submission not found") if row.status != "pending_review": @@ -956,7 +996,7 @@ async def approve_guardrail_submission( ) now = datetime.now(timezone.utc) - await GuardrailsRepository(prisma_client).table.update( + await _guardrails_table(prisma_client).update( where={"guardrail_id": guardrail_id}, data={"status": "active", "reviewed_at": now, "updated_at": now}, ) @@ -1026,7 +1066,7 @@ async def reject_guardrail_submission( raise HTTPException(status_code=500, detail="Prisma client not initialized") try: - row = await GuardrailsRepository(prisma_client).table.find_unique(where={"guardrail_id": guardrail_id}) + row = await _guardrails_table(prisma_client).find_unique(where={"guardrail_id": guardrail_id}) if row is None: raise HTTPException(status_code=404, detail="Guardrail submission not found") if row.status != "pending_review": @@ -1036,7 +1076,7 @@ async def reject_guardrail_submission( ) now = datetime.now(timezone.utc) - await GuardrailsRepository(prisma_client).table.update( + await _guardrails_table(prisma_client).update( where={"guardrail_id": guardrail_id}, data={"status": "rejected", "reviewed_at": now, "updated_at": now}, ) @@ -1886,13 +1926,13 @@ class TestCustomCodeGuardrailRequest(BaseModel): custom_code: str """The Python-like code containing the apply_guardrail function.""" - test_input: Dict[str, Any] + test_input: Dict[str, object] """The test input to pass to the guardrail. Should contain 'texts', optionally 'images', 'tools', etc.""" input_type: str = "request" """Whether this is a 'request' or 'response' input type.""" - request_data: Optional[Dict[str, Any]] = None + request_data: Optional[Dict[str, object]] = None """Optional mock request_data (model, user_id, team_id, metadata, etc.).""" @@ -1902,7 +1942,7 @@ class TestCustomCodeGuardrailResponse(BaseModel): success: bool """Whether the test executed successfully (no errors).""" - result: Optional[Dict[str, Any]] = None + result: Optional[Dict[str, object]] = None """The guardrail result: action (allow/block/modify), reason, modified_texts, etc.""" error: Optional[str] = None @@ -2006,7 +2046,7 @@ async def test_custom_code_guardrail( exec_globals = build_sandbox_globals() try: - compiled = compile_sandboxed(request.custom_code) + compiled: CodeType = compile_sandboxed(request.custom_code) exec(compiled, exec_globals) # noqa: S102 except SyntaxError as e: return TestCustomCodeGuardrailResponse( @@ -2030,7 +2070,7 @@ async def test_custom_code_guardrail( error_type="compilation", ) - apply_fn = exec_globals["apply_guardrail"] + apply_fn: object = exec_globals["apply_guardrail"] if not callable(apply_fn): return TestCustomCodeGuardrailResponse( success=False, @@ -2055,7 +2095,7 @@ async def test_custom_code_guardrail( # Step 4: Execute the function with timeout protection - def execute_guardrail(): + def execute_guardrail() -> object: return apply_fn(test_inputs, safe_request_data, request.input_type) try: diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 80d9ee21a44..1a0978c8eec 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -15,9 +15,9 @@ These are members of a Team on LiteLLM import asyncio import json import traceback -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from datetime import datetime, timezone -from typing import Any, Optional, cast +from typing import Any, cast import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Request, status @@ -28,6 +28,10 @@ from litellm._uuid import uuid from litellm.proxy._types import * from litellm.proxy.auth.auth_checks import get_team_object, get_user_object from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.user_api_key_cache import ( + object_permission_cache_key, + user_object_permission_id_cache_key, +) from litellm.proxy.hooks.user_management_event_hooks import UserManagementEventHooks from litellm.proxy.management_endpoints.common_daily_activity import ( DailySpendRecord, @@ -45,10 +49,6 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( generate_key_helper_fn, prepare_metadata_fields, ) -from litellm.proxy.common_utils.user_api_key_cache import ( - object_permission_cache_key, - user_object_permission_id_cache_key, -) from litellm.proxy.management_helpers.object_permission_utils import ( _set_object_permission, handle_update_object_permission_common, @@ -82,11 +82,74 @@ from litellm.types.proxy.management_endpoints.scim_v2 import ( ) if TYPE_CHECKING: + from prisma import models as prisma_models + from prisma import types as prisma_types + from prisma.actions import ( + LiteLLM_InvitationLinkActions, + LiteLLM_OrganizationMembershipActions, + LiteLLM_TeamMembershipActions, + LiteLLM_TeamTableActions, + LiteLLM_UserTableActions, + LiteLLM_VerificationTokenActions, + ) + + 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() +def _user_table( + prisma_client: "PrismaClient | None", +) -> "LiteLLM_UserTableActions[prisma_models.LiteLLM_UserTable]": + user_table: LiteLLM_UserTableActions[prisma_models.LiteLLM_UserTable] = UserRepository(prisma_client).table + return user_table + + +def _team_table( + prisma_client: "PrismaClient | None", +) -> "LiteLLM_TeamTableActions[prisma_models.LiteLLM_TeamTable]": + team_table: LiteLLM_TeamTableActions[prisma_models.LiteLLM_TeamTable] = TeamRepository(prisma_client).table + return team_table + + +def _verification_token_table( + prisma_client: "PrismaClient | None", +) -> "LiteLLM_VerificationTokenActions[prisma_models.LiteLLM_VerificationToken]": + token_table: LiteLLM_VerificationTokenActions[prisma_models.LiteLLM_VerificationToken] = ( + VerificationTokenRepository(prisma_client).table + ) + return token_table + + +def _organization_membership_table( + prisma_client: "PrismaClient | None", +) -> "LiteLLM_OrganizationMembershipActions[prisma_models.LiteLLM_OrganizationMembership]": + membership_table: LiteLLM_OrganizationMembershipActions[prisma_models.LiteLLM_OrganizationMembership] = ( + OrganizationMembershipRepository(prisma_client).table + ) + return membership_table + + +def _invitation_link_table( + prisma_client: "PrismaClient | None", +) -> "LiteLLM_InvitationLinkActions[prisma_models.LiteLLM_InvitationLink]": + invitation_table: LiteLLM_InvitationLinkActions[prisma_models.LiteLLM_InvitationLink] = InvitationLinkRepository( + prisma_client + ).table + return invitation_table + + +def _team_membership_table( + prisma_client: "PrismaClient | None", +) -> "LiteLLM_TeamMembershipActions[prisma_models.LiteLLM_TeamMembership]": + team_membership_table: LiteLLM_TeamMembershipActions[prisma_models.LiteLLM_TeamMembership] = ( + TeamMembershipRepository(prisma_client).table + ) + return team_membership_table + + def _hash_password_in_dict(data: dict) -> None: """Hash password field in-place if present.""" if "password" in data and data["password"] is not None: @@ -138,7 +201,7 @@ def _update_internal_new_user_params(data_json: dict, data: NewUserRequest) -> d async def _check_duplicate_user_field( field_name: str, field_value: str | None, - prisma_client: Any, + prisma_client: "PrismaClient | None", *, case_insensitive: bool = False, label: str | None = None, @@ -177,7 +240,7 @@ async def _check_duplicate_user_field( ) -async def _check_duplicate_user_email(user_email: str | None, prisma_client: Any) -> None: +async def _check_duplicate_user_email(user_email: str | None, prisma_client: "PrismaClient | None") -> None: """ Helper function to check if a user email already exists in the database. """ @@ -190,7 +253,7 @@ async def _check_duplicate_user_email(user_email: str | None, prisma_client: Any ) -async def _check_duplicate_user_id(user_id: str | None, prisma_client: Any) -> None: +async def _check_duplicate_user_id(user_id: str | None, prisma_client: "PrismaClient | None") -> None: """ Helper function to check if a user id already exists in the database. """ @@ -724,8 +787,8 @@ _SCIM_DIRECTORY_METADATA_KEYS = frozenset( def _redact_scim_enterprise_metadata( - metadata: dict[str, Any] | None, -) -> dict[str, Any] | None: + metadata: dict[str, object] | None, +) -> dict[str, object] | None: """SCIM enterprise attributes, entitlements, and roles are persisted in user metadata so reporting can group on them, but they are directory-only fields that generic user-info endpoints must not surface; SCIM clients read them @@ -845,7 +908,7 @@ async def user_info( async def _check_user_info_v2_access( user_api_key_dict: UserAPIKeyAuth, target_user_id: str, -) -> Optional["LiteLLM_UserTable"]: +) -> "prisma_models.LiteLLM_UserTable | None": """ Check if the caller is allowed to access the target user's info. @@ -867,7 +930,7 @@ async def _check_user_info_v2_access( # Helper: fetch the target user row (reused across branches). object_permission is included so # callers can read the user's MCP/vector-store entitlements without a second round trip. async def _fetch_target_user(): - return await UserRepository(prisma_client).table.find_unique( + return await _user_table(prisma_client).find_unique( where={"user_id": target_user_id}, include={"object_permission": True} ) @@ -882,9 +945,7 @@ async def _check_user_info_v2_access( # 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() @@ -892,7 +953,7 @@ async def _check_user_info_v2_access( 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.model_validate(team.model_dump()) if _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team_obj): @@ -1160,7 +1221,7 @@ 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.model_validate(updated_user_row.model_dump(exclude_none=True)) asyncio.create_task( @@ -1207,7 +1268,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. @@ -1295,13 +1356,9 @@ async def _update_single_user_helper( existing_user_row: BaseModel | None = 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) @@ -1690,7 +1747,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( @@ -1805,9 +1862,9 @@ async def bulk_user_update( async def get_user_key_counts( - prisma_client, + prisma_client: "PrismaClient | None", user_ids: list[str] | None = None, -): +) -> Mapping[str, int]: """ Helper function to get the count of keys for each user using Prisma's count method. @@ -1823,7 +1880,7 @@ async def get_user_key_counts( if not user_ids or len(user_ids) == 0: return {} - result = {} + result: dict[str, int] = {} # Get count for each user_id individually for user_id in user_ids: @@ -1876,9 +1933,9 @@ def _validate_sort_params(sort_by: str | None, sort_order: str) -> dict[str, str async def _authorize_user_list_request( user_api_key_dict: UserAPIKeyAuth, organization_ids: str | None, - prisma_client: Any, - user_api_key_cache: Any, - proxy_logging_obj: Any, + prisma_client: "PrismaClient | None", + user_api_key_cache: "UserApiKeyCache", + proxy_logging_obj: "ProxyLogging | None", ) -> str | None: """ Authorize the /user/list request and return the (possibly scoped) organization_ids string. @@ -2016,7 +2073,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 @@ -2064,7 +2121,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[prisma_models.LiteLLM_UserTable] | None = await UserRepository(prisma_client).table.find_many( where=where_conditions, skip=skip, take=page_size, @@ -2072,7 +2129,7 @@ async def get_users( ) # Get total count of user rows - total_count = await UserRepository(prisma_client).table.count(where=where_conditions) + total_count: int = await UserRepository(prisma_client).table.count(where=where_conditions) # Get key count for each user if users is not None: @@ -2168,7 +2225,7 @@ async def delete_user( caller_admin_org_ids: set = set() if not caller_is_proxy_admin: 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, "user_role": LitellmUserRoles.ORG_ADMIN.value, @@ -2188,7 +2245,7 @@ async def delete_user( # an N+1 DB call when delete_user is called with a large user_ids list. target_org_ids_by_user: dict[str, set] = {} if not caller_is_proxy_admin: - all_target_memberships = await OrganizationMembershipRepository(prisma_client).table.find_many( + all_target_memberships = await _organization_membership_table(prisma_client).find_many( where={"user_id": {"in": data.user_ids}} ) for m in all_target_memberships: @@ -2276,10 +2333,10 @@ async def delete_user( # 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}}, @@ -2290,13 +2347,13 @@ async def delete_user( ) ## DELETE ASSOCIATED ORGANIZATION MEMBERSHIPS - await OrganizationMembershipRepository(prisma_client).table.delete_many(where={"user_id": {"in": data.user_ids}}) + await _organization_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 @@ -2348,9 +2405,9 @@ async def add_internal_user_to_organization( async def _resolve_org_filter_for_user_search( user_api_key_dict: UserAPIKeyAuth, team_id: str | None, - prisma_client: Any, - user_api_key_cache: Any, - proxy_logging_obj: Any, + prisma_client: "PrismaClient | None", + user_api_key_cache: "UserApiKeyCache", + proxy_logging_obj: "ProxyLogging | None", ) -> list[str] | None: """ Return a list of org IDs to filter by, or ``None`` for no filter. @@ -2414,9 +2471,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 | None", + user_api_key_cache: "UserApiKeyCache", + proxy_logging_obj: "ProxyLogging | None", ) -> 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 @@ -2506,7 +2563,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: prisma_types.LiteLLM_UserTableWhereInput = {} if user_id: where_conditions["user_id"] = { @@ -2525,7 +2582,7 @@ async def ui_view_users( where_conditions["organization_memberships"] = {"some": {"organization_id": {"in": org_filter_ids}}} # Query users with pagination and filters - users: list[BaseModel] | None = await UserRepository(prisma_client).table.find_many( + users = await _user_table(prisma_client).find_many( where=where_conditions, skip=skip, take=page_size, @@ -2557,7 +2614,7 @@ async def _resolve_user_email_metadata( } 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 a94a75fdfa3..b910eee7130 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -18,9 +18,9 @@ import os 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 +from typing import Any, Callable, Dict, List, Literal, Optional, Protocol, Tuple, TypeVar, cast import fastapi import yaml @@ -37,6 +37,7 @@ from litellm.constants import ( ) from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.models.credentials import CredentialItem from litellm.proxy._experimental.mcp_server.db import ( rotate_mcp_server_credentials_master_key, rotate_mcp_user_credentials_master_key, @@ -101,8 +102,9 @@ from litellm.proxy.utils import ( handle_exception_on_proxy, is_valid_api_key, ) +from litellm.repositories.base_repository import BaseRepository from litellm.repositories.budget_repository import BudgetRepository -from litellm.repositories.config_repository import ConfigRepository +from litellm.repositories.config_repository import ConfigParam, ConfigRepository from litellm.repositories.credentials_repository import CredentialsRepository from litellm.repositories.model_repository import ModelRepository from litellm.repositories.table_repositories import ( @@ -131,6 +133,68 @@ from litellm.types.utils import ( TeamUIKeyGenerationConfig, ) +_PrismaRowT = TypeVar("_PrismaRowT") +_RepositoryModelT = TypeVar("_RepositoryModelT", bound=BaseModel) + + +class _PrismaTableActions(Protocol[_PrismaRowT]): + """Typed view of the Prisma table actions a repository exposes through its untyped ``table``.""" + + async def find_unique( + self, + *, + where: Mapping[str, object], + include: Mapping[str, object] | None = None, + ) -> _PrismaRowT | None: ... + + async def find_first( + self, + *, + where: Mapping[str, object], + include: Mapping[str, object] | None = None, + ) -> _PrismaRowT | None: ... + + async def find_many( + self, + *, + where: Mapping[str, object] | None = None, + include: Mapping[str, object] | None = None, + order: Mapping[str, object] | None = None, + skip: int | None = None, + take: int | None = None, + ) -> list[_PrismaRowT]: ... + + async def count(self, *, where: Mapping[str, object] | None = None) -> int: ... + + async def create_many(self, *, data: Sequence[Mapping[str, object]]) -> int: ... + + async def update( + self, + *, + where: Mapping[str, object], + data: Mapping[str, object], + ) -> _PrismaRowT | None: ... + + +def _prisma_table( + repository: BaseRepository[_RepositoryModelT], +) -> _PrismaTableActions[_RepositoryModelT]: + return repository.table + + +def _deleted_verification_token_table( + prisma_client: PrismaClient, +) -> _PrismaTableActions[LiteLLM_DeletedVerificationToken]: + return DeletedVerificationTokenRepository(prisma_client).table + + +def _credentials_table(prisma_client: PrismaClient) -> _PrismaTableActions[CredentialItem]: + return CredentialsRepository(prisma_client).table + + +def _config_table(prisma_client: PrismaClient) -> _PrismaTableActions[ConfigParam]: + return ConfigRepository(prisma_client).table + async def _check_custom_key_allowed(custom_key_value: Optional[str]) -> None: """Raise 403 if custom API keys are disabled and a custom key was provided.""" @@ -490,7 +554,7 @@ _NON_ADMIN_SAFE_ALLOWED_ROUTES_PRESETS = frozenset({"llm_api_routes", "info_rout def _validate_caller_can_change_key_ownership( data: Optional[BaseModel], - existing_key_row: Any, + existing_key_row: LiteLLM_VerificationToken, user_api_key_dict: UserAPIKeyAuth, ) -> None: """ @@ -670,7 +734,7 @@ async def validate_team_id_used_in_service_account_request( ) # check if team_id exists in the database - team = await TeamRepository(prisma_client).table.find_unique( + team = await _prisma_table(TeamRepository(prisma_client)).find_unique( where={"team_id": team_id}, ) if team is None: @@ -1261,7 +1325,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 = await _prisma_table(VerificationTokenRepository(prisma_client)).find_many( where={"team_id": team_table.team_id}, ) # Exclude the key being updated to avoid double-counting its limits. @@ -1405,7 +1469,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 = await _prisma_table(UserRepository(prisma_client)).find_unique( where={"user_id": user_api_key_dict.user_id}, include={"organization_memberships": True}, ) @@ -1443,7 +1507,7 @@ async def _check_org_key_limits( # get all organization keys # calculate allocated tpm/rpm limit # check if specified tpm/rpm limit is greater than allocated tpm/rpm limit - keys = await VerificationTokenRepository(prisma_client).table.find_many( + keys = await _prisma_table(VerificationTokenRepository(prisma_client)).find_many( where={"organization_id": org_table.organization_id}, ) # Exclude the key being updated to avoid double-counting its limits. @@ -2048,9 +2112,9 @@ async def _get_and_validate_existing_key( if token is not None: hashed_token = _hash_token_if_needed(token=token) - existing_key_row: LiteLLM_VerificationToken | None = await VerificationTokenRepository( - prisma_client - ).table.find_unique(where={"token": hashed_token}) + existing_key_row: LiteLLM_VerificationToken | None = await _prisma_table( + VerificationTokenRepository(prisma_client) + ).find_unique(where={"token": hashed_token}) if existing_key_row is None: raise ProxyException( @@ -2070,7 +2134,7 @@ async def _get_and_validate_existing_key( code=status.HTTP_400_BAD_REQUEST, ) - rows: list[LiteLLM_VerificationToken] = await VerificationTokenRepository(prisma_client).table.find_many( + rows: list[LiteLLM_VerificationToken] = await _prisma_table(VerificationTokenRepository(prisma_client)).find_many( where={"key_alias": key_alias}, take=2 ) @@ -2112,7 +2176,7 @@ 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, @@ -2265,9 +2329,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.""" @@ -2302,12 +2366,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. @@ -2939,7 +3003,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() @@ -3416,7 +3480,7 @@ async def info_key_fn_v2( # Resolve key_aliases to tokens so we never pass token=None (unbounded query) tokens_to_query = list(data.keys) if data.keys else [] if data.key_aliases: - alias_rows = await VerificationTokenRepository(prisma_client).table.find_many( + alias_rows = await _prisma_table(VerificationTokenRepository(prisma_client)).find_many( where={"key_alias": {"in": data.key_aliases}}, include={"litellm_budget_table": True}, ) @@ -4088,7 +4152,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 [] @@ -4141,13 +4205,13 @@ def _transform_verification_tokens_to_deleted_records( async def _save_deleted_verification_token_records( - records: List[Dict[str, Any]], + records: Sequence[Mapping[str, object]], prisma_client: PrismaClient, ) -> None: """Save deleted verification token records to the database.""" if not records: return - await DeletedVerificationTokenRepository(prisma_client).table.create_many(data=records) + await _deleted_verification_token_table(prisma_client).create_many(data=records) async def _persist_deleted_verification_tokens( @@ -4175,7 +4239,7 @@ 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( + _keys_being_deleted = await _prisma_table(VerificationTokenRepository(prisma_client)).find_many( where={"key_alias": {"in": key_aliases}} ) @@ -4212,7 +4276,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: Optional[List] = await _prisma_table(ModelRepository(prisma_client)).find_many() except Exception: models = None # 2. process model table @@ -4242,7 +4306,7 @@ async def _rotate_master_key( ) # 3. process config table try: - config = await ConfigRepository(prisma_client).table.find_many() + config = await _config_table(prisma_client).find_many() except Exception: config = None @@ -4263,7 +4327,7 @@ async def _rotate_master_key( ) if encrypted_env_vars: - await ConfigRepository(prisma_client).table.update( + await _config_table(prisma_client).update( where={"param_name": "environment_variables"}, data={"param_value": prisma.Json(encrypted_env_vars)}, # type: ignore[attr-defined] ) @@ -4307,7 +4371,7 @@ async def _rotate_master_key( # 5. process credentials table try: - credentials = await CredentialsRepository(prisma_client).table.find_many() + credentials = await _credentials_table(prisma_client).find_many() except Exception: credentials = None if credentials: @@ -4330,7 +4394,7 @@ async def _rotate_master_key( _cred_data["credential_info"] = prisma.Json( # type: ignore[attr-defined] _cred_data["credential_info"] ) - await CredentialsRepository(prisma_client).table.update( + await _credentials_table(prisma_client).update( where={"credential_name": cred.credential_name}, data={ **_cred_data, @@ -4772,7 +4836,7 @@ 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 = await _prisma_table(VerificationTokenRepository(prisma_client)).find_unique( where={"token": hashed_api_key}, ) if _key_in_db is None: @@ -4976,7 +5040,7 @@ async def reset_key_spend_fn( else: hashed_api_key = hash_token(key) - _key_in_db = await VerificationTokenRepository(prisma_client).table.find_unique( + _key_in_db = await _prisma_table(VerificationTokenRepository(prisma_client)).find_unique( where={"token": hashed_api_key}, include={"litellm_budget_table": True}, ) @@ -4996,7 +5060,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 = await _prisma_table(VerificationTokenRepository(prisma_client)).update( where={"token": hashed_api_key}, data={"spend": reset_to}, ) @@ -5067,7 +5131,7 @@ async def validate_key_list_check( param="user_id", code=status.HTTP_403_FORBIDDEN, ) - complete_user_info_db_obj: Optional[BaseModel] = await UserRepository(prisma_client).table.find_unique( + complete_user_info_db_obj: Optional[BaseModel] = await _prisma_table(UserRepository(prisma_client)).find_unique( where={"user_id": user_api_key_dict.user_id}, include={"organization_memberships": True}, ) @@ -5421,8 +5485,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 @@ -5435,7 +5499,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 = await _prisma_table(UserRepository(prisma_client)).find_unique( + where={"user_id": user_api_key_dict.user_id} + ) if user_row is not None: user_teams = getattr(user_row, "teams", []) or [] @@ -5493,7 +5559,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 != ''", @@ -5601,7 +5667,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}}]} @@ -5848,11 +5914,11 @@ async def _list_key_helper( # Get total count of keys if use_deleted_table: - total_count = await DeletedVerificationTokenRepository(prisma_client).table.count( + total_count = await _deleted_verification_token_table(prisma_client).count( where=where # type: ignore ) else: - total_count = await VerificationTokenRepository(prisma_client).table.count( + total_count = await _prisma_table(VerificationTokenRepository(prisma_client)).count( where=where # type: ignore ) @@ -5931,8 +5997,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 | None, user_api_key_cache: UserApiKeyCache, route: str, ) -> None: @@ -5951,7 +6017,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 = await _prisma_table(VerificationTokenRepository(prisma_client)).find_unique( + where={"token": hashed_token} + ) if target_key_row is None: raise HTTPException( status_code=404, @@ -6047,7 +6115,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 = await _prisma_table(VerificationTokenRepository(prisma_client)).find_unique( + where={"token": hashed_token} + ) if existing_record is None: raise ProxyException( message="Key not found.", @@ -6077,7 +6147,7 @@ async def block_key( ) ) - record = await VerificationTokenRepository(prisma_client).table.update( + record = await _prisma_table(VerificationTokenRepository(prisma_client)).update( where={"token": hashed_token}, data={"blocked": True}, # type: ignore ) @@ -6158,7 +6228,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 = await _prisma_table(VerificationTokenRepository(prisma_client)).find_unique( + where={"token": hashed_token} + ) if existing_record is None: raise ProxyException( message="Key not found.", @@ -6188,7 +6260,7 @@ async def unblock_key( ) ) - record = await VerificationTokenRepository(prisma_client).table.update( + record = await _prisma_table(VerificationTokenRepository(prisma_client)).update( where={"token": hashed_token}, data={"blocked": False}, # type: ignore ) @@ -6443,7 +6515,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: """ @@ -6459,12 +6531,12 @@ async def _enforce_unique_key_alias( ProxyException: If key alias already exists on a different key """ if key_alias is not None and prisma_client is not None: - where_clause: dict[str, Any] = {"key_alias": key_alias} + where_clause: dict[str, object] = {"key_alias": key_alias} if existing_key_token: # Exclude the current key from the uniqueness check where_clause["NOT"] = {"token": existing_key_token} - existing_key = await VerificationTokenRepository(prisma_client).table.find_first(where=where_clause) + existing_key = await _prisma_table(VerificationTokenRepository(prisma_client)).find_first(where=where_clause) if existing_key is not None: raise ProxyException( message=f"Key with alias '{key_alias}' already exists. Unique key aliases across all keys are required.", diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 41904e3883b..4885ec42578 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -16,7 +16,7 @@ import json from collections.abc import Mapping, Sequence from typing import Any, Dict, List, Literal, Optional, Set, Tuple, Union, cast -from fastapi import APIRouter, Depends, HTTPException, Header, Request, status +from fastapi import APIRouter, Depends, Header, HTTPException, Request, status from pydantic import BaseModel, ConfigDict, Field from litellm._logging import verbose_proxy_logger @@ -39,6 +39,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.management_endpoints.common_utils import _is_user_team_admin from litellm.proxy.management_endpoints.team_endpoints import ( _refresh_cached_team, @@ -49,18 +50,18 @@ from litellm.proxy.management_endpoints.team_endpoints import ( update_team as _legacy_update_team, ) from litellm.proxy.management_helpers.audit_logs import create_object_audit_log -from litellm.proxy.utils import PrismaClient +from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.repositories.model_repository import ModelRepository from litellm.repositories.table_repositories import ModelTableRepository from litellm.repositories.team_repository import TeamRepository from litellm.router import Router -from litellm.types.proxy.management_endpoints.model_management_endpoints import ( - UpdateUsefulLinksRequest, -) from litellm.router_utils.auto_router_model_naming import ( STRATEGY_ROUTER_PARAM_FIELDS, validate_strategy_router_model_write, ) +from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + UpdateUsefulLinksRequest, +) from litellm.types.router import ( SPECIAL_MODEL_INFO_PARAMS, Deployment, @@ -843,8 +844,8 @@ async def _get_team_public_model_names( async def _remove_unbacked_team_models( model_params: Deployment, prisma_client: PrismaClient, - user_api_key_cache: Any, - proxy_logging_obj: Any, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, llm_router: Router | None = None, ) -> None: """ @@ -904,7 +905,7 @@ async def _remove_unbacked_team_models( if existing_team_row is None: return - updated_team_row = await prisma_client.db.litellm_teamtable.update( + updated_team_row: LiteLLM_TeamTable = await prisma_client.db.litellm_teamtable.update( where={"team_id": team_id}, data={"models": [model for model in existing_team_row.models if model not in names_to_remove]}, include={"object_permission": True}, # type: ignore diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py index e44299d018d..f1f71bf6458 100644 --- a/litellm/proxy/management_endpoints/scim/scim_v2.py +++ b/litellm/proxy/management_endpoints/scim/scim_v2.py @@ -7,7 +7,18 @@ This is an enterprise feature and requires a premium license. import re from collections.abc import Mapping, Sequence from itertools import chain -from typing import Any, Dict, Iterable, List, NamedTuple, Optional, Set, Tuple +from typing import ( + TYPE_CHECKING, + Dict, + Iterable, + List, + NamedTuple, + Optional, + Protocol, + Set, + Tuple, + overload, +) from fastapi import ( APIRouter, @@ -69,13 +80,95 @@ from litellm.repositories.verification_token_repository import ( ) from litellm.types.proxy.management_endpoints.scim_v2 import * +if TYPE_CHECKING: + from prisma.models import LiteLLM_VerificationToken as PrismaVerificationToken + + +class _UserTableClient(Protocol): + async def find_first(self, where: Mapping[str, object]) -> LiteLLM_UserTable | None: ... + + async def find_unique(self, where: Mapping[str, object]) -> LiteLLM_UserTable | None: ... + + async def find_many( + self, + where: Mapping[str, object] | None = None, + skip: int | None = None, + take: int | None = None, + order: Mapping[str, str] | None = None, + ) -> Sequence[LiteLLM_UserTable]: ... + + async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> LiteLLM_UserTable: ... + + async def delete(self, where: Mapping[str, object]) -> LiteLLM_UserTable | None: ... + + async def count(self, where: Mapping[str, object] | None = None) -> int: ... + + +class _TeamTableClient(Protocol): + async def find_unique(self, where: Mapping[str, object]) -> LiteLLM_TeamTable | None: ... + + async def find_many( + self, + where: Mapping[str, object] | None = None, + skip: int | None = None, + take: int | None = None, + order: Mapping[str, str] | None = None, + ) -> Sequence[LiteLLM_TeamTable]: ... + + async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> LiteLLM_TeamTable: ... + + async def delete(self, where: Mapping[str, object]) -> LiteLLM_TeamTable | None: ... + + async def count(self, where: Mapping[str, object] | None = None) -> int: ... + + +class _VerificationTokenTableClient(Protocol): + async def find_many(self, where: Mapping[str, object] | None = None) -> "Sequence[PrismaVerificationToken]": ... + + async def update( + self, where: Mapping[str, object], data: Mapping[str, object] + ) -> "PrismaVerificationToken | None": ... + + +class _UserReferencingTableClient(Protocol): + async def delete_many(self, where: Mapping[str, object]) -> int: ... + + +@overload +def _table(repository: UserRepository) -> _UserTableClient: ... + + +@overload +def _table(repository: TeamRepository) -> _TeamTableClient: ... + + +@overload +def _table(repository: VerificationTokenRepository) -> _VerificationTokenTableClient: ... + + +@overload +def _table( + repository: InvitationLinkRepository | OrganizationMembershipRepository | TeamMembershipRepository, +) -> _UserReferencingTableClient: ... + + +def _table( + repository: UserRepository + | TeamRepository + | VerificationTokenRepository + | InvitationLinkRepository + | OrganizationMembershipRepository + | TeamMembershipRepository, +) -> object: + return repository.table + class UserProvisionerHelpers: """Helper methods for user provisioning operations.""" @staticmethod async def handle_existing_user_by_email( - prisma_client, + prisma_client: PrismaClient, new_user_request: NewUserRequest, admin_group: Optional[str] = None, ) -> Optional[SCIMUser]: @@ -97,7 +190,7 @@ class UserProvisionerHelpers: if not new_user_request.user_email: return None - existing_user = await UserRepository(prisma_client).table.find_first( + existing_user = await _table(UserRepository(prisma_client)).find_first( where={"user_email": new_user_request.user_email} ) @@ -107,7 +200,7 @@ class UserProvisionerHelpers: new_teams = list(dict.fromkeys(new_user_request.teams or [])) if new_user_request.user_id != existing_user.user_id: - await UserRepository(prisma_client).table.update( + await _table(UserRepository(prisma_client)).update( where={"user_id": existing_user.user_id}, data={"user_id": new_user_request.user_id}, ) @@ -119,7 +212,7 @@ class UserProvisionerHelpers: raise_on_error=True, ) - updated_user = await UserRepository(prisma_client).table.update( + updated_user = await _table(UserRepository(prisma_client)).update( where={"user_id": new_user_request.user_id}, data={ "user_email": new_user_request.user_email, @@ -177,11 +270,11 @@ async def _get_prisma_client_or_raise_exception(): return prisma_client -async def _check_user_exists(user_id: str): +async def _check_user_exists(user_id: str) -> LiteLLM_UserTable: """Check if user exists and return user, raise 404 if not found.""" prisma_client = await _get_prisma_client_or_raise_exception() - user = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id}) + user = await _table(UserRepository(prisma_client)).find_unique(where={"user_id": user_id}) if not user: raise HTTPException(status_code=404, detail={"error": f"User not found with ID: {user_id}"}) @@ -189,11 +282,11 @@ async def _check_user_exists(user_id: str): return user -async def _check_team_exists(team_id: str): +async def _check_team_exists(team_id: str) -> LiteLLM_TeamTable: """Check if team exists and return team, raise 404 if not found.""" prisma_client = await _get_prisma_client_or_raise_exception() - team = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id}) + team = await _table(TeamRepository(prisma_client)).find_unique(where={"team_id": team_id}) if not team: raise HTTPException(status_code=404, detail={"error": f"Group not found with ID: {team_id}"}) @@ -236,9 +329,9 @@ def _build_scim_metadata( enterprise: Optional[SCIMEnterpriseUser] = None, entitlements: list[SCIMMultiValuedAttribute] | None = None, roles: list[SCIMMultiValuedAttribute] | None = None, -) -> Dict[str, Any]: +) -> dict[str, object]: """Build metadata dictionary with SCIM data.""" - metadata: Dict[str, Any] = { + metadata: dict[str, object] = { "scim_metadata": LiteLLM_UserScimMetadata( givenName=given_name, familyName=family_name, @@ -338,13 +431,15 @@ def _resolve_scim_user_role( return default_role -async def _scim_groups_from_team_ids(prisma_client: Any, team_ids: list[str]) -> list[SCIMUserGroup]: +async def _scim_groups_from_team_ids(prisma_client: PrismaClient, team_ids: list[str]) -> list[SCIMUserGroup]: """ Build SCIMUserGroup objects from team ids, populating display from each team's alias so admin-group matching by display name works the same way it does on PUT (where SCIM groups carry display names natively). """ - teams = [await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id}) for team_id in team_ids] + teams = [ + await _table(TeamRepository(prisma_client)).find_unique(where={"team_id": team_id}) for team_id in team_ids + ] return [ SCIMUserGroup( value=team_id, @@ -354,7 +449,7 @@ async def _scim_groups_from_team_ids(prisma_client: Any, team_ids: list[str]) -> ] -async def _recompute_scim_member_roles(prisma_client: Any, user_ids: Iterable[str]) -> None: +async def _recompute_scim_member_roles(prisma_client: PrismaClient, user_ids: Iterable[str]) -> None: """ Recompute and persist each user's global proxy role from their resulting team membership. No-op unless scim_admin_group is configured, so a SCIM group write @@ -367,7 +462,7 @@ async def _recompute_scim_member_roles(prisma_client: Any, user_ids: Iterable[st default_role = _default_scim_user_role() for user_id in user_ids: - user = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id}) + user = await _table(UserRepository(prisma_client)).find_unique(where={"user_id": user_id}) if user is None: continue resolved_role = _resolve_scim_user_role( @@ -375,7 +470,7 @@ async def _recompute_scim_member_roles(prisma_client: Any, user_ids: Iterable[st admin_group, default_role, ) - await UserRepository(prisma_client).table.update( + await _table(UserRepository(prisma_client)).update( where={"user_id": user_id}, data={"user_role": resolved_role}, ) @@ -471,7 +566,7 @@ async def _classify_group_member(member: SCIMMember, prisma_client: PrismaClient if member_type == "group": return _SkippedGroupMember(value=value, reason="nested_group") - user = await UserRepository(prisma_client).table.find_unique(where={"user_id": value}) + user = await _table(UserRepository(prisma_client)).find_unique(where={"user_id": value}) if user is not None: return _ResolvedUserMember(user_id=value) @@ -479,7 +574,7 @@ async def _classify_group_member(member: SCIMMember, prisma_client: PrismaClient return _SkippedGroupMember(value=value, reason="non_user_type") if member_type is None: - team = await TeamRepository(prisma_client).table.find_unique(where={"team_id": value}) + team = await _table(TeamRepository(prisma_client)).find_unique(where={"team_id": value}) if team is not None and _team_metadata_has_scim_provenance(team.metadata): return _SkippedGroupMember(value=value, reason="existing_team") @@ -619,7 +714,7 @@ async def _get_team_members_display(member_ids: List[str]) -> List[SCIMMember]: members: List[SCIMMember] = [] for member_id in member_ids: - user = await UserRepository(prisma_client).table.find_unique(where={"user_id": member_id}) + user = await _table(UserRepository(prisma_client)).find_unique(where={"user_id": member_id}) if user: display_name = user.user_email or user.user_id members.append(SCIMMember(value=user.user_id, display=display_name, type="User")) @@ -652,7 +747,7 @@ async def _handle_team_membership_changes( SCIM_BLOCKED_METADATA_KEY = "scim_blocked" -def _key_was_scim_blocked(metadata: Any) -> bool: +def _key_was_scim_blocked(metadata: object) -> bool: """True if a verification token carries the SCIM-block marker in metadata.""" return isinstance(metadata, dict) and metadata.get(SCIM_BLOCKED_METADATA_KEY) is True @@ -676,7 +771,7 @@ async def _set_user_keys_blocked(user_id: str, blocked: bool) -> int: # `blocked` is a nullable column with no default, so existing rows # typically hold NULL; treat NULL as "not blocked" since SQL equality # on NULL would otherwise silently skip them. - candidates = await VerificationTokenRepository(prisma_client).table.find_many( + candidates = await _table(VerificationTokenRepository(prisma_client)).find_many( where={ "user_id": user_id, "OR": [{"blocked": False}, {"blocked": None}], @@ -684,7 +779,7 @@ async def _set_user_keys_blocked(user_id: str, blocked: bool) -> int: ) affected_keys = candidates else: - candidates = await VerificationTokenRepository(prisma_client).table.find_many( + candidates = await _table(VerificationTokenRepository(prisma_client)).find_many( where={"user_id": user_id, "blocked": True}, ) affected_keys = [k for k in candidates if _key_was_scim_blocked(k.metadata)] @@ -693,12 +788,12 @@ async def _set_user_keys_blocked(user_id: str, blocked: bool) -> int: return 0 for key_row in affected_keys: - current_metadata: Dict[str, Any] = dict(key_row.metadata) if isinstance(key_row.metadata, dict) else {} + current_metadata: dict[str, object] = dict(key_row.metadata) if isinstance(key_row.metadata, dict) else {} if blocked: new_metadata = {**current_metadata, SCIM_BLOCKED_METADATA_KEY: True} else: new_metadata = {k: v for k, v in current_metadata.items() if k != SCIM_BLOCKED_METADATA_KEY} - await VerificationTokenRepository(prisma_client).table.update( + await _table(VerificationTokenRepository(prisma_client)).update( where={"token": key_row.token}, data={"blocked": blocked, "metadata": safe_dumps(new_metadata)}, ) @@ -719,14 +814,14 @@ async def _set_user_keys_blocked(user_id: str, blocked: bool) -> int: return len(affected_keys) -async def _delete_rows_referencing_user(prisma_client: Any, *, user_id: str) -> None: +async def _delete_rows_referencing_user(prisma_client: PrismaClient, *, user_id: str) -> None: """Drop rows whose foreign keys reference ``LiteLLM_UserTable.user_id``. Required before deleting the user row itself, otherwise Postgres rejects the user delete with an FK constraint violation (e.g. ``LiteLLM_InvitationLink_user_id_fkey``). """ - await InvitationLinkRepository(prisma_client).table.delete_many( + await _table(InvitationLinkRepository(prisma_client)).delete_many( where={ "OR": [ {"user_id": user_id}, @@ -735,11 +830,11 @@ async def _delete_rows_referencing_user(prisma_client: Any, *, user_id: str) -> ] } ) - await OrganizationMembershipRepository(prisma_client).table.delete_many(where={"user_id": user_id}) - await TeamMembershipRepository(prisma_client).table.delete_many(where={"user_id": user_id}) + await _table(OrganizationMembershipRepository(prisma_client)).delete_many(where={"user_id": user_id}) + await _table(TeamMembershipRepository(prisma_client)).delete_many(where={"user_id": user_id}) -def _scim_active_value(metadata: Optional[Dict[str, Any]]) -> Optional[bool]: +def _scim_active_value(metadata: Optional[Mapping[str, object]]) -> Optional[bool]: """Read the SCIM active flag from a user's metadata dict, if present.""" if not metadata: return None @@ -749,6 +844,12 @@ def _scim_active_value(metadata: Optional[Dict[str, Any]]) -> Optional[bool]: return bool(value) +def _user_scim_active(user: LiteLLM_UserTable) -> Optional[bool]: + """Read the SCIM active flag off a user row's metadata, if present.""" + metadata: dict[str, object] | None = user.metadata + return _scim_active_value(metadata) + + async def _create_user_if_not_exists(user_id: str, created_via: str = "scim_group") -> Optional[NewUserResponse]: """ Helper function to create a user if they don't exist. @@ -820,7 +921,7 @@ async def set_scim_content_type(response: Response): response.headers["Content-Type"] = "application/scim+json" -def _get_resource_types(base_url: str = "/scim/v2") -> list: +def _get_resource_types(base_url: str = "/scim/v2") -> Sequence[SCIMResourceType]: """Return the list of SCIM ResourceType definitions per RFC 7643 Section 6.""" return [ SCIMResourceType( @@ -848,7 +949,7 @@ def _get_resource_types(base_url: str = "/scim/v2") -> list: ] -def _get_schemas() -> list: +def _get_schemas() -> Sequence[SCIMSchema]: """Return the list of SCIM Schema definitions per RFC 7643 Section 7.""" return [ SCIMSchema( @@ -1241,7 +1342,7 @@ async def get_users( try: prisma_client = await _get_prisma_client_or_raise_exception() # Parse filter if provided (basic support) - where_conditions: Dict[str, Any] = {} + where_conditions: dict[str, object] = {} if filter: # Okta locates users by userName before deprovisioning. LiteLLM # exposes SCIM userName from user_email, while older SCIM-created @@ -1258,7 +1359,7 @@ async def get_users( where_conditions["user_email"] = filter_value # Get users from database - users: List[LiteLLM_UserTable] = await UserRepository(prisma_client).table.find_many( + users: Sequence[LiteLLM_UserTable] = await _table(UserRepository(prisma_client)).find_many( where=where_conditions, skip=(startIndex - 1), take=count, @@ -1266,7 +1367,7 @@ async def get_users( ) # Get total count for pagination - total_count = await UserRepository(prisma_client).table.count(where=where_conditions) + total_count = await _table(UserRepository(prisma_client)).count(where=where_conditions) # Convert to SCIM format scim_users: List[SCIMUser] = [] @@ -1330,7 +1431,7 @@ async def create_user( # Check if user already exists if user.userName: - existing_user = await UserRepository(prisma_client).table.find_unique(where={"user_id": user.userName}) + existing_user = await _table(UserRepository(prisma_client)).find_unique(where={"user_id": user.userName}) if existing_user: raise HTTPException( status_code=409, @@ -1406,7 +1507,7 @@ async def update_user( prisma_client = await _get_prisma_client_or_raise_exception() existing_user = await _check_user_exists(user_id) - prev_active = _scim_active_value(existing_user.metadata) + prev_active = _user_scim_active(existing_user) user_data = _extract_scim_user_data(user) @@ -1447,7 +1548,7 @@ async def update_user( user.groups or [], admin_group, _default_scim_user_role() ) - updated_user = await UserRepository(prisma_client).table.update( + updated_user = await _table(UserRepository(prisma_client)).update( where={"user_id": user_id}, data=update_data, ) @@ -1483,19 +1584,20 @@ async def delete_user( existing_user = await _check_user_exists(user_id) # Get teams user belongs to - teams = [] - if existing_user.teams: - for team_id in existing_user.teams: - team = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id}) - if team: - teams.append(team) + found_teams = tuple( + [ + await _table(TeamRepository(prisma_client)).find_unique(where={"team_id": team_id}) + for team_id in existing_user.teams or [] + ] + ) + teams = tuple(team for team in found_teams if team) # Remove user from all teams for team in teams: - current_members = team.members or [] + current_members: Sequence[str] = team.members or [] if user_id in current_members: new_members = [m for m in current_members if m != user_id] - await TeamRepository(prisma_client).table.update( + await _table(TeamRepository(prisma_client)).update( where={"team_id": team.team_id}, data={"members": new_members} ) @@ -1511,7 +1613,7 @@ async def delete_user( await _delete_rows_referencing_user(prisma_client, user_id=user_id) # Delete user - await UserRepository(prisma_client).table.delete(where={"user_id": user_id}) + await _table(UserRepository(prisma_client)).delete(where={"user_id": user_id}) return Response(status_code=204) except Exception as e: @@ -1581,7 +1683,7 @@ def _extract_ids_from_path_filter(path: str | None, attribute: str) -> List[str] return [extracted] if extracted else [] -def _handle_displayname_update(op_type: str, value: Any, update_data: Dict[str, Any]) -> None: +def _handle_displayname_update(op_type: str, value: object, update_data: dict[str, object]) -> None: """Handle displayname updates.""" if op_type == "remove": update_data["user_alias"] = None @@ -1589,7 +1691,7 @@ def _handle_displayname_update(op_type: str, value: Any, update_data: Dict[str, update_data["user_alias"] = str(value) -def _handle_externalid_update(op_type: str, value: Any, update_data: Dict[str, Any]) -> None: +def _handle_externalid_update(op_type: str, value: object, update_data: dict[str, object]) -> None: """Handle externalid updates.""" if op_type == "remove": update_data["sso_user_id"] = None @@ -1597,7 +1699,7 @@ def _handle_externalid_update(op_type: str, value: Any, update_data: Dict[str, A update_data["sso_user_id"] = str(value) -def _handle_active_update(op_type: str, value: Any, metadata: Dict[str, Any]) -> None: +def _handle_active_update(op_type: str, value: object, metadata: dict[str, object]) -> None: """Handle active status updates.""" if op_type == "remove": metadata.pop("scim_active", None) @@ -1610,7 +1712,7 @@ def _handle_active_update(op_type: str, value: Any, metadata: Dict[str, Any]) -> metadata["scim_active"] = bool_val -def _handle_name_update(path: str, op_type: str, value: Any, scim_metadata: Dict[str, Any]) -> None: +def _handle_name_update(path: str, op_type: str, value: object, scim_metadata: dict[str, object]) -> None: """Handle name field updates (givenName, familyName).""" if path == "name.givenname": if op_type == "remove": @@ -1624,7 +1726,7 @@ def _handle_name_update(path: str, op_type: str, value: Any, scim_metadata: Dict scim_metadata["familyName"] = str(value) -def _handle_group_operations(op_type: str, value: Any, teams_set: Set[str], path: str | None) -> Set[str] | None: +def _handle_group_operations(op_type: str, value: object, teams_set: Set[str], path: str | None) -> Set[str] | None: """Handle group/team membership operations.""" group_values = _extract_group_values(value) if not group_values and value is None: @@ -1644,7 +1746,7 @@ def _multi_valued_attribute_base(path: str) -> str: return path.split("[", 1)[0].split(".", 1)[0] -def _handle_multi_valued_attribute_update(path: str, op_type: str, value: Any, metadata: dict[str, Any]) -> None: +def _handle_multi_valued_attribute_update(path: str, op_type: str, value: object, metadata: dict[str, object]) -> None: """Handle add/replace/remove for the entitlements and roles multi-valued attributes.""" base = _multi_valued_attribute_base(path) metadata_key = SCIM_MULTI_VALUED_ATTRIBUTE_METADATA_KEYS[base] @@ -1681,7 +1783,7 @@ def _handle_multi_valued_attribute_update(path: str, op_type: str, value: Any, m metadata[metadata_key] = dumped -def _handle_generic_metadata(path: str, op_type: str, value: Any, metadata: Dict[str, Any]) -> None: +def _handle_generic_metadata(path: str, op_type: str, value: object, metadata: dict[str, object]) -> None: """Handle generic metadata operations for unknown paths.""" if op_type == "remove": metadata.pop(path, None) @@ -1692,9 +1794,9 @@ def _handle_generic_metadata(path: str, op_type: str, value: Any, metadata: Dict def _apply_patch_ops( existing_user: LiteLLM_UserTable, patch_ops: SCIMPatchOp, -) -> Tuple[Dict[str, Any], Set[str]]: +) -> Tuple[dict[str, object], Set[str]]: """Apply patch operations and return update data and final team set.""" - update_data: Dict[str, Any] = {} + update_data: dict[str, object] = {} metadata = existing_user.metadata or {} scim_metadata = metadata.get("scim_metadata", {}) @@ -1843,14 +1945,15 @@ async def patch_user( prisma_client = await _get_prisma_client_or_raise_exception() existing_user = await _check_user_exists(user_id) - prev_active = _scim_active_value(existing_user.metadata) + prev_active = _user_scim_active(existing_user) update_data, final_team_set = _apply_patch_ops( existing_user=existing_user, patch_ops=patch_ops, ) - new_active = _scim_active_value(update_data.get("metadata")) + patched_metadata = update_data.get("metadata") + new_active = _scim_active_value(patched_metadata if isinstance(patched_metadata, Mapping) else None) # Handle team membership changes await _handle_team_membership_changes( @@ -1875,7 +1978,7 @@ async def patch_user( update_data["metadata"] = safe_dumps(update_data["metadata"]) - updated_user = await UserRepository(prisma_client).table.update( + updated_user = await _table(UserRepository(prisma_client)).update( where={"user_id": user_id}, data=update_data, ) @@ -1891,6 +1994,12 @@ async def patch_user( raise handle_exception_on_proxy(e) +class _TeamWhereConditions(TypedDict, total=False): + """The team columns SCIM GET /Groups can filter on, as Prisma where-conditions.""" + + team_alias: str + + # Group Endpoints @scim_router.get( "/Groups", @@ -1915,7 +2024,7 @@ async def get_groups( try: prisma_client = await _get_prisma_client_or_raise_exception() # Parse filter if provided (basic support) - where_conditions = {} + where_conditions: _TeamWhereConditions = {} if filter: # Very basic filter support - only handling displayName eq if "displayName eq" in filter: @@ -1923,7 +2032,7 @@ async def get_groups( where_conditions["team_alias"] = team_alias # Get teams from database - teams = await TeamRepository(prisma_client).table.find_many( + teams = await _table(TeamRepository(prisma_client)).find_many( where=where_conditions, skip=(startIndex - 1), take=count, @@ -1931,10 +2040,10 @@ async def get_groups( ) # Get total count for pagination - total_count = await TeamRepository(prisma_client).table.count(where=where_conditions) + total_count = await _table(TeamRepository(prisma_client)).count(where=where_conditions) # Convert to SCIM format - scim_groups = [] + scim_groups: List[SCIMGroup] = [] for team in teams: # Get team members with display names. members_with_roles is the # source of truth; the legacy `members` column is not populated by @@ -2018,7 +2127,7 @@ async def create_group( team_id = group.id or group.externalId or str(uuid.uuid4()) # Check if team already exists - existing_team = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id}) + existing_team = await _table(TeamRepository(prisma_client)).find_unique(where={"team_id": team_id}) if existing_team: raise HTTPException( @@ -2091,7 +2200,7 @@ async def update_group( } # Update team in database - updated_team = await TeamRepository(prisma_client).table.update( + updated_team = await _table(TeamRepository(prisma_client)).update( where={"team_id": group_id}, data=update_data, ) @@ -2145,19 +2254,19 @@ async def delete_group( # For each member, remove this team from their teams list for member_id in member_ids: - user = await UserRepository(prisma_client).table.find_unique(where={"user_id": member_id}) + user = await _table(UserRepository(prisma_client)).find_unique(where={"user_id": member_id}) if user: current_teams = user.teams or [] if group_id in current_teams: new_teams = [t for t in current_teams if t != group_id] - await UserRepository(prisma_client).table.update( + await _table(UserRepository(prisma_client)).update( where={"user_id": member_id}, data={"teams": new_teams} ) await _recompute_scim_member_roles(prisma_client, member_ids) # Delete team - await TeamRepository(prisma_client).table.delete(where={"team_id": group_id}) + await _table(TeamRepository(prisma_client)).delete(where={"team_id": group_id}) return Response(status_code=204) @@ -2166,8 +2275,8 @@ async def delete_group( async def _process_group_patch_operations( - patch_ops: SCIMPatchOp, existing_team, prisma_client -) -> Tuple[Dict[str, Any], Set[str], Set[str] | None]: + patch_ops: SCIMPatchOp, existing_team: LiteLLM_TeamTable, prisma_client: PrismaClient +) -> Tuple[dict[str, object], Set[str], Set[str] | None]: """Process patch operations for a group and return update data, final members and, when the request contained a member ``replace`` op, the absolute target roster it declared (``None`` otherwise). @@ -2183,7 +2292,7 @@ async def _process_group_patch_operations( have admitted - the phantom users this endpoint used to create for nested groups - impossible to clean up. """ - update_data: Dict[str, Any] = {} + update_data: dict[str, object] = {} # Create a fresh copy of existing metadata to avoid Prisma issues metadata = {**(existing_team.metadata or {}), SCIM_MANAGED_TEAM_METADATA_KEY: True} @@ -2251,7 +2360,7 @@ async def _process_group_patch_operations( return update_data, final_members, replace_target -async def _apply_group_patch_updates(group_id: str, update_data: Dict[str, Any], prisma_client): +async def _apply_group_patch_updates(group_id: str, update_data: dict[str, object], prisma_client: PrismaClient): """Apply the group's metadata/displayName patch updates to the database. Membership itself is not written here; it is reconciled onto the source of @@ -2330,7 +2439,7 @@ async def patch_group( # Apply the metadata/displayName updates to the database updated_team = await _apply_group_patch_updates(group_id, update_data, prisma_client) - refreshed_team = await TeamRepository(prisma_client).table.find_unique(where={"team_id": group_id}) + refreshed_team = await _table(TeamRepository(prisma_client)).find_unique(where={"team_id": group_id}) refreshed_current = ( set( await _get_team_member_user_ids_from_team(LiteLLM_TeamTable.model_validate(refreshed_team.model_dump())) @@ -2356,7 +2465,7 @@ async def patch_group( ) # Refresh team one more time to get final state after membership changes - final_team = await TeamRepository(prisma_client).table.find_unique(where={"team_id": group_id}) + final_team = await _table(TeamRepository(prisma_client)).find_unique(where={"team_id": group_id}) if final_team: updated_team = final_team diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index fa01f43d049..95c89c24851 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -14,7 +14,19 @@ import json import math import traceback from datetime import datetime, timezone -from typing import Annotated, Any, Dict, List, Mapping, Optional, Tuple, Union, cast +from typing import ( + Annotated, + Dict, + List, + Mapping, + Optional, + Protocol, + Sequence, + Tuple, + TypeVar, + Union, + cast, +) import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Request, status @@ -30,11 +42,14 @@ from litellm.proxy._types import ( BlockTeamRequest, CommonProxyErrors, DeleteTeamRequest, + LiteLLM_AccessGroupTable, LiteLLM_AuditLogs, + LiteLLM_BudgetTableFull, LiteLLM_DeletedTeamTable, LiteLLM_ManagementEndpoint_MetadataFields, LiteLLM_ManagementEndpoint_MetadataFields_Premium, LiteLLM_ModelTable, + LiteLLM_OrganizationMembershipTable, LiteLLM_OrganizationTable, LiteLLM_OrganizationTableWithMembers, LiteLLM_TeamMembership, @@ -78,6 +93,7 @@ from litellm.proxy.auth.auth_checks import ( from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.callback_utils import encrypt_callback_vars from litellm.proxy.common_utils.json_merge_patch import apply_json_merge_patch +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.management_endpoints.common_utils import ( _check_passthrough_routes_caller_permission, _is_user_org_admin_for_team, @@ -106,7 +122,7 @@ from litellm.proxy.management_helpers.utils import ( add_new_member, management_endpoint_wrapper, ) -from litellm.proxy.utils import PrismaClient, handle_exception_on_proxy +from litellm.proxy.utils import PrismaClient, ProxyLogging, handle_exception_on_proxy from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.table_repositories import ( @@ -141,6 +157,127 @@ from litellm.types.proxy.management_endpoints.team_endpoints import ( router = APIRouter() +_DbRecordT = TypeVar("_DbRecordT") + + +class _PrismaTableActions(Protocol[_DbRecordT]): + async def find_unique( + self, + where: Mapping[str, object], + include: Mapping[str, bool] | None = None, + ) -> _DbRecordT | None: ... + + async def find_first( + self, + where: Mapping[str, object] | None = None, + order: Mapping[str, str] | None = None, + ) -> _DbRecordT | None: ... + + async def find_many( + self, + where: Mapping[str, object] | None = None, + include: Mapping[str, bool] | None = None, + order: Mapping[str, str] | None = None, + skip: int | None = None, + take: int | None = None, + cursor: Mapping[str, object] | None = None, + ) -> list[_DbRecordT]: ... + + async def create( + self, + data: Mapping[str, object], + include: Mapping[str, bool] | None = None, + ) -> _DbRecordT: ... + + 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, bool] | None = None, + ) -> _DbRecordT: ... + + async def update_many( + self, + where: Mapping[str, object], + data: Mapping[str, object], + ) -> int: ... + + async def upsert( + self, + where: Mapping[str, object], + data: Mapping[str, Mapping[str, object]], + ) -> _DbRecordT: ... + + 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_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_TeamTable]": + team_table: _PrismaTableActions[LiteLLM_TeamTable] = TeamRepository(prisma_client).table + return team_table + + +def _team_membership_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_TeamMembership]": + membership_table: _PrismaTableActions[LiteLLM_TeamMembership] = TeamMembershipRepository(prisma_client).table + return membership_table + + +def _user_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_UserTable]": + user_table: _PrismaTableActions[LiteLLM_UserTable] = UserRepository(prisma_client).table + return user_table + + +def _model_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_ModelTable]": + model_table: _PrismaTableActions[LiteLLM_ModelTable] = ModelTableRepository(prisma_client).table + return model_table + + +def _org_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_OrganizationTable]": + org_table: _PrismaTableActions[LiteLLM_OrganizationTable] = OrganizationRepository(prisma_client).table + return org_table + + +def _org_membership_db( + prisma_client: PrismaClient | None, +) -> "_PrismaTableActions[LiteLLM_OrganizationMembershipTable]": + org_membership_table: _PrismaTableActions[LiteLLM_OrganizationMembershipTable] = OrganizationMembershipRepository( + prisma_client + ).table + return org_membership_table + + +def _budget_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_BudgetTableFull]": + budget_table: _PrismaTableActions[LiteLLM_BudgetTableFull] = BudgetRepository(prisma_client).table + return budget_table + + +def _deleted_team_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_DeletedTeamTable]": + deleted_team_table: _PrismaTableActions[LiteLLM_DeletedTeamTable] = DeletedTeamRepository(prisma_client).table + return deleted_team_table + + +def _access_group_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_AccessGroupTable]": + access_group_table: _PrismaTableActions[LiteLLM_AccessGroupTable] = AccessGroupRepository(prisma_client).table + return access_group_table + + +def _tokens_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_VerificationToken]": + tokens_table: _PrismaTableActions[LiteLLM_VerificationToken] = VerificationTokenRepository(prisma_client).table + return tokens_table + def _sanitize_for_log(value: object) -> str: """Strip CR/LF from user-controlled values to prevent log injection.""" @@ -152,9 +289,9 @@ def _sanitize_for_log(value: object) -> str: async def _refresh_cached_team( - team_row: Any, - user_api_key_cache: Any, - proxy_logging_obj: Any, + team_row: LiteLLM_TeamTable, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, ) -> None: """ Refresh the in-memory cached team object after a DB write. @@ -396,7 +533,7 @@ class TeamMemberBudgetHandler: @staticmethod async def backfill_team_member_budget_entries( team_id: str, - members_with_roles: List[Union[Member, dict]], + members_with_roles: Sequence[Union[Member, dict[str, object]]], team_member_budget_id: str, prisma_client: PrismaClient, ) -> None: @@ -415,7 +552,7 @@ class TeamMemberBudgetHandler: return # Batch-fetch existing memberships for this team (avoids N+1 queries) - existing_memberships = await TeamMembershipRepository(prisma_client).table.find_many(where={"team_id": team_id}) + existing_memberships = await _team_membership_db(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. @@ -448,7 +585,7 @@ class TeamMemberBudgetHandler: # Heal existing membership rows that predate the team_member_budget # configuration: populate budget_id where it is currently NULL. # Rows with an explicit budget_id (per-member override) are left alone. - updated = await TeamMembershipRepository(prisma_client).table.update_many( + updated = await _team_membership_db(prisma_client).update_many( where={"team_id": team_id, "budget_id": None}, data={"budget_id": team_member_budget_id}, ) @@ -461,7 +598,7 @@ class TeamMemberBudgetHandler: ) -def _get_default_team_param(field: str) -> Any: +def _get_default_team_param(field: str) -> object: """ Returns a default value for the given field from litellm.default_team_params config. Returns None if no default is configured. @@ -504,7 +641,7 @@ 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_db(prisma_client).find_many( where=where_obj, include={"litellm_budget_table": True}, ) @@ -766,7 +903,7 @@ 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_db(prisma_client).find_many( where={"organization_id": org_table.organization_id}, ) @@ -791,7 +928,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. @@ -1052,7 +1189,7 @@ async def new_team( ) # Check if license is over limit - total_teams = await TeamRepository(prisma_client).table.count() + total_teams = await _team_db(prisma_client).count() if total_teams and _license_check.is_team_count_over_limit(team_count=total_teams): raise HTTPException( status_code=403, @@ -1154,7 +1291,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_db(prisma_client).create( {**litellm_modeltable.json(exclude_none=True)} # type: ignore ) # type: ignore @@ -1358,11 +1495,11 @@ async def _create_team_update_audit_log( async def _update_model_table( data: UpdateTeamRequest, - model_id: Optional[str], + model_id: Optional[int], prisma_client: PrismaClient, user_api_key_dict: UserAPIKeyAuth, litellm_proxy_admin_name: str, -) -> Optional[str]: +) -> Optional[int]: """ Upsert model table and return the model id """ @@ -1375,11 +1512,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_db(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_db(prisma_client).upsert( where={"id": model_id}, data={ "update": {**litellm_modeltable.json(exclude_none=True)}, # type: ignore @@ -1395,7 +1532,7 @@ async def _update_model_table( async def _auto_add_team_members_to_organization( team: LiteLLM_TeamTable, organization: LiteLLM_OrganizationTableWithMembers, - prisma_client: Any, + prisma_client: PrismaClient, ) -> None: """ When moving a team to an org, ensure all team members are also org members. @@ -1433,11 +1570,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_OrganizationTable: """ Fetch and validate an organization for team update operations. @@ -1456,7 +1593,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 _org_db(prisma_client).find_unique( where={"organization_id": organization_id}, include={"litellm_budget_table": True, "members": True, "teams": True}, ) @@ -1758,7 +1895,7 @@ async def update_team( ): # Is the caller org_admin of the destination org? caller_memberships = ( - await OrganizationMembershipRepository(prisma_client).table.find_many( + await _org_membership_db(prisma_client).find_many( where={ "user_id": user_api_key_dict.user_id, "organization_id": data.organization_id, @@ -2005,7 +2142,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_db(prisma_client).find_unique(where={"team_id": team_id}) if existing_team_row is None: raise HTTPException( status_code=404, @@ -2487,7 +2624,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_db(prisma_client).find_first( where={"user_email": {"equals": member.user_email, "mode": "insensitive"}} ) @@ -2516,7 +2653,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_db(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 @@ -2707,7 +2844,7 @@ 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_db(prisma_client).find_unique(where={"team_id": data.team_id}) if _existing_team_row is None: raise HTTPException( @@ -2745,7 +2882,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_db(prisma_client).update( where={ "team_id": data.team_id, }, @@ -2835,7 +2972,7 @@ _MEMBER_BUDGET_PATCH_FIELDS = { } -def _build_member_budget_patch(data: TeamMemberUpdateRequest) -> Dict[str, Any]: +def _build_member_budget_patch(data: TeamMemberUpdateRequest) -> Dict[str, object]: """Map the budget fields the request actually set (merge-patch: a sent value updates, an explicit null clears, an absent field is left untouched) to their budget-table columns.""" @@ -2911,7 +3048,7 @@ 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_db(prisma_client).find_unique(where={"team_id": data.team_id}) if _existing_team_row is None: raise HTTPException( @@ -3007,7 +3144,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_db(prisma_client).update( where={"team_id": data.team_id}, data={"members_with_roles": json.dumps(_db_team_members)}, # type: ignore ) @@ -3138,7 +3275,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_db(prisma_client).find_many(order={"created_at": "desc"}) data.members = [ Member( user_id=user.user_id, @@ -3254,9 +3391,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: Optional[BaseModel] = await _team_db(prisma_client).find_unique(where={"team_id": team_id}) if team_row_base is None: raise Exception except Exception: @@ -3379,7 +3514,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]]: +) -> List[Dict[str, object]]: """Transform teams into deleted team records ready for persistence.""" if not teams: return [] @@ -3424,13 +3559,13 @@ def _transform_teams_to_deleted_records( async def _save_deleted_team_records( - records: List[Dict[str, Any]], + records: List[Dict[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_db(prisma_client).create_many(data=records) async def _persist_deleted_team_records( @@ -3506,9 +3641,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_db(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( @@ -3518,7 +3651,7 @@ async def _add_team_member_budget_table( return team_info_response_object -async def _resolve_team_access_group_resources(_team_info: Any) -> None: +async def _resolve_team_access_group_resources(_team_info: TeamInfoResponseObjectTeamTable) -> None: """Populate access_group_models / mcp_server_ids / agent_ids on the team info response by resolving inherited resources from its access groups.""" if not _team_info.access_group_ids: @@ -3572,7 +3705,7 @@ async def team_info( ) try: - team_info: Optional[BaseModel] = await TeamRepository(prisma_client).table.find_unique( + team_info: Optional[BaseModel] = await _team_db(prisma_client).find_unique( where={"team_id": team_id}, include={"litellm_model_table": True, "object_permission": True}, ) @@ -3819,7 +3952,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_db(prisma_client).find_unique(where={"team_id": data.team_id}) if existing_team is None: raise HTTPException( status_code=404, @@ -3832,7 +3965,7 @@ async def block_team( user_api_key_dict=user_api_key_dict, ) - record = await TeamRepository(prisma_client).table.update( + record = await _team_db(prisma_client).update( where={"team_id": data.team_id}, data={"blocked": True}, # type: ignore ) @@ -3868,7 +4001,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_db(prisma_client).find_unique(where={"team_id": data.team_id}) if existing_team is None: raise HTTPException( status_code=404, @@ -3881,7 +4014,7 @@ async def unblock_team( user_api_key_dict=user_api_key_dict, ) - record = await TeamRepository(prisma_client).table.update( + record = await _team_db(prisma_client).update( where={"team_id": data.team_id}, data={"blocked": False}, # type: ignore ) @@ -3915,7 +4048,7 @@ 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_db(prisma_client).find_unique(where={"user_id": user_api_key_dict.user_id}) if user_info is None: raise HTTPException( status_code=404, @@ -3925,7 +4058,7 @@ async def list_available_teams( 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_db(prisma_client).find_many(where={"team_id": {"in": available_teams}}) available_teams_correct_type = [LiteLLM_TeamTable.model_validate(team.model_dump()) for team in available_teams_db] @@ -3934,9 +4067,9 @@ async def list_available_teams( 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. @@ -3976,16 +4109,16 @@ async def _build_team_list_where_conditions( search: Optional[str] = None, search_team_id_match: TeamIdSearchMatch = "exact", org_admin_org_ids: Optional[List[str]] = None, - user_api_key_cache: Optional[Any] = None, - proxy_logging_obj: Optional[Any] = None, -) -> Optional[Dict[str, Any]]: + user_api_key_cache: Optional[UserApiKeyCache] = None, + proxy_logging_obj: Optional[ProxyLogging] = None, +) -> Optional[Dict[str, object]]: """ Build where conditions for team list query. Returns None when the query is guaranteed to yield no results (e.g. user has no team memberships), allowing the caller to skip the DB round-trip. """ - where_conditions: Dict[str, Any] = {} + where_conditions: Dict[str, object] = {} if team_id: where_conditions["team_id"] = team_id @@ -4067,7 +4200,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_db(_prisma_client).find_many( where={"access_group_id": {"in": unique_ids}}, ) @@ -4115,8 +4248,8 @@ def _convert_teams_to_response_models( async def _get_keys_count_by_team( - prisma_client: Any, - teams: list, + prisma_client: PrismaClient, + teams: Sequence[LiteLLM_TeamTable], ) -> Dict[str, int]: """Aggregate virtual-key counts per team for the given page of teams. @@ -4140,9 +4273,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. @@ -4341,23 +4474,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_db(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_db(prisma_client).count(where=where_conditions) else: - teams = await TeamRepository(prisma_client).table.find_many( + teams = await _team_db(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_db(prisma_client).count(where=where_conditions) # Calculate total pages total_pages = -(-total_count // page_size) # Ceiling division @@ -4400,9 +4533,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. @@ -4574,7 +4707,7 @@ 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_db(prisma_client).count() # Get paginated teams teams = await TeamRepository(prisma_client).table.find_many( @@ -4710,7 +4843,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_db(prisma_client).find_unique(where={"team_id": data.team_id}) if team_row is None: raise HTTPException( @@ -4756,7 +4889,7 @@ async def team_model_add( # the writer and lets Prisma bump updated_at. # `include` mirrors the relations the auth path consumes off the cached # team object so that `_refresh_cached_team` doesn't null them out. - updated_team = await TeamRepository(prisma_client).table.update( + updated_team = await _team_db(prisma_client).update( where={"team_id": data.team_id}, data={"updated_at": datetime.now(timezone.utc)}, include={"object_permission": True}, # type: ignore @@ -4810,7 +4943,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_db(prisma_client).find_unique(where={"team_id": data.team_id}) if team_row is None: raise HTTPException( @@ -4973,7 +5106,7 @@ async def update_team_member_permissions( }, ) # Update the team member permissions - updated_team = await TeamRepository(prisma_client).table.update( + updated_team = await _team_db(prisma_client).update( where={"team_id": data.team_id}, data={"team_member_permissions": data.team_member_permissions}, ) @@ -5043,7 +5176,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: Sequence[LiteLLM_TeamTable], permissions_to_add: set) -> int: """Compute merged permissions and batch-write updates. Returns count of teams updated.""" updates = [] for team in teams: @@ -5065,9 +5198,11 @@ async def _compute_and_batch_updates(prisma_client, teams, permissions_to_add: s return len(updates) -async def _append_permissions_to_specific_teams(prisma_client, team_ids: List[str], permissions_to_add: set) -> int: +async def _append_permissions_to_specific_teams( + prisma_client: PrismaClient, team_ids: List[str], permissions_to_add: set +) -> int: """Fetch specific teams by ID and append permissions.""" - teams = await TeamRepository(prisma_client).table.find_many( + teams = await _team_db(prisma_client).find_many( where={"team_id": {"in": team_ids}}, ) @@ -5082,7 +5217,7 @@ async def _append_permissions_to_specific_teams(prisma_client, team_ids: List[st return await _compute_and_batch_updates(prisma_client, teams, permissions_to_add) -async def _append_permissions_to_all_teams(prisma_client, permissions_to_add: set) -> int: +async def _append_permissions_to_all_teams(prisma_client: PrismaClient, permissions_to_add: set) -> int: """Paginated read + batched write across all teams.""" teams_updated = 0 cursor = None @@ -5228,9 +5363,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( - where={"user_id": user_api_key_dict.user_id} - ) + user_keys = await _tokens_db(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] # If user has no API keys, return empty result if not user_api_keys: diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index f7e429ddb31..bc37058689e 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -11,7 +11,10 @@ from typing import ( Literal, Mapping, NamedTuple, + Protocol, Sequence, + TypedDict, + TypeVar, Union, ) @@ -48,6 +51,204 @@ router = APIRouter() SPEND_LOGS_PAGINATION_COUNT_CAP = 10000 +_RowT = TypeVar("_RowT") + + +class _SupportsModelDump(Protocol): + def model_dump(self) -> Mapping[str, object]: ... + + +class _SpendLogOwnershipRow(Protocol): + user: str | None + team_id: str | None + + +class _ActivityRow(TypedDict): + date: str + api_requests: int + total_tokens: int + + +class _ActivityModelRow(TypedDict): + model_group: str + date: str + api_requests: int + total_tokens: int + + +class _DeploymentExceptionsRow(TypedDict): + api_base: str + date: str + num_rate_limit_exceptions: int + + +class _ExceptionsRow(TypedDict): + date: str + num_rate_limit_exceptions: int + + +class _ModelIdSpendRow(TypedDict): + model_id: str + spend: float + + +class _TagNameRow(TypedDict): + individual_request_tag: str + + +class _TeamSpendRow(TypedDict): + team_alias: str | None + total_spend: float + + +class _TagSpendRow(TypedDict): + individual_request_tag: str + total_spend: float + + +class _SpendLogsCountRow(TypedDict): + total_count: int + + +class _PgClassRow(TypedDict): + relname: str + relkind: str + + +class _TotalSpendRow(TypedDict): + total_spend: float + + +class _TeamDailySpendRow(TypedDict): + team_alias: str | None + spend_date: str | None + total_spend: float + + +class _EndUserRow(TypedDict): + end_user: str | None + + +class _DailyTagSpendRow(TypedDict): + individual_request_tag: str + log_count: int + total_spend: float + + +class _SessionCountAggregate(TypedDict): + session_id: int + + +class _SessionCountRow(TypedDict): + session_id: str + _count: _SessionCountAggregate + + +class _SessionSpendRow(TypedDict): + session_id: str + session_total_spend: float + mcp_tool_call_count: int + mcp_tool_call_spend: float + + +async def _query_raw(prisma_client: PrismaClient, sql_query: str, *args: object) -> Sequence[_RowT]: + """Run a raw read query and return its rows as the row type the caller declares.""" + return await prisma_client.db.query_raw(sql_query, *args) + + +async def _query_raw_or_none(prisma_client: PrismaClient, sql_query: str, *args: object) -> Sequence[_RowT] | None: + """``_query_raw`` for the call sites that guard the result against ``None``.""" + return await _query_raw(prisma_client, sql_query, *args) + + +class _SpendLogsTable(Protocol): + """The subset of the Prisma spend-logs table API this module uses.""" + + async def find_many( + self, *, where: Mapping[str, object], order: Mapping[str, str] + ) -> Sequence[_SupportsModelDump]: ... + + async def find_unique( + self, *, where: Mapping[str, object], include: None = None + ) -> _SpendLogOwnershipRow | None: ... + + async def count(self, *, where: Mapping[str, object]) -> int: ... + + async def group_by( + self, *, by: Sequence[str], where: Mapping[str, object], count: Mapping[str, bool] + ) -> Sequence[_SessionCountRow]: ... + + +class _TeamTable(Protocol): + """The subset of the Prisma team table API this module uses.""" + + async def find_unique(self, *, where: Mapping[str, object]) -> _SupportsModelDump | None: ... + + async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_SupportsModelDump]: ... + + async def update_many(self, *, data: Mapping[str, float], where: Mapping[str, object]) -> int: ... + + +class _VerificationTokenTable(Protocol): + """The subset of the Prisma verification token table API this module uses.""" + + async def update_many(self, *, data: Mapping[str, float], where: Mapping[str, object]) -> int: ... + + +def _spend_logs_table(prisma_client: PrismaClient) -> _SpendLogsTable: + return SpendLogsRepository(prisma_client).table + + +def _team_table(prisma_client: PrismaClient) -> _TeamTable: + return TeamRepository(prisma_client).table + + +def _verification_token_table(prisma_client: PrismaClient) -> _VerificationTokenTable: + return VerificationTokenRepository(prisma_client).table + + +async def _find_spend_logs( + prisma_client: PrismaClient, + where: Mapping[str, object], + order: Mapping[str, str], +) -> Sequence[_SupportsModelDump]: + """Read spend log rows as Prisma model instances.""" + return await _spend_logs_table(prisma_client).find_many(where=where, order=order) + + +async def _find_spend_log_row(prisma_client: PrismaClient, request_id: str) -> _SpendLogOwnershipRow | None: + """Read the single spend log row identified by ``request_id``.""" + return await _spend_logs_table(prisma_client).find_unique( + where={"request_id": request_id}, + include=None, + ) + + +async def _count_spend_logs(prisma_client: PrismaClient, where: Mapping[str, object]) -> int: + """Count the spend log rows matching ``where``.""" + return await _spend_logs_table(prisma_client).count(where=where) + + +async def _count_logs_per_session( + prisma_client: PrismaClient, session_ids: Sequence[str | None] +) -> Sequence[_SessionCountRow]: + """Count spend log rows per session for the given session ids.""" + return await _spend_logs_table(prisma_client).group_by( + by=["session_id"], + where={"session_id": {"in": session_ids}}, + count={"session_id": True}, + ) + + +async def _find_team_row(prisma_client: PrismaClient, team_id: str) -> _SupportsModelDump | None: + """Read a single team row as a Prisma model instance.""" + return await _team_table(prisma_client).find_unique(where={"team_id": team_id}) + + +async def _find_team_rows(prisma_client: PrismaClient, team_ids: Sequence[str]) -> Sequence[_SupportsModelDump]: + """Read team rows as Prisma model instances.""" + return await _team_table(prisma_client).find_many(where={"team_id": {"in": team_ids}}) + @router.get( "/spend/keys", @@ -281,7 +482,9 @@ async def get_global_activity_internal_user( AND "user" = $3 GROUP BY date_trunc('day', "startTime") """ - db_response = await prisma_client.db.query_raw(sql_query, start_date, end_date, user_id) + db_response: Sequence[_ActivityRow] | None = await _query_raw_or_none( + prisma_client, sql_query, start_date, end_date, user_id + ) return db_response @@ -345,6 +548,7 @@ async def get_global_activity( "Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" ) + db_response: Sequence[_ActivityRow] | None if ( user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER or user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY @@ -361,7 +565,7 @@ async def get_global_activity( AND "startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC') GROUP BY date_trunc('day', "startTime") """ - db_response = await prisma_client.db.query_raw(sql_query, start_date_obj, end_date_obj) + db_response = await _query_raw_or_none(prisma_client, sql_query, start_date_obj, end_date_obj) if db_response is None: return [] @@ -420,7 +624,9 @@ async def get_global_activity_model_internal_user( AND "user" = $3 GROUP BY model_group, date_trunc('day', "startTime") """ - db_response = await prisma_client.db.query_raw(sql_query, start_date, end_date, user_id) + db_response: Sequence[_ActivityModelRow] | None = await _query_raw_or_none( + prisma_client, sql_query, start_date, end_date, user_id + ) return db_response @@ -507,6 +713,7 @@ async def get_global_activity_model( "Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" ) + db_response: Sequence[_ActivityModelRow] | None if ( user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER or user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY @@ -524,7 +731,7 @@ async def get_global_activity_model( AND "startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC') GROUP BY model_group, date_trunc('day', "startTime") """ - db_response = await prisma_client.db.query_raw(sql_query, start_date_obj, end_date_obj) + db_response = await _query_raw_or_none(prisma_client, sql_query, start_date_obj, end_date_obj) if db_response is None: return [] @@ -672,7 +879,9 @@ async def get_global_activity_exceptions_per_deployment( ORDER BY date; """ - db_response = await prisma_client.db.query_raw(sql_query, start_date_obj, end_date_obj, model_group) + db_response: Sequence[_DeploymentExceptionsRow] | None = await _query_raw_or_none( + prisma_client, sql_query, start_date_obj, end_date_obj, model_group + ) if db_response is None: return [] @@ -795,7 +1004,9 @@ async def get_global_activity_exceptions( ORDER BY date; """ - db_response = await prisma_client.db.query_raw(sql_query, start_date_obj, end_date_obj, model_group) + db_response: Sequence[_ExceptionsRow] | None = await _query_raw_or_none( + prisma_client, sql_query, start_date_obj, end_date_obj, model_group + ) if db_response is None: return [] @@ -883,6 +1094,7 @@ async def get_global_spend_provider( "Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" ) + db_response: Sequence[_ModelIdSpendRow] | None if ( user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER or user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY @@ -902,7 +1114,7 @@ async def get_global_spend_provider( AND "user" = $3 GROUP BY model_id """ - db_response = await prisma_client.db.query_raw(sql_query, start_date_obj, end_date_obj, user_id) + db_response = await _query_raw_or_none(prisma_client, sql_query, start_date_obj, end_date_obj, user_id) else: sql_query = """ SELECT @@ -914,7 +1126,7 @@ async def get_global_spend_provider( AND length(model_id) > 0 GROUP BY model_id """ - db_response = await prisma_client.db.query_raw(sql_query, start_date_obj, end_date_obj) + db_response = await _query_raw_or_none(prisma_client, sql_query, start_date_obj, end_date_obj) if db_response is None: return [] @@ -1042,6 +1254,7 @@ async def get_global_spend_report( if premium_user is not True: verbose_proxy_logger.debug("accessing /spend/report but not a premium user") raise ValueError("/spend/report endpoint " + CommonProxyErrors.not_premium_user.value) + db_response: Sequence[Mapping[str, object]] | None if api_key is not None: verbose_proxy_logger.debug("Getting /spend for api_key: [set=%s]", api_key is not None) if api_key.startswith("sk-"): @@ -1082,7 +1295,7 @@ async def get_global_spend_report( ORDER BY total_cost DESC; """ - db_response = await prisma_client.db.query_raw(sql_query, start_date_obj, end_date_obj, api_key) + db_response = await _query_raw_or_none(prisma_client, sql_query, start_date_obj, end_date_obj, api_key) if db_response is None: return [] @@ -1125,7 +1338,9 @@ async def get_global_spend_report( ORDER BY total_cost DESC; """ - db_response = await prisma_client.db.query_raw(sql_query, start_date_obj, end_date_obj, internal_user_id) + db_response = await _query_raw_or_none( + prisma_client, sql_query, start_date_obj, end_date_obj, internal_user_id + ) if db_response is None: return [] @@ -1190,7 +1405,7 @@ async def get_global_spend_report( group_by_day; """ - db_response = await prisma_client.db.query_raw(sql_query, start_date_obj, end_date_obj) + db_response = await _query_raw_or_none(prisma_client, sql_query, start_date_obj, end_date_obj) if db_response is None: return [] @@ -1231,7 +1446,7 @@ async def get_global_spend_report( ORDER BY total_cost DESC; """ - db_response = await prisma_client.db.query_raw(sql_query, start_date_obj, end_date_obj) + db_response = await _query_raw_or_none(prisma_client, sql_query, start_date_obj, end_date_obj) if db_response is None: return [] @@ -1268,7 +1483,7 @@ async def global_get_all_tag_names(): FROM "LiteLLM_SpendLogs"; """ - db_response = await prisma_client.db.query_raw(sql_query) + db_response: Sequence[_TagNameRow] | None = await _query_raw_or_none(prisma_client, sql_query) if db_response is None: return [] @@ -1415,7 +1630,9 @@ async def _get_spend_report_for_time_range( ORDER BY total_spend DESC; """ - response = await prisma_client.db.query_raw(sql_query, start_date_obj, end_date_obj) + response: Sequence[_TeamSpendRow] | None = await _query_raw_or_none( + prisma_client, sql_query, start_date_obj, end_date_obj + ) # get spend per tag for today sql_query = """ @@ -1429,7 +1646,9 @@ async def _get_spend_report_for_time_range( ORDER BY total_spend DESC; """ - spend_per_tag = await prisma_client.db.query_raw(sql_query, start_date_obj, end_date_obj) + spend_per_tag: Sequence[_TagSpendRow] | None = await _query_raw_or_none( + prisma_client, sql_query, start_date_obj, end_date_obj + ) return response, spend_per_tag except Exception as e: @@ -1894,7 +2113,7 @@ async def ui_view_spend_logs( # (messages, response, proxy_server_request can be hundreds of KB per row). # These are only needed in the detail endpoint /spend/logs/ui/{request_id}. sql_conditions: List[str] = [] - sql_params: List[Any] = [] + sql_params: list[object] = [] p = 1 # parameter index counter # Date range. Wrap the param side with `AT TIME ZONE 'UTC'` so comparison @@ -2002,7 +2221,9 @@ async def ui_view_spend_logs( LIMIT ${p} ) AS bounded_matches """ - count_rows = await prisma_client.db.query_raw(count_query, *sql_params, SPEND_LOGS_PAGINATION_COUNT_CAP + 1) + count_rows: Sequence[_SpendLogsCountRow] | None = await _query_raw_or_none( + prisma_client, count_query, *sql_params, SPEND_LOGS_PAGINATION_COUNT_CAP + 1 + ) raw_total = int(count_rows[0]["total_count"]) if count_rows else 0 total_is_capped = raw_total > SPEND_LOGS_PAGINATION_COUNT_CAP total_records = SPEND_LOGS_PAGINATION_COUNT_CAP if total_is_capped else raw_total @@ -2067,7 +2288,7 @@ def _spend_log_field_has_content(value: Union[str, list, dict] | None) -> bool: return True -def _hydrate_spend_log_metadata(rows: Sequence[Any]) -> None: +def _hydrate_spend_log_metadata(rows: Sequence[Mapping[str, object]]) -> None: """Re-hydrate the JSONB ``metadata`` column returned by ``query_raw`` as a string. The Prisma serialiser bypasses the model-layer JSON hydration we get on the ORM @@ -2227,7 +2448,9 @@ async def ui_view_request_response_for_request_id( WHERE request_id = $1 LIMIT 1 """ - db_result = await prisma_client.db.query_raw(sql_query, request_id) + db_result: Sequence[Mapping[str, object]] | None = await _query_raw_or_none( + prisma_client, sql_query, request_id + ) if db_result and len(db_result) > 0: resolved = await _resolve_request_response_payload(db_result[0], cold_storage_handler=ColdStorageHandler()) return resolved._asdict() @@ -2359,11 +2582,10 @@ async def view_spend_logs( # Check if user wants unsummarized data if not summarize: # Return filtered individual log entries (similar to UI endpoint) - data = await SpendLogsRepository(prisma_client).table.find_many( - where=filter_query, # type: ignore - order={ - "startTime": "desc", - }, + data = await _find_spend_logs( + prisma_client, + where=filter_query, + order={"startTime": "desc"}, ) return data @@ -2421,7 +2643,7 @@ async def view_spend_logs( return response else: - scoped_filter: Dict[str, Any] = {} + scoped_filter: dict[str, str] = {} if api_key is not None and isinstance(api_key, str): if api_key.startswith("sk-"): hashed_token = prisma_client.hash_token(token=api_key) @@ -2437,8 +2659,9 @@ async def view_spend_logs( spend_logs = await prisma_client.get_data(table_name="spend", query_type="find_all") return spend_logs - data = await SpendLogsRepository(prisma_client).table.find_many( - where=scoped_filter, # type: ignore + data = await _find_spend_logs( + prisma_client, + where=scoped_filter, order={"startTime": "desc"}, ) return data @@ -2489,8 +2712,8 @@ async def global_spend_reset(): code=status.HTTP_401_UNAUTHORIZED, ) - await VerificationTokenRepository(prisma_client).table.update_many(data={"spend": 0.0}, where={}) - await TeamRepository(prisma_client).table.update_many(data={"spend": 0.0}, where={}) + await _verification_token_table(prisma_client).update_many(data={"spend": 0.0}, where={}) + await _team_table(prisma_client).update_many(data={"spend": 0.0}, where={}) return { "message": "Spend for all API Keys and Teams reset successfully", @@ -2533,7 +2756,7 @@ async def global_spend_refresh(): WHERE relname = 'MonthlyGlobalSpend'; """ try: - resp = await prisma_client.db.query_raw(sql_query) + resp: Sequence[_PgClassRow] = await _query_raw(prisma_client, sql_query) return resp[0]["relkind"] == "m" except Exception: @@ -2562,7 +2785,7 @@ async def global_spend_refresh(): }, ) await new_client.db.connect() - await new_client.db.query_raw(sql_query) + await _query_raw(new_client, sql_query) verbose_proxy_logger.info("MonthlyGlobalSpend view refreshed") return { "message": "MonthlyGlobalSpend view refreshed", @@ -2601,13 +2824,13 @@ async def global_spend_for_internal_user( ORDER BY "date"; """ - response = await prisma_client.db.query_raw(sql_query, api_key, user_id) + response: Sequence[Mapping[str, object]] = await _query_raw(prisma_client, sql_query, api_key, user_id) return response sql_query = """SELECT * FROM "MonthlyGlobalSpendPerUserPerKey" WHERE "user" = $1 ORDER BY "date";""" - response = await prisma_client.db.query_raw(sql_query, user_id) + response = await _query_raw(prisma_client, sql_query, user_id) return response except Exception as e: @@ -2652,6 +2875,7 @@ async def global_spend_logs( code=status.HTTP_500_INTERNAL_SERVER_ERROR, ) + response: Sequence[Mapping[str, object]] if ( user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER or user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY @@ -2669,7 +2893,7 @@ async def global_spend_logs( if api_key is None: sql_query = """SELECT * FROM "MonthlyGlobalSpend" ORDER BY "date";""" - response = await prisma_client.db.query_raw(query=sql_query) + response = await _query_raw(prisma_client, sql_query) return response else: @@ -2679,7 +2903,7 @@ async def global_spend_logs( ORDER BY "date"; """ - response = await prisma_client.db.query_raw(sql_query, api_key) + response = await _query_raw(prisma_client, sql_query, api_key) return response @@ -2726,7 +2950,7 @@ async def global_spend(): if prisma_client is None: raise HTTPException(status_code=500, detail={"error": "No db connected"}) sql_query = """SELECT SUM(spend) as total_spend FROM "MonthlyGlobalSpend";""" - response = await prisma_client.db.query_raw(query=sql_query) + response: Sequence[_TotalSpendRow] | None = await _query_raw_or_none(prisma_client, sql_query) if response is not None: if isinstance(response, list) and len(response) > 0: total_spend = response[0].get("total_spend", 0.0) @@ -2791,7 +3015,7 @@ async def global_spend_key_internal_user(user_api_key_dict: UserAPIKeyAuth, limi """ - response = await prisma_client.db.query_raw(sql_query, user_id, limit) + response: Sequence[Mapping[str, object]] = await _query_raw(prisma_client, sql_query, user_id, limit) return response @@ -2816,6 +3040,7 @@ async def global_spend_keys( """ from litellm.proxy.proxy_server import prisma_client + response: Sequence[Mapping[str, object]] if ( user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER or user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY @@ -2828,14 +3053,14 @@ async def global_spend_keys( sql_query = """SELECT * FROM "Last30dKeysBySpend";""" if limit is None: - response = await prisma_client.db.query_raw(sql_query) + response = await _query_raw(prisma_client, sql_query) return response try: limit = int(limit) if limit < 1: raise ValueError("Limit must be greater than 0") sql_query = """SELECT * FROM "Last30dKeysBySpend" LIMIT $1 ;""" - response = await prisma_client.db.query_raw(sql_query, limit) + response = await _query_raw(prisma_client, sql_query, limit) except ValueError as e: raise HTTPException(status_code=422, detail={"error": f"Invalid limit: {limit}, error: {e}"}) from e @@ -2875,7 +3100,7 @@ async def global_spend_per_team(): ORDER BY spend_date; """ - response = await prisma_client.db.query_raw(query=sql_query) + response: Sequence[_TeamDailySpendRow] = await _query_raw(prisma_client, sql_query) # transform the response for the Admin UI spend_by_date = {} @@ -2952,7 +3177,7 @@ async def global_view_all_end_users(): SELECT DISTINCT end_user FROM "LiteLLM_SpendLogs" """ - db_response = await prisma_client.db.query_raw(query=sql_query) + db_response: Sequence[_EndUserRow] | None = await _query_raw_or_none(prisma_client, sql_query) if db_response is None: return [] @@ -3009,7 +3234,9 @@ GROUP BY end_user ORDER BY total_spend DESC LIMIT 100 """ - response = await prisma_client.db.query_raw(sql_query, startTime, endTime, selected_api_key) + response: Sequence[Mapping[str, object]] = await _query_raw( + prisma_client, sql_query, startTime, endTime, selected_api_key + ) return response @@ -3040,7 +3267,7 @@ async def global_spend_models_internal_user(user_api_key_dict: UserAPIKeyAuth, l LIMIT $2; """ - response = await prisma_client.db.query_raw(sql_query, user_id, limit) + response: Sequence[Mapping[str, object]] = await _query_raw(prisma_client, sql_query, user_id, limit) return response @@ -3077,7 +3304,7 @@ async def global_spend_models( sql_query = """SELECT * FROM "Last30dModelsBySpend" LIMIT $1 ;""" - response = await prisma_client.db.query_raw(sql_query, int(limit)) + response: Sequence[Mapping[str, object]] = await _query_raw(prisma_client, sql_query, int(limit)) return response @@ -3169,14 +3396,17 @@ async def provider_budgets() -> ProviderBudgetResponse: async def get_spend_by_tags(prisma_client: PrismaClient, start_date=None, end_date=None): - response = await prisma_client.db.query_raw(""" + response: Sequence[Mapping[str, object]] = await _query_raw( + prisma_client, + """ SELECT jsonb_array_elements_text(request_tags) AS individual_request_tag, COUNT(*) AS log_count, SUM(spend) AS total_spend FROM "LiteLLM_SpendLogs" GROUP BY individual_request_tag; - """) + """, + ) return response @@ -3203,7 +3433,7 @@ async def ui_get_spend_by_tags( if prisma_client is None: raise HTTPException(status_code=500, detail={"error": "No db connected"}) - response = None + response: Sequence[_DailyTagSpendRow] | None = None if tags_list is None or (isinstance(tags_list, list) and "all-tags" in tags_list): # Get spend for all tags sql_query = """ @@ -3216,7 +3446,8 @@ async def ui_get_spend_by_tags( WHERE spend_date >= $1::date AND spend_date <= $2::date ORDER BY total_spend DESC; """ - response = await prisma_client.db.query_raw( + response = await _query_raw( + prisma_client, sql_query, start_date, end_date, @@ -3234,7 +3465,8 @@ async def ui_get_spend_by_tags( GROUP BY individual_request_tag ORDER BY total_spend DESC; """ - response = await prisma_client.db.query_raw( + response = await _query_raw( + prisma_client, sql_query, start_date, end_date, @@ -3353,7 +3585,7 @@ async def ui_view_session_spend_logs( skip = (page - 1) * page_size # Get total count for pagination metadata - total_records = await SpendLogsRepository(prisma_client).table.count(where=where_conditions) + total_records = await _count_spend_logs(prisma_client, where_conditions) # Query with raw SQL to exclude heavy columns (messages, response, proxy_server_request) sql_query = f""" @@ -3370,7 +3602,9 @@ async def ui_view_session_spend_logs( ORDER BY "startTime" DESC LIMIT $2 OFFSET $3 """ - result = await prisma_client.db.query_raw(sql_query, session_id, page_size, skip, *scope_params) + result: Sequence[Mapping[str, object]] = await _query_raw( + prisma_client, sql_query, session_id, page_size, skip, *scope_params + ) _hydrate_spend_log_metadata(result) total_pages = (total_records + page_size - 1) // page_size @@ -3434,7 +3668,7 @@ async def _build_ui_spend_logs_response( """ count_map: dict[str, int] = {} if enrich_session_counts: - session_ids = list( + session_ids: Sequence[str | None] = list( { (row.get("session_id") if isinstance(row, dict) else getattr(row, "session_id", None)) for row in data @@ -3446,11 +3680,7 @@ async def _build_ui_spend_logs_response( # is bounded by page_size (typically 25-50 distinct session IDs). # If performance degrades at scale, consider short-lived caching or # folding the count into the main query via a window function. - counts = await SpendLogsRepository(prisma_client).table.group_by( - by=["session_id"], - where={"session_id": {"in": session_ids}}, - count={"session_id": True}, - ) + counts = await _count_logs_per_session(prisma_client, session_ids) count_map = {r["session_id"]: r["_count"]["session_id"] for r in counts if r.get("session_id")} session_spend_map: dict[str, dict[str, Union[int, float]]] = {} @@ -3461,14 +3691,15 @@ async def _build_ui_spend_logs_response( # Collect api_keys already present in the authorized page rows so the # aggregate is scoped to the same ownership as the main query — prevents # cross-tenant disclosure via a colliding session_id. - authorized_api_keys = list( + authorized_api_keys: Sequence[str | None] = list( { (row.get("api_key") if isinstance(row, dict) else getattr(row, "api_key", None)) for row in data if (row.get("api_key") if isinstance(row, dict) else getattr(row, "api_key", None)) } ) - rows = await prisma_client.db.query_raw( + rows: Sequence[_SessionSpendRow] = await _query_raw( + prisma_client, """ SELECT session_id, COALESCE(SUM(spend), 0)::double precision AS session_total_spend, @@ -3531,7 +3762,7 @@ async def _build_ui_spend_logs_response( } -def _build_status_filter_condition(status_filter: str | None) -> Dict[str, Any]: +def _build_status_filter_condition(status_filter: str | None) -> Mapping[str, object]: """ Helper function to build the status filter condition for database queries. @@ -3539,7 +3770,7 @@ def _build_status_filter_condition(status_filter: str | None) -> Dict[str, Any]: status_filter (Optional[str]): The status to filter by. Can be "success" or "failure". Returns: - Dict[str, Any]: A dictionary containing the status filter condition. + Mapping[str, object]: A mapping containing the status filter condition. """ if status_filter is None: return {} @@ -3568,7 +3799,7 @@ def _is_admin_view_safe(user_api_key_dict: UserAPIKeyAuth) -> bool: async def _can_team_member_view_log( - prisma_client, + prisma_client: PrismaClient, user_api_key_dict: UserAPIKeyAuth, team_id: str | None, ) -> bool: @@ -3584,7 +3815,7 @@ async def _can_team_member_view_log( if team_id is None: return False - team_row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id}) + team_row = await _find_team_row(prisma_client, team_id) if team_row is None: return False team_obj = LiteLLM_TeamTable.model_validate(team_row.model_dump()) @@ -3614,7 +3845,7 @@ def _can_user_view_spend_log(user_api_key_dict: UserAPIKeyAuth) -> bool: async def _assert_user_can_view_request_id( - prisma_client, + prisma_client: PrismaClient, user_api_key_dict: UserAPIKeyAuth, request_id: str, ) -> None: @@ -3624,10 +3855,7 @@ async def _assert_user_can_view_request_id( permitted teams (admin or ``/spend/logs`` permission). Raises HTTP 403 if not. """ - row = await SpendLogsRepository(prisma_client).table.find_unique( - where={"request_id": request_id}, - include=None, - ) + row = await _find_spend_log_row(prisma_client, request_id) if row is None: return @@ -3650,7 +3878,7 @@ async def _assert_user_can_view_request_id( async def _get_permitted_team_ids_for_spend_logs( - prisma_client, + prisma_client: PrismaClient, user_api_key_dict: UserAPIKeyAuth, ) -> List[str]: """ @@ -3675,7 +3903,7 @@ async def _get_permitted_team_ids_for_spend_logs( if user_obj is None or not user_obj.teams: return [] - team_rows = await TeamRepository(prisma_client).table.find_many(where={"team_id": {"in": user_obj.teams}}) + team_rows = await _find_team_rows(prisma_client, user_obj.teams) permitted: List[str] = [] for team_row in team_rows: diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 9d30e40cd54..02c4712828d 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -1,5 +1,6 @@ import re import traceback +from collections.abc import Mapping, Sequence from datetime import datetime from typing import ( TYPE_CHECKING, @@ -22,7 +23,11 @@ from litellm.proxy._experimental.mcp_server.utils import ( ) from litellm.responses.main import aresponses from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator -from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.types.llms.openai import ( + ResponseInputParam, + ResponsesAPIResponse, + ResponsesAPIStreamingResponse, +) from litellm.types.utils import ( CallTypes, Choices, @@ -32,8 +37,10 @@ from litellm.types.utils import ( from litellm.utils import Rules, function_setup if TYPE_CHECKING: + from mcp.types import CallToolResult from mcp.types import Tool as MCPTool + from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.utils import ProxyLogging else: MCPTool = Any @@ -94,7 +101,7 @@ class LiteLLM_Proxy_MCP_Handler: @staticmethod def _parse_mcp_tools( - tools: Optional[Iterable[ToolParam]], + tools: Iterable[Mapping[str, object]] | None, ) -> Tuple[List[ToolParam], List[Any]]: """ Parse tools and separate MCP tools with litellm_proxy from other tools. @@ -134,8 +141,8 @@ class LiteLLM_Proxy_MCP_Handler: async def _apply_toolset_permissions( resolved_toolset_ids: List[str], resolved_mcp_servers: List[str], - user_api_key_auth: Any, - ) -> Any: + user_api_key_auth: "UserAPIKeyAuth", + ) -> "UserAPIKeyAuth": """Apply resolved toolset permissions to user_api_key_auth and return updated auth.""" from litellm.proxy._types import LiteLLM_ObjectPermissionTable @@ -174,8 +181,8 @@ class LiteLLM_Proxy_MCP_Handler: @staticmethod async def _get_mcp_tools_from_manager( - user_api_key_auth: Any, - mcp_tools_with_litellm_proxy: Optional[Iterable[ToolParam]], + user_api_key_auth: "UserAPIKeyAuth | None", + mcp_tools_with_litellm_proxy: Iterable[Mapping[str, object]] | None, litellm_trace_id: Optional[str] = None, mcp_auth_header: Optional[str] = None, mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, @@ -330,7 +337,7 @@ class LiteLLM_Proxy_MCP_Handler: @staticmethod def _filter_mcp_tools_by_allowed_tools( - mcp_tools: List[MCPTool], mcp_tools_with_litellm_proxy: List[ToolParam] + mcp_tools: List[MCPTool], mcp_tools_with_litellm_proxy: Sequence[Mapping[str, object]] ) -> List[MCPTool]: """Filter MCP tools based on allowed_tools parameter from the original tool configs.""" # Collect all allowed tool names from all MCP tool configs @@ -368,8 +375,8 @@ class LiteLLM_Proxy_MCP_Handler: @staticmethod async def _process_mcp_tools_to_openai_format( - user_api_key_auth: Any, - mcp_tools_with_litellm_proxy: List[ToolParam], + user_api_key_auth: "UserAPIKeyAuth | None", + mcp_tools_with_litellm_proxy: Sequence[Mapping[str, object]], litellm_trace_id: Optional[str] = None, request_tags: Optional[list[str]] = None, ) -> tuple[List[Any], dict[str, str]]: @@ -402,12 +409,12 @@ class LiteLLM_Proxy_MCP_Handler: @staticmethod async def _process_mcp_tools_without_openai_transform( user_api_key_auth: Any, - mcp_tools_with_litellm_proxy: List[ToolParam], + mcp_tools_with_litellm_proxy: Sequence[Mapping[str, object]], litellm_trace_id: Optional[str] = None, mcp_auth_header: Optional[str] = None, mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, request_tags: Optional[list[str]] = None, - ) -> tuple[List[Any], dict[str, str]]: + ) -> tuple[List[MCPTool], dict[str, str]]: """ Process MCP tools through filtering and deduplication pipeline without OpenAI transformation. This is useful for cases where we need the original MCP tool objects (e.g., for events). @@ -453,7 +460,7 @@ class LiteLLM_Proxy_MCP_Handler: @staticmethod def _transform_mcp_tools_to_openai( - mcp_tools: List[Any], + mcp_tools: Sequence[MCPTool], target_format: Literal["responses", "chat"] = "responses", ) -> List[Any]: """Transform MCP tools to OpenAI-compatible format.""" @@ -464,7 +471,6 @@ class LiteLLM_Proxy_MCP_Handler: openai_tools: List[Any] = [] for mcp_tool in mcp_tools: - openai_tool: Any if target_format == "chat": openai_tool = transform_mcp_tool_to_openai_tool(mcp_tool) else: @@ -475,7 +481,7 @@ class LiteLLM_Proxy_MCP_Handler: @staticmethod def _should_auto_execute_tools( - mcp_tools_with_litellm_proxy: Union[List[Dict[str, Any]], List[ToolParam]], + mcp_tools_with_litellm_proxy: Sequence[Mapping[str, object]], ) -> bool: """Check if we should auto-execute tool calls. @@ -514,9 +520,9 @@ class LiteLLM_Proxy_MCP_Handler: return tool_calls @staticmethod - def _extract_tool_calls_from_chat_response(response: ModelResponse) -> List[Any]: + def _extract_tool_calls_from_chat_response(response: ModelResponse) -> list[object]: """Extract tool calls from a chat completion response.""" - tool_calls: List[Any] = [] + tool_calls: list[object] = [] try: for choice in response.choices: @@ -583,7 +589,7 @@ class LiteLLM_Proxy_MCP_Handler: return tool_arguments or {} @staticmethod - def _parse_mcp_result(result: Any) -> str: + def _parse_mcp_result(result: "CallToolResult") -> str: """Parse MCP tool call result and extract meaningful content.""" if not result or not hasattr(result, "content") or not result.content: return "Tool executed successfully" @@ -626,7 +632,7 @@ class LiteLLM_Proxy_MCP_Handler: @staticmethod async def _execute_tool_calls( tool_server_map: dict[str, str], - tool_calls: List[Any], + tool_calls: Sequence[object], user_api_key_auth: Any, mcp_auth_header: Optional[str] = None, mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, @@ -908,7 +914,7 @@ class LiteLLM_Proxy_MCP_Handler: def _create_follow_up_messages_for_chat( original_messages: List[Any], response: ModelResponse, - tool_results: List[Dict[str, Any]], + tool_results: Sequence[Mapping[str, object]], ) -> List[Any]: """Create follow-up chat messages that include tool execution results.""" from copy import deepcopy @@ -952,8 +958,8 @@ class LiteLLM_Proxy_MCP_Handler: @staticmethod def _create_follow_up_input( response: ResponsesAPIResponse, - tool_results: List[Dict[str, Any]], - original_input: Any = None, + tool_results: Sequence[Mapping[str, object]], + original_input: str | ResponseInputParam | None = None, ) -> List[Any]: """Create follow-up input with tool results in proper format.""" follow_up_input: List[Any] = [] @@ -1049,7 +1055,7 @@ class LiteLLM_Proxy_MCP_Handler: *, proxy_logging_obj: Optional["ProxyLogging"], user_api_key_auth: Any, - request_data: Dict[str, Any], + request_data: dict[str, object], error: Exception, ) -> None: """Log MCP tool failures via proxy logging hooks.""" @@ -1071,11 +1077,11 @@ class LiteLLM_Proxy_MCP_Handler: @staticmethod def _create_mcp_streaming_response( - input: Union[str, Any], + input: str | ResponseInputParam, model: str, - all_tools: Optional[List[Any]], - mcp_tools_with_litellm_proxy: List[Any], - mcp_discovery_events: List[Any], + all_tools: Sequence[object] | None, + mcp_tools_with_litellm_proxy: list[Mapping[str, object]], + mcp_discovery_events: list[ResponsesAPIStreamingResponse], call_params: Dict[str, Any], previous_response_id: Optional[str], tool_server_map: dict[str, str], @@ -1116,9 +1122,9 @@ class LiteLLM_Proxy_MCP_Handler: @staticmethod def _build_request_params( - input: Union[str, Any], + input: str | ResponseInputParam, model: str, - all_tools: Optional[List[Any]], + all_tools: Sequence[object] | None, call_params: Dict[str, Any], previous_response_id: Optional[str], **kwargs, @@ -1149,7 +1155,9 @@ class LiteLLM_Proxy_MCP_Handler: return request_params @staticmethod - def _create_tool_execution_events(tool_calls: List[Any], tool_results: List[Dict[str, Any]]) -> List[Any]: + def _create_tool_execution_events( + tool_calls: Sequence[object], tool_results: List[Dict[str, Any]] + ) -> list[ResponsesAPIStreamingResponse]: """ Create MCP tool execution events for streaming. @@ -1163,7 +1171,7 @@ class LiteLLM_Proxy_MCP_Handler: from litellm._uuid import uuid from litellm.responses.mcp.mcp_streaming_iterator import create_mcp_call_events - tool_execution_events: List[Any] = [] + tool_execution_events: list[ResponsesAPIStreamingResponse] = [] # Create events for each tool execution for tool_result in tool_results: @@ -1233,8 +1241,8 @@ class LiteLLM_Proxy_MCP_Handler: @staticmethod def _add_mcp_output_elements_to_response( response: ResponsesAPIResponse, - mcp_tools_fetched: List[Any], - tool_results: List[Dict[str, Any]], + mcp_tools_fetched: Sequence[object], + tool_results: Sequence[Mapping[str, object]], ) -> ResponsesAPIResponse: """Add custom output elements to the final response for MCP tool execution.""" # Import the required classes for creating output items diff --git a/litellm/videos/main.py b/litellm/videos/main.py index 6d81fec36b3..cc3a15f9cd2 100644 --- a/litellm/videos/main.py +++ b/litellm/videos/main.py @@ -2,7 +2,7 @@ import asyncio import contextvars import json from functools import partial -from typing import Any, Coroutine, Dict, List, Literal, Optional, Union, overload +from typing import Coroutine, Dict, List, Literal, Optional, Union, overload import litellm from litellm.constants import DEFAULT_VIDEO_ENDPOINT_MODEL @@ -40,9 +40,9 @@ async def avideo_generation( custom_llm_provider=None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_query: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, **kwargs, ) -> VideoObject: """ @@ -126,13 +126,13 @@ def video_generation( user: Optional[str] = None, timeout: int = 600, custom_llm_provider: Optional[str] = None, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_query: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, *, avideo_generation: Literal[True], - **kwargs: Any, -) -> Coroutine[Any, Any, VideoObject]: + **kwargs: object, +) -> Coroutine[object, object, VideoObject]: ... @@ -146,12 +146,12 @@ def video_generation( user: Optional[str] = None, timeout: int = 600, custom_llm_provider: Optional[str] = None, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_query: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, *, avideo_generation: Literal[False] = False, - **kwargs: Any, + **kwargs: object, ) -> VideoObject: ... @@ -170,13 +170,13 @@ def video_generation( custom_llm_provider=None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_query: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, **kwargs, ) -> Union[ VideoObject, - Coroutine[Any, Any, VideoObject], + Coroutine[object, object, VideoObject], ]: """ Maps the https://api.openai.com/v1/videos endpoint. @@ -277,13 +277,13 @@ def video_content( variant: Optional[str] = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_query: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, **kwargs, ) -> Union[ bytes, - Coroutine[Any, Any, bytes], + Coroutine[object, object, bytes], ]: """ Download video content from OpenAI's video API. @@ -390,9 +390,9 @@ async def avideo_content( variant: Optional[str] = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_query: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, **kwargs, ) -> bytes: """ @@ -461,9 +461,9 @@ async def avideo_remix( custom_llm_provider=None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_query: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, **kwargs, ) -> VideoObject: """ @@ -528,13 +528,13 @@ def video_remix( prompt: str, timeout: int = 600, custom_llm_provider: Optional[str] = None, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_query: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, *, avideo_remix: Literal[True], - **kwargs: Any, -) -> Coroutine[Any, Any, VideoObject]: + **kwargs: object, +) -> Coroutine[object, object, VideoObject]: ... @@ -544,12 +544,12 @@ def video_remix( prompt: str, timeout: int = 600, custom_llm_provider: Optional[str] = None, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_query: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, *, avideo_remix: Literal[False] = False, - **kwargs: Any, + **kwargs: object, ) -> VideoObject: ... @@ -564,13 +564,13 @@ def video_remix( custom_llm_provider=None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_query: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, **kwargs, ) -> Union[ VideoObject, - Coroutine[Any, Any, VideoObject], + Coroutine[object, object, VideoObject], ]: """ Maps the https://api.openai.com/v1/videos/{video_id}/remix endpoint. @@ -668,9 +668,9 @@ async def avideo_list( custom_llm_provider=None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_query: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, **kwargs, ) -> List[VideoObject]: """ @@ -744,13 +744,13 @@ def video_list( order: Optional[str] = None, timeout: int = 600, custom_llm_provider: Optional[str] = None, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_query: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, *, avideo_list: Literal[True], - **kwargs: Any, -) -> Coroutine[Any, Any, List[VideoObject]]: + **kwargs: object, +) -> Coroutine[object, object, List[VideoObject]]: ... @@ -761,12 +761,12 @@ def video_list( order: Optional[str] = None, timeout: int = 600, custom_llm_provider: Optional[str] = None, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_query: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, *, avideo_list: Literal[False] = False, - **kwargs: Any, + **kwargs: object, ) -> List[VideoObject]: ... @@ -782,13 +782,13 @@ def video_list( custom_llm_provider=None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_query: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, **kwargs, ) -> Union[ List[VideoObject], - Coroutine[Any, Any, List[VideoObject]], + Coroutine[object, object, List[VideoObject]], ]: """ Maps the https://api.openai.com/v1/videos endpoint. @@ -882,9 +882,9 @@ async def avideo_status( custom_llm_provider=None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_query: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, **kwargs, ) -> VideoObject: """ @@ -947,13 +947,13 @@ def video_status( video_id: str, timeout: int = 600, custom_llm_provider: Optional[str] = None, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_query: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, *, avideo_status: Literal[True], - **kwargs: Any, -) -> Coroutine[Any, Any, VideoObject]: + **kwargs: object, +) -> Coroutine[object, object, VideoObject]: ... # Overload for when avideo_status=False (returns VideoObject) @@ -962,12 +962,12 @@ def video_status( video_id: str, timeout: int = 600, custom_llm_provider: Optional[str] = None, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_query: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, *, avideo_status: Literal[False] = False, - **kwargs: Any, + **kwargs: object, ) -> VideoObject: ... @@ -981,13 +981,13 @@ def video_status( custom_llm_provider=None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_query: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, **kwargs, ) -> Union[ VideoObject, - Coroutine[Any, Any, VideoObject], + Coroutine[object, object, VideoObject], ]: """ Retrieve video status from OpenAI's video API. @@ -1097,12 +1097,12 @@ def video_status( @client async def avideo_create_character( name: str, - video: Any, + video: FileTypes, timeout=600, custom_llm_provider=None, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_query: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, **kwargs, ) -> CharacterObject: """ @@ -1152,14 +1152,14 @@ async def avideo_create_character( @client def video_create_character( name: str, - video: Any, + video: FileTypes, timeout=600, custom_llm_provider=None, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_query: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, **kwargs, -) -> Union[CharacterObject, Coroutine[Any, Any, CharacterObject]]: +) -> Union[CharacterObject, Coroutine[object, object, CharacterObject]]: """ Create a character from an uploaded video file. Maps to POST /v1/videos/characters @@ -1230,9 +1230,9 @@ async def avideo_get_character( character_id: str, timeout=600, custom_llm_provider=None, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_query: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, **kwargs, ) -> CharacterObject: """ @@ -1280,11 +1280,11 @@ def video_get_character( character_id: str, timeout=600, custom_llm_provider=None, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_query: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, **kwargs, -) -> Union[CharacterObject, Coroutine[Any, Any, CharacterObject]]: +) -> Union[CharacterObject, Coroutine[object, object, CharacterObject]]: """ Retrieve a character by ID. Maps to GET /v1/videos/characters/{character_id} @@ -1355,9 +1355,9 @@ async def avideo_edit( prompt: str, timeout=600, custom_llm_provider=None, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_query: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, **kwargs, ) -> VideoObject: """ @@ -1407,11 +1407,11 @@ def video_edit( prompt: str, timeout=600, custom_llm_provider=None, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_query: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, **kwargs, -) -> Union[VideoObject, Coroutine[Any, Any, VideoObject]]: +) -> Union[VideoObject, Coroutine[object, object, VideoObject]]: """ Create a video edit job. Maps to POST /v1/videos/edits @@ -1486,9 +1486,9 @@ async def avideo_extension( seconds: str, timeout=600, custom_llm_provider=None, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_query: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, **kwargs, ) -> VideoObject: """ @@ -1540,11 +1540,11 @@ def video_extension( seconds: str, timeout=600, custom_llm_provider=None, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, object]] = None, + extra_query: Optional[Dict[str, object]] = None, + extra_body: Optional[Dict[str, object]] = None, **kwargs, -) -> Union[VideoObject, Coroutine[Any, Any, VideoObject]]: +) -> Union[VideoObject, Coroutine[object, object, VideoObject]]: """ Create a video extension. Maps to POST /v1/videos/extensions diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 6fb3ed748b6..b8650eea7aa 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,6 +1,6 @@ { "ANN001": { - "limit": 3118 + "limit": 3104 }, "ANN002": { "limit": 69 @@ -9,10 +9,10 @@ "limit": 831 }, "ANN201": { - "limit": 2138 + "limit": 2137 }, "ANN202": { - "limit": 944 + "limit": 941 }, "ANN204": { "limit": 724 @@ -24,7 +24,7 @@ "limit": 130 }, "ANN401": { - "limit": 2009 + "limit": 1851 }, "ASYNC230": { "limit": 14 @@ -123,7 +123,7 @@ "limit": 52 }, "I001": { - "limit": 270 + "limit": 261 }, "LOG015": { "limit": 8 @@ -222,7 +222,7 @@ "limit": 38 }, "RET504": { - "limit": 716 + "limit": 702 }, "RUF010": { "limit": 874 @@ -306,7 +306,7 @@ "limit": 9 }, "TID251": { - "limit": 2652 + "limit": 2649 }, "TRY002": { "limit": 547 @@ -324,7 +324,7 @@ "limit": 879 }, "UP006": { - "limit": 12135 + "limit": 12050 }, "UP007": { "limit": 2526 @@ -360,9 +360,9 @@ "limit": 4 }, "UP037": { - "limit": 105 + "limit": 104 }, "UP045": { - "limit": 17805 + "limit": 17793 } } diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 05499e83c42..ff037a2872e 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 23250 + "limit": 23191 }, "LIT002": { - "limit": 27277 + "limit": 27276 }, "LIT003": { "limit": 292 @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1108 + "limit": 1106 }, "LIT007": { "limit": 0 @@ -24,6 +24,6 @@ "limit": 1004 }, "LIT009": { - "limit": 2473 + "limit": 2467 } }