scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-09 22:31:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 86

Merge pull request #36722 from BerriAI/litellm_decrease_anys_fable6

Browse source 
chore(typing): clear 1.2k basedpyright Any errors across 16 hotspot files
This commit is contained in:
Mateo Wang 2026-08-30 10:09:02 -07:00 • committed by GitHub
commit 8a156ed42d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
16 changed files with 1212 additions and 662 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

24
basedpyright-code-budget.json
View file

@ -1,9 +1,9 @@
{
"reportAny": {
"limit": 16389
"limit": 16171
},
"reportArgumentType": {
"limit": 2229
"limit": 2226
},
"reportAssignmentType": {
"limit": 319
@ -18,13 +18,13 @@
"limit": 40
},
"reportDeprecated": {
"limit": 212
"limit": 211
},
"reportDuplicateImport": {
"limit": 19
},
"reportExplicitAny": {
"limit": 5242
"limit": 5199
},
"reportFunctionMemberAccess": {
"limit": 7
@ -54,10 +54,10 @@
"limit": 0
},
"reportMissingParameterType": {
"limit": 5614
"limit": 5611
},
"reportMissingTypeArgument": {
"limit": 15356
"limit": 15350
},
"reportMissingTypeStubs": {
"limit": 40
@ -93,25 +93,25 @@
"limit": 181
},
"reportTypedDictNotRequiredAccess": {
"limit": 25
"limit": 24
},
"reportUndefinedVariable": {
"limit": 0
},
"reportUnknownArgumentType": {
"limit": 44389
"limit": 44368
},
"reportUnknownLambdaType": {
"limit": 109
},
"reportUnknownMemberType": {
"limit": 38500
"limit": 38468
},
"reportUnknownParameterType": {
"limit": 19673
"limit": 19665
},
"reportUnknownVariableType": {
"limit": 30092
"limit": 30066
},
"reportUnnecessaryCast": {
"limit": 111
@ -123,7 +123,7 @@
"limit": 5
},
"reportUnnecessaryIsInstance": {
"limit": 829
"limit": 828
},
"reportUntypedBaseClass": {
"limit": 0

167
litellm/google_genai/adapters/transformation.py
View file

@ -1,8 +1,9 @@
import json
from collections.abc import AsyncIterator, Iterator, Sequence
from typing import Any, Final, TypedDict, cast
from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence
from types import MappingProxyType
from typing import Any, Final, TypeAlias, cast
from typing_extensions import ReadOnly
from typing_extensions import ReadOnly, TypedDict
from litellm import verbose_logger
from litellm.litellm_core_utils.json_validation_rule import normalize_tool_schema
@ -11,7 +12,6 @@ from litellm.types.llms.openai import (
ChatCompletionAssistantMessage,
ChatCompletionAssistantToolCall,
ChatCompletionImageObject,
ChatCompletionRequest,
ChatCompletionSystemMessage,
ChatCompletionTextObject,
ChatCompletionToolCallFunctionChunk,
@ -23,35 +23,63 @@ from litellm.types.llms.openai import (
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import (
AdapterCompletionStreamWrapper,
ChatCompletionDeltaCustomToolCall,
ChatCompletionDeltaToolCall,
ChatCompletionMessageCustomToolCall,
ChatCompletionMessageToolCall,
Choices,
Delta,
Function,
Message,
ModelResponse,
ModelResponseStream,
StreamingChoices,
Usage,
)
class _GenAITextPart(TypedDict, total=False):
text: ReadOnly[str]
_JsonDict: TypeAlias = dict[str, object]
_JsonDictList: TypeAlias = list[_JsonDict]
class _GenAISystemInstruction(TypedDict, total=False):
parts: ReadOnly[list[_GenAITextPart]]
class _ToolCallAccumulator(TypedDict):
name: ReadOnly[str]
arguments: ReadOnly[str]
class _GenAIFunctionCall(TypedDict):
name: ReadOnly[str]
args: ReadOnly[Mapping[str, object]]
class _GenAIPart(TypedDict, total=False):
text: ReadOnly[str]
functionCall: ReadOnly[dict[str, object]]
functionCall: ReadOnly[_GenAIFunctionCall]
class _GenAIFunctionResponse(TypedDict, total=False):
name: ReadOnly[str]
response: ReadOnly[object]
class _GenAIRequestFunctionCall(TypedDict, total=False):
name: ReadOnly[str]
args: ReadOnly[Mapping[str, object]]
class _GenAIContentPart(TypedDict, total=False):
text: ReadOnly[str]
inline_data: ReadOnly[Mapping[str, str]]
functionResponse: ReadOnly[_GenAIFunctionResponse]
functionCall: ReadOnly[_GenAIRequestFunctionCall]
class _GenAIFunctionDeclaration(TypedDict, total=False):
name: ReadOnly[str]
description: ReadOnly[str]
parametersJsonSchema: ReadOnly[dict[str, object]]
parametersJsonSchema: ReadOnly[object]
class _GenAITool(TypedDict, total=False):
functionDeclarations: ReadOnly[list[_GenAIFunctionDeclaration]]
functionDeclarations: ReadOnly[Sequence[_GenAIFunctionDeclaration]]
class _GenAIFunctionCallingConfig(TypedDict, total=False):
@ -62,9 +90,11 @@ class _GenAIToolConfig(TypedDict, total=False):
functionCallingConfig: ReadOnly[_GenAIFunctionCallingConfig]
def _decode_tool_call_arguments(raw_arguments: str) -> object:
"""Decode a tool call's JSON-encoded arguments into the value Google GenAI expects."""
return json.loads(raw_arguments)
class _GenAISystemInstruction(TypedDict, total=False):
parts: ReadOnly[Sequence[Mapping[str, str]]]
_EMPTY_STR_MAPPING: Final[Mapping[str, str]] = MappingProxyType({})
class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
@ -74,12 +104,12 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
"""
sent_first_chunk: bool = False
# State tracking for accumulating partial tool calls
accumulated_tool_calls: dict[int, dict[str, str]]
_parse_accumulated_args: Callable[[str], Mapping[str, object]] = staticmethod(json.loads)
def __init__(self, completion_stream: object):
self.sent_first_chunk = False
self.accumulated_tool_calls = {}
# State tracking for accumulating partial tool calls
self.accumulated_tool_calls = dict[int, _ToolCallAccumulator]()
self._returned_response = False
super().__init__(completion_stream)
@ -124,7 +154,7 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
# After the stream is exhausted, check for any remaining accumulated tool calls
if self.accumulated_tool_calls:
try:
parts: Final[list[_GenAIPart]] = []
parts: Final = list[_GenAIPart]()
for (
tool_call_index,
tool_call_data,
@ -132,7 +162,9 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
try:
# For tool calls with no arguments, accumulated_args will be "", which is not valid JSON.
# We default to an empty JSON object in this case.
parsed_args = _decode_tool_call_arguments(tool_call_data["arguments"] or "{}")
parsed_args: Mapping[str, object] = self._parse_accumulated_args(
tool_call_data["arguments"] or "{}"
)
function_call_part: _GenAIPart = {
"functionCall": {
"name": tool_call_data["name"] or "undefined_tool_name",
@ -149,7 +181,7 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
tool_call_data["arguments"],
)
if parts:
final_chunk: Final[dict[str, object]] = {
final_chunk: Final = {
"candidates": [
{
"content": {"parts": parts, "role": "model"},
@ -211,14 +243,16 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
class GoogleGenAIAdapter:
"""Adapter for transforming Google GenAI generate_content requests to/from litellm.completion format"""
_parse_tool_call_args: Callable[[str], Mapping[str, object]] = staticmethod(json.loads)
def __init__(self) -> None:
pass
def translate_generate_content_to_completion(
self,
model: str,
contents: list[dict[str, Any]] | dict[str, Any],
config: dict[str, Any] | None = None,
contents: _JsonDictList | _JsonDict,
config: Mapping[str, object] | None = None,
litellm_params: GenericLiteLLMParams | None = None,
**kwargs,
) -> dict[str, Any]:
@ -250,7 +284,7 @@ class GoogleGenAIAdapter:
messages: Final = self._transform_contents_to_messages(contents_list, system_instruction=system_instruction)
# Create base request as dict (which is compatible with ChatCompletionRequest)
completion_request: Final[ChatCompletionRequest] = {
completion_request: Final[_JsonDict] = {
"model": model,
"messages": messages,
}
@ -312,9 +346,9 @@ class GoogleGenAIAdapter:
def _add_generic_litellm_params_to_request(
self,
completion_request_dict: dict[str, object],
completion_request_dict: _JsonDict,
litellm_params: GenericLiteLLMParams | None = None,
) -> dict[str, object]:
) -> _JsonDict:
"""Add generic litellm params to request. e.g add api_base, api_key, api_version, etc.
Args:
@ -326,7 +360,7 @@ class GoogleGenAIAdapter:
"""
allowed_fields: Final = GenericLiteLLMParams.model_fields.keys()
if litellm_params:
litellm_dict: Final = litellm_params.model_dump(exclude_none=True)
litellm_dict: Final[_JsonDict] = litellm_params.model_dump(exclude_none=True)
for key, value in litellm_dict.items():
if key in allowed_fields:
completion_request_dict[key] = value
@ -346,12 +380,12 @@ class GoogleGenAIAdapter:
tools: Sequence[_GenAITool],
) -> list[ChatCompletionToolParam]:
"""Transform Google GenAI tools to OpenAI tools format"""
openai_tools: Final[list[dict[str, object]]] = []
openai_tools: Final = list[_JsonDict]()
for tool in tools:
if "functionDeclarations" in tool:
for func_decl in tool["functionDeclarations"]:
function_chunk: dict[str, object] = {
function_chunk: _JsonDict = {
"name": func_decl.get("name", ""),
}
@ -360,7 +394,7 @@ class GoogleGenAIAdapter:
if "parametersJsonSchema" in func_decl:
function_chunk["parameters"] = func_decl["parametersJsonSchema"]
openai_tool: dict[str, object] = {"type": "function", "function": function_chunk}
openai_tool: _JsonDict = {"type": "function", "function": function_chunk}
openai_tools.append(openai_tool)
# normalize the tool schemas
@ -391,13 +425,13 @@ class GoogleGenAIAdapter:
# Handle system instruction
if system_instruction:
system_parts: Final = system_instruction.get("parts", [])
system_parts: Final[Sequence[Mapping[str, str]]] = system_instruction.get("parts", [])
if system_parts and "text" in system_parts[0]:
messages.append(ChatCompletionSystemMessage(role="system", content=system_parts[0]["text"]))
for content in contents:
role = content.get("role", "user")
parts = content.get("parts", [])
parts: Sequence[_GenAIContentPart | str | None] = content.get("parts", [])
if role == "user":
# Handle user messages with potential function responses
@ -500,7 +534,7 @@ class GoogleGenAIAdapter:
def translate_completion_to_generate_content(
self,
response: ModelResponse,
) -> dict[str, object]:
) -> _JsonDict:
"""
Transform litellm completion response to Google GenAI generate_content format
@ -523,13 +557,13 @@ class GoogleGenAIAdapter:
parts = self._transform_openai_message_to_google_genai_parts(choice.message)
else:
# Fallback for generic choice objects
message_content = getattr(choice, "message", {}).get("content", "") or getattr(choice, "delta", {}).get(
"content", ""
)
message_content: str = getattr(choice, "message", _EMPTY_STR_MAPPING).get("content", "") or getattr(
choice, "delta", _EMPTY_STR_MAPPING
).get("content", "")
parts = [{"text": message_content}] if message_content else []
# Create Google GenAI format response
generate_content_response: Final[dict[str, object]] = {
generate_content_response: Final[_JsonDict] = {
"candidates": [
{
"content": {"parts": parts, "role": "model"},
@ -563,7 +597,7 @@ class GoogleGenAIAdapter:
self,
response: ModelResponse | ModelResponseStream,
wrapper: GoogleGenAIStreamWrapper,
) -> dict[str, object] | None:
) -> Mapping[str, object] | None:
"""
Transform streaming litellm completion chunk to Google GenAI generate_content format
@ -590,7 +624,7 @@ class GoogleGenAIAdapter:
finish_reason: str | None = getattr(choice, "finish_reason", None)
else:
# Fallback for generic choice objects
message_content: Final = getattr(choice, "delta", {}).get("content", "")
message_content: Final[str] = getattr(choice, "delta", _EMPTY_STR_MAPPING).get("content", "")
parts = [{"text": message_content}] if message_content else []
finish_reason = getattr(choice, "finish_reason", None)
@ -599,7 +633,7 @@ class GoogleGenAIAdapter:
return None
# Create Google GenAI streaming format response
streaming_chunk: Final[dict[str, object]] = {
streaming_chunk: Final[_JsonDict] = {
"candidates": [
{
"content": {"parts": parts, "role": "model"},
@ -635,10 +669,10 @@ class GoogleGenAIAdapter:
def _transform_openai_message_to_google_genai_parts(
self,
message: Any,
) -> list[_GenAIPart]:
message: Message,
) -> Sequence[_GenAIPart]:
"""Transform OpenAI message to Google GenAI parts format"""
parts: Final[list[_GenAIPart]] = []
parts: Final = list[_GenAIPart]()
# Add text content if present
if hasattr(message, "content") and message.content:
@ -646,20 +680,22 @@ class GoogleGenAIAdapter:
# Add tool calls if present
if hasattr(message, "tool_calls") and message.tool_calls:
for tool_call in message.tool_calls:
if hasattr(tool_call, "function") and tool_call.function:
tool_calls: Final[Sequence[ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall]] = (
message.tool_calls
)
for tool_call in tool_calls:
function: Function | None = getattr(tool_call, "function", None)
if function:
try:
args = (
_decode_tool_call_arguments(tool_call.function.arguments)
if tool_call.function.arguments
else {}
args: Mapping[str, object] = (
self._parse_tool_call_args(function.arguments) if function.arguments else {}
)
except json.JSONDecodeError:
args = {}
function_call_part: _GenAIPart = {
"functionCall": {
"name": tool_call.function.name or "undefined_tool_name",
"name": function.name or "undefined_tool_name",
"args": args,
}
}
@ -668,21 +704,23 @@ class GoogleGenAIAdapter:
return parts if parts else [{"text": ""}]
def _transform_openai_delta_to_google_genai_parts_with_accumulation(
self, delta: Any, wrapper: GoogleGenAIStreamWrapper
) -> list[_GenAIPart]:
self, delta: Delta, wrapper: GoogleGenAIStreamWrapper
) -> Sequence[_GenAIPart]:
"""Transforms OpenAI delta to Google GenAI parts, accumulating streaming tool calls."""
# 1. Initialize wrapper state if it doesn't exist
if not hasattr(wrapper, "accumulated_tool_calls"):
wrapper.accumulated_tool_calls = {}
parts: Final[list[_GenAIPart]] = []
parts: Final = list[_GenAIPart]()
if hasattr(delta, "content") and delta.content:
parts.append({"text": delta.content})
# 2. Ensure tool_calls is iterable
tool_calls: Final = delta.tool_calls or []
tool_calls: Final[Sequence[ChatCompletionDeltaToolCall | ChatCompletionDeltaCustomToolCall]] = (
delta.tool_calls or []
)
for tool_call in tool_calls:
if not hasattr(tool_call, "function"):
@ -701,19 +739,20 @@ class GoogleGenAIAdapter:
}
# Accumulate name and arguments
function_name = getattr(tool_call.function, "name", None)
args_chunk = getattr(tool_call.function, "arguments", None)
delta_function: Function | None = getattr(tool_call, "function", None)
function_name: str | None = getattr(delta_function, "name", None)
args_chunk: str | None = getattr(delta_function, "arguments", None)
# Optimization: Skip chunks that have no new data
if not function_name and not args_chunk:
verbose_logger.debug("Skipping empty tool call chunk for index: %s", tool_call_index)
continue
if function_name:
wrapper.accumulated_tool_calls[tool_call_index]["name"] = function_name
if args_chunk:
wrapper.accumulated_tool_calls[tool_call_index]["arguments"] += args_chunk
previous_data: _ToolCallAccumulator = wrapper.accumulated_tool_calls[tool_call_index]
wrapper.accumulated_tool_calls[tool_call_index] = _ToolCallAccumulator(
name=function_name or previous_data["name"],
arguments=previous_data["arguments"] + (args_chunk or ""),
)
# Attempt to parse and emit a complete tool call
accumulated_data = wrapper.accumulated_tool_calls[tool_call_index]
@ -723,7 +762,7 @@ class GoogleGenAIAdapter:
# 5. Attempt to parse arguments even if name hasn't arrived.
try:
# Attempt to parse the accumulated arguments string
parsed_args = _decode_tool_call_arguments(accumulated_args)
parsed_args: Mapping[str, object] = self._parse_tool_call_args(accumulated_args)
# If parsing succeeds, but we don't have a name yet, wait.
# The part will be created by a later chunk that brings the name.
@ -757,7 +796,7 @@ class GoogleGenAIAdapter:
return mapping.get(finish_reason, "STOP")
def _map_usage(self, usage: Usage | None) -> dict[str, int]:
def _map_usage(self, usage: object) -> Mapping[str, int]:
"""Map OpenAI usage to Google GenAI usage format"""
return {
"promptTokenCount": getattr(usage, "prompt_tokens", 0) or 0,

211
litellm/integrations/rubrik.py
View file

@ -8,11 +8,12 @@ import uuid
from collections import Counter
from collections.abc import Awaitable, Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict
from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, TypedDict, overload
import httpx
from typing_extensions import Never, ReadOnly
from typing_extensions import Never, ReadOnly, Required
from litellm._logging import verbose_logger
from litellm.integrations.custom_batch_logger import CustomBatchLogger
@ -30,6 +31,7 @@ from litellm.llms.custom_httpx.http_handler import (
httpxSpecialProvider,
)
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk
from litellm.types.utils import (
ChatCompletionMessageToolCall,
Function,
@ -52,17 +54,102 @@ _DROP_WARNING_INTERVAL_SECONDS: Final = 60.0
_EMPTY_MAPPING: Final[Mapping[str, Never]] = MappingProxyType({})
class _ServiceToolCall(TypedDict):
id: ReadOnly[str]
class _ModerationToolCall(TypedDict, total=False):
id: ReadOnly[Required[str]]
class _ServiceMessage(TypedDict, total=False):
class _ModerationMessage(TypedDict, total=False):
content: ReadOnly[str | None]
tool_calls: ReadOnly[Sequence[_ModerationToolCall] | None]
class _ModerationChoice(TypedDict, total=False):
message: ReadOnly[_ModerationMessage | None]
class _ModerationResponse(TypedDict, total=False):
choices: ReadOnly[Sequence[_ModerationChoice]]
class _LogEventKwargs(TypedDict, total=False):
standard_logging_object: ReadOnly[Required[StandardLoggingPayload]]
litellm_call_id: ReadOnly[str]
class _HasCallId(Protocol):
def get(self, key: Literal["litellm_call_id"], /) -> str | None: ...
class _HasModelAttr(Protocol):
model: str | None
class _ResponseSource(Protocol):
def get(self, key: Literal["response"], /) -> "_HasModelAttr | None": ...
class _ModelSource(Protocol):
def get(self, key: Literal["model"], default: str, /) -> str: ...
class _FallbackSource(Protocol):
@overload
def get(self, key: Literal["start_time"], /) -> datetime | None: ...
@overload
def get(self, key: str, /) -> object | None: ...
class _RequestContextSource(Protocol):
@overload
def get(self, key: Literal["optional_params"], /) -> Mapping[str, object] | None: ...
@overload
def get(self, key: str, /) -> object | None: ...
def __contains__(self, key: object, /) -> bool: ...
def __getitem__(self, key: str, /) -> object: ...
class _ToolCallLike(Protocol):
id: str | None
type: str | None
function: Function
class _ModerationSourceToolCall(TypedDict, total=False):
function: ReadOnly[Mapping[str, object] | None]
class _ModerationSourceMessage(TypedDict, total=False):
role: ReadOnly[str]
function_call: ReadOnly[Mapping[str, object] | None]
tool_calls: ReadOnly[Sequence[_ModerationSourceToolCall | None] | None]
class _FlattenedModerationMessage(TypedDict):
role: ReadOnly[str | None]
content: ReadOnly[str]
tool_calls: ReadOnly[Sequence[_ServiceToolCall]]
class _ServiceChoice(TypedDict, total=False):
message: ReadOnly[_ServiceMessage]
class _CorrelatablePayload(TypedDict):
id: str # writable-ok: _apply_correlation_id overwrites the provider id on a deep-copied payload
class _SystemPromptCarrier(TypedDict, total=False):
messages: object # writable-ok: _prepend_system_prompt rebinds messages on the copied payload by design
class _BlockFailurePayload(TypedDict, total=False):
id: object # writable-ok: correlation id is pinned after copying the base payload
model: ReadOnly[object]
model_group: ReadOnly[object]
model_id: ReadOnly[str]
model_parameters: ReadOnly[object]
startTime: ReadOnly[float | None]
endTime: ReadOnly[float | None]
completionStartTime: ReadOnly[float | None]
messages: object # writable-ok: passed to _prepend_system_prompt, which rebinds messages
metadata: ReadOnly[StandardLoggingUserAPIKeyMetadata]
response: str # writable-ok: block failure text replaces the copied response
status: ReadOnly[str]
class _MalformedToolBlockingResponseError(Exception):
@ -385,7 +472,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
@staticmethod
def _stash_block_context(
logging_obj: Optional["LiteLLMLoggingObj"],
request_data: dict,
request_data: dict[str, object],
) -> None:
"""Stash signals so the deferred success-event skips this request and
``async_post_call_failure_hook`` can build the failure payload.
@ -414,12 +501,16 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
request_data["_rubrik_logging_obj"] = logging_obj
@staticmethod
def _normalize_tool_calls(tool_calls: Sequence[object]) -> tuple[ChatCompletionMessageToolCall, ...]:
def _normalize_tool_calls(
tool_calls: Sequence[ChatCompletionToolCallChunk | ChatCompletionMessageToolCall | _ToolCallLike],
) -> tuple[ChatCompletionMessageToolCall, ...]:
"""Convert tool_calls from inputs to ChatCompletionMessageToolCall objects."""
return tuple(RubrikLogger._normalize_tool_call(tc) for tc in tool_calls)
@staticmethod
def _normalize_tool_call(tc: Any) -> ChatCompletionMessageToolCall:
def _normalize_tool_call(
tc: ChatCompletionToolCallChunk | ChatCompletionMessageToolCall | _ToolCallLike,
) -> ChatCompletionMessageToolCall:
if isinstance(tc, ChatCompletionMessageToolCall):
return tc
if isinstance(tc, dict):
@ -460,12 +551,15 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
``content`` is sent so the webhook can moderate the response text;
``None`` when the assistant produced no text (tool-call-only response).
"""
message: Final[dict[str, object]] = {
message: Final[Mapping[str, object]] = {
"role": "assistant",
"content": content or None,
**(
{"tool_calls": tuple(tc.model_dump(exclude_none=True) for tc in tool_calls)}
if tool_calls
else _EMPTY_MAPPING
),
}
if tool_calls:
message["tool_calls"] = tuple(tc.model_dump(exclude_none=True) for tc in tool_calls)
return {
"id": request_id or f"chatcmpl-{uuid.uuid4()}",
"object": "chat.completion",
@ -481,7 +575,9 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
}
@staticmethod
def _flatten_messages_for_moderation(messages: Sequence[object] | None) -> tuple[Mapping[str, Any], ...]:
def _flatten_messages_for_moderation(
messages: Sequence[AllMessageValues | None] | None,
) -> tuple[_FlattenedModerationMessage, ...]:
"""Collapse each message's content to a plain string for the webhook.
litellm normalizes Anthropic ``/v1/messages`` requests to OpenAI shape,
@ -502,7 +598,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
)
@staticmethod
def _moderation_text_parts(message: Mapping[str, Any]) -> tuple[str, ...]:
def _moderation_text_parts(message: _ModerationSourceMessage) -> tuple[str, ...]:
"""Every attacker-controlled text segment of a message: its content plus
the arguments of any tool call or deprecated function call."""
fc: Final = message.get("function_call")
@ -530,16 +626,8 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
``/v1/messages`` requests too. Optional fields are sent only when
present so the payload stays clean.
"""
payload: Final[dict[str, object]] = {
"model": inputs.get("model") or request_data.get("model") or "",
"messages": RubrikLogger._flatten_messages_for_moderation(inputs.get("structured_messages")),
}
tools: Final = inputs.get("tools")
if tools is not None:
payload["tools"] = tools
user: Final = request_data.get("user")
if user:
payload["user"] = user
# Fall back to litellm_call_id, the stable cross-provider join key the
# response/tool path uses (see _correlation_id). LiteLLM does not
# populate request_data["correlation_key"]; it carries litellm_call_id.
@ -547,14 +635,18 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
# when correlation_key is empty, so without this the block fires but no
# log is ever written. An explicit correlation_key still wins.
correlation_key: Final = request_data.get("correlation_key") or request_data.get("litellm_call_id")
if correlation_key:
payload["correlation_key"] = correlation_key
return payload
return {
"model": inputs.get("model") or request_data.get("model") or "",
"messages": RubrikLogger._flatten_messages_for_moderation(inputs.get("structured_messages")),
**({"tools": tools} if tools is not None else _EMPTY_MAPPING),
**({"user": user} if user else _EMPTY_MAPPING),
**({"correlation_key": correlation_key} if correlation_key else _EMPTY_MAPPING),
}
@staticmethod
def _extract_request_data(
call_details: Mapping[str, Any],
request_data: Mapping[str, object] | None,
call_details: _RequestContextSource,
request_data: _RequestContextSource | None,
) -> Mapping[str, object]:
"""Extract original request data from model_call_details for the
response moderation service envelope.
@ -590,7 +682,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
}
@staticmethod
def _sanitize_proxy_server_request(proxy_server_request: object) -> object:
def _sanitize_proxy_server_request(proxy_server_request: Mapping[str, object] | str | None) -> object:
"""Allowlist only routing fields (``url``, ``method``) when forwarding
``proxy_server_request`` to an external webhook, dropping inbound
``headers`` (Authorization, Cookie, x-api-key, ...) and the raw
@ -600,18 +692,19 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
return {key: proxy_server_request[key] for key in ("url", "method") if key in proxy_server_request}
@staticmethod
def _resolve_model(request_data: Mapping[str, object], call_details: Mapping[str, str]) -> str:
def _resolve_model(request_data: _ResponseSource, call_details: _ModelSource) -> str:
"""Get the model name for the ModifyResponseException."""
response: Final = request_data.get("response")
if response and hasattr(response, "model"):
response_model: Final[str | None] = getattr(response, "model", None)
return response_model or "unknown"
return response.model or "unknown"
return call_details.get("model", "unknown")
# -- Logging hooks ---------------------------------------------------------
@staticmethod
def _correlation_id(call_details: Mapping[str, str], request_data: Mapping[str, str] | None = None) -> str | None:
def _correlation_id(
call_details: _HasCallId | _LogEventKwargs, request_data: _HasCallId | None = None
) -> str | None:
"""The id that joins a blocked request's two S3 logs by filename: the
moderation (``_blocking``) log and the failure (response) log.
@ -625,7 +718,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
return call_details.get("litellm_call_id") or (request_data or _EMPTY_MAPPING).get("litellm_call_id")
@classmethod
def _apply_correlation_id(cls, payload: dict[str, object], source: Mapping[str, str]) -> None:
def _apply_correlation_id(cls, payload: _CorrelatablePayload, source: _HasCallId | _LogEventKwargs) -> None:
"""Pin ``payload["id"]`` to ``litellm_call_id`` in place so this log
shares its S3 filename id with the moderation (``_blocking``) and
failure logs for the same request -- for every provider.
@ -645,7 +738,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
payload["id"] = correlated
@staticmethod
def _prepend_system_prompt(payload: dict[str, object], source: Mapping[str, object]) -> None:
def _prepend_system_prompt(payload: _SystemPromptCarrier, source: Mapping[str, object]) -> None:
"""Prepend ``source["system"]`` onto ``payload["messages"]``.
Builds a NEW messages list rather than mutating ``payload["messages"]``
@ -673,9 +766,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
exc_info=True,
)
async def _prepare_log_payload(
self, kwargs: Mapping[str, object], event_type: str
) -> StandardLoggingPayload | None:
async def _prepare_log_payload(self, kwargs: _LogEventKwargs, event_type: str) -> StandardLoggingPayload | None:
"""Shared logic for success logging (sampled)."""
if random.random() > self.sampling_rate:
verbose_logger.debug("Skipping Rubrik %s logging (sampling_rate=%s)", event_type, self.sampling_rate)
@ -684,12 +775,12 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
# Deep-copy so mutations don't affect other callbacks sharing this object
standard_logging_payload: Final[StandardLoggingPayload] = safe_deep_copy(kwargs["standard_logging_object"])
self._apply_correlation_id(standard_logging_payload, kwargs) # pyright: ignore[reportArgumentType] # StandardLoggingPayload is dict[str,Any] at runtime
self._apply_correlation_id(standard_logging_payload, kwargs)
self._prepend_system_prompt(standard_logging_payload, kwargs) # pyright: ignore[reportArgumentType] # StandardLoggingPayload is dict[str,Any] at runtime
return standard_logging_payload
async def _append_and_maybe_flush(self, payload) -> None:
async def _append_and_maybe_flush(self, payload: Mapping[str, object]) -> None:
self._ensure_periodic_flush_task()
self.log_queue.append(payload)
self._enforce_max_queue_size()
@ -714,7 +805,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
self._dropped_since_warning = 0
self._last_drop_warning_time = now
async def _enqueue_log_event(self, kwargs: Mapping[str, object], event_type: str):
async def _enqueue_log_event(self, kwargs: _LogEventKwargs, event_type: str):
try:
payload: Final = await self._prepare_log_payload(kwargs, event_type)
if payload is None:
@ -835,7 +926,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
logging_obj: "LiteLLMLoggingObj",
exception: "ModifyResponseException",
user_api_key_dict: "UserAPIKeyAuth",
) -> StandardLoggingPayload:
) -> _BlockFailurePayload:
"""Build a failure-style payload using the exception text as response.
Blocked-tool events are security-relevant and **bypass sampling**:
@ -877,9 +968,9 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
call_details: Final = logging_obj.model_call_details
exception_text: Final = f"{type(exception).__name__}: {exception.message}"
base: Final = call_details.get("standard_logging_object")
base: Final[StandardLoggingPayload | None] = call_details.get("standard_logging_object")
if base is not None:
payload: dict[str, object] = safe_deep_copy(base)
payload: _BlockFailurePayload = self._copy_block_payload_base(base)
else:
verbose_logger.debug(
"Rubrik: standard_logging_object not yet on model_call_details "
@ -901,6 +992,10 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
return payload
@staticmethod
def _copy_block_payload_base(base: StandardLoggingPayload) -> _BlockFailurePayload:
return safe_deep_copy(base)
@staticmethod
def _caller_metadata(user_api_key_dict: "UserAPIKeyAuth") -> StandardLoggingUserAPIKeyMetadata:
"""Identify the caller whose request was blocked.
@ -923,9 +1018,9 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
@classmethod
def _build_fallback_payload(
cls,
call_details: Mapping[str, Any],
call_details: _FallbackSource,
user_api_key_dict: "UserAPIKeyAuth",
) -> dict[str, object]:
) -> _BlockFailurePayload:
# Convert datetime to a Unix float so json.dumps can serialize it.
# httpx's json= parameter uses stdlib json.dumps with no custom encoder.
_raw_start: Final = call_details.get("start_time")
@ -959,7 +1054,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
response: Final = await self.async_httpx_client.post(
url=self.logging_endpoint,
json=data,
headers=self._headers,
headers=dict(self._headers),
)
response.raise_for_status()
except httpx.HTTPStatusError as e:
@ -1013,7 +1108,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
# -- Webhook services ------------------------------------------------------
async def _post_json(self, endpoint: str, payload: Mapping[str, object], service_name: str) -> Mapping[str, Any]:
async def _post_json(self, endpoint: str, payload: Mapping[str, object], service_name: str) -> _ModerationResponse:
"""POST ``payload`` to a Rubrik webhook and return its dict response.
Raises:
@ -1023,11 +1118,11 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
verbose_logger.debug("Sending request to %s: %s", service_name, endpoint)
http_response: Final = await self.moderation_client.post(
endpoint,
json=payload,
headers=self._headers,
json=dict(payload),
headers=dict(self._headers),
)
http_response.raise_for_status()
result: Final[object] = http_response.json()
result: Final[_ModerationResponse | None] = http_response.json()
if not isinstance(result, dict):
raise TypeError(
f"{service_name} returned non-dict JSON "
@ -1040,7 +1135,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
self,
response_data: Mapping[str, object],
request_data: Mapping[str, object],
) -> Mapping[str, Any]:
) -> _ModerationResponse:
"""Post the ``{request, response}`` envelope to the after_completion
webhook and return its (possibly rewritten) response.
@ -1056,7 +1151,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
"Response moderation service",
)
async def _post_to_prompt_moderation_endpoint(self, payload: Mapping[str, object]) -> Mapping[str, Any]:
async def _post_to_prompt_moderation_endpoint(self, payload: Mapping[str, object]) -> _ModerationResponse:
"""Post a bare OpenAI request to the before_prompt webhook.
Returns ``{}`` (passthrough) or a synthetic chat.completion (block).
@ -1064,14 +1159,14 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
return await self._post_json(self.prompt_moderation_endpoint, payload, "Prompt moderation service")
@staticmethod
def _extract_prompt_refusal(service_response: Mapping[str, Any]) -> str | None:
def _extract_prompt_refusal(service_response: _ModerationResponse) -> str | None:
"""Return the refusal text when the prompt was blocked, else None.
The before_prompt webhook returns ``{}`` (passthrough) or a synthetic
chat.completion whose ``choices[0].message.content`` is the refusal
explanation.
"""
choices: Final[Sequence[_ServiceChoice] | None] = service_response.get("choices")
choices: Final = service_response.get("choices")
if not choices:
return None
message: Final = choices[0].get("message") or _EMPTY_MAPPING
@ -1080,7 +1175,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
@staticmethod
def _extract_response_block(
service_response: Mapping[str, Any],
service_response: _ModerationResponse,
all_tool_calls: Sequence[ChatCompletionMessageToolCall],
sent_content: str,
) -> BlockedResponseResult | None:
@ -1103,7 +1198,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
Expects service_response in OpenAI chat completion format:
{"choices": [{"message": {"tool_calls": [...], "content": "..."}}]}
"""
choices: Final[Sequence[_ServiceChoice]] = service_response.get("choices") or ()
choices: Final = service_response.get("choices") or ()
if not choices:
raise _MalformedToolBlockingResponseError("Response moderation service returned empty response")

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

@ -10,7 +10,7 @@ import asyncio
import math
import uuid
from collections.abc import AsyncIterator, Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, TypedDict, TypeVar, cast
from typing import TYPE_CHECKING, Any, Final, Literal, Never, TypedDict, TypeVar, cast
from typing_extensions import ReadOnly
@ -46,7 +46,13 @@ from litellm.types.integrations.websearch_interception import (
AnthropicServerToolUseBlock,
WebSearchInterceptionConfig,
)
from litellm.types.llms.openai import AllMessageValues
from litellm.types.llms.anthropic import AnthropicThinkingParam
from litellm.types.llms.openai import (
AllMessageValues,
ChatCompletionAudioParam,
ChatCompletionPredictionContentParam,
OpenAIWebSearchOptions,
)
from litellm.types.utils import (
AgenticLoopParams,
CallTypes,
@ -56,6 +62,8 @@ from litellm.types.utils import (
from litellm.utils import ProviderConfigManager
if TYPE_CHECKING:
from aiohttp import ClientSession
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.anthropic_messages.transformation import (
BaseAnthropicMessagesConfig,
@ -94,23 +102,98 @@ class _WebSearchSettingsView(TypedDict):
websearch_interception_params: WebSearchInterceptionConfig
class _SearchToolLitellmParams(TypedDict, total=False):
search_provider: ReadOnly[str | None]
class _SearchToolConfig(TypedDict, total=False):
search_tool_name: str
litellm_params: Mapping[str, object] | None
litellm_params: ReadOnly[_SearchToolLitellmParams | None]
class _DeploymentKwargsView(TypedDict):
"""Typed reads of the untyped request kwargs seen by the deployment hook."""
class _LitellmParamsProviderView(TypedDict, total=False):
custom_llm_provider: ReadOnly[str]
litellm_params: ReadOnly[Mapping[str, object]]
class _DeploymentCallKwargsView(TypedDict):
custom_llm_provider: ReadOnly[str]
litellm_params: ReadOnly[_LitellmParamsProviderView]
model: ReadOnly[str]
class _UserAuthView(TypedDict):
"""Typed read of the optional team attached to the caller's auth object."""
class _AcreateNamedParams(TypedDict, total=False):
metadata: ReadOnly[Never]
stop_sequences: ReadOnly[Never]
stream: ReadOnly[bool | None]
system: ReadOnly[str | None]
temperature: ReadOnly[float | None]
thinking: ReadOnly[Never]
tool_choice: ReadOnly[Never]
tools: ReadOnly[Never]
top_k: ReadOnly[int | None]
top_p: ReadOnly[float | None]
container: ReadOnly[Never]
team_id: ReadOnly[str | None]
class _AsearchNamedParams(TypedDict, total=False):
max_results: ReadOnly[int | None]
search_domain_filter: ReadOnly[Never]
max_tokens_per_page: ReadOnly[int | None]
country: ReadOnly[str | None]
api_key: ReadOnly[str | None]
api_base: ReadOnly[str | None]
timeout: ReadOnly[float | None]
extra_headers: ReadOnly[Never]
class _AcompletionNamedParams(TypedDict, total=False):
functions: ReadOnly[Never]
function_call: ReadOnly[str | None]
timeout: ReadOnly[float | None]
temperature: ReadOnly[float | None]
top_p: ReadOnly[float | None]
n: ReadOnly[int | None]
stream: ReadOnly[bool | None]
stream_options: ReadOnly[Never]
stop: ReadOnly[Never]
max_tokens: ReadOnly[int | None]
max_completion_tokens: ReadOnly[int | None]
modalities: ReadOnly[Never]
prediction: ReadOnly[ChatCompletionPredictionContentParam | None]
audio: ReadOnly[ChatCompletionAudioParam | None]
presence_penalty: ReadOnly[float | None]
frequency_penalty: ReadOnly[float | None]
logit_bias: ReadOnly[Never]
user: ReadOnly[str | None]
response_format: ReadOnly[Never]
seed: ReadOnly[int | None]
tools: ReadOnly[Never]
tool_choice: ReadOnly[Never]
parallel_tool_calls: ReadOnly[bool | None]
logprobs: ReadOnly[bool | None]
top_logprobs: ReadOnly[int | None]
deployment_id: ReadOnly[str | None]
reasoning_effort: ReadOnly[Literal["none", "minimal", "low", "medium", "high", "xhigh", "default"] | None]
verbosity: ReadOnly[Literal["low", "medium", "high"] | None]
safety_identifier: ReadOnly[str | None]
service_tier: ReadOnly[str | None]
store: ReadOnly[bool | None]
prompt_cache_key: ReadOnly[str | None]
base_url: ReadOnly[str | None]
api_version: ReadOnly[str | None]
api_key: ReadOnly[str | None]
model_list: ReadOnly[Never]
extra_headers: ReadOnly[Never]
thinking: ReadOnly[AnthropicThinkingParam | None]
web_search_options: ReadOnly[OpenAIWebSearchOptions | None]
include_server_side_tool_invocations: ReadOnly[bool | None]
shared_session: ReadOnly["ClientSession | None"]
enable_json_schema_validation: ReadOnly[bool | None]
_NO_ACREATE_NAMED: Final[_AcreateNamedParams] = {}
_NO_ASEARCH_NAMED: Final[_AsearchNamedParams] = {}
_NO_ACOMPLETION_NAMED: Final[_AcompletionNamedParams] = {}
class WebSearchInterceptionLogger(CustomLogger):
@ -312,17 +395,17 @@ class WebSearchInterceptionLogger(CustomLogger):
"""
# Check if this is for an enabled provider
# Try top-level kwargs first, then nested litellm_params, then derive from model name
kwargs_view: Final[_DeploymentKwargsView] = {
call_kwargs_view: Final[_DeploymentCallKwargsView] = {
"custom_llm_provider": kwargs.get("custom_llm_provider", ""),
"litellm_params": kwargs.get("litellm_params", {}),
"model": kwargs.get("model", ""),
}
custom_llm_provider = kwargs_view["custom_llm_provider"] or kwargs_view["litellm_params"].get(
custom_llm_provider = call_kwargs_view["custom_llm_provider"] or call_kwargs_view["litellm_params"].get(
"custom_llm_provider", ""
)
if not custom_llm_provider:
try:
_, custom_llm_provider, _, _ = litellm.get_llm_provider(model=kwargs_view["model"])
_, custom_llm_provider, _, _ = litellm.get_llm_provider(model=call_kwargs_view["model"])
except Exception:
custom_llm_provider = ""
if custom_llm_provider not in self.enabled_providers:
@ -1218,10 +1301,10 @@ class WebSearchInterceptionLogger(CustomLogger):
messages: list[dict],
tool_calls: list[dict],
thinking_blocks: list[dict],
anthropic_messages_optional_request_params: dict,
anthropic_messages_optional_request_params: Mapping[str, object],
logging_obj: "LiteLLMLoggingObj | None",
stream: bool,
kwargs: dict,
kwargs: Mapping[str, object],
) -> "AnthropicMessagesResponse | AsyncIterator[object]":
"""Legacy path: execute search + build patch + run follow-up call."""
request_patch, structured_results = await self._build_anthropic_request_patch(
@ -1229,9 +1312,9 @@ class WebSearchInterceptionLogger(CustomLogger):
messages=messages,
tool_calls=tool_calls,
thinking_blocks=thinking_blocks,
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
anthropic_messages_optional_request_params=dict[str, object](anthropic_messages_optional_request_params),
logging_obj=logging_obj,
kwargs=kwargs,
kwargs=dict[str, object](kwargs),
)
if request_patch.messages is None:
raise ValueError("WebSearchInterception: missing follow-up messages")
@ -1246,12 +1329,14 @@ class WebSearchInterceptionLogger(CustomLogger):
if max_tokens is None:
max_tokens = cast(int, kwargs.get("max_tokens", 1024))
patch_kwargs: Final = dict[str, object](request_patch.kwargs)
response: AnthropicMessagesResponse | AsyncIterator[object] = await anthropic_messages.acreate(
max_tokens=max_tokens,
messages=request_patch.messages,
model=request_patch.model or model,
**_NO_ACREATE_NAMED,
**optional_params,
**request_patch.kwargs,
**patch_kwargs,
)
# Legacy path: the new path goes through the typed plan + core
@ -1393,12 +1478,13 @@ class WebSearchInterceptionLogger(CustomLogger):
search_tool: Final = self._select_search_tool_from_router(llm_router=llm_router)
search_provider: str | None = None
search_litellm_params: dict[str, Any] = {}
search_litellm_params: Mapping[str, object] = {}
search_tool_name: Final = self._selected_search_tool_name(search_tool=search_tool)
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_provider = search_litellm_params.get("search_provider")
tool_params: Final[_SearchToolLitellmParams] = search_tool.get("litellm_params", {}) or {}
search_litellm_params = dict[str, object](tool_params)
search_provider = tool_params.get("search_provider")
# Fallback to perplexity if no router or no search tools configured
if not search_provider:
@ -1426,12 +1512,15 @@ class WebSearchInterceptionLogger(CustomLogger):
if key != "search_provider" and value is not None
}
result: Final = (
await litellm.asearch(query=query, search_provider=search_provider, **search_kwargs)
await litellm.asearch(
query=query, search_provider=search_provider, **_NO_ASEARCH_NAMED, **search_kwargs
)
if search_metadata is None
else await litellm.asearch(
query=query,
search_provider=search_provider,
litellm_metadata=search_metadata,
**_NO_ASEARCH_NAMED,
**search_kwargs,
)
)
@ -1471,8 +1560,7 @@ class WebSearchInterceptionLogger(CustomLogger):
valid_token=user_api_key_auth,
)
auth_view: Final[_UserAuthView] = {"team_id": getattr(user_api_key_auth, "team_id", None)}
team_id: Final = auth_view["team_id"]
team_id: Final[str | None] = getattr(user_api_key_auth, "team_id", None)
if team_id:
from litellm.proxy.proxy_server import (
prisma_client,
@ -1587,10 +1675,10 @@ class WebSearchInterceptionLogger(CustomLogger):
model: str,
messages: list[dict],
tool_calls: list[dict],
optional_params: dict,
optional_params: Mapping[str, object],
logging_obj: "LiteLLMLoggingObj | None",
stream: bool,
kwargs: dict,
kwargs: Mapping[str, object],
response_format: str = "openai",
) -> "ModelResponse | CustomStreamWrapper":
"""Legacy path: execute search + build patch + run follow-up call."""
@ -1598,8 +1686,8 @@ class WebSearchInterceptionLogger(CustomLogger):
model=model,
messages=messages,
tool_calls=tool_calls,
optional_params=optional_params,
kwargs=kwargs,
optional_params=dict[str, object](optional_params),
kwargs=dict[str, object](kwargs),
response_format=response_format,
)
if request_patch.messages is None:
@ -1607,11 +1695,13 @@ class WebSearchInterceptionLogger(CustomLogger):
params: Final = dict(optional_params)
params.update(request_patch.optional_params)
params.pop("tool_choice", None)
patch_kwargs: Final = dict[str, object](request_patch.kwargs)
return await litellm.acompletion(
model=request_patch.model or model,
messages=request_patch.messages,
**_NO_ACOMPLETION_NAMED,
**params,
**request_patch.kwargs,
**patch_kwargs,
)
async def _build_chat_completion_request_patch(

171
litellm/llms/anthropic/chat/guardrail_translation/handler.py
View file

@ -16,9 +16,9 @@ import json
from collections.abc import Mapping, Sequence
from copy import deepcopy
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Final, cast
from typing import TYPE_CHECKING, Any, Final, Protocol, cast, overload, runtime_checkable
from typing_extensions import assert_never
from typing_extensions import ReadOnly, TypedDict, assert_never
from litellm._logging import verbose_proxy_logger
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
@ -98,6 +98,48 @@ InputWriteBackTarget = (
)
class _SSEDelta(TypedDict, total=False):
type: ReadOnly[str]
text: ReadOnly[str]
stop_reason: ReadOnly[str | None]
class _SSEEventData(TypedDict, total=False):
delta: ReadOnly[_SSEDelta]
def _as_str_mapping(value: Mapping[str, object]) -> Mapping[str, object]:
return value
def _content_block_at(blocks: Sequence[object], index: int) -> object:
return blocks[index]
@runtime_checkable
class _ModelDumpBlock(Protocol):
def model_dump(self) -> Mapping[str, object]: ...
@runtime_checkable
class _TextAttrBlock(Protocol):
text: str
class _WritableMessage(Protocol):
@overload
def get(self, key: str, /) -> object | None: ...
@overload
def get(self, key: str, default: object, /) -> object: ...
def __setitem__(self, key: str, value: object, /) -> None: ...
def _as_writable(value: _WritableMessage) -> _WritableMessage:
return value
@dataclass(frozen=True, slots=True)
class ScannedText:
text: str
@ -126,7 +168,7 @@ class AnthropicMessagesHandler(BaseTranslation):
@staticmethod
def _build_streaming_usage_response(
responses_so_far: list[object],
responses_so_far: Sequence[object],
request_data: dict | None,
) -> ModelResponse | None:
chunks: Final = tuple(response for response in responses_so_far if isinstance(response, (str, bytes)))
@ -144,7 +186,7 @@ class AnthropicMessagesHandler(BaseTranslation):
self,
exc: "ModifyResponseException",
stream_started: bool = False,
responses_so_far: list[object] | None = None,
responses_so_far: Sequence[object] | None = None,
) -> list[bytes]:
"""
Build an Anthropic SSE sequence delivering the guardrail block message
@ -162,7 +204,7 @@ class AnthropicMessagesHandler(BaseTranslation):
would make Anthropic clients reject the stream.
"""
if stream_started:
return self._block_continuation_chunks(exc, responses_so_far or [])
return list(self._block_continuation_chunks(exc, responses_so_far or []))
return self._standalone_block_chunks(exc)
def _standalone_block_chunks(self, exc: "ModifyResponseException") -> list[bytes]:
@ -187,7 +229,9 @@ class AnthropicMessagesHandler(BaseTranslation):
)
return list(FakeAnthropicMessagesStreamIterator(response=block_response))
def _block_continuation_chunks(self, exc: "ModifyResponseException", responses_so_far: list[object]) -> list[bytes]:
def _block_continuation_chunks(
self, exc: "ModifyResponseException", responses_so_far: Sequence[object]
) -> Sequence[bytes]:
"""Continue an already-started message: close the open content block,
append the block message as a new text block, then end the message --
without a second message_start."""
@ -199,7 +243,7 @@ class AnthropicMessagesHandler(BaseTranslation):
def _sse(event_type: str, payload: dict) -> bytes:
return f"event: {event_type}\ndata: {json.dumps(payload)}\n\n".encode()
output_tokens: Final = blocked_response_usage(getattr(exc, "original_response", None))["output_tokens"]
output_tokens: Final = blocked_response_usage(getattr(exc, "original_response", None)).get("output_tokens", 0)
open_index, max_index = self._content_block_state(responses_so_far)
new_index: Final = (max_index + 1) if max_index is not None else 0
chunks: list[bytes] = []
@ -237,7 +281,7 @@ class AnthropicMessagesHandler(BaseTranslation):
@staticmethod
def _content_block_state(
responses_so_far: list[object],
responses_so_far: Sequence[object],
) -> tuple[int | None, int | None]:
"""From the SSE chunks already sent to the client, return (open
content-block index or None, highest content-block index seen or None).
@ -263,7 +307,20 @@ class AnthropicMessagesHandler(BaseTranslation):
return open_index, max_index
@staticmethod
def _iter_sse_events(item: object) -> list[dict[str, object]]:
def _parse_sse_data_line(raw_line: str) -> tuple[Mapping[str, object], ...]:
line: Final = raw_line.strip()
if not line.startswith("data:"):
return ()
try:
parsed: Final[object] = json.loads(line[len("data:") :].strip())
except json.JSONDecodeError:
return ()
if not isinstance(parsed, dict):
return ()
return (_as_str_mapping(parsed),)
@staticmethod
def _iter_sse_events(item: object) -> Sequence[Mapping[str, object]]:
"""Yield the event-data dicts in one stream chunk.
Handles both formats this stream can carry (see
@ -271,24 +328,15 @@ class AnthropicMessagesHandler(BaseTranslation):
several events separated by a blank line -- and an already-parsed event
``dict``."""
if isinstance(item, dict):
return [item]
return (_as_str_mapping(item),)
if not isinstance(item, (bytes, bytearray)):
return []
events: Final[list[dict[str, object]]] = []
for block in item.decode("utf-8", errors="replace").split("\n\n"):
for line in block.split("\n"):
line = line.strip()
if not line.startswith("data:"):
continue
try:
parsed: str | int | float | bool | None | Sequence[object] | Mapping[str, object] = json.loads(
line[len("data:") :].strip()
)
except json.JSONDecodeError:
continue
if isinstance(parsed, dict):
events.append(parsed)
return events
return ()
return tuple(
event
for block in item.decode("utf-8", errors="replace").split("\n\n")
for line in block.split("\n")
for event in AnthropicMessagesHandler._parse_sse_data_line(line)
)
def _translate_to_openai(self, data: dict) -> ChatCompletionRequest:
"""Translate Anthropic request to OpenAI chat completion format."""
@ -321,7 +369,7 @@ class AnthropicMessagesHandler(BaseTranslation):
data: dict,
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
) -> Any:
) -> Mapping[str, object]:
"""
Process input messages by applying guardrails to text content.
"""
@ -481,7 +529,7 @@ class AnthropicMessagesHandler(BaseTranslation):
@staticmethod
def _openai_system_message_to_anthropic(
message: dict[str, object],
message: Mapping[str, object],
) -> dict[str, object] | None: # mutable-ok: API message payload
"""Convert an OpenAI system message to the client's Anthropic-shaped entry."""
content: Final = message.get("content")
@ -561,7 +609,7 @@ class AnthropicMessagesHandler(BaseTranslation):
@staticmethod
def _defer_systems_inside_tool_exchanges(
structured_messages: list, # mutable-ok: API message payload
structured_messages: Sequence[Mapping[str, object]],
) -> list:
"""Hold a system row until the tool exchange around it completes so the call/result pair converts together."""
from litellm.litellm_core_utils.prompt_templates.factory import group_tool_exchanges
@ -755,7 +803,7 @@ class AnthropicMessagesHandler(BaseTranslation):
if scan_only_tool_results:
return EMPTY_EXTRACTED_INPUT
text_str: Final = content_item.get("text", None)
text_str: Final[str | None] = content_item.get("text")
return ExtractedInput(
scanned=(
() if text_str is None else (ScannedText(text_str, ContentBlockTextTarget(msg_idx, content_idx)),)
@ -805,7 +853,7 @@ class AnthropicMessagesHandler(BaseTranslation):
async def _apply_guardrail_responses_to_input(
self,
messages: list[dict[str, object]],
messages: Sequence[_WritableMessage],
responses: list[str],
scanned: tuple[ScannedText, ...],
) -> None:
@ -931,7 +979,7 @@ class AnthropicMessagesHandler(BaseTranslation):
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
user_api_key_dict: "UserAPIKeyAuth | None" = None,
request_data: dict | None = None,
) -> list[Any]:
) -> Sequence[object]:
"""
Process output streaming response by applying guardrails to text content.
@ -1027,7 +1075,7 @@ class AnthropicMessagesHandler(BaseTranslation):
return request_data
@staticmethod
def _get_response_content(response: object) -> list[Any]:
def _get_response_content(response: object) -> Sequence[object]:
"""Extract content list from a dict or object response."""
if isinstance(response, dict):
return response.get("content", []) or []
@ -1037,7 +1085,7 @@ class AnthropicMessagesHandler(BaseTranslation):
def _extract_from_content_blocks(
self,
response_content: list[Any],
response_content: Sequence[object],
texts_to_check: list[str],
images_to_check: list[str],
task_mappings: list[tuple[int, int | None]],
@ -1045,21 +1093,10 @@ class AnthropicMessagesHandler(BaseTranslation):
) -> None:
"""Extract text, images, and tool calls from content blocks."""
for content_idx, content_block in enumerate(response_content):
block_dict: dict[str, object] = {}
if isinstance(content_block, dict):
block_type = content_block.get("type")
block_dict = cast(dict[str, object], content_block)
elif hasattr(content_block, "type"):
block_type = getattr(content_block, "type", None)
if hasattr(content_block, "model_dump"):
block_dict = content_block.model_dump()
else:
block_dict = {
"type": block_type,
"text": getattr(content_block, "text", None),
}
else:
fields = self._output_block_fields(content_block)
if fields is None:
continue
block_type, block_dict = fields
if block_type in ["text", "tool_use"]:
self._extract_output_text_and_images(
@ -1071,6 +1108,21 @@ class AnthropicMessagesHandler(BaseTranslation):
tool_calls_to_check=tool_calls_to_check,
)
@staticmethod
def _output_block_fields(content_block: object) -> "tuple[object, Mapping[str, object]] | None":
if isinstance(content_block, dict):
block_dict: Final = _as_str_mapping(content_block)
return block_dict.get("type"), block_dict
if not hasattr(content_block, "type"):
return None
block_type: Final = getattr(content_block, "type", None)
if isinstance(content_block, _ModelDumpBlock):
return block_type, content_block.model_dump()
return block_type, {
"type": block_type,
"text": getattr(content_block, "text", None),
}
@staticmethod
def _build_guardrail_inputs(
texts_to_check: list[str],
@ -1093,7 +1145,7 @@ class AnthropicMessagesHandler(BaseTranslation):
inputs["model"] = response_model
return inputs
def get_streaming_string_so_far(self, responses_so_far: list[Any]) -> str:
def get_streaming_string_so_far(self, responses_so_far: Sequence[object]) -> str:
"""
Parse streaming responses and extract accumulated text content.
@ -1164,7 +1216,7 @@ class AnthropicMessagesHandler(BaseTranslation):
# Only process content_block_delta events
if event_type == "content_block_delta" and data_line:
try:
data = json.loads(data_line)
data: _SSEEventData = json.loads(data_line)
delta = data.get("delta", {})
if delta.get("type") == "text_delta":
text += delta.get("text", "")
@ -1176,7 +1228,7 @@ class AnthropicMessagesHandler(BaseTranslation):
return text
def _check_streaming_has_ended(self, responses_so_far: list[Any]) -> bool:
def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool:
"""
Check if streaming response has ended by looking for non-null stop_reason.
@ -1227,7 +1279,7 @@ class AnthropicMessagesHandler(BaseTranslation):
# Check for message_delta event with stop_reason
if event_type == "message_delta" and data_line:
try:
data = json.loads(data_line)
data: _SSEEventData = json.loads(data_line)
delta = data.get("delta", {})
stop_reason = delta.get("stop_reason")
if stop_reason is not None:
@ -1271,7 +1323,7 @@ class AnthropicMessagesHandler(BaseTranslation):
def _extract_output_text_and_images(
self,
content_block: dict[str, object],
content_block: Mapping[str, object],
content_idx: int,
texts_to_check: list[str],
images_to_check: list[str],
@ -1294,7 +1346,7 @@ class AnthropicMessagesHandler(BaseTranslation):
task_mappings.append((content_idx, None))
# Extract tool calls
elif content_type == "tool_use":
elif content_type == "tool_use" and isinstance(content_block, dict):
tool_call: Final = AnthropicConfig.convert_tool_use_to_openai_format(
anthropic_tool_content=content_block,
index=content_idx,
@ -1319,7 +1371,7 @@ class AnthropicMessagesHandler(BaseTranslation):
content_idx = cast(int, mapping[0])
# Handle both dict and object responses
response_content: list[Any] = []
response_content: Sequence[object] = []
if isinstance(response, dict):
response_content = response.get("content", []) or []
elif hasattr(response, "content"):
@ -1335,14 +1387,15 @@ class AnthropicMessagesHandler(BaseTranslation):
if content_idx >= len(response_content):
continue
content_block = response_content[content_idx]
content_block = _content_block_at(response_content, content_idx)
# Verify it's a text block and update the text field
# Handle both dict and Pydantic object content blocks
if isinstance(content_block, dict):
if content_block.get("type") == "text":
cast(dict[str, object], content_block)["text"] = guardrail_response
block = _as_writable(content_block)
if block.get("type") == "text":
block["text"] = guardrail_response
elif hasattr(content_block, "type") and getattr(content_block, "type", None) == "text":
# Update Pydantic object's text attribute
if hasattr(content_block, "text"):
if isinstance(content_block, _TextAttrBlock):
content_block.text = guardrail_response

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

@ -13,10 +13,10 @@ Mirrors Anthropic's native ``compact_20260112`` for non-Anthropic providers:
"""
import re
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, Literal, NotRequired, Optional, TypedDict, Union, cast
from collections.abc import Awaitable, Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, TypeVar, Union, cast
from typing_extensions import ReadOnly
from typing_extensions import NotRequired, ReadOnly, TypedDict, Unpack
import litellm
from litellm._logging import verbose_logger
@ -29,6 +29,7 @@ from litellm.types.llms.anthropic import (
if TYPE_CHECKING:
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.hooks.parallel_request_limiter_v3 import RateLimitDescriptor, RateLimitResponse
from litellm.router import Router
from litellm.types.llms.anthropic import (
AllAnthropicPassThroughMessageValues,
@ -84,6 +85,77 @@ _PROPAGATED_METADATA_KEYS: Final = (
_SUMMARY_TAG_RE: Final = re.compile(r"<summary>(.*?)</summary>", re.IGNORECASE | re.DOTALL)
_MsgT: Final = TypeVar("_MsgT", bound=Mapping[str, object])
def _as_object(value: object) -> object:
return value
def _is_tool_result_block(block: object) -> bool:
return isinstance(block, dict) and block.get("type") in ("tool_result",)
class _SummaryCallKwargs(TypedDict):
model: ReadOnly[str]
max_tokens: ReadOnly[int]
timeout: ReadOnly[float]
litellm_metadata: ReadOnly[Mapping[str, object]]
user: ReadOnly[NotRequired[str]]
allowed_model_region: ReadOnly[NotRequired[str]]
class _SummaryOptionalKwargs(TypedDict, total=False):
user: ReadOnly[str]
allowed_model_region: ReadOnly[str]
class _SummaryAcompletion(Protocol):
def __call__(
self,
*,
messages: Sequence[Mapping[str, object]],
**kwargs: Unpack[_SummaryCallKwargs], # kwargs-ok: forwarded verbatim to acompletion, which owns them
) -> "Awaitable[ModelResponse | CustomStreamWrapper]": ...
class _CreateRateLimitDescriptors(Protocol):
def __call__(
self,
*,
user_api_key_dict: "UserAPIKeyAuth",
data: Mapping[str, str],
rpm_limit_type: object,
tpm_limit_type: object,
model_has_failures: bool,
) -> "Sequence[RateLimitDescriptor]": ...
class _AddModelRateLimitDescriptor(Protocol):
def __call__(
self,
*,
user_api_key_dict: "UserAPIKeyAuth",
requested_model: str,
descriptors: "Sequence[RateLimitDescriptor]",
) -> None: ...
class _CreateOrgRateLimitDescriptors(Protocol):
def __call__(
self, user_api_key_dict: "UserAPIKeyAuth", requested_model: str | None = None
) -> "Sequence[RateLimitDescriptor]": ...
class _ShouldRateLimit(Protocol):
def __call__(
self,
*,
descriptors: "Sequence[RateLimitDescriptor]",
parent_otel_span: object,
read_only: bool,
) -> "Awaitable[RateLimitResponse]": ...
def _read_summary_model_setting() -> str | None:
"""Look up the configured summarization model from proxy general_settings."""
@ -159,11 +231,11 @@ async def _check_summary_model_access(
return True
key_models: Final = list(getattr(user_api_key_auth, "models", None) or [])
team_id: Final = getattr(user_api_key_auth, "team_id", None)
team_id: Final[str | None] = getattr(user_api_key_auth, "team_id", None)
team_model_aliases: Final = getattr(user_api_key_auth, "team_model_aliases", None)
team_models: Final = list(getattr(user_api_key_auth, "team_models", None) or [])
user_id: Final = getattr(user_api_key_auth, "user_id", None)
project_id: Final = getattr(user_api_key_auth, "project_id", None)
user_id: Final[str | None] = getattr(user_api_key_auth, "user_id", None)
project_id: Final[str | None] = getattr(user_api_key_auth, "project_id", None)
checks: Final[tuple[tuple[Literal["key", "team"], list[str]], ...]] = (
("key", key_models),
@ -372,7 +444,7 @@ async def _check_summary_model_budget(
return False
end_user_model_max_budget: Final = getattr(user_api_key_auth, "end_user_model_max_budget", None)
end_user_id: Final = getattr(user_api_key_auth, "end_user_id", None)
end_user_id: Final[str | None] = getattr(user_api_key_auth, "end_user_id", None)
if isinstance(end_user_model_max_budget, dict) and end_user_model_max_budget and end_user_id is not None:
try:
await model_max_budget_limiter.is_end_user_within_model_budget(
@ -424,40 +496,57 @@ async def _check_summary_model_rate_limit(
except Exception:
return True
limiter: Final = getattr(proxy_logging_obj, "max_parallel_request_limiter", None)
limiter: Final[object] = getattr(proxy_logging_obj, "max_parallel_request_limiter", None)
should_rate_limit_check: Final[_ShouldRateLimit | None] = getattr(limiter, "should_rate_limit", None)
create_descriptors: Final[_CreateRateLimitDescriptors | None] = getattr(
limiter, "_create_rate_limit_descriptors", None
)
add_team_descriptor: Final[_AddModelRateLimitDescriptor | None] = getattr(
limiter, "_add_team_model_rate_limit_descriptor_from_metadata", None
)
add_project_descriptor: Final[_AddModelRateLimitDescriptor | None] = getattr(
limiter, "_add_project_model_rate_limit_descriptor_from_metadata", None
)
create_org_descriptors: Final[_CreateOrgRateLimitDescriptors | None] = getattr(
limiter, "create_organization_rate_limit_descriptor", None
)
if (
limiter is None
or not hasattr(limiter, "should_rate_limit")
or not hasattr(limiter, "_create_rate_limit_descriptors")
or should_rate_limit_check is None
or create_descriptors is None
or add_team_descriptor is None
or add_project_descriptor is None
or create_org_descriptors is None
):
return True
try:
metadata: Final = getattr(user_api_key_auth, "metadata", None) or {}
metadata: Final[Mapping[str, object]] = getattr(user_api_key_auth, "metadata", None) or {}
data: Final = {"model": summary_model}
descriptors: Final = limiter._create_rate_limit_descriptors(
base_descriptors: Final = create_descriptors(
user_api_key_dict=user_api_key_auth,
data=data,
rpm_limit_type=metadata.get("rpm_limit_type"),
tpm_limit_type=metadata.get("tpm_limit_type"),
model_has_failures=False,
)
limiter._add_team_model_rate_limit_descriptor_from_metadata(
add_team_descriptor(
user_api_key_dict=user_api_key_auth,
requested_model=summary_model,
descriptors=descriptors,
descriptors=base_descriptors,
)
limiter._add_project_model_rate_limit_descriptor_from_metadata(
add_project_descriptor(
user_api_key_dict=user_api_key_auth,
requested_model=summary_model,
descriptors=descriptors,
descriptors=base_descriptors,
)
descriptors.extend(limiter.create_organization_rate_limit_descriptor(user_api_key_auth, summary_model))
descriptors: Final = (*base_descriptors, *create_org_descriptors(user_api_key_auth, summary_model))
if not descriptors:
return True
response: Final = await limiter.should_rate_limit(
parent_otel_span: Final[object] = getattr(user_api_key_auth, "parent_otel_span", None)
response: Final[RateLimitResponse] = await should_rate_limit_check(
descriptors=descriptors,
parent_otel_span=getattr(user_api_key_auth, "parent_otel_span", None),
parent_otel_span=parent_otel_span,
read_only=True,
)
except Exception as e:
@ -471,7 +560,7 @@ async def _check_summary_model_rate_limit(
def _find_latest_compaction_index(
messages: list[dict[str, object]],
messages: Sequence[Mapping[str, object]],
) -> tuple[int | None, int | None]:
"""Return (message_index, block_index) of the most recent compaction block.
@ -490,8 +579,8 @@ def _find_latest_compaction_index(
def _slice_around_compaction_block(
messages: list[dict[str, Any]],
) -> tuple[list[dict[str, object]], dict[str, object] | None]:
messages: Sequence[_MsgT],
) -> tuple[Sequence[_MsgT | dict[str, object]], dict[str, object] | None]:
"""Apply Anthropic's "drop everything before the compaction block" rule.
Returns ``(sliced_messages_with_compaction_block, compaction_block_dict)``
@ -506,19 +595,21 @@ def _slice_around_compaction_block(
original_msg: Final = messages[msg_idx]
original_content: Final = original_msg["content"]
compaction_block: Final = cast(dict[str, object], original_content[blk_idx])
if not isinstance(original_content, list):
return messages, None
original_blocks: Final = cast("Sequence[dict[str, object]]", original_content)
compaction_block: Final = original_blocks[blk_idx]
# Per Anthropic's contract everything before the compaction block is
# dropped, including earlier blocks within the same assistant message.
sliced_content: Final = list(original_content[blk_idx:])
sliced_content: Final = list(original_blocks[blk_idx:])
sliced_messages: Final[list[dict[str, object]]] = [{**original_msg, "content": sliced_content}]
sliced_messages.extend(messages[msg_idx + 1 :])
sliced_messages: Final = [{**original_msg, "content": sliced_content}, *messages[msg_idx + 1 :]]
return sliced_messages, compaction_block
def _strip_compaction_blocks(
messages: list[dict[str, object]],
messages: Sequence[dict[str, object]],
) -> list[dict[str, object]]:
"""Drop any ``compaction`` content blocks from messages.
@ -625,7 +716,7 @@ def _propagate_metadata(
def _count_effective_tokens(
model: str,
effective_messages: list[dict[str, object]],
effective_messages: Sequence[dict[str, object]],
compaction_block: CompactionBlock | None,
tools: list[dict[str, object]] | None,
system: str | list[dict[str, object]] | None = None,
@ -704,17 +795,18 @@ def _system_to_text(
return ""
if isinstance(system, str):
return system
parts: Final[list[str]] = []
for block in system:
if isinstance(block, dict) and block.get("type") == "text":
text = block.get("text")
if isinstance(text, str) and text:
parts.append(text)
return "\n".join(parts)
return "\n".join(
text
for block in system
if isinstance(block, dict)
and block.get("type") == "text"
and isinstance(text := block.get("text"), str)
and text
)
def _select_last_user_question(
messages: list[dict[str, object]],
messages: Sequence[dict[str, object]],
) -> list[dict[str, object]]:
"""Pick the most recent ``user`` turn that is a real question.
@ -729,16 +821,18 @@ def _select_last_user_question(
turns, or contained no user turns at all). The downstream call always
needs a non-empty user message.
"""
blocks: Sequence[object]
for msg in reversed(messages):
if msg.get("role") != "user":
continue
content = msg.get("content")
if isinstance(content, list):
filtered = [blk for blk in content if not (isinstance(blk, dict) and blk.get("type") == "tool_result")]
blocks = [*map(_as_object, content)]
filtered = [blk for blk in blocks if not _is_tool_result_block(blk)]
if not filtered:
# Purely tool_result — skip and look for an earlier turn.
continue
if len(filtered) < len(content):
if len(filtered) < len(blocks):
return [{**msg, "content": filtered}]
return [msg]
return [
@ -761,7 +855,7 @@ def _extract_summary_text(raw: str | None) -> str | None:
def _system_to_openai_message(
system: str | list[dict[str, Any]] | None,
) -> dict[str, object] | None:
) -> Mapping[str, object] | None:
"""Translate Anthropic-shaped ``system`` to an OpenAI system message.
Accepts a bare string or a list of Anthropic content blocks; returns
@ -772,17 +866,19 @@ def _system_to_openai_message(
if isinstance(system, str):
return {"role": "system", "content": system} if system else None
if isinstance(system, list):
parts = [block.get("text", "") for block in system if isinstance(block, dict) and block.get("type") == "text"]
parts: Final[tuple[str, ...]] = tuple(
block.get("text", "") for block in system if isinstance(block, dict) and block.get("type") == "text"
)
joined: Final = "\n\n".join(part for part in parts if part)
return {"role": "system", "content": joined} if joined else None
return None
def _build_summary_messages(
effective_messages: list[dict[str, object]],
effective_messages: Sequence[dict[str, object]],
prompt: str,
system: str | list[dict[str, object]] | None = None,
) -> list[dict[str, object]]:
) -> Sequence[Mapping[str, object]]:
"""Build the OpenAI-shape message list for the summary call.
The caller's ``system`` prompt is prepended (the default summarization
@ -810,7 +906,7 @@ def _build_summary_messages(
)
openai_messages = stripped
summary_messages: Final[list[dict[str, object]]] = []
summary_messages: Final[list[Mapping[str, object]]] = []
system_message: Final = _system_to_openai_message(system)
if system_message is not None:
summary_messages.append(system_message)
@ -845,35 +941,17 @@ def _append_text_to_content(content: object, extra_text: str) -> object:
if isinstance(content, str):
return f"{content}\n\n{extra_text}"
if isinstance(content, list):
appended: Final[list[object]] = [*content, {"type": "text", "text": extra_text}]
appended: Final[Sequence[object]] = [*map(_as_object, content), {"type": "text", "text": extra_text}]
return appended
return [content, {"type": "text", "text": extra_text}]
class _SummaryCallUserKwarg(TypedDict, total=False):
user: ReadOnly[object]
class _SummaryCallRegionKwarg(TypedDict, total=False):
allowed_model_region: ReadOnly[str]
class _SummaryCallKwargs(TypedDict):
model: ReadOnly[str]
messages: ReadOnly[list[dict[str, object]]]
max_tokens: ReadOnly[int]
timeout: ReadOnly[float]
litellm_metadata: ReadOnly[Mapping[str, object]]
user: NotRequired[ReadOnly[object]]
allowed_model_region: NotRequired[ReadOnly[str]]
async def _call_summary_model(
*,
summary_model: str,
summary_messages: list[dict[str, object]],
summary_messages: Sequence[Mapping[str, object]],
metadata: Mapping[str, object],
llm_router: Any,
llm_router: object,
allowed_model_region: str | None = None,
max_tokens: int = COMPACT_SUMMARY_MAX_TOKENS,
) -> Union["ModelResponse", "CustomStreamWrapper"]:
@ -909,28 +987,37 @@ async def _call_summary_model(
# than from ``litellm_metadata``, so without it the summary tokens would not
# debit the caller's end-user counters.
end_user_id: Final = metadata.get("user_api_key_end_user_id")
user_kwargs: Final = (
_SummaryOptionalKwargs(user=end_user_id)
if isinstance(end_user_id, str) and end_user_id
else _SummaryOptionalKwargs()
)
region_kwargs: Final = (
_SummaryOptionalKwargs(allowed_model_region=allowed_model_region)
if allowed_model_region is not None
else _SummaryOptionalKwargs()
)
call_kwargs: Final[_SummaryCallKwargs] = {
"model": summary_model,
"messages": summary_messages,
"max_tokens": max_tokens,
"timeout": COMPACT_SUMMARY_TIMEOUT_SECONDS,
"litellm_metadata": metadata,
**(_SummaryCallUserKwarg(user=end_user_id) if end_user_id else _SummaryCallUserKwarg()),
**(
_SummaryCallRegionKwarg(allowed_model_region=allowed_model_region)
if allowed_model_region is not None
else _SummaryCallRegionKwarg()
),
**user_kwargs,
**region_kwargs,
}
if llm_router is not None and hasattr(llm_router, "acompletion"):
return await llm_router.acompletion(**call_kwargs)
return await litellm.acompletion(**call_kwargs)
router_acompletion: Final[_SummaryAcompletion | None] = getattr(llm_router, "acompletion", None)
if llm_router is not None and router_acompletion is not None:
return await router_acompletion(messages=summary_messages, **call_kwargs)
return await litellm.acompletion(messages=[*summary_messages], **call_kwargs)
def _extract_response_text(response: Any) -> str | None:
def _extract_response_text(response: object) -> str | None:
try:
choice: Final = response.choices[0]
message: Final = choice.message
choices: Final[Sequence[object] | None] = getattr(response, "choices", None)
if choices is None:
return None
choice: Final = choices[0]
message: Final = getattr(choice, "message", None)
content: Final = getattr(message, "content", None)
if isinstance(content, str):
return content
@ -946,7 +1033,7 @@ def _extract_response_text(response: Any) -> str | None:
def _extract_usage(response: object) -> tuple[int, int]:
usage: Final = getattr(response, "usage", None)
usage: Final[object] = getattr(response, "usage", None)
if usage is None:
return 0, 0
return (

67
litellm/llms/runwayml/videos/transformation.py
View file

@ -117,6 +117,10 @@ class RunwayMLVideoConfig(BaseVideoConfig):
def __init__(self):
super().__init__()
@staticmethod
def _parse_task_response(raw_response: httpx.Response) -> _RunwayTaskResponse:
return raw_response.json()
def get_supported_openai_params(self, model: str) -> list:
"""
Get the list of supported OpenAI parameters for video generation.
@ -141,7 +145,7 @@ class RunwayMLVideoConfig(BaseVideoConfig):
video_create_optional_params: VideoCreateOptionalRequestParams,
model: str,
drop_params: bool,
) -> dict:
) -> dict[str, object]:
"""
Map OpenAI parameters to RunwayML format.
@ -151,37 +155,44 @@ class RunwayMLVideoConfig(BaseVideoConfig):
- size -> ratio (convert "WIDTHxHEIGHT" to "WIDTH:HEIGHT")
- seconds -> duration (convert to integer)
"""
mapped_params: Final[dict[str, object]] = {}
supported_openai_params: Final = self.get_supported_openai_params(model)
return {
**self._prompt_image_param(video_create_optional_params),
**self._ratio_param(video_create_optional_params),
**self._duration_param(video_create_optional_params),
# Pass through other parameters that aren't OpenAI-specific
**{key: value for key, value in video_create_optional_params.items() if key not in supported_openai_params},
}
@staticmethod
def _prompt_image_param(video_create_optional_params: VideoCreateOptionalRequestParams) -> Mapping[str, object]:
# Handle input_reference parameter - map to promptImage
# RunwayML supports URLs and data URIs directly
if "input_reference" in video_create_optional_params:
input_reference: Final = video_create_optional_params["input_reference"]
# RunwayML supports URLs and data URIs directly
mapped_params["promptImage"] = input_reference
return {"promptImage": video_create_optional_params["input_reference"]}
return {}
@staticmethod
def _ratio_param(video_create_optional_params: VideoCreateOptionalRequestParams) -> Mapping[str, str]:
# Handle size parameter - convert "1280x720" to "1280:720"
if "size" in video_create_optional_params:
size: Final = video_create_optional_params["size"]
if isinstance(size, str) and "x" in size:
mapped_params["ratio"] = size.replace("x", ":")
return {"ratio": size.replace("x", ":")}
return {}
@staticmethod
def _duration_param(video_create_optional_params: VideoCreateOptionalRequestParams) -> Mapping[str, int]:
# Handle seconds parameter - convert to integer
if "seconds" in video_create_optional_params:
seconds: Final = video_create_optional_params["seconds"]
if seconds is not None:
try:
mapped_params["duration"] = int(float(seconds)) if isinstance(seconds, str) else int(seconds)
return {"duration": int(float(seconds)) if isinstance(seconds, str) else int(seconds)}
except (ValueError, TypeError):
# If conversion fails, use default duration
pass
# Pass through other parameters that aren't OpenAI-specific
supported_openai_params: Final = self.get_supported_openai_params(model)
for key, value in video_create_optional_params.items():
if key not in supported_openai_params:
mapped_params[key] = value
return mapped_params
return {}
def validate_environment(
self,
@ -236,7 +247,7 @@ class RunwayMLVideoConfig(BaseVideoConfig):
model: str,
prompt: str,
api_base: str,
video_create_optional_request_params: dict,
video_create_optional_request_params: dict[str, object],
litellm_params: GenericLiteLLMParams,
headers: dict,
) -> tuple[dict, RequestFiles, str]:
@ -406,20 +417,18 @@ class RunwayMLVideoConfig(BaseVideoConfig):
# Get task status to retrieve video URL
url: Final = f"{api_base}/tasks/{encoded_video_id}"
params: Final[dict[str, str]] = {}
return url, dict[str, str]()
return url, params
def _extract_video_url_from_response(self, response_data: dict[str, Any]) -> str:
def _extract_video_url_from_response(self, response_data: _RunwayTaskResponse) -> str:
"""
Helper method to extract video URL from RunwayML response.
Shared between sync and async transforms.
"""
# Extract video URL from the output field
video_url = None
if "output" in response_data and response_data["output"]:
output: Final = response_data["output"]
video_url = output[0] if isinstance(output, list) else output
raw_output: Final = response_data.get("output")
if raw_output:
video_url = raw_output if isinstance(raw_output, str) else raw_output[0]
if not video_url:
# Check if the video generation failed or is still processing
@ -453,7 +462,7 @@ class RunwayMLVideoConfig(BaseVideoConfig):
"output":["https://dnznrvs05pmza.cloudfront.net/.../video.mp4?_jwt=..."]
}
"""
response_data: Final = raw_response.json()
response_data: Final[_RunwayTaskResponse] = self._parse_task_response(raw_response)
video_url: Final = self._extract_video_url_from_response(response_data)
# Download the video from the CloudFront URL synchronously
@ -482,7 +491,7 @@ class RunwayMLVideoConfig(BaseVideoConfig):
"output":["https://dnznrvs05pmza.cloudfront.net/.../video.mp4?_jwt=..."]
}
"""
response_data: Final = raw_response.json()
response_data: Final[_RunwayTaskResponse] = self._parse_task_response(raw_response)
video_url: Final = self._extract_video_url_from_response(response_data)
# Download the video from the CloudFront URL asynchronously
@ -564,9 +573,7 @@ class RunwayMLVideoConfig(BaseVideoConfig):
# Construct the URL for task cancellation
url: Final = f"{api_base}/tasks/{encoded_video_id}/cancel"
data: Final[dict[str, str]] = {}
return url, data
return url, dict[str, str]()
def transform_video_delete_response(
self,
@ -604,9 +611,7 @@ class RunwayMLVideoConfig(BaseVideoConfig):
url: Final = f"{api_base}/tasks/{encoded_video_id}"
# Empty dict for GET request (no body)
data: Final[dict[str, str]] = {}
return url, data
return url, dict[str, str]()
def transform_video_status_retrieve_response(
self,

101
litellm/proxy/db/tool_registry_writer.py
View file

@ -6,9 +6,11 @@ Admins use the management endpoints to read and update input_policy / output_pol
"""
import uuid
from collections.abc import Mapping
from collections.abc import Mapping, Sequence
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any, Final
from typing import TYPE_CHECKING, Final, Protocol
from pydantic import TypeAdapter
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import ToolDiscoveryQueueItem
@ -27,6 +29,13 @@ if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
class _ModelDumpMethod(Protocol):
def __call__(self) -> Mapping: ...
_ROW_DICT: Final = TypeAdapter(dict)
def _tool_table_actions(prisma_client: "PrismaClient") -> "TableActions[prisma_db_models.LiteLLM_ToolTable]":
table: Final[TableActions[prisma_db_models.LiteLLM_ToolTable]] = ToolRepository(prisma_client).table
return table
@ -41,33 +50,35 @@ def _object_permission_table_actions(
return table
def _row_to_model(row: dict | Any) -> LiteLLM_ToolTableRow:
def _row_to_model(row: object) -> LiteLLM_ToolTableRow:
"""Convert a Prisma model instance or dict to LiteLLM_ToolTableRow."""
model_dump: Final = getattr(row, "model_dump", None)
model_dump: Final[_ModelDumpMethod | None] = getattr(row, "model_dump", None)
if callable(model_dump):
row = model_dump()
elif not isinstance(row, dict):
row = {
k: getattr(row, k, None)
for k in (
"tool_id",
"tool_name",
"origin",
"input_policy",
"output_policy",
"call_count",
"assignments",
"key_hash",
"team_id",
"key_alias",
"user_agent",
"last_used_at",
"created_at",
"updated_at",
"created_by",
"updated_by",
)
}
row = _ROW_DICT.validate_python(
{
k: getattr(row, k, None)
for k in (
"tool_id",
"tool_name",
"origin",
"input_policy",
"output_policy",
"call_count",
"assignments",
"key_hash",
"team_id",
"key_alias",
"user_agent",
"last_used_at",
"created_at",
"updated_at",
"created_by",
"updated_by",
)
}
)
return LiteLLM_ToolTableRow(
tool_id=row.get("tool_id", ""),
tool_name=row.get("tool_name", ""),
@ -190,7 +201,7 @@ async def update_tool_policy(
_updated_by: Final = updated_by or "system"
now: Final = datetime.now(timezone.utc)
create_data: Final[dict[str, object]] = {
create_data: Final[Mapping[str, str | datetime]] = {
"tool_id": str(uuid.uuid4()),
"tool_name": tool_name,
"input_policy": input_policy or "untrusted",
@ -200,14 +211,16 @@ async def update_tool_policy(
"created_at": now,
"updated_at": now,
}
update_data: Final[dict[str, object]] = {
"updated_by": _updated_by,
"updated_at": now,
update_data: Final[Mapping[str, str | datetime]] = {
key: value
for key, value in (
("updated_by", _updated_by),
("updated_at", now),
("input_policy", input_policy),
("output_policy", output_policy),
)
if value is not None
}
if input_policy is not None:
update_data["input_policy"] = input_policy
if output_policy is not None:
update_data["output_policy"] = output_policy
await _tool_table_actions(prisma_client).upsert(
where={"tool_name": tool_name},
@ -338,7 +351,7 @@ class ToolPolicyRegistry:
self._blocked_tools_by_op_id = {}
for row in perms:
op_id = getattr(row, "object_permission_id", None)
blocked = getattr(row, "blocked_tools", None) or []
blocked: Sequence[str] = getattr(row, "blocked_tools", None) or []
if op_id:
self._blocked_tools_by_op_id[op_id] = list(blocked)
@ -370,10 +383,12 @@ class ToolPolicyRegistry:
"""
if not tool_names:
return {}
blocked: Final[set[str]] = set()
for op_id in (object_permission_id, team_object_permission_id):
if op_id and op_id.strip():
blocked.update(self._blocked_tools_by_op_id.get(op_id.strip(), []))
blocked: Final[frozenset[str]] = frozenset(
tool
for op_id in (object_permission_id, team_object_permission_id)
if op_id and op_id.strip()
for tool in self._blocked_tools_by_op_id.get(op_id.strip(), [])
)
result: Final[dict[str, str]] = {}
for name in tool_names:
if name in blocked:
@ -408,13 +423,12 @@ async def add_tool_to_object_permission_blocked(
)
if row is None:
return False
current: Final = list(getattr(row, "blocked_tools", []) or [])
current: Final[Sequence[str]] = getattr(row, "blocked_tools", []) or []
if tool_name in current:
return True
current.append(tool_name)
await _object_permission_table_actions(prisma_client).update(
where={"object_permission_id": object_permission_id},
data={"blocked_tools": current},
data={"blocked_tools": [*current, tool_name]},
)
return True
except Exception as e:
@ -436,13 +450,12 @@ async def remove_tool_from_object_permission_blocked(
)
if row is None:
return False
current = list(getattr(row, "blocked_tools", []) or [])
current: Final[Sequence[str]] = getattr(row, "blocked_tools", []) or []
if tool_name not in current:
return False
current = [t for t in current if t != tool_name]
await _object_permission_table_actions(prisma_client).update(
where={"object_permission_id": object_permission_id},
data={"blocked_tools": current},
data={"blocked_tools": [t for t in current if t != tool_name]},
)
return True
except Exception as e:

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

@ -17,7 +17,7 @@ import json
import traceback
from collections.abc import Awaitable, Mapping, Sequence
from datetime import datetime, timezone
from typing import Any, Final, Literal, cast
from typing import Any, Final, Literal, Protocol, cast, overload
import fastapi
from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
@ -735,10 +735,44 @@ def _enforce_user_info_access(user_id: str | None, user_api_key_dict: UserAPIKey
)
async def _get_user_info_teams(
prisma_client: Any,
class _UserInfoDataClient(Protocol):
@overload
async def get_data(self, *, user_id: str) -> "prisma_models.LiteLLM_UserTable | None": ...
@overload
async def get_data(
self,
*,
user_id: str | None,
table_name: Literal["key"],
query_type: Literal["find_all"],
) -> "Sequence[LiteLLM_VerificationToken] | None": ...
@overload
async def get_data(
self,
*,
team_id_list: list[str],
table_name: Literal["team"],
query_type: Literal["find_all"],
) -> "Sequence[TeamListResponseObject] | None": ...
async def _get_user_info_keys(
prisma_client: "_UserInfoDataClient",
user_id: str | None,
user_info: Any | None,
) -> "Sequence[LiteLLM_VerificationToken] | None":
return await prisma_client.get_data(
user_id=user_id,
table_name="key",
query_type="find_all",
)
async def _get_user_info_teams(
prisma_client: "_UserInfoDataClient",
user_id: str | None,
user_info: "prisma_models.LiteLLM_UserTable",
user_api_key_dict: UserAPIKeyAuth,
) -> tuple[list[TeamListResponseObject], list[TeamListResponseObject] | None]:
"""Fetch and merge teams from membership + user.teams field."""
@ -759,7 +793,7 @@ async def _get_user_info_teams(
team_list = teams_1
team_id_list = [team.team_id for team in teams_1]
teams_2: list[TeamListResponseObject] | None = None
teams_2: Sequence[TeamListResponseObject] | None = None
target_team_ids: Final = getattr(user_info, "teams", None)
if target_team_ids and isinstance(target_team_ids, list):
@ -769,8 +803,8 @@ async def _get_user_info_teams(
query_type="find_all",
)
elif user_api_key_dict.user_id is not None and user_id is None:
caller_user_info: Final[object] = await prisma_client.get_data(user_id=user_api_key_dict.user_id)
caller_team_ids: Final = getattr(caller_user_info, "teams", None)
caller_user_info: Final = await prisma_client.get_data(user_id=user_api_key_dict.user_id)
caller_team_ids: Final = caller_user_info.teams if caller_user_info is not None else None
if caller_team_ids:
teams_2 = await prisma_client.get_data(
team_id_list=caller_team_ids,
@ -807,7 +841,7 @@ def _redact_scim_enterprise_metadata(
def _build_user_info_response(
user_id: str | None,
user_info: Any | None,
keys: list[LiteLLM_VerificationToken] | None,
keys: Sequence[LiteLLM_VerificationToken] | None,
team_list: list[TeamListResponseObject],
teams_1: list[TeamListResponseObject] | None,
model_max_budget_usage: dict[str, dict[str, object]] | None = None,
@ -894,11 +928,7 @@ async def user_info(
)
## GET ALL KEYS ##
keys: Final = await prisma_client.get_data(
user_id=user_id,
table_name="key",
query_type="find_all",
)
keys: Final = await _get_user_info_keys(prisma_client, user_id)
response_data: Final = _build_user_info_response(
user_id=user_id,
@ -1077,6 +1107,12 @@ async def user_info_v2(
raise handle_exception_on_proxy(e)
async def _fetch_admin_teams_and_keys_rows(
prisma_client: "PrismaClient", sql_query: str
) -> Sequence[Mapping[str, Sequence[Mapping[str, object]] | None]]:
return await prisma_client.db.query_raw(sql_query)
async def _get_user_info_for_proxy_admin(user_api_key_dict: UserAPIKeyAuth):
"""
Admin UI Endpoint - Returns All Teams and Keys when Proxy Admin is querying
@ -1100,22 +1136,25 @@ async def _get_user_info_for_proxy_admin(user_api_key_dict: UserAPIKeyAuth):
"Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys"
)
results: Final = await prisma_client.db.query_raw(sql_query)
results: Final = await _fetch_admin_teams_and_keys_rows(prisma_client, sql_query)
verbose_proxy_logger.debug("results_keys: %s", results)
_keys_in_db: Final[Sequence[dict[str, object]]] = results[0]["keys"] or []
_keys_in_db: Final[Sequence[Mapping[str, object]]] = results[0]["keys"] or []
# cast all keys to LiteLLM_VerificationToken
keys_in_db: Final = []
for key in _keys_in_db:
if key.get("models") is None:
key["models"] = []
keys_in_db.append(LiteLLM_VerificationToken.model_validate(key))
key_payload = dict[str, object](key)
if key_payload.get("models") is None:
key_payload["models"] = []
keys_in_db.append(LiteLLM_VerificationToken.model_validate(key_payload))
# cast all teams to LiteLLM_TeamTable
_teams_in_db: list[LiteLLM_TeamTable] = results[0]["teams"] or []
_teams_in_db = [LiteLLM_TeamTable.model_validate(team) for team in _teams_in_db]
_teams_in_db.sort(key=lambda x: getattr(x, "team_alias", "") or "")
_teams_rows: Final[Sequence[Mapping[str, object]]] = results[0]["teams"] or []
_teams_in_db: Final = sorted(
(LiteLLM_TeamTable.model_validate(team) for team in _teams_rows),
key=lambda x: getattr(x, "team_alias", "") or "",
)
returned_keys: Final = _process_keys_for_user_info(keys=keys_in_db, all_teams=_teams_in_db)
# Get admin's own user_id and user_info
@ -1140,7 +1179,7 @@ async def _get_user_info_for_proxy_admin(user_api_key_dict: UserAPIKeyAuth):
def _process_keys_for_user_info(
keys: list[LiteLLM_VerificationToken] | None,
keys: Sequence[LiteLLM_VerificationToken] | None,
all_teams: list[LiteLLM_TeamTable] | list[TeamListResponseObject] | None,
):
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
@ -1231,7 +1270,7 @@ def _update_internal_user_params(data_json: dict, data: UpdateUserRequest | Upda
async def _schedule_user_update_audit_log(
response: dict[str, Any],
response: Mapping[str, object],
existing_user_row: BaseModel | None,
litellm_changed_by: str | None,
user_api_key_dict: UserAPIKeyAuth,

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

@ -18,7 +18,8 @@ import os
import re
import secrets
import traceback
from collections.abc import Awaitable, Callable, Mapping, Sequence
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
from contextlib import AbstractAsyncContextManager
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, TypeVar, cast
@ -230,6 +231,54 @@ def _config_table(prisma_client: PrismaClient) -> _ConfigTableActions:
)
class _CustomKeyHooksModule(Protocol):
user_custom_key_generate: Callable[..., Awaitable[Mapping[str, object]]] | None
user_custom_key_update: Callable[..., Awaitable[Mapping[str, object]]] | None
def _custom_key_generate_hook(
hooks: _CustomKeyHooksModule,
) -> Callable[..., Awaitable[Mapping[str, object]]] | None:
return hooks.user_custom_key_generate
def _custom_key_update_hook(
hooks: _CustomKeyHooksModule,
) -> Callable[..., Awaitable[Mapping[str, object]]] | None:
return hooks.user_custom_key_update
class _LegacyDumpable(Protocol):
def dict(self) -> Mapping[str, object]: ...
def _legacy_model_dict(row: _LegacyDumpable) -> Mapping[str, object]:
return row.dict()
def _as_object_dict(values: Mapping[str, object]) -> Mapping[str, object]:
return values
def _model_items(model: BaseModel) -> Iterator[tuple[str, object]]:
return iter(model)
class _EnvVarsParam(Protocol):
@property
def param_value(self) -> Mapping[str, str] | None: ...
def _env_vars_param_value(param: _EnvVarsParam) -> Mapping[str, str] | None:
return param.param_value
def _tx_tables_context(
open_tx: Callable[[], AbstractAsyncContextManager[_TxTables]],
) -> AbstractAsyncContextManager[_TxTables]:
return open_tx()
async def _check_custom_key_allowed(custom_key_value: str | None) -> None:
"""Raise 403 if custom API keys are disabled and a custom key was provided."""
if custom_key_value is None:
@ -910,7 +959,7 @@ async def _common_key_generation_helper(
# check if user set default key/generate params on config.yaml
if litellm.default_key_generate_params is not None:
for elem in data:
for elem in _model_items(data):
key, value = elem
if (
value is None
@ -1692,11 +1741,11 @@ async def generate_key_fn(
- user_id: (str) Unique user id - used for tracking spend across multiple keys for same user id.
"""
try:
from litellm.proxy import proxy_server
from litellm.proxy._types import CommonProxyErrors
from litellm.proxy.proxy_server import (
prisma_client,
user_api_key_cache,
user_custom_key_generate,
)
if prisma_client is None:
@ -1723,7 +1772,7 @@ async def generate_key_fn(
)
custom_key_generate_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = (
user_custom_key_generate
_custom_key_generate_hook(proxy_server)
)
if custom_key_generate_hook is not None:
if inspect.iscoroutinefunction(custom_key_generate_hook):
@ -1892,11 +1941,11 @@ async def generate_service_account_key_fn(
- user_id: (str) Unique user id - used for tracking spend across multiple keys for same user id.
"""
from litellm.proxy import proxy_server
from litellm.proxy._types import CommonProxyErrors
from litellm.proxy.proxy_server import (
prisma_client,
user_api_key_cache,
user_custom_key_generate,
)
if prisma_client is None:
@ -1924,7 +1973,9 @@ async def generate_service_account_key_fn(
verbose_proxy_logger.debug("entered /key/generate")
custom_key_generate_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = user_custom_key_generate
custom_key_generate_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = _custom_key_generate_hook(
proxy_server
)
if custom_key_generate_hook is not None:
if inspect.iscoroutinefunction(custom_key_generate_hook):
result: Final = await custom_key_generate_hook(data)
@ -1998,7 +2049,7 @@ def prepare_metadata_fields(data: BaseModel, non_default_values: dict, existing_
)
casted_metadata[reserved_field] = existing_value
data_json: Final[Mapping[str, object]] = data.model_dump(exclude_unset=True, exclude_none=True)
data_json: Final = _as_object_dict(data.model_dump(exclude_unset=True, exclude_none=True))
try:
for k, v in data_json.items():
@ -2805,13 +2856,13 @@ async def update_key_fn(
}'
```
"""
from litellm.proxy import proxy_server
from litellm.proxy.proxy_server import (
llm_router,
premium_user,
prisma_client,
proxy_logging_obj,
user_api_key_cache,
user_custom_key_update,
)
try:
@ -2842,7 +2893,9 @@ async def update_key_fn(
)
# Custom key update hook
custom_key_update_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = user_custom_key_update
custom_key_update_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = _custom_key_update_hook(
proxy_server
)
if custom_key_update_hook is not None:
if inspect.iscoroutinefunction(custom_key_update_hook):
result: Final = await custom_key_update_hook(data)
@ -3004,14 +3057,16 @@ async def bulk_update_keys(
}'
```
"""
from litellm.proxy import proxy_server
from litellm.proxy.proxy_server import (
llm_router,
prisma_client,
proxy_logging_obj,
user_api_key_cache,
user_custom_key_update,
)
custom_key_update_hook: Final = _custom_key_update_hook(proxy_server)
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value:
raise HTTPException(
status_code=403,
@ -3057,7 +3112,7 @@ async def bulk_update_keys(
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
llm_router=llm_router,
user_custom_key_update=user_custom_key_update,
user_custom_key_update=custom_key_update_hook,
)
successful_updates.append(
@ -3135,7 +3190,7 @@ def _build_failed_team_key_update(
if hasattr(existing_key_row, "model_dump"):
key_info = existing_key_row.model_dump()
elif hasattr(existing_key_row, "dict"):
key_info = existing_key_row.dict()
key_info = dict[str, object](_legacy_model_dict(existing_key_row))
if key_info:
key_info.pop("token", None)
@ -3166,14 +3221,16 @@ async def bulk_update_team_keys(
Callable by proxy admins, or by team admins with `KEY_UPDATE` permission.
"""
from litellm.proxy import proxy_server
from litellm.proxy.proxy_server import (
llm_router,
prisma_client,
proxy_logging_obj,
user_api_key_cache,
user_custom_key_update,
)
custom_key_update_hook: Final = _custom_key_update_hook(proxy_server)
if prisma_client is None:
raise HTTPException(
status_code=500,
@ -3302,7 +3359,7 @@ async def bulk_update_team_keys(
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
llm_router=llm_router,
user_custom_key_update=user_custom_key_update,
user_custom_key_update=custom_key_update_hook,
existing_key_row=existing_by_token[db_token],
)
@ -4301,7 +4358,7 @@ def _transform_verification_tokens_to_deleted_records(
"litellm_changed_by": litellm_changed_by,
}
)
record = deleted_record.model_dump()
record = dict[str, object](_as_object_dict(deleted_record.model_dump()))
# Map org_id to organization_id (model uses org_id, but schema expects organization_id)
org_id_value: object = record.pop("org_id", None)
@ -4437,13 +4494,12 @@ async def _rotate_master_key(
should_create_model_in_db=False,
)
if new_model:
_dumped = new_model.model_dump(exclude_none=True)
_dumped = dict[str, object](_as_object_dict(new_model.model_dump(exclude_none=True)))
_dumped["litellm_params"] = prisma.Json(_dumped["litellm_params"])
_dumped["model_info"] = prisma.Json(_dumped["model_info"])
new_models.append(_dumped)
verbose_proxy_logger.debug("Resetting proxy model table")
async with prisma_client.db.tx() as tx_ctx:
tx: Final[_TxTables] = tx_ctx
async with _tx_tables_context(prisma_client.db.tx) as tx:
await tx.litellm_proxymodeltable.delete_many()
verbose_proxy_logger.debug("Creating %s models", len(new_models))
await tx.litellm_proxymodeltable.create_many(
@ -4458,14 +4514,14 @@ async def _rotate_master_key(
if config:
"""If environment_variables is found, decrypt it and encrypt it with the new master key"""
environment_variables_dict = {}
environment_variables_dict: Mapping[str, str] | None = {}
for c in config:
if c.param_name == "environment_variables":
environment_variables_dict = c.param_value
environment_variables_dict = _env_vars_param_value(c)
if environment_variables_dict:
decrypted_env_vars: Final = proxy_config._decrypt_and_set_db_env_variables(
environment_variables=environment_variables_dict
environment_variables=dict[str, str](environment_variables_dict)
)
encrypted_env_vars: Final = proxy_config._encrypt_env_variables(
environment_variables=decrypted_env_vars,
@ -4531,7 +4587,7 @@ async def _rotate_master_key(
updated_patch=decrypted_cred,
new_encryption_key=new_master_key,
)
_cred_data = encrypted_cred.model_dump(exclude_none=True)
_cred_data = dict[str, object](_as_object_dict(encrypted_cred.model_dump(exclude_none=True)))
if "credential_values" in _cred_data:
_cred_data["credential_values"] = prisma.Json(_cred_data["credential_values"])
if "credential_info" in _cred_data:

4
litellm/proxy/prompts/prompt_endpoints.py
View file

@ -6,7 +6,7 @@ import tempfile
from collections.abc import Awaitable, Mapping, Sequence
from datetime import datetime
from pathlib import Path
from typing import TYPE_CHECKING, Any, Final, Protocol, cast
from typing import TYPE_CHECKING, Final, Protocol, cast
from fastapi import (
APIRouter,
@ -1317,7 +1317,7 @@ async def test_prompt(
async def convert_prompt_file_to_json(
file: UploadFile = File(...),
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
) -> dict[str, Any]:
) -> Mapping[str, object]:
"""
Convert a .prompt file to JSON format.

110
litellm/proxy/response_polling/background_streaming.py
View file

@ -10,12 +10,12 @@ https://platform.openai.com/docs/api-reference/responses-streaming
import asyncio
import json
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Final, TypedDict
from collections.abc import Callable, Mapping, Sequence
from typing import TYPE_CHECKING, Final, TypeAlias
from fastapi import Request, Response
from fastapi.responses import StreamingResponse
from typing_extensions import ReadOnly
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
@ -29,53 +29,59 @@ if TYPE_CHECKING:
from litellm.router import Router
class _StreamContentPart(TypedDict, total=False):
text: ReadOnly[str]
_JsonDict: TypeAlias = dict[str, object]
_JsonList: TypeAlias = list[object]
class _StreamOutputItem(TypedDict, total=False):
class _OutputItem(TypedDict, total=False):
id: ReadOnly[str]
content: ReadOnly[Sequence[_StreamContentPart | None]]
content: ReadOnly[Sequence[object]]
class _StreamTerminalResponse(TypedDict, total=False):
"""Fields of the ``response`` payload carried by a terminal streaming event."""
class _TerminalResponse(TypedDict, total=False):
status: ReadOnly[ResponsesAPIStatus]
error: ReadOnly[_JsonDict]
usage: ReadOnly[_JsonDict]
reasoning: ReadOnly[_JsonDict]
tool_choice: ReadOnly[object]
tools: ReadOnly[_JsonList]
model: ReadOnly[str]
instructions: ReadOnly[str]
temperature: ReadOnly[float]
top_p: ReadOnly[float]
max_output_tokens: ReadOnly[int]
previous_response_id: ReadOnly[str]
text: ReadOnly[_JsonDict]
truncation: ReadOnly[str]
parallel_tool_calls: ReadOnly[bool]
user: ReadOnly[str]
store: ReadOnly[bool]
output: ReadOnly[Sequence[_StreamOutputItem]]
incomplete_details: ReadOnly[_JsonDict]
output: ReadOnly[Sequence[_OutputItem]]
class _StreamEvent(TypedDict, total=False):
"""One decoded ``data:`` frame of an OpenAI Responses streaming body."""
type: ReadOnly[str]
item: ReadOnly[_StreamOutputItem]
item: ReadOnly[_OutputItem]
item_id: ReadOnly[str]
part: ReadOnly[_StreamContentPart]
content_index: ReadOnly[int]
delta: ReadOnly[str]
response: ReadOnly[_StreamTerminalResponse]
part: ReadOnly[object]
response: ReadOnly[_TerminalResponse]
class _StreamEventParser:
parse: Callable[[str], _StreamEvent] = staticmethod(json.loads)
async def background_streaming_task(
polling_id: str,
data,
data: dict,
polling_handler: ResponsePollingHandler,
request: Request,
fastapi_response: Response,
user_api_key_dict: UserAPIKeyAuth,
general_settings,
general_settings: dict,
llm_router: "Router | None",
proxy_config: "ProxyConfig",
proxy_logging_obj: "ProxyLogging",
@ -138,9 +144,10 @@ async def background_streaming_task(
# Process streaming response following OpenAI events format
# https://platform.openai.com/docs/api-reference/responses-streaming
output_items: Final[dict[str, _StreamOutputItem]] = {} # Track output items by ID
# Track accumulated text deltas by (item_id, content_index)
accumulated_text: Final[dict[tuple[str, int], str]] = {}
output_items: Final = dict[str, _OutputItem]() # Track output items by ID
accumulated_text: Final = dict[
tuple[str, int], str
]() # Track accumulated text deltas by (item_id, content_index)
# ResponsesAPIResponse fields to extract from response.completed
usage_data = None
@ -210,7 +217,7 @@ async def background_streaming_task(
break
try:
event: _StreamEvent = json.loads(chunk_data)
event: _StreamEvent = _StreamEventParser.parse(chunk_data)
event_type = event.get("type", "")
# Process different event types based on OpenAI streaming spec
@ -229,19 +236,18 @@ async def background_streaming_task(
if item_id and item_id in output_items:
# Update the output item with new content
current_item = output_items[item_id]
appended_item: _StreamOutputItem = {
**current_item,
"content": (*current_item.get("content", ()), content_part),
added_item = output_items[item_id]
output_items[item_id] = {
**added_item,
"content": (*added_item.get("content", ()), content_part),
}
output_items[item_id] = appended_item
state_dirty = True
elif event_type == "response.output_text.delta":
# Text delta - accumulate text content
# https://platform.openai.com/docs/api-reference/responses-streaming/response-text-delta
item_id = event.get("item_id")
content_index: int = event.get("content_index", 0)
content_index = event.get("content_index", 0)
delta = event.get("delta", "")
if item_id and item_id in output_items:
@ -252,24 +258,14 @@ async def background_streaming_task(
accumulated_text[key] += delta
# Update the content in output_items
current_item = output_items[item_id]
content_list: Sequence[_StreamContentPart | None] = current_item.get("content", ())
if content_index < len(content_list):
# Update existing content part with accumulated text
content_entry = content_list[content_index]
if isinstance(content_entry, dict):
delta_part: _StreamContentPart = {
**content_entry,
"text": accumulated_text[key],
}
delta_item: _StreamOutputItem = {
**current_item,
"content": tuple(
delta_part if index == content_index else entry
for index, entry in enumerate(content_list)
),
}
output_items[item_id] = delta_item
delta_item = output_items[item_id]
if "content" in delta_item:
content_list = delta_item["content"]
if content_index < len(content_list):
# Update existing content part with accumulated text
content_entry = content_list[content_index]
if isinstance(content_entry, dict):
content_entry["text"] = accumulated_text[key]
state_dirty = True
elif event_type == "response.content_part.done":
@ -280,17 +276,17 @@ async def background_streaming_task(
if item_id and item_id in output_items:
# Update with final content from event
current_item = output_items[item_id]
content_list = current_item.get("content", ())
if content_index < len(content_list):
finalized_item: _StreamOutputItem = {
**current_item,
"content": tuple(
content_part if index == content_index else entry
for index, entry in enumerate(content_list)
),
}
output_items[item_id] = finalized_item
done_item = output_items[item_id]
if "content" in done_item:
content_list = done_item["content"]
if content_index < len(content_list):
output_items[item_id] = {
**done_item,
"content": tuple(
content_part if part_index == content_index else existing_part
for part_index, existing_part in enumerate(content_list)
),
}
state_dirty = True
elif event_type == "response.output_item.done":

205
litellm/responses/streaming_iterator.py
View file

@ -5,11 +5,11 @@ import json
import time
import traceback
import uuid
from collections.abc import Awaitable, Callable, Mapping, Sequence
from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
from datetime import datetime
from functools import lru_cache
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, runtime_checkable
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, overload, runtime_checkable
import httpx
from openai._streaming import SSEDecoder
@ -42,27 +42,14 @@ from litellm.types.utils import CallTypes
from litellm.utils import async_post_call_success_deployment_hook
if TYPE_CHECKING:
from litellm.caching.caching_handler import LLMCachingHandler
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.responses.streaming_websocket import (
PresidioGuardrailCallback,
ResponsesBackendWebSocket,
ResponsesClientWebSocket,
)
class _StreamCachingHandler(Protocol):
"""The ``_llm_caching_handler`` attached to a logging object, as this module uses it."""
original_function: Callable[..., object]
def _should_store_result_in_cache(
self, original_function: Callable[..., object], kwargs: Mapping[str, object]
) -> bool: ...
class PiiUnmaskingGuardrailCallback(PresidioGuardrailCallback, Protocol):
"""Guardrail callback that can also reverse its own masking, selected by
``llm_http_handler`` on exactly this attribute."""
def _unmask_pii_text(self, text: str, pii_tokens: Mapping[str, str]) -> str: ...
from litellm.types.router import LiteLLM_Params
class ProjectQuotaCallback(Protocol):
@ -94,6 +81,60 @@ def _is_str_mapping(value: object) -> TypeIs[dict[str, str]]: # guard-ok: verif
return _is_json_object(value) and all(isinstance(item, str) for item in value.values())
class _MutableJsonObject(Protocol):
@overload
def get(self, key: str, /) -> object | None: ...
@overload
def get(self, key: str, default: object, /) -> object: ...
def __getitem__(self, key: str, /) -> object: ...
def __setitem__(self, key: str, value: object, /) -> None: ...
def __contains__(self, key: object, /) -> bool: ...
def items(self) -> Iterable[tuple[str, object]]: ...
class _GetsLitellmParams(Protocol):
def __call__(self, key: str, default: Mapping[str, object], /) -> LiteLLM_Params: ...
class _PopsOptionalStr(Protocol):
def __call__(self, key: str, default: None, /) -> str | None: ...
class _UnmasksPiiText(Protocol):
def __call__(self, text: str, pii_tokens: Mapping[str, str]) -> str: ...
class _ShouldStoreResultInCache(Protocol):
def __call__(self, *, original_function: Callable[..., object] | None, kwargs: Mapping[str, object]) -> bool: ...
class _PostStreamingDeploymentHook(Protocol):
def __call__(
self,
*,
request_data: Mapping[str, object],
response_chunk: ResponsesAPIStreamingResponse,
call_type: CallTypes | None,
) -> Awaitable[ResponsesAPIStreamingResponse | None]: ...
@runtime_checkable
class _HasPostStreamingDeploymentHook(Protocol):
async_post_call_streaming_deployment_hook: _PostStreamingDeploymentHook
def _typed_gets_litellm_params(fn: _GetsLitellmParams) -> _GetsLitellmParams:
return fn
def _typed_pops_optional_str(fn: _PopsOptionalStr) -> _PopsOptionalStr:
return fn
_SHOULD_STORE_RESULT_IN_CACHE_ATTR: Final = "_should_store_result_in_cache"
_UNMASK_PII_TEXT_ATTR: Final = "_unmask_pii_text"
def _load_json_object(payload: str | bytes) -> dict[str, object]:
"""Parse a JSON payload that the caller consumes as an object."""
return json.loads(payload)
@ -220,7 +261,7 @@ class BaseResponsesAPIStreamingIterator:
# This matches the stream wrapper in litellm/litellm_core_utils/streaming_handler.py
_api_base: Final = get_api_base(
model=model or "",
optional_params=self.logging_obj.model_call_details.get("litellm_params", {}),
optional_params=_typed_gets_litellm_params(self.logging_obj.model_call_details.get)("litellm_params", {}),
)
self._hidden_params: dict[str, object] = {
"model_id": _model_id_from_metadata(litellm_metadata),
@ -549,7 +590,7 @@ class BaseResponsesAPIStreamingIterator:
if response_obj is None:
return
caching_handler: Final[_StreamCachingHandler | None] = getattr(self.logging_obj, "_llm_caching_handler", None)
caching_handler: Final[LLMCachingHandler | None] = getattr(self.logging_obj, "_llm_caching_handler", None)
if caching_handler is None:
return
@ -567,8 +608,11 @@ class BaseResponsesAPIStreamingIterator:
if preset_cache_key is not None:
request_kwargs["cache_key"] = preset_cache_key
if not caching_handler._should_store_result_in_cache( # pyright: ignore[reportPrivateUsage] # no public API
original_function=caching_handler.original_function,
should_store_result_in_cache: Final[_ShouldStoreResultInCache] = getattr(
caching_handler, _SHOULD_STORE_RESULT_IN_CACHE_ATTR
)
if not should_store_result_in_cache(
original_function=getattr(caching_handler, "original_function", None),
kwargs=request_kwargs,
):
return
@ -624,12 +668,15 @@ class BaseResponsesAPIStreamingIterator:
typed_call_type = None
request_data: Final = self.request_data or getattr(self.logging_obj, "model_call_details", {})
callbacks: Final = getattr(litellm, "callbacks", None) or []
callbacks: Final[Sequence[object]] = getattr(litellm, "callbacks", None) or []
hooks_ran = False
for callback in callbacks:
if hasattr(callback, "async_post_call_streaming_deployment_hook"):
if isinstance(callback, _HasPostStreamingDeploymentHook):
hooks_ran = True
result = await callback.async_post_call_streaming_deployment_hook(
post_streaming_hook: _PostStreamingDeploymentHook = (
callback.async_post_call_streaming_deployment_hook
)
result = await post_streaming_hook(
request_data=request_data,
response_chunk=chunk,
call_type=typed_call_type,
@ -1083,7 +1130,7 @@ class _HasModelDumpJson(Protocol):
def model_dump_json(self, *, exclude_none: bool = ...) -> str: ...
def _dump_response_object(obj: object) -> dict[str, Any]:
def _dump_response_object(obj: object) -> Mapping[str, object]:
if isinstance(obj, _HasModelDump):
return obj.model_dump()
if _is_json_object(obj):
@ -1113,21 +1160,20 @@ def _build_content_part_done_event(
item_id: str,
output_index: int,
content_index: int,
part_payload: dict[str, Any],
part_payload: Mapping[str, object],
) -> ResponsesAPIStreamingResponse | None:
openai_types: Final = _get_openai_response_types()
part_type: Final = part_payload.get("type")
part: PART_UNION_TYPES
if part_type == "output_text":
annotations: Final = [
openai_types.BaseLiteLLMOpenAIResponseObject(**annotation)
for annotation in part_payload.get("annotations", []) or []
]
part = openai_types.ContentPartDonePartOutputText(
type="output_text",
text=str(part_payload.get("text") or ""),
annotations=annotations,
logprobs=part_payload.get("logprobs"),
raw_annotations: Final[object] = part_payload.get("annotations", []) or []
part = openai_types.ContentPartDonePartOutputText.model_validate(
{
"type": "output_text",
"text": str(part_payload.get("text") or ""),
"annotations": raw_annotations,
"logprobs": part_payload.get("logprobs"),
}
)
elif part_type == "refusal":
part = openai_types.ContentPartDonePartRefusal(
@ -1157,7 +1203,7 @@ def _add_text_like_part_events(
item_id: str,
output_index: int,
content_index: int,
part_payload: dict[str, Any],
part_payload: Mapping[str, object],
chunk_size: int,
) -> None:
openai_types: Final = _get_openai_response_types()
@ -1174,16 +1220,19 @@ def _add_text_like_part_events(
delta=text[i : i + chunk_size],
)
)
annotations_payload: Final[Sequence[dict[str, object]]] = part_payload.get("annotations", []) or []
for annotation_index, annotation in enumerate(annotations_payload):
raw_annotation_items: Final = part_payload.get("annotations")
annotation_items: Final[Sequence[object]] = raw_annotation_items if _is_json_array(raw_annotation_items) else []
for annotation_index, annotation in enumerate(annotation_items):
events.append(
openai_types.OutputTextAnnotationAddedEvent(
type=openai_types.ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED,
item_id=item_id,
output_index=output_index,
content_index=content_index,
annotation_index=annotation_index,
annotation=annotation,
openai_types.OutputTextAnnotationAddedEvent.model_validate(
{
"type": openai_types.ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED,
"item_id": item_id,
"output_index": output_index,
"content_index": content_index,
"annotation_index": annotation_index,
"annotation": annotation,
}
)
)
events.append(
@ -1256,7 +1305,8 @@ def _build_synthetic_response_events(
)
if item_type == "message":
content_parts: Sequence[object] = output_item_payload.get("content", []) or []
raw_content_parts = output_item_payload.get("content")
content_parts: Sequence[object] = raw_content_parts if _is_json_array(raw_content_parts) else []
for content_index, part in enumerate(content_parts):
part_payload = _dump_response_object(part)
events.append(
@ -1304,8 +1354,9 @@ def _build_synthetic_response_events(
)
)
elif item_type == "reasoning":
summaries: Sequence[object] = output_item_payload.get("summary", []) or []
for summary_index, summary in enumerate(summaries):
raw_summary_items = output_item_payload.get("summary")
summary_items: Sequence[object] = raw_summary_items if _is_json_array(raw_summary_items) else []
for summary_index, summary in enumerate(summary_items):
summary_payload = _dump_response_object(summary)
summary_text = str(summary_payload.get("text") or "")
for i in range(0, len(summary_text), chunk_size):
@ -1476,7 +1527,7 @@ class ResponsesWebSocketStreaming:
user_api_key_dict: UserAPIKeyAuth | None = None,
request_data: dict[str, object] | None = None,
first_message: str | None = None,
guardrail_callbacks: list[PiiUnmaskingGuardrailCallback] | None = None,
guardrail_callbacks: Sequence[PresidioGuardrailCallback] | None = None,
output_guardrail_callbacks: list[PresidioGuardrailCallback] | None = None,
quota_callbacks: Sequence[ProjectQuotaCallback] | None = None,
authorized_model: str | None = None,
@ -1486,17 +1537,17 @@ class ResponsesWebSocketStreaming:
self.logging_obj = logging_obj
self.user_api_key_dict = user_api_key_dict
self.request_data: dict[str, object] = request_data or {}
self.messages: list[dict[str, object]] = []
self.messages: list[_MutableJsonObject] = []
self.input_messages: list[dict[str, object]] = []
self.first_message = first_message
self.guardrail_callbacks: list[PiiUnmaskingGuardrailCallback] = guardrail_callbacks or []
self.guardrail_callbacks: Sequence[PresidioGuardrailCallback] = guardrail_callbacks or []
self.output_guardrail_callbacks: list[PresidioGuardrailCallback] = output_guardrail_callbacks or []
self.quota_callbacks: tuple[ProjectQuotaCallback, ...] = tuple(quota_callbacks) if quota_callbacks else ()
# Model name authorized at connection time; enforced on every
# response.create frame to prevent deployment-substitution attacks.
self.authorized_model: str | None = authorized_model
def _should_store_event(self, event_obj: Mapping[str, object]) -> bool:
def _should_store_event(self, event_obj: _MutableJsonObject) -> bool:
return event_obj.get("type") in RESPONSES_WS_LOGGED_EVENT_TYPES
def _store_event(self, event: str | bytes | dict[str, object]) -> None:
@ -1610,7 +1661,7 @@ class ResponsesWebSocketStreaming:
finally:
await self._log_messages()
def _enforce_authorized_model(self, msg_obj: dict[str, object]) -> bool:
def _enforce_authorized_model(self, msg_obj: _MutableJsonObject) -> bool:
"""
Overwrite any ``model`` field in a ``response.create`` frame with the
connection-authorized model to prevent deployment-substitution attacks.
@ -1679,7 +1730,7 @@ class ResponsesWebSocketStreaming:
# forwarded unmasked regardless of where the client places it.
nested_candidate = msg_obj.get("response")
nested_response = nested_candidate if _is_json_object(nested_candidate) else None
text_containers: list[tuple[dict[str, object], str]] = []
text_containers: list[tuple[_MutableJsonObject, str]] = []
for container in (msg_obj, nested_response):
if container is None:
continue
@ -1786,6 +1837,7 @@ class ResponsesWebSocketStreaming:
return response_str
cb: Final = self.guardrail_callbacks[0]
unmask_pii_text: Final[_UnmasksPiiText] = getattr(cb, _UNMASK_PII_TEXT_ATTR)
event_type: Final = evt_obj.get("type")
if event_type == "response.completed":
@ -1805,9 +1857,7 @@ class ResponsesWebSocketStreaming:
continue
text = content_block.get("text")
if isinstance(text, str):
unmasked = cb._unmask_pii_text( # pyright: ignore[reportPrivateUsage] # no public unmasker
text, pii_tokens
)
unmasked = unmask_pii_text(text, pii_tokens)
if unmasked != text:
content_block["text"] = unmasked
modified = True
@ -1816,9 +1866,7 @@ class ResponsesWebSocketStreaming:
if event_type in self._DELTA_EVENT_TYPES:
delta: Final = evt_obj.get("delta")
if isinstance(delta, str):
unmasked = cb._unmask_pii_text( # pyright: ignore[reportPrivateUsage] # no public unmasker
delta, pii_tokens
)
unmasked = unmask_pii_text(delta, pii_tokens)
if unmasked != delta:
evt_obj["delta"] = unmasked
return json.dumps(evt_obj)
@ -2020,7 +2068,7 @@ class ManagedResponsesWebSocketHandler:
model: str,
logging_obj: LiteLLMLoggingObj,
user_api_key_dict: UserAPIKeyAuth | None = None,
litellm_metadata: dict[str, Any] | None = None,
litellm_metadata: Mapping[str, object] | None = None,
api_key: str | None = None,
api_base: str | None = None,
timeout: float | None = None,
@ -2033,10 +2081,11 @@ class ManagedResponsesWebSocketHandler:
self.model = model
self.logging_obj = logging_obj
self.user_api_key_dict = user_api_key_dict
self.litellm_metadata: dict[str, Any] = litellm_metadata or {}
self.model_group: str | None = self.litellm_metadata.get("model_group") or self.litellm_metadata.get(
self.litellm_metadata: Mapping[str, object] = litellm_metadata or {}
raw_model_group: Final = self.litellm_metadata.get("model_group") or self.litellm_metadata.get(
"deployment_model_name"
)
self.model_group: str | None = raw_model_group if isinstance(raw_model_group, str) else None
self.api_key = api_key
self.api_base = api_base
self.timeout = timeout
@ -2057,7 +2106,7 @@ class ManagedResponsesWebSocketHandler:
# ------------------------------------------------------------------
@staticmethod
def _serialize_chunk(chunk: Any) -> str | None:
def _serialize_chunk(chunk: object) -> str | None:
"""Serialize a streaming chunk to a JSON string for WebSocket transmission."""
try:
if isinstance(chunk, _HasModelDumpJson):
@ -2100,7 +2149,7 @@ class ManagedResponsesWebSocketHandler:
self._session_history[response_id] = messages
@staticmethod
def _extract_response_id(completed_event: dict[str, object]) -> str | None:
def _extract_response_id(completed_event: _MutableJsonObject) -> str | None:
"""
Pull the raw (decoded) response ID out of a ``response.completed`` event.
Returns *None* if the event doesn't contain a usable ID.
@ -2115,7 +2164,7 @@ class ManagedResponsesWebSocketHandler:
@staticmethod
def _extract_output_messages(
completed_event: dict[str, object],
completed_event: _MutableJsonObject,
) -> list[dict[str, object]]:
"""
Convert the output items in a ``response.completed`` event into
@ -2172,7 +2221,7 @@ class ManagedResponsesWebSocketHandler:
# _process_response_create sub-methods
# ------------------------------------------------------------------
async def _parse_message(self, raw_message: str) -> dict[str, object] | None:
async def _parse_message(self, raw_message: str) -> _MutableJsonObject | None:
"""Parse raw WS text; return the message dict or None (JSON error / ignored type)."""
try:
msg_obj: Final = _load_json_object(raw_message)
@ -2185,7 +2234,7 @@ class ManagedResponsesWebSocketHandler:
return msg_obj
@staticmethod
def _is_warmup_frame(msg_obj: dict[str, object]) -> bool:
def _is_warmup_frame(msg_obj: _MutableJsonObject) -> bool:
"""Return True for a response.create whose generate flag is false."""
nested: Final = msg_obj.get("response")
source: Final = nested if _is_json_object(nested) and nested else msg_obj
@ -2201,13 +2250,13 @@ class ManagedResponsesWebSocketHandler:
return str(raw_id).startswith(_WARMUP_RESPONSE_ID_PREFIX)
@staticmethod
def _warmup_source_params(msg_obj: dict[str, object]) -> dict[str, object]:
def _warmup_source_params(msg_obj: _MutableJsonObject) -> dict[str, object]:
nested: Final = msg_obj.get("response")
if _is_json_object(nested) and nested:
return nested
return {k: v for k, v in msg_obj.items() if k != "type"}
def _build_warmup_response(self, msg_obj: dict[str, object]) -> dict[str, object]:
def _build_warmup_response(self, msg_obj: _MutableJsonObject) -> dict[str, object]:
"""Build a minimal completed Responses API object for a warmup ack."""
source: Final = self._warmup_source_params(msg_obj)
wire_model: Final = source.get("model") or self.model_group or self.model
@ -2225,7 +2274,7 @@ class ManagedResponsesWebSocketHandler:
},
}
async def _send_warmup_ack(self, msg_obj: dict[str, object]) -> None:
async def _send_warmup_ack(self, msg_obj: _MutableJsonObject) -> None:
"""
Acknowledge a generate=false prewarm without calling the provider.
@ -2248,7 +2297,7 @@ class ManagedResponsesWebSocketHandler:
await self.websocket.send_text(serialized)
@staticmethod
def _build_base_call_kwargs(msg_obj: dict[str, object]) -> dict[str, Any]:
def _build_base_call_kwargs(msg_obj: _MutableJsonObject) -> dict[str, Any]:
"""
Extract Responses API params from the event, handling both wire formats:
Nested: {"type": "response.create", "response": {"input": [...], ...}}
@ -2357,7 +2406,7 @@ class ManagedResponsesWebSocketHandler:
call_kwargs.setdefault("litellm_params", {})
call_kwargs["litellm_params"]["proxy_server_request"] = proxy_server_request
async def _stream_and_forward(self, model: str, call_kwargs: dict[str, Any]) -> dict[str, object] | None:
async def _stream_and_forward(self, model: str, call_kwargs: dict[str, Any]) -> _MutableJsonObject | None:
"""
Stream ``litellm.aresponses`` and forward every chunk over the WebSocket.
@ -2365,7 +2414,7 @@ class ManagedResponsesWebSocketHandler:
directly (before serialization) to avoid a redundant JSON round-trip on
every chunk. Returns the completed event dict, or ``None``.
"""
completed_event: dict[str, object] | None = (
completed_event: _MutableJsonObject | None = (
None # rebind-ok: captures the completed event once the stream yields it
)
stream_response: Final = await litellm.aresponses(model=model, **call_kwargs)
@ -2391,7 +2440,7 @@ class ManagedResponsesWebSocketHandler:
def _save_turn_history(
self,
completed_event: dict[str, object] | None,
completed_event: _MutableJsonObject | None,
prior_history: list[dict[str, object]],
current_messages: list[dict[str, object]],
) -> None:
@ -2464,12 +2513,14 @@ class ManagedResponsesWebSocketHandler:
# reuse the router-resolved self.model; passing the alias raw to
# litellm.aresponses fails in get_llm_provider. A genuinely different
# provider-prefixed per-frame model is still honored.
requested_model: Final[str | None] = call_kwargs.pop("model", None)
requested_model: Final[str | None] = _typed_pops_optional_str(call_kwargs.pop)("model", None)
model: Final[str] = (
self.model if requested_model is None or requested_model == self.model_group else requested_model
)
previous_response_id: Final[str | None] = call_kwargs.pop("previous_response_id", None)
previous_response_id: Final[str | None] = _typed_pops_optional_str(call_kwargs.pop)(
"previous_response_id", None
)
current_messages: Final = self._input_to_messages(call_kwargs.get("input"))
# Fetch history once; reused in both _apply_history and _save_turn_history

212
litellm/router_strategy/tag_based_routing.py
View file

@ -8,16 +8,14 @@ Use this to route requests between Teams
"""
import re
from collections.abc import Iterable, Mapping, Sequence
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict
from typing_extensions import ReadOnly
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, overload
from litellm._logging import verbose_logger
from litellm.constants import CONSUMED_REQUEST_TAGS_METADATA_KEY
from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs
from litellm.types.router import ConsumedRequestTagsStamp, RouterErrors
from litellm.types.router import ConsumedRequestTagsStamp, DeploymentTypedDict, RouterErrors
if TYPE_CHECKING:
from litellm.router import Router as _Router
@ -27,34 +25,63 @@ else:
LitellmRouter = Any
class _TagRoutingLitellmParams(TypedDict, total=False):
tags: ReadOnly[Sequence[str] | None]
tag_regex: ReadOnly[Sequence[str] | None]
class _TagLitellmParamsLike(Protocol):
@overload
def get(self, key: Literal["tags"], /) -> Sequence[str] | None: ...
@overload
def get(self, key: Literal["tags"], default: Sequence[str], /) -> Sequence[str]: ...
@overload
def get(self, key: Literal["tag_regex"], /) -> Sequence[str] | None: ...
class _TagRoutingDeployment(TypedDict, total=False):
model_name: ReadOnly[str]
litellm_params: ReadOnly[_TagRoutingLitellmParams]
model_info: ReadOnly[Mapping[str, object] | None]
class _ModelInfoLike(Protocol):
@overload
def get(self, key: Literal["allow_fail_open"], /) -> bool | None: ...
@overload
def get(self, key: Literal["enable_tag_filtering"], /) -> bool | None: ...
class _TagRoutingMatchStamp(TypedDict):
matched_deployment: ReadOnly[str | None]
matched_via: ReadOnly[str]
matched_value: ReadOnly[str]
request_tags: ReadOnly[Sequence[str]]
user_agent: ReadOnly[str]
class _DeploymentLike(Protocol):
@overload
def get(self, key: Literal["litellm_params"], default: Mapping[str, object], /) -> _TagLitellmParamsLike: ...
@overload
def get(self, key: Literal["model_info"], /) -> _ModelInfoLike | None: ...
@overload
def get(self, key: Literal["model_name"], /) -> object: ...
class _TagRoutingMetadata(TypedDict, total=False):
tags: ReadOnly[Sequence[str] | None]
inherited_tags: ReadOnly[Sequence[str] | None]
user_agent: ReadOnly[str]
tag_routing: ReadOnly[_TagRoutingMatchStamp]
_consumed_request_tags: ReadOnly[object]
class _MetadataLike(Protocol):
@overload
def get(self, key: Literal["tags"], /) -> Sequence[str] | None: ...
@overload
def get(self, key: Literal["tags"], default: Sequence[str], /) -> Sequence[str] | None: ...
@overload
def get(self, key: Literal["user_agent"], default: str, /) -> str: ...
@overload
def get(self, key: Literal["inherited_tags"], /) -> object: ...
def __contains__(self, key: object, /) -> bool: ...
def __setitem__(self, key: Literal["tag_routing"], value: Mapping[str, object], /) -> None: ...
_EMPTY_MODEL_INFO: Final[Mapping[str, object]] = MappingProxyType({})
class _NestedLitellmParamsLike(Protocol):
def get(
self, key: Literal["metadata", "litellm_metadata"], default: Mapping[str, object], /
) -> _MetadataLike | None: ...
class _RequestKwargsLike(Protocol):
@overload
def get(self, key: Literal["enable_tag_filtering"], /) -> bool | None: ...
@overload
def get(self, key: Literal["metadata", "litellm_metadata"], /) -> _MetadataLike | None: ...
def __contains__(self, key: object, /) -> bool: ...
@overload
def __getitem__(self, key: Literal["metadata", "litellm_metadata"], /) -> _MetadataLike: ...
@overload
def __getitem__(self, key: Literal["litellm_params"], /) -> _NestedLitellmParamsLike: ...
_DeploymentPool = Sequence[_DeploymentLike] | Mapping[_DeploymentLike, object]
def _is_valid_deployment_tag_regex(
@ -109,11 +136,11 @@ def is_valid_deployment_tag(
def _match_deployment(
deployment: _TagRoutingDeployment,
request_tags: Sequence[str] | None,
header_strings: Sequence[str],
deployment: _DeploymentLike,
request_tags: list[str] | None,
header_strings: list[str],
match_any: bool,
) -> Mapping[str, str] | None:
) -> dict[str, str] | None:
"""
Determine whether *deployment* matches the current request.
@ -198,38 +225,38 @@ def _split_tags(tags: Sequence[str]) -> tuple[tuple[str, ...], list[str], tuple[
def _exclude_deployments(
deployments: Iterable[_TagRoutingDeployment],
deployments: _DeploymentPool,
excluded_set: frozenset[str],
) -> list[_TagRoutingDeployment]:
) -> Sequence[_DeploymentLike]:
if not excluded_set:
return list(deployments)
return [d for d in deployments if not excluded_set.intersection(d.get("litellm_params", {}).get("tags") or [])]
def _require_all_tags(
deployments: Iterable[_TagRoutingDeployment],
deployments: _DeploymentPool,
required_set: frozenset[str],
) -> tuple[_TagRoutingDeployment, ...]:
) -> tuple[_DeploymentLike, ...]:
if not required_set:
return tuple(deployments)
return tuple(d for d in deployments if required_set.issubset(d.get("litellm_params", {}).get("tags") or []))
def _default_tagged_pool(
deployments: Iterable[_TagRoutingDeployment],
) -> tuple[_TagRoutingDeployment, ...]:
deployments: _DeploymentPool,
) -> tuple[_DeploymentLike, ...]:
defaults: Final = tuple(d for d in deployments if "default" in (d.get("litellm_params", {}).get("tags") or []))
return defaults if defaults else tuple(deployments)
def _known_tag_values(deployments: Iterable[_TagRoutingDeployment]) -> frozenset[str]:
def _known_tag_values(deployments: _DeploymentPool) -> frozenset[str]:
return frozenset(
tag for d in deployments for tag in (d.get("litellm_params", _TagRoutingLitellmParams()).get("tags") or ())
tag for d in deployments for tag in (d.get("litellm_params", MappingProxyType({})).get("tags") or ())
)
def _unknown_required_tag_hides_an_answer(
healthy_deployments: Iterable[_TagRoutingDeployment],
healthy_deployments: _DeploymentPool,
excluded_set: frozenset[str],
required_set: frozenset[str],
routing_confirmed: frozenset[str],
@ -253,23 +280,23 @@ def _unknown_required_tag_hides_an_answer(
def _chain_allows_fail_open(
healthy_deployments: Iterable[_TagRoutingDeployment],
healthy_deployments: _DeploymentPool,
excluded_set: frozenset[str],
required_set: frozenset[str],
routing_confirmed: frozenset[str],
) -> bool:
if _unknown_required_tag_hides_an_answer(healthy_deployments, excluded_set, required_set, routing_confirmed):
return False
return any((d.get("model_info") or _EMPTY_MODEL_INFO).get("allow_fail_open") is True for d in healthy_deployments)
return any((d.get("model_info") or {}).get("allow_fail_open") is True for d in healthy_deployments)
def _trusted_only_pool(
healthy_deployments: Iterable[_TagRoutingDeployment],
healthy_deployments: _DeploymentPool,
excluded_set: frozenset[str],
required_set: frozenset[str],
inherited_excluded_set: frozenset[str] | None,
inherited_required_set: frozenset[str] | None,
) -> tuple[_TagRoutingDeployment, ...]:
) -> tuple[_DeploymentLike, ...]:
# inherited_*_set is None only when this request carries no origin information
# at all (e.g. direct SDK Router usage, bypassing the proxy layer that
# populates metadata.inherited_tags) -- treat every constraint as
@ -296,8 +323,8 @@ def _trusted_only_pool(
def _resolve_or_fail_open(
pool: Sequence[_TagRoutingDeployment],
healthy_deployments: Iterable[_TagRoutingDeployment],
pool: Sequence[_DeploymentLike],
healthy_deployments: _DeploymentPool,
excluded_set: frozenset[str],
required_set: frozenset[str],
inherited_excluded_set: frozenset[str] | None,
@ -305,7 +332,7 @@ def _resolve_or_fail_open(
routing_confirmed: frozenset[str],
model: str,
request_tags: object,
) -> tuple[_TagRoutingDeployment, ...]:
) -> tuple[_DeploymentLike, ...]:
if pool:
return tuple(pool)
if _chain_allows_fail_open(healthy_deployments, excluded_set, required_set, routing_confirmed):
@ -325,7 +352,7 @@ def _resolve_or_fail_open(
def _resolve_constraint_only_pool(
healthy_deployments: Iterable[_TagRoutingDeployment],
healthy_deployments: _DeploymentPool,
excluded_set: frozenset[str],
required_set: frozenset[str],
inherited_excluded_set: frozenset[str] | None,
@ -333,7 +360,7 @@ def _resolve_constraint_only_pool(
routing_confirmed: frozenset[str],
model: str,
request_tags: object,
) -> tuple[_TagRoutingDeployment, ...]:
) -> tuple[_DeploymentLike, ...]:
pool: Final = (
_require_all_tags(_exclude_deployments(healthy_deployments, excluded_set), required_set)
if required_set
@ -355,8 +382,8 @@ def _resolve_constraint_only_pool(
def _all_deployments_or_fallback(
llm_router_instance: LitellmRouter,
model: str,
fallback: Iterable[_TagRoutingDeployment],
) -> Iterable[_TagRoutingDeployment]:
fallback: _DeploymentPool,
) -> Sequence[_DeploymentLike | DeploymentTypedDict] | Mapping[_DeploymentLike, object]:
try:
return llm_router_instance._get_all_deployments(model_name=model)
except Exception: # noqa: BLE001 # fail safe toward today's healthy-only behavior on lookup errors
@ -366,8 +393,8 @@ def _all_deployments_or_fallback(
def _chain_tag_filtering_override(
llm_router_instance: LitellmRouter,
model: str,
healthy_deployments: Iterable[_TagRoutingDeployment],
) -> object:
healthy_deployments: _DeploymentPool,
) -> bool | None:
# Resolved from every deployment configured for this model group, not just the
# ones that survived cooldown/health filtering (async_get_healthy_deployments
# filters cooldowns before calling get_deployments_for_tag) -- otherwise the
@ -379,14 +406,14 @@ def _chain_tag_filtering_override(
# than crashing the request.
all_deployments: Final = _all_deployments_or_fallback(llm_router_instance, model, healthy_deployments)
for d in all_deployments:
value = (d.get("model_info") or _EMPTY_MODEL_INFO).get("enable_tag_filtering")
value = (d.get("model_info") or MappingProxyType({})).get("enable_tag_filtering")
if value is not None:
return value
return None
def _inherited_constraint_sets(
inherited_tags: Sequence[str] | None, routing_prefix: str
inherited_tags: object, routing_prefix: str
) -> tuple[frozenset[str] | None, frozenset[str] | None]:
# None means no origin information is available at all (e.g. this request
# bypassed the proxy layer that populates metadata.inherited_tags, as direct
@ -417,43 +444,42 @@ def _tag_known_to_group(
if tag_set & routing_confirmed:
return True
try:
all_deployments: Final[Sequence[_TagRoutingDeployment]] = llm_router_instance._get_all_deployments(
model_name=model
)
all_deployments: Final = llm_router_instance._get_all_deployments(model_name=model)
except Exception: # noqa: BLE001 # fail safe toward "unrecognized" so lookup errors preserve the existing silent-fallback behavior
return False
return any(
tag_set.intersection(d.get("litellm_params", _TagRoutingLitellmParams()).get("tags") or ())
for d in all_deployments
tag_set.intersection(d.get("litellm_params", MappingProxyType({})).get("tags") or ()) for d in all_deployments
)
def _request_tags_after_router_consumption(metadata: _TagRoutingMetadata, model: str) -> Sequence[str] | None:
def _request_tags_after_router_consumption(metadata: object, model: str) -> Sequence[str] | None:
# The pre-routing hook stamps which tags selected the router it rewrote the request
# to: those tags already did their job and must not also constrain deployment choice
# inside the routed group. The request's other tags still apply there, on top of the
# inherited_tags snapshot that keeps key/team policy applying. Every other model
# group keeps the full list.
stamp: Final = metadata.get(CONSUMED_REQUEST_TAGS_METADATA_KEY)
if not isinstance(metadata, Mapping):
return None
typed_metadata: Final[Mapping[str, object]] = metadata
request_tags: Final = _tags_in_metadata(typed_metadata)
stamp: Final = typed_metadata.get(CONSUMED_REQUEST_TAGS_METADATA_KEY)
if not isinstance(stamp, ConsumedRequestTagsStamp) or stamp.model_group != model:
return metadata.get("tags")
request_tags: Final = metadata.get("tags")
leftover: Final = tuple(
tag for tag in (request_tags if isinstance(request_tags, (list, tuple)) else ()) if tag not in stamp.tags
)
inherited_tags: Final = metadata.get("inherited_tags")
return request_tags
leftover: Final = tuple(tag for tag in request_tags if tag not in stamp.tags)
inherited_tags: Final = typed_metadata.get("inherited_tags")
if not isinstance(inherited_tags, (list, tuple)):
return leftover or None
return tuple(dict.fromkeys((*leftover, *inherited_tags)))
typed_inherited_tags: Final[Sequence[object]] = inherited_tags
return tuple(dict.fromkeys((*leftover, *(tag for tag in typed_inherited_tags if isinstance(tag, str)))))
async def get_deployments_for_tag(
llm_router_instance: LitellmRouter,
model: str, # used to raise the correct error
healthy_deployments: list[Any] | dict[Any, Any],
request_kwargs: dict[Any, Any] | None = None,
healthy_deployments: _DeploymentPool,
request_kwargs: _RequestKwargsLike | None = None,
metadata_variable_name: Literal["metadata", "litellm_metadata"] = "metadata",
):
) -> _DeploymentPool:
"""
Returns a list of deployments that match the requested model and tags in the request.
@ -486,8 +512,7 @@ async def get_deployments_for_tag(
verbose_logger.debug("request metadata: %s", request_kwargs.get(metadata_variable_name))
if metadata_variable_name in request_kwargs:
metadata: Final[_TagRoutingMetadata] = request_kwargs[metadata_variable_name]
stampable_metadata: Final[dict[str, object]] = request_kwargs[metadata_variable_name]
metadata: Final = request_kwargs[metadata_variable_name]
request_tags: Final = _request_tags_after_router_consumption(metadata, model)
match_any: Final = llm_router_instance.tag_filtering_match_any
routing_prefix: Final = llm_router_instance.tag_routing_prefix or ""
@ -532,25 +557,25 @@ async def get_deployments_for_tag(
request_tags,
)
new_healthy_deployments: Final[list[_TagRoutingDeployment]] = []
default_deployments: Final[list[_TagRoutingDeployment]] = []
if has_positive_filter:
verbose_logger.debug(
"get_deployments_for_tag routing: request_tags=%s user_agent=%s",
request_tags,
user_agent,
)
for deployment in candidates:
deployment_tags = deployment.get("litellm_params", {}).get("tags")
match_result = _match_deployment(
deployment=deployment,
request_tags=positive_tags,
header_strings=header_strings,
match_any=match_any,
deployment_matches: Final = tuple(
(
deployment,
_match_deployment(
deployment=deployment,
request_tags=positive_tags,
header_strings=header_strings,
match_any=match_any,
),
)
for deployment in candidates
)
for deployment, match_result in deployment_matches:
if match_result is not None:
verbose_logger.debug(
"tag routing match: deployment=%s matched_via=%s matched_value=%s",
@ -559,17 +584,17 @@ async def get_deployments_for_tag(
match_result["matched_value"],
)
if "tag_routing" not in metadata:
stampable_metadata["tag_routing"] = {
metadata["tag_routing"] = {
"matched_deployment": deployment.get("model_name"),
"matched_via": match_result["matched_via"],
"matched_value": match_result["matched_value"],
"request_tags": request_tags or [],
"user_agent": user_agent,
}
new_healthy_deployments.append(deployment)
if deployment_tags and "default" in deployment_tags:
default_deployments.append(deployment)
new_healthy_deployments: Final = [d for d, result in deployment_matches if result is not None]
default_deployments: Final = [
d for d, _ in deployment_matches if "default" in (d.get("litellm_params", {}).get("tags") or ())
]
if len(new_healthy_deployments) == 0 and len(default_deployments) == 0:
return _resolve_or_fail_open(
@ -604,10 +629,11 @@ async def get_deployments_for_tag(
return new_healthy_deployments if len(new_healthy_deployments) > 0 else default_deployments
# for Untagged requests use default deployments if set
_default_deployments_with_tags: Final[list[_TagRoutingDeployment]] = []
for deployment in healthy_deployments:
if "default" in deployment.get("litellm_params", {}).get("tags", []):
_default_deployments_with_tags.append(deployment)
_default_deployments_with_tags: Final = [
deployment
for deployment in healthy_deployments
if "default" in deployment.get("litellm_params", {}).get("tags", [])
]
if len(_default_deployments_with_tags) > 0:
return _default_deployments_with_tags

14
ruff-strict-budget.json
View file

@ -1,6 +1,6 @@
{
"ANN001": {
"limit": 2998
"limit": 2995
},
"ANN002": {
"limit": 71
@ -9,7 +9,7 @@
"limit": 809
},
"ANN201": {
"limit": 2003
"limit": 2002
},
"ANN202": {
"limit": 845
@ -24,7 +24,7 @@
"limit": 133
},
"ANN401": {
"limit": 597
"limit": 587
},
"ASYNC230": {
"limit": 11
@ -123,7 +123,7 @@
"limit": 12
},
"PERF403": {
"limit": 34
"limit": 33
},
"PIE804": {
"limit": 18
@ -177,7 +177,7 @@
"limit": 8
},
"RUF019": {
"limit": 32
"limit": 31
},
"RUF046": {
"limit": 4
@ -198,7 +198,7 @@
"limit": 56
},
"SIM102": {
"limit": 315
"limit": 314
},
"SIM103": {
"limit": 119
@ -231,7 +231,7 @@
"limit": 5
},
"TID251": {
"limit": 1111
"limit": 1108
},
"TRY002": {
"limit": 524

12
type-discipline-budget.json
View file

@ -1,9 +1,9 @@
{
"LIT001": {
"limit": 22642
"limit": 22521
},
"LIT002": {
"limit": 26834
"limit": 26820
},
"LIT003": {
"limit": 269
@ -15,7 +15,7 @@
"limit": 0
},
"LIT006": {
"limit": 1041
"limit": 1039
},
"LIT007": {
"limit": 0
@ -27,12 +27,12 @@
"limit": 0
},
"LIT010": {
"limit": 16550
"limit": 16546
},
"LIT011": {
"limit": 5576
"limit": 5575
},
"LIT012": {
"limit": 4496
"limit": 4495
}
}