scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 156

chore(typing): clear basedpyright Any errors in streaming, scim, and mcp sampling

Browse source 
Replaces Any-typed seams with real types (Pydantic model_validate,
TypedDicts, Protocols, domain-model conversions, precise return types)
across six files, cutting whole-tree reportAny from 21518 to 20786 and
reportExplicitAny from 7258 to 7115 with no regression in any other
basedpyright rule.

Per-file reportAny + reportExplicitAny:
- litellm_core_utils/streaming_handler.py: 407 -> 204
- responses/streaming_iterator.py: 304 -> 151
- proxy/management_endpoints/scim/scim_v2.py: 233 -> 2
- proxy/_experimental/mcp_server/sampling_handler.py: 226 -> 121
- integrations/websearch_interception/handler.py: 189 -> 85
- llms/anthropic/experimental_pass_through/context_management/editors/compact.py: 179 -> 98

llm_http_handler.py and proxy_server.py rank higher by raw count but
were left for a follow-up pass given their size (13k and 17k lines)
and blast radius as core request-handling hot paths.

Adds UP037 to ruff.toml's lint.external list so the base ruff config
stops flagging as unused the noqa: UP037 suppression on scim_v2.py's
delete_user teams annotation, which must stay quoted since
PrismaTeamTable is a TYPE_CHECKING-only import.
This commit is contained in:
mateo-berri 2026-07-31 11:49:12 +00:00
parent 3c2264cfac
commit 3bf65116d8
No known key found for this signature in database
11 changed files with 693 additions and 440 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

16
basedpyright-code-budget.json
View file

@ -1,6 +1,6 @@
{
"reportAny": {
"limit": 31903
"limit": 28975
},
"reportArgumentType": {
"limit": 2645
@ -24,7 +24,7 @@
"limit": 42
},
"reportExplicitAny": {
"limit": 10214
"limit": 9640
},
"reportFunctionMemberAccess": {
"limit": 11
@ -54,10 +54,10 @@
"limit": 0
},
"reportMissingParameterType": {
"limit": 5869
"limit": 5829
},
"reportMissingTypeArgument": {
"limit": 15861
"limit": 15839
},
"reportMissingTypeStubs": {
"limit": 41
@ -99,19 +99,19 @@
"limit": 0
},
"reportUnknownArgumentType": {
"limit": 45366
"limit": 45365
},
"reportUnknownLambdaType": {
"limit": 113
},
"reportUnknownMemberType": {
"limit": 40477
"limit": 40309
},
"reportUnknownParameterType": {
"limit": 20338
"limit": 20275
},
"reportUnknownVariableType": {
"limit": 32047
"limit": 31867
},
"reportUnnecessaryCast": {
"limit": 177

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

@ -9,7 +9,7 @@ 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 typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast
import litellm
from litellm._logging import verbose_logger
@ -39,8 +39,17 @@ from litellm.types.integrations.custom_logger import (
AgenticLoopRequestPatch,
)
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import CallTypes, LlmProviders
from litellm.utils import ProviderConfigManager
from litellm.types.router import SearchToolTypedDict
from litellm.types.utils import CallTypes, LlmProviders, ModelResponse
from litellm.utils import CustomStreamWrapper, 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.router import Router
# Key used to flag, on per-request kwargs, that the originating client sent
# an Anthropic-native ``web_search_*`` tool — meaning the final response
@ -94,8 +103,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: dict[str, object] | None = None,
) -> dict[str, object] | None:
"""
Short-circuit web-search-only requests by executing the search directly.
@ -188,7 +197,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 +219,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 +237,7 @@ 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: CallTypes | None) -> dict | None:
"""
Pre-call hook to convert native Anthropic web_search tools to regular tools.
@ -297,7 +306,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: dict[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
@ -362,15 +371,17 @@ class WebSearchInterceptionLogger(CustomLogger):
)
@staticmethod
def _tool_name(tool: dict[str, Any]) -> Optional[str]:
def _tool_name(tool: dict[str, object]) -> Optional[str]:
"""Effective tool name, handling OpenAI ``function`` wrapper shape."""
fn = tool.get("function")
if tool.get("type") == "function" and isinstance(fn, dict):
return fn.get("name")
return tool.get("name")
name = fn.get("name")
return name if isinstance(name, str) else None
name = tool.get("name")
return name if isinstance(name, str) else None
@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: object, 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 +479,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 +589,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 +647,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 +698,13 @@ class WebSearchInterceptionLogger(CustomLogger):
tools: Dict,
model: str,
messages: List[Dict],
response: Any,
anthropic_messages_provider_config: Any,
response: object,
anthropic_messages_provider_config: Optional["BaseAnthropicMessagesConfig"],
anthropic_messages_optional_request_params: Dict,
logging_obj: Any,
logging_obj: "LiteLLMLoggingObj",
stream: bool,
kwargs: Dict,
) -> Any:
) -> object:
"""
Execute agentic loop with WebSearch execution for Anthropic Messages API.
@ -721,10 +732,10 @@ class WebSearchInterceptionLogger(CustomLogger):
tools: Dict,
model: str,
messages: List[Dict],
response: Any,
anthropic_messages_provider_config: Any,
response: object,
anthropic_messages_provider_config: Optional["BaseAnthropicMessagesConfig"],
anthropic_messages_optional_request_params: Dict,
logging_obj: Any,
logging_obj: "LiteLLMLoggingObj",
stream: bool,
kwargs: Dict,
) -> AgenticLoopPlan:
@ -764,7 +775,7 @@ class WebSearchInterceptionLogger(CustomLogger):
kwargs=kwargs,
)
metadata: Dict[str, Any] = {
metadata: Dict[str, object] = {
"tool_type": "websearch",
"response_format": "anthropic",
}
@ -787,10 +798,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 +821,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 +836,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: object, native_blocks: List[Dict[str, object]]) -> object:
"""Prepend native blocks to response content, dict or object form."""
if not native_blocks:
return response
@ -835,7 +846,7 @@ class WebSearchInterceptionLogger(CustomLogger):
return response
existing = getattr(response, "content", None) or []
try:
response.content = list(native_blocks) + list(existing)
setattr(response, "content", list(native_blocks) + list(existing))
except (AttributeError, TypeError):
# Object refused write — fall through and leave the response
# untouched rather than crash the request.
@ -849,12 +860,12 @@ class WebSearchInterceptionLogger(CustomLogger):
tools: Dict,
model: str,
messages: List[Dict],
response: Any,
response: object,
optional_params: Dict,
logging_obj: Any,
logging_obj: "LiteLLMLoggingObj",
stream: bool,
kwargs: Dict,
) -> Any:
) -> object:
"""
Execute agentic loop with WebSearch execution for Chat Completions API.
@ -884,9 +895,9 @@ class WebSearchInterceptionLogger(CustomLogger):
tools: Dict,
model: str,
messages: List[Dict],
response: Any,
response: object,
optional_params: Dict,
logging_obj: Any,
logging_obj: "LiteLLMLoggingObj",
stream: bool,
kwargs: Dict,
) -> AgenticLoopPlan:
@ -911,9 +922,9 @@ class WebSearchInterceptionLogger(CustomLogger):
tools: dict,
model: str,
messages: list[dict],
response: Any,
response: object,
optional_params: dict,
logging_obj: Any,
logging_obj: "LiteLLMLoggingObj",
stream: bool,
kwargs: dict,
) -> AgenticLoopPlan:
@ -1023,7 +1034,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 +1102,10 @@ class WebSearchInterceptionLogger(CustomLogger):
tool_calls: List[Dict],
thinking_blocks: List[Dict],
anthropic_messages_optional_request_params: Dict,
logging_obj: Any,
logging_obj: Optional["LiteLLMLoggingObj"],
stream: bool,
kwargs: Dict,
) -> Any:
) -> object:
"""Legacy path: execute search + build patch + run follow-up call."""
request_patch, structured_results = await self._build_anthropic_request_patch(
model=model,
@ -1145,7 +1156,7 @@ class WebSearchInterceptionLogger(CustomLogger):
tool_calls: List[Dict],
thinking_blocks: List[Dict],
anthropic_messages_optional_request_params: Dict,
logging_obj: Any,
logging_obj: Optional["LiteLLMLoggingObj"],
kwargs: Dict,
) -> Tuple[AgenticLoopRequestPatch, List[Optional[SearchResponse]]]:
"""
@ -1205,7 +1216,9 @@ class WebSearchInterceptionLogger(CustomLogger):
follow_up_messages = messages + [assistant_message, cast(Dict, user_message)]
# Correlation context for structured logging
_call_id = getattr(logging_obj, "litellm_call_id", None) or kwargs.get("litellm_call_id", "unknown")
_call_id = (logging_obj.litellm_call_id if logging_obj is not None else None) or kwargs.get(
"litellm_call_id", "unknown"
)
full_model_name = model # safe default before try block
@ -1238,7 +1251,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: dict[str, object] | None = None
) -> Tuple[str, Optional[SearchResponse]]:
"""
Execute a single web search using router's search tools.
@ -1263,10 +1276,10 @@ class WebSearchInterceptionLogger(CustomLogger):
search_tool = self._select_search_tool_from_router(llm_router=llm_router)
search_provider: Optional[str] = None
search_litellm_params: dict[str, Any] = {}
search_litellm_params: Dict[str, Any] = {}
if search_tool is not None:
await self._authorize_search_tool(search_tool=search_tool, kwargs=kwargs)
search_litellm_params = dict(search_tool.get("litellm_params", {}) or {})
search_litellm_params = dict(search_tool["litellm_params"])
search_provider = search_litellm_params.get("search_provider")
# Fallback to perplexity if no router or no search tools configured
@ -1300,11 +1313,11 @@ class WebSearchInterceptionLogger(CustomLogger):
async def _authorize_search_tool(
self,
search_tool: dict[str, Any],
kwargs: Optional[dict[str, Any]],
search_tool: SearchToolTypedDict,
kwargs: dict[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:
if not search_tool_name:
return
user_api_key_auth = self._get_user_api_key_auth_from_kwargs(kwargs)
@ -1322,7 +1335,7 @@ class WebSearchInterceptionLogger(CustomLogger):
valid_token=user_api_key_auth,
)
team_id = getattr(user_api_key_auth, "team_id", None)
team_id = user_api_key_auth.team_id
if team_id:
from litellm.proxy.proxy_server import (
prisma_client,
@ -1334,7 +1347,7 @@ class WebSearchInterceptionLogger(CustomLogger):
team_id=team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=getattr(user_api_key_auth, "parent_otel_span", None),
parent_otel_span=user_api_key_auth.parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
await can_team_call_search_tool(
@ -1343,7 +1356,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: dict[str, object] | None) -> Optional["UserAPIKeyAuth"]:
if not kwargs:
return None
@ -1363,21 +1376,20 @@ class WebSearchInterceptionLogger(CustomLogger):
return None
def _select_search_tool_from_router(self, llm_router: Any) -> Optional[dict[str, Any]]:
if llm_router is None or not hasattr(llm_router, "search_tools"):
def _select_search_tool_from_router(self, llm_router: Optional["Router"]) -> Optional[SearchToolTypedDict]:
if llm_router is None:
return None
search_tools = list(getattr(llm_router, "search_tools") or [])
return self._select_search_tool_from_list(search_tools=search_tools, source="router")
return self._select_search_tool_from_list(search_tools=list(llm_router.search_tools), source="router")
def _select_search_tool_from_list(
self,
search_tools: list[dict[str, Any]],
search_tools: list[SearchToolTypedDict],
source: str,
) -> Optional[dict[str, Any]]:
) -> Optional[SearchToolTypedDict]:
if self.search_tool_name:
matching_tools = [tool for tool in search_tools if tool.get("search_tool_name") == self.search_tool_name]
if matching_tools:
search_provider = (matching_tools[0].get("litellm_params", {}) or {}).get("search_provider")
search_provider = matching_tools[0]["litellm_params"].get("search_provider")
verbose_logger.debug(
f"WebSearchInterception: Found search tool '{self.search_tool_name}' "
f"from {source} with provider '{search_provider}'"
@ -1390,7 +1402,7 @@ class WebSearchInterceptionLogger(CustomLogger):
if search_tools:
first_tool = search_tools[0]
search_provider = (first_tool.get("litellm_params", {}) or {}).get("search_provider")
search_provider = first_tool["litellm_params"].get("search_provider")
verbose_logger.debug(
f"WebSearchInterception: Using first available search tool from {source} "
f"with provider '{search_provider}'"
@ -1405,11 +1417,11 @@ class WebSearchInterceptionLogger(CustomLogger):
messages: List[Dict],
tool_calls: List[Dict],
optional_params: Dict,
logging_obj: Any,
logging_obj: Optional["LiteLLMLoggingObj"],
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,

57
litellm/litellm_core_utils/streaming_handler.py
View file

@ -22,7 +22,7 @@ from typing import (
import anyio
import httpx
from pydantic import BaseModel
from pydantic import BaseModel, TypeAdapter, ValidationError
import litellm
from litellm import verbose_logger
@ -62,6 +62,21 @@ _SYNC_ITER_EXHAUSTED = object()
_GCHUNK_FIELDS: frozenset = frozenset(GChunk.__annotations__)
_STR_KEYED_DICT_ADAPTER: TypeAdapter[Dict[str, object]] = TypeAdapter(Dict[str, object])
def _as_str_keyed_dict(raw: object) -> Dict[str, object]:
"""
Best-effort coercion of a loosely-typed value (e.g. from a Dict[str, Any]
call site, or a test double that doesn't behave like a real dict) into a
plain string-keyed dict, matching the historical `**raw or {}`-style
tolerance of this code.
"""
try:
return _STR_KEYED_DICT_ADAPTER.validate_python(raw)
except ValidationError:
return {}
def _next_sync_or_exhausted(it: Any) -> Any:
"""
@ -77,7 +92,7 @@ def _next_sync_or_exhausted(it: Any) -> Any:
return _SYNC_ITER_EXHAUSTED
def is_async_iterable(obj: Any) -> bool:
def is_async_iterable(obj: object) -> bool:
"""
Check if an object is an async iterable (can be used with 'async for').
@ -105,7 +120,7 @@ class _ProviderChunkParsed:
@dataclass(frozen=True, slots=True)
class _ProviderChunkEarlyReturn:
value: Any
value: ModelResponseStream | None
_ProviderChunkResult = Union[_ProviderChunkParsed, _ProviderChunkEarlyReturn]
@ -116,12 +131,14 @@ class CustomStreamWrapper:
self,
completion_stream,
model,
logging_obj: Any,
logging_obj: Optional[LiteLLMLoggingObject],
custom_llm_provider: Optional[str] = None,
stream_options=None,
make_call: Optional[Callable] = None,
_response_headers: Optional[dict] = None,
):
if logging_obj is None:
raise ValueError("CustomStreamWrapper requires a logging_obj")
self.model = model
self.make_call = make_call
self.custom_llm_provider = custom_llm_provider
@ -131,9 +148,8 @@ class CustomStreamWrapper:
self.sent_last_chunk = False
self._stream_created_time: float = time.time()
litellm_params: GenericLiteLLMParams = GenericLiteLLMParams(
**self.logging_obj.model_call_details.get("litellm_params", {})
)
_litellm_params_dict = _as_str_keyed_dict(self.logging_obj.model_call_details.get("litellm_params", {}))
litellm_params: GenericLiteLLMParams = GenericLiteLLMParams.model_validate(_litellm_params_dict)
self.merge_reasoning_content_in_choices: bool = litellm_params.merge_reasoning_content_in_choices or False
self.sent_first_thinking_block = False
self.sent_last_thinking_block = False
@ -158,7 +174,7 @@ class CustomStreamWrapper:
_api_base = get_api_base(
model=model or "",
optional_params=self.logging_obj.model_call_details.get("litellm_params", {}),
optional_params=_litellm_params_dict,
)
self._hidden_params = {
@ -195,7 +211,7 @@ class CustomStreamWrapper:
# Snapshot assumes self._hidden_params is populated from litellm_params
# at init and never mutated during the stream. If that ever changes,
# this cache must be removed.
self._base_hidden_params: Dict[str, Any] = {
self._base_hidden_params: Dict[str, object] = {
**self._hidden_params,
"response_cost": None,
}
@ -246,14 +262,13 @@ class CustomStreamWrapper:
def check_send_stream_usage(self, stream_options: Optional[dict]):
return stream_options is not None and stream_options.get("include_usage", False) is True
def check_is_function_call(self, logging_obj) -> bool:
def check_is_function_call(self, logging_obj: LiteLLMLoggingObject) -> bool:
from litellm.litellm_core_utils.prompt_templates.common_utils import (
is_function_call,
)
if hasattr(logging_obj, "optional_params") and isinstance(logging_obj.optional_params, dict):
if is_function_call(logging_obj.optional_params):
return True
if hasattr(logging_obj, "optional_params"):
return is_function_call(logging_obj.optional_params)
return False
@ -670,14 +685,14 @@ class CustomStreamWrapper:
_logging_obj_llm_provider = self._cached_logging_llm_provider
if chunk is None:
args: Dict[str, Any] = {"model": _model}
args: Dict[str, object] = {"model": _model}
else:
chunk.pop("model", None)
args = {"model": _model}
if chunk:
args.update({k: v for k, v in chunk.items() if k != "stream"})
model_response = ModelResponseStream(**args)
model_response = ModelResponseStream.model_validate(args)
if self.response_id is not None:
model_response.id = self.response_id
if self.system_fingerprint is not None:
@ -817,7 +832,7 @@ class CustomStreamWrapper:
_initial_delta = model_response.choices[0].delta.model_dump()
_initial_delta.pop("role", None)
model_response.choices[0].delta = Delta(**_initial_delta)
model_response.choices[0].delta = Delta.model_validate(_initial_delta)
return model_response
def _has_special_delta_content(self, model_response: ModelResponseStream) -> bool:
@ -915,7 +930,7 @@ class CustomStreamWrapper:
choice_json.pop(
"finish_reason", None
) # for mistral etc. which return a value in their last chunk (not-openai compatible).
choices.append(StreamingChoices(**choice_json))
choices.append(StreamingChoices.model_validate(choice_json))
except Exception:
choices.append(StreamingChoices())
setattr(model_response, "choices", choices)
@ -946,7 +961,7 @@ class CustomStreamWrapper:
self.sent_first_chunk = True
if response_obj.get("provider_specific_fields") is not None:
completion_obj["provider_specific_fields"] = response_obj["provider_specific_fields"]
model_response.choices[0].delta = Delta(**completion_obj)
model_response.choices[0].delta = Delta.model_validate(completion_obj)
_index: Optional[int] = completion_obj.get("index")
if _index is not None:
model_response.choices[0].index = _index
@ -1443,7 +1458,7 @@ class CustomStreamWrapper:
):
# if function returned but type set to None - mistral's api returns type: None
tool["type"] = "function"
model_response.choices[0].delta = Delta(**_json_delta)
model_response.choices[0].delta = Delta.model_validate(_json_delta)
except Exception as e:
verbose_logger.exception(
"litellm.CustomStreamWrapper.chunk_creator(): Exception occured - {}".format(str(e))
@ -1458,7 +1473,7 @@ class CustomStreamWrapper:
if original_chunk.choices[0].delta is None
else dict(original_chunk.choices[0].delta)
)
model_response.choices[0].delta = Delta(**delta)
model_response.choices[0].delta = Delta.model_validate(delta)
except Exception:
model_response.choices[0].delta = Delta()
else:
@ -1672,7 +1687,7 @@ class CustomStreamWrapper:
else:
asyncio.run(self.logging_obj.async_success_handler(processed_chunk, None, None, cache_hit))
## SYNC LOGGING — only for sync SDK entrypoints; async proxy paths export via async_success_handler
litellm_params = self.logging_obj.model_call_details.get("litellm_params", {})
litellm_params = _as_str_keyed_dict(self.logging_obj.model_call_details.get("litellm_params", {}))
if self.logging_obj._is_sync_litellm_request(litellm_params):
self.logging_obj.success_handler(processed_chunk, None, None, cache_hit)

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

@ -13,7 +13,7 @@ Mirrors Anthropic's native ``compact_20260112`` for non-Anthropic providers:
"""
import re
from typing import Any, Dict, List, Literal, Optional, Tuple, Union, cast
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, TypedDict, Union, cast
import litellm
from litellm._logging import verbose_logger
@ -22,6 +22,7 @@ from litellm.types.llms.anthropic import (
CompactionBlock,
UsageIteration,
)
from litellm.types.llms.openai import ChatCompletionToolParam
from ..constants import (
COMPACT_DEFAULT_INSTRUCTIONS,
@ -38,6 +39,10 @@ from ..constants import (
from ..errors import AnthropicContextManagementError
from ..result import PolyfillResult
if TYPE_CHECKING:
from litellm.proxy._types import UserAPIKeyAuth
from litellm.router import Router
# Auth metadata fields propagated from the parent request to the summary call
# so the summary's spend is attributed to the same scopes. The list mirrors the
# fields populated by
@ -98,9 +103,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 +299,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 +362,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
@ -616,18 +621,18 @@ def _count_effective_tokens(
"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[Union[List[ChatCompletionToolParam], List[Dict[str, Any]]]] = 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)
openai_tools = translated_tools
except Exception as e:
verbose_logger.debug(
"compact_20260112: anthropic→openai tools translation failed "
@ -638,7 +643,7 @@ def _count_effective_tokens(
total = litellm.token_counter(
model=model,
messages=cast(Any, openai_shape),
messages=openai_shape,
tools=cast(Any, openai_tools),
)
if compaction_block is not None:
@ -761,7 +766,7 @@ 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]] = []
system_message = _system_to_openai_message(system)
@ -802,6 +807,11 @@ def _append_text_to_content(content: Any, extra_text: str) -> Any:
return [content, {"type": "text", "text": extra_text}]
class _OptionalSummaryCallKwargs(TypedDict, total=False):
user: str
allowed_model_region: str
async def _call_summary_model(
*,
summary_model: str,
@ -838,25 +848,33 @@ async def _call_summary_model(
# the parent ``/v1/messages`` request. On timeout the caller catches the
# exception and surfaces ``applied_edits[0].error = "summary_call_failed"``,
# forwarding the request without compaction rather than hanging.
call_kwargs: Dict[str, Any] = {
"model": summary_model,
"messages": summary_messages,
"max_tokens": max_tokens,
"timeout": COMPACT_SUMMARY_TIMEOUT_SECONDS,
"litellm_metadata": metadata,
}
# The end-user id must also travel as the top-level ``user`` kwarg: legacy
# limiter hooks and prometheus end-user tracking read it from there rather
# than from ``litellm_metadata``, so without it the summary tokens would not
# debit the caller's end-user counters.
end_user_id = metadata.get("user_api_key_end_user_id")
if end_user_id:
call_kwargs["user"] = end_user_id
raw_end_user_id = metadata.get("user_api_key_end_user_id")
optional_kwargs: _OptionalSummaryCallKwargs = {}
if isinstance(raw_end_user_id, str) and raw_end_user_id:
optional_kwargs["user"] = raw_end_user_id
if allowed_model_region is not None:
call_kwargs["allowed_model_region"] = allowed_model_region
optional_kwargs["allowed_model_region"] = allowed_model_region
if llm_router is not None and hasattr(llm_router, "acompletion"):
return await llm_router.acompletion(**call_kwargs)
return await litellm.acompletion(**call_kwargs)
return await llm_router.acompletion(
model=summary_model,
messages=summary_messages,
max_tokens=max_tokens,
timeout=COMPACT_SUMMARY_TIMEOUT_SECONDS,
litellm_metadata=metadata,
**optional_kwargs,
)
return await litellm.acompletion(
model=summary_model,
messages=summary_messages,
max_tokens=max_tokens,
timeout=COMPACT_SUMMARY_TIMEOUT_SECONDS,
litellm_metadata=metadata,
**optional_kwargs,
)
def _extract_response_text(response: Any) -> Optional[str]:
@ -941,8 +959,8 @@ async def apply_compact_20260112(
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,
llm_router: Optional["Router"] = None,
user_api_key_auth: Optional["UserAPIKeyAuth"] = None,
) -> PolyfillResult:
"""Apply ``compact_20260112``; return a ``PolyfillResult``.

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

@ -10,16 +10,21 @@ MCP Spec Reference:
https://modelcontextprotocol.io/specification/2025-11-25/client/sampling
"""
from typing import Any, Dict, List, Optional, Union
import typing
from typing import Any, Dict, List, Literal, Optional, Tuple, TypedDict, Union
from starlette.types import Scope
if typing.TYPE_CHECKING:
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.utils import ProxyLogging
from litellm.types.llms.openai import ChatCompletionToolParam
from litellm.types.utils import ModelResponse
from fastapi import HTTPException, Request
from litellm._logging import verbose_logger
from fastapi import HTTPException
# Guard imports that require the mcp package
try:
from mcp.types import (
@ -29,12 +34,15 @@ try:
ErrorData,
ModelPreferences,
SamplingMessage,
SamplingMessageContentBlock,
TextContent,
Tool,
ToolChoice,
ToolUseContent,
)
MCPSamplingContent = Union[SamplingMessageContentBlock, List[SamplingMessageContentBlock]]
MCP_SAMPLING_AVAILABLE = True
except ImportError as _sampling_import_err:
MCP_SAMPLING_AVAILABLE = False
@ -65,7 +73,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
@ -153,6 +161,13 @@ def _has_priorities(model_preferences: "ModelPreferences") -> bool:
)
class _ScoredModel(TypedDict):
name: str
cost: float
max_output: float
output_tps: float
def _select_model_by_priority(
model_names: List[str],
model_preferences: "ModelPreferences",
@ -188,7 +203,7 @@ def _select_model_by_priority(
intel_weight = 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)
@ -257,7 +272,7 @@ def _select_model_by_priority(
def _convert_mcp_content_to_openai(
content: Any,
content: "MCPSamplingContent",
) -> Union[str, Dict[str, Any], List[Dict[str, Any]]]:
"""
Convert MCP SamplingMessage content to OpenAI message content format.
@ -282,7 +297,7 @@ def _convert_mcp_content_to_openai(
def _convert_single_content(
content: Any,
content: "SamplingMessageContentBlock",
) -> Union[Dict[str, Any], List[Dict[str, Any]]]:
"""Convert a single MCP content item to OpenAI format.
@ -294,19 +309,14 @@ def _convert_single_content(
"""
import json
content_type = getattr(content, "type", None)
if content_type == "text":
if content.type == "text":
return {"type": "text", "text": content.text}
elif content_type == "image":
data = getattr(content, "data", "")
mime_type = getattr(content, "mimeType", "image/png")
elif content.type == "image":
return {
"type": "image_url",
"image_url": {"url": f"data:{mime_type};base64,{data}"},
"image_url": {"url": f"data:{content.mimeType};base64,{content.data}"},
}
elif content_type == "audio":
data = getattr(content, "data", "")
mime_type = getattr(content, "mimeType", "audio/wav")
elif content.type == "audio":
# Map MIME type to OpenAI audio format
format_map = {
"audio/wav": "wav",
@ -315,40 +325,35 @@ def _convert_single_content(
"audio/flac": "flac",
"audio/ogg": "ogg",
}
audio_format = format_map.get(mime_type, "wav")
audio_format = format_map.get(content.mimeType, "wav")
return {
"type": "input_audio",
"input_audio": {"data": data, "format": audio_format},
"input_audio": {"data": content.data, "format": audio_format},
}
elif content_type == "tool_use":
elif content.type == "tool_use":
# ToolUseContent → proper OpenAI function-call representation.
# The ``_marker_type`` key lets the message-level converter
# hoist this into the ``tool_calls`` array on the assistant
# message instead of embedding it inline as a content part.
return {
"_marker_type": "tool_use",
"id": getattr(content, "id", f"call_{id(content)}"),
"id": content.id,
"type": "function",
"function": {
"name": getattr(content, "name", ""),
"arguments": json.dumps(getattr(content, "input", {}), default=str),
"name": content.name,
"arguments": json.dumps(content.input, default=str),
},
}
elif content_type == "tool_result":
elif content.type == "tool_result":
# ToolResultContent → proper OpenAI tool-role message.
# 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", [])
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 ""
else:
result_text = str(nested_content)
text_parts = [c.text for c in content.content if c.type == "text"]
result_text = "\n".join(text_parts) if text_parts else ""
return {
"_marker_type": "tool_result",
"role": "tool",
"tool_call_id": tool_use_id,
"tool_call_id": content.toolUseId,
"content": result_text,
}
# Fallback: treat as text
@ -442,69 +447,60 @@ def _convert_mcp_messages_to_openai(
return openai_messages
def _has_tool_use(content: Any) -> bool:
def _has_tool_use(content: "MCPSamplingContent") -> 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"
return any(c.type == "tool_use" for c in content)
return content.type == "tool_use"
def _has_tool_result(content: Any) -> bool:
def _has_tool_result(content: "MCPSamplingContent") -> 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"
return any(c.type == "tool_result" for c in content)
return content.type == "tool_result"
def _extract_tool_calls(content: Any) -> List[Dict[str, Any]]:
def _extract_tool_calls(content: "MCPSamplingContent") -> List[Dict[str, Any]]:
"""Extract OpenAI-format tool_calls from MCP ToolUseContent."""
import json
items = content if isinstance(content, list) else [content]
tool_calls = []
for item in items:
if getattr(item, "type", None) == "tool_use":
if item.type == "tool_use":
tool_calls.append(
{
"id": getattr(item, "id", f"call_{id(item)}"),
"id": item.id,
"type": "function",
"function": {
"name": getattr(item, "name", ""),
"arguments": json.dumps(getattr(item, "input", {}), default=str),
"name": item.name,
"arguments": json.dumps(item.input, default=str),
},
}
)
return tool_calls
def _extract_text_parts(content: Any) -> Optional[str]:
def _extract_text_parts(content: "MCPSamplingContent") -> Optional[str]:
"""Extract text parts from mixed content."""
items = content if isinstance(content, list) else [content]
texts = []
for item in items:
if getattr(item, "type", None) == "text":
texts.append(getattr(item, "text", ""))
texts = [item.text for item in items if item.type == "text"]
return "\n".join(texts) if texts else None
def _extract_tool_results(content: Any) -> List[Dict[str, Any]]:
def _extract_tool_results(content: "MCPSamplingContent") -> List[Dict[str, Any]]:
"""Extract OpenAI-format tool messages from MCP ToolResultContent."""
items = content if isinstance(content, list) else [content]
results = []
for item in items:
if getattr(item, "type", None) == "tool_result":
tool_use_id = getattr(item, "toolUseId", "")
# Extract text from nested content
nested_content = 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 ""
else:
result_text = str(nested_content)
if item.type == "tool_result":
text_parts = [c.text for c in item.content if c.type == "text"]
result_text = "\n".join(text_parts) if text_parts else ""
results.append(
{
"role": "tool",
"tool_call_id": tool_use_id,
"tool_call_id": item.toolUseId,
"content": result_text,
}
)
@ -513,7 +509,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]]]:
) -> Optional[List["ChatCompletionToolParam"]]:
"""
Convert MCP Tool definitions to OpenAI function calling format.
MCP Tool: {name, description, inputSchema}
@ -521,9 +517,9 @@ def _convert_mcp_tools_to_openai(
"""
if not tools:
return None
openai_tools = []
openai_tools: List["ChatCompletionToolParam"] = []
for tool in tools:
openai_tool = {
openai_tool: "ChatCompletionToolParam" = {
"type": "function",
"function": {
"name": tool.name,
@ -541,7 +537,7 @@ def _convert_mcp_tools_to_openai(
def _convert_mcp_tool_choice_to_openai(
tool_choice: Optional["ToolChoice"],
) -> Optional[Union[str, Dict[str, Any]]]:
) -> Optional[Literal["auto", "required", "none"]]:
"""
Convert MCP ToolChoice to OpenAI tool_choice format.
MCP: {mode: "auto"} | {mode: "required"} | {mode: "none"}
@ -549,7 +545,7 @@ def _convert_mcp_tool_choice_to_openai(
"""
if not tool_choice:
return None
mode = getattr(tool_choice, "mode", "auto")
mode = tool_choice.mode or "auto"
if mode == "auto":
return "auto"
elif mode == "required":
@ -560,7 +556,7 @@ def _convert_mcp_tool_choice_to_openai(
def _convert_openai_response_to_mcp_result(
response: Any,
response: "ModelResponse",
model_name: str,
) -> Union["CreateMessageResult", "CreateMessageResultWithTools", "ErrorData"]:
"""
@ -586,19 +582,19 @@ def _convert_openai_response_to_mcp_result(
choice = response.choices[0]
message = choice.message
# Determine stop reason
finish_reason = getattr(choice, "finish_reason", "stop")
finish_reason = choice.finish_reason
if finish_reason == "tool_calls":
stop_reason = "toolUse"
elif finish_reason == "length":
stop_reason = "maxTokens"
else:
stop_reason = "endTurn"
actual_model = getattr(response, "model", model_name) or model_name
actual_model = response.model or model_name
# Check if response has tool calls
tool_calls = getattr(message, "tool_calls", None)
tool_calls = message.tool_calls
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))
@ -606,17 +602,15 @@ def _convert_openai_response_to_mcp_result(
for tc in tool_calls:
import json
tool_input = tc.function.arguments
if isinstance(tool_input, str):
try:
tool_input = json.loads(tool_input)
except (json.JSONDecodeError, TypeError):
tool_input = {"raw": tool_input}
try:
tool_input = json.loads(tc.function.arguments)
except (json.JSONDecodeError, TypeError):
tool_input = {"raw": tc.function.arguments}
content_parts.append(
ToolUseContent(
type="tool_use",
id=tc.id,
name=tc.function.name,
name=tc.function.name or "",
input=tool_input,
)
)
@ -636,7 +630,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: Optional["UserAPIKeyAuth"]) -> Optional["ErrorData"]:
"""Enforce model-permission checks for MCP sampling requests.
Runs the same authorization checks as ``/chat/completions``:
@ -649,9 +643,9 @@ async def _check_model_access(model: str, user_api_key_auth: Any) -> Optional["E
if user_api_key_auth is None:
return None
_api_key = getattr(user_api_key_auth, "api_key", None)
_token = getattr(user_api_key_auth, "token", None)
_user_role = getattr(user_api_key_auth, "user_role", None)
_api_key = user_api_key_auth.api_key
_token = user_api_key_auth.token
_user_role = user_api_key_auth.user_role
_has_real_credential = bool(_api_key) or bool(_token)
_is_admin = _user_role in ("proxy_admin", "proxy_admin_viewer") if _user_role else False
@ -678,14 +672,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:
@ -695,21 +689,25 @@ async def _check_model_access(model: str, user_api_key_auth: Any) -> Optional["E
await can_key_call_model(
model=model,
llm_model_list=getattr(litellm, "model_list", None),
llm_model_list=litellm.model_list,
valid_token=user_api_key_auth,
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 = user_api_key_auth.team_id
_user_id = user_api_key_auth.user_id
_project_id = user_api_key_auth.project_id
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]
@ -731,7 +729,7 @@ async def _check_model_access(model: str, user_api_key_auth: Any) -> Optional["E
model=model,
team_object=team_obj,
llm_router=_llm_router,
team_model_aliases=getattr(user_api_key_auth, "team_model_aliases", None),
team_model_aliases=user_api_key_auth.team_model_aliases,
)
if _user_id and _proxy_logging_obj:
await _check_team_member_model_access(
@ -799,7 +797,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 +809,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 = user_api_key_auth.team_id
_user_id = user_api_key_auth.user_id
team_obj = None
if _team_id and _prisma_client and _user_api_key_cache:
@ -935,7 +941,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
@ -958,10 +964,8 @@ def _build_sampling_request(
original headers don't already carry it, as a fallback for
IP attribution.
"""
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).
@ -1016,7 +1020,7 @@ def _build_sampling_request(
if client_ip:
_client_tuple = (client_ip, 0)
scope: Dict[str, Any] = {
scope: Scope = {
"type": "http",
"method": "POST",
"path": "/mcp/sampling/createMessage",
@ -1035,7 +1039,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]:
@ -1065,7 +1069,7 @@ async def _build_completion_kwargs(
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
from litellm.proxy.proxy_server import proxy_config
completion_kwargs["user"] = getattr(user_api_key_auth, "user_id", None)
completion_kwargs["user"] = user_api_key_auth.user_id
_dummy_request = _build_sampling_request(raw_headers=raw_headers, client_ip=client_ip)
completion_kwargs = await add_litellm_data_to_request(
data=completion_kwargs,
@ -1078,7 +1082,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 +1115,10 @@ async def _run_guardrails_and_call_llm(
async def handle_sampling_create_message(
context: Any,
context: object,
params: "CreateMessageRequestParams",
default_model: Optional[str] = None,
user_api_key_auth: Optional[Any] = None,
user_api_key_auth: Optional["UserAPIKeyAuth"] = None,
raw_headers: Optional[Dict[str, str]] = None,
client_ip: Optional[str] = None,
) -> Union["CreateMessageResult", "CreateMessageResultWithTools", "ErrorData"]:
@ -1214,7 +1218,6 @@ async def handle_sampling_create_message(
RateLimitError,
ServiceUnavailableError,
)
from litellm.proxy._types import ProxyException
if isinstance(

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

@ -4,10 +4,23 @@
This is an enterprise feature and requires a premium license.
"""
import json
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,
TypeVar,
overload,
)
from fastapi import (
APIRouter,
@ -69,13 +82,129 @@ from litellm.repositories.verification_token_repository import (
)
from litellm.types.proxy.management_endpoints.scim_v2 import *
if TYPE_CHECKING:
from prisma.models import LiteLLM_TeamTable as PrismaTeamTable
from prisma.models import LiteLLM_UserTable as PrismaUserTable
from prisma.models import LiteLLM_VerificationToken as PrismaVerificationToken
class _UserTableClient(Protocol):
async def find_first(self, where: Mapping[str, object]) -> "PrismaUserTable | None": ...
async def find_unique(self, where: Mapping[str, object]) -> "PrismaUserTable | None": ...
async def find_many(
self,
where: Mapping[str, object] | None = None,
skip: int | None = None,
take: int | None = None,
order: Mapping[str, object] | None = None,
) -> "Sequence[PrismaUserTable]": ...
async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> "PrismaUserTable": ...
async def delete(self, where: Mapping[str, object]) -> "PrismaUserTable | None": ...
async def count(self, where: Mapping[str, object] | None = None) -> int: ...
class _TeamTableClient(Protocol):
async def find_unique(self, where: Mapping[str, object]) -> "PrismaTeamTable | None": ...
async def find_many(
self,
where: Mapping[str, object] | None = None,
skip: int | None = None,
take: int | None = None,
order: Mapping[str, object] | None = None,
) -> "Sequence[PrismaTeamTable]": ...
async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> "PrismaTeamTable": ...
async def delete(self, where: Mapping[str, object]) -> "PrismaTeamTable | None": ...
async def count(self, where: Mapping[str, object] | None = None) -> int: ...
class _VerificationTokenTableClient(Protocol):
async def find_many(self, where: Mapping[str, object]) -> "Sequence[PrismaVerificationToken]": ...
async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> "PrismaVerificationToken": ...
class _DeleteManyTableClient(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,
) -> "_DeleteManyTableClient": ...
def _table(
repository: (
UserRepository
| TeamRepository
| VerificationTokenRepository
| InvitationLinkRepository
| OrganizationMembershipRepository
| TeamMembershipRepository
),
) -> object:
prisma_table: object = repository.table
return prisma_table
def _with_decoded_json_fields(data: Mapping[str, object], fields: tuple[str, ...]) -> Mapping[str, object]:
return {
key: (json.loads(value) if key in fields and isinstance(value, str) else value) for key, value in data.items()
}
_DomainModel = TypeVar("_DomainModel", bound=BaseModel)
def _construct_from_data(model_class: type[_DomainModel], data: Mapping[str, object]) -> _DomainModel:
instance = model_class.model_construct()
for key, value in data.items():
if key in model_class.model_fields:
setattr(instance, key, value)
return instance
@overload
def _to_domain_user(row: "PrismaUserTable") -> LiteLLM_UserTable: ...
@overload
def _to_domain_user(row: "PrismaUserTable | None") -> Optional[LiteLLM_UserTable]: ...
def _to_domain_user(row: "PrismaUserTable | None") -> Optional[LiteLLM_UserTable]:
if row is None:
return None
data = _with_decoded_json_fields(row, ("metadata",)) if isinstance(row, dict) else vars(row)
return _construct_from_data(LiteLLM_UserTable, data)
@overload
def _to_domain_team(row: "PrismaTeamTable") -> LiteLLM_TeamTable: ...
@overload
def _to_domain_team(row: "PrismaTeamTable | None") -> Optional[LiteLLM_TeamTable]: ...
def _to_domain_team(row: "PrismaTeamTable | None") -> Optional[LiteLLM_TeamTable]:
if row is None:
return None
data = _with_decoded_json_fields(row, ("metadata", "members_with_roles")) if isinstance(row, dict) else vars(row)
return _construct_from_data(LiteLLM_TeamTable, data)
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 +226,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 +236,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 +248,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,
@ -130,7 +259,7 @@ class UserProvisionerHelpers:
},
)
return await ScimTransformations.transform_litellm_user_to_scim_user(updated_user)
return await ScimTransformations.transform_litellm_user_to_scim_user(_to_domain_user(updated_user))
class ScimUserData(TypedDict):
@ -177,11 +306,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) -> "PrismaUserTable":
"""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 +318,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) -> "PrismaTeamTable":
"""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 +365,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 +467,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 +485,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 +498,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 +506,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 +602,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 +610,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 +750,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,9 +783,10 @@ 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
fields = _json_object_fields(metadata)
return fields is not None and fields.get(SCIM_BLOCKED_METADATA_KEY) is True
async def _set_user_keys_blocked(user_id: str, blocked: bool) -> int:
@ -676,7 +808,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 +816,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 +825,13 @@ 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 {}
key_row_fields = _json_object_fields(key_row.metadata)
current_metadata: Dict[str, object] = dict(key_row_fields) if key_row_fields is not None 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 +852,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,15 +868,16 @@ 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: object) -> Optional[bool]:
"""Read the SCIM active flag from a user's metadata dict, if present."""
if not metadata:
fields = _json_object_fields(metadata)
if fields is None:
return None
value = metadata.get("scim_active")
value = fields.get("scim_active")
if value is None:
return None
return bool(value)
@ -1241,7 +1375,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,15 +1392,18 @@ 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(
where=where_conditions,
skip=(startIndex - 1),
take=count,
order={"created_at": "desc"},
)
users: List[LiteLLM_UserTable] = [
_to_domain_user(row)
for row in await _table(UserRepository(prisma_client)).find_many(
where=where_conditions,
skip=(startIndex - 1),
take=count,
order={"created_at": "desc"},
)
]
# 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] = []
@ -1302,7 +1439,7 @@ async def get_user(
user = await _check_user_exists(user_id)
# Convert to SCIM format
scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(user)
scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(_to_domain_user(user))
return scim_user
except Exception as e:
@ -1330,7 +1467,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,
@ -1447,9 +1584,11 @@ async def update_user(
user.groups or [], admin_group, _default_scim_user_role()
)
updated_user = await UserRepository(prisma_client).table.update(
where={"user_id": user_id},
data=update_data,
updated_user = _to_domain_user(
await _table(UserRepository(prisma_client)).update(
where={"user_id": user_id},
data=update_data,
)
)
if client_set_active:
@ -1482,11 +1621,11 @@ async def delete_user(
prisma_client = await _get_prisma_client_or_raise_exception()
existing_user = await _check_user_exists(user_id)
# Get teams user belongs to
teams = []
# Get teams user belongs to; PrismaTeamTable is TYPE_CHECKING-only, so keep the quotes
teams: List["PrismaTeamTable"] = [] # noqa: UP037 # avoids a runtime NameError
if existing_user.teams:
for team_id in existing_user.teams:
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 team:
teams.append(team)
@ -1495,11 +1634,11 @@ async def delete_user(
current_members = 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}
)
team_row = LiteLLM_TeamTable.model_validate(team.model_dump())
team_row = _to_domain_team(team)
if any(member.user_id == user_id for member in team_row.members_with_roles or []):
await team_member_delete(
data=TeamMemberDeleteRequest(team_id=team_row.team_id, user_id=user_id),
@ -1511,7 +1650,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 +1720,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 +1728,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,20 +1736,16 @@ 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)
else:
bool_val = value
if isinstance(value, str):
bool_val = value.lower() == "true"
else:
bool_val = bool(value)
bool_val = value.lower() == "true" if isinstance(value, str) else bool(value)
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 +1759,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 +1779,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 +1816,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 +1827,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", {})
@ -1703,12 +1838,13 @@ def _apply_patch_ops(
for op in patch_ops.Operations:
path = (op.path or "").lower()
value = op.value
value: object = op.value
op_type = op.op
# Handle SCIM operations without path where value contains the fields
if not path and isinstance(value, dict):
for key, val in value.items():
fields = _json_object_fields(value) if not path else None
if fields is not None:
for key, val in fields.items():
key_lower = key.lower()
if key_lower == "active":
_handle_active_update(op_type, val, metadata)
@ -1718,8 +1854,9 @@ def _apply_patch_ops(
_handle_externalid_update(op_type, val, update_data)
elif key_lower in SCIM_MULTI_VALUED_ATTRIBUTE_METADATA_KEYS:
_handle_multi_valued_attribute_update(key_lower, op_type, val, metadata)
elif key_lower == "name" and isinstance(val, dict):
for name_key, name_val in val.items():
elif key_lower == "name":
name_fields = _json_object_fields(val)
for name_key, name_val in (name_fields or {}).items():
name_key_lower = name_key.lower()
if name_key_lower in ("givenname", "familyname"):
_handle_name_update(
@ -1846,7 +1983,7 @@ async def patch_user(
prev_active = _scim_active_value(existing_user.metadata)
update_data, final_team_set = _apply_patch_ops(
existing_user=existing_user,
existing_user=_to_domain_user(existing_user),
patch_ops=patch_ops,
)
@ -1875,9 +2012,11 @@ async def patch_user(
update_data["metadata"] = safe_dumps(update_data["metadata"])
updated_user = await UserRepository(prisma_client).table.update(
where={"user_id": user_id},
data=update_data,
updated_user = _to_domain_user(
await _table(UserRepository(prisma_client)).update(
where={"user_id": user_id},
data=update_data,
)
)
if new_active is not None and new_active != (True if prev_active is None else prev_active):
@ -1923,18 +2062,21 @@ async def get_groups(
where_conditions["team_alias"] = team_alias
# Get teams from database
teams = await TeamRepository(prisma_client).table.find_many(
where=where_conditions,
skip=(startIndex - 1),
take=count,
order={"created_at": "desc"},
)
teams = [
_to_domain_team(row)
for row in await _table(TeamRepository(prisma_client)).find_many(
where=where_conditions,
skip=(startIndex - 1),
take=count,
order={"created_at": "desc"},
)
]
# 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
@ -1942,7 +2084,7 @@ async def get_groups(
# list to the IdP and trigger repeated re-provisioning.
members = await _get_team_members_display(await _get_team_member_user_ids_from_team(team))
verbose_proxy_logger.debug(f"SCIM GET GROUPS members: {members}")
team_alias = getattr(team, "team_alias", team.team_id)
team_alias = team.team_alias or team.team_id
team_created_at = team.created_at.isoformat() if team.created_at else None
team_updated_at = team.updated_at.isoformat() if team.updated_at else None
@ -1987,7 +2129,7 @@ async def get_group(
try:
team = await _check_team_exists(group_id)
scim_group = await ScimTransformations.transform_litellm_team_to_scim_group(team)
scim_group = await ScimTransformations.transform_litellm_team_to_scim_group(_to_domain_team(team))
verbose_proxy_logger.debug(f"SCIM GET GROUP response: {scim_group}")
return scim_group
@ -2018,7 +2160,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(
@ -2078,9 +2220,8 @@ async def update_group(
verbose_proxy_logger.debug(f"SCIM PUT GROUP created_users: {len(member_result.created_users)}")
# Prepare update data
existing_metadata = existing_team.metadata if existing_team.metadata else {}
updated_metadata = {
**existing_metadata,
updated_metadata: Mapping[str, object] = {
**(_json_object_fields(existing_team.metadata) or {}),
SCIM_TEAM_DATA_METADATA_KEY: group.model_dump(),
SCIM_MANAGED_TEAM_METADATA_KEY: True,
}
@ -2091,13 +2232,15 @@ async def update_group(
}
# Update team in database
updated_team = await TeamRepository(prisma_client).table.update(
where={"team_id": group_id},
data=update_data,
updated_team = _to_domain_team(
await _table(TeamRepository(prisma_client)).update(
where={"team_id": group_id},
data=update_data,
)
)
# Handle user-team relationship changes
current_members = set(await _get_team_member_user_ids_from_team(existing_team))
current_members = set(await _get_team_member_user_ids_from_team(_to_domain_team(existing_team)))
verbose_proxy_logger.debug(f"SCIM PUT GROUP current_members: {current_members}")
final_members = set(member_result.all_member_ids)
verbose_proxy_logger.debug(f"SCIM PUT GROUP final_members: {final_members}")
@ -2141,23 +2284,23 @@ async def delete_group(
prisma_client = await _get_prisma_client_or_raise_exception()
existing_team = await _check_team_exists(group_id)
member_ids = await _get_team_member_user_ids_from_team(existing_team)
member_ids = await _get_team_member_user_ids_from_team(_to_domain_team(existing_team))
# 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 +2309,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 +2326,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}
@ -2199,7 +2342,7 @@ async def _process_group_patch_operations(
# Process each patch operation
for op in patch_ops.Operations:
path = (op.path or "").lower()
value = op.value
value: object = op.value
op_type = op.op
if path == "displayname":
@ -2251,7 +2394,9 @@ 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
) -> "PrismaTeamTable | None":
"""Apply the group's metadata/displayName patch updates to the database.
Membership itself is not written here; it is reconciled onto the source of
@ -2264,11 +2409,11 @@ async def _apply_group_patch_updates(group_id: str, update_data: Dict[str, Any],
update_data["metadata"] = safe_dumps(update_data["metadata"])
if update_data:
return await TeamRepository(prisma_client).table.update(
return await _table(TeamRepository(prisma_client)).update(
where={"team_id": group_id},
data=update_data,
)
return await TeamRepository(prisma_client).table.find_unique(where={"team_id": group_id})
return await _table(TeamRepository(prisma_client)).find_unique(where={"team_id": group_id})
async def _handle_group_membership_changes(group_id: str, current_members: Set[str], final_members: Set[str]):
@ -2316,7 +2461,7 @@ async def patch_group(
try:
prisma_client = await _get_prisma_client_or_raise_exception()
existing_team = await _check_team_exists(group_id)
existing_team = _to_domain_team(await _check_team_exists(group_id))
# Process patch operations
update_data, final_members, replace_target = await _process_group_patch_operations(
@ -2328,13 +2473,11 @@ async def patch_group(
intended_remove = snapshot_members - final_members
# Apply the metadata/displayName updates to the database
updated_team = await _apply_group_patch_updates(group_id, update_data, prisma_client)
updated_team = _to_domain_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()))
)
set(await _get_team_member_user_ids_from_team(_to_domain_team(refreshed_team)))
if refreshed_team
else snapshot_members
)
@ -2356,14 +2499,18 @@ 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
updated_team = _to_domain_team(final_team)
if updated_team is None:
raise HTTPException(
status_code=404,
detail={"error": f"Group not found with ID: {group_id}"},
)
# Convert to SCIM format and return
scim_group = await ScimTransformations.transform_litellm_team_to_scim_group(
LiteLLM_TeamTable.model_validate(updated_team.model_dump())
)
scim_group = await ScimTransformations.transform_litellm_team_to_scim_group(updated_team)
return scim_group
except Exception as e:

221
litellm/responses/streaming_iterator.py
View file

@ -8,7 +8,17 @@ import uuid
from datetime import datetime
from functools import lru_cache
from types import MappingProxyType
from typing import Any, Dict, List, Literal, Mapping, Optional
from typing import (
TYPE_CHECKING,
Any,
Dict,
List,
Literal,
Mapping,
Optional,
Protocol,
Union,
)
import httpx
from openai._streaming import SSEDecoder
@ -29,10 +39,52 @@ from litellm.litellm_core_utils.llm_response_utils.response_metadata import (
from litellm.litellm_core_utils.thread_pool_executor import executor
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
from litellm.responses.utils import ResponseAPILoggingUtils, ResponsesAPIRequestUtils
from litellm.types.llms.openai import ResponsesAPIStreamEvents
from litellm.types.guardrails import PresidioPerRequestConfig
from litellm.types.llms.openai import (
ContentPartDoneEvent,
ContentPartDonePartOutputText,
ContentPartDonePartReasoningText,
ContentPartDonePartRefusal,
ResponseCompletedEvent,
ResponseCreatedEvent,
ResponseFailedEvent,
ResponseIncompleteEvent,
ResponseInProgressEvent,
ResponsesAPIResponse,
ResponsesAPIStreamEvents,
ResponsesAPIStreamingResponse,
)
from litellm.types.utils import CallTypes
from litellm.utils import async_post_call_success_deployment_hook
if TYPE_CHECKING:
from fastapi import WebSocket
from websockets.asyncio.client import ClientConnection
from litellm.proxy._types import UserAPIKeyAuth
_RESPONSE_EVENT_TYPES_WITH_RESPONSE = (
ResponseCreatedEvent,
ResponseInProgressEvent,
ResponseCompletedEvent,
ResponseFailedEvent,
ResponseIncompleteEvent,
)
class _PIIMaskingGuardrail(Protocol):
async def check_pii(
self,
text: str,
output_parse_pii: bool,
presidio_config: Optional[PresidioPerRequestConfig],
request_data: Mapping[str, Any],
) -> str: ...
def get_presidio_settings_from_request_data(
self, data: Mapping[str, Any]
) -> Optional[PresidioPerRequestConfig]: ...
@lru_cache(maxsize=1)
def _get_openai_response_types():
@ -41,7 +93,7 @@ def _get_openai_response_types():
return openai_types
def _log_background_task_failure(task: "asyncio.Task[Any]", *, task_name: str) -> None:
def _log_background_task_failure(task: asyncio.Task[object], *, task_name: str) -> None:
if task.cancelled():
return
exception = task.exception()
@ -130,7 +182,7 @@ class BaseResponsesAPIStreamingIterator:
self.logging_obj = logging_obj
self.finished = False
self.responses_api_provider_config = responses_api_provider_config
self.completed_response: Optional[Any] = None
self.completed_response: Optional[ResponsesAPIStreamingResponse] = None
self.start_time = getattr(logging_obj, "start_time", datetime.now())
self._failure_handled = False # Track if failure handler has been called
self._yielded_first_chunk = False
@ -175,7 +227,7 @@ class BaseResponsesAPIStreamingIterator:
llm_provider=self.custom_llm_provider or "",
)
def _process_chunk(self, chunk) -> Optional[Any]:
def _process_chunk(self, chunk: str) -> Optional[ResponsesAPIStreamingResponse]:
"""Process a single chunk of data from the stream"""
if not chunk:
return None
@ -195,7 +247,7 @@ class BaseResponsesAPIStreamingIterator:
try:
# Parse the JSON chunk
parsed_chunk = json.loads(chunk)
parsed_chunk: object = json.loads(chunk)
# Format as ResponsesAPIStreamingResponse
if isinstance(parsed_chunk, dict):
@ -210,15 +262,16 @@ class BaseResponsesAPIStreamingIterator:
# Only when the SSE JSON carries a response body (delta events do not).
# Using getattr(..., "response") alone is unsafe with Mocks: they synthesize a
# truthy child Mock for any attribute, which breaks tests and is wrong on stream.
if "response" in parsed_chunk:
response_object = getattr(openai_responses_api_chunk, "response", None)
if response_object is not None:
response = ResponsesAPIRequestUtils._update_responses_api_response_id_with_model_id(
responses_api_response=response_object,
if "response" in parsed_chunk and isinstance(
openai_responses_api_chunk, _RESPONSE_EVENT_TYPES_WITH_RESPONSE
):
openai_responses_api_chunk.response = (
ResponsesAPIRequestUtils._update_responses_api_response_id_with_model_id(
responses_api_response=openai_responses_api_chunk.response,
litellm_metadata=self.litellm_metadata,
custom_llm_provider=self.custom_llm_provider,
)
setattr(openai_responses_api_chunk, "response", response)
)
# Encode container_id on streaming events so proxy/UI follow-ups route correctly
_event_type = getattr(openai_responses_api_chunk, "type", None)
@ -297,20 +350,21 @@ class BaseResponsesAPIStreamingIterator:
):
self.completed_response = openai_responses_api_chunk
# Add cost to usage object if include_cost_in_streaming_usage is True
if litellm.include_cost_in_streaming_usage and self.logging_obj is not None:
response_obj: Optional[Any] = getattr(openai_responses_api_chunk, "response", None)
if response_obj:
usage_obj: Optional[Any] = getattr(response_obj, "usage", None)
if usage_obj is not None:
try:
cost: Optional[float] = self.logging_obj._response_cost_calculator(
result=response_obj
)
if cost is not None:
setattr(usage_obj, "cost", cost)
except Exception:
# Best-effort usage cost annotation should not break stream replay.
pass
if (
litellm.include_cost_in_streaming_usage
and self.logging_obj is not None
and isinstance(openai_responses_api_chunk, _RESPONSE_EVENT_TYPES_WITH_RESPONSE)
):
response_obj = openai_responses_api_chunk.response
usage_obj = response_obj.usage
if usage_obj is not None:
try:
cost: Optional[float] = self.logging_obj._response_cost_calculator(result=response_obj)
if cost is not None:
usage_obj.cost = cost
except Exception:
# Best-effort usage cost annotation should not break stream replay.
pass
if _chunk_type == openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED:
self._handle_logging_failed_response()
@ -391,8 +445,8 @@ class BaseResponsesAPIStreamingIterator:
async_failure_handler / failure_handler so logging integrations correctly
record the call as failed.
"""
response_obj = getattr(self.completed_response, "response", None) if self.completed_response else None
error_info = getattr(response_obj, "error", None) if response_obj else None
response_obj = getattr(self.completed_response, "response", None)
error_info = getattr(response_obj, "error", None)
error_message, error_type, error_code = _error_event_fields(error_info)
self._record_failed_response_usage(response_obj)
exception = litellm.APIError(
@ -403,10 +457,10 @@ class BaseResponsesAPIStreamingIterator:
)
self._handle_failure(exception)
def _record_failed_response_usage(self, response_obj: Optional[Any]) -> None:
def _record_failed_response_usage(self, response_obj: Optional[ResponsesAPIResponse]) -> None:
if response_obj is None or self.logging_obj is None:
return
usage_obj = getattr(response_obj, "usage", None)
usage_obj = response_obj.usage
if usage_obj is None:
return
try:
@ -453,7 +507,7 @@ class BaseResponsesAPIStreamingIterator:
is_pre_first_chunk=not self._yielded_first_chunk,
)
def _get_completed_response_object(self) -> Optional[Any]:
def _get_completed_response_object(self) -> Optional[ResponsesAPIResponse]:
openai_types = _get_openai_response_types()
completed_response = self.completed_response
if isinstance(completed_response, openai_types.ResponsesAPIResponse):
@ -529,7 +583,9 @@ class BaseResponsesAPIStreamingIterator:
self._completed_response_cached = True
async def _call_post_streaming_deployment_hook(self, chunk):
async def _call_post_streaming_deployment_hook(
self, chunk: ResponsesAPIStreamingResponse
) -> ResponsesAPIStreamingResponse:
"""
Allow callbacks to modify streaming chunks before returning (parity with chat).
"""
@ -566,7 +622,9 @@ class BaseResponsesAPIStreamingIterator:
except Exception:
return chunk
async def call_post_streaming_hooks_for_testing(self, chunk):
async def call_post_streaming_hooks_for_testing(
self, chunk: ResponsesAPIStreamingResponse
) -> ResponsesAPIStreamingResponse:
"""
Helper to invoke streaming deployment hooks explicitly (used in tests).
"""
@ -668,7 +726,9 @@ class BaseResponsesAPIStreamingIterator:
pass
async def call_post_streaming_hooks_for_testing(iterator, chunk):
async def call_post_streaming_hooks_for_testing(
iterator: BaseResponsesAPIStreamingIterator, chunk: ResponsesAPIStreamingResponse
) -> ResponsesAPIStreamingResponse:
"""
Module-level helper for tests to ensure hooks can be invoked even if the iterator is wrapped.
"""
@ -709,7 +769,7 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
def __aiter__(self):
return self
async def __anext__(self) -> Any:
async def __anext__(self) -> ResponsesAPIStreamingResponse:
try:
self._check_max_streaming_duration()
while True:
@ -791,7 +851,7 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
def __iter__(self):
return self
def __next__(self):
def __next__(self) -> ResponsesAPIStreamingResponse:
try:
self._check_max_streaming_duration()
while True:
@ -882,7 +942,7 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
def _set_events_from_response(
self,
transformed: Any,
transformed: ResponsesAPIResponse,
logging_obj: LiteLLMLoggingObj,
) -> None:
self._events = _build_synthetic_response_events(
@ -896,7 +956,7 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
def __aiter__(self):
return self
async def __anext__(self) -> Any:
async def __anext__(self) -> ResponsesAPIStreamingResponse:
if self._idx >= len(self._events):
raise StopAsyncIteration
evt = self._events[self._idx]
@ -910,7 +970,7 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
def __iter__(self):
return self
def __next__(self) -> Any:
def __next__(self) -> ResponsesAPIStreamingResponse:
if self._idx >= len(self._events):
raise StopIteration
evt = self._events[self._idx]
@ -925,7 +985,7 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
def __init__(
self,
response: Any,
response: ResponsesAPIResponse,
logging_obj: LiteLLMLoggingObj,
request_data: Optional[Dict[str, Any]] = None,
call_type: Optional[str] = None,
@ -933,7 +993,7 @@ class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
BaseResponsesAPIStreamingIterator.__init__(
self,
response=httpx.Response(200),
model=getattr(response, "model", ""),
model=response.model or "",
responses_api_provider_config=None,
logging_obj=logging_obj,
litellm_metadata=None,
@ -943,13 +1003,13 @@ class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
)
self._completed_response_cache_hit = True
self._persist_completed_response_before_logging = False
self._events: List[Any] = []
self._events: List[ResponsesAPIStreamingResponse] = []
self._idx = 0
self._set_events_from_response(transformed=response, logging_obj=logging_obj)
def _set_events_from_response(
self,
transformed: Any,
transformed: ResponsesAPIResponse,
logging_obj: LiteLLMLoggingObj,
) -> None:
self._events = _build_synthetic_response_events(
@ -963,7 +1023,7 @@ class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
def __aiter__(self):
return self
async def __anext__(self) -> Any:
async def __anext__(self) -> ResponsesAPIStreamingResponse:
if self._idx >= len(self._events):
raise StopAsyncIteration
evt = self._events[self._idx]
@ -977,7 +1037,7 @@ class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
def __iter__(self):
return self
def __next__(self) -> Any:
def __next__(self) -> ResponsesAPIStreamingResponse:
if self._idx >= len(self._events):
raise StopIteration
evt = self._events[self._idx]
@ -1002,8 +1062,8 @@ def _build_response_status_event(
"response.created",
"response.in_progress",
],
transformed: Any,
) -> Any:
transformed: ResponsesAPIResponse,
) -> Union[ResponseCreatedEvent, ResponseInProgressEvent]:
openai_types = _get_openai_response_types()
in_progress_response = transformed.model_copy(
deep=True,
@ -1020,10 +1080,10 @@ def _build_content_part_done_event(
output_index: int,
content_index: int,
part_payload: Dict[str, Any],
) -> Optional[Any]:
) -> Optional[ContentPartDoneEvent]:
openai_types = _get_openai_response_types()
part_type = part_payload.get("type")
part: Any
part: Union[ContentPartDonePartOutputText, ContentPartDonePartRefusal, ContentPartDonePartReasoningText]
if part_type == "output_text":
annotations = [
openai_types.BaseLiteLLMOpenAIResponseObject(**annotation)
@ -1059,7 +1119,7 @@ def _build_content_part_done_event(
def _add_text_like_part_events(
*,
events: List[Any],
events: List[ResponsesAPIStreamingResponse],
item_id: str,
output_index: int,
content_index: int,
@ -1125,28 +1185,27 @@ def _add_text_like_part_events(
def _build_synthetic_response_events(
*,
transformed: Any,
transformed: ResponsesAPIResponse,
logging_obj: LiteLLMLoggingObj,
chunk_size: int,
) -> List[Any]:
) -> List[ResponsesAPIStreamingResponse]:
openai_types = _get_openai_response_types()
if litellm.include_cost_in_streaming_usage and logging_obj is not None:
usage_obj: Optional[Any] = getattr(transformed, "usage", None)
usage_obj = transformed.usage
if usage_obj is not None:
try:
cost: Optional[float] = logging_obj._response_cost_calculator(result=transformed)
if cost is not None:
setattr(usage_obj, "cost", cost)
usage_obj.cost = cost
except Exception:
pass
events: List[Any] = [
_build_response_status_event(openai_types.ResponsesAPIStreamEvents.RESPONSE_CREATED, transformed),
_build_response_status_event(openai_types.ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS, transformed),
]
events: List[ResponsesAPIStreamingResponse] = []
events.append(_build_response_status_event(openai_types.ResponsesAPIStreamEvents.RESPONSE_CREATED, transformed))
events.append(_build_response_status_event(openai_types.ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS, transformed))
sequence_number = 0
for output_index, output_item in enumerate(getattr(transformed, "output", []) or []):
for output_index, output_item in enumerate(transformed.output or []):
output_item_payload = _dump_response_object(output_item)
item_id = str(output_item_payload.get("id") or transformed.id)
item_type = output_item_payload.get("type")
@ -1294,34 +1353,34 @@ class ResponsesWebSocketStreaming:
def __init__(
self,
websocket: Any,
backend_ws: Any,
websocket: WebSocket,
backend_ws: ClientConnection,
logging_obj: LiteLLMLoggingObj,
user_api_key_dict: Optional[Any] = None,
request_data: Optional[Dict] = None,
user_api_key_dict: Optional[UserAPIKeyAuth] = None,
request_data: Optional[Dict[str, Any]] = None,
first_message: Optional[str] = None,
guardrail_callbacks: Optional[List[Any]] = None,
output_guardrail_callbacks: Optional[List[Any]] = None,
guardrail_callbacks: Optional[List[_PIIMaskingGuardrail]] = None,
output_guardrail_callbacks: Optional[List[_PIIMaskingGuardrail]] = None,
authorized_model: Optional[str] = None,
):
self.websocket = websocket
self.backend_ws = backend_ws
self.logging_obj = logging_obj
self.user_api_key_dict = user_api_key_dict
self.request_data: Dict = request_data or {}
self.messages: list[Dict] = []
self.request_data: Dict[str, Any] = request_data or {}
self.messages: list[Mapping[str, Any]] = []
self.input_messages: list[Dict[str, str]] = []
self.first_message = first_message
self.guardrail_callbacks: List[Any] = guardrail_callbacks or []
self.output_guardrail_callbacks: List[Any] = output_guardrail_callbacks or []
self.guardrail_callbacks: List[_PIIMaskingGuardrail] = guardrail_callbacks or []
self.output_guardrail_callbacks: List[_PIIMaskingGuardrail] = output_guardrail_callbacks or []
# Model name authorized at connection time; enforced on every
# response.create frame to prevent deployment-substitution attacks.
self.authorized_model: Optional[str] = authorized_model
def _should_store_event(self, event_obj: dict) -> bool:
def _should_store_event(self, event_obj: Mapping[str, Any]) -> bool:
return event_obj.get("type") in RESPONSES_WS_LOGGED_EVENT_TYPES
def _store_event(self, event: Any) -> None:
def _store_event(self, event: Union[str, bytes, Mapping[str, Any]]) -> None:
if isinstance(event, bytes):
event = event.decode("utf-8")
if isinstance(event, str):
@ -1335,15 +1394,13 @@ class ResponsesWebSocketStreaming:
if self._should_store_event(event_obj):
self.messages.append(event_obj)
def _collect_input_from_client_event(self, message: Any) -> None:
def _collect_input_from_client_event(self, message: Union[str, Mapping[str, Any]]) -> None:
"""Extract user input content from response.create for logging."""
try:
if isinstance(message, str):
msg_obj = json.loads(message)
elif isinstance(message, dict):
msg_obj = message
else:
return
msg_obj = message
if msg_obj.get("type") != "response.create":
return
@ -1370,7 +1427,7 @@ class ResponsesWebSocketStreaming:
except (json.JSONDecodeError, AttributeError, TypeError):
pass
def _store_input(self, message: Any) -> None:
def _store_input(self, message: Union[str, Mapping[str, Any]]) -> None:
self._collect_input_from_client_event(message)
if self.logging_obj:
self.logging_obj.pre_call(input=message, api_key="")
@ -1620,7 +1677,7 @@ class ResponsesWebSocketStreaming:
continue
text = content_block.get("text")
if isinstance(text, str):
unmasked = cb._unmask_pii_text(text, pii_tokens)
unmasked = getattr(cb, "_unmask_pii_text")(text, pii_tokens)
if unmasked != text:
content_block["text"] = unmasked
modified = True
@ -1629,7 +1686,7 @@ class ResponsesWebSocketStreaming:
if event_type in self._DELTA_EVENT_TYPES:
delta = evt_obj.get("delta")
if isinstance(delta, str):
unmasked = cb._unmask_pii_text(delta, pii_tokens)
unmasked = getattr(cb, "_unmask_pii_text")(delta, pii_tokens)
if unmasked != delta:
evt_obj["delta"] = unmasked
return json.dumps(evt_obj)
@ -1795,10 +1852,10 @@ class ManagedResponsesWebSocketHandler:
def __init__(
self,
websocket: Any,
websocket: WebSocket,
model: str,
logging_obj: "LiteLLMLoggingObj",
user_api_key_dict: Optional[Any] = None,
logging_obj: LiteLLMLoggingObj,
user_api_key_dict: Optional[UserAPIKeyAuth] = None,
litellm_metadata: Optional[Dict[str, Any]] = None,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
@ -1927,7 +1984,7 @@ class ManagedResponsesWebSocketHandler:
return messages
@staticmethod
def _input_to_messages(input_val: Any) -> List[Dict[str, Any]]:
def _input_to_messages(input_val: object) -> List[Dict[str, Any]]:
"""
Normalise the ``input`` field of a ``response.create`` event to a list
of Responses API message dicts.

24
ruff-strict-budget.json
View file

@ -1,6 +1,6 @@
{
"ANN001": {
"limit": 3118
"limit": 3078
},
"ANN002": {
"limit": 69
@ -9,13 +9,13 @@
"limit": 831
},
"ANN201": {
"limit": 2138
"limit": 2130
},
"ANN202": {
"limit": 944
"limit": 928
},
"ANN204": {
"limit": 724
"limit": 720
},
"ANN205": {
"limit": 127
@ -24,7 +24,7 @@
"limit": 130
},
"ANN401": {
"limit": 2010
"limit": 1650
},
"ASYNC230": {
"limit": 14
@ -42,7 +42,7 @@
"limit": 84
},
"B010": {
"limit": 194
"limit": 186
},
"B018": {
"limit": 5
@ -123,7 +123,7 @@
"limit": 52
},
"I001": {
"limit": 270
"limit": 258
},
"LOG015": {
"limit": 8
@ -135,7 +135,7 @@
"limit": 30
},
"PERF401": {
"limit": 142
"limit": 138
},
"PERF402": {
"limit": 9
@ -264,7 +264,7 @@
"limit": 61
},
"SIM102": {
"limit": 324
"limit": 320
},
"SIM103": {
"limit": 129
@ -306,7 +306,7 @@
"limit": 9
},
"TID251": {
"limit": 2652
"limit": 2648
},
"TRY002": {
"limit": 547
@ -324,7 +324,7 @@
"limit": 879
},
"UP006": {
"limit": 12138
"limit": 12134
},
"UP007": {
"limit": 2526
@ -363,6 +363,6 @@
"limit": 105
},
"UP045": {
"limit": 17805
"limit": 17794
}
}

2
ruff.toml
View file

@ -6,7 +6,7 @@ lint.extend-select = ["T20", "PGH004", "RUF008", "RUF009", "RUF100"]
# litellm's own ruff config both rely on suppressions this config can't see.
lint.external = [
# Enforced by the strict-rule gate (scripts/ruff_strict_gate.py + ruff-strict.toml)
"C901",
"C901", "UP037",
# Enforced by upstream litellm's ruff config, but not run in this repo's CI
"PLC0415", "E402", "BLE001", "ARG002", "S102", "S324", "S606", "D401", "F403", "F405",
]

3
tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py
View file

@ -24,6 +24,7 @@ from litellm.proxy.management_endpoints.scim.scim_v2 import (
_process_group_patch_operations,
_recompute_scim_member_roles,
_resolve_group_member_ids,
_to_domain_user,
create_group,
create_user,
delete_group,
@ -627,7 +628,7 @@ async def test_handle_existing_user_by_email_existing_user_updated(mocker):
raise_on_error=True,
)
mock_transform.assert_called_once_with(updated_user)
mock_transform.assert_called_once_with(_to_domain_user(updated_user))
@pytest.mark.asyncio

6
type-discipline-budget.json
View file

@ -1,9 +1,9 @@
{
"LIT001": {
"limit": 23253
"limit": 23227
},
"LIT002": {
"limit": 27427
"limit": 27399
},
"LIT003": {
"limit": 292
@ -15,7 +15,7 @@
"limit": 0
},
"LIT006": {
"limit": 1108
"limit": 1092
},
"LIT007": {
"limit": 0