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

Merge pull request #35452 from BerriAI/litellm_decrease_anys_fable2

Browse source 
chore(typing): clear 2.4k basedpyright errors across 15 Any hotspot files
This commit is contained in:
Mateo Wang 2026-08-01 09:56:46 -07:00 • committed by GitHub
commit 2b3070890a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
18 changed files with 1903 additions and 1042 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
basedpyright-code-budget.json
View file

@ -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

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

@ -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,

134
litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py
View file

@ -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(

558
litellm/llms/custom_httpx/llm_http_handler.py
View file

File diff suppressed because it is too large Load diff

61
litellm/llms/openai/videos/transformation.py
View file

@ -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

188
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

@ -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

209
litellm/proxy/_experimental/mcp_server/sampling_handler.py
View file

@ -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(

96
litellm/proxy/guardrails/guardrail_endpoints.py
View file

@ -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:

159
litellm/proxy/management_endpoints/internal_user_endpoints.py
View file

@ -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}

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

@ -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.",

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

@ -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

255
litellm/proxy/management_endpoints/scim/scim_v2.py
View file

@ -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

297
litellm/proxy/management_endpoints/team_endpoints.py
View file

@ -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:

368
litellm/proxy/spend_tracking/spend_management_endpoints.py
View file

@ -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:

72
litellm/responses/mcp/litellm_proxy_mcp_handler.py
View file

@ -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

204
litellm/videos/main.py
View file

@ -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

20
ruff-strict-budget.json
View file

@ -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
}
}

8
type-discipline-budget.json
View file

@ -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
}
}