scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

chore(typing): clear 1.2k basedpyright Any errors across 16 hotspot files

Browse source 
Replace Any-typed seams with real types in the files carrying the highest
remaining reportAny/reportExplicitAny density: Literal-keyed structural
Protocols for deployment dicts in tag-based routing, typed Prisma table
wrappers and row protocols in the key and internal-user management
endpoints, TypedDict views for websearch interception kwargs, typed
streaming state in the responses iterator and background polling, and
concrete request/response types in the google_genai, vertex_ai files,
runwayml, rubrik, anthropic context-management, and guardrail translation
modules. Mutable annotations introduced along the way were rewritten as
read-only views (Mapping/Sequence/tuple) built functionally.

No casts, no type: ignore, no noqa, no suppression comments, no new Any
annotations, no behavior changes. Whole-tree basedpyright: reportAny
15,496 -> 14,523, reportExplicitAny 5,356 -> 5,102, all rules
145,547 -> 143,989, with no rule increased repo-wide or per-file.
Budgets ratcheted: basedpyright -1,545, ruff-strict -73,
type-discipline -237.
This commit is contained in:
mateo-berri 2026-08-12 19:41:30 -07:00
parent 0ca0fa22b8
commit cb65bf08b8
19 changed files with 2261 additions and 943 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

22
basedpyright-code-budget.json
View file

@ -1,9 +1,9 @@
{
"reportAny": {
"limit": 23919
"limit": 21974
},
"reportArgumentType": {
"limit": 2580
"limit": 2575
},
"reportAssignmentType": {
"limit": 323
@ -24,7 +24,7 @@
"limit": 19
},
"reportExplicitAny": {
"limit": 7573
"limit": 7068
},
"reportFunctionMemberAccess": {
"limit": 7
@ -54,10 +54,10 @@
"limit": 0
},
"reportMissingParameterType": {
"limit": 5719
"limit": 5697
},
"reportMissingTypeArgument": {
"limit": 15657
"limit": 15627
},
"reportMissingTypeStubs": {
"limit": 40
@ -99,19 +99,19 @@
"limit": 0
},
"reportUnknownArgumentType": {
"limit": 44832
"limit": 44549
},
"reportUnknownLambdaType": {
"limit": 113
"limit": 112
},
"reportUnknownMemberType": {
"limit": 39269
"limit": 39156
},
"reportUnknownParameterType": {
"limit": 19988
"limit": 19951
},
"reportUnknownVariableType": {
"limit": 30923
"limit": 30798
},
"reportUnnecessaryCast": {
"limit": 118
@ -123,7 +123,7 @@
"limit": 5
},
"reportUnnecessaryIsInstance": {
"limit": 853
"limit": 852
},
"reportUntypedBaseClass": {
"limit": 0

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

@ -1,6 +1,9 @@
import json
from collections.abc import AsyncIterator, Iterator
from typing import Any, Final, 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 TypedDict
from litellm import verbose_logger
from litellm.litellm_core_utils.json_validation_rule import normalize_tool_schema
@ -9,7 +12,6 @@ from litellm.types.llms.openai import (
ChatCompletionAssistantMessage,
ChatCompletionAssistantToolCall,
ChatCompletionImageObject,
ChatCompletionRequest,
ChatCompletionSystemMessage,
ChatCompletionTextObject,
ChatCompletionToolCallFunctionChunk,
@ -21,12 +23,79 @@ 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,
)
_JsonDict: TypeAlias = dict[str, object]
_JsonDictList: TypeAlias = list[_JsonDict]
class _ToolCallAccumulator(TypedDict):
name: str
arguments: str
class _GenAIFunctionCall(TypedDict):
name: str
args: Mapping[str, object]
class _GenAIPart(TypedDict, total=False):
text: str
functionCall: _GenAIFunctionCall
class _GenAIFunctionResponse(TypedDict, total=False):
name: str
response: object
class _GenAIRequestFunctionCall(TypedDict, total=False):
name: str
args: Mapping[str, object]
class _GenAIContentPart(TypedDict, total=False):
text: str
inline_data: Mapping[str, str]
functionResponse: _GenAIFunctionResponse
functionCall: _GenAIRequestFunctionCall
class _GenAIFunctionDeclaration(TypedDict, total=False):
name: str
description: str
parametersJsonSchema: object
class _GenAITool(TypedDict, total=False):
functionDeclarations: Sequence[_GenAIFunctionDeclaration]
class _GenAIFunctionCallingConfig(TypedDict, total=False):
mode: str
class _GenAIToolConfig(TypedDict, total=False):
functionCallingConfig: _GenAIFunctionCallingConfig
class _GenAISystemInstruction(TypedDict, total=False):
parts: Sequence[Mapping[str, str]]
_EMPTY_STR_MAPPING: Final[Mapping[str, str]] = MappingProxyType({})
class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
"""
@ -35,12 +104,12 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
"""
sent_first_chunk: bool = False
# State tracking for accumulating partial tool calls
accumulated_tool_calls: dict[str, dict[str, Any]]
_parse_accumulated_args: Callable[[str], Mapping[str, object]] = staticmethod(json.loads)
def __init__(self, completion_stream: Any):
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)
@ -85,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 = []
parts: Final = list[_GenAIPart]()
for (
tool_call_index,
tool_call_data,
@ -93,8 +162,10 @@ 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 = json.loads(tool_call_data["arguments"] or "{}")
function_call_part = {
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",
"args": parsed_args,
@ -172,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]:
@ -211,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,
}
@ -273,9 +346,9 @@ class GoogleGenAIAdapter:
def _add_generic_litellm_params_to_request(
self,
completion_request_dict: dict[str, Any],
completion_request_dict: _JsonDict,
litellm_params: GenericLiteLLMParams | None = None,
) -> dict:
) -> _JsonDict:
"""Add generic litellm params to request. e.g add api_base, api_key, api_version, etc.
Args:
@ -287,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
@ -295,7 +368,7 @@ class GoogleGenAIAdapter:
def translate_completion_output_params_streaming(
self,
completion_stream: Any,
completion_stream: object,
) -> AsyncIterator[bytes] | None:
"""Transform streaming completion output to Google GenAI format"""
google_genai_wrapper: Final = GoogleGenAIStreamWrapper(completion_stream=completion_stream)
@ -304,15 +377,15 @@ class GoogleGenAIAdapter:
def _transform_google_genai_tools_to_openai(
self,
tools: list[dict[str, Any]],
tools: Sequence[_GenAITool],
) -> list[ChatCompletionToolParam]:
"""Transform Google GenAI tools to OpenAI tools format"""
openai_tools: Final[list[dict[str, Any]]] = []
openai_tools: Final = list[_JsonDict]()
for tool in tools:
if "functionDeclarations" in tool:
for func_decl in tool["functionDeclarations"]:
function_chunk: dict[str, Any] = {
function_chunk: _JsonDict = {
"name": func_decl.get("name", ""),
}
@ -321,7 +394,7 @@ class GoogleGenAIAdapter:
if "parametersJsonSchema" in func_decl:
function_chunk["parameters"] = func_decl["parametersJsonSchema"]
openai_tool = {"type": "function", "function": function_chunk}
openai_tool: _JsonDict = {"type": "function", "function": function_chunk}
openai_tools.append(openai_tool)
# normalize the tool schemas
@ -331,7 +404,7 @@ class GoogleGenAIAdapter:
def _transform_google_genai_tool_config_to_openai(
self,
tool_config: dict[str, Any],
tool_config: _GenAIToolConfig,
) -> ChatCompletionToolChoiceValues | None:
"""Transform Google GenAI tool_config to OpenAI tool_choice"""
function_calling_config: Final = tool_config.get("functionCallingConfig", {})
@ -345,20 +418,20 @@ class GoogleGenAIAdapter:
def _transform_contents_to_messages(
self,
contents: list[dict[str, Any]],
system_instruction: dict[str, Any] | None = None,
system_instruction: _GenAISystemInstruction | None = None,
) -> list[AllMessageValues]:
"""Transform Google GenAI contents to OpenAI messages format"""
messages: Final[list[AllMessageValues]] = []
# 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
@ -461,7 +534,7 @@ class GoogleGenAIAdapter:
def translate_completion_to_generate_content(
self,
response: ModelResponse,
) -> dict[str, Any]:
) -> _JsonDict:
"""
Transform litellm completion response to Google GenAI generate_content format
@ -484,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, Any]] = {
generate_content_response: Final[_JsonDict] = {
"candidates": [
{
"content": {"parts": parts, "role": "model"},
@ -524,7 +597,7 @@ class GoogleGenAIAdapter:
self,
response: ModelResponse | ModelResponseStream,
wrapper: GoogleGenAIStreamWrapper,
) -> dict[str, Any] | None:
) -> Mapping[str, object] | None:
"""
Transform streaming litellm completion chunk to Google GenAI generate_content format
@ -548,10 +621,10 @@ class GoogleGenAIAdapter:
parts = self._transform_openai_delta_to_google_genai_parts_with_accumulation(choice.delta, wrapper)
else:
parts = []
finish_reason = getattr(choice, "finish_reason", None)
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)
@ -560,7 +633,7 @@ class GoogleGenAIAdapter:
return None
# Create Google GenAI streaming format response
streaming_chunk: Final[dict[str, Any]] = {
streaming_chunk: Final[_JsonDict] = {
"candidates": [
{
"content": {"parts": parts, "role": "model"},
@ -596,10 +669,10 @@ class GoogleGenAIAdapter:
def _transform_openai_message_to_google_genai_parts(
self,
message: Any,
) -> list[dict[str, Any]]:
message: Message,
) -> Sequence[_GenAIPart]:
"""Transform OpenAI message to Google GenAI parts format"""
parts: Final[list[dict[str, Any]]] = []
parts: Final = list[_GenAIPart]()
# Add text content if present
if hasattr(message, "content") and message.content:
@ -607,16 +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 = json.loads(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 = {
function_call_part: _GenAIPart = {
"functionCall": {
"name": tool_call.function.name or "undefined_tool_name",
"name": function.name or "undefined_tool_name",
"args": args,
}
}
@ -625,28 +704,30 @@ class GoogleGenAIAdapter:
return parts if parts else [{"text": ""}]
def _transform_openai_delta_to_google_genai_parts_with_accumulation(
self, delta: Any, wrapper: GoogleGenAIStreamWrapper
) -> list[dict[str, Any]]:
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[dict[str, Any]]] = []
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"):
continue
# 3. Use `index` as the primary key for accumulation
tool_call_index = getattr(tool_call, "index", None)
tool_call_index: int | None = getattr(tool_call, "index", None)
if tool_call_index is None:
continue # Index is essential for tracking streaming tool calls
@ -658,8 +739,9 @@ 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:
@ -680,13 +762,13 @@ class GoogleGenAIAdapter:
# 5. Attempt to parse arguments even if name hasn't arrived.
try:
# Attempt to parse the accumulated arguments string
parsed_args = json.loads(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.
if accumulated_name:
# If successful, create the part and clean up
function_call_part = {"functionCall": {"name": accumulated_name, "args": parsed_args}}
function_call_part: _GenAIPart = {"functionCall": {"name": accumulated_name, "args": parsed_args}}
parts.append(function_call_part)
# Remove the completed tool call from the accumulator
@ -714,7 +796,7 @@ class GoogleGenAIAdapter:
return mapping.get(finish_reason, "STOP")
def _map_usage(self, usage: Any) -> 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,

228
litellm/integrations/rubrik.py
View file

@ -6,12 +6,14 @@ import random
import time
import uuid
from collections import Counter
from collections.abc import Mapping, Sequence
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
from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, TypedDict, overload
import httpx
from typing_extensions import Never, Required
from litellm._logging import verbose_logger
from litellm.integrations.custom_batch_logger import CustomBatchLogger
@ -29,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,
@ -48,7 +51,105 @@ _WEBHOOK_PATH_PROMPT_MODERATION: Final = "/v1/before_prompt/openai/v1"
_WEBHOOK_PATH_LOGGING_BATCH: Final = "/v1/litellm/batch"
_MAX_QUEUE_SIZE: Final = 10_000
_DROP_WARNING_INTERVAL_SECONDS: Final = 60.0
_EMPTY_MAPPING: Final[Mapping[str, Any]] = MappingProxyType({})
_EMPTY_MAPPING: Final[Mapping[str, Never]] = MappingProxyType({})
class _ModerationToolCall(TypedDict, total=False):
id: Required[str]
class _ModerationMessage(TypedDict, total=False):
content: str | None
tool_calls: Sequence[_ModerationToolCall] | None
class _ModerationChoice(TypedDict, total=False):
message: _ModerationMessage | None
class _ModerationResponse(TypedDict, total=False):
choices: Sequence[_ModerationChoice]
class _LogEventKwargs(TypedDict, total=False):
standard_logging_object: Required[StandardLoggingPayload]
litellm_call_id: 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: Mapping[str, object] | None
class _ModerationSourceMessage(TypedDict, total=False):
role: str
function_call: Mapping[str, object] | None
tool_calls: Sequence[_ModerationSourceToolCall | None] | None
class _FlattenedModerationMessage(TypedDict):
role: str | None
content: str
class _CorrelatablePayload(TypedDict):
id: str
class _SystemPromptCarrier(TypedDict, total=False):
messages: object
class _BlockFailurePayload(TypedDict, total=False):
id: object
model: object
model_group: object
model_id: str
model_parameters: object
startTime: float | None
endTime: float | None
completionStartTime: float | None
messages: object
metadata: StandardLoggingUserAPIKeyMetadata
response: str
status: str
class _MalformedToolBlockingResponseError(Exception):
@ -143,7 +244,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
else {"Content-Type": "application/json"}
)
self._periodic_flush_task: asyncio.Task[Any] | None = self._start_periodic_flush_task()
self._periodic_flush_task: asyncio.Task[None] | None = self._start_periodic_flush_task()
@classmethod
def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]:
@ -191,7 +292,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
params={"timeout": httpx.Timeout(5.0, connect=2.0)},
)
def _start_periodic_flush_task(self) -> asyncio.Task[Any] | None:
def _start_periodic_flush_task(self) -> asyncio.Task[None] | None:
"""Start the periodic flush task only when an event loop is already running."""
try:
loop: Final = asyncio.get_running_loop()
@ -212,7 +313,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
Closing them here would close the shared connection pool for every
other logger instance; let LiteLLM manage their lifecycle instead.
"""
task: Final = getattr(self, "_periodic_flush_task", None)
task: Final[asyncio.Task[None] | None] = getattr(self, "_periodic_flush_task", None)
if task is not None:
task.cancel()
@ -253,7 +354,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
@staticmethod
async def _guarded(
coro: Any,
coro: Awaitable[GenericGuardrailAPIInputs],
inputs: GenericGuardrailAPIInputs,
label: str,
) -> GenericGuardrailAPIInputs:
@ -371,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.
@ -400,12 +501,16 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
request_data["_rubrik_logging_obj"] = logging_obj
@staticmethod
def _normalize_tool_calls(tool_calls: Any) -> 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):
@ -427,7 +532,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
raise TypeError(f"Cannot normalize tool_call of type {type(tc).__name__}: {tc!r}")
@staticmethod
def _join_texts(texts: Any) -> str:
def _join_texts(texts: Sequence[str] | None) -> str:
"""Join response text segments into the single content string the
webhook evaluates. Empty when there is no assistant text."""
if not texts:
@ -439,19 +544,22 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
tool_calls: Sequence[ChatCompletionMessageToolCall],
content: str,
request_id: str | None,
) -> Mapping[str, Any]:
) -> Mapping[str, object]:
"""Build an OpenAI ChatCompletion-format dict (assistant text + tool
calls) for the after_completion webhook.
``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, Any]] = {
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",
@ -467,7 +575,9 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
}
@staticmethod
def _flatten_messages_for_moderation(messages: Any) -> 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,
@ -488,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")
@ -506,8 +616,8 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
@staticmethod
def _build_prompt_moderation_payload(
inputs: GenericGuardrailAPIInputs,
request_data: Mapping[str, Any],
) -> Mapping[str, Any]:
request_data: Mapping[str, object],
) -> Mapping[str, object]:
"""Build the bare OpenAI request the before_prompt webhook consumes.
Unlike the after_completion envelope, this endpoint takes a raw OpenAI
@ -516,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, Any]] = {
"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.
@ -533,15 +635,19 @@ 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, Any] | None,
) -> Mapping[str, Any]:
call_details: _RequestContextSource,
request_data: _RequestContextSource | None,
) -> Mapping[str, object]:
"""Extract original request data from model_call_details for the
response moderation service envelope.
@ -576,7 +682,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
}
@staticmethod
def _sanitize_proxy_server_request(proxy_server_request: Any) -> Any:
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
@ -586,7 +692,7 @@ 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, Any], call_details: Mapping[str, Any]) -> 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"):
@ -596,7 +702,9 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
# -- Logging hooks ---------------------------------------------------------
@staticmethod
def _correlation_id(call_details: Mapping[str, Any], request_data: Mapping[str, Any] | 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.
@ -610,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, Any], source: Mapping[str, Any]) -> 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.
@ -630,7 +738,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
payload["id"] = correlated
@staticmethod
def _prepend_system_prompt(payload: dict[str, Any], source: Mapping[str, Any]) -> 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"]``
@ -658,7 +766,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
exc_info=True,
)
async def _prepare_log_payload(self, kwargs: Mapping[str, Any], 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)
@ -667,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()
@ -697,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, Any], 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:
@ -818,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**:
@ -860,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 = 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 "
@ -884,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.
@ -906,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, Any]:
) -> _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")
@ -942,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:
@ -996,7 +1108,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
# -- Webhook services ------------------------------------------------------
async def _post_json(self, endpoint: str, payload: Mapping[str, Any], 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:
@ -1006,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 = http_response.json()
result: Final[_ModerationResponse | None] = http_response.json()
if not isinstance(result, dict):
raise TypeError(
f"{service_name} returned non-dict JSON "
@ -1021,9 +1133,9 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
async def _post_to_response_moderation_endpoint(
self,
response_data: Mapping[str, Any],
request_data: Mapping[str, Any],
) -> Mapping[str, Any]:
response_data: Mapping[str, object],
request_data: Mapping[str, object],
) -> _ModerationResponse:
"""Post the ``{request, response}`` envelope to the after_completion
webhook and return its (possibly rewritten) response.
@ -1039,7 +1151,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
"Response moderation service",
)
async def _post_to_prompt_moderation_endpoint(self, payload: Mapping[str, Any]) -> 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).
@ -1047,7 +1159,7 @@ 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
@ -1063,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:

157
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, cast
from typing import TYPE_CHECKING, Any, Final, Literal, Never, TypedDict, TypeVar, cast
import litellm
from litellm._logging import verbose_logger
@ -41,7 +41,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,
@ -51,6 +57,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,
@ -72,6 +80,8 @@ WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY: Final = "_websearch_interception_emit_native_b
# ``web_search_tool_result`` blocks to inject into the final response.
WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY: Final = "websearch_native_blocks"
_ResponseT = TypeVar("_ResponseT")
class _PlanMetadataView(TypedDict):
websearch_native_blocks: Sequence[Mapping[str, object]] | None
@ -85,9 +95,96 @@ class _WebSearchSettingsView(TypedDict):
websearch_interception_params: WebSearchInterceptionConfig
class _SearchToolLitellmParams(TypedDict, total=False):
search_provider: str | None
class _SearchToolConfig(TypedDict, total=False):
search_tool_name: str
litellm_params: Mapping[str, object] | None
litellm_params: _SearchToolLitellmParams | None
class _LitellmParamsProviderView(TypedDict, total=False):
custom_llm_provider: str
class _DeploymentCallKwargsView(TypedDict):
custom_llm_provider: str
litellm_params: _LitellmParamsProviderView
model: str
class _AcreateNamedParams(TypedDict, total=False):
metadata: Never
stop_sequences: Never
stream: bool | None
system: str | None
temperature: float | None
thinking: Never
tool_choice: Never
tools: Never
top_k: int | None
top_p: float | None
container: Never
class _AsearchNamedParams(TypedDict, total=False):
max_results: int | None
search_domain_filter: Never
max_tokens_per_page: int | None
country: str | None
api_key: str | None
api_base: str | None
timeout: float | None
extra_headers: Never
class _AcompletionNamedParams(TypedDict, total=False):
functions: Never
function_call: str | None
timeout: float | None
temperature: float | None
top_p: float | None
n: int | None
stream: bool | None
stream_options: Never
stop: Never
max_tokens: int | None
max_completion_tokens: int | None
modalities: Never
prediction: ChatCompletionPredictionContentParam | None
audio: ChatCompletionAudioParam | None
presence_penalty: float | None
frequency_penalty: float | None
logit_bias: Never
user: str | None
response_format: Never
seed: int | None
tools: Never
tool_choice: Never
parallel_tool_calls: bool | None
logprobs: bool | None
top_logprobs: int | None
deployment_id: str | None
reasoning_effort: Literal["none", "minimal", "low", "medium", "high", "xhigh", "default"] | None
verbosity: Literal["low", "medium", "high"] | None
safety_identifier: str | None
service_tier: str | None
base_url: str | None
api_version: str | None
api_key: str | None
model_list: Never
extra_headers: Never
thinking: AnthropicThinkingParam | None
web_search_options: OpenAIWebSearchOptions | None
include_server_side_tool_invocations: bool | None
shared_session: "ClientSession | None"
enable_json_schema_validation: bool | None
_NO_ACREATE_NAMED: Final[_AcreateNamedParams] = {}
_NO_ASEARCH_NAMED: Final[_AsearchNamedParams] = {}
_NO_ACOMPLETION_NAMED: Final[_AcompletionNamedParams] = {}
class WebSearchInterceptionLogger(CustomLogger):
@ -275,12 +372,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
custom_llm_provider = kwargs.get("custom_llm_provider", "") or kwargs.get("litellm_params", {}).get(
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 = 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.get("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:
@ -903,17 +1005,18 @@ class WebSearchInterceptionLogger(CustomLogger):
)
@staticmethod
def _inject_native_blocks(response: Any, native_blocks: Sequence[Mapping[str, object]]) -> Any:
def _inject_native_blocks(response: _ResponseT, native_blocks: Sequence[Mapping[str, object]]) -> _ResponseT:
"""Prepend native blocks to response content, dict or object form."""
if not native_blocks:
return response
if isinstance(response, dict):
existing = response.get("content") or []
existing: Sequence[object] = response.get("content") or []
response["content"] = list(native_blocks) + list(existing)
return response
existing = getattr(response, "content", None) or []
content_attribute: Final = "content"
try:
response.content = list(native_blocks) + list(existing)
setattr(response, content_attribute, list(native_blocks) + list(existing))
except (AttributeError, TypeError):
# Object refused write — fall through and leave the response
# untouched rather than crash the request.
@ -1169,10 +1272,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(
@ -1180,9 +1283,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")
@ -1197,12 +1300,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
@ -1344,12 +1449,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:
@ -1377,12 +1483,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,
)
)
@ -1422,7 +1531,7 @@ class WebSearchInterceptionLogger(CustomLogger):
valid_token=user_api_key_auth,
)
team_id: Final = getattr(user_api_key_auth, "team_id", None)
team_id: Final[str | None] = getattr(user_api_key_auth, "team_id", None)
if team_id:
from litellm.proxy.proxy_server import (
prisma_client,
@ -1537,10 +1646,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."""
@ -1548,8 +1657,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:
@ -1557,11 +1666,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(

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

@ -13,12 +13,12 @@ Pattern Overview:
"""
import json
from collections.abc import Mapping
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 TypedDict, assert_never
from litellm._logging import verbose_proxy_logger
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
@ -61,6 +61,7 @@ if TYPE_CHECKING:
ModifyResponseException,
)
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.llms.anthropic_messages.anthropic_response import (
AnthropicMessagesResponse,
)
@ -95,6 +96,48 @@ InputWriteBackTarget = (
)
class _SSEDelta(TypedDict, total=False):
type: str
text: str
stop_reason: str | None
class _SSEEventData(TypedDict, total=False):
delta: _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
@ -123,7 +166,7 @@ class AnthropicMessagesHandler(BaseTranslation):
@staticmethod
def _build_streaming_usage_response(
responses_so_far: list[Any],
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)))
@ -141,7 +184,7 @@ class AnthropicMessagesHandler(BaseTranslation):
self,
exc: "ModifyResponseException",
stream_started: bool = False,
responses_so_far: list[Any] | None = None,
responses_so_far: Sequence[object] | None = None,
) -> list[bytes]:
"""
Build an Anthropic SSE sequence delivering the guardrail block message
@ -159,7 +202,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]:
@ -184,7 +227,9 @@ class AnthropicMessagesHandler(BaseTranslation):
)
return list(FakeAnthropicMessagesStreamIterator(response=block_response))
def _block_continuation_chunks(self, exc: "ModifyResponseException", responses_so_far: list[Any]) -> 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."""
@ -234,7 +279,7 @@ class AnthropicMessagesHandler(BaseTranslation):
@staticmethod
def _content_block_state(
responses_so_far: list[Any],
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).
@ -260,7 +305,20 @@ class AnthropicMessagesHandler(BaseTranslation):
return open_index, max_index
@staticmethod
def _iter_sse_events(item: Any) -> list[dict]:
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
@ -268,22 +326,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]] = []
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 = 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."""
@ -315,8 +366,8 @@ class AnthropicMessagesHandler(BaseTranslation):
self,
data: dict,
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: Any | None = None,
) -> Any:
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
) -> Mapping[str, object]:
"""
Process input messages by applying guardrails to text content.
"""
@ -467,8 +518,8 @@ class AnthropicMessagesHandler(BaseTranslation):
@staticmethod
def _openai_system_message_to_anthropic(
message: dict[str, Any],
) -> dict[str, Any] | None: # mutable-ok: API message payload
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")
if isinstance(content, str):
@ -477,14 +528,14 @@ class AnthropicMessagesHandler(BaseTranslation):
) # mutable-ok: API message payload
if not isinstance(content, list):
return None
blocks: Final[list[dict[str, Any]]] = [] # mutable-ok: API message payload
blocks: Final[list[dict[str, object]]] = [] # mutable-ok: API message payload
for block in content:
if not isinstance(block, dict) or block.get("type") != "text":
continue
text = block.get("text")
if not isinstance(text, str) or not text:
continue
anthropic_block: dict[str, Any] = { # mutable-ok: API message payload
anthropic_block: dict[str, object] = { # mutable-ok: API message payload
"type": "text",
"text": text,
} # mutable-ok: API message payload
@ -514,7 +565,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
@ -602,7 +653,7 @@ class AnthropicMessagesHandler(BaseTranslation):
@staticmethod
def _extract_midturn_system_text(
message: dict[str, Any], # mutable-ok: API message payload
message: Mapping[str, object],
msg_idx: int,
) -> ExtractedInput:
"""Match the adapter's filtering so positional guardrail write-back stays aligned."""
@ -636,7 +687,7 @@ class AnthropicMessagesHandler(BaseTranslation):
@classmethod
def _extract_input_text_and_images(
cls,
message: dict[str, Any],
message: Mapping[str, object],
msg_idx: int,
skip_system_message: bool = False,
skip_tool_message: bool = False,
@ -696,7 +747,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)),)
@ -707,7 +758,7 @@ class AnthropicMessagesHandler(BaseTranslation):
@classmethod
def _extract_tool_result(
cls,
content_item: Mapping[str, Any],
content_item: Mapping[str, object],
msg_idx: int,
content_idx: int,
) -> ExtractedInput:
@ -736,7 +787,7 @@ class AnthropicMessagesHandler(BaseTranslation):
)
@staticmethod
def _image_sources(block: Mapping[str, Any]) -> tuple[str, ...]:
def _image_sources(block: Mapping[str, object]) -> tuple[str, ...]:
source: Final = block.get("source")
if not isinstance(source, Mapping):
return ()
@ -746,7 +797,7 @@ class AnthropicMessagesHandler(BaseTranslation):
async def _apply_guardrail_responses_to_input(
self,
messages: list[dict[str, Any]],
messages: Sequence[_WritableMessage],
responses: list[str],
scanned: tuple[ScannedText, ...],
) -> None:
@ -788,10 +839,10 @@ class AnthropicMessagesHandler(BaseTranslation):
self,
response: "AnthropicMessagesResponse",
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: Any | None = None,
user_api_key_dict: Any | None = None,
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
user_api_key_dict: "UserAPIKeyAuth | None" = None,
request_data: dict | None = None,
) -> Any:
) -> "AnthropicMessagesResponse":
"""
Process output response by applying guardrails to text content and tool calls.
@ -869,10 +920,10 @@ class AnthropicMessagesHandler(BaseTranslation):
self,
responses_so_far: list[Any],
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: Any | None = None,
user_api_key_dict: Any | None = None,
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.
@ -950,8 +1001,8 @@ class AnthropicMessagesHandler(BaseTranslation):
def _prepare_request_data(
self,
request_data: dict | None,
response: Any,
user_api_key_dict: Any | None,
response: object,
user_api_key_dict: "UserAPIKeyAuth | None",
key: str,
) -> dict:
"""Ensure request_data has the response/responses_so_far key and metadata."""
@ -968,7 +1019,7 @@ class AnthropicMessagesHandler(BaseTranslation):
return request_data
@staticmethod
def _get_response_content(response: Any) -> 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 []
@ -978,7 +1029,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]],
@ -986,21 +1037,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, Any] = {}
if isinstance(content_block, dict):
block_type = content_block.get("type")
block_dict = cast(dict[str, Any], 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(
@ -1012,12 +1052,27 @@ 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],
images_to_check: list[str],
tool_calls_to_check: list["ChatCompletionToolCallChunk"],
response: Any,
response: object,
) -> "GenericGuardrailAPIInputs":
"""Build GenericGuardrailAPIInputs with optional images, tool calls, model."""
inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check)
@ -1034,7 +1089,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.
@ -1105,7 +1160,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", "")
@ -1117,7 +1172,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.
@ -1168,7 +1223,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:
@ -1212,7 +1267,7 @@ class AnthropicMessagesHandler(BaseTranslation):
def _extract_output_text_and_images(
self,
content_block: dict[str, Any],
content_block: Mapping[str, object],
content_idx: int,
texts_to_check: list[str],
images_to_check: list[str],
@ -1235,7 +1290,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,
@ -1260,7 +1315,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"):
@ -1276,14 +1331,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, Any], 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

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

@ -13,8 +13,10 @@ Mirrors Anthropic's native ``compact_20260112`` for non-Anthropic providers:
"""
import re
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, 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 NotRequired, TypedDict, Unpack
import litellm
from litellm._logging import verbose_logger
@ -27,11 +29,11 @@ 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,
AllAnthropicToolsValues,
AnthopicMessagesAssistantMessageParam,
AnthropicMessagesUserMessageParam,
)
from litellm.types.llms.openai import ChatCompletionToolParam
from litellm.types.utils import ModelResponse
@ -82,6 +84,69 @@ _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: str
max_tokens: int
timeout: float
litellm_metadata: Mapping[str, object]
user: NotRequired[str]
allowed_model_region: NotRequired[str]
class _SummaryAcompletion(Protocol):
def __call__(
self, *, messages: Sequence[Mapping[str, object]], **kwargs: Unpack[_SummaryCallKwargs]
) -> "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."""
@ -157,11 +222,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),
@ -347,7 +412,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(
@ -399,40 +464,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:
@ -446,7 +528,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.
@ -465,8 +547,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)``
@ -481,19 +563,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.
@ -600,7 +684,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,
@ -623,7 +707,7 @@ def _count_effective_tokens(
try:
openai_shape = adapter.translate_anthropic_messages_to_openai(
messages=cast(
"list[AnthropicMessagesUserMessageParam | AnthopicMessagesAssistantMessageParam]",
"list[AllAnthropicPassThroughMessageValues]",
messages_without_compaction,
)
)
@ -679,17 +763,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.
@ -704,16 +789,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 [
@ -736,7 +823,7 @@ def _extract_summary_text(raw: str | None) -> str | None:
def _system_to_openai_message(
system: str | list[dict[str, Any]] | None,
) -> dict[str, Any] | 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
@ -747,17 +834,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
@ -773,7 +862,7 @@ def _build_summary_messages(
try:
openai_messages = LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai(
messages=cast(
"list[AnthropicMessagesUserMessageParam | AnthopicMessagesAssistantMessageParam]",
"list[AllAnthropicPassThroughMessageValues]",
stripped,
)
)
@ -785,7 +874,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)
@ -809,7 +898,7 @@ def _is_user_message(msg: object) -> bool:
return isinstance(msg, dict) and msg.get("role") == "user"
def _append_text_to_content(content: Any, extra_text: str) -> Any:
def _append_text_to_content(content: object, extra_text: str) -> object:
"""Append ``extra_text`` to an OpenAI-shape message ``content`` field.
Handles the two common shapes: ``str`` and ``list`` of content parts.
@ -820,16 +909,17 @@ def _append_text_to_content(content: Any, extra_text: str) -> Any:
if isinstance(content, str):
return f"{content}\n\n{extra_text}"
if isinstance(content, list):
return [*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}]
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"]:
@ -860,9 +950,8 @@ async def _call_summary_model(
# the parent ``/v1/messages`` request. On timeout the caller catches the
# exception and surfaces ``applied_edits[0].error = "summary_call_failed"``,
# forwarding the request without compaction rather than hanging.
call_kwargs: Final[dict[str, Any]] = {
call_kwargs: Final[_SummaryCallKwargs] = {
"model": summary_model,
"messages": summary_messages,
"max_tokens": max_tokens,
"timeout": COMPACT_SUMMARY_TIMEOUT_SECONDS,
"litellm_metadata": metadata,
@ -872,19 +961,23 @@ 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")
if end_user_id:
if isinstance(end_user_id, str) and end_user_id:
call_kwargs["user"] = end_user_id
if allowed_model_region is not None:
call_kwargs["allowed_model_region"] = allowed_model_region
if llm_router is not None and hasattr(llm_router, "acompletion"):
return await llm_router.acompletion(**call_kwargs)
return await litellm.acompletion(**call_kwargs)
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
@ -900,7 +993,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 (

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

@ -1,8 +1,10 @@
from collections.abc import Mapping, Sequence
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final
from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict
import httpx
from httpx._types import RequestFiles
from httpx._types import FileTypes, RequestFiles
from typing_extensions import NotRequired
import litellm
from litellm.constants import RUNWAYML_DEFAULT_API_VERSION
@ -31,6 +33,31 @@ else:
LiteLLMLoggingObj = Any
class _RunwayTaskResponse(TypedDict, total=False):
id: str
status: str
createdAt: str
completedAt: str
output: Sequence[str] | str
progress: int
failureCode: str
failure: str
class _RunwayVideoData(TypedDict):
id: str
object: Literal["video"]
status: str
created_at: int
output_url: NotRequired[str]
completed_at: NotRequired[int]
progress: NotRequired[int]
error: NotRequired[Mapping[str, str]]
model: NotRequired[str]
size: NotRequired[str]
seconds: NotRequired[str]
class RunwayMLVideoConfig(BaseVideoConfig):
"""
Configuration class for RunwayML video generation.
@ -44,6 +71,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.
@ -68,7 +99,7 @@ class RunwayMLVideoConfig(BaseVideoConfig):
video_create_optional_params: VideoCreateOptionalRequestParams,
model: str,
drop_params: bool,
) -> dict:
) -> dict[str, object]:
"""
Map OpenAI parameters to RunwayML format.
@ -78,37 +109,44 @@ class RunwayMLVideoConfig(BaseVideoConfig):
- size -> ratio (convert "WIDTHxHEIGHT" to "WIDTH:HEIGHT")
- seconds -> duration (convert to integer)
"""
mapped_params: Final[dict[str, Any]] = {}
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,
@ -163,7 +201,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]:
@ -179,17 +217,15 @@ class RunwayMLVideoConfig(BaseVideoConfig):
"duration": 5
}
"""
# Build the request data
request_data: Final[dict[str, Any]] = {
# Build the request data with the mapped parameters merged in
request_data: Final = {
"model": model,
"promptText": prompt,
**video_create_optional_request_params,
}
# Add mapped parameters
request_data.update(video_create_optional_request_params)
# RunwayML uses JSON body, no files multipart
files_list: Final[list[tuple[str, Any]]] = []
files_list: Final[Sequence[tuple[str, FileTypes]]] = []
# Append the specific endpoint for video generation
full_api_base: Final = f"{api_base}/image_to_video"
@ -216,10 +252,10 @@ class RunwayMLVideoConfig(BaseVideoConfig):
We map this to OpenAI VideoObject format.
"""
response_data: Final = raw_response.json()
response_data: Final[_RunwayTaskResponse] = self._parse_task_response(raw_response)
# Map RunwayML task response to VideoObject format
video_data: Final[dict[str, Any]] = {
video_data: Final[_RunwayVideoData] = {
"id": response_data.get("id", ""),
"object": "video",
"status": self._map_runway_status(response_data.get("status", "pending")),
@ -229,9 +265,8 @@ class RunwayMLVideoConfig(BaseVideoConfig):
# Add optional fields if present
if "output" in response_data and response_data["output"]:
# RunwayML returns output as array of URLs when task succeeds
video_data["output_url"] = (
response_data["output"][0] if isinstance(response_data["output"], list) else response_data["output"]
)
output: Final = response_data["output"]
video_data["output_url"] = output if isinstance(output, str) else output[0]
if "completedAt" in response_data:
video_data["completed_at"] = self._parse_runway_timestamp(response_data.get("completedAt"))
@ -254,7 +289,7 @@ class RunwayMLVideoConfig(BaseVideoConfig):
if "duration" in request_data:
video_data["seconds"] = str(request_data["duration"])
video_obj: Final = VideoObject(**video_data)
video_obj: Final = VideoObject.model_validate(video_data)
if custom_llm_provider and video_obj.id:
video_obj.id = encode_video_id_with_provider(video_obj.id, custom_llm_provider, model)
@ -326,20 +361,18 @@ class RunwayMLVideoConfig(BaseVideoConfig):
# Get task status to retrieve video URL
url: Final = f"{api_base}/tasks/{encoded_video_id}"
params: Final[dict[str, Any]] = {}
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
@ -373,7 +406,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
@ -402,7 +435,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
@ -421,7 +454,7 @@ class RunwayMLVideoConfig(BaseVideoConfig):
api_base: str,
litellm_params: GenericLiteLLMParams,
headers: dict,
extra_body: dict[str, Any] | None = None,
extra_body: Mapping[str, object] | None = None,
) -> tuple[str, dict]:
"""
Transform the video remix request for RunwayML API.
@ -448,7 +481,7 @@ class RunwayMLVideoConfig(BaseVideoConfig):
after: str | None = None,
limit: int | None = None,
order: str | None = None,
extra_query: dict[str, Any] | None = None,
extra_query: Mapping[str, object] | None = None,
) -> tuple[str, dict]:
"""
Transform the video list request for RunwayML API.
@ -484,9 +517,7 @@ class RunwayMLVideoConfig(BaseVideoConfig):
# Construct the URL for task cancellation
url: Final = f"{api_base}/tasks/{encoded_video_id}/cancel"
data: Final[dict[str, Any]] = {}
return url, data
return url, dict[str, str]()
def transform_video_delete_response(
self,
@ -494,7 +525,7 @@ class RunwayMLVideoConfig(BaseVideoConfig):
logging_obj: LiteLLMLoggingObj,
) -> VideoObject:
"""Transform the RunwayML video delete/cancel response."""
response_data: Final = raw_response.json()
response_data: Final[_RunwayTaskResponse] = self._parse_task_response(raw_response)
video_obj: Final = VideoObject(
id=response_data.get("id", ""),
@ -524,9 +555,7 @@ class RunwayMLVideoConfig(BaseVideoConfig):
url: Final = f"{api_base}/tasks/{encoded_video_id}"
# Empty dict for GET request (no body)
data: Final[dict[str, Any]] = {}
return url, data
return url, dict[str, str]()
def transform_video_status_retrieve_response(
self,
@ -537,10 +566,10 @@ class RunwayMLVideoConfig(BaseVideoConfig):
"""
Transform the RunwayML video status retrieve response.
"""
response_data: Final = raw_response.json()
response_data: Final[_RunwayTaskResponse] = self._parse_task_response(raw_response)
# Map RunwayML task response to VideoObject format
video_data: Final[dict[str, Any]] = {
video_data: Final[_RunwayVideoData] = {
"id": response_data.get("id", ""),
"object": "video",
"status": self._map_runway_status(response_data.get("status", "pending")),
@ -549,9 +578,8 @@ class RunwayMLVideoConfig(BaseVideoConfig):
# Add optional fields if present
if "output" in response_data and response_data["output"]:
video_data["output_url"] = (
response_data["output"][0] if isinstance(response_data["output"], list) else response_data["output"]
)
output: Final = response_data["output"]
video_data["output_url"] = output if isinstance(output, str) else output[0]
if "completedAt" in response_data:
video_data["completed_at"] = self._parse_runway_timestamp(response_data.get("completedAt"))
@ -565,14 +593,14 @@ class RunwayMLVideoConfig(BaseVideoConfig):
"message": response_data.get("failure", "Video generation failed"),
}
video_obj: Final = VideoObject(**video_data)
video_obj: Final = VideoObject.model_validate(video_data)
if custom_llm_provider and video_obj.id:
video_obj.id = encode_video_id_with_provider(video_obj.id, custom_llm_provider, None)
return video_obj
def transform_video_create_character_request(self, name, video, api_base, litellm_params, headers):
def transform_video_create_character_request(self, name, video: object, api_base, litellm_params, headers):
raise NotImplementedError("video create character is not supported for RunwayML")
def transform_video_create_character_response(self, raw_response, logging_obj):

114
litellm/llms/vertex_ai/files/transformation.py
View file

@ -5,12 +5,13 @@ import json
import os
import re
import time
from collections.abc import Callable, Iterable, Iterator
from typing import Any, Final
from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
from typing import Final, TypedDict
import httpx
from httpx import Headers, Response
from openai.types.file_deleted import FileDeleted
from typing_extensions import Required
import litellm
from litellm._uuid import uuid
@ -50,9 +51,10 @@ from litellm.types.llms.openai import (
HttpxBinaryResponseContent,
OpenAICreateFileRequestOptionalParams,
OpenAIFileObject,
OpenAIFilesPurpose,
PathLike,
)
from litellm.types.llms.vertex_ai import GcsBucketResponse
from litellm.types.llms.vertex_ai import GcsBucketResponse, GenerateContentResponseBody
from litellm.types.utils import LlmProviders, ModelResponse
from ..common_utils import VertexAIError
@ -62,6 +64,47 @@ _GCP_LABEL_VALUE_MAX_LEN: Final = 63
_CUSTOM_ID_RAW_LABEL_PREFIX: Final = "b32_"
class _OpenAIBatchRequestBody(TypedDict, total=False):
model: str
messages: Sequence[AllMessageValues]
class _OpenAIBatchJsonlEntry(TypedDict, total=False):
custom_id: Required[object]
body: _OpenAIBatchRequestBody
class _VertexBatchOutputRequest(TypedDict, total=False):
labels: Mapping[str, str]
class _VertexBatchResponse(GenerateContentResponseBody, total=False):
modelVersion: str
class _VertexBatchOutputRow(TypedDict, total=False):
request: _VertexBatchOutputRequest
status: str
processed_time: str
response: _VertexBatchResponse
class _GcsObjectMetadata(TypedDict, total=False):
purpose: OpenAIFilesPurpose
class _GcsObjectResponse(GcsBucketResponse, total=False):
metadata: _GcsObjectMetadata
def _parse_gcs_object_response(raw_response: Response) -> _GcsObjectResponse:
return raw_response.json()
def _parse_vertex_batch_output_row(line: str) -> _VertexBatchOutputRow:
return json.loads(line)
def _sanitize_gcp_label_value(value: str) -> str:
"""
Sanitize a string to meet GCP label value constraints.
@ -106,7 +149,7 @@ def _decode_gcp_label_value_chunks(values: list[str]) -> str | None:
return None
def _set_litellm_batch_custom_id_labels(labels: dict[str, str], custom_id: Any) -> None:
def _litellm_batch_custom_id_labels(custom_id: object) -> Mapping[str, str]:
"""
Store OpenAI batch custom_id for Vertex batch correlation.
@ -115,15 +158,19 @@ def _set_litellm_batch_custom_id_labels(labels: dict[str, str], custom_id: Any)
round-trip correlation in batch output transforms.
"""
custom_id_str: Final = str(custom_id)
labels["litellm_custom_id"] = _sanitize_gcp_label_value(custom_id_str)
raw_label_chunks: Final = _encode_gcp_label_value_chunks(custom_id_str)
labels["litellm_custom_id_raw"] = raw_label_chunks[0]
for index, raw_label_chunk in enumerate(raw_label_chunks[1:], start=1):
labels[f"litellm_custom_id_raw_{index}"] = raw_label_chunk
return {
"litellm_custom_id": _sanitize_gcp_label_value(custom_id_str),
"litellm_custom_id_raw": raw_label_chunks[0],
**{
f"litellm_custom_id_raw_{index}": raw_label_chunk
for index, raw_label_chunk in enumerate(raw_label_chunks[1:], start=1)
},
}
def _get_litellm_batch_custom_id_from_labels(labels: dict[str, Any]) -> str:
"""Prefer encoded custom_id when present (see _set_litellm_batch_custom_id_labels)."""
def _get_litellm_batch_custom_id_from_labels(labels: Mapping[str, str]) -> str:
"""Prefer encoded custom_id when present (see _litellm_batch_custom_id_labels)."""
raw: Final = labels.get("litellm_custom_id_raw")
if raw:
raw_chunks: Final = [str(raw)]
@ -141,9 +188,9 @@ def _get_litellm_batch_custom_id_from_labels(labels: dict[str, Any]) -> str:
def _openai_batch_jsonl_entry_to_vertex_wrapped_request(
openai_entry: dict[str, Any],
map_openai_to_vertex_params: Callable[[dict[str, Any]], dict[str, Any]],
) -> dict[str, Any]:
openai_entry: _OpenAIBatchJsonlEntry,
map_openai_to_vertex_params: Callable[[_OpenAIBatchRequestBody], Mapping[str, object]],
) -> Mapping[str, object]:
"""
Transforms a single OpenAI JSONL batch entry into its Vertex wrapped request.
@ -151,11 +198,11 @@ def _openai_batch_jsonl_entry_to_vertex_wrapped_request(
Example Vertex jsonl
{"request":{"contents": [{"role": "user", "parts": [{"text": "What is the relation between the following video and image samples?"}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/video/animals.mp4", "mimeType": "video/mp4"}}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/image/cricket.jpeg", "mimeType": "image/jpeg"}}]}]}}
"""
openai_request_body: Final = openai_entry.get("body") or {}
openai_request_body: Final[_OpenAIBatchRequestBody] = openai_entry.get("body") or {}
vertex_request_body: Final = _transform_request_body(
messages=openai_request_body.get("messages", []),
messages=[*openai_request_body.get("messages", [])],
model=openai_request_body.get("model", ""),
optional_params=map_openai_to_vertex_params(openai_request_body),
optional_params=dict(map_openai_to_vertex_params(openai_request_body)),
custom_llm_provider="vertex_ai",
litellm_params={},
cached_content=None,
@ -163,9 +210,10 @@ def _openai_batch_jsonl_entry_to_vertex_wrapped_request(
custom_id: Final = openai_entry.get("custom_id")
if custom_id is not None:
if "labels" not in vertex_request_body:
vertex_request_body["labels"] = {}
_set_litellm_batch_custom_id_labels(vertex_request_body["labels"], custom_id)
vertex_request_body["labels"] = {
**vertex_request_body.get("labels", {}),
**_litellm_batch_custom_id_labels(custom_id),
}
return {"request": vertex_request_body}
@ -186,7 +234,7 @@ def _iter_openai_jsonl_lines(openai_file_content: FileTypes) -> Iterator[str]:
``str.splitlines()`` + ``line.strip()`` for ``\\n`` / ``\\r\\n`` delimited
JSONL.
"""
content: Any = openai_file_content
content: FileTypes | str = openai_file_content
if isinstance(content, tuple):
content = content[1]
@ -241,7 +289,7 @@ def _iter_openai_jsonl_lines(openai_file_content: FileTypes) -> Iterator[str]:
def _iter_openai_jsonl_entries(
openai_file_content: FileTypes,
) -> Iterator[dict[str, Any]]:
) -> Iterator[_OpenAIBatchJsonlEntry]:
for line in _iter_openai_jsonl_lines(openai_file_content):
yield json.loads(line)
@ -257,7 +305,7 @@ class _OpenAIToVertexBatchUploadStream(BaseFileUploadStream):
def __init__(
self,
openai_file_content: FileTypes,
map_openai_to_vertex_params: Callable[[dict[str, Any]], dict[str, Any]],
map_openai_to_vertex_params: Callable[[_OpenAIBatchRequestBody], Mapping[str, object]],
) -> None:
self._openai_file_content = openai_file_content
self._map_openai_to_vertex_params = map_openai_to_vertex_params
@ -308,7 +356,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
def _get_gcs_object_name_from_batch_jsonl(
self,
openai_jsonl_content: list[dict[str, Any]],
openai_jsonl_content: Sequence[_OpenAIBatchJsonlEntry],
) -> str:
"""
Gets a unique GCS object name for the VertexAI batch prediction job
@ -396,8 +444,8 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
def _map_openai_to_vertex_params(
self,
openai_request_body: dict[str, Any],
) -> dict[str, Any]:
openai_request_body: _OpenAIBatchRequestBody,
) -> Mapping[str, object]:
"""
wrapper to call VertexGeminiConfig.map_openai_params
"""
@ -409,7 +457,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
_model: Final = openai_request_body.get("model", "")
vertex_params: Final = config.map_openai_params(
model=_model,
non_default_params=openai_request_body,
non_default_params=dict(openai_request_body),
optional_params={},
drop_params=False,
)
@ -463,10 +511,10 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
"""
Transform VertexAI File upload response into OpenAI-style FileObject
"""
response_json: Final = raw_response.json()
response_json: Final = _parse_gcs_object_response(raw_response)
try:
response_object: Final = GcsBucketResponse(**response_json)
response_object: Final = _GcsObjectResponse(**response_json)
except Exception as e:
raise VertexAIError(
status_code=raw_response.status_code,
@ -523,7 +571,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
) -> OpenAIFileObject:
response_json: Final = raw_response.json()
response_json: Final = _parse_gcs_object_response(raw_response)
gcs_id = response_json.get("id", "")
gcs_id = "/".join(gcs_id.split("/")[:-1]) if gcs_id else ""
return OpenAIFileObject(
@ -682,7 +730,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
# discriminating fields. Anything else (e.g. a binary file whose
# first line is not valid UTF-8/JSON) raises and falls through to the
# passthrough below, leaving the content untouched.
first_row: Final = json.loads(first_line)
first_row: Final = _parse_vertex_batch_output_row(first_line)
is_vertex_batch_output: Final = (
"request" in first_row
and "response" in first_row
@ -723,7 +771,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
for line in itertools.chain([first_line], lines):
try:
openai_output = self._transform_single_vertex_batch_output_to_openai(
vertex_output=json.loads(line),
vertex_output=_parse_vertex_batch_output_row(line),
vertex_gemini_config=vertex_gemini_config,
logging_obj=batch_transform_logging_obj,
mock_httpx_response=mock_httpx_response,
@ -742,11 +790,11 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
def _transform_single_vertex_batch_output_to_openai(
self,
vertex_output: dict[str, Any],
vertex_output: _VertexBatchOutputRow,
vertex_gemini_config: VertexGeminiConfig,
logging_obj: Logging,
mock_httpx_response: httpx.Response,
) -> dict[str, Any]:
) -> Mapping[str, object]:
"""
Transform a single Vertex AI batch output line to OpenAI format.
Uses the existing VertexGeminiConfig transformation for the response.

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

@ -6,8 +6,11 @@ Admins use the management endpoints to read and update input_policy / output_pol
"""
import uuid
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
@ -23,33 +26,109 @@ if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
def _row_to_model(row: dict | Any) -> LiteLLM_ToolTableRow:
class _ToolTableRecord(Protocol):
tool_name: str
input_policy: str | None
output_policy: str | None
class _TokenRelationRecord(Protocol):
token: str | None
key_alias: str | None
class _TeamRelationRecord(Protocol):
team_id: str | None
team_alias: str | None
class _ObjectPermissionRecord(Protocol):
object_permission_id: str
blocked_tools: Sequence[str] | None
verification_tokens: Sequence[_TokenRelationRecord] | None
teams: Sequence[_TeamRelationRecord] | None
class _ToolTable(Protocol):
async def find_many(
self,
*,
where: Mapping[str, object] | None = None,
order: Mapping[str, str] | None = None,
) -> Sequence[_ToolTableRecord]: ...
async def find_unique(self, *, where: Mapping[str, object]) -> _ToolTableRecord | None: ...
async def upsert(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> object: ...
class _ObjectPermissionTable(Protocol):
async def find_many(
self,
*,
where: Mapping[str, object] | None = None,
include: Mapping[str, bool] | None = None,
) -> Sequence[_ObjectPermissionRecord]: ...
async def find_unique(self, *, where: Mapping[str, object]) -> _ObjectPermissionRecord | None: ...
async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> object: ...
class _ModelDumpMethod(Protocol):
def __call__(self) -> Mapping: ...
_ROW_DICT: Final = TypeAdapter(dict)
class _ToolTableHolder(Protocol):
@property
def table(self) -> _ToolTable: ...
class _ObjectPermissionTableHolder(Protocol):
@property
def table(self) -> _ObjectPermissionTable: ...
def _tool_table(repo: _ToolTableHolder) -> _ToolTable:
return repo.table
def _object_permission_table(repo: _ObjectPermissionTableHolder) -> _ObjectPermissionTable:
return repo.table
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", ""),
@ -87,7 +166,7 @@ async def batch_upsert_tools(
if not data:
return
now: Final = datetime.now(timezone.utc)
table: Final = ToolRepository(prisma_client).table
table: Final = _tool_table(ToolRepository(prisma_client))
for item in data:
tool_name = item.get("tool_name", "")
origin = item.get("origin") or "user_defined"
@ -132,8 +211,8 @@ async def list_tools(
) -> list[LiteLLM_ToolTableRow]:
"""Return all tools, optionally filtered by input_policy."""
try:
where: Final = {"input_policy": input_policy} if input_policy is not None else {}
rows: Final = await ToolRepository(prisma_client).table.find_many(
where: Final[Mapping[str, str]] = {"input_policy": input_policy} if input_policy is not None else {}
rows: Final = await _tool_table(ToolRepository(prisma_client)).find_many(
where=where,
order={"created_at": "desc"},
)
@ -149,7 +228,7 @@ async def get_tool(
) -> LiteLLM_ToolTableRow | None:
"""Return a single tool row by tool_name."""
try:
row: Final = await ToolRepository(prisma_client).table.find_unique(
row: Final = await _tool_table(ToolRepository(prisma_client)).find_unique(
where={"tool_name": tool_name},
)
if row is None:
@ -172,7 +251,7 @@ async def update_tool_policy(
_updated_by: Final = updated_by or "system"
now: Final = datetime.now(timezone.utc)
create_data: Final[dict] = {
create_data: Final[Mapping[str, str | datetime]] = {
"tool_id": str(uuid.uuid4()),
"tool_name": tool_name,
"input_policy": input_policy or "untrusted",
@ -182,16 +261,18 @@ async def update_tool_policy(
"created_at": now,
"updated_at": now,
}
update_data: Final[dict] = {
"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 ToolRepository(prisma_client).table.upsert(
await _tool_table(ToolRepository(prisma_client)).upsert(
where={"tool_name": tool_name},
data={
"create": create_data,
@ -214,7 +295,7 @@ async def get_tools_by_names(
if not tool_names:
return {}
try:
rows: Final = await ToolRepository(prisma_client).table.find_many(
rows: Final = await _tool_table(ToolRepository(prisma_client)).find_many(
where={"tool_name": {"in": tool_names}},
)
return {
@ -239,7 +320,7 @@ async def list_overrides_for_tool(
"""
out: Final[list[ToolPolicyOverrideRow]] = []
try:
perms: Final = await ObjectPermissionRepository(prisma_client).table.find_many(
perms: Final = await _object_permission_table(ObjectPermissionRepository(prisma_client)).find_many(
where={"blocked_tools": {"has": tool_name}},
include={
"verification_tokens": True,
@ -248,8 +329,8 @@ async def list_overrides_for_tool(
)
for perm in perms:
op_id = getattr(perm, "object_permission_id", None) or ""
tokens = getattr(perm, "verification_tokens", []) or []
teams = getattr(perm, "teams", []) or []
tokens: Sequence[_TokenRelationRecord] = getattr(perm, "verification_tokens", []) or []
teams: Sequence[_TeamRelationRecord] = getattr(perm, "teams", []) or []
for t in tokens:
out.append(
ToolPolicyOverrideRow(
@ -302,7 +383,7 @@ class ToolPolicyRegistry:
try:
tools: Final = await call_with_db_reconnect_retry(
prisma_client,
lambda: ToolRepository(prisma_client).table.find_many(),
lambda: _tool_table(ToolRepository(prisma_client)).find_many(),
reason="sync_tool_policy_from_db_tools_lookup_failure",
)
self._tool_input_policies = {
@ -314,13 +395,13 @@ class ToolPolicyRegistry:
perms: Final = await call_with_db_reconnect_retry(
prisma_client,
lambda: ObjectPermissionRepository(prisma_client).table.find_many(),
lambda: _object_permission_table(ObjectPermissionRepository(prisma_client)).find_many(),
reason="sync_tool_policy_from_db_perms_lookup_failure",
)
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)
@ -352,10 +433,12 @@ class ToolPolicyRegistry:
"""
if not tool_names:
return {}
blocked: Final[set] = 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:
@ -385,18 +468,17 @@ async def add_tool_to_object_permission_blocked(
if not object_permission_id or not tool_name:
return False
try:
row: Final = await ObjectPermissionRepository(prisma_client).table.find_unique(
row: Final = await _object_permission_table(ObjectPermissionRepository(prisma_client)).find_unique(
where={"object_permission_id": object_permission_id},
)
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 ObjectPermissionRepository(prisma_client).table.update(
await _object_permission_table(ObjectPermissionRepository(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:
@ -413,18 +495,17 @@ async def remove_tool_from_object_permission_blocked(
if not object_permission_id or not tool_name:
return False
try:
row: Final = await ObjectPermissionRepository(prisma_client).table.find_unique(
row: Final = await _object_permission_table(ObjectPermissionRepository(prisma_client)).find_unique(
where={"object_permission_id": object_permission_id},
)
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 ObjectPermissionRepository(prisma_client).table.update(
await _object_permission_table(ObjectPermissionRepository(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:

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

@ -17,7 +17,7 @@ import json
import traceback
from collections.abc import Mapping, Sequence
from datetime import datetime, timezone
from typing import Any, Final, Literal, cast
from typing import Any, Final, Literal, Protocol, TypeAlias, TypeVar, cast, overload
import fastapi
from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
@ -85,14 +85,6 @@ from litellm.types.proxy.management_endpoints.scim_v2 import (
if TYPE_CHECKING:
from prisma import models as prisma_models
from prisma import types as prisma_types
from prisma.actions import (
LiteLLM_InvitationLinkActions,
LiteLLM_OrganizationMembershipActions,
LiteLLM_TeamMembershipActions,
LiteLLM_TeamTableActions,
LiteLLM_UserTableActions,
LiteLLM_VerificationTokenActions,
)
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.proxy_server import PrismaClient
@ -100,55 +92,151 @@ if TYPE_CHECKING:
router: Final = APIRouter()
_PrismaTableT = TypeVar("_PrismaTableT", covariant=True)
class _TableActions(Protocol[_PrismaTableT]):
async def find_unique(
self,
*,
where: Mapping[str, object],
include: Mapping[str, object] | None = None,
) -> "_PrismaTableT | None": ...
async def find_first(self, *, where: Mapping[str, object]) -> "_PrismaTableT | None": ...
async def find_many(
self,
*,
where: Mapping[str, object] | None = None,
order: Mapping[str, object] | Sequence[Mapping[str, object]] | None = None,
skip: int | None = None,
take: int | None = None,
) -> "Sequence[_PrismaTableT]": ...
async def create(self, *, data: Mapping[str, object]) -> "_PrismaTableT": ...
async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> "_PrismaTableT | None": ...
async def delete_many(self, *, where: Mapping[str, object] | None = None) -> int: ...
class _PrismaTableHolder(Protocol[_PrismaTableT]):
@property
def table(self) -> "_TableActions[_PrismaTableT]": ...
def _typed_table(holder: "_PrismaTableHolder[_PrismaTableT]") -> "_TableActions[_PrismaTableT]":
return holder.table
class _LenientTableActions(Protocol[_PrismaTableT]):
async def find_first(self, *, where: Mapping[str, object]) -> "_PrismaTableT | None": ...
async def find_many(
self,
*,
where: Mapping[str, object] | None = None,
order: Mapping[str, object] | Sequence[Mapping[str, object]] | None = None,
skip: int | None = None,
take: int | None = None,
) -> "Sequence[_PrismaTableT] | None": ...
async def count(self, *, where: Mapping[str, object] | None = None) -> int: ...
async def update_many(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> int: ...
class _LenientTableHolder(Protocol[_PrismaTableT]):
@property
def table(self) -> "_LenientTableActions[_PrismaTableT]": ...
def _lenient_table(holder: "_LenientTableHolder[_PrismaTableT]") -> "_LenientTableActions[_PrismaTableT]":
return holder.table
class _UserDeleteRow(Protocol):
user_id: str
user_email: str | None
@property
def teams(self) -> Sequence[str]: ...
def json(self, *, exclude_none: bool) -> str: ...
class _TeamCleanupRow(Protocol):
team_id: str
members_with_roles: str
def model_dump(self) -> Mapping[str, object]: ...
def _user_table(
prisma_client: "PrismaClient | None",
) -> "LiteLLM_UserTableActions[prisma_models.LiteLLM_UserTable]":
user_table: Final[LiteLLM_UserTableActions[prisma_models.LiteLLM_UserTable]] = UserRepository(prisma_client).table
return user_table
) -> "_TableActions[prisma_models.LiteLLM_UserTable]":
return _typed_table(UserRepository(prisma_client))
def _user_table_lenient(
prisma_client: "PrismaClient | None",
) -> "_LenientTableActions[prisma_models.LiteLLM_UserTable]":
return _lenient_table(UserRepository(prisma_client))
def _user_delete_table(
prisma_client: "PrismaClient | None",
) -> "_TableActions[_UserDeleteRow]":
return _typed_table(UserRepository(prisma_client))
def _team_table(
prisma_client: "PrismaClient | None",
) -> "LiteLLM_TeamTableActions[prisma_models.LiteLLM_TeamTable]":
team_table: Final[LiteLLM_TeamTableActions[prisma_models.LiteLLM_TeamTable]] = TeamRepository(prisma_client).table
return team_table
) -> "_TableActions[prisma_models.LiteLLM_TeamTable]":
return _typed_table(TeamRepository(prisma_client))
def _team_cleanup_table(
prisma_client: "PrismaClient | None",
) -> "_TableActions[_TeamCleanupRow]":
return _typed_table(TeamRepository(prisma_client))
def _verification_token_table(
prisma_client: "PrismaClient | None",
) -> "LiteLLM_VerificationTokenActions[prisma_models.LiteLLM_VerificationToken]":
token_table: Final[LiteLLM_VerificationTokenActions[prisma_models.LiteLLM_VerificationToken]] = (
VerificationTokenRepository(prisma_client).table
)
return token_table
) -> "_TableActions[prisma_models.LiteLLM_VerificationToken]":
return _typed_table(VerificationTokenRepository(prisma_client))
def _verification_token_table_lenient(
prisma_client: "PrismaClient | None",
) -> "_LenientTableActions[prisma_models.LiteLLM_VerificationToken]":
return _lenient_table(VerificationTokenRepository(prisma_client))
def _organization_table(
prisma_client: "PrismaClient | None",
) -> "_TableActions[prisma_models.LiteLLM_OrganizationTable]":
return _typed_table(OrganizationRepository(prisma_client))
def _organization_membership_table(
prisma_client: "PrismaClient | None",
) -> "LiteLLM_OrganizationMembershipActions[prisma_models.LiteLLM_OrganizationMembership]":
membership_table: Final[LiteLLM_OrganizationMembershipActions[prisma_models.LiteLLM_OrganizationMembership]] = (
OrganizationMembershipRepository(prisma_client).table
)
return membership_table
) -> "_TableActions[prisma_models.LiteLLM_OrganizationMembership]":
return _typed_table(OrganizationMembershipRepository(prisma_client))
def _invitation_link_table(
prisma_client: "PrismaClient | None",
) -> "LiteLLM_InvitationLinkActions[prisma_models.LiteLLM_InvitationLink]":
invitation_table: LiteLLM_InvitationLinkActions[prisma_models.LiteLLM_InvitationLink] = InvitationLinkRepository(
prisma_client
).table
return invitation_table
) -> "_TableActions[prisma_models.LiteLLM_InvitationLink]":
return _typed_table(InvitationLinkRepository(prisma_client))
def _team_membership_table(
prisma_client: "PrismaClient | None",
) -> "LiteLLM_TeamMembershipActions[prisma_models.LiteLLM_TeamMembership]":
team_membership_table: Final[LiteLLM_TeamMembershipActions[prisma_models.LiteLLM_TeamMembership]] = (
TeamMembershipRepository(prisma_client).table
)
return team_membership_table
) -> "_TableActions[prisma_models.LiteLLM_TeamMembership]":
return _typed_table(TeamMembershipRepository(prisma_client))
def _hash_password_in_dict(data: dict) -> None:
@ -234,7 +322,7 @@ async def _check_duplicate_user_field(
if case_insensitive:
where_clause[field_name]["mode"] = "insensitive"
existing_user: Final = await UserRepository(prisma_client).table.find_first(where=where_clause)
existing_user: Final = await _user_table_lenient(prisma_client).find_first(where=where_clause)
if existing_user is not None:
existing_value: Final = getattr(existing_user, field_name, value)
@ -650,7 +738,7 @@ async def ui_get_available_role(
def get_team_from_list(
team_list: list[LiteLLM_TeamTable] | list[TeamListResponseObject] | None,
team_list: Sequence[LiteLLM_TeamTable] | Sequence[TeamListResponseObject] | None,
team_id: str,
) -> LiteLLM_TeamTable | LiteLLM_TeamMembership | None:
if team_list is None:
@ -732,18 +820,59 @@ def _enforce_user_info_access(user_id: str | None, user_api_key_dict: UserAPIKey
)
async def _get_user_info_teams(
prisma_client: Any,
_TeamIdList: TypeAlias = list[str]
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: _TeamIdList,
table_name: Literal["team"],
query_type: Literal["find_all"],
) -> "Sequence[TeamListResponseObject] | None": ...
async def _get_user_info_row(
prisma_client: "_UserInfoDataClient",
user_id: str,
) -> "prisma_models.LiteLLM_UserTable | None":
return await prisma_client.get_data(user_id=user_id)
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[Any], list[Any] | None]:
) -> tuple[Sequence[TeamListResponseObject], Sequence[TeamListResponseObject] | None]:
"""Fetch and merge teams from membership + user.teams field."""
from litellm.proxy.management_endpoints.team_endpoints import list_team
team_list: list[Any] = []
team_id_list: list[str] = []
teams_1: Final = await list_team(
http_request=Request(
scope={"type": "http", "path": "/user/info"},
@ -752,11 +881,10 @@ async def _get_user_info_teams(
user_api_key_dict=user_api_key_dict,
)
if teams_1 is not None and isinstance(teams_1, list):
team_list = teams_1
team_id_list = [team.team_id for team in teams_1]
team_list: Final = teams_1 if teams_1 is not None and isinstance(teams_1, list) else list[TeamListResponseObject]()
team_id_list: Final = [team.team_id for team in team_list]
teams_2: list[Any] | 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):
@ -767,7 +895,7 @@ async def _get_user_info_teams(
)
elif user_api_key_dict.user_id is not None and user_id is None:
caller_user_info: Final = await prisma_client.get_data(user_id=user_api_key_dict.user_id)
caller_team_ids: Final = getattr(caller_user_info, "teams", None)
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,
@ -804,9 +932,9 @@ def _redact_scim_enterprise_metadata(
def _build_user_info_response(
user_id: str | None,
user_info: Any | None,
keys: list[LiteLLM_VerificationToken] | None,
team_list: list[Any],
teams_1: list[Any] | None,
keys: Sequence[LiteLLM_VerificationToken] | None,
team_list: Sequence[TeamListResponseObject],
teams_1: Sequence[TeamListResponseObject] | None,
) -> UserInfoResponse:
"""Create UserInfoResponse while filtering sensitive fields."""
if user_info is None and keys is not None:
@ -814,7 +942,7 @@ def _build_user_info_response(
user_info = {"spend": spend}
returned_keys: Final = _process_keys_for_user_info(keys=keys, all_teams=teams_1)
team_list.sort(key=lambda x: getattr(x, "team_alias", "") or "")
sorted_team_list: Final = sorted(team_list, key=lambda x: getattr(x, "team_alias", "") or "")
_user_info: Final = user_info.model_dump() if isinstance(user_info, BaseModel) else user_info
if isinstance(_user_info, dict):
@ -825,7 +953,7 @@ def _build_user_info_response(
user_id=user_id,
user_info=_user_info,
keys=returned_keys,
teams=team_list,
teams=sorted_team_list,
)
@ -870,9 +998,9 @@ async def user_info(
user_id = user_api_key_dict.user_id
## GET USER ROW ##
user_info = None
user_info: prisma_models.LiteLLM_UserTable | None = None
if user_id is not None:
user_info = await prisma_client.get_data(user_id=user_id)
user_info = await _get_user_info_row(prisma_client, user_id)
if user_info is None:
raise HTTPException(
@ -888,11 +1016,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,
@ -1058,6 +1182,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
@ -1081,22 +1211,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[list] = 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 = 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
@ -1121,8 +1254,8 @@ async def _get_user_info_for_proxy_admin(user_api_key_dict: UserAPIKeyAuth):
def _process_keys_for_user_info(
keys: list[LiteLLM_VerificationToken] | None,
all_teams: list[LiteLLM_TeamTable] | list[TeamListResponseObject] | None,
keys: Sequence[LiteLLM_VerificationToken] | None,
all_teams: Sequence[LiteLLM_TeamTable] | Sequence[TeamListResponseObject] | None,
):
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
from litellm.proxy.proxy_server import general_settings, litellm_master_key_hash
@ -1212,7 +1345,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,
@ -1768,7 +1901,10 @@ async def bulk_user_update(
# Apply update transformations (reuse existing logic)
data_json: Final[dict] = data.user_updates.model_dump(exclude_unset=True)
non_default_values: Final = _update_internal_user_params(data_json=data_json, data=data.user_updates)
_raw_update_values: Final[Mapping[str, object]] = _update_internal_user_params(
data_json=data_json, data=data.user_updates
)
non_default_values: Final = dict[str, object](_raw_update_values)
# Remove user identification fields since we're updating by user_id
non_default_values.pop("user_id", None)
@ -1780,7 +1916,7 @@ async def bulk_user_update(
try:
# Perform bulk database update
await UserRepository(prisma_client).table.update_many(
await _user_table_lenient(prisma_client).update_many(
where={},
data=non_default_values, # Update all users
)
@ -1885,7 +2021,7 @@ async def get_user_key_counts(
# Get count for each user_id individually
for user_id in user_ids:
count = await VerificationTokenRepository(prisma_client).table.count(
count = await _verification_token_table_lenient(prisma_client).count(
where={
"user_id": user_id,
"OR": [
@ -2122,7 +2258,7 @@ async def get_users(
_validate_sort_params(sort_by, sort_order) if sort_by is not None and isinstance(sort_by, str) else None
)
users: Sequence[prisma_models.LiteLLM_UserTable] | None = await UserRepository(prisma_client).table.find_many(
users: Sequence[prisma_models.LiteLLM_UserTable] | None = await _user_table_lenient(prisma_client).find_many(
where=where_conditions,
skip=skip,
take=page_size,
@ -2130,7 +2266,7 @@ async def get_users(
)
# Get total count of user rows
total_count: Final[int] = await UserRepository(prisma_client).table.count(where=where_conditions)
total_count: Final[int] = await _user_table_lenient(prisma_client).count(where=where_conditions)
# Get key count for each user
if users is not None:
@ -2256,7 +2392,7 @@ async def delete_user(
# check that all teams passed exist
for user_id in data.user_ids:
user_row = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id})
user_row = await _user_delete_table(prisma_client).find_unique(where={"user_id": user_id})
if user_row is None:
raise HTTPException(
@ -2308,8 +2444,8 @@ async def delete_user(
)
## CLEANUP MEMBERS_WITH_ROLES
fetch_all_teams = await TeamRepository(prisma_client).table.find_many(where={"team_id": {"in": user_row.teams}})
teams_to_update = []
fetch_all_teams = await _team_cleanup_table(prisma_client).find_many(where={"team_id": {"in": user_row.teams}})
teams_to_update = list[_TeamCleanupRow]()
for team in fetch_all_teams:
is_member_in_team, new_team_members = _cleanup_members_with_roles(
existing_team_row=LiteLLM_TeamTable.model_validate(team.model_dump()),
@ -2327,7 +2463,7 @@ async def delete_user(
## update teams
for team in teams_to_update:
await TeamRepository(prisma_client).table.update(
await _team_cleanup_table(prisma_client).update(
where={"team_id": team.team_id},
data={"members_with_roles": team.members_with_roles},
)
@ -2382,14 +2518,14 @@ async def add_internal_user_to_organization(
try:
# Check if organization_id exists
organization_row: Final = await OrganizationRepository(prisma_client).table.find_unique(
organization_row: Final = await _organization_table(prisma_client).find_unique(
where={"organization_id": organization_id}
)
if organization_row is None:
raise Exception(f"Organization not found, passed organization_id={organization_id}")
# Create a new organization membership entry
new_membership: Final = await OrganizationMembershipRepository(prisma_client).table.create(
new_membership: Final = await _organization_membership_table(prisma_client).create(
data={
"user_id": user_id,
"organization_id": organization_id,

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

@ -18,9 +18,10 @@ 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 Any, Final, Literal, Optional, Protocol, TypeVar, cast
from typing import Any, Final, Literal, Optional, Protocol, TypeAlias, TypeVar, cast
import fastapi
import yaml
@ -88,6 +89,7 @@ from litellm.proxy.management_endpoints.common_utils import (
from litellm.proxy.management_endpoints.model_management_endpoints import (
_add_model_to_db,
)
from litellm.proxy.management_helpers import object_permission_utils
from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at
from litellm.proxy.management_helpers.object_permission_utils import (
_set_object_permission,
@ -144,6 +146,7 @@ from litellm.types.utils import (
)
_PrismaRowT = TypeVar("_PrismaRowT")
_PrismaRowCoT: Final = TypeVar("_PrismaRowCoT", covariant=True)
_RepositoryModelT = TypeVar("_RepositoryModelT", bound=BaseModel)
@ -169,7 +172,7 @@ class _PrismaTableActions(Protocol[_PrismaRowT]):
*,
where: Mapping[str, object] | None = None,
include: Mapping[str, object] | None = None,
order: Mapping[str, object] | None = None,
order: Mapping[str, object] | Sequence[Mapping[str, object]] | None = None,
skip: int | None = None,
take: int | None = None,
) -> list[_PrismaRowT]: ...
@ -189,17 +192,73 @@ class _PrismaTableActions(Protocol[_PrismaRowT]):
data: Mapping[str, object],
) -> _PrismaRowT | None: ...
async def upsert(
self,
*,
where: Mapping[str, object],
data: Mapping[str, object],
) -> _PrismaRowT: ...
class _UserRowLike(Protocol):
user_id: str | None
user_email: str | None
user_alias: str | None
def model_dump(self) -> Mapping[str, object]: ...
class _PrismaTableHolder(Protocol[_PrismaRowT]):
@property
def table(self) -> _PrismaTableActions[_PrismaRowT]: ...
def _typed_table(holder: _PrismaTableHolder[_PrismaRowT]) -> _PrismaTableActions[_PrismaRowT]:
return holder.table
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()
class _PrismaTableLenient(Protocol[_PrismaRowCoT]):
async def find_unique(self, *, where: Mapping[str, object]) -> _PrismaRowCoT: ...
async def find_many(self, *, where: Mapping[str, object] | None = None) -> Sequence[_PrismaRowCoT] | None: ...
class _PrismaTableLenientHolder(Protocol[_PrismaRowCoT]):
@property
def table(self) -> _PrismaTableLenient[_PrismaRowCoT]: ...
def _lenient_table(holder: _PrismaTableLenientHolder[_PrismaRowCoT]) -> _PrismaTableLenient[_PrismaRowCoT]:
return holder.table
def _prisma_table_lenient(
repository: BaseRepository[_RepositoryModelT],
) -> _PrismaTableLenient[_RepositoryModelT]:
return _lenient_table(repository)
def _jsonify_for_db(client: PrismaClient, data: Mapping[str, object]) -> Mapping[str, object]:
return client.jsonify_object(dict[str, object](data))
class _TxTables(Protocol):
litellm_proxymodeltable: _PrismaTableActions[object]
@ -207,21 +266,82 @@ class _TxTables(Protocol):
def _prisma_table(
repository: BaseRepository[_RepositoryModelT],
) -> _PrismaTableActions[_RepositoryModelT]:
return repository.table
return _typed_table(repository)
def _deleted_verification_token_table(
prisma_client: PrismaClient,
) -> _PrismaTableActions[LiteLLM_DeletedVerificationToken]:
return DeletedVerificationTokenRepository(prisma_client).table
return _typed_table(DeletedVerificationTokenRepository(prisma_client))
def _credentials_table(prisma_client: PrismaClient) -> _PrismaTableActions[CredentialItem]:
return CredentialsRepository(prisma_client).table
return _typed_table(CredentialsRepository(prisma_client))
def _config_table(prisma_client: PrismaClient) -> _PrismaTableActions[ConfigParam]:
return ConfigRepository(prisma_client).table
return _typed_table(ConfigRepository(prisma_client))
def _deprecated_verification_token_table(prisma_client: PrismaClient) -> _PrismaTableActions[object]:
return _typed_table(DeprecatedVerificationTokenRepository(prisma_client))
_StringList: TypeAlias = list[str]
class _CreatedUserRow(Protocol):
models: _StringList
def _created_user_row(user_row: "_CreatedUserRow | None") -> "_CreatedUserRow | None":
return user_row
async def _query_raw_text_rows(prisma_client: PrismaClient, sql: str, *params: object) -> Sequence[Mapping[str, str]]:
return await prisma_client.db.query_raw(sql, *params)
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 _SpendCache(Protocol):
async def async_get_cache(self, key: str) -> float | None: ...
def _spend_cache(cache: _SpendCache) -> _SpendCache:
return cache
class _ObjectPermissionUtils(Protocol):
@property
def attach_object_permission_to_dict(
self,
) -> Callable[..., Awaitable[Mapping[str, object]]]: ...
def _object_permission_utils(module: _ObjectPermissionUtils) -> _ObjectPermissionUtils:
return module
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:
@ -886,7 +1006,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 and key in [
"max_budget",
@ -984,9 +1104,9 @@ async def _common_key_generation_helper(
soft_budget=data.soft_budget,
model_max_budget=data.model_max_budget or {},
)
new_budget: Final = prisma_client.jsonify_object(budget_row.json(exclude_none=True))
new_budget: Final = _jsonify_for_db(prisma_client, budget_row.json(exclude_none=True))
_budget: Final = await BudgetRepository(prisma_client).table.create(
_budget: Final = await _prisma_table(BudgetRepository(prisma_client)).create(
data={
**new_budget,
"created_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
@ -1655,11 +1775,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:
@ -1686,7 +1806,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):
@ -1855,11 +1975,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:
@ -1887,7 +2007,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)
@ -1959,7 +2081,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():
@ -2737,13 +2859,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:
@ -2774,7 +2896,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)
@ -2927,14 +3051,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,
@ -2980,7 +3106,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(
@ -3058,7 +3184,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)
@ -3089,14 +3215,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,
@ -3223,7 +3351,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],
)
@ -3435,7 +3563,7 @@ async def _get_model_max_budget_current_spend(
virtual_key_model_spend_cache_key = (
f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{api_key_hash}:{model}:{budget_config.budget_duration}"
)
current_spend: float | None = await user_api_key_cache.async_get_cache(
current_spend: float | None = await _spend_cache(user_api_key_cache).async_get_cache(
key=virtual_key_model_spend_cache_key,
)
if current_spend is None:
@ -3444,7 +3572,7 @@ async def _get_model_max_budget_current_spend(
f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:"
f"{api_key_hash}:{model_without_prefix}:{budget_config.budget_duration}"
)
current_spend = await user_api_key_cache.async_get_cache(
current_spend = await _spend_cache(user_api_key_cache).async_get_cache(
key=virtual_key_model_spend_cache_key,
)
try:
@ -3635,7 +3763,7 @@ async def info_key_fn(
hashed_key: str | None = key
if key is not None:
hashed_key = _hash_token_if_needed(token=key)
key_info = await VerificationTokenRepository(prisma_client).table.find_unique(
key_info = await _prisma_table(VerificationTokenRepository(prisma_client)).find_unique(
where={"token": hashed_key},
include={"litellm_budget_table": True},
)
@ -3945,7 +4073,7 @@ async def generate_key_helper_fn(
if table_name is None or table_name == "user": # do not auto-create users for `/key/generate`
## CREATE USER (If necessary)
if query_type == "insert_data":
user_row = await prisma_client.insert_data(data=user_data, table_name="user")
user_row = _created_user_row(await prisma_client.insert_data(data=user_data, table_name="user"))
if user_row is None:
raise Exception("Failed to create user")
@ -4232,7 +4360,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)
@ -4355,13 +4483,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(
@ -4376,14 +4503,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,
@ -4449,7 +4576,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:
@ -4605,7 +4732,7 @@ async def _insert_deprecated_key(
try:
revoke_at: Final = datetime.now(timezone.utc) + timedelta(seconds=grace_seconds)
await DeprecatedVerificationTokenRepository(prisma_client).table.upsert(
await _deprecated_verification_token_table(prisma_client).upsert(
where={"token": old_token_hash},
data={
"create": {
@ -4704,11 +4831,13 @@ async def _execute_virtual_key_regeneration(
grace_period=data.grace_period if data else None,
)
updated_token: Final[Mapping[str, object] | None] = await VerificationTokenRepository(prisma_client).table.update(
updated_token: Final = await _prisma_table(VerificationTokenRepository(prisma_client)).update(
where={"token": hashed_api_key},
data=with_settings_updated_at(jsonified_update_data),
)
updated_token_dict: Final[dict[str, object]] = dict(updated_token) if updated_token is not None else {}
updated_token_dict: Final = (
dict[str, object](_as_object_dict(dict(updated_token))) if updated_token is not None else dict[str, object]()
)
updated_token_dict["key"] = new_token
updated_token_dict["token_id"] = updated_token_dict.pop("token")
@ -5247,7 +5376,7 @@ async def validate_key_list_check(
if key_hash:
try:
key_info: Final = await VerificationTokenRepository(prisma_client).table.find_unique(
key_info: Final = await _prisma_table_lenient(VerificationTokenRepository(prisma_client)).find_unique(
where={"token": key_hash},
)
except Exception:
@ -5278,7 +5407,7 @@ async def _fetch_user_team_objects(
if complete_user_info is None or not complete_user_info.teams:
return []
teams: Final[list[BaseModel] | None] = await TeamRepository(prisma_client).table.find_many(
teams: Final = await _prisma_table_lenient(TeamRepository(prisma_client)).find_many(
where={"team_id": {"in": complete_user_info.teams}}
)
if teams is None:
@ -5653,7 +5782,7 @@ async def key_aliases(
where_sql: Final = " AND ".join(where_parts)
count_sql: Final = f'SELECT COUNT(*) AS count FROM "LiteLLM_VerificationToken" WHERE {where_sql}'
count_rows: Final[Sequence[Mapping[str, int]]] = await prisma_client.db.query_raw(count_sql, *query_params)
count_rows: Final = await _query_raw_text_rows(prisma_client, count_sql, *query_params)
total_count: Final = int(count_rows[0]["count"]) if count_rows else 0
aliases_params: Final = query_params + [size, (page - 1) * size]
@ -5666,7 +5795,7 @@ async def key_aliases(
f" ORDER BY key_alias ASC"
f" LIMIT ${limit_idx} OFFSET ${offset_idx}"
)
alias_rows: Final[Sequence[Mapping[str, str]]] = await prisma_client.db.query_raw(aliases_sql, *aliases_params)
alias_rows: Final = await _query_raw_text_rows(prisma_client, aliases_sql, *aliases_params)
aliases: Final[list[str]] = [row["key_alias"] for row in alias_rows if row.get("key_alias")]
total_pages: Final = -(-total_count // size) if total_count > 0 else 0
@ -5952,7 +6081,7 @@ async def _list_key_helper(
# Fetch keys with pagination
if use_deleted_table:
keys = await DeletedVerificationTokenRepository(prisma_client).table.find_many(
keys = await _deleted_verification_token_table(prisma_client).find_many(
where=where,
skip=skip,
take=size,
@ -5966,7 +6095,7 @@ async def _list_key_helper(
),
)
else:
keys = await VerificationTokenRepository(prisma_client).table.find_many(
keys = await _prisma_table(VerificationTokenRepository(prisma_client)).find_many(
where=where,
skip=skip,
take=size,
@ -5995,13 +6124,13 @@ async def _list_key_helper(
total_pages: Final = -(-total_count // size) # Ceiling division
# Fetch user information if expand includes "user"
user_map = {}
user_map: Mapping[str, LiteLLM_UserTable] = {}
if expand and "user" in expand:
user_ids: Final = [key.user_id for key in keys if key.user_id]
created_by_ids: Final = [key.created_by for key in keys if key.created_by]
all_ids: Final = list(set(user_ids + created_by_ids)) # Remove duplicates
if all_ids:
users: Final[Sequence[_UserRowLike]] = await UserRepository(prisma_client).table.find_many(
users: Final = await _prisma_table(UserRepository(prisma_client)).find_many(
where={"user_id": {"in": all_ids}}
)
user_map = {user.user_id: user for user in users}
@ -6014,10 +6143,14 @@ async def _list_key_helper(
key_dict = key.model_dump()
except Exception:
# Fallback for Pydantic v1 compatibility
key_dict = key.dict()
key_dict = dict[str, object](_legacy_model_dict(key))
# Attach object_permission if object_permission_id is set (only for non-deleted keys)
if not use_deleted_table:
key_dict = await attach_object_permission_to_dict(key_dict, prisma_client)
key_dict = dict[str, object](
await _object_permission_utils(object_permission_utils).attach_object_permission_to_dict(
key_dict, prisma_client
)
)
# Include user information if expand includes "user"
if expand and "user" in expand:
@ -6025,7 +6158,7 @@ async def _list_key_helper(
try:
key_dict["user"] = user_map[key.user_id].model_dump()
except Exception:
key_dict["user"] = user_map[key.user_id].dict()
key_dict["user"] = _legacy_model_dict(user_map[key.user_id])
if key.created_by and key.created_by in user_map:
created_by_user = user_map[key.created_by]
key_dict["created_by_user"] = {
@ -6039,7 +6172,7 @@ async def _list_key_helper(
# Use deleted key type to preserve deleted_at, deleted_by, etc.
key_list.append(LiteLLM_DeletedVerificationToken.model_validate(key_dict))
else:
key_list.append(UserAPIKeyAuth(**key_dict)) # Return full key object
key_list.append(UserAPIKeyAuth.model_validate(key_dict)) # Return full key object
else:
_token = key_dict.get("token")
key_list.append(cast(str, _token)) # Return only the token

213
litellm/proxy/management_endpoints/tool_management_endpoints.py
View file

@ -10,8 +10,9 @@ POST /v1/tool/policy - Update the input_policy / output_policy for a
"""
import uuid
from collections.abc import Mapping, Sequence
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Annotated, Any, Final
from typing import TYPE_CHECKING, Annotated, Final, Protocol
from fastapi import APIRouter, Depends, HTTPException, Query
from pydantic import BaseModel, Field, TypeAdapter
@ -49,6 +50,144 @@ from litellm.types.tool_management import (
ToolUsageLogsResponse,
)
class _DailyToolSpendRecord(Protocol):
date: str
tool_name: str
spend: float
request_count: int
class _SpendLogToolIndexRecord(Protocol):
request_id: str
class _SpendLogRecord(Protocol):
request_id: str
startTime: datetime
model: str | None
spend: float | None
total_tokens: int | None
messages: object
proxy_server_request: object
class _VerificationTokenRecord(Protocol):
object_permission_id: str | None
class _TeamRecord(Protocol):
object_permission_id: str | None
class _DailyToolSpendTable(Protocol):
async def group_by(
self,
*,
by: Sequence[str],
sum: Mapping[str, bool],
where: Mapping[str, object],
order: Mapping[str, object],
take: int,
) -> Sequence[object] | None: ...
async def find_many(
self,
*,
where: Mapping[str, object],
order: Sequence[Mapping[str, str]],
) -> Sequence[_DailyToolSpendRecord]: ...
class _SpendLogToolIndexTable(Protocol):
async def count(self, *, where: Mapping[str, object]) -> int: ...
async def find_many(
self,
*,
where: Mapping[str, object],
order: Mapping[str, str],
skip: int,
take: int,
) -> Sequence[_SpendLogToolIndexRecord]: ...
class _SpendLogsTable(Protocol):
async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_SpendLogRecord]: ...
class _VerificationTokenTable(Protocol):
async def find_unique(self, *, where: Mapping[str, object]) -> _VerificationTokenRecord | None: ...
async def update_many(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> int: ...
class _TeamTable(Protocol):
async def find_unique(self, *, where: Mapping[str, object]) -> _TeamRecord | None: ...
async def update_many(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> int: ...
class _ObjectPermissionTable(Protocol):
async def create(self, *, data: Mapping[str, str | Sequence[str]]) -> object: ...
async def delete(self, *, where: Mapping[str, object]) -> object: ...
class _DailyToolSpendTableHolder(Protocol):
@property
def table(self) -> _DailyToolSpendTable: ...
class _SpendLogToolIndexTableHolder(Protocol):
@property
def table(self) -> _SpendLogToolIndexTable: ...
class _SpendLogsTableHolder(Protocol):
@property
def table(self) -> _SpendLogsTable: ...
class _VerificationTokenTableHolder(Protocol):
@property
def table(self) -> _VerificationTokenTable: ...
class _TeamTableHolder(Protocol):
@property
def table(self) -> _TeamTable: ...
class _ObjectPermissionTableHolder(Protocol):
@property
def table(self) -> _ObjectPermissionTable: ...
def _daily_tool_spend_table(repo: _DailyToolSpendTableHolder) -> _DailyToolSpendTable:
return repo.table
def _spend_log_tool_index_table(repo: _SpendLogToolIndexTableHolder) -> _SpendLogToolIndexTable:
return repo.table
def _spend_logs_table(repo: _SpendLogsTableHolder) -> _SpendLogsTable:
return repo.table
def _verification_token_table(repo: _VerificationTokenTableHolder) -> _VerificationTokenTable:
return repo.table
def _team_table(repo: _TeamTableHolder) -> _TeamTable:
return repo.table
def _object_permission_table(repo: _ObjectPermissionTableHolder) -> _ObjectPermissionTable:
return repo.table
router: Final = APIRouter()
TOOL_POLICY_OPTIONS: Final = ToolPolicyOptionsResponse(
@ -154,6 +293,7 @@ class _TopToolRow(BaseModel):
_TOP_TOOL_ROWS: Final = TypeAdapter(list[_TopToolRow])
_PARSED_JSON: Final = TypeAdapter(object)
@router.get(
@ -201,7 +341,7 @@ async def get_tool_spend(
end_str: Final = end_day.strftime("%Y-%m-%d")
date_window: Final = {"date": {"gte": start_str, "lte": end_str}}
table: Final = DailyToolSpendRepository(prisma_client).table
table: Final = _daily_tool_spend_table(DailyToolSpendRepository(prisma_client))
top_tools: Final = _TOP_TOOL_ROWS.validate_python(
await table.group_by(
by=["tool_name"],
@ -222,7 +362,7 @@ async def get_tool_spend(
for row in top_tools
]
daily_rows: Final = (
daily_rows: Final[Sequence[_DailyToolSpendRecord]] = (
await table.find_many(
where={**date_window, "tool_name": {"in": [row.tool_name for row in top_tools]}},
order=[{"date": "asc"}, {"spend": "desc"}],
@ -270,23 +410,23 @@ async def get_tool_detail(
raise HTTPException(status_code=500, detail=str(e))
def _input_snippet_for_tool_log(sl: Any, max_len: int = 200) -> str | None:
def _input_snippet_for_tool_log(sl: _SpendLogRecord | None, max_len: int = 200) -> str | None:
"""Short snippet from messages or proxy_server_request for tool usage log row."""
if sl is None:
return None
messages: Final = getattr(sl, "messages", None)
messages: Final[object] = getattr(sl, "messages", None)
if messages is not None:
s = _snippet_str(messages, max_len)
if s:
return s
psr = getattr(sl, "proxy_server_request", None)
psr: object = getattr(sl, "proxy_server_request", None)
if not psr:
return None
if isinstance(psr, str):
import json
try:
psr = json.loads(psr)
psr = _PARSED_JSON.validate_python(json.loads(psr))
except Exception:
return _snippet_str(psr, max_len)
if isinstance(psr, dict):
@ -299,7 +439,7 @@ def _input_snippet_for_tool_log(sl: Any, max_len: int = 200) -> str | None:
return _snippet_str(psr, max_len)
def _snippet_str(text: Any, max_len: int = 200) -> str | None:
def _snippet_str(text: object, max_len: int = 200) -> str | None:
if text is None:
return None
if isinstance(text, str):
@ -344,10 +484,9 @@ async def get_tool_usage_logs(
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
try:
where: Final[dict] = {"tool_name": tool_name}
start_time_filter: datetime | None = None
end_time_filter: datetime | None = None
if start_date or end_date:
start_time_filter: datetime | None = None
end_time_filter: datetime | None = None
if start_date:
try:
start_time_filter = datetime.strptime(start_date + "T00:00:00", "%Y-%m-%dT%H:%M:%S").replace(
@ -362,15 +501,15 @@ async def get_tool_usage_logs(
)
except ValueError:
pass
if start_time_filter is not None or end_time_filter is not None:
where["start_time"] = {}
if start_time_filter is not None:
where["start_time"]["gte"] = start_time_filter
if end_time_filter is not None:
where["start_time"]["lte"] = end_time_filter
start_time_range: Final[Mapping[str, datetime]] = {
key: value for key, value in (("gte", start_time_filter), ("lte", end_time_filter)) if value is not None
}
where: Final[Mapping[str, str | Mapping[str, datetime]]] = (
{"tool_name": tool_name, "start_time": start_time_range} if start_time_range else {"tool_name": tool_name}
)
total: Final = await SpendLogToolIndexRepository(prisma_client).table.count(where=where)
index_rows: Final = await SpendLogToolIndexRepository(prisma_client).table.find_many(
total: Final = await _spend_log_tool_index_table(SpendLogToolIndexRepository(prisma_client)).count(where=where)
index_rows: Final = await _spend_log_tool_index_table(SpendLogToolIndexRepository(prisma_client)).find_many(
where=where,
order={"start_time": "desc"},
skip=(page - 1) * page_size,
@ -380,7 +519,9 @@ async def get_tool_usage_logs(
if not request_ids:
return ToolUsageLogsResponse(logs=[], total=total, page=page, page_size=page_size)
spend_logs = await SpendLogsRepository(prisma_client).table.find_many(where={"request_id": {"in": request_ids}})
spend_logs = await _spend_logs_table(SpendLogsRepository(prisma_client)).find_many(
where={"request_id": {"in": request_ids}}
)
log_by_id: Final = {s.request_id: s for s in spend_logs}
logs_out: Final[list[ToolUsageLogEntry]] = []
@ -449,23 +590,29 @@ async def _resolve_key_hash_to_object_permission_id(
hashed: Final = key_hash if "sk-" not in (key_hash or "") else hash_token(key_hash)
if not hashed:
return None
row = await VerificationTokenRepository(prisma_client).table.find_unique(where={"token": hashed})
row = await _verification_token_table(VerificationTokenRepository(prisma_client)).find_unique(
where={"token": hashed}
)
if row is None:
return None
op_id: Final = getattr(row, "object_permission_id", None)
op_id: Final[str | None] = getattr(row, "object_permission_id", None)
if op_id:
return op_id
new_id: Final = str(uuid.uuid4())
await ObjectPermissionRepository(prisma_client).table.create(
await _object_permission_table(ObjectPermissionRepository(prisma_client)).create(
data={"object_permission_id": new_id, "blocked_tools": []}
)
updated_count: Final = await VerificationTokenRepository(prisma_client).table.update_many(
updated_count: Final = await _verification_token_table(VerificationTokenRepository(prisma_client)).update_many(
where={"token": hashed, "object_permission_id": None},
data={"object_permission_id": new_id},
)
if updated_count == 0:
await ObjectPermissionRepository(prisma_client).table.delete(where={"object_permission_id": new_id})
row = await VerificationTokenRepository(prisma_client).table.find_unique(where={"token": hashed})
await _object_permission_table(ObjectPermissionRepository(prisma_client)).delete(
where={"object_permission_id": new_id}
)
row = await _verification_token_table(VerificationTokenRepository(prisma_client)).find_unique(
where={"token": hashed}
)
return getattr(row, "object_permission_id", None) if row else None
return new_id
@ -478,23 +625,25 @@ async def _resolve_team_id_to_object_permission_id(
if not team_id or not team_id.strip():
return None
team_id_clean: Final = team_id.strip()
row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id_clean})
row = await _team_table(TeamRepository(prisma_client)).find_unique(where={"team_id": team_id_clean})
if row is None:
return None
op_id: Final = getattr(row, "object_permission_id", None)
op_id: Final[str | None] = getattr(row, "object_permission_id", None)
if op_id:
return op_id
new_id: Final = str(uuid.uuid4())
await ObjectPermissionRepository(prisma_client).table.create(
await _object_permission_table(ObjectPermissionRepository(prisma_client)).create(
data={"object_permission_id": new_id, "blocked_tools": []}
)
updated_count: Final = await TeamRepository(prisma_client).table.update_many(
updated_count: Final = await _team_table(TeamRepository(prisma_client)).update_many(
where={"team_id": team_id_clean, "object_permission_id": None},
data={"object_permission_id": new_id},
)
if updated_count == 0:
await ObjectPermissionRepository(prisma_client).table.delete(where={"object_permission_id": new_id})
row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id_clean})
await _object_permission_table(ObjectPermissionRepository(prisma_client)).delete(
where={"object_permission_id": new_id}
)
row = await _team_table(TeamRepository(prisma_client)).find_unique(where={"team_id": team_id_clean})
return getattr(row, "object_permission_id", None) if row else None
return new_id

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

@ -3,8 +3,9 @@ CRUD ENDPOINTS FOR PROMPTS
"""
import tempfile
from collections.abc import Mapping, Sequence
from pathlib import Path
from typing import Any, Final, cast
from typing import TYPE_CHECKING, Final, Protocol, cast
from fastapi import (
APIRouter,
@ -15,7 +16,7 @@ from fastapi import (
Response,
UploadFile,
)
from pydantic import BaseModel
from pydantic import BaseModel, TypeAdapter
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import (
@ -38,8 +39,47 @@ from litellm.types.prompts.init_prompts import (
)
from litellm.types.proxy.prompt_endpoints import TestPromptRequest
if TYPE_CHECKING:
from litellm.proxy.prompts.prompt_registry import InMemoryPromptRegistry
from litellm.proxy.utils import PrismaClient
class _PromptRecord(Protocol):
id: str
version: int
environment: str | None
class _PromptTable(Protocol):
async def find_many(
self,
*,
where: Mapping[str, object],
order: Mapping[str, str] | None = None,
take: int | None = None,
distinct: Sequence[str] | None = None,
) -> Sequence[_PromptRecord]: ...
async def create(self, *, data: Mapping[str, str | int | None]) -> _PromptRecord: ...
async def update(self, *, where: Mapping[str, object], data: Mapping[str, str]) -> _PromptRecord: ...
async def delete_many(self, *, where: Mapping[str, str]) -> object: ...
class _PromptTableHolder(Protocol):
@property
def table(self) -> _PromptTable: ...
def _prompt_table(repo: _PromptTableHolder) -> _PromptTable:
return repo.table
router: Final = APIRouter()
_PARSED_VALUE: Final = TypeAdapter(object)
def get_base_prompt_id(prompt_id: str) -> str:
"""
@ -132,7 +172,7 @@ def construct_versioned_prompt_id(prompt_id: str, version: int | None = None) ->
return f"{base_id}.v{version}"
def get_latest_version_prompt_id(prompt_id: str, all_prompt_ids: dict[str, Any]) -> str:
def get_latest_version_prompt_id(prompt_id: str, all_prompt_ids: Mapping[str, object]) -> str:
"""
Find the latest version of a prompt from available prompt IDs.
@ -198,7 +238,9 @@ def get_latest_prompt_versions(prompts: list[PromptSpec]) -> list[PromptSpec]:
return list(latest_prompts.values())
async def get_next_version_for_prompt(prisma_client, prompt_id: str, environment: str = "development") -> int:
async def get_next_version_for_prompt(
prisma_client: "PrismaClient", prompt_id: str, environment: str = "development"
) -> int:
"""
Get the next version number for a prompt in a specific environment.
@ -210,7 +252,7 @@ async def get_next_version_for_prompt(prisma_client, prompt_id: str, environment
Returns:
Next version number (1 if no versions exist, max_version + 1 otherwise)
"""
existing_prompts: Final = await PromptRepository(prisma_client).table.find_many(
existing_prompts: Final = await _prompt_table(PromptRepository(prisma_client)).find_many(
where={"prompt_id": prompt_id, "environment": environment}
)
@ -431,10 +473,10 @@ async def get_prompt_versions(
# Query DB for versions
versioned_prompts: Final = []
if prisma_client is not None:
where_clause: Final[dict[str, Any]] = {"prompt_id": base_prompt_id}
if environment:
where_clause["environment"] = environment
db_prompts: Final = await PromptRepository(prisma_client).table.find_many(
where_clause: Final[Mapping[str, str]] = (
{"prompt_id": base_prompt_id, "environment": environment} if environment else {"prompt_id": base_prompt_id}
)
db_prompts: Final = await _prompt_table(PromptRepository(prisma_client)).find_many(
where=where_clause,
order={"version": "desc"},
)
@ -590,7 +632,7 @@ async def get_prompt_info(
# Query all environments this prompt exists in (lightweight: distinct on environment)
all_environments: list[str] = []
if prisma_client is not None:
all_prompt_rows: Final = await PromptRepository(prisma_client).table.find_many(
all_prompt_rows: Final = await _prompt_table(PromptRepository(prisma_client)).find_many(
where={"prompt_id": base_prompt_id},
distinct=["environment"],
)
@ -602,13 +644,16 @@ async def get_prompt_info(
prompt_spec = None
requested_version: Final = get_version_number(prompt_id=prompt_id) if prompt_id != base_prompt_id else None
if environment and prisma_client is not None:
where_clause: Final[dict[str, Any]] = {
"prompt_id": base_prompt_id,
"environment": environment,
where_clause: Final[Mapping[str, str | int]] = {
key: value
for key, value in (
("prompt_id", base_prompt_id),
("environment", environment),
("version", requested_version),
)
if value is not None
}
if requested_version is not None:
where_clause["version"] = requested_version
env_prompts: Final = await PromptRepository(prisma_client).table.find_many(
env_prompts: Final = await _prompt_table(PromptRepository(prisma_client)).find_many(
where=where_clause,
order={"version": "desc"},
take=1,
@ -721,7 +766,7 @@ async def create_prompt(
)
# Store prompt in db with version
prompt_db_entry: Final = await PromptRepository(prisma_client).table.create(
prompt_db_entry: Final = await _prompt_table(PromptRepository(prisma_client)).create(
data={
"prompt_id": request.prompt_id,
"version": new_version,
@ -811,7 +856,9 @@ async def update_prompt(
)
# Check if any version of this prompt exists (in any environment)
existing_prompts = await PromptRepository(prisma_client).table.find_many(where={"prompt_id": base_prompt_id})
existing_prompts = await _prompt_table(PromptRepository(prisma_client)).find_many(
where={"prompt_id": base_prompt_id}
)
if not existing_prompts:
raise HTTPException(
@ -835,7 +882,7 @@ async def update_prompt(
)
# Store new version in db
prompt_db_entry: Final = await PromptRepository(prisma_client).table.create(
prompt_db_entry: Final = await _prompt_table(PromptRepository(prisma_client)).create(
data={
"prompt_id": base_prompt_id,
"version": new_version,
@ -936,12 +983,12 @@ async def delete_prompt(
base_prompt_id: Final = get_base_prompt_id(prompt_id=prompt_id)
# Build delete filter; scope to environment if provided
delete_where: Final[dict[str, Any]] = {"prompt_id": base_prompt_id}
if environment:
delete_where["environment"] = environment
delete_where: Final[Mapping[str, str]] = (
{"prompt_id": base_prompt_id, "environment": environment} if environment else {"prompt_id": base_prompt_id}
)
# Delete versions from the database (scoped to environment if provided)
await PromptRepository(prisma_client).table.delete_many(where=delete_where)
await _prompt_table(PromptRepository(prisma_client)).delete_many(where=delete_where)
# Remove matching prompts from memory — scope to environment if provided
if environment:
@ -967,7 +1014,9 @@ async def delete_prompt(
raise HTTPException(status_code=500, detail=str(e))
def _reload_prompt_in_registry(registry: Any, versioned_id: str, updated_prompt_spec: PromptSpec) -> PromptSpec:
def _reload_prompt_in_registry(
registry: "InMemoryPromptRegistry", versioned_id: str, updated_prompt_spec: PromptSpec
) -> PromptSpec:
"""Remove stale entry and re-initialize the prompt in the in-memory registry."""
if versioned_id in registry.IN_MEMORY_PROMPTS:
del registry.IN_MEMORY_PROMPTS[versioned_id]
@ -1033,14 +1082,13 @@ async def patch_prompt(
requested_version: Final = get_version_number(prompt_id=prompt_id) if prompt_id != base_prompt_id else None
# Build query to find the exact row by composite unique key
find_where: Final[dict[str, Any]] = {
"prompt_id": base_prompt_id,
"environment": env,
find_where: Final[Mapping[str, str | int]] = {
key: value
for key, value in (("prompt_id", base_prompt_id), ("environment", env), ("version", requested_version))
if value is not None
}
if requested_version is not None:
find_where["version"] = requested_version
db_rows: Final = await PromptRepository(prisma_client).table.find_many(
db_rows: Final = await _prompt_table(PromptRepository(prisma_client)).find_many(
where=find_where,
order={"version": "desc"},
take=1,
@ -1084,15 +1132,18 @@ async def patch_prompt(
raise HTTPException(status_code=400, detail="litellm_params cannot be None")
# Build update data dict
update_data: Final[dict[str, Any]] = {
"litellm_params": updated_litellm_params.model_dump_json(),
"prompt_info": updated_prompt_info.model_dump_json(),
update_data: Final[Mapping[str, str]] = {
key: value
for key, value in (
("litellm_params", updated_litellm_params.model_dump_json()),
("prompt_info", updated_prompt_info.model_dump_json()),
("created_by", user_api_key_dict.user_id),
)
if value
}
if user_api_key_dict.user_id:
update_data["created_by"] = user_api_key_dict.user_id
# Update by primary key (id) to target exactly one row
updated_prompt_db_entry: Final = await PromptRepository(prisma_client).table.update(
updated_prompt_db_entry: Final = await _prompt_table(PromptRepository(prisma_client)).update(
where={"id": target_row.id},
data=update_data,
)
@ -1216,23 +1267,25 @@ async def test_prompt(
# Use ProxyBaseLLMRequestProcessing to go through all proxy logic
base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data)
result: Final = await base_llm_response_processor.base_process_llm_request(
request=fastapi_request,
fastapi_response=fastapi_response,
user_api_key_dict=user_api_key_dict,
route_type="acompletion",
proxy_logging_obj=proxy_logging_obj,
llm_router=llm_router,
general_settings=general_settings,
proxy_config=proxy_config,
select_data_generator=select_data_generator,
model=None,
user_model=user_model,
user_temperature=user_temperature,
user_request_timeout=user_request_timeout,
user_max_tokens=user_max_tokens,
user_api_base=user_api_base,
version=version,
result: Final = _PARSED_VALUE.validate_python(
await base_llm_response_processor.base_process_llm_request(
request=fastapi_request,
fastapi_response=fastapi_response,
user_api_key_dict=user_api_key_dict,
route_type="acompletion",
proxy_logging_obj=proxy_logging_obj,
llm_router=llm_router,
general_settings=general_settings,
proxy_config=proxy_config,
select_data_generator=select_data_generator,
model=None,
user_model=user_model,
user_temperature=user_temperature,
user_request_timeout=user_request_timeout,
user_max_tokens=user_max_tokens,
user_api_base=user_api_base,
version=version,
)
)
if isinstance(result, BaseModel):
@ -1257,7 +1310,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.

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

@ -10,9 +10,12 @@ https://platform.openai.com/docs/api-reference/responses-streaming
import asyncio
import json
from typing import Any, Final, cast
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 TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
@ -20,6 +23,56 @@ from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessin
from litellm.proxy.response_polling.polling_handler import ResponsePollingHandler
from litellm.types.llms.openai import ResponsesAPIStatus
if TYPE_CHECKING:
from litellm.proxy.proxy_server import ProxyConfig
from litellm.proxy.utils import ProxyLogging
from litellm.router import Router
_JsonDict: TypeAlias = dict[str, object]
_JsonList: TypeAlias = list[object]
class _OutputItem(TypedDict, total=False):
id: str
content: Sequence[object]
class _TerminalResponse(TypedDict, total=False):
status: ResponsesAPIStatus
error: _JsonDict
usage: _JsonDict
reasoning: _JsonDict
tool_choice: object
tools: _JsonList
model: str
instructions: str
temperature: float
top_p: float
max_output_tokens: int
previous_response_id: str
text: _JsonDict
truncation: str
parallel_tool_calls: bool
user: str
store: bool
incomplete_details: _JsonDict
output: Sequence[_OutputItem]
class _StreamEvent(TypedDict, total=False):
type: str
item: _OutputItem
item_id: str
content_index: int
delta: str
part: object
response: _TerminalResponse
class _StreamEventParser:
parse: Callable[[str], _StreamEvent] = staticmethod(json.loads)
async def background_streaming_task(
polling_id: str,
@ -29,16 +82,16 @@ async def background_streaming_task(
fastapi_response: Response,
user_api_key_dict: UserAPIKeyAuth,
general_settings: dict,
llm_router,
proxy_config,
proxy_logging_obj,
llm_router: "Router | None",
proxy_config: "ProxyConfig",
proxy_logging_obj: "ProxyLogging",
select_data_generator,
user_model,
user_temperature,
user_request_timeout,
user_max_tokens,
user_api_base,
version,
user_temperature: float | None,
user_request_timeout: float | None,
user_max_tokens: int | None,
user_api_base: str | None,
version: str | None,
):
"""
Background task to stream response and update cache
@ -69,7 +122,7 @@ async def background_streaming_task(
# Make streaming request.
# Pre-call checks (rate limits, guardrails, budget) were already run
# before polling ID creation, so skip them here to avoid double-counting.
response: Final = await processor.base_process_llm_request(
response: Final[StreamingResponse] = await processor.base_process_llm_request(
request=request,
fastapi_response=fastapi_response,
user_api_key_dict=user_api_key_dict,
@ -91,8 +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, dict[str, Any]]] = {} # Track output items by ID
accumulated_text: Final = {} # Track accumulated text deltas by (item_id, content_index)
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
@ -121,7 +176,7 @@ async def background_streaming_task(
None # Will be set by response.completed/failed/incomplete/cancelled
)
terminal_error = None
_event_to_status: Final = {
_event_to_status: Final[Mapping[str, ResponsesAPIStatus]] = {
"response.completed": "completed",
"response.failed": "failed",
"response.incomplete": "incomplete",
@ -162,7 +217,7 @@ async def background_streaming_task(
break
try:
event = json.loads(chunk_data)
event: _StreamEvent = _StreamEventParser.parse(chunk_data)
event_type = event.get("type", "")
# Process different event types based on OpenAI streaming spec
@ -181,9 +236,8 @@ async def background_streaming_task(
if item_id and item_id in output_items:
# Update the output item with new content
if "content" not in output_items[item_id]:
output_items[item_id]["content"] = []
output_items[item_id]["content"].append(content_part)
added_item = output_items[item_id]
added_item["content"] = (*added_item.get("content", ()), content_part)
state_dirty = True
elif event_type == "response.output_text.delta":
@ -201,12 +255,14 @@ async def background_streaming_task(
accumulated_text[key] += delta
# Update the content in output_items
if "content" in output_items[item_id]:
content_list = output_items[item_id]["content"]
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
if isinstance(content_list[content_index], dict):
content_list[content_index]["text"] = accumulated_text[key]
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":
@ -217,10 +273,14 @@ async def background_streaming_task(
if item_id and item_id in output_items:
# Update with final content from event
if "content" in output_items[item_id]:
content_list = output_items[item_id]["content"]
done_item = output_items[item_id]
if "content" in done_item:
content_list = done_item["content"]
if content_index < len(content_list):
content_list[content_index] = content_part
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":
@ -248,12 +308,9 @@ async def background_streaming_task(
# Terminal event - extract all ResponsesAPIResponse fields
# https://platform.openai.com/docs/api-reference/responses-streaming
response_data = event.get("response", {})
terminal_status = cast(
ResponsesAPIStatus,
response_data.get(
"status",
_event_to_status.get(event_type, "completed"),
),
terminal_status = response_data.get(
"status",
_event_to_status.get(event_type, "completed"),
)
# Extract error for failed and incomplete responses

286
litellm/responses/file_search/emulated_handler.py
View file

@ -14,16 +14,19 @@ Flow:
import json
import time
import uuid
from collections.abc import Iterable
from typing import Any, Final, cast
from collections.abc import Iterable, Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, TypedDict, cast
from litellm._internal_context import is_internal_call
from litellm._logging import verbose_logger
from litellm.types.llms.openai import ResponseOutputItem, ResponsesAPIResponse
from litellm.types.llms.openai import ResponsesAPIResponse
from litellm.types.vector_stores import VectorStoreSearchResult
if TYPE_CHECKING:
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
# Keep ToolParam broad so we stay compatible with both dict and Pydantic forms
ToolParam = Any
ToolParam = object
FILE_SEARCH_FUNCTION_NAME: Final = "litellm_file_search"
@ -35,7 +38,7 @@ FILE_SEARCH_FUNCTION_NAME: Final = "litellm_file_search"
def should_use_emulated_file_search(
tools: Iterable[ToolParam] | None,
provider_config: Any, # BaseResponsesAPIConfig
provider_config: "BaseResponsesAPIConfig | None",
) -> bool:
"""Return True when there is a file_search tool and the provider can't handle it natively."""
if not tools:
@ -51,7 +54,7 @@ def should_use_emulated_file_search(
# ---------------------------------------------------------------------------
def _build_function_tool(vector_store_ids: list[str]) -> dict[str, Any]:
def _build_function_tool(vector_store_ids: Sequence[str]) -> Mapping[str, object]:
"""
Create a Responses API function-tool definition that describes file search.
The function accepts one or more natural-language queries (like OpenAI's native
@ -94,27 +97,26 @@ def _build_function_tool(vector_store_ids: list[str]) -> dict[str, Any]:
}
def _file_search_tool_vector_store_ids(tool: object) -> Sequence[str] | None:
if not (isinstance(tool, dict) and tool.get("type") == "file_search"):
return None
return tool.get("vector_store_ids") or []
def _replace_file_search_tools(
tools: Iterable[ToolParam] | None,
) -> tuple[list[dict[str, Any]], list[str]]:
) -> tuple[Sequence[object], Sequence[str]]:
"""
Replace all file_search tools with a single function tool.
Returns:
(new_tools_list, all_vector_store_ids)
"""
non_file_search: Final[list[dict[str, Any]]] = []
vector_store_ids: Final[list[str]] = []
for tool in tools or []:
if isinstance(tool, dict) and tool.get("type") == "file_search":
ids = tool.get("vector_store_ids") or []
vector_store_ids.extend(ids)
else:
non_file_search.append(tool)
ids_and_tools: Final = tuple((_file_search_tool_vector_store_ids(tool), tool) for tool in tools or ())
# Deduplicate while preserving order
unique_ids: Final[list[str]] = list(dict.fromkeys(vector_store_ids))
unique_ids: Final = list(dict.fromkeys(vs_id for ids, _ in ids_and_tools if ids is not None for vs_id in ids))
non_file_search: Final = [tool for ids, tool in ids_and_tools if ids is None]
if unique_ids:
non_file_search.append(_build_function_tool(unique_ids))
@ -127,9 +129,9 @@ def _replace_file_search_tools(
async def _run_vector_searches(
queries: list[str],
vector_store_ids: list[str],
) -> tuple[list[str], list[VectorStoreSearchResult]]:
queries: Sequence[str],
vector_store_ids: Sequence[str],
) -> tuple[Sequence[str], Sequence[VectorStoreSearchResult]]:
"""
Run `asearch` against all vector stores for all queries and collect results.
@ -172,7 +174,7 @@ async def _run_vector_searches(
# ---------------------------------------------------------------------------
def _get_field(result: Any, key: str, default: Any = None) -> Any:
def _get_field(result: object, key: str, default: object = None) -> object:
"""Read a field from either a dict/TypedDict or an attribute-based object."""
if isinstance(result, dict):
return result.get(key, default)
@ -180,7 +182,7 @@ def _get_field(result: Any, key: str, default: Any = None) -> Any:
def _format_search_results_as_tool_output(
results: list[VectorStoreSearchResult],
results: Sequence[VectorStoreSearchResult],
) -> str:
"""Serialize search results into a string to pass back as the tool's output."""
if not results:
@ -191,7 +193,8 @@ def _format_search_results_as_tool_output(
score = _get_field(result, "score")
file_id = _get_field(result, "file_id")
filename = _get_field(result, "filename")
content_items = _get_field(result, "content") or []
raw_content = _get_field(result, "content")
content_items = raw_content if isinstance(raw_content, list) else []
text_chunks = [c.get("text", "") if isinstance(c, dict) else getattr(c, "text", "") for c in content_items]
text = " ".join(t for t in text_chunks if t)
@ -209,9 +212,24 @@ def _format_search_results_as_tool_output(
return "\n\n".join(parts)
def _format_result_for_include(result: VectorStoreSearchResult) -> Mapping[str, object]:
file_id: Final = _get_field(result, "file_id") or ""
raw_content: Final = _get_field(result, "content")
content_items: Final = raw_content if isinstance(raw_content, list) else []
text_chunks: Final = [c.get("text", "") if isinstance(c, dict) else getattr(c, "text", "") for c in content_items]
text: Final = " ".join(t for t in text_chunks if t)
return {
"file_id": file_id,
"filename": _get_field(result, "filename") or "",
"score": _get_field(result, "score"),
"text": text,
"attributes": _get_field(result, "attributes") or {},
}
def _build_search_results_for_include(
results: list[VectorStoreSearchResult],
) -> list[dict[str, Any]]:
results: Sequence[VectorStoreSearchResult],
) -> Sequence[Mapping[str, object]]:
"""
Convert VectorStoreSearchResult objects to the format expected in
file_search_call.search_results (mirrors OpenAI's include= format).
@ -220,30 +238,15 @@ def _build_search_results_for_include(
behaviour of OpenAI's native file_search which surfaces every relevant
chunk even when multiple chunks originate from the same document.
"""
formatted: Final[list[dict[str, Any]]] = []
for result in results:
file_id = _get_field(result, "file_id") or ""
content_items = _get_field(result, "content") or []
text_chunks = [c.get("text", "") if isinstance(c, dict) else getattr(c, "text", "") for c in content_items]
text = " ".join(t for t in text_chunks if t)
formatted.append(
{
"file_id": file_id,
"filename": _get_field(result, "filename") or "",
"score": _get_field(result, "score"),
"text": text,
"attributes": _get_field(result, "attributes") or {},
}
)
return formatted
return [_format_result_for_include(result) for result in results]
def _build_file_search_call_output(
call_id: str,
queries: list[str],
results: list[VectorStoreSearchResult] | None = None,
queries: Sequence[str],
results: Sequence[VectorStoreSearchResult] | None = None,
include_search_results: bool = False,
) -> dict[str, Any]:
) -> Mapping[str, object]:
"""Build the file_search_call output item (mirrors OpenAI's format).
Args:
@ -266,39 +269,34 @@ def _build_file_search_call_output(
def _build_file_citation_annotations(
results: list[VectorStoreSearchResult],
results: Sequence[VectorStoreSearchResult],
text: str,
) -> list[dict[str, Any]]:
) -> Sequence[Mapping[str, object]]:
"""
Build file_citation annotations for the text.
Each result with a file_id gets a citation at the end of the text.
"""
annotations: Final[list[dict[str, Any]]] = []
index: Final = len(text) # cite at end of text block
seen_file_ids: Final[set] = set()
id_filename_pairs: Final = tuple(
(_get_field(result, "file_id"), _get_field(result, "filename")) for result in results
)
first_filename_by_id: Final = {file_id: filename for file_id, filename in reversed(id_filename_pairs) if file_id}
for result in results:
file_id = _get_field(result, "file_id")
filename = _get_field(result, "filename")
if not file_id or file_id in seen_file_ids:
continue
seen_file_ids.add(file_id)
annotations.append(
{
"type": "file_citation",
"index": index,
"file_id": file_id,
"filename": filename or "",
}
)
return annotations
return [
{
"type": "file_citation",
"index": index,
"file_id": file_id,
"filename": first_filename_by_id[file_id] or "",
}
for file_id in dict.fromkeys(file_id for file_id, _ in id_filename_pairs if file_id)
]
def _build_message_output(
response_text: str,
results: list[VectorStoreSearchResult],
) -> dict[str, Any]:
results: Sequence[VectorStoreSearchResult],
) -> Mapping[str, object]:
"""Build the message output item with optional file_citation annotations."""
annotations: Final = _build_file_citation_annotations(results, response_text)
return {
@ -330,8 +328,8 @@ def _extract_text_from_responses_output(response: ResponsesAPIResponse) -> str:
def _synthesize_responses_api_response(
original_response: ResponsesAPIResponse,
file_search_call_output: dict[str, Any],
message_output: dict[str, Any],
file_search_call_output: Mapping[str, object],
message_output: Mapping[str, object],
first_response: ResponsesAPIResponse | None = None,
) -> ResponsesAPIResponse:
"""
@ -343,21 +341,20 @@ def _synthesize_responses_api_response(
synthesized _hidden_params so that billing callbacks see the total cost of
both provider calls that the emulated flow makes.
"""
synthesized_output: Final[list[dict[str, Any]]] = [file_search_call_output, message_output]
synthesized: Final = ResponsesAPIResponse(
id=getattr(original_response, "id", f"resp_{uuid.uuid4().hex}"),
object="response",
created_at=getattr(original_response, "created_at", int(time.time())),
status="completed",
model=getattr(original_response, "model", ""),
output=cast(list[ResponseOutputItem | dict[str, Any]], synthesized_output),
output=[dict(file_search_call_output), dict(message_output)],
usage=getattr(original_response, "usage", None),
error=None,
)
if hasattr(original_response, "_hidden_params"):
hidden: Final = dict(getattr(original_response, "_hidden_params") or {})
if first_response is not None and hasattr(first_response, "_hidden_params"):
first_hidden: Final = getattr(first_response, "_hidden_params") or {}
first_hidden: Final[object] = getattr(first_response, "_hidden_params", None) or {}
first_cost: Final = (
first_hidden.get("response_cost")
if isinstance(first_hidden, dict)
@ -382,9 +379,10 @@ async def _call_aresponses(input, model, tools, **kwargs): # pragma: no cover
def _prepare_emulated_file_search_call(
kwargs: dict[str, Any],
) -> tuple[bool, dict[str, Any]]:
include_items: Final[list[str]] = list(kwargs.get("include") or [])
kwargs: Mapping[str, object],
) -> tuple[bool, Mapping[str, object]]:
raw_include: Final = kwargs.get("include")
include_items: Final[Sequence[str]] = raw_include if isinstance(raw_include, list) else []
include_search_results: Final = "file_search_call.results" in include_items
original_stream: Final = kwargs.get("stream")
@ -398,7 +396,7 @@ def _prepare_emulated_file_search_call(
return include_search_results, updated_kwargs
def _extract_tool_call_fields(tool_call: Any, fallback_call_id: str) -> tuple[str, str]:
def _extract_tool_call_fields(tool_call: object, fallback_call_id: str) -> tuple[str, str]:
"""Extract (call_id, raw_arguments_string) from a dict or Pydantic tool_call item."""
if isinstance(tool_call, dict):
call_id = str(tool_call.get("call_id") or tool_call.get("id") or fallback_call_id)
@ -410,7 +408,13 @@ def _extract_tool_call_fields(tool_call: Any, fallback_call_id: str) -> tuple[st
return call_id, raw_args
def _resolve_queries_from_args(args: dict[str, Any], input: Any) -> list[str]:
class _FileSearchArguments(TypedDict, total=False):
queries: Sequence[str]
query: str
vector_store_id: str
def _resolve_queries_from_args(args: _FileSearchArguments, input: object) -> Sequence[str]:
"""Pull the queries list out of parsed tool-call arguments, with backward-compat fallbacks."""
queries_from_call: Final = args.get("queries")
if not queries_from_call:
@ -422,76 +426,96 @@ def _resolve_queries_from_args(args: dict[str, Any], input: Any) -> list[str]:
return queries_from_call
async def _execute_file_search_tool_calls(
file_search_calls: list[Any],
all_vs_ids: list[str],
input: Any,
def _parse_file_search_arguments(raw_args: str) -> _FileSearchArguments:
if not isinstance(raw_args, str):
return raw_args
try:
return json.loads(raw_args)
except json.JSONDecodeError:
return {}
async def _execute_single_file_search_call(
tool_call: object,
all_vs_ids: Sequence[str],
input: object,
file_search_call_id: str,
) -> tuple[list[dict[str, Any]], list[str], list[VectorStoreSearchResult]]:
) -> tuple[Mapping[str, object], Sequence[str], Sequence[VectorStoreSearchResult]]:
call_id, raw_args = _extract_tool_call_fields(tool_call, fallback_call_id=file_search_call_id)
args: Final = _parse_file_search_arguments(raw_args)
queries_from_call: Final = _resolve_queries_from_args(args, input)
vs_id_arg: Final = args.get("vector_store_id")
vs_ids_for_call: Final = [vs_id_arg] if vs_id_arg else all_vs_ids
queries, results = await _run_vector_searches(
queries=queries_from_call,
vector_store_ids=vs_ids_for_call,
)
return (
{
"type": "function_call_output",
"call_id": call_id,
"output": _format_search_results_as_tool_output(results),
},
queries,
results,
)
async def _execute_file_search_tool_calls(
file_search_calls: Sequence[object],
all_vs_ids: Sequence[str],
input: object,
file_search_call_id: str,
) -> tuple[Sequence[Mapping[str, object]], Sequence[str], Sequence[VectorStoreSearchResult]]:
"""Run the vector search for each file_search tool_call and collect results."""
tool_results: Final[list[dict[str, Any]]] = []
all_queries: Final[list[str]] = []
all_results: Final[list[VectorStoreSearchResult]] = []
per_call: Final = tuple(
[
await _execute_single_file_search_call(
tool_call=tool_call,
all_vs_ids=all_vs_ids,
input=input,
file_search_call_id=file_search_call_id,
)
for tool_call in file_search_calls
]
)
for tool_call in file_search_calls:
call_id, raw_args = _extract_tool_call_fields(tool_call, fallback_call_id=file_search_call_id)
try:
args = json.loads(raw_args) if isinstance(raw_args, str) else raw_args
except json.JSONDecodeError:
args = {}
queries_from_call = _resolve_queries_from_args(args, input)
vs_id_arg = args.get("vector_store_id")
vs_ids_for_call = [vs_id_arg] if vs_id_arg else all_vs_ids
queries, results = await _run_vector_searches(
queries=queries_from_call,
vector_store_ids=vs_ids_for_call,
)
all_queries.extend(queries)
all_results.extend(results)
tool_results.append(
{
"type": "function_call_output",
"call_id": call_id,
"output": _format_search_results_as_tool_output(results),
}
)
return tool_results, all_queries, all_results
return (
[tool_result for tool_result, _, _ in per_call],
[query for _, queries, _ in per_call for query in queries],
[result for _, _, results in per_call for result in results],
)
def _build_follow_up_input(
input: Any,
input: object,
first_response: ResponsesAPIResponse,
tool_results: list[dict[str, Any]],
) -> list[Any]:
tool_results: Sequence[Mapping[str, object]],
) -> Sequence[object]:
"""Assemble the follow-up call input: original messages + first-response output + tool results.
Including all output items (text blocks, reasoning, non-file-search calls) ensures providers
like Anthropic that emit text before the tool call have complete conversation context.
Serializes Pydantic model instances to plain dicts so the transformation layer can call .get().
"""
original_input_items: Final = (
list(input) if isinstance(input, (list, tuple)) else [{"role": "user", "content": str(input)}]
original_input_items: Final[tuple[object, ...]] = (
tuple(input) if isinstance(input, (list, tuple)) else ({"role": "user", "content": str(input)},)
)
first_response_output_items: Final[tuple[object, ...]] = tuple(
_item
if isinstance(_item, dict)
else (_item.model_dump(exclude_none=True) if hasattr(_item, "model_dump") else _item)
for _item in first_response.output
)
first_response_output_items: Final[list[Any]] = []
for _item in first_response.output:
if isinstance(_item, dict):
first_response_output_items.append(_item)
elif hasattr(_item, "model_dump"):
first_response_output_items.append(_item.model_dump(exclude_none=True))
else:
first_response_output_items.append(_item)
return original_input_items + first_response_output_items + tool_results
return [*original_input_items, *first_response_output_items, *tool_results]
async def aresponses_with_emulated_file_search(
input: Any,
input: object,
model: str,
tools: Iterable[ToolParam] | None = None,
# Pass-through params — forwarded as-is to the underlying aresponses call
@ -504,7 +528,7 @@ async def aresponses_with_emulated_file_search(
runs vector search, and synthesizes an OpenAI-format response.
"""
# Determine whether caller wants search_results populated in the output.
_include_search_results, kwargs = _prepare_emulated_file_search_call(kwargs=kwargs)
_include_search_results, call_kwargs = _prepare_emulated_file_search_call(kwargs=kwargs)
# 1. Replace file_search tools with function tool
transformed_tools, all_vs_ids = _replace_file_search_tools(tools)
@ -521,7 +545,7 @@ async def aresponses_with_emulated_file_search(
input=input,
model=model,
tools=transformed_tools or None,
**kwargs,
**call_kwargs,
),
)
finally:
@ -585,7 +609,7 @@ async def aresponses_with_emulated_file_search(
input=follow_up_input,
model=model,
tools=None, # no tools needed for the answer step
**kwargs,
**call_kwargs,
),
)
finally:

227
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
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,12 +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,
)
from litellm.types.router import LiteLLM_Params
@lru_cache(maxsize=1)
@ -69,6 +71,79 @@ 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 _LoadsJsonValue(Protocol):
def __call__(self, s: str | bytes, /) -> object: ...
class _LoadsJsonDict(Protocol):
def __call__(self, s: str | bytes, /) -> _MutableJsonObject: ...
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_loads_json_value(fn: _LoadsJsonValue) -> _LoadsJsonValue:
return fn
def _typed_loads_json_dict(fn: _LoadsJsonDict) -> _LoadsJsonDict:
return fn
def _typed_gets_litellm_params(fn: _GetsLitellmParams) -> _GetsLitellmParams:
return fn
def _typed_pops_optional_str(fn: _PopsOptionalStr) -> _PopsOptionalStr:
return fn
_LOADS_JSON_VALUE: Final = _typed_loads_json_value(json.loads)
_LOADS_JSON_DICT: Final = _typed_loads_json_dict(json.loads)
_SHOULD_STORE_RESULT_IN_CACHE_ATTR: Final = "_should_store_result_in_cache"
_UNMASK_PII_TEXT_ATTR: Final = "_unmask_pii_text"
def _model_id_from_metadata(litellm_metadata: dict[str, object] | None) -> str | None:
model_info: Final = litellm_metadata.get("model_info") if litellm_metadata else None
model_id: Final = model_info.get("id") if _is_json_object(model_info) else None
@ -185,7 +260,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),
@ -228,7 +303,7 @@ class BaseResponsesAPIStreamingIterator:
try:
# Parse the JSON chunk
parsed_chunk: Final = json.loads(chunk)
parsed_chunk: Final = _LOADS_JSON_VALUE(chunk)
# Format as ResponsesAPIStreamingResponse
if isinstance(parsed_chunk, dict):
@ -514,7 +589,7 @@ class BaseResponsesAPIStreamingIterator:
if response_obj is None:
return
caching_handler: Final = 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
@ -532,8 +607,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(
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
@ -586,12 +664,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,
@ -1043,8 +1124,8 @@ class _HasModelDumpJson(Protocol):
def model_dump_json(self, *, exclude_none: bool = ...) -> str: ...
def _dump_response_object(obj: Any) -> dict[str, Any]:
if hasattr(obj, "model_dump"):
def _dump_response_object(obj: object) -> Mapping[str, object]:
if isinstance(obj, _HasModelDump):
return obj.model_dump()
if _is_json_object(obj):
return obj
@ -1073,21 +1154,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(
@ -1117,7 +1197,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()
@ -1134,15 +1214,19 @@ def _add_text_like_part_events(
delta=text[i : i + chunk_size],
)
)
for annotation_index, annotation in enumerate(part_payload.get("annotations", []) or []):
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(
@ -1200,7 +1284,8 @@ def _build_synthetic_response_events(
]
sequence_number = 0
for output_index, output_item in enumerate(getattr(transformed, "output", []) or []):
output_items: Final[Sequence[object]] = getattr(transformed, "output", []) or []
for output_index, output_item in enumerate(output_items):
output_item_payload = _dump_response_object(output_item)
item_id = str(output_item_payload.get("id") or transformed.id)
item_type = output_item_payload.get("type")
@ -1214,7 +1299,9 @@ def _build_synthetic_response_events(
)
if item_type == "message":
for content_index, part in enumerate(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(
openai_types.ContentPartAddedEvent(
@ -1261,7 +1348,9 @@ def _build_synthetic_response_events(
)
)
elif item_type == "reasoning":
for summary_index, summary in enumerate(output_item_payload.get("summary", []) or []):
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):
@ -1354,7 +1443,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[Any] | None = None,
guardrail_callbacks: Sequence[PresidioGuardrailCallback] | None = None,
output_guardrail_callbacks: list[PresidioGuardrailCallback] | None = None,
authorized_model: str | None = None,
):
@ -1363,16 +1452,16 @@ 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[Any] = guardrail_callbacks or []
self.guardrail_callbacks: Sequence[PresidioGuardrailCallback] = guardrail_callbacks or []
self.output_guardrail_callbacks: list[PresidioGuardrailCallback] = output_guardrail_callbacks or []
# 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:
@ -1380,7 +1469,7 @@ class ResponsesWebSocketStreaming:
event = event.decode("utf-8")
if isinstance(event, str):
try:
event_obj = json.loads(event)
event_obj = _LOADS_JSON_DICT(event)
except (json.JSONDecodeError, TypeError):
return
else:
@ -1393,7 +1482,7 @@ class ResponsesWebSocketStreaming:
"""Extract user input content from response.create for logging."""
try:
if isinstance(message, str):
msg_obj = json.loads(message)
msg_obj = _LOADS_JSON_DICT(message)
elif _is_json_object(message):
msg_obj = message
else:
@ -1463,7 +1552,7 @@ class ResponsesWebSocketStreaming:
# masked response.completed.
if self.output_guardrail_callbacks:
try:
_evt_type = json.loads(response_str).get("type")
_evt_type = _LOADS_JSON_DICT(response_str).get("type")
except (json.JSONDecodeError, TypeError):
_evt_type = None
if _evt_type in self._DELTA_EVENT_TYPES or _evt_type in self._OUTPUT_DONE_EVENT_TYPES:
@ -1485,7 +1574,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.
@ -1527,7 +1616,7 @@ class ResponsesWebSocketStreaming:
Non-``response.create`` messages are returned unchanged.
"""
try:
msg_obj: Final = json.loads(message)
msg_obj: Final = _LOADS_JSON_DICT(message)
except (json.JSONDecodeError, TypeError):
return message
@ -1553,7 +1642,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
@ -1655,11 +1744,12 @@ class ResponsesWebSocketStreaming:
return response_str
try:
evt_obj: Final = json.loads(response_str)
evt_obj: Final = _LOADS_JSON_DICT(response_str)
except (json.JSONDecodeError, TypeError):
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":
@ -1679,7 +1769,7 @@ class ResponsesWebSocketStreaming:
continue
text = content_block.get("text")
if isinstance(text, str):
unmasked = cb._unmask_pii_text(text, pii_tokens)
unmasked = unmask_pii_text(text, pii_tokens)
if unmasked != text:
content_block["text"] = unmasked
modified = True
@ -1688,7 +1778,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(delta, pii_tokens)
unmasked = unmask_pii_text(delta, pii_tokens)
if unmasked != delta:
evt_obj["delta"] = unmasked
return json.dumps(evt_obj)
@ -1711,7 +1801,7 @@ class ResponsesWebSocketStreaming:
return response_str
try:
evt_obj: Final[Mapping[str, object]] = json.loads(response_str)
evt_obj: Final = _LOADS_JSON_DICT(response_str)
except (json.JSONDecodeError, TypeError):
return response_str
@ -1859,7 +1949,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,
@ -1871,10 +1961,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
@ -1894,7 +1985,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):
@ -1937,7 +2028,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.
@ -1952,7 +2043,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
@ -2009,10 +2100,10 @@ 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 = json.loads(raw_message)
msg_obj: Final = _LOADS_JSON_DICT(raw_message)
except json.JSONDecodeError:
await self._send_error("Invalid JSON in response.create event", "invalid_request_error")
return None
@ -2022,7 +2113,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
@ -2038,13 +2129,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
@ -2062,7 +2153,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.
@ -2085,7 +2176,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": [...], ...}}
@ -2194,7 +2285,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.
@ -2202,7 +2293,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)
@ -2216,7 +2307,7 @@ class ManagedResponsesWebSocketHandler:
continue
if chunk_type == "response.completed" and completed_event is None:
try:
completed_event = json.loads(serialized)
completed_event = _LOADS_JSON_DICT(serialized)
except Exception:
pass
try:
@ -2228,7 +2319,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:
@ -2293,13 +2384,15 @@ 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 = call_kwargs.pop("model", None)
requested_model: Final = _typed_pops_optional_str(call_kwargs.pop)("model", None)
if requested_model is None or requested_model == self.model_group:
model = self.model
else:
model = 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

174
litellm/router_strategy/tag_based_routing.py
View file

@ -10,10 +10,10 @@ Use this to route requests between Teams
import re
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, overload
from litellm._logging import verbose_logger
from litellm.types.router import RouterErrors
from litellm.types.router import DeploymentTypedDict, RouterErrors
if TYPE_CHECKING:
from litellm.router import Router as _Router
@ -23,9 +23,68 @@ else:
LitellmRouter = Any
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 _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 _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 _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: ...
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(
tag_regexes: list[str],
header_strings: list[str],
tag_regexes: Sequence[str],
header_strings: Sequence[str],
) -> str | None:
"""
Test compiled regex patterns against "Header-Name: value" strings.
@ -46,7 +105,9 @@ def _is_valid_deployment_tag_regex(
return None
def is_valid_deployment_tag(deployment_tags: list[str], request_tags: list[str], match_any: bool = True) -> bool:
def is_valid_deployment_tag(
deployment_tags: Sequence[str], request_tags: Sequence[str], match_any: bool = True
) -> bool:
"""
Check if a tag is valid, the matching can be either any or all based on `match_any` flag
"""
@ -73,7 +134,7 @@ def is_valid_deployment_tag(deployment_tags: list[str], request_tags: list[str],
def _match_deployment(
deployment: Any,
deployment: _DeploymentLike,
request_tags: list[str] | None,
header_strings: list[str],
match_any: bool,
@ -90,8 +151,8 @@ def _match_deployment(
ran and failed, so the regex cannot override strict-tag policy.
"""
litellm_params: Final = deployment.get("litellm_params", {})
deployment_tags: Final[list[str] | None] = litellm_params.get("tags")
deployment_tag_regex: Final[list[str] | None] = litellm_params.get("tag_regex")
deployment_tags: Final[Sequence[str] | None] = litellm_params.get("tags")
deployment_tag_regex: Final[Sequence[str] | None] = litellm_params.get("tag_regex")
# 1. Exact tag match (existing behaviour).
if deployment_tags and request_tags:
@ -162,38 +223,38 @@ def _split_tags(tags: Sequence[str]) -> tuple[tuple[str, ...], list[str], tuple[
def _exclude_deployments(
deployments: Sequence[Any] | Mapping[Any, Any],
deployments: _DeploymentPool,
excluded_set: frozenset[str],
) -> list[Any]:
) -> 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: Sequence[Any] | Mapping[Any, Any],
deployments: _DeploymentPool,
required_set: frozenset[str],
) -> tuple[Any, ...]:
) -> 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: Sequence[Any] | Mapping[Any, Any],
) -> tuple[Any, ...]:
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: Sequence[Any] | Mapping[Any, Any]) -> frozenset[str]:
def _known_tag_values(deployments: _DeploymentPool) -> frozenset[str]:
return frozenset(
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: Sequence[Any] | Mapping[Any, Any],
healthy_deployments: _DeploymentPool,
excluded_set: frozenset[str],
required_set: frozenset[str],
routing_confirmed: frozenset[str],
@ -217,7 +278,7 @@ def _unknown_required_tag_hides_an_answer(
def _chain_allows_fail_open(
healthy_deployments: Sequence[Any] | Mapping[Any, Any],
healthy_deployments: _DeploymentPool,
excluded_set: frozenset[str],
required_set: frozenset[str],
routing_confirmed: frozenset[str],
@ -228,12 +289,12 @@ def _chain_allows_fail_open(
def _trusted_only_pool(
healthy_deployments: Sequence[Any] | Mapping[Any, Any],
healthy_deployments: _DeploymentPool,
excluded_set: frozenset[str],
required_set: frozenset[str],
inherited_excluded_set: frozenset[str] | None,
inherited_required_set: frozenset[str] | None,
) -> tuple[Any, ...]:
) -> 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
@ -260,8 +321,8 @@ def _trusted_only_pool(
def _resolve_or_fail_open(
pool: Sequence[Any],
healthy_deployments: Sequence[Any] | Mapping[Any, Any],
pool: Sequence[_DeploymentLike],
healthy_deployments: _DeploymentPool,
excluded_set: frozenset[str],
required_set: frozenset[str],
inherited_excluded_set: frozenset[str] | None,
@ -269,7 +330,7 @@ def _resolve_or_fail_open(
routing_confirmed: frozenset[str],
model: str,
request_tags: object,
) -> tuple[Any, ...]:
) -> tuple[_DeploymentLike, ...]:
if pool:
return tuple(pool)
if _chain_allows_fail_open(healthy_deployments, excluded_set, required_set, routing_confirmed):
@ -289,7 +350,7 @@ def _resolve_or_fail_open(
def _resolve_constraint_only_pool(
healthy_deployments: Sequence[Any] | Mapping[Any, Any],
healthy_deployments: _DeploymentPool,
excluded_set: frozenset[str],
required_set: frozenset[str],
inherited_excluded_set: frozenset[str] | None,
@ -297,7 +358,7 @@ def _resolve_constraint_only_pool(
routing_confirmed: frozenset[str],
model: str,
request_tags: object,
) -> tuple[Any, ...]:
) -> tuple[_DeploymentLike, ...]:
pool: Final = (
_require_all_tags(_exclude_deployments(healthy_deployments, excluded_set), required_set)
if required_set
@ -319,8 +380,8 @@ def _resolve_constraint_only_pool(
def _all_deployments_or_fallback(
llm_router_instance: LitellmRouter,
model: str,
fallback: Sequence[Any] | Mapping[Any, Any],
) -> Sequence[Any] | Mapping[Any, Any]:
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
@ -330,7 +391,7 @@ def _all_deployments_or_fallback(
def _chain_tag_filtering_override(
llm_router_instance: LitellmRouter,
model: str,
healthy_deployments: Sequence[Any] | Mapping[Any, Any],
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
@ -392,10 +453,10 @@ def _tag_known_to_group(
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.
@ -473,25 +534,25 @@ async def get_deployments_for_tag(
request_tags,
)
new_healthy_deployments: Final[list[Any]] = []
default_deployments: Final[list[Any]] = []
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",
@ -507,10 +568,10 @@ async def get_deployments_for_tag(
"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(
@ -545,10 +606,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 = []
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
@ -562,7 +624,7 @@ async def get_deployments_for_tag(
def _get_tags_from_request_kwargs(
request_kwargs: dict[Any, Any] | None = None,
request_kwargs: _RequestKwargsLike | None = None,
metadata_variable_name: Literal["metadata", "litellm_metadata"] = "metadata",
) -> list[str]:
"""
@ -577,12 +639,12 @@ def _get_tags_from_request_kwargs(
if request_kwargs is None:
return []
if metadata_variable_name in request_kwargs:
metadata: Final = request_kwargs[metadata_variable_name] or {}
metadata: Final[_MetadataLike] = request_kwargs[metadata_variable_name] or {}
tags = metadata.get("tags", [])
return tags if tags is not None else []
return list(tags) if tags is not None else []
elif "litellm_params" in request_kwargs:
litellm_params: Final = request_kwargs["litellm_params"] or {}
_metadata: Final = litellm_params.get(metadata_variable_name, {}) or {}
litellm_params: Final[_NestedLitellmParamsLike] = request_kwargs["litellm_params"] or {}
_metadata: Final[_MetadataLike] = litellm_params.get(metadata_variable_name, {}) or {}
tags = _metadata.get("tags", [])
return tags if tags is not None else []
return list(tags) if tags is not None else []
return []

16
ruff-strict-budget.json
View file

@ -1,6 +1,6 @@
{
"ANN001": {
"limit": 3058
"limit": 3036
},
"ANN002": {
"limit": 71
@ -9,7 +9,7 @@
"limit": 827
},
"ANN201": {
"limit": 2022
"limit": 2020
},
"ANN202": {
"limit": 855
@ -24,7 +24,7 @@
"limit": 133
},
"ANN401": {
"limit": 1384
"limit": 1286
},
"ASYNC230": {
"limit": 11
@ -39,7 +39,7 @@
"limit": 505
},
"B009": {
"limit": 64
"limit": 63
},
"B010": {
"limit": 190
@ -123,7 +123,7 @@
"limit": 12
},
"PERF403": {
"limit": 34
"limit": 33
},
"PIE804": {
"limit": 18
@ -180,7 +180,7 @@
"limit": 8
},
"RUF019": {
"limit": 38
"limit": 36
},
"RUF046": {
"limit": 4
@ -201,7 +201,7 @@
"limit": 58
},
"SIM102": {
"limit": 321
"limit": 318
},
"SIM103": {
"limit": 119
@ -234,7 +234,7 @@
"limit": 5
},
"TID251": {
"limit": 1224
"limit": 1214
},
"TRY002": {
"limit": 524

10
type-discipline-budget.json
View file

@ -1,9 +1,9 @@
{
"LIT001": {
"limit": 23003
"limit": 22780
},
"LIT002": {
"limit": 27146
"limit": 27144
},
"LIT003": {
"limit": 269
@ -15,7 +15,7 @@
"limit": 0
},
"LIT006": {
"limit": 1077
"limit": 1069
},
"LIT007": {
"limit": 0
@ -27,9 +27,9 @@
"limit": 0
},
"LIT010": {
"limit": 16731
"limit": 16725
},
"LIT011": {
"limit": 5596
"limit": 5590
}
}