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

chore(typing): clear basedpyright Any errors in streaming, passthrough, and proxy hot paths

Browse source
This commit is contained in:
mateo-berri 2026-08-04 04:30:21 +00:00
parent 491eda319c
commit a38ce9e585
No known key found for this signature in database
20 changed files with 2961 additions and 1799 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

34
basedpyright-code-budget.json
View file

@ -1,15 +1,15 @@
{
"reportAny": {
"limit": 29813
"limit": 28296
},
"reportArgumentType": {
"limit": 2645
"limit": 2634
},
"reportAssignmentType": {
"limit": 329
},
"reportAttributeAccessIssue": {
"limit": 516
"limit": 511
},
"reportCallIssue": {
"limit": 123
@ -24,7 +24,7 @@
"limit": 42
},
"reportExplicitAny": {
"limit": 9473
"limit": 9093
},
"reportFunctionMemberAccess": {
"limit": 11
@ -54,13 +54,13 @@
"limit": 0
},
"reportMissingParameterType": {
"limit": 5855
"limit": 5829
},
"reportMissingTypeArgument": {
"limit": 15852
"limit": 15785
},
"reportMissingTypeStubs": {
"limit": 41
"limit": 40
},
"reportOperatorIssue": {
"limit": 0
@ -72,7 +72,7 @@
"limit": 0
},
"reportOptionalMemberAccess": {
"limit": 1079
"limit": 1071
},
"reportOptionalOperand": {
"limit": 0
@ -84,13 +84,13 @@
"limit": 77
},
"reportPrivateUsage": {
"limit": 2437
"limit": 2436
},
"reportRedeclaration": {
"limit": 12
},
"reportReturnType": {
"limit": 219
"limit": 218
},
"reportTypedDictNotRequiredAccess": {
"limit": 27
@ -99,31 +99,31 @@
"limit": 0
},
"reportUnknownArgumentType": {
"limit": 45324
"limit": 45178
},
"reportUnknownLambdaType": {
"limit": 113
},
"reportUnknownMemberType": {
"limit": 40452
"limit": 40239
},
"reportUnknownParameterType": {
"limit": 20309
"limit": 20229
},
"reportUnknownVariableType": {
"limit": 31978
"limit": 31822
},
"reportUnnecessaryCast": {
"limit": 177
"limit": 172
},
"reportUnnecessaryComparison": {
"limit": 1021
"limit": 1019
},
"reportUnnecessaryContains": {
"limit": 7
},
"reportUnnecessaryIsInstance": {
"limit": 1204
"limit": 1199
},
"reportUntypedBaseClass": {
"limit": 165

452
litellm/litellm_core_utils/streaming_chunk_builder_utils.py
View file

@ -1,17 +1,24 @@
import base64
import time
from typing import TYPE_CHECKING, Any, Union, cast
from collections.abc import Sequence
from typing import TYPE_CHECKING, Literal
from pydantic import TypeAdapter
from typing_extensions import Required, TypedDict
from litellm._logging import verbose_logger
from litellm.types.llms.openai import (
AllMessageValues,
ChatCompletionAssistantContentValue,
ChatCompletionAudioDelta,
ChatCompletionRedactedThinkingBlock,
ChatCompletionThinkingBlock,
)
from litellm.types.utils import (
CacheCreationTokenDetails,
ChatCompletionAudioResponse,
ChatCompletionDeltaToolCall,
ChatCompletionMessageToolCall,
Choices,
CompletionTokensDetails,
CompletionTokensDetailsWrapper,
Function,
@ -20,57 +27,158 @@ from litellm.types.utils import (
ModelResponseStream,
PromptTokensDetailsWrapper,
ServerToolUse,
StreamingChoices,
Usage,
)
from litellm.utils import print_verbose, token_counter
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.types.litellm_core_utils.streaming_chunk_builder_utils import (
UsagePerChunk,
)
from litellm.types.llms.openai import (
ChatCompletionRedactedThinkingBlock,
ChatCompletionThinkingBlock,
)
RawHiddenParamsDict = dict[str, object]
class ChunkHiddenParamsDict(TypedDict, total=False):
created_at: int | float
custom_llm_provider: str | None
class UsageDict(TypedDict, total=False):
prompt_tokens: int
completion_tokens: int
total_tokens: int
reasoning_tokens: int | None
cache_creation_input_tokens: int | None
cache_read_input_tokens: int | None
cost: float | None
prompt_tokens_details: dict[str, object] | None
completion_tokens_details: dict[str, object] | None
server_tool_use: dict[str, object] | None
class ToolCallFunctionDict(TypedDict, total=False):
name: str | None
arguments: str
provider_specific_fields: dict[str, object] | None
class ToolCallDict(TypedDict, total=False):
id: str | None
type: str | None
index: int
function: ToolCallFunctionDict | Function | None
provider_specific_fields: dict[str, object] | None
class DeltaDict(TypedDict, total=False):
role: str | None
content: str | None
reasoning_content: str | None
function_call: FunctionCall | None
tool_calls: list[ToolCallDict | ChatCompletionDeltaToolCall] | None
audio: ChatCompletionAudioDelta | None
thinking_blocks: list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock] | None
class ChoiceDict(TypedDict, total=False):
index: int
delta: Required[DeltaDict]
finish_reason: str | None
class ChunkDict(TypedDict, total=False):
id: Required[str]
object: Required[str]
created: Required[int]
model: Required[str]
choices: Required[list[ChoiceDict | StreamingChoices]]
system_fingerprint: str | None
usage: Usage | UsageDict | None
_hidden_params: RawHiddenParamsDict
class AccumulatedToolCallDict(TypedDict):
id: str | None
name: str | None
type: str | None
arguments: list[str]
provider_specific_fields: dict[str, object] | None
class UsageChunkCalculationDict(TypedDict):
prompt_tokens: int | None
completion_tokens: int | None
cache_creation_input_tokens: int | None
cache_read_input_tokens: int | None
completion_tokens_details: CompletionTokensDetails | None
prompt_tokens_details: PromptTokensDetailsWrapper | None
cost: float | None
class CompletionTokensDetailsDumpDict(TypedDict, total=False):
accepted_prediction_tokens: int | None
audio_tokens: int | None
reasoning_tokens: int | None
rejected_prediction_tokens: int | None
text_tokens: int | None
image_tokens: int | None
video_tokens: int | None
class UsageDumpDict(TypedDict, total=False):
prompt_tokens: int
completion_tokens: int
total_tokens: int
completion_tokens_details: dict[str, object] | None
prompt_tokens_details: dict[str, object] | None
server_tool_use: dict[str, object] | None
cost: float | None
cache_creation_input_tokens: int
cache_read_input_tokens: int
_CHUNK_HIDDEN_PARAMS_ADAPTER: TypeAdapter[ChunkHiddenParamsDict] = TypeAdapter(ChunkHiddenParamsDict)
_AUDIO_DELTA_ADAPTER: TypeAdapter[ChatCompletionAudioDelta | None] = TypeAdapter(ChatCompletionAudioDelta | None)
_COMPLETION_TOKENS_DETAILS_DUMP_ADAPTER: TypeAdapter[CompletionTokensDetailsDumpDict] = TypeAdapter(
CompletionTokensDetailsDumpDict
)
_USAGE_DUMP_ADAPTER: TypeAdapter[UsageDumpDict] = TypeAdapter(UsageDumpDict)
class ChunkProcessor:
def __init__(self, chunks: list, messages: list | None = None):
def __init__(self, chunks: list[ChunkDict], messages: list[AllMessageValues] | None = None):
self.chunks = self._sort_chunks(chunks)
self.messages = messages
self.first_chunk = chunks[0]
def _sort_chunks(self, chunks: list) -> list:
@staticmethod
def _chunk_hidden_params(chunk: "ChunkDict | ModelResponseStream") -> ChunkHiddenParamsDict:
candidate: object = (
chunk.get("_hidden_params", {}) if isinstance(chunk, dict) else getattr(chunk, "_hidden_params", {})
)
if isinstance(candidate, dict):
return _CHUNK_HIDDEN_PARAMS_ADAPTER.validate_python(candidate)
return {}
@staticmethod
def _chunk_created_at(chunk: "ChunkDict | ModelResponseStream") -> int | float:
return ChunkProcessor._chunk_hidden_params(chunk).get("created_at", float("inf"))
def _sort_chunks(self, chunks: list[ChunkDict]) -> list[ChunkDict]:
if not chunks:
return []
first_chunk = chunks[0]
first_hidden_params: dict[str, Any] = {}
if isinstance(first_chunk, dict):
candidate = first_chunk.get("_hidden_params", {})
if isinstance(candidate, dict):
first_hidden_params = candidate
else:
candidate = getattr(first_chunk, "_hidden_params", {})
if isinstance(candidate, dict):
first_hidden_params = candidate
first_hidden_params = self._chunk_hidden_params(chunks[0])
if first_hidden_params.get("created_at"):
def _created_at(chunk: Any) -> int | float:
if isinstance(chunk, dict):
params = chunk.get("_hidden_params", {})
else:
params = getattr(chunk, "_hidden_params", {})
if isinstance(params, dict):
return cast(int | float, params.get("created_at", float("inf")))
return float("inf")
return sorted(chunks, key=_created_at)
return sorted(chunks, key=self._chunk_created_at)
return chunks
def update_model_response_with_hidden_params(
self, model_response: ModelResponse, chunk: dict[str, Any] | None = None
self, model_response: ModelResponse, chunk: ChunkDict | None = None
) -> ModelResponse:
if chunk is None:
return model_response
@ -82,17 +190,17 @@ class ChunkProcessor:
@staticmethod
def apply_provider_assembled_streaming_metadata(
response: ModelResponse,
chunks: list[Any],
logging_obj: Any | None = None,
chunks: list[ChunkDict],
logging_obj: "Logging | None" = None,
) -> None:
if not chunks:
return
model = getattr(response, "model", None)
model: str | None = getattr(response, "model", None)
if not model:
return
custom_llm_provider = None
custom_llm_provider: str | None = None
if logging_obj is not None:
custom_llm_provider = logging_obj.model_call_details.get("custom_llm_provider")
@ -126,7 +234,7 @@ class ChunkProcessor:
)
@staticmethod
def _get_chunk_id(chunks: list[dict[str, Any]]) -> str:
def _get_chunk_id(chunks: Sequence[ChunkDict]) -> str:
"""
Chunks:
[{"id": ""}, {"id": "1"}, {"id": "1"}]
@ -137,7 +245,7 @@ class ChunkProcessor:
return ""
@staticmethod
def _get_model_from_chunks(chunks: list[dict[str, Any]], first_chunk_model: str) -> str:
def _get_model_from_chunks(chunks: Sequence[ChunkDict], first_chunk_model: str) -> str:
"""
Get the actual model from chunks, preferring a model that differs from the first chunk.
@ -153,7 +261,7 @@ class ChunkProcessor:
# Fall back to first chunk's model if no different model found
return first_chunk_model
def build_base_response(self, chunks: list[dict[str, Any]]) -> ModelResponse:
def build_base_response(self, chunks: Sequence[ChunkDict]) -> ModelResponse:
chunk = self.first_chunk
id = ChunkProcessor._get_chunk_id(chunks)
object = chunk["object"]
@ -164,15 +272,19 @@ class ChunkProcessor:
system_fingerprint = chunk.get("system_fingerprint", None)
first_chunk_with_choices = next((c for c in chunks if c.get("choices")), chunk)
role = first_chunk_with_choices["choices"][0]["delta"]["role"]
first_choice = first_chunk_with_choices["choices"][0]
if isinstance(first_choice, dict):
role = first_choice["delta"].get("role")
else:
role = first_choice.delta.role
finish_reason = "stop"
for chunk in chunks:
if "choices" in chunk and len(chunk["choices"]) > 0:
chunk_finish_reason = None
if hasattr(chunk["choices"][0], "finish_reason"):
chunk_finish_reason = chunk["choices"][0].finish_reason
elif "finish_reason" in chunk["choices"][0]:
chunk_finish_reason = chunk["choices"][0]["finish_reason"]
finish_reason_choice = chunk["choices"][0]
if isinstance(finish_reason_choice, dict):
chunk_finish_reason = finish_reason_choice.get("finish_reason")
else:
chunk_finish_reason = finish_reason_choice.finish_reason
if chunk_finish_reason is not None:
finish_reason = chunk_finish_reason
@ -202,15 +314,18 @@ class ChunkProcessor:
response = self.update_model_response_with_hidden_params(model_response=response, chunk=chunk)
return response
def get_combined_tool_content(self, tool_call_chunks: list[dict[str, Any]]) -> list[ChatCompletionMessageToolCall]:
def get_combined_tool_content(self, tool_call_chunks: Sequence[ChunkDict]) -> list[ChatCompletionMessageToolCall]:
tool_calls_list: list[ChatCompletionMessageToolCall] = []
tool_call_map: dict[int, dict[str, Any]] = {} # Map to store tool calls by index
tool_call_map: dict[int, AccumulatedToolCallDict] = {} # Map to store tool calls by index
for chunk in tool_call_chunks:
choices = chunk["choices"]
for choice in choices:
delta = choice.get("delta", {})
tool_calls = delta.get("tool_calls", [])
tool_calls: Sequence[ToolCallDict | ChatCompletionDeltaToolCall]
if isinstance(choice, dict):
tool_calls = choice.get("delta", {}).get("tool_calls") or []
else:
tool_calls = choice.delta.tool_calls or []
for tool_call in tool_calls:
# Handle both dict and object formats
@ -222,7 +337,7 @@ class ChunkProcessor:
if isinstance(tool_call, dict):
has_function = "function" in tool_call and tool_call["function"] is not None
else:
has_function = hasattr(tool_call, "function") and tool_call.function is not None
has_function = getattr(tool_call, "function", None) is not None
if not has_function:
continue
@ -244,22 +359,26 @@ class ChunkProcessor:
# Extract id, type, and function data (handle both dict and object)
if isinstance(tool_call, dict):
if tool_call.get("id"):
tool_call_map[index]["id"] = tool_call["id"]
if tool_call.get("type"):
tool_call_map[index]["type"] = tool_call["type"]
tool_call_id = tool_call.get("id")
if tool_call_id:
tool_call_map[index]["id"] = tool_call_id
tool_call_type = tool_call.get("type")
if tool_call_type:
tool_call_map[index]["type"] = tool_call_type
function = tool_call.get("function", {})
if isinstance(function, dict):
if function.get("name"):
tool_call_map[index]["name"] = function["name"]
if function.get("arguments"):
tool_call_map[index]["arguments"].append(function["arguments"])
function_name = function.get("name")
if function_name:
tool_call_map[index]["name"] = function_name
function_arguments = function.get("arguments")
if function_arguments:
tool_call_map[index]["arguments"].append(function_arguments)
else:
# function is an object
if hasattr(function, "name") and function.name:
if function is not None and hasattr(function, "name") and function.name:
tool_call_map[index]["name"] = function.name
if hasattr(function, "arguments") and function.arguments:
if function is not None and hasattr(function, "arguments") and function.arguments:
tool_call_map[index]["arguments"].append(function.arguments)
else:
# tool_call is an object
@ -274,27 +393,24 @@ class ChunkProcessor:
tool_call_map[index]["arguments"].append(tool_call.function.arguments)
# Preserve provider_specific_fields from streaming chunks
provider_fields = None
provider_fields: dict[str, object] | None = None
if isinstance(tool_call, dict):
provider_fields = tool_call.get("provider_specific_fields")
if not provider_fields and isinstance(tool_call.get("function"), dict):
provider_fields = tool_call["function"].get("provider_specific_fields")
if not provider_fields:
function_value = tool_call.get("function")
if isinstance(function_value, dict):
provider_fields = function_value.get("provider_specific_fields")
else:
if hasattr(tool_call, "provider_specific_fields") and tool_call.provider_specific_fields:
provider_fields = tool_call.provider_specific_fields
elif (
hasattr(tool_call, "function")
and hasattr(tool_call.function, "provider_specific_fields")
and tool_call.function.provider_specific_fields
):
provider_fields = tool_call.function.provider_specific_fields
object_provider_fields = getattr(tool_call, "provider_specific_fields", None) or getattr(
getattr(tool_call, "function", None), "provider_specific_fields", None
)
if isinstance(object_provider_fields, dict) and object_provider_fields:
provider_fields = object_provider_fields
if provider_fields:
if isinstance(provider_fields, dict) and provider_fields:
# Merge provider_specific_fields if multiple chunks have them
if tool_call_map[index]["provider_specific_fields"] is None:
tool_call_map[index]["provider_specific_fields"] = {}
if isinstance(provider_fields, dict):
tool_call_map[index]["provider_specific_fields"].update(provider_fields)
existing_fields = tool_call_map[index]["provider_specific_fields"] or {}
tool_call_map[index]["provider_specific_fields"] = {**existing_fields, **provider_fields}
# Convert the map to a list of tool calls
for index in sorted(tool_call_map.keys()):
@ -308,39 +424,48 @@ class ChunkProcessor:
name=tool_call_data["name"],
)
# Prepare params for ChatCompletionMessageToolCall
tool_call_params = {
"id": tool_call_data["id"],
"function": function,
"type": tool_call_data["type"] or "function",
}
# Add provider_specific_fields if present (for thought signatures in Gemini 3)
if tool_call_data.get("provider_specific_fields"):
tool_call_params["provider_specific_fields"] = tool_call_data["provider_specific_fields"]
tool_call = ChatCompletionMessageToolCall(**tool_call_params)
tool_calls_list.append(tool_call)
provider_specific_fields = tool_call_data["provider_specific_fields"]
if provider_specific_fields:
tool_calls_list.append(
ChatCompletionMessageToolCall(
id=tool_call_data["id"],
function=function,
type=tool_call_data["type"] or "function",
provider_specific_fields=provider_specific_fields,
)
)
else:
tool_calls_list.append(
ChatCompletionMessageToolCall(
id=tool_call_data["id"],
function=function,
type=tool_call_data["type"] or "function",
)
)
return tool_calls_list
def get_combined_function_call_content(self, function_call_chunks: list[dict[str, Any]]) -> FunctionCall:
argument_list = []
delta = function_call_chunks[0]["choices"][0]["delta"]
function_call = delta.get("function_call", "")
function_call_name = function_call.name
def get_combined_function_call_content(self, function_call_chunks: Sequence[ChunkDict]) -> FunctionCall:
argument_list: list[str] = []
first_choice = function_call_chunks[0]["choices"][0]
if isinstance(first_choice, dict):
first_function_call = first_choice["delta"].get("function_call")
else:
first_function_call = getattr(first_choice.delta, "function_call", None)
function_call_name = first_function_call.name
for chunk in function_call_chunks:
choices = chunk["choices"]
for choice in choices:
delta = choice.get("delta", {})
function_call = delta.get("function_call", "")
if isinstance(choice, dict):
function_call = choice.get("delta", {}).get("function_call")
else:
function_call = getattr(choice.delta, "function_call", None)
# Check if a function call is present
if function_call:
# Now, function_call is expected to be a dictionary
arguments = function_call.arguments
argument_list.append(arguments)
argument_list.append(function_call.arguments)
combined_arguments = "".join(argument_list)
@ -350,14 +475,16 @@ class ChunkProcessor:
)
def get_combined_content(
self, chunks: list[dict[str, Any]], delta_key: str = "content"
self, chunks: Sequence[ChunkDict], delta_key: Literal["content", "reasoning_content"] = "content"
) -> ChatCompletionAssistantContentValue:
content_list: list[str] = []
for chunk in chunks:
choices = chunk["choices"]
for choice in choices:
delta = choice.get("delta", {})
content = delta.get(delta_key, "")
if isinstance(choice, dict):
content = choice.get("delta", {}).get(delta_key, "")
else:
content = getattr(choice.delta, delta_key, "")
if content is None:
continue # openai v1.0.0 sets content = None for chunks
content_list.append(content)
@ -369,13 +496,8 @@ class ChunkProcessor:
return combined_content
def get_combined_thinking_content(
self, chunks: list[dict[str, Any]]
) -> list[Union["ChatCompletionThinkingBlock", "ChatCompletionRedactedThinkingBlock"]] | None:
from litellm.types.llms.openai import (
ChatCompletionRedactedThinkingBlock,
ChatCompletionThinkingBlock,
)
self, chunks: Sequence[ChunkDict]
) -> list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock] | None:
thinking_blocks: list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock] = []
current_thinking_text_parts: list[str] = []
current_signature: str | None = None
@ -396,12 +518,13 @@ class ChunkProcessor:
for chunk in chunks:
choices = chunk["choices"]
for choice in choices:
delta = choice.get("delta", {})
thinking = delta.get("thinking_blocks", None)
if thinking and isinstance(thinking, list):
if isinstance(choice, dict):
thinking = choice.get("delta", {}).get("thinking_blocks", None)
else:
thinking = getattr(choice.delta, "thinking_blocks", None)
if isinstance(thinking, list) and thinking:
for thinking_block in thinking:
thinking_type = thinking_block.get("type", None)
if thinking_type and thinking_type == "redacted_thinking":
if thinking_block.get("type", None) == "redacted_thinking":
_flush_thinking_block()
redacted_data = thinking_block.get("data", None)
if redacted_data:
@ -426,10 +549,10 @@ class ChunkProcessor:
return thinking_blocks
return None
def get_combined_reasoning_content(self, chunks: list[dict[str, Any]]) -> ChatCompletionAssistantContentValue:
def get_combined_reasoning_content(self, chunks: Sequence[ChunkDict]) -> ChatCompletionAssistantContentValue:
return self.get_combined_content(chunks, delta_key="reasoning_content")
def get_combined_audio_content(self, chunks: list[dict[str, Any]]) -> ChatCompletionAudioResponse:
def get_combined_audio_content(self, chunks: Sequence[ChunkDict]) -> ChatCompletionAudioResponse:
base64_data_list: list[str] = []
transcript_list: list[str] = []
expires_at: int | None = None
@ -438,8 +561,12 @@ class ChunkProcessor:
for chunk in chunks:
choices = chunk["choices"]
for choice in choices:
delta = choice.get("delta") or {}
audio: ChatCompletionAudioDelta | None = delta.get("audio")
audio: ChatCompletionAudioDelta | None
if isinstance(choice, dict):
delta: DeltaDict = choice.get("delta") or {}
audio = delta.get("audio")
else:
audio = _AUDIO_DELTA_ADAPTER.validate_python(getattr(choice.delta, "audio", None))
if audio is not None:
for k, v in audio.items():
if k == "data" and v is not None and isinstance(v, str):
@ -459,9 +586,9 @@ class ChunkProcessor:
id=id,
)
def _usage_chunk_calculation_helper(self, usage_chunk: Usage) -> dict:
prompt_tokens = 0
completion_tokens = 0
def _usage_chunk_calculation_helper(self, usage_chunk: Usage) -> UsageChunkCalculationDict:
prompt_tokens: int = 0
completion_tokens: int = 0
## anthropic prompt caching information ##
cache_creation_input_tokens: int | None = None
cache_read_input_tokens: int | None = None
@ -503,30 +630,30 @@ class ChunkProcessor:
def count_reasoning_tokens(self, response: ModelResponse) -> int | None:
reasoning_tokens: int | None = None
for choice in response.choices:
if (
hasattr(cast(Choices, choice).message, "reasoning_content")
and cast(Choices, choice).message.reasoning_content is not None
):
if hasattr(choice.message, "reasoning_content") and choice.message.reasoning_content is not None:
if reasoning_tokens is None:
reasoning_tokens = 0
reasoning_tokens += token_counter(
text=cast(Choices, choice).message.reasoning_content,
text=choice.message.reasoning_content,
count_response_tokens=True,
)
return reasoning_tokens
@staticmethod
def _extract_usage_chunk(chunk: dict[str, Any] | ModelResponse | ModelResponseStream) -> Usage | None:
usage_chunk: Usage | dict[str, Any] | None = None
if hasattr(chunk, "usage") and chunk.usage is not None:
usage_chunk = chunk.usage
elif "usage" in chunk:
usage_chunk = chunk["usage"]
elif (isinstance(chunk, ModelResponse) or isinstance(chunk, ModelResponseStream)) and hasattr(
chunk, "_hidden_params"
):
usage_chunk = chunk._hidden_params.get("usage", None)
def _extract_usage_chunk(chunk: ChunkDict | ModelResponse | ModelResponseStream) -> Usage | None:
usage_chunk: Usage | UsageDict | None = None
if isinstance(chunk, dict):
usage_chunk = chunk.get("usage")
elif hasattr(chunk, "usage"):
usage_chunk = getattr(chunk, "usage", None)
elif hasattr(chunk, "_hidden_params"):
hidden_usage: object = chunk._hidden_params.get("usage", None)
if isinstance(hidden_usage, Usage):
return hidden_usage
if isinstance(hidden_usage, dict):
return Usage(**hidden_usage)
return None
if isinstance(usage_chunk, dict):
return Usage(**usage_chunk)
@ -534,7 +661,7 @@ class ChunkProcessor:
def _calculate_usage_per_chunk(
self,
chunks: list[dict[str, Any] | ModelResponse],
chunks: Sequence[ChunkDict],
) -> "UsagePerChunk":
from litellm.types.litellm_core_utils.streaming_chunk_builder_utils import (
UsagePerChunk,
@ -601,24 +728,15 @@ class ChunkProcessor:
server_tool_use = usage_chunk.server_tool_use
else:
server_tool_use = ServerToolUse.model_validate(usage_chunk.server_tool_use)
if (
usage_chunk_dict["prompt_tokens_details"] is not None
and getattr(
usage_chunk_dict["prompt_tokens_details"],
"web_search_requests",
None,
)
is not None
):
web_search_requests = getattr(
usage_chunk_dict["prompt_tokens_details"],
"web_search_requests",
)
prompt_tokens_details = cast(
PromptTokensDetailsWrapper | None,
web_search_requests_value: int | None = getattr(
usage_chunk_dict["prompt_tokens_details"],
"web_search_requests",
None,
)
if usage_chunk_dict["prompt_tokens_details"] is not None and web_search_requests_value is not None:
web_search_requests = web_search_requests_value
prompt_tokens_details = usage_chunk_dict["prompt_tokens_details"]
cache_creation_token_details = self._capture_cache_creation_token_details(
prompt_tokens_details, cache_creation_token_details
@ -654,9 +772,10 @@ class ChunkProcessor:
prompt_tokens_details: PromptTokensDetailsWrapper | None,
current: CacheCreationTokenDetails | None,
) -> CacheCreationTokenDetails | None:
incoming = cast(
CacheCreationTokenDetails | None,
getattr(prompt_tokens_details, "cache_creation_token_details", None),
if prompt_tokens_details is None:
return current
incoming: CacheCreationTokenDetails | None = getattr(
prompt_tokens_details, "cache_creation_token_details", None
)
if incoming is not None:
return incoming
@ -669,9 +788,8 @@ class ChunkProcessor:
) -> PromptTokensDetailsWrapper | None:
if prompt_tokens_details is None or cache_creation_token_details is None:
return prompt_tokens_details
existing = cast(
CacheCreationTokenDetails | None,
getattr(prompt_tokens_details, "cache_creation_token_details", None),
existing: CacheCreationTokenDetails | None = getattr(
prompt_tokens_details, "cache_creation_token_details", None
)
if existing is not None:
return prompt_tokens_details
@ -679,7 +797,7 @@ class ChunkProcessor:
@staticmethod
def _reset_anthropic_cursor_completion_tokens(
chunks: list[dict[str, Any] | ModelResponse],
chunks: Sequence[ChunkDict],
completion_tokens: int,
completion_usage_updates: int,
) -> int:
@ -704,13 +822,7 @@ class ChunkProcessor:
custom_llm_provider: str | None = None
if chunks:
first_chunk = chunks[0]
if isinstance(first_chunk, dict):
hp = first_chunk.get("_hidden_params")
else:
hp = getattr(first_chunk, "_hidden_params", None)
if isinstance(hp, dict):
custom_llm_provider = hp.get("custom_llm_provider")
custom_llm_provider = ChunkProcessor._chunk_hidden_params(chunks[0]).get("custom_llm_provider")
if custom_llm_provider == "anthropic" and completion_tokens == 1:
return 0
@ -718,10 +830,10 @@ class ChunkProcessor:
def calculate_usage(
self,
chunks: list[dict[str, Any] | ModelResponse],
chunks: Sequence[ChunkDict],
model: str,
completion_output: str,
messages: list | None = None,
messages: list[AllMessageValues] | None = None,
reasoning_tokens: int | None = None,
) -> Usage:
"""
@ -773,7 +885,7 @@ class ChunkProcessor:
if completion_tokens_details is not None:
if isinstance(completion_tokens_details, CompletionTokensDetails):
returned_usage.completion_tokens_details = CompletionTokensDetailsWrapper(
**completion_tokens_details.model_dump()
**_COMPLETION_TOKENS_DETAILS_DUMP_ADAPTER.validate_python(completion_tokens_details.model_dump())
)
else:
returned_usage.completion_tokens_details = completion_tokens_details
@ -806,12 +918,12 @@ class ChunkProcessor:
# Return a new usage object with the new values
returned_usage = Usage(**returned_usage.model_dump())
returned_usage = Usage(**_USAGE_DUMP_ADAPTER.validate_python(returned_usage.model_dump()))
return returned_usage
def concatenate_base64_list(base64_strings: list[str]) -> str:
def concatenate_base64_list(base64_strings: Sequence[str]) -> str:
"""
Concatenates a list of base64-encoded strings.

640
litellm/litellm_core_utils/streaming_handler.py
View file

File diff suppressed because it is too large Load diff

233
litellm/litellm_core_utils/streaming_handler_types.py Normal file
View file

@ -0,0 +1,233 @@
from collections.abc import Mapping, Sequence
from typing import Protocol, Union, runtime_checkable
from pydantic import BaseModel
from typing_extensions import ReadOnly, Required, TypedDict
from litellm.types.llms.openai import (
ChatCompletionToolCallChunk,
ChatCompletionToolCallFunctionChunk,
)
from litellm.types.utils import (
ChatCompletionDeltaToolCall,
Function,
ModelResponseStream,
)
class CompletionUsageLike(Protocol):
prompt_tokens: int
completion_tokens: int
total_tokens: int
class ParsedProviderChunk(TypedDict, total=False):
text: ReadOnly[str | None]
is_finished: ReadOnly[bool]
finish_reason: ReadOnly[str | None]
logprobs: ReadOnly[object]
original_chunk: ReadOnly[ModelResponseStream | None]
usage: ReadOnly[object]
tool_use: ReadOnly[ChatCompletionToolCallChunk | None]
tool_calls: ReadOnly[list[ChatCompletionDeltaToolCall] | None]
provider_specific_fields: ReadOnly[dict[str, object] | None]
prompt_tokens: ReadOnly[int]
completion_tokens: ReadOnly[int]
index: ReadOnly[int]
class SimpleParsedChunk(ParsedProviderChunk):
text: Required[ReadOnly[str | None]]
is_finished: Required[ReadOnly[bool]]
finish_reason: Required[ReadOnly[str | None]]
class TextCompletionParsedChunk(SimpleParsedChunk):
usage: Required[ReadOnly[CompletionUsageLike | None]]
class OpenAIChatParsedChunk(SimpleParsedChunk):
logprobs: Required[ReadOnly[object]]
original_chunk: Required[ReadOnly[ModelResponseStream | None]]
usage: Required[ReadOnly[object]]
class StreamingCompletionObj(TypedDict, total=False):
content: Required[str | None]
role: str
tool_calls: list[ChatCompletionToolCallChunk] | list[ChatCompletionDeltaToolCall] | None
function_call: ChatCompletionToolCallFunctionChunk | Function | None
provider_specific_fields: dict[str, object] | None
index: int
class TextCompletionChoiceLike(Protocol):
@property
def text(self) -> str: ...
@property
def finish_reason(self) -> str | None: ...
@runtime_checkable
class HasTextChoices(Protocol):
@property
def choices(self) -> Sequence[TextCompletionChoiceLike]: ...
@runtime_checkable
class HasCompletionUsage(Protocol):
usage: CompletionUsageLike | None
@runtime_checkable
class HasFunctionsAttr(Protocol):
functions: object
@runtime_checkable
class HasStatusCode(Protocol):
@property
def status_code(self) -> int: ...
VertexArgValue = Union[
str,
int,
float,
bool,
None,
Sequence["VertexArgValue"],
Mapping[str, "VertexArgValue"],
]
class VertexFunctionCallLike(Protocol):
@property
def args(self) -> Mapping[str, VertexArgValue]: ...
@property
def name(self) -> str: ...
class VertexPartLike(Protocol):
@property
def function_call(self) -> VertexFunctionCallLike: ...
class VertexContentLike(Protocol):
@property
def parts(self) -> Sequence[VertexPartLike]: ...
class VertexEnumNameLike(Protocol):
@property
def name(self) -> str: ...
class VertexCandidateLike(Protocol):
@property
def content(self) -> VertexContentLike: ...
@property
def finish_reason(self) -> VertexEnumNameLike: ...
@runtime_checkable
class VertexProtoChunkLike(Protocol):
@property
def candidates(self) -> Sequence[VertexCandidateLike]: ...
class VllmCompletionOutputLike(Protocol):
@property
def text(self) -> str: ...
@runtime_checkable
class VllmRequestOutputLike(Protocol):
@property
def outputs(self) -> Sequence[VllmCompletionOutputLike]: ...
@runtime_checkable
class HasModelAttr(Protocol):
model: str | None
class PredibaseStreamToken(BaseModel):
text: str | None = None
class PredibaseStreamDetails(BaseModel):
finish_reason: str | None = None
class PredibaseStreamChunk(BaseModel):
token: PredibaseStreamToken | None = None
details: PredibaseStreamDetails | None = None
generated_text: str | None = None
error: str | None = None
class AI21StreamData(BaseModel):
text: str
class AI21StreamCompletion(BaseModel):
data: AI21StreamData
class AI21StreamChunk(BaseModel):
completions: list[AI21StreamCompletion]
class MaritalkStreamChunk(BaseModel):
answer: str
class NlpCloudStreamChunk(BaseModel):
generated_text: str
class AlephAlphaStreamCompletion(BaseModel):
completion: str
class AlephAlphaStreamChunk(BaseModel):
completions: list[AlephAlphaStreamCompletion]
class AzureStreamDelta(BaseModel):
content: str | None = ""
class AzureStreamChoice(BaseModel):
delta: AzureStreamDelta | None = None
finish_reason: str | None = None
class AzureStreamChunk(BaseModel):
choices: list[AzureStreamChoice]
class BasetenStreamToken(BaseModel):
text: str | None = None
class BasetenModelOutput(BaseModel):
data: list[str] | None = None
class BasetenStreamChunk(BaseModel):
token: BasetenStreamToken | None = None
model_output: BasetenModelOutput | str | None = None
completion: str | None = None
class TritonStreamChunk(BaseModel):
text_output: str | None = ""
stop_reason: str | None = None
is_finished: bool = False
input_token_count: int = 0
generated_token_count: int = 0

177
litellm/llms/anthropic/experimental_pass_through/adapters/handler.py
View file

@ -1,9 +1,11 @@
from collections.abc import AsyncIterator, Coroutine, Iterator
from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence
from typing import (
Any,
cast,
)
from pydantic import TypeAdapter, ValidationError
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.asyncify import run_async_function
@ -27,8 +29,27 @@ from litellm.utils import get_model_info
# Anthropic-only keys already mapped by the translator; strip on extra_kwargs re-merge.
ANTHROPIC_ONLY_REQUEST_KEYS: frozenset[str] = frozenset({"output_config"})
ContextManagementSpec = dict[str, object] | list[dict[str, object]] | None
def _messages_have_compaction_block(messages: list[dict]) -> bool:
_CONTEXT_MANAGEMENT_SPEC_ADAPTER: TypeAdapter[ContextManagementSpec] = TypeAdapter(ContextManagementSpec)
def _validate_context_management_spec(raw: object) -> ContextManagementSpec:
"""Validate the raw ``context_management`` request value into a typed spec.
Shapes the dispatcher cannot use (non-dict scalars, lists with non-dict
entries, non-string keys) normalize to ``None``, matching the dispatcher's
own isinstance-based rejection of malformed specs.
"""
if raw is None:
return None
try:
return _CONTEXT_MANAGEMENT_SPEC_ADAPTER.validate_python(raw, strict=True)
except ValidationError:
return None
def _messages_have_compaction_block(messages: Sequence[Mapping[str, object]]) -> bool:
"""Return True when any message carries a ``compaction`` content block."""
for msg in messages:
content = msg.get("content")
@ -40,7 +61,7 @@ def _messages_have_compaction_block(messages: list[dict]) -> bool:
return False
def _extract_proxy_litellm_metadata(kwargs: dict[str, Any]) -> dict[str, Any] | None:
def _extract_proxy_litellm_metadata(kwargs: Mapping[str, object]) -> dict[str, object] | None:
"""Return ``kwargs["litellm_metadata"]`` when it's a dict; ``None`` otherwise.
The proxy attaches its auth/spend-attribution fields (``user_api_key``,
@ -61,14 +82,14 @@ def _extract_proxy_litellm_metadata(kwargs: dict[str, Any]) -> dict[str, Any] |
async def _prepare_context_managed_request(
*,
model: str,
messages: list[dict],
tools: list[dict] | None,
system: Any | None,
context_management_spec: Any,
litellm_metadata: dict | None,
additional_drop_params: list[str] | None,
llm_router: Any,
user_api_key_auth: Any = None,
messages: list[dict[str, object]],
tools: list[dict[str, object]] | None,
system: str | list[dict[str, object]] | None,
context_management_spec: ContextManagementSpec,
litellm_metadata: dict[str, object] | None,
additional_drop_params: Sequence[str] | None,
llm_router: object,
user_api_key_auth: object = None,
) -> PolyfillResult | None:
"""Apply client compaction history, then optional context_management polyfill."""
from litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact import (
@ -88,11 +109,11 @@ async def _prepare_context_managed_request(
if polyfill_will_run:
history_result: PolyfillResult | None = None
working_messages: list[dict] = messages
working_system: Any | None = system
working_messages: list[dict[str, object]] = messages
working_system: str | list[dict[str, object]] | None = system
else:
history_result = apply_client_compaction_block_history(
messages=cast(list[dict[str, Any]], messages),
messages=messages,
system=system,
)
working_messages = history_result.messages if history_result is not None else messages
@ -122,7 +143,7 @@ async def _prepare_context_managed_request(
# to non-Anthropic backends that would reject them.
if polyfill_will_run and history_result is None:
history_result = apply_client_compaction_block_history(
messages=cast(list[dict[str, Any]], messages),
messages=messages,
system=system,
)
return history_result
@ -130,8 +151,8 @@ async def _prepare_context_managed_request(
def _polyfill_will_run(
*,
context_management_spec: Any,
additional_drop_params: list[str] | None,
context_management_spec: ContextManagementSpec,
additional_drop_params: Sequence[str] | None,
) -> bool:
"""Return True when ``compact_20260112`` will run via the polyfill dispatcher.
@ -157,8 +178,8 @@ def _polyfill_will_run(
def _spec_has_non_compact_edits(
*,
context_management_spec: Any,
additional_drop_params: list[str] | None,
context_management_spec: ContextManagementSpec,
additional_drop_params: Sequence[str] | None,
) -> bool:
"""Return True when the spec includes edits other than ``compact_20260112``.
@ -184,7 +205,7 @@ def _spec_has_non_compact_edits(
)
def _context_management_explicitly_dropped(additional_drop_params: list[str] | None) -> bool:
def _context_management_explicitly_dropped(additional_drop_params: Sequence[str] | None) -> bool:
"""True when the caller opted out of context_management via ``additional_drop_params``.
``drop_params`` deliberately does NOT gate the polyfill: ``context_management``
@ -198,9 +219,9 @@ def _context_management_explicitly_dropped(additional_drop_params: list[str] | N
def _normalize_spec_edits(
*,
context_management_spec: Any,
additional_drop_params: list[str] | None,
) -> list[dict[str, Any]] | None:
context_management_spec: ContextManagementSpec,
additional_drop_params: Sequence[str] | None,
) -> Sequence[Mapping[str, object]] | None:
"""Return the normalized ``edits`` list, or ``None`` if the polyfill won't run.
Delegates spec-shape normalization to the dispatcher's ``_normalize_spec``
@ -225,14 +246,14 @@ def _normalize_spec_edits(
async def _run_polyfill_if_enabled(
*,
model: str,
messages: list[dict],
tools: list[dict] | None,
system: Any | None,
context_management_spec: Any,
litellm_metadata: dict | None,
additional_drop_params: list[str] | None,
llm_router: Any,
user_api_key_auth: Any = None,
messages: list[dict[str, object]],
tools: list[dict[str, object]] | None,
system: str | list[dict[str, object]] | None,
context_management_spec: ContextManagementSpec,
litellm_metadata: dict[str, object] | None,
additional_drop_params: Sequence[str] | None,
llm_router: object,
user_api_key_auth: object = None,
) -> PolyfillResult | None:
"""Run the async context_management polyfill if a spec is present.
@ -293,9 +314,9 @@ ANTHROPIC_ADAPTER = AnthropicAdapter()
class LiteLLMMessagesToCompletionTransformationHandler:
@staticmethod
def _route_openai_thinking_to_responses_api_if_needed(
completion_kwargs: dict[str, Any],
completion_kwargs: dict[str, object],
*,
thinking: dict[str, Any] | None,
thinking: Mapping[str, object] | None,
) -> None:
"""
When users call `litellm.anthropic.messages.*` with a non-Anthropic model and
@ -308,8 +329,9 @@ class LiteLLMMessagesToCompletionTransformationHandler:
If the user provides a `summary` field in the thinking dict, it is passed
through to the OpenAI reasoning params (opt-in per OpenAI spec).
"""
custom_llm_provider = completion_kwargs.get("custom_llm_provider")
if custom_llm_provider is None:
raw_provider = completion_kwargs.get("custom_llm_provider")
custom_llm_provider: str | None
if raw_provider is None:
try:
_, inferred_provider, _, _ = litellm.utils.get_llm_provider(
model=cast(str, completion_kwargs.get("model"))
@ -317,6 +339,8 @@ class LiteLLMMessagesToCompletionTransformationHandler:
custom_llm_provider = inferred_provider
except Exception:
custom_llm_provider = None
else:
custom_llm_provider = raw_provider if isinstance(raw_provider, str) else None
if custom_llm_provider != "openai":
return
@ -342,7 +366,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
reasoning_effort = completion_kwargs.get("reasoning_effort")
summary = thinking.get("summary")
if isinstance(reasoning_effort, str) and reasoning_effort:
reasoning_dict: dict[str, Any] = {"effort": reasoning_effort}
reasoning_dict: dict[str, object] = {"effort": reasoning_effort}
if summary:
reasoning_dict["summary"] = summary
elif auto_summary:
@ -358,7 +382,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
@staticmethod
def _normalize_reasoning_effort(
completion_kwargs: dict[str, Any],
completion_kwargs: dict[str, object],
) -> None:
"""
Normalize reasoning_effort values based on target model capabilities.
@ -375,7 +399,8 @@ class LiteLLMMessagesToCompletionTransformationHandler:
return
model = cast(str, completion_kwargs.get("model", ""))
custom_llm_provider = completion_kwargs.get("custom_llm_provider")
raw_provider = completion_kwargs.get("custom_llm_provider")
custom_llm_provider = raw_provider if isinstance(raw_provider, str) else None
if isinstance(reasoning_effort, str):
normalized = normalize_reasoning_effort_value(
@ -396,21 +421,21 @@ class LiteLLMMessagesToCompletionTransformationHandler:
def _prepare_completion_kwargs(
*,
max_tokens: int,
messages: list[dict],
messages: Sequence[Mapping[str, object]],
model: str,
metadata: dict | None = None,
stop_sequences: list[str] | None = None,
metadata: Mapping[str, object] | None = None,
stop_sequences: Sequence[str] | None = None,
stream: bool | None = False,
system: str | list[dict[str, Any]] | None = None,
system: str | Sequence[Mapping[str, object]] | None = None,
temperature: float | None = None,
thinking: dict | None = None,
tool_choice: dict | None = None,
tools: list[dict] | None = None,
thinking: Mapping[str, object] | None = None,
tool_choice: Mapping[str, object] | None = None,
tools: Sequence[Mapping[str, object]] | None = None,
top_k: int | None = None,
top_p: float | None = None,
output_format: dict | None = None,
extra_kwargs: dict[str, Any] | None = None,
) -> tuple[dict[str, Any], dict[str, str]]:
output_format: Mapping[str, object] | None = None,
extra_kwargs: Mapping[str, object] | None = None,
) -> tuple[Mapping[str, Any], dict[str, str]]: # any-ok: kwargs spread into litellm.completion's open params
"""Prepare kwargs for litellm.completion/acompletion.
Returns:
@ -422,7 +447,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
Logging as LiteLLMLoggingObject,
)
request_data = {
request_data: dict[str, object] = {
"model": model,
"messages": messages,
"max_tokens": max_tokens,
@ -467,7 +492,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
if openai_request is None:
raise ValueError("Failed to translate request to OpenAI format")
completion_kwargs: dict[str, Any] = dict(openai_request)
completion_kwargs: dict[str, Any] = dict(openai_request) # any-ok: heterogeneous per-provider kwargs
if stream:
completion_kwargs["stream"] = stream
@ -517,24 +542,25 @@ class LiteLLMMessagesToCompletionTransformationHandler:
@staticmethod
async def async_anthropic_messages_handler(
max_tokens: int,
messages: list[dict],
messages: list[dict[str, object]],
model: str,
metadata: dict | None = None,
stop_sequences: list[str] | None = None,
metadata: Mapping[str, object] | None = None,
stop_sequences: Sequence[str] | None = None,
stream: bool | None = False,
system: str | None = None,
temperature: float | None = None,
thinking: dict | None = None,
tool_choice: dict | None = None,
tools: list[dict] | None = None,
thinking: Mapping[str, object] | None = None,
tool_choice: Mapping[str, object] | None = None,
tools: list[dict[str, object]] | None = None,
top_k: int | None = None,
top_p: float | None = None,
output_format: dict | None = None,
**kwargs,
) -> AnthropicMessagesResponse | AsyncIterator[Any] | Iterator[bytes]:
output_format: Mapping[str, object] | None = None,
**kwargs: object,
) -> AnthropicMessagesResponse | AsyncIterator[bytes] | Iterator[bytes]:
"""Handle non-Anthropic models asynchronously using the adapter"""
context_management = kwargs.pop("context_management", None)
additional_drop_params: list[str] | None = kwargs.get("additional_drop_params", None)
context_management = _validate_context_management_spec(kwargs.pop("context_management", None))
raw_drop_params = kwargs.get("additional_drop_params", None)
additional_drop_params: list[str] | None = raw_drop_params if isinstance(raw_drop_params, list) else None
litellm_router = kwargs.pop("litellm_router", None)
if litellm_router is None:
try:
@ -585,7 +611,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
extra_kwargs=kwargs,
)
completion_response = await litellm.acompletion(**completion_kwargs)
completion_response = await litellm.acompletion(**completion_kwargs) # any-ok: open per-provider kwargs
if stream:
transformed_stream = ANTHROPIC_ADAPTER.translate_completion_output_params_streaming(
@ -611,26 +637,26 @@ class LiteLLMMessagesToCompletionTransformationHandler:
@staticmethod
def anthropic_messages_handler(
max_tokens: int,
messages: list[dict],
messages: list[dict[str, object]],
model: str,
metadata: dict | None = None,
stop_sequences: list[str] | None = None,
metadata: Mapping[str, object] | None = None,
stop_sequences: Sequence[str] | None = None,
stream: bool | None = False,
system: str | None = None,
temperature: float | None = None,
thinking: dict | None = None,
tool_choice: dict | None = None,
tools: list[dict] | None = None,
thinking: Mapping[str, object] | None = None,
tool_choice: Mapping[str, object] | None = None,
tools: list[dict[str, object]] | None = None,
top_k: int | None = None,
top_p: float | None = None,
output_format: dict | None = None,
output_format: Mapping[str, object] | None = None,
_is_async: bool = False,
**kwargs,
**kwargs: object,
) -> (
AnthropicMessagesResponse
| Iterator[bytes]
| AsyncIterator[Any]
| Coroutine[Any, Any, AnthropicMessagesResponse | AsyncIterator[Any] | Iterator[bytes]]
| AsyncIterator[bytes]
| Coroutine[None, None, AnthropicMessagesResponse | AsyncIterator[bytes] | Iterator[bytes]]
):
"""Handle non-Anthropic models using the adapter."""
if _is_async is True:
@ -657,8 +683,9 @@ class LiteLLMMessagesToCompletionTransformationHandler:
# ``clear_tool_uses_20250919``. The dispatcher is async (so the
# ``compact_20260112`` editor can ``await`` the summarization model);
# bridge to it via ``run_async_function``.
context_management = kwargs.pop("context_management", None)
additional_drop_params: list[str] | None = kwargs.get("additional_drop_params", None)
context_management = _validate_context_management_spec(kwargs.pop("context_management", None))
raw_drop_params = kwargs.get("additional_drop_params", None)
additional_drop_params: list[str] | None = raw_drop_params if isinstance(raw_drop_params, list) else None
# Deliberately do NOT auto-attach the proxy ``llm_router`` here:
# ``run_async_function`` spawns a new event loop in a worker thread
# to bridge to the async dispatcher, but the proxy router's httpx
@ -722,7 +749,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
extra_kwargs=kwargs,
)
completion_response = litellm.completion(**completion_kwargs)
completion_response = litellm.completion(**completion_kwargs) # any-ok: open per-provider kwargs
if stream:
transformed_stream = ANTHROPIC_ADAPTER.translate_completion_output_params_streaming(

277
litellm/proxy/common_request_processing.py
View file

@ -4,13 +4,17 @@ import logging
import math
import time
import traceback
from collections.abc import AsyncGenerator, Callable, Mapping
from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Coroutine, Mapping
from datetime import datetime
from functools import lru_cache
from typing import (
TYPE_CHECKING,
Any,
Literal,
Protocol,
TypedDict,
TypeVar,
runtime_checkable,
)
import anyio
@ -18,6 +22,7 @@ import httpx
import orjson
from fastapi import HTTPException, Request, status
from fastapi.responses import JSONResponse, Response, StreamingResponse
from pydantic import TypeAdapter
from starlette.types import Receive, Scope, Send
import litellm
@ -58,10 +63,76 @@ from litellm.types.router import RouterRateLimitError
from litellm.types.utils import ServerToolUse
# Type alias for streaming chunk serializer (chunk after hooks + cost injection -> wire format)
StreamChunkSerializer = Callable[[Any], str]
StreamChunkSerializer = Callable[[Any], str] # any-ok: chunks are heterogeneous provider frames forwarded verbatim
# Type alias for streaming error serializer (ProxyException -> wire format)
StreamErrorSerializer = Callable[[ProxyException], str]
_T = TypeVar("_T")
_OBJECT_ADAPTER: TypeAdapter[object] = TypeAdapter(object)
@runtime_checkable
class _AsyncHttpResponseLike(Protocol):
status_code: int
headers: httpx.Headers
async def aread(self) -> bytes: ...
class _ModelInfoIdPayload(TypedDict, total=False):
id: str
class _ModelInfoMetadataPayload(TypedDict, total=False):
model_info: _ModelInfoIdPayload
class _ModelInfoLitellmParamsPayload(TypedDict, total=False):
model_info: _ModelInfoIdPayload
metadata: _ModelInfoMetadataPayload
class _LoggingKwargsPayload(TypedDict, total=False):
litellm_params: _ModelInfoLitellmParamsPayload
class _StreamUsageKwargsPayload(TypedDict, total=False):
prompt_tokens: int
completion_tokens: int
total_tokens: int
completion_tokens_details: dict[str, object]
prompt_tokens_details: dict[str, object]
server_tool_use: ServerToolUse
cache_creation_input_tokens: int
cache_read_input_tokens: int
async def _await_str_stream_coroutine(
coro: "Coroutine[None, None, AsyncGenerator[str, None]]",
) -> AsyncGenerator[str, None]:
return await coro
async def _await_bytes_stream_coroutine(
coro: "Coroutine[None, None, AsyncIterator[bytes]]",
) -> AsyncIterator[bytes]:
return await coro
@runtime_checkable
class _BytesAsyncIterator(Protocol):
def __aiter__(self) -> AsyncIterator[bytes]: ...
def __anext__(self) -> Awaitable[bytes]: ...
def _as_bytes_async_iterator(response: object) -> AsyncIterator[bytes]:
if isinstance(response, _BytesAsyncIterator):
return response
raise TypeError(f"expected an async iterator of bytes, got {type(response).__name__}")
if TYPE_CHECKING:
from litellm.proxy.proxy_server import ProxyConfig as _ProxyConfig
@ -232,7 +303,7 @@ async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, respons
return True
async def _cancel_pending_gather_tasks(tasks: list["asyncio.Task[Any]"]) -> None:
async def _cancel_pending_gather_tasks(tasks: list["asyncio.Task[object]"]) -> None:
pending_tasks = [task for task in tasks if not task.done()]
for task in pending_tasks:
task.cancel()
@ -309,8 +380,8 @@ def _stream_usage_tracking_updates(
def _serialize_http_exception_detail(
detail: Any,
) -> tuple[str, dict | None]:
detail: object,
) -> tuple[str, dict[str, object] | None]:
"""
Convert an HTTPException.detail value into (message, structured_fields)
for ProxyException / SSE error frames.
@ -339,7 +410,7 @@ def _serialize_http_exception_detail(
return str(detail), None
def _collect_response_file_search_vector_store_ids(data: dict[str, Any]) -> set[str]:
def _collect_response_file_search_vector_store_ids(data: dict[str, object]) -> set[str]:
vector_store_ids: set[str] = set()
tools = data.get("tools")
if not isinstance(tools, list):
@ -366,7 +437,7 @@ def _collect_response_file_search_vector_store_ids(data: dict[str, Any]) -> set[
async def _authorize_response_file_search_vector_stores(
data: dict[str, Any],
data: dict[str, object],
user_api_key_dict: UserAPIKeyAuth,
) -> None:
vector_store_ids = _collect_response_file_search_vector_store_ids(data)
@ -392,7 +463,7 @@ async def _parse_event_data_for_error(event_line: str | bytes) -> int | None:
if not json_str or json_str == "[DONE]": # handle empty data or [DONE] message
return None
try:
data = orjson.loads(json_str)
data = _OBJECT_ADAPTER.validate_python(orjson.loads(json_str))
if isinstance(data, dict) and "error" in data and isinstance(data["error"], dict):
error_code_raw = data["error"].get("code")
error_code: int | None = None
@ -445,7 +516,7 @@ def _extract_error_from_sse_chunk(event_line: str | bytes) -> dict:
return default_error
try:
data = orjson.loads(json_str)
data = _OBJECT_ADAPTER.validate_python(orjson.loads(json_str))
if isinstance(data, dict) and "error" in data:
error_obj = data["error"]
if isinstance(error_obj, dict):
@ -596,7 +667,7 @@ async def create_response(
try:
# Handle coroutine that returns a generator
if asyncio.iscoroutine(generator):
generator = await generator
generator = await _await_str_stream_coroutine(generator)
# Now get the first chunk from the actual generator
first_chunk_value = await _buffer_first_chunk_honoring_disconnect(generator, request)
@ -668,13 +739,13 @@ async def create_response(
existing_fields = getattr(e, "provider_specific_fields", None) or {}
if structured_fields:
merged_fields: dict | None = {**existing_fields, **structured_fields}
merged_fields: dict[str, object] | None = {**existing_fields, **structured_fields}
else:
merged_fields = existing_fields or None
# Match ProxyException.to_dict() shape so streaming and non-streaming
# error frames are byte-identical.
error_obj: dict[str, Any] = {
error_obj: dict[str, object] = {
"message": message,
"type": getattr(e, "type", "None"),
"param": getattr(e, "param", "None"),
@ -740,7 +811,7 @@ def _is_azure_model_router_request(model: str) -> bool:
def _override_openai_response_model(
*,
response_obj: Any,
response_obj: object,
requested_model: str,
log_context: str,
return_raw_model_name: bool = False,
@ -913,7 +984,7 @@ def _log_llm_api_exception(e: Exception) -> None:
async def _cancel_llm_call_on_client_disconnect(
request: Request,
llm_api_call: "asyncio.Future[Any]",
llm_api_call: "asyncio.Future[_T]",
disconnect_event: asyncio.Event,
) -> None:
try:
@ -932,8 +1003,8 @@ async def _cancel_llm_call_on_client_disconnect(
async def _await_llm_call_cancelling_on_disconnect(
request: Request,
llm_api_call: "asyncio.Future[Any]",
) -> Any:
llm_api_call: "asyncio.Future[_T]",
) -> _T:
disconnect_event = asyncio.Event()
monitor = asyncio.create_task(_cancel_llm_call_on_client_disconnect(request, llm_api_call, disconnect_event))
try:
@ -1054,7 +1125,7 @@ class ProxyBaseLLMRequestProcessing:
@staticmethod
async def build_litellm_proxy_success_headers_from_llm_response(
*,
response: Any,
response: object,
request_data: dict,
request: Request,
user_api_key_dict: UserAPIKeyAuth,
@ -1309,7 +1380,7 @@ class ProxyBaseLLMRequestProcessing:
## LOGGING OBJECT ## - initialize logging object for logging success/failure events for call
## IMPORTANT Note: - initialize this before running pre-call checks. Ensures we log rejected requests to langfuse.
logging_obj, self.data = litellm.utils.function_setup(
logging_obj, self.data = litellm.utils.function_setup( # any-ok: function_setup (litellm/utils.py) is untyped
original_function=route_type,
rules_obj=litellm.utils.Rules(),
start_time=start_time,
@ -1357,7 +1428,7 @@ class ProxyBaseLLMRequestProcessing:
self.data["router_settings_override"] = router_settings
if "messages" in self.data and self.data["messages"]:
logging_obj.update_messages(self.data["messages"])
logging_obj.update_messages(self.data["messages"]) # any-ok: logging_obj/messages come from untyped seams
return self.data, logging_obj
@ -1509,7 +1580,7 @@ class ProxyBaseLLMRequestProcessing:
@staticmethod
def _response_cost_from_logging_obj(
*,
response: Any,
response: Any, # any-ok: forwarded verbatim to the legacy calculator union in litellm_logging.py
logging_obj: LiteLLMLoggingObj,
) -> float | str:
"""
@ -1523,7 +1594,7 @@ class ProxyBaseLLMRequestProcessing:
stored_cost = logging_obj.model_call_details.get("response_cost")
if isinstance(stored_cost, (int, float)):
return float(stored_cost)
recomputed_cost = logging_obj._response_cost_calculator(result=response)
recomputed_cost = logging_obj._response_cost_calculator(result=response) # any-ok: untyped response seam
return recomputed_cost if isinstance(recomputed_cost, (int, float)) else ""
def _debug_log_request_payload(self) -> None:
@ -1654,7 +1725,7 @@ class ProxyBaseLLMRequestProcessing:
is_streaming_request: bool | None = False,
contents: list | None = None, # Add contents parameter
skip_pre_call_logic: bool = False,
) -> Any:
) -> object:
"""
Common request processing logic for both chat completions and responses API endpoints
"""
@ -1664,7 +1735,7 @@ class ProxyBaseLLMRequestProcessing:
self._debug_log_request_payload()
if skip_pre_call_logic:
logging_obj = self.data.get("litellm_logging_obj")
logging_obj: LiteLLMLoggingObj | None = self.data.get("litellm_logging_obj")
if logging_obj is None:
raise ValueError(
"skip_pre_call_logic=True requires litellm_logging_obj to be set in data. "
@ -1737,7 +1808,9 @@ class ProxyBaseLLMRequestProcessing:
llm_call_task = asyncio.create_task(llm_call)
tasks.append(llm_call_task)
llm_responses = asyncio.gather(*tasks) # run the moderation check in parallel to the actual llm api call
llm_responses: asyncio.Future[list[object]] = asyncio.gather(
*tasks
) # run the moderation check in parallel to the actual llm api call
try:
if general_settings.get("cancel_on_disconnect", False):
@ -1747,7 +1820,7 @@ class ProxyBaseLLMRequestProcessing:
finally:
await _cancel_pending_gather_tasks(tasks)
response = responses[1]
response: object = responses[1]
_exception_raised = False
try:
@ -1830,7 +1903,9 @@ class ProxyBaseLLMRequestProcessing:
_captured_user_api_key_dict = user_api_key_dict
_captured_logging_obj = logging_obj
async def _on_deferred_stream_complete(assembled_response, cache_hit):
async def _on_deferred_stream_complete(
assembled_response: ModelResponse, cache_hit: bool | None
) -> None:
await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails(
captured_data=_captured_data,
captured_user_api_key_dict=_captured_user_api_key_dict,
@ -1845,17 +1920,15 @@ class ProxyBaseLLMRequestProcessing:
# Check if response is an async generator
if self._is_streaming_response(response):
if asyncio.iscoroutine(response):
generator = await response
generator = await _await_bytes_stream_coroutine(response)
else:
generator = response
generator = _as_bytes_async_iterator(response)
if (
self._has_post_call_guardrails_for_passthrough()
and self._passthrough_endpoint_has_stream_guardrail_handler()
):
body_bytes = b"".join(
[chunk async for chunk in generator] # type: ignore[union-attr]
)
body_bytes = b"".join([chunk async for chunk in generator])
modified_bytes = await self._handle_event_stream_allm_passthrough_route(
body_bytes=body_bytes,
proxy_logging_obj=proxy_logging_obj,
@ -1874,7 +1947,7 @@ class ProxyBaseLLMRequestProcessing:
# For passthrough routes, stream directly without error parsing
# since we're dealing with raw binary data (e.g., AWS event streams)
return StreamingResponse(
content=generator, # type: ignore[arg-type]
content=generator,
status_code=status.HTTP_200_OK,
headers=custom_headers,
)
@ -1888,11 +1961,12 @@ class ProxyBaseLLMRequestProcessing:
)
if _early is not None:
return _early
return StreamingResponse(
content=response.aiter_bytes(), # type: ignore[union-attr]
status_code=response.status_code, # type: ignore[union-attr]
headers=custom_headers,
)
if isinstance(response, httpx.Response):
return StreamingResponse(
content=response.aiter_bytes(),
status_code=response.status_code,
headers=custom_headers,
)
elif route_type == "anthropic_messages":
# Check if response is actually a streaming response (async generator)
# Non-streaming responses (dict) should be returned directly
@ -1914,7 +1988,7 @@ class ProxyBaseLLMRequestProcessing:
)
# Non-streaming response - fall through to normal response handling
elif select_data_generator:
selected_data_generator = select_data_generator(
selected_data_generator: AsyncGenerator[str, None] = select_data_generator(
response=response,
user_api_key_dict=user_api_key_dict,
request_data=self.data,
@ -1978,10 +2052,12 @@ class ProxyBaseLLMRequestProcessing:
if _early is not None:
return _early
response = await proxy_logging_obj.post_call_success_hook(
data=self.data,
user_api_key_dict=user_api_key_dict,
response=response, # type: ignore[arg-type]
response = _OBJECT_ADAPTER.validate_python(
await proxy_logging_obj.post_call_success_hook(
data=self.data,
user_api_key_dict=user_api_key_dict,
response=response, # type: ignore[arg-type]
)
)
except Exception:
_exception_raised = True
@ -2080,7 +2156,7 @@ class ProxyBaseLLMRequestProcessing:
@staticmethod
async def _record_container_owners_from_responses_if_needed(
response: Any,
response: object,
user_api_key_dict: UserAPIKeyAuth,
) -> None:
"""Register code-interpreter containers so follow-up file APIs pass ownership checks."""
@ -2103,7 +2179,7 @@ class ProxyBaseLLMRequestProcessing:
)
@staticmethod
def _extract_completed_responses_response(stream_response: Any) -> Any:
def _extract_completed_responses_response(stream_response: object) -> object:
"""Pull the assembled ``ResponsesAPIResponse`` off a streaming iterator.
``ResponsesAPIStreamingIterator`` stores the terminal stream event
@ -2113,20 +2189,20 @@ class ProxyBaseLLMRequestProcessing:
``ResponsesAPIResponse`` directly. Handle both shapes so the
container-ownership recording path can walk ``.output`` either way.
"""
completed = getattr(stream_response, "completed_response", None)
completed: object = getattr(stream_response, "completed_response", None)
if completed is None:
return None
response_obj = getattr(completed, "response", None)
response_obj: object = getattr(completed, "response", None)
if response_obj is not None:
return response_obj
return completed
@staticmethod
async def _wrap_responses_stream_for_container_ownership(
original_stream_response: Any,
wrapped_generator: Any,
original_stream_response: object,
wrapped_generator: AsyncGenerator[str, None],
user_api_key_dict: UserAPIKeyAuth,
):
) -> AsyncGenerator[str, None]:
"""Forward SSE chunks, then record container ownership at stream end.
Streaming ``/v1/responses`` short-circuits out of
@ -2222,6 +2298,9 @@ class ProxyBaseLLMRequestProcessing:
if isinstance(result, Response):
return result
if not isinstance(result, httpx.Response):
return result
content = await result.aread()
return Response(
content=content,
@ -2232,7 +2311,7 @@ class ProxyBaseLLMRequestProcessing:
),
)
def _is_streaming_response(self, response: Any) -> bool:
def _is_streaming_response(self, response: object) -> bool:
"""
Check if the response object is actually a streaming response by inspecting its type.
@ -2340,7 +2419,7 @@ class ProxyBaseLLMRequestProcessing:
async def _handle_non_streaming_allm_passthrough_route(
self,
response: Any,
response: object,
proxy_logging_obj: "ProxyLogging",
user_api_key_dict: "UserAPIKeyAuth",
custom_headers: dict,
@ -2358,12 +2437,12 @@ class ProxyBaseLLMRequestProcessing:
HttpPassThroughEndpointHelpers,
)
try:
response_status: int = response.status_code # type: ignore[union-attr]
content_type: str = response.headers.get("content-type", "") # type: ignore[union-attr]
except AttributeError:
if not isinstance(response, _AsyncHttpResponseLike):
return None
response_status: int = response.status_code
content_type: str = dict(response.headers).get("content-type", "")
if response_status >= 300:
return None
@ -2374,7 +2453,7 @@ class ProxyBaseLLMRequestProcessing:
return None
response_headers = HttpPassThroughEndpointHelpers.get_response_headers(
headers=response.headers, # type: ignore[union-attr]
headers=response.headers,
custom_headers=custom_headers,
)
callback_headers = await proxy_logging_obj.post_call_response_headers_hook(
@ -2387,7 +2466,7 @@ class ProxyBaseLLMRequestProcessing:
response_headers.update(callback_headers)
if is_event_stream:
body_bytes = await response.aread() # type: ignore[union-attr]
body_bytes = await response.aread()
modified_bytes = await self._handle_event_stream_allm_passthrough_route(
body_bytes=body_bytes,
proxy_logging_obj=proxy_logging_obj,
@ -2400,9 +2479,9 @@ class ProxyBaseLLMRequestProcessing:
headers=response_headers,
)
body_bytes = await response.aread() # type: ignore[union-attr]
body_bytes = await response.aread()
try:
parsed = _json.loads(body_bytes)
parsed = _json.loads(body_bytes) # any-ok: provider-shaped JSON consumed by duck-typed guardrail hooks
except (_json.JSONDecodeError, UnicodeDecodeError):
return Response(
content=body_bytes,
@ -2413,7 +2492,7 @@ class ProxyBaseLLMRequestProcessing:
processed = await proxy_logging_obj.post_call_success_hook(
data=self.data,
user_api_key_dict=user_api_key_dict,
response=parsed,
response=parsed, # any-ok: hook contract declares the OpenAI model union; passthrough bodies are dicts
)
if isinstance(processed, dict):
content = _json.dumps(processed).encode()
@ -2449,7 +2528,7 @@ class ProxyBaseLLMRequestProcessing:
@staticmethod
def _flush_deferred_async_logging(
logging_obj: Any,
logging_obj: object,
exception_raised: bool,
) -> None:
"""
@ -2474,10 +2553,10 @@ class ProxyBaseLLMRequestProcessing:
Extracted as a static method so tests can exercise the production
gating logic directly rather than reimplementing the finally block.
"""
_enqueue_fn = getattr(logging_obj, "_enqueue_deferred_logging", None)
_enqueue_fn: Callable[[], None] | None = getattr(logging_obj, "_enqueue_deferred_logging", None)
if _enqueue_fn is None:
return
logging_obj._enqueue_deferred_logging = None # type: ignore[union-attr]
setattr(logging_obj, "_enqueue_deferred_logging", None)
if exception_raised:
return
try:
@ -2489,9 +2568,9 @@ class ProxyBaseLLMRequestProcessing:
async def _run_deferred_stream_guardrails(
captured_data: dict,
captured_user_api_key_dict: "UserAPIKeyAuth",
captured_logging_obj: Any,
assembled_response: Any,
cache_hit: Any,
captured_logging_obj: LiteLLMLoggingObj,
assembled_response: ModelResponse,
cache_hit: bool | None,
) -> None:
"""
Run non-streaming post-call guardrail hooks on an assembled streaming
@ -2526,7 +2605,7 @@ class ProxyBaseLLMRequestProcessing:
):
continue
try:
guardrail_result = None
guardrail_result: ModelResponse | None = None
if "apply_guardrail" in type(cb).__dict__:
# Skip — apply_guardrail guardrails already ran via
# unified_guardrail's end-of-stream block in the
@ -2541,7 +2620,7 @@ class ProxyBaseLLMRequestProcessing:
# here would duplicate the scan and can spuriously block the guardrail that already passed / failed.
continue
else:
guardrail_result = await cb.async_post_call_success_hook(
guardrail_result = await cb.async_post_call_success_hook( # any-ok: hook declared -> Any
user_api_key_dict=captured_user_api_key_dict,
data=guardrail_data,
response=_response,
@ -2640,9 +2719,9 @@ class ProxyBaseLLMRequestProcessing:
headers = getattr(e, "headers", None) or {}
if not headers:
# Try to get headers from e.response.headers (httpx.Response)
_response = getattr(e, "response", None)
_response: httpx.Response | None = getattr(e, "response", None)
if _response is not None:
_response_headers = getattr(_response, "headers", None)
_response_headers: httpx.Headers | None = getattr(_response, "headers", None)
if _response_headers:
headers = get_response_headers(dict(_response_headers))
headers.update(custom_headers)
@ -2736,7 +2815,7 @@ class ProxyBaseLLMRequestProcessing:
#########################################################
@staticmethod
def return_sse_chunk(chunk: Any) -> str:
def return_sse_chunk(chunk: Any) -> str: # any-ok: chunks are provider frames (str/bytes/dict) forwarded verbatim
"""
Helper function to format streaming chunks for Anthropic API format
@ -2751,13 +2830,13 @@ class ProxyBaseLLMRequestProcessing:
chunk_str = safe_dumps(chunk)
return f"{STREAM_SSE_DATA_PREFIX}{chunk_str}\n\n"
else:
return chunk
return chunk # any-ok: non-dict frames (str/bytes) pass through unchanged
@staticmethod
async def _finalize_streaming_generator_cleanup(
request: Request | None,
request_data: dict,
response: Any,
response: object,
stream_completed: bool = False,
client_disconnected: bool = False,
user_api_key_dict: UserAPIKeyAuth | None = None,
@ -2794,9 +2873,10 @@ class ProxyBaseLLMRequestProcessing:
):
await proxy_logging_obj._arelease_max_parallel_requests_on_disconnect(user_api_key_dict)
if hasattr(response, "aclose"):
aclose_fn: Callable[[], Awaitable[object]] | None = getattr(response, "aclose", None)
if aclose_fn is not None:
try:
await response.aclose()
await aclose_fn()
except BaseException as e: # noqa: BLE001
verbose_proxy_logger.debug(
"async_streaming_data_generator: error closing response stream: %s",
@ -2805,7 +2885,7 @@ class ProxyBaseLLMRequestProcessing:
@staticmethod
async def async_streaming_data_generator(
response: Any,
response: object,
user_api_key_dict: UserAPIKeyAuth,
request_data: dict,
proxy_logging_obj: ProxyLogging,
@ -2859,7 +2939,7 @@ class ProxyBaseLLMRequestProcessing:
str_so_far += response_str
elif hasattr(chunk, "model_dump"):
try:
d = chunk.model_dump(mode="json", exclude_none=True)
d: dict[str, object] = chunk.model_dump(mode="json", exclude_none=True)
if isinstance(d, dict):
str_so_far += str(d.get("content", ""))
except Exception:
@ -2936,7 +3016,7 @@ class ProxyBaseLLMRequestProcessing:
@staticmethod
def async_sse_data_generator(
response: Any,
response: object,
user_api_key_dict: UserAPIKeyAuth,
request_data: dict,
proxy_logging_obj: ProxyLogging,
@ -2963,7 +3043,7 @@ class ProxyBaseLLMRequestProcessing:
)
@staticmethod
def _process_chunk_with_cost_injection(chunk: Any, model_name: str) -> Any:
def _process_chunk_with_cost_injection(chunk: object, model_name: str) -> object:
"""
Process a streaming chunk and inject cost information if enabled.
@ -3023,12 +3103,13 @@ class ProxyBaseLLMRequestProcessing:
if stripped_ln.startswith("data:"):
json_part = stripped_ln.split("data:", 1)[1].strip()
if json_part and json_part != "[DONE]":
obj = json.loads(json_part)
maybe_modified = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(obj, model_name)
if maybe_modified is not None:
# Replace just this line with updated JSON using safe_dumps
lines[idx] = f"data: {safe_dumps(maybe_modified)}"
return "\n".join(lines)
obj = _OBJECT_ADAPTER.validate_python(json.loads(json_part))
if isinstance(obj, dict):
maybe_modified = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(obj, model_name)
if maybe_modified is not None:
# Replace just this line with updated JSON using safe_dumps
lines[idx] = f"data: {safe_dumps(maybe_modified)}"
return "\n".join(lines)
return None
except Exception:
return None
@ -3054,13 +3135,13 @@ class ProxyBaseLLMRequestProcessing:
)
# Extract additional usage fields
cache_creation_input_tokens = _usage.get("cache_creation_input_tokens")
cache_read_input_tokens = _usage.get("cache_read_input_tokens")
web_search_requests = _usage.get("web_search_requests")
completion_tokens_details = _usage.get("completion_tokens_details")
prompt_tokens_details = _usage.get("prompt_tokens_details")
cache_creation_input_tokens: int | None = _usage.get("cache_creation_input_tokens")
cache_read_input_tokens: int | None = _usage.get("cache_read_input_tokens")
web_search_requests: int | None = _usage.get("web_search_requests")
completion_tokens_details: dict[str, object] | None = _usage.get("completion_tokens_details")
prompt_tokens_details: dict[str, object] | None = _usage.get("prompt_tokens_details")
usage_kwargs: dict[str, Any] = {
usage_kwargs: _StreamUsageKwargsPayload = {
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"total_tokens": total_tokens,
@ -3120,18 +3201,18 @@ class ProxyBaseLLMRequestProcessing:
# 2. Fallback to kwargs (initial)
if not model_id:
_kwargs = getattr(_logging_obj, "kwargs", None)
_kwargs: _LoggingKwargsPayload | None = getattr(_logging_obj, "kwargs", None)
if _kwargs:
litellm_params = _kwargs.get("litellm_params", {})
kwargs_litellm_params = _kwargs.get("litellm_params", {})
# First check direct model_info path
model_info = litellm_params.get("model_info") or {}
model_id = model_info.get("id", None)
kwargs_model_info: _ModelInfoIdPayload = kwargs_litellm_params.get("model_info") or {}
model_id = kwargs_model_info.get("id", None)
# Fallback to nested metadata path
if not model_id:
metadata = litellm_params.get("metadata") or {}
model_info = metadata.get("model_info") or {}
model_id = model_info.get("id", None)
kwargs_metadata: _ModelInfoMetadataPayload = kwargs_litellm_params.get("metadata") or {}
kwargs_model_info = kwargs_metadata.get("model_info") or {}
model_id = kwargs_model_info.get("id", None)
# 3. Final fallback to self.data["litellm_metadata"] (for routes like /v1/responses that populate data before error)
if not model_id:

455
litellm/proxy/hooks/parallel_request_limiter_v3.py
View file

@ -8,7 +8,7 @@ import asyncio
import binascii
import os
import uuid
from collections.abc import Callable
from collections.abc import Awaitable, Callable, Mapping, Sequence
from contextvars import ContextVar
from dataclasses import dataclass, field
from datetime import datetime
@ -16,11 +16,12 @@ from typing import (
TYPE_CHECKING,
Any,
Literal,
Protocol,
TypedDict,
Union,
cast,
)
from pydantic import TypeAdapter
from litellm import DualCache
from litellm._logging import verbose_proxy_logger
from litellm.constants import DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE
@ -54,9 +55,10 @@ if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache
from litellm.types.agents import AgentResponse
from litellm.types.caching import RedisPipelineIncrementOperation
Span = Union[_Span, Any]
Span = _Span
InternalUsageCache = _InternalUsageCache
else:
Span = Any
@ -342,6 +344,71 @@ class RateLimitResponseWithDescriptors(TypedDict):
response: RateLimitResponse
class _WindowLimitMetadata(TypedDict):
requests_limit: int | None
tokens_limit: int | None
window_size: int
descriptor_key: str
class _AtomicCounterMeta(TypedDict):
descriptor_key: str
current_limit: int
rate_limit_type: Literal["requests", "tokens"]
window_key: str
counter_key: str
increment: int
ttl: int
window_size: int
class _AtomicCounterState(TypedDict):
window_expired: bool
current: int
class _TokenUsageDict(TypedDict, total=False):
prompt_tokens: int | None
completion_tokens: int | None
total_tokens: int | None
prompt_tokens_details: Mapping[str, int | None] | None
class _RedisLuaScript(Protocol):
def __call__(self, keys: Sequence[str], args: Sequence[int | float | str]) -> Awaitable[Sequence[object]]: ...
class _CallTypeRateLimiter(Protocol):
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict[str, object],
call_type: str,
) -> Exception | str | dict[str, object] | None: ...
_CACHED_VALUE_ADAPTER: TypeAdapter[object] = TypeAdapter(object)
def _as_int(value: object) -> int:
if isinstance(value, (int, float, str, bytes)):
return int(value)
raise TypeError(f"Expected an int-convertible rate limiter value, got {type(value).__name__}")
def _str_or_none(value: object) -> str | None:
return value if isinstance(value, str) else None
def _object_dict_or_none(value: object) -> dict[object, object] | None:
return value if isinstance(value, dict) else None
def _str_object_dict_or_none(value: object) -> dict[str, object] | None:
return value if isinstance(value, dict) else None
@dataclass(slots=True)
class RequestRateLimiterStash:
"""
@ -391,7 +458,7 @@ def get_or_create_request_stash() -> RequestRateLimiterStash:
return stash
def claim_request_stash_for_data(data: dict) -> RequestRateLimiterStash:
def claim_request_stash_for_data(data: Mapping) -> RequestRateLimiterStash:
stash = get_or_create_request_stash()
owner_call_id = data.get("litellm_call_id")
if isinstance(owner_call_id, str):
@ -424,23 +491,23 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
self.internal_usage_cache = internal_usage_cache
self._time_provider = time_provider or datetime.now
if self.internal_usage_cache.dual_cache.redis_cache is not None:
self.batch_rate_limiter_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script(
BATCH_RATE_LIMITER_SCRIPT
self.batch_rate_limiter_script: _RedisLuaScript | None = (
self.internal_usage_cache.dual_cache.redis_cache.async_register_script(BATCH_RATE_LIMITER_SCRIPT)
)
self.token_increment_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script(
TOKEN_INCREMENT_SCRIPT
self.token_increment_script: _RedisLuaScript | None = (
self.internal_usage_cache.dual_cache.redis_cache.async_register_script(TOKEN_INCREMENT_SCRIPT)
)
self.check_and_increment_by_n_script = (
self.check_and_increment_by_n_script: _RedisLuaScript | None = (
self.internal_usage_cache.dual_cache.redis_cache.async_register_script(CHECK_AND_INCREMENT_BY_N_SCRIPT)
)
self.parallel_acquire_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script(
PARALLEL_ACQUIRE_SCRIPT
self.parallel_acquire_script: _RedisLuaScript | None = (
self.internal_usage_cache.dual_cache.redis_cache.async_register_script(PARALLEL_ACQUIRE_SCRIPT)
)
self.parallel_release_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script(
PARALLEL_RELEASE_SCRIPT
self.parallel_release_script: _RedisLuaScript | None = (
self.internal_usage_cache.dual_cache.redis_cache.async_register_script(PARALLEL_RELEASE_SCRIPT)
)
self.parallel_count_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script(
PARALLEL_COUNT_SCRIPT
self.parallel_count_script: _RedisLuaScript | None = (
self.internal_usage_cache.dual_cache.redis_cache.async_register_script(PARALLEL_COUNT_SCRIPT)
)
else:
self.batch_rate_limiter_script = None
@ -459,7 +526,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
self.tpm_reservation_enabled = os.getenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", "true").lower() == "true"
# Batch rate limiter (lazy loaded)
self._batch_rate_limiter: Any | None = None
self._batch_rate_limiter: _CallTypeRateLimiter | None = None
# Serializes multi-phase check+increment sequences (batch + dynamic
# limiters) within this process to close the TOCTOU window between
@ -477,7 +544,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
# one round-trip.
self._check_and_increment_lock = asyncio.Lock()
def _get_batch_rate_limiter(self) -> Any | None:
def _get_batch_rate_limiter(self) -> _CallTypeRateLimiter | None:
"""Get or lazy-load the batch rate limiter."""
if self._batch_rate_limiter is None:
try:
@ -497,6 +564,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
"""Return the current time for rate limiting calculations."""
return self._time_provider()
async def _get_local_cache_value(self, key: str, parent_otel_span: Span | None) -> object:
return _CACHED_VALUE_ADAPTER.validate_python(
await self.internal_usage_cache.async_get_cache(
key=key,
litellm_parent_otel_span=parent_otel_span,
local_only=True,
)
)
@staticmethod
def _no_max_tokens_output_floor(
min_configured_tpm_limit: int | None,
@ -514,7 +590,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
def _estimate_tokens_for_request(
self,
data: dict,
data: Mapping,
model: str | None = None,
min_configured_tpm_limit: int | None = None,
) -> int:
@ -601,15 +677,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
async def in_memory_cache_sliding_window(
self,
keys: list[str],
keys: Sequence[str],
now_int: int,
window_size: int,
) -> list[Any]:
) -> Sequence[object]:
"""
Implement sliding window rate limiting logic using in-memory cache operations.
This follows the same logic as the Redis Lua script but uses async cache operations.
"""
results: list[Any] = []
results: list[object] = []
# Process each window/counter pair
for i in range(0, len(keys), 2):
@ -618,14 +694,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
increment_value = 1
# Get the window start time
window_start = await self.internal_usage_cache.async_get_cache(
key=window_key,
litellm_parent_otel_span=None,
local_only=True,
)
window_start = await self._get_local_cache_value(key=window_key, parent_otel_span=None)
# Check if window exists and is valid
if window_start is None or (now_int - int(window_start)) >= window_size:
if window_start is None or (now_int - _as_int(window_start)) >= window_size:
# Reset window and counter
await self.internal_usage_cache.async_set_cache(
key=window_key,
@ -645,12 +717,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
results.append(increment_value) # counter
else:
# Increment the counter
current_counter = await self.internal_usage_cache.async_get_cache(
key=counter_key,
litellm_parent_otel_span=None,
local_only=True,
)
new_counter_value = (int(current_counter) if current_counter is not None else 0) + increment_value
current_counter = await self._get_local_cache_value(key=counter_key, parent_otel_span=None)
new_counter_value = (_as_int(current_counter) if current_counter is not None else 0) + increment_value
await self.internal_usage_cache.async_set_cache(
key=counter_key,
value=new_counter_value,
@ -678,9 +746,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
def is_cache_list_over_limit(
self,
keys_to_fetch: list[str],
cache_values: list[Any],
key_metadata: dict[str, Any],
keys_to_fetch: Sequence[str],
cache_values: Sequence[object],
key_metadata: Mapping[str, _WindowLimitMetadata],
) -> RateLimitResponse:
"""
Check if the cache values are over the limit.
@ -709,12 +777,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
if current_limit is None or rate_limit_type is None:
continue
if counter_value is not None and int(counter_value) > current_limit:
if counter_value is not None and _as_int(counter_value) > current_limit:
overall_code = "OVER_LIMIT"
item_code = "OVER_LIMIT"
# Only compute limit_remaining if current_limit is not None
limit_remaining = current_limit - int(counter_value) if counter_value is not None else current_limit
limit_remaining = current_limit - _as_int(counter_value) if counter_value is not None else current_limit
statuses.append(
{
@ -783,7 +851,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
self,
keys_to_fetch: list[str],
now_int: int,
) -> list[Any]:
) -> Sequence[object]:
"""
Execute Redis operations grouped by hash tag for cluster compatibility.
@ -792,13 +860,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
now_int: int - Current timestamp
Returns:
List[Any] - List of cache values
List of cache values
"""
if self.batch_rate_limiter_script is None:
return []
key_groups = self._group_keys_by_hash_tag(keys_to_fetch)
all_cache_values = []
all_cache_values: list[object] = []
for hash_tag, group_keys in key_groups.items():
try:
@ -821,7 +889,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
async def should_rate_limit(
self,
descriptors: list[RateLimitDescriptor],
descriptors: Sequence[RateLimitDescriptor],
parent_otel_span: Span | None = None,
read_only: bool = False,
skip_tpm_check: bool = False,
@ -866,7 +934,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
windowed_response = RateLimitResponse(overall_code="OK", statuses=[])
if keys_to_fetch:
## CHECK IN-MEMORY CACHE
cache_values = await self.internal_usage_cache.async_batch_get_cache(
cache_values: Sequence[object] | None = await self.internal_usage_cache.async_batch_get_cache(
keys=keys_to_fetch,
parent_otel_span=parent_otel_span,
local_only=True,
@ -888,9 +956,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
# For keys that don't exist yet, set them to 0
if cache_values is None:
cache_values = []
for _ in keys_to_fetch:
cache_values.append(str(now_int) if _.endswith(":window") else 0)
cache_values = [str(now_int) if key.endswith(":window") else 0 for key in keys_to_fetch]
elif self.batch_rate_limiter_script is not None:
# NORMAL MODE: Increment counters in Redis
# Group keys by hash tag for Redis cluster compatibility
@ -947,16 +1013,16 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
def _collect_windowed_keys_and_gauges(
self,
descriptors: list[RateLimitDescriptor],
descriptors: Sequence[RateLimitDescriptor],
skip_tpm_check: bool,
) -> tuple[list[str], dict[str, dict[str, Any]], list[ParallelRequestGauge]]:
) -> tuple[list[str], dict[str, _WindowLimitMetadata], list[ParallelRequestGauge]]:
"""
Split descriptors into the windowed (window_key, counter_key) fetch
list with its per-window metadata, and the concurrency gauges for
descriptors carrying a max_parallel_requests limit.
"""
keys_to_fetch: list[str] = []
key_metadata: dict[str, dict[str, Any]] = {}
key_metadata: dict[str, _WindowLimitMetadata] = {}
gauges: list[ParallelRequestGauge] = []
for descriptor in descriptors:
descriptor_key = descriptor["key"]
@ -1012,7 +1078,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
descriptor_key=gauge["descriptor_key"],
)
def _gauge_in_flight_from_cache_value(self, raw_value: Any) -> int:
def _gauge_in_flight_from_cache_value(self, raw_value: object) -> int:
"""
In-flight count from a cached gauge value: a dict of slot_id ->
acquire timestamp when the in-memory registry is authoritative, or
@ -1020,14 +1086,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
"""
if raw_value is None:
return 0
if isinstance(raw_value, dict):
slot_acquire_times = _object_dict_or_none(raw_value)
if slot_acquire_times is not None:
cutoff = self._get_current_time().timestamp() - PARALLEL_REQUEST_SLOT_TTL_SECONDS
return sum(1 for ts in raw_value.values() if isinstance(ts, (int, float)) and ts >= cutoff)
return max(0, int(raw_value))
return sum(1 for ts in slot_acquire_times.values() if isinstance(ts, (int, float)) and ts >= cutoff)
return max(0, _as_int(raw_value))
async def _check_parallel_request_gauges(
self,
gauges: list[ParallelRequestGauge],
gauges: Sequence[ParallelRequestGauge],
slot_id: str,
parent_otel_span: Span | None = None,
read_only: bool = False,
@ -1053,7 +1120,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
keys=gauge_keys,
args=[PARALLEL_REQUEST_SLOT_TTL_SECONDS for _ in gauges],
)
counts = [max(0, int(value)) for value in raw_counts]
counts = [max(0, _as_int(value)) for value in raw_counts]
except Exception as e: # noqa: BLE001 - any Redis/Lua failure degrades to the local mirror, never a 500
verbose_proxy_logger.warning(f"parallel_count_script failed, using local mirror: {e!s}")
counts = await self._read_local_gauge_counts(gauge_keys, parent_otel_span)
@ -1088,22 +1155,22 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
verbose_proxy_logger.warning(f"parallel_acquire_script failed, falling back to in-memory gauge: {e!s}")
async with self._check_and_increment_lock:
return await self._acquire_parallel_slots_in_memory(gauges, slot_id, parent_otel_span)
if int(raw[0]) == 1:
gauge = gauges[int(raw[1]) - 1]
if _as_int(raw[0]) == 1:
gauge = gauges[_as_int(raw[1]) - 1]
return RateLimitResponse(
overall_code="OVER_LIMIT",
statuses=[self._gauge_status(gauge, int(raw[2]), "OVER_LIMIT")],
statuses=[self._gauge_status(gauge, _as_int(raw[2]), "OVER_LIMIT")],
)
statuses = []
for gauge, in_flight in zip(gauges, raw[1:]):
await self.internal_usage_cache.async_set_cache(
key=gauge["counter_key"],
value=int(in_flight),
value=_as_int(in_flight),
ttl=PARALLEL_REQUEST_SLOT_TTL_SECONDS,
litellm_parent_otel_span=parent_otel_span,
local_only=True,
)
statuses.append(self._gauge_status(gauge, int(in_flight), "OK"))
statuses.append(self._gauge_status(gauge, _as_int(in_flight), "OK"))
return RateLimitResponse(overall_code="OK", statuses=statuses)
async with self._check_and_increment_lock:
@ -1113,8 +1180,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
self,
gauge_keys: list[str],
parent_otel_span: Span | None = None,
) -> list[int]:
values = await self.internal_usage_cache.async_batch_get_cache(
) -> Sequence[int]:
values: Sequence[object] | None = await self.internal_usage_cache.async_batch_get_cache(
keys=gauge_keys,
parent_otel_span=parent_otel_span,
local_only=True,
@ -1125,7 +1192,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
async def _acquire_parallel_slots_in_memory(
self,
gauges: list[ParallelRequestGauge],
gauges: Sequence[ParallelRequestGauge],
slot_id: str,
parent_otel_span: Span | None = None,
) -> RateLimitResponse:
@ -1143,14 +1210,16 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
cutoff = now - PARALLEL_REQUEST_SLOT_TTL_SECONDS
states: list[tuple[dict[str, float] | None, int]] = []
for gauge in gauges:
raw_value = await self.internal_usage_cache.async_get_cache(
raw_value = await self._get_local_cache_value(
key=gauge["counter_key"],
litellm_parent_otel_span=parent_otel_span,
local_only=True,
parent_otel_span=parent_otel_span,
)
if isinstance(raw_value, dict):
slot_acquire_times = _object_dict_or_none(raw_value)
if slot_acquire_times is not None:
registry: dict[str, float] | None = {
key: float(ts) for key, ts in raw_value.items() if isinstance(ts, (int, float)) and ts >= cutoff
key: float(ts)
for key, ts in slot_acquire_times.items()
if isinstance(key, str) and isinstance(ts, (int, float)) and ts >= cutoff
}
in_flight = len(registry or {})
elif raw_value is None:
@ -1158,7 +1227,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
in_flight = 0
else:
registry = None
in_flight = max(0, int(raw_value))
in_flight = max(0, _as_int(raw_value))
if in_flight + 1 > gauge["limit"]:
return RateLimitResponse(
overall_code="OVER_LIMIT",
@ -1205,7 +1274,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
for counter_key, remaining in zip(counter_keys, raw):
await self.internal_usage_cache.async_set_cache(
key=counter_key,
value=max(0, int(remaining)),
value=max(0, _as_int(remaining)),
ttl=PARALLEL_REQUEST_SLOT_TTL_SECONDS,
litellm_parent_otel_span=parent_otel_span,
local_only=True,
@ -1218,19 +1287,21 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
async with self._check_and_increment_lock:
for counter_key in counter_keys:
raw_value = await self.internal_usage_cache.async_get_cache(
raw_value = await self._get_local_cache_value(
key=counter_key,
litellm_parent_otel_span=parent_otel_span,
local_only=True,
parent_otel_span=parent_otel_span,
)
if isinstance(raw_value, dict):
if slot_id not in raw_value:
slot_acquire_times = _object_dict_or_none(raw_value)
if slot_acquire_times is not None:
if slot_id not in slot_acquire_times:
continue
new_value: dict[str, float] | int = {key: ts for key, ts in raw_value.items() if key != slot_id}
new_value: dict[object, object] | int = {
key: ts for key, ts in slot_acquire_times.items() if key != slot_id
}
elif raw_value is None:
continue
else:
new_value = max(0, int(raw_value) - 1)
new_value = max(0, _as_int(raw_value) - 1)
await self.internal_usage_cache.async_set_cache(
key=counter_key,
value=new_value,
@ -1241,8 +1312,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
async def atomic_check_and_increment_by_n(
self,
descriptors: list[RateLimitDescriptor],
increments: list[dict[Literal["requests", "tokens"], int]],
descriptors: Sequence[RateLimitDescriptor],
increments: Sequence[Mapping[Literal["requests", "tokens"], int]],
parent_otel_span: Span | None = None,
) -> RateLimitResponse:
"""
@ -1277,7 +1348,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
# Build per-descriptor (keys, args, meta) groups. All keys within a
# group share the descriptor's {key:value} hash tag, so a single Lua
# call per group never triggers CROSSSLOT on Redis Cluster.
descriptor_groups: list[tuple[list[str], list[Any], list[dict[str, Any]]]] = []
descriptor_groups: list[tuple[Sequence[str], Sequence[int], Sequence[_AtomicCounterMeta]]] = []
for descriptor, increment_amounts in zip(descriptors, increments):
keys, args, meta = self._build_descriptor_atomic_payload(
descriptor=descriptor,
@ -1300,7 +1371,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
parent_otel_span=parent_otel_span,
)
flat_meta: list[dict[str, Any]] = [m for _keys, _args, group_meta in descriptor_groups for m in group_meta]
flat_meta: Sequence[_AtomicCounterMeta] = [
m for _keys, _args, group_meta in descriptor_groups for m in group_meta
]
async with self._check_and_increment_lock:
return await self._atomic_check_and_increment_in_memory(
per_counter_meta=flat_meta,
@ -1310,8 +1383,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
def _build_descriptor_atomic_payload(
self,
descriptor: RateLimitDescriptor,
increment_amounts: dict[Literal["requests", "tokens"], int],
) -> tuple[list[str], list[Any], list[dict[str, Any]]]:
increment_amounts: Mapping[Literal["requests", "tokens"], int],
) -> tuple[list[str], list[int], list[_AtomicCounterMeta]]:
"""
Build (KEYS, ARGV, per-counter meta) for a single descriptor's Lua
call. All keys returned share the descriptor's {key:value} hash tag.
@ -1325,11 +1398,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
window_key = f"{{{descriptor_key}:{descriptor_value}}}:window"
keys: list[str] = []
args: list[Any] = []
meta: list[dict[str, Any]] = []
args: list[int] = []
meta: list[_AtomicCounterMeta] = []
for rate_limit_type in ("requests", "tokens"):
rlt: Literal["requests", "tokens"] = cast(Literal["requests", "tokens"], rate_limit_type)
rlt: Literal["requests", "tokens"] = rate_limit_type
if rlt == "requests":
limit_value = rate_limit.get("requests_per_unit")
inc_amount = int(increment_amounts.get("requests", 0) or 0)
@ -1365,7 +1438,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
async def _atomic_lua_per_descriptor(
self,
descriptor_groups: list[tuple[list[str], list[Any], list[dict[str, Any]]]],
descriptor_groups: Sequence[tuple[Sequence[str], Sequence[int], Sequence[_AtomicCounterMeta]]],
parent_otel_span: Span | None = None,
) -> RateLimitResponse:
"""
@ -1374,7 +1447,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
descriptor i, refund descriptors 0..i-1's increments. On Lua failure
mid-loop, refund applied increments and fall back to in-memory.
"""
applied: list[list[dict[str, Any]]] = []
applied: list[Sequence[_AtomicCounterMeta]] = []
statuses: list[RateLimitStatus] = []
for _idx, (keys, args, meta) in enumerate(descriptor_groups):
@ -1397,7 +1470,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
f"{self.window_size}s)."
)
await self._refund_applied_descriptor_groups(applied)
flat_meta: list[dict[str, Any]] = [m for _k, _a, group_meta in descriptor_groups for m in group_meta]
flat_meta: Sequence[_AtomicCounterMeta] = [
m for _k, _a, group_meta in descriptor_groups for m in group_meta
]
async with self._check_and_increment_lock:
return await self._atomic_check_and_increment_in_memory(
per_counter_meta=flat_meta,
@ -1415,7 +1490,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
async def _refund_applied_descriptor_groups(
self,
applied: list[list[dict[str, Any]]],
applied: Sequence[Sequence[_AtomicCounterMeta]],
) -> None:
"""
Decrement counters for descriptor groups already applied via Lua.
@ -1441,8 +1516,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
def _build_atomic_response(
self,
raw: list[Any],
per_counter_meta: list[dict[str, Any]],
raw: Sequence[object],
per_counter_meta: Sequence[_AtomicCounterMeta],
) -> RateLimitResponse:
"""Convert Lua script return value to RateLimitResponse.
@ -1458,12 +1533,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
if not raw:
return RateLimitResponse(overall_code="OK", statuses=[])
status_code = int(raw[0])
status_code = _as_int(raw[0])
if status_code == 1:
# Over limit: { 1, counter_index (1-based), current_counter, limit }
descriptor_index = int(raw[1]) - 1
current_counter = int(raw[2])
limit = int(raw[3])
descriptor_index = _as_int(raw[1]) - 1
current_counter = _as_int(raw[2])
limit = _as_int(raw[3])
meta = per_counter_meta[descriptor_index]
return RateLimitResponse(
overall_code="OVER_LIMIT",
@ -1484,7 +1559,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
RateLimitStatus(
code="OK",
current_limit=meta["current_limit"],
limit_remaining=max(0, meta["current_limit"] - int(new_counter)),
limit_remaining=max(0, meta["current_limit"] - _as_int(new_counter)),
rate_limit_type=meta["rate_limit_type"],
descriptor_key=meta["descriptor_key"],
)
@ -1493,7 +1568,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
async def _atomic_check_and_increment_in_memory(
self,
per_counter_meta: list[dict[str, Any]],
per_counter_meta: Sequence[_AtomicCounterMeta],
parent_otel_span: Span | None = None,
) -> RateLimitResponse:
"""In-memory all-or-nothing check-and-increment. Caller holds lock.
@ -1508,27 +1583,23 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
now_int = int(self._get_current_time().timestamp())
# Pass 1: read state, validate.
descriptor_state: list[dict[str, Any]] = []
descriptor_state: list[_AtomicCounterState] = []
for meta in per_counter_meta:
window_size = meta["window_size"]
window_start = await self.internal_usage_cache.async_get_cache(
window_start = await self._get_local_cache_value(
key=meta["window_key"],
litellm_parent_otel_span=parent_otel_span,
local_only=True,
parent_otel_span=parent_otel_span,
)
window_expired = window_start is None or (now_int - int(window_start)) >= window_size
current_counter = (
0
window_expired = window_start is None or (now_int - _as_int(window_start)) >= window_size
counter_raw = (
None
if window_expired
else int(
await self.internal_usage_cache.async_get_cache(
key=meta["counter_key"],
litellm_parent_otel_span=parent_otel_span,
local_only=True,
)
or 0
else await self._get_local_cache_value(
key=meta["counter_key"],
parent_otel_span=parent_otel_span,
)
)
current_counter = 0 if window_expired else _as_int(counter_raw or 0)
over_limit = (
current_counter + meta["increment"] > meta["current_limit"]
if meta["increment"] > 0
@ -1581,7 +1652,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
async def reserve_tpm_tokens(
self,
descriptors: list[RateLimitDescriptor],
descriptors: Sequence[RateLimitDescriptor],
estimated_tokens: int,
parent_otel_span: Span | None = None,
) -> RateLimitResponse:
@ -1595,7 +1666,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
atomicity (Lua on Redis, asyncio-locked DualCache otherwise) to the
shared primitive.
"""
tpm_descriptors: list[RateLimitDescriptor] = [
tpm_descriptors: Sequence[RateLimitDescriptor] = [
RateLimitDescriptor(
key=d["key"],
value=d["value"],
@ -1610,7 +1681,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
if not tpm_descriptors:
return RateLimitResponse(overall_code="OK", statuses=[])
increments: list[dict[Literal["requests", "tokens"], int]] = [
increments: Sequence[Mapping[Literal["requests", "tokens"], int]] = [
{"tokens": estimated_tokens} for _ in tpm_descriptors
]
return await self.atomic_check_and_increment_by_n(
@ -1621,7 +1692,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
def create_organization_rate_limit_descriptor(
self, user_api_key_dict: UserAPIKeyAuth, requested_model: str | None = None
) -> list[RateLimitDescriptor]:
) -> Sequence[RateLimitDescriptor]:
descriptors: list[RateLimitDescriptor] = []
# Global org rate limits
@ -1920,7 +1991,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
"""
return rpm_limit_type == "dynamic" or tpm_limit_type == "dynamic"
def _get_agent_from_registry(self, agent_id: str) -> Any | None:
def _get_agent_from_registry(self, agent_id: str) -> "AgentResponse | None":
"""Look up an agent from the in-memory registry by ID."""
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
@ -1931,7 +2002,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
Resolve the agent_id from either the API key or request metadata.
Key-level agent_id takes precedence over metadata/header-supplied agent_id.
"""
key_agent_id = getattr(user_api_key_dict, "agent_id", None)
key_agent_id_raw: object = getattr(user_api_key_dict, "agent_id", None)
key_agent_id = _str_or_none(key_agent_id_raw)
if key_agent_id:
return key_agent_id
metadata = data.get("metadata") or {}
@ -1969,8 +2041,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
if agent is None:
return descriptors
agent_rpm = getattr(agent, "rpm_limit", None)
agent_tpm = getattr(agent, "tpm_limit", None)
agent_rpm = agent.rpm_limit
agent_tpm = agent.tpm_limit
if agent_rpm is not None or agent_tpm is not None:
descriptors.append(
RateLimitDescriptor(
@ -1984,8 +2056,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
)
)
session_rpm = getattr(agent, "session_rpm_limit", None)
session_tpm = getattr(agent, "session_tpm_limit", None)
session_rpm = agent.session_rpm_limit
session_tpm = agent.session_tpm_limit
if session_rpm is not None or session_tpm is not None:
session_id = self._get_session_id_from_data(data)
if session_id is not None:
@ -2244,7 +2316,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
# Fail safe: enforce limits if we can't check
return True
def get_rate_limiter_for_call_type(self, call_type: str) -> Any | None:
def get_rate_limiter_for_call_type(self, call_type: str) -> _CallTypeRateLimiter | None:
"""Get the rate limiter for the call type."""
if call_type == "acreate_batch":
batch_limiter = self._get_batch_rate_limiter()
@ -2371,7 +2443,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
self,
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
data: dict[str, object],
call_type: str,
):
"""
@ -2402,7 +2474,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
# For dynamic mode, check if the model has recent failures
model_has_failures = False
requested_model = data.get("model", None)
requested_model = _str_or_none(data.get("model", None))
if (
self._is_dynamic_rate_limiting_enabled(
@ -2604,7 +2676,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
return pipeline_operations
def _get_total_tokens_from_usage(
self, usage: Any | None, rate_limit_type: Literal["output", "input", "total"]
self, usage: "Usage | _TokenUsageDict | None", rate_limit_type: Literal["output", "input", "total"]
) -> int:
"""
Get total tokens from response usage for rate limiting.
@ -2616,34 +2688,33 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
total_tokens = 0
cached_tokens = 0
if usage:
if isinstance(usage, Usage):
if rate_limit_type == "output":
total_tokens = usage.completion_tokens or 0
elif rate_limit_type == "input":
total_tokens = usage.prompt_tokens or 0
elif rate_limit_type == "total":
total_tokens = usage.total_tokens or 0
if isinstance(usage, Usage):
if rate_limit_type == "output":
total_tokens = usage.completion_tokens or 0
elif rate_limit_type == "input":
total_tokens = usage.prompt_tokens or 0
elif rate_limit_type == "total":
total_tokens = usage.total_tokens or 0
# Get cached tokens to exclude from input/total
if rate_limit_type in ("input", "total"):
if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details is not None:
cached_tokens = getattr(usage.prompt_tokens_details, "cached_tokens", 0) or 0
# Get cached tokens to exclude from input/total
if rate_limit_type in ("input", "total"):
if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details is not None:
cached_tokens = getattr(usage.prompt_tokens_details, "cached_tokens", 0) or 0
elif isinstance(usage, dict):
# Responses API usage comes as a dict
if rate_limit_type == "output":
total_tokens = usage.get("completion_tokens", 0) or 0
elif rate_limit_type == "input":
total_tokens = usage.get("prompt_tokens", 0) or 0
elif rate_limit_type == "total":
total_tokens = usage.get("total_tokens", 0) or 0
elif isinstance(usage, dict) and usage:
# Responses API usage comes as a dict
if rate_limit_type == "output":
total_tokens = usage.get("completion_tokens", 0) or 0
elif rate_limit_type == "input":
total_tokens = usage.get("prompt_tokens", 0) or 0
elif rate_limit_type == "total":
total_tokens = usage.get("total_tokens", 0) or 0
# Get cached tokens from dict
if rate_limit_type in ("input", "total"):
prompt_details = usage.get("prompt_tokens_details") or {}
if isinstance(prompt_details, dict):
cached_tokens = prompt_details.get("cached_tokens", 0) or 0
# Get cached tokens from dict
if rate_limit_type in ("input", "total"):
prompt_details = usage.get("prompt_tokens_details") or {}
if isinstance(prompt_details, dict):
cached_tokens = prompt_details.get("cached_tokens", 0) or 0
# Subtract cached tokens for input/total (providers don't count them)
if cached_tokens > 0:
@ -2652,7 +2723,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
return total_tokens
@staticmethod
def _aggregate_only_total_tokens(usage: Usage | dict | None) -> int:
def _aggregate_only_total_tokens(usage: "Usage | _TokenUsageDict | None") -> int:
"""Total for usage that carries no input/output split, else 0.
A source that can only report one number for the whole request (a
@ -2767,15 +2838,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
@staticmethod
def _merge_ratelimit_statuses_into_additional_headers(
additional_headers: dict[str, Any],
additional_headers: Mapping[str, object],
statuses: list[RateLimitStatus],
) -> dict[str, Any]:
) -> dict[str, object]:
"""
Return ``additional_headers`` extended with
``x-ratelimit-{descriptor_key}-{remaining|limit}-{rate_limit_type}``
entries. Non-mutating so callers pick their own target dict.
"""
merged: dict[str, Any] = dict(additional_headers)
merged: dict[str, object] = dict(additional_headers)
for status in statuses:
prefix = f"x-ratelimit-{status['descriptor_key']}"
merged[f"{prefix}-remaining-{status['rate_limit_type']}"] = status["limit_remaining"]
@ -2784,8 +2855,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
def _collect_tpm_scope_targets(
self,
standard_logging_metadata: dict[str, Any],
kwargs: Any,
standard_logging_metadata: Mapping[str, object],
kwargs: Mapping[str, object],
model_group: str | None,
) -> list[tuple[str, str]]:
"""
@ -2795,16 +2866,18 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
the emitter; this helper just lists the candidate scopes so callers
can split reserved-vs-unreserved.
"""
user_api_key = standard_logging_metadata.get("user_api_key_hash")
user_api_key_user_id = standard_logging_metadata.get("user_api_key_user_id")
user_api_key_team_id = standard_logging_metadata.get("user_api_key_team_id")
user_api_key_organization_id = standard_logging_metadata.get("user_api_key_org_id")
user_api_key_project_id = standard_logging_metadata.get("user_api_key_project_id")
user_api_key_end_user_id = (
user_api_key = _str_or_none(standard_logging_metadata.get("user_api_key_hash"))
user_api_key_user_id = _str_or_none(standard_logging_metadata.get("user_api_key_user_id"))
user_api_key_team_id = _str_or_none(standard_logging_metadata.get("user_api_key_team_id"))
user_api_key_organization_id = _str_or_none(standard_logging_metadata.get("user_api_key_org_id"))
user_api_key_project_id = _str_or_none(standard_logging_metadata.get("user_api_key_project_id"))
user_api_key_end_user_id = _str_or_none(
kwargs.get("user") if isinstance(kwargs, dict) else None
) or standard_logging_metadata.get("user_api_key_end_user_id")
agent_id = standard_logging_metadata.get("agent_id")
session_id = standard_logging_metadata.get("session_id") or standard_logging_metadata.get("trace_id")
) or _str_or_none(standard_logging_metadata.get("user_api_key_end_user_id"))
agent_id = _str_or_none(standard_logging_metadata.get("agent_id"))
session_id = _str_or_none(standard_logging_metadata.get("session_id")) or _str_or_none(
standard_logging_metadata.get("trace_id")
)
targets: list[tuple[str, str]] = []
if user_api_key:
@ -2882,8 +2955,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
def _build_success_event_pipeline_operations(
self,
kwargs: Any,
response_obj: Any,
kwargs: dict[str, object],
response_obj: object,
rate_limit_type: Literal["output", "input", "total"],
) -> list[RedisPipelineIncrementOperation]:
"""Build Redis pipeline increment ops for TPM / parallel-request counters."""
@ -2893,8 +2966,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
# Get metadata from standard_logging_object - this correctly handles both
# 'metadata' and 'litellm_metadata' fields from litellm_params
standard_logging_object = kwargs.get("standard_logging_object") or {}
standard_logging_metadata = standard_logging_object.get("metadata") or {}
standard_logging_object = _str_object_dict_or_none(kwargs.get("standard_logging_object")) or {}
standard_logging_metadata = _str_object_dict_or_none(standard_logging_object.get("metadata")) or {}
model_group = get_model_group_from_litellm_kwargs(kwargs)
@ -2903,7 +2976,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
# than parsed out of the body) carry their usage in
# ``combined_usage_object`` instead, and would otherwise never charge
# the TPM window.
_usage: Usage | dict | None = None
_usage: Usage | _TokenUsageDict | None = None
if isinstance(
response_obj,
(
@ -2967,7 +3040,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
return pipeline_operations
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
async def async_log_success_event(
self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object
):
"""
Update TPM usage on successful API calls by incrementing counters using pipeline
"""
@ -3007,10 +3082,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
async def async_logging_hook(
self,
kwargs: dict,
result: Any,
kwargs: dict[str, object],
result: object,
call_type: str,
) -> tuple[dict, Any]:
) -> tuple[dict[str, object], object]:
"""
Mirror the pre-call rate-limit snapshot into the SLP so streaming
success callbacks see the same ``x-ratelimit-*`` headers the
@ -3027,8 +3102,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
def _mirror_ratelimit_response_into_logging_payload(
self,
kwargs: Any,
response_obj: Any,
kwargs: object,
response_obj: object,
) -> None:
"""
Copy the stashed ``RateLimitResponse`` into the SLP's
@ -3056,7 +3131,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
)
standard_logging_object["hidden_params"] = hidden_params
response_hidden = getattr(response_obj, "_hidden_params", None)
response_hidden: object = getattr(response_obj, "_hidden_params", None)
if isinstance(response_hidden, dict):
existing = response_hidden.get("additional_headers")
response_hidden["additional_headers"] = self._merge_ratelimit_statuses_into_additional_headers(
@ -3064,7 +3139,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
statuses=statuses,
)
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
async def async_log_failure_event(
self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object
):
"""
On failure: decrement max_parallel_requests and refund the upfront
TPM reservation only against the scopes the reservation actually
@ -3162,19 +3239,17 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
if litellm_proxy_rate_limit_response is not None:
# Update response headers
if hasattr(response, "_hidden_params"):
_hidden_params = getattr(response, "_hidden_params")
_hidden_params_raw: object = getattr(response, "_hidden_params", None)
if isinstance(_hidden_params_raw, BaseModel):
_hidden_params: dict[str, object] | None = _hidden_params_raw.model_dump()
else:
_hidden_params = None
if _hidden_params is not None and (
isinstance(_hidden_params, BaseModel) or isinstance(_hidden_params, dict)
):
if isinstance(_hidden_params, BaseModel):
_hidden_params = _hidden_params.model_dump()
_hidden_params = _str_object_dict_or_none(_hidden_params_raw)
if _hidden_params is not None:
existing_headers = _str_object_dict_or_none(_hidden_params.get("additional_headers"))
_additional_headers = self._merge_ratelimit_statuses_into_additional_headers(
additional_headers=_hidden_params.get("additional_headers", {}) or {},
additional_headers=existing_headers or {},
statuses=litellm_proxy_rate_limit_response["statuses"],
)

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

@ -13,8 +13,8 @@ model/{model_id}/update - PATCH endpoint for model update.
import asyncio
import datetime
import json
from collections.abc import Mapping, Sequence
from typing import Any, Literal, cast
from collections.abc import Callable, Coroutine, Mapping, Sequence
from typing import TYPE_CHECKING, Literal, Protocol, TypeVar, cast
from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
from pydantic import BaseModel, ConfigDict, Field
@ -76,8 +76,99 @@ from litellm.types.router import (
)
from litellm.utils import get_utc_datetime
if TYPE_CHECKING:
from prisma.models import LiteLLM_ProxyModelTable as PrismaProxyModelTable
from prisma.types import LiteLLM_ProxyModelTableWhereInput
router = APIRouter()
_TableT_co = TypeVar("_TableT_co", covariant=True)
class _RepositoryTable(Protocol[_TableT_co]):
@property
def table(self) -> _TableT_co: ...
class _ProxyModelRow(Protocol):
@property
def model_id(self) -> str: ...
@property
def litellm_params(self) -> Mapping[str, object]: ...
def model_dump(self, *, exclude_none: bool = ...) -> dict[str, object]: ...
def model_dump_json(self, *, exclude_none: bool = ...) -> str: ...
class _ProxyModelQueryClient(Protocol):
async def find_many(self, *, where: "LiteLLM_ProxyModelTableWhereInput") -> "Sequence[PrismaProxyModelTable]": ...
class _ProxyModelTableClient(_ProxyModelQueryClient, Protocol):
async def find_unique(self, *, where: Mapping[str, object]) -> "_ProxyModelRow | None": ...
async def create(self, *, data: Mapping[str, object]) -> LiteLLM_ProxyModelTable: ...
async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> "_ProxyModelRow": ...
async def delete(self, *, where: Mapping[str, object]) -> "_ProxyModelRow | None": ...
class _TeamRow(Protocol):
@property
def team_id(self) -> str: ...
@property
def models(self) -> Sequence[str]: ...
def model_dump(self) -> dict[str, object]: ...
class _TeamTableClient(Protocol):
async def find_unique(self, *, where: Mapping[str, object]) -> "_TeamRow | None": ...
async def update(
self,
*,
where: Mapping[str, object],
data: Mapping[str, object],
include: Mapping[str, object] | None = ...,
) -> LiteLLM_TeamTable: ...
class _ModelAliasRow(Protocol):
@property
def id(self) -> int: ...
@property
def model_aliases(self) -> dict[str, str]: ...
@property
def team(self) -> "_TeamRow | None": ...
class _ModelAliasTableClient(Protocol):
async def find_many(self, *, include: Mapping[str, object]) -> "Sequence[_ModelAliasRow]": ...
async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> "_ModelAliasRow": ...
def _proxy_model_table(repository: "_RepositoryTable[_ProxyModelTableClient]") -> "_ProxyModelTableClient":
return repository.table
def _team_table(repository: "_RepositoryTable[_TeamTableClient]") -> "_TeamTableClient":
return repository.table
def _model_alias_table(repository: "_RepositoryTable[_ModelAliasTableClient]") -> "_ModelAliasTableClient":
return repository.table
_loads_json: Callable[[str], object] = json.loads
async def update_team(*args, **kwargs):
"""
@ -96,15 +187,12 @@ class UpdatePublicModelGroupsRequest(BaseModel):
async def get_db_model(model_id: str, prisma_client: PrismaClient) -> Deployment | None:
db_model = cast(
BaseModel | None,
await ModelRepository(prisma_client).table.find_unique(where={"model_id": model_id}),
)
db_model = await _proxy_model_table(ModelRepository(prisma_client)).find_unique(where={"model_id": model_id})
if not db_model:
return None
deployment_pydantic_obj = Deployment(**db_model.model_dump(exclude_none=True))
deployment_pydantic_obj = Deployment.model_validate(db_model.model_dump(exclude_none=True))
return deployment_pydantic_obj
@ -151,7 +239,9 @@ def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> Pr
merged_deployment_dict = DeploymentTypedDict(
model_name=db_model.model_name,
litellm_params=LiteLLMParamsTypedDict(
**db_model.litellm_params.model_dump(exclude_none=True) # type: ignore
**db_model.litellm_params.model_dump(
exclude_none=True
) # any-ok: raw pydantic dump is Any-valued; TypedDict validation would drop extra="allow" params and reject encrypted values
),
model_info=db_model.model_info.model_dump(exclude_none=True),
)
@ -163,7 +253,12 @@ def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> Pr
if updated_patch.litellm_params:
# Encrypt any sensitive values
encrypted_params = {
k: encrypt_value_helper(v) for k, v in updated_patch.litellm_params.model_dump(exclude_none=True).items()
k: encrypt_value_helper(
v
) # any-ok: v is Any from pydantic model_dump; encrypt_value_helper handles non-str at runtime
for k, v in updated_patch.litellm_params.model_dump(
exclude_none=True
).items() # any-ok: same pydantic model_dump seam
}
merged_deployment_dict["litellm_params"].update(encrypted_params) # type: ignore
@ -322,7 +417,7 @@ async def patch_model(
update_data["updated_at"] = cast(str, get_utc_datetime())
# Perform partial update
updated_model = await ModelRepository(prisma_client).table.update(
updated_model = await _proxy_model_table(ModelRepository(prisma_client)).update(
where={"model_id": model_id},
data=update_data,
)
@ -425,7 +520,7 @@ async def _set_model_blocked_status(
param=None,
)
updated_model = await ModelRepository(prisma_client).table.update(
updated_model = await _proxy_model_table(ModelRepository(prisma_client)).update(
where={"model_id": data.model_id},
data={
"blocked": blocked,
@ -459,7 +554,7 @@ async def _set_model_blocked_status(
still_desired=still_desired_ids,
)
return updated_model
return LiteLLM_ProxyModelTable.model_validate(updated_model.model_dump())
except Exception as e:
verbose_proxy_logger.exception(f"Error in model {action}: {e!s}")
@ -549,27 +644,30 @@ async def _add_model_to_db(
# encrypt litellm params #
_litellm_params_dict = model_params.litellm_params.dict(exclude_none=True)
_original_litellm_model_name = model_params.litellm_params.model
for k, v in _litellm_params_dict.items():
encrypted_value = encrypt_value_helper(value=v, new_encryption_key=new_encryption_key)
for (
k,
v,
) in (
_litellm_params_dict.items()
): # any-ok: v is Any from pydantic dict(); encrypt_value_helper handles non-str at runtime
encrypted_value = encrypt_value_helper(
value=v, new_encryption_key=new_encryption_key
) # any-ok: same pydantic dict() seam
model_params.litellm_params[k] = encrypted_value
_data: dict = {
_data: dict[str, str | None] = {
"model_id": model_params.model_info.id,
"model_name": model_params.model_name,
"litellm_params": model_params.litellm_params.model_dump_json(exclude_none=True), # type: ignore
"model_info": model_params.model_info.model_dump_json( # type: ignore
exclude_none=True
),
"litellm_params": model_params.litellm_params.model_dump_json(exclude_none=True),
"model_info": model_params.model_info.model_dump_json(exclude_none=True),
"created_by": user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME,
"updated_by": user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME,
}
if model_params.model_info.id is not None:
_data["model_id"] = model_params.model_info.id
if should_create_model_in_db:
model_response = await ModelRepository(prisma_client).table.create(
data=_data # type: ignore
)
model_response = await _proxy_model_table(ModelRepository(prisma_client)).create(data=_data)
else:
model_response = LiteLLM_ProxyModelTable(**_data)
model_response = LiteLLM_ProxyModelTable.model_validate(_data)
return model_response
@ -757,8 +855,8 @@ async def _setup_new_team_model_assignment(
async def _get_team_deployments(
team_id: str, prisma_client: PrismaClient, table: Any | None = None
) -> list[LiteLLM_ProxyModelTable]:
team_id: str, prisma_client: PrismaClient, table: "_ProxyModelQueryClient | None" = None
) -> "list[PrismaProxyModelTable]":
"""
Fetch all deployments for a given team_id from the database.
@ -773,7 +871,7 @@ async def _get_team_deployments(
existing transaction.
"""
prefix = f"model_name_{team_id}_"
table = table or ModelRepository(prisma_client).table
table = table or _proxy_model_table(ModelRepository(prisma_client))
response = await table.find_many(
where={
"model_name": {"startswith": prefix},
@ -783,7 +881,7 @@ async def _get_team_deployments(
return []
# Confirm team_id in model_info (defensive check)
result = []
result: list[PrismaProxyModelTable] = []
for row in response:
model_info = model_info_as_mapping(row.model_info)
if model_info is not None and model_info.get("team_id") == team_id:
@ -794,7 +892,7 @@ async def _get_team_deployments(
async def delete_team_models(
team_ids: list[str],
prisma_client: PrismaClient,
llm_router: Any | None,
llm_router: Router | None,
) -> list[str]:
"""
Delete every BYOK model owned by the given teams, from the DB and the router.
@ -808,7 +906,7 @@ async def delete_team_models(
Returns the model_ids that were deleted.
"""
deleted_model_ids: list[str] = []
async with prisma_client.db.tx() as tx:
async with prisma_client.tx() as tx:
for team_id in team_ids:
rows = await _get_team_deployments(team_id, prisma_client, table=tx.litellm_proxymodeltable)
model_ids = [row.model_id for row in rows]
@ -908,14 +1006,14 @@ async def _remove_unbacked_team_models(
if not names_to_remove:
return
existing_team_row = await prisma_client.db.litellm_teamtable.find_unique(where={"team_id": team_id})
existing_team_row = await _team_table(TeamRepository(prisma_client)).find_unique(where={"team_id": team_id})
if existing_team_row is None:
return
updated_team_row: LiteLLM_TeamTable = await prisma_client.db.litellm_teamtable.update(
updated_team_row: LiteLLM_TeamTable = await _team_table(TeamRepository(prisma_client)).update(
where={"team_id": team_id},
data={"models": [model for model in existing_team_row.models if model not in names_to_remove]},
include={"object_permission": True}, # type: ignore
include={"object_permission": True},
)
await _refresh_cached_team(
team_row=updated_team_row,
@ -941,7 +1039,7 @@ async def _update_existing_team_model_assignment(
"""
def _get_team_public_model_name(
model_info: dict | str | None,
model_info: object,
) -> str | None:
parsed = model_info_as_mapping(model_info)
if parsed is None:
@ -1050,7 +1148,7 @@ class ModelManagementAuthChecks:
detail={"error": CommonProxyErrors.not_premium_user.value},
)
_existing_team_row = await TeamRepository(prisma_client).table.find_unique(
_existing_team_row = await _team_table(TeamRepository(prisma_client)).find_unique(
where={"team_id": model_params.model_info.team_id}
)
@ -1079,7 +1177,7 @@ class ModelManagementAuthChecks:
) -> Literal[True]:
## Check team model auth
if model_params.model_info is not None and model_params.model_info.team_id is not None:
team_obj_row = await TeamRepository(prisma_client).table.find_unique(
team_obj_row = await _team_table(TeamRepository(prisma_client)).find_unique(
where={"team_id": model_params.model_info.team_id}
)
if team_obj_row is None:
@ -1157,14 +1255,16 @@ async def delete_model(
},
)
model_in_db = await ModelRepository(prisma_client).table.find_unique(where={"model_id": model_info.id})
model_in_db = await _proxy_model_table(ModelRepository(prisma_client)).find_unique(
where={"model_id": model_info.id}
)
if model_in_db is None:
raise HTTPException(
status_code=400,
detail={"error": f"Model with id={model_info.id} not found in db"},
)
model_params = Deployment(**model_in_db.model_dump())
model_params = Deployment.model_validate(model_in_db.model_dump())
await ModelManagementAuthChecks.can_user_make_model_call(
model_params=model_params,
user_api_key_dict=user_api_key_dict,
@ -1180,7 +1280,7 @@ async def delete_model(
- store keys separately
"""
# encrypt litellm params #
result = await ModelRepository(prisma_client).table.delete(where={"model_id": model_info.id})
result = await _proxy_model_table(ModelRepository(prisma_client)).delete(where={"model_id": model_info.id})
if result is None:
raise HTTPException(
@ -1253,9 +1353,9 @@ async def delete_team_model_alias(
Returns:
- List of team id + model alias pairs that were removed
"""
team_model_aliases = await ModelTableRepository(prisma_client).table.find_many(include={"team": True})
tasks = []
removed_model_aliases = []
team_model_aliases = await _model_alias_table(ModelTableRepository(prisma_client)).find_many(include={"team": True})
tasks: list[Coroutine[None, None, _ModelAliasRow]] = []
removed_model_aliases: list[tuple[str, str]] = []
for team_model_alias in team_model_aliases:
model_aliases = team_model_alias.model_aliases # {"alias": "public model name"}
id = team_model_alias.id
@ -1266,7 +1366,7 @@ async def delete_team_model_alias(
removed_model_aliases.append((team_model_alias.team.team_id, key))
del model_aliases[key]
tasks.append(
ModelTableRepository(prisma_client).table.update(
_model_alias_table(ModelTableRepository(prisma_client)).update(
where={"id": id},
data={"model_aliases": json.dumps(model_aliases)},
)
@ -1481,7 +1581,7 @@ async def update_model(
)
_model_id = None
_model_info = getattr(model_params, "model_info", None)
_model_info = model_params.model_info
if _model_info is None:
raise Exception("model_info not provided")
@ -1489,7 +1589,9 @@ async def update_model(
if _model_id is None:
raise Exception("model_info.id not provided")
_existing_litellm_params = await ModelRepository(prisma_client).table.find_unique(where={"model_id": _model_id})
_existing_litellm_params = await _proxy_model_table(ModelRepository(prisma_client)).find_unique(
where={"model_id": _model_id}
)
if _existing_litellm_params is None:
if llm_router is not None and llm_router.get_deployment(model_id=_model_id) is not None:
@ -1499,7 +1601,7 @@ async def update_model(
)
else:
raise Exception("model not found")
deployment = Deployment(**_existing_litellm_params.model_dump())
deployment = Deployment.model_validate(_existing_litellm_params.model_dump())
await ModelManagementAuthChecks.can_user_make_model_call(
model_params=deployment,
@ -1523,15 +1625,23 @@ async def update_model(
_new_litellm_params_dict = model_params.litellm_params.dict(exclude_none=True)
### ENCRYPT PARAMS ###
for k, v in _new_litellm_params_dict.items():
encrypted_value = encrypt_value_helper(value=v)
for (
k,
v,
) in (
_new_litellm_params_dict.items()
): # any-ok: v is Any from pydantic dict(); encrypt_value_helper handles non-str at runtime
encrypted_value = encrypt_value_helper(value=v) # any-ok: same pydantic dict() seam
model_params.litellm_params[k] = encrypted_value
### MERGE WITH EXISTING DATA ###
merged_dictionary = {}
_mp = model_params.litellm_params.dict()
for key, value in _mp.items():
for (
key,
value,
) in _mp.items(): # any-ok: value is Any from pydantic dict(); merged verbatim into the stored params
if value is not None:
merged_dictionary[key] = value
elif key in _existing_litellm_params_dict and _existing_litellm_params_dict[key] is not None:
@ -1539,13 +1649,13 @@ async def update_model(
else:
pass
_data: dict = {
"litellm_params": json.dumps(merged_dictionary), # type: ignore
_data: dict[str, str] = {
"litellm_params": json.dumps(merged_dictionary),
"updated_by": user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME,
}
model_response = await ModelRepository(prisma_client).table.update(
model_response = await _proxy_model_table(ModelRepository(prisma_client)).update(
where={"model_id": _model_id},
data=_data, # type: ignore
data=_data,
)
# Clear cache and reload models (uses config setting or defaults to preserving config models for DB updates)
@ -1787,7 +1897,7 @@ def model_info_as_mapping(model_info: object) -> Mapping[str, object] | None:
if not isinstance(model_info, str):
return None
try:
parsed = json.loads(model_info)
parsed = _loads_json(model_info)
except (TypeError, ValueError):
return None
return parsed if isinstance(parsed, Mapping) else None

2
litellm/proxy/management_endpoints/ui_sso.py
View file

@ -2383,7 +2383,7 @@ async def insert_sso_user(
)
if result_openid and hasattr(result_openid, "provider"):
new_user_request.metadata = {"auth_provider": getattr(result_openid, "provider")}
new_user_request.metadata = {"auth_provider": result_openid.provider}
response = await new_user(
data=new_user_request,

2
litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
View file

@ -725,7 +725,7 @@ async def handle_bedrock_passthrough_router_model(
user_max_tokens: int | None,
user_api_base: str | None,
version: str | None,
) -> Response | StreamingResponse:
) -> object:
"""
Handle Bedrock passthrough for router models (models defined in config.yaml).

355
litellm/proxy/pass_through_endpoints/managed_id_rewriter.py
View file

@ -32,10 +32,12 @@ from __future__ import annotations
import json
import re
from typing import Any
from collections.abc import Mapping
from typing import TYPE_CHECKING, Literal, TypeVar
from urllib.parse import quote, unquote
from fastapi import HTTPException
from pydantic import JsonValue, TypeAdapter
from litellm._logging import verbose_proxy_logger
from litellm.llms.base_llm.managed_resources.isolation import (
@ -43,13 +45,26 @@ from litellm.llms.base_llm.managed_resources.isolation import (
can_access_resource,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.repositories.table_repositories import (
ManagedFileRepository,
ManagedObjectRepository,
)
from litellm.types.llms.openai import OpenAIFileObject
from .managed_id_codec import ManagedIdPayload, decode, is_managed, new_managed_id
from .managed_id_rewriter_types import (
ManagedFileRowLike,
ManagedObjectRowLike,
PassthroughListResponse,
as_file_row,
as_file_rows,
as_managed_file_reader,
as_managed_file_writer,
as_object_row,
as_object_rows,
file_table_actions,
object_table_actions,
)
if TYPE_CHECKING:
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy.utils import PrismaClient
# ---------------------------------------------------------------------------
# Field map
@ -172,7 +187,7 @@ class _RawIdGuardBudget:
def __init__(self, limit: int = _MAX_RAW_ID_GUARD_LOOKUPS) -> None:
self._remaining = limit
self._seen: set = set()
self._seen: set[str] = set()
def reserve(self, raw_id: str) -> bool:
"""Return True when a guard lookup for *raw_id* should run. Returns
@ -196,8 +211,10 @@ class _RawIdGuardBudget:
# to the upstream provider, so each caller only sees IDs they own.
# ---------------------------------------------------------------------------
_ListResourceKind = Literal["files", "batches"]
# Maps (provider, canonical_path) -> "files" | "batches"
_LIST_ROUTE_TABLE: dict[tuple[str, str], str] = {
_LIST_ROUTE_TABLE: dict[tuple[str, str], _ListResourceKind] = {
("openai", "/v1/files"): "files",
("openai", "/v1/batches"): "batches",
("azure", "/v1/files"): "files",
@ -263,8 +280,8 @@ async def _resolve_one(
managed_id: str,
provider: str,
user_api_key_dict: UserAPIKeyAuth,
prisma_client: Any,
managed_files_hook: Any,
prisma_client: PrismaClient | None,
managed_files_hook: CustomLogger | None,
) -> str:
"""
Resolve a single value that may be a passthrough managed ID.
@ -307,7 +324,7 @@ async def _resolve_one(
# File table — use hook's internal cache for speed when available
if managed_files_hook is not None:
try:
file_row = await managed_files_hook.get_unified_file_id(
file_row = await as_managed_file_reader(managed_files_hook).get_unified_file_id(
managed_id,
litellm_parent_otel_span=None,
)
@ -322,8 +339,8 @@ async def _resolve_one(
)
if not found and prisma_client is not None:
try:
db_row = await ManagedFileRepository(prisma_client).table.find_first(
where={"unified_file_id": managed_id}
db_row = as_file_row(
await file_table_actions(prisma_client).find_first(where={"unified_file_id": managed_id})
)
if db_row is not None:
row_created_by = db_row.created_by
@ -338,8 +355,8 @@ async def _resolve_one(
# Object table (batches, responses)
if prisma_client is not None:
try:
obj_row = await ManagedObjectRepository(prisma_client).table.find_first(
where={"unified_object_id": managed_id}
obj_row = as_object_row(
await object_table_actions(prisma_client).find_first(where={"unified_object_id": managed_id})
)
if obj_row is not None:
row_created_by = obj_row.created_by
@ -372,7 +389,7 @@ async def _guard_raw_provider_id(
raw_id: str,
provider: str,
user_api_key_dict: UserAPIKeyAuth,
prisma_client: Any,
prisma_client: PrismaClient | None,
budget: _RawIdGuardBudget | None = None,
) -> None:
"""Deny a raw provider ID that maps to a managed resource the caller does
@ -398,15 +415,15 @@ async def _guard_raw_provider_id(
# id and scope to the current provider in the application layer (same as
# _mint_or_reuse_file's dedup).
try:
candidates = await ManagedFileRepository(prisma_client).table.find_many(
where={"flat_model_file_ids": {"has": raw_id}},
candidates = as_file_rows(
await file_table_actions(prisma_client).find_many(
where={"flat_model_file_ids": {"has": raw_id}},
)
)
except Exception:
verbose_proxy_logger.debug("managed_id_rewriter: raw file-id guard lookup failed", exc_info=True)
return
provider_rows = [
row for row in (candidates or []) if _managed_id_matches_provider(row.unified_file_id, provider)
]
provider_rows = [row for row in candidates if _managed_id_matches_provider(row.unified_file_id, provider)]
if provider_rows and not any(
can_access_resource(user_api_key_dict, row.created_by, row.team_id) for row in provider_rows
):
@ -419,8 +436,10 @@ async def _guard_raw_provider_id(
# Object rows store model_object_id as "passthrough:{provider}:{raw}", so
# the lookup is exact and already provider-scoped.
try:
existing = await ManagedObjectRepository(prisma_client).table.find_first(
where={"model_object_id": f"passthrough:{provider}:{raw_id}"}
existing = as_object_row(
await object_table_actions(prisma_client).find_first(
where={"model_object_id": f"passthrough:{provider}:{raw_id}"}
)
)
except Exception:
verbose_proxy_logger.debug("managed_id_rewriter: raw object-id guard lookup failed", exc_info=True)
@ -434,7 +453,7 @@ async def _guard_raw_provider_id(
# ---------------------------------------------------------------------------
def _build_managed_file_object(snapshot: dict[str, Any] | None, managed_id: str) -> OpenAIFileObject | None:
def _build_managed_file_object(snapshot: Mapping[str, object] | None, managed_id: str) -> OpenAIFileObject | None:
"""Build an ``OpenAIFileObject`` (with the managed ID swapped in) from an
upstream file response so the DB-served list returns the same metadata as a
direct file GET. Returns ``None`` when no usable snapshot is available, in
@ -442,7 +461,7 @@ def _build_managed_file_object(snapshot: dict[str, Any] | None, managed_id: str)
if not snapshot:
return None
try:
return OpenAIFileObject(**{**snapshot, "id": managed_id})
return OpenAIFileObject.model_validate({**snapshot, "id": managed_id})
except Exception:
verbose_proxy_logger.debug(
"managed_id_rewriter: file object snapshot incomplete; storing file row without list metadata",
@ -455,9 +474,9 @@ async def _mint_or_reuse_file(
raw_id: str,
provider: str,
user_api_key_dict: UserAPIKeyAuth,
prisma_client: Any,
managed_files_hook: Any,
file_object_snapshot: dict[str, Any] | None = None,
prisma_client: PrismaClient | None,
managed_files_hook: CustomLogger | None,
file_object_snapshot: Mapping[str, object] | None = None,
is_create_route: bool = True,
) -> str:
"""Return an existing managed file ID or mint + store a new one."""
@ -479,16 +498,16 @@ async def _mint_or_reuse_file(
# reuse a stable row instead of minting duplicate rows on every call.
if prisma_client is not None:
try:
candidates = await ManagedFileRepository(prisma_client).table.find_many(
where={"flat_model_file_ids": {"has": raw_id}},
order={"created_at": "asc"},
candidates = as_file_rows(
await file_table_actions(prisma_client).find_many(
where={"flat_model_file_ids": {"has": raw_id}},
order={"created_at": "asc"},
)
)
except Exception:
candidates = []
verbose_proxy_logger.debug("managed_id_rewriter: file dedup lookup failed", exc_info=True)
provider_rows = [
row for row in (candidates or []) if _managed_id_matches_provider(row.unified_file_id, provider)
]
provider_rows = [row for row in candidates if _managed_id_matches_provider(row.unified_file_id, provider)]
owned_row = next(
(row for row in provider_rows if can_access_resource(user_api_key_dict, row.created_by, row.team_id)),
None,
@ -525,7 +544,7 @@ async def _mint_or_reuse_file(
)
if managed_files_hook is not None:
try:
await managed_files_hook.store_unified_file_id(
await as_managed_file_writer(managed_files_hook).store_unified_file_id(
file_id=managed_id,
file_object=_build_managed_file_object(file_object_snapshot, managed_id),
litellm_parent_otel_span=None,
@ -551,14 +570,15 @@ async def _mint_or_reuse_object(
raw_id: str,
provider: str,
file_purpose: str,
body_snapshot: dict,
body_snapshot: dict[str, object],
user_api_key_dict: UserAPIKeyAuth,
prisma_client: Any,
prisma_client: PrismaClient | None,
is_create_route: bool,
) -> str:
"""Return an existing managed object ID (batch/response) or mint + store one."""
if prisma_client is None:
return raw_id
client = prisma_client
# Namespace raw_id with provider so two providers that happen to issue
# the same raw batch/response ID get distinct rows. The @unique constraint
@ -569,7 +589,7 @@ async def _mint_or_reuse_object(
# f"{purpose}:{provider}:{raw_id}" for the same reason.
namespaced_model_object_id = f"passthrough:{provider}:{raw_id}"
async def _reuse_existing(existing: Any, refresh_snapshot: bool) -> str:
async def _reuse_existing(existing: ManagedObjectRowLike, refresh_snapshot: bool) -> str:
"""Resolve an already-persisted namespaced row: enforce the access
check, optionally refresh the snapshot, and return its managed ID."""
if not can_access_resource(user_api_key_dict, existing.created_by, existing.team_id):
@ -598,7 +618,7 @@ async def _mint_or_reuse_object(
# the batch's latest state (e.g. output_file_id / error_file_id that
# were null at creation but populated once the batch completed).
try:
await ManagedObjectRepository(prisma_client).table.update(
await object_table_actions(client).update(
where={"unified_object_id": existing.unified_object_id},
data={
"file_object": json.dumps(body_snapshot),
@ -618,8 +638,8 @@ async def _mint_or_reuse_object(
# Dedup: look up by the namespaced key — guaranteed unique per provider.
try:
existing = await ManagedObjectRepository(prisma_client).table.find_first(
where={"model_object_id": namespaced_model_object_id}
existing = as_object_row(
await object_table_actions(client).find_first(where={"model_object_id": namespaced_model_object_id})
)
except Exception:
verbose_proxy_logger.debug("managed_id_rewriter: object dedup lookup failed", exc_info=True)
@ -635,7 +655,7 @@ async def _mint_or_reuse_object(
raw_id.split("_", 1)[0],
)
try:
await ManagedObjectRepository(prisma_client).table.upsert(
await object_table_actions(client).upsert(
where={"unified_object_id": managed_id},
data={
"create": {
@ -659,8 +679,8 @@ async def _mint_or_reuse_object(
# the winner's managed ID so both callers converge on one ID instead of
# the loser silently keeping the raw id.
try:
raced = await ManagedObjectRepository(prisma_client).table.find_first(
where={"model_object_id": namespaced_model_object_id}
raced = as_object_row(
await object_table_actions(client).find_first(where={"model_object_id": namespaced_model_object_id})
)
except Exception:
raced = None
@ -677,14 +697,17 @@ async def _mint_or_reuse_object(
return managed_id
_BODY_DICT_ADAPTER: TypeAdapter[dict[str, object]] = TypeAdapter(dict[str, object])
async def rewrite_response_ids(
provider: str,
method: str,
route: str,
body: dict,
user_api_key_dict: UserAPIKeyAuth,
prisma_client: Any,
managed_files_hook: Any,
prisma_client: PrismaClient | None,
managed_files_hook: CustomLogger | None,
) -> dict:
"""
Mint managed IDs for raw provider values listed in
@ -714,7 +737,8 @@ async def rewrite_response_ids(
# creates may degrade to a raw id on a cross-owner collision.
is_create_route = "{" not in canonical
mutated = dict(body) # shallow copy; only return if something changed
typed_body = _BODY_DICT_ADAPTER.validate_python(body)
mutated = dict(typed_body) # shallow copy; only return if something changed
changed = False
def _record(field_name: str, raw_value: str, managed_id: str) -> None:
@ -746,7 +770,7 @@ async def rewrite_response_ids(
managed_files_hook,
# The file's own ``id`` carries the full upstream metadata; nested
# references do not, so only the former is persisted as a snapshot.
file_object_snapshot=body if field_name == "id" else None,
file_object_snapshot=typed_body if field_name == "id" else None,
is_create_route=is_create_route,
)
_record(field_name, raw_value, managed_id)
@ -795,21 +819,24 @@ def is_passthrough_list_route(provider: str, method: str, route: str) -> bool:
return (provider, canonical) in _LIST_ROUTE_TABLE
def _parse_file_object(file_object: Any) -> Any:
_FILE_OBJECT_DICT_ADAPTER: TypeAdapter[dict[str, JsonValue]] = TypeAdapter(dict[str, JsonValue])
def _parse_file_object(file_object: object) -> dict[str, JsonValue] | None:
"""Prisma may return ``Json`` columns as either a parsed dict or the raw
JSON string (depending on driver / row source). Mirror the handling used
elsewhere (see ``openai_files_endpoints/common_utils.py``) so callers can
treat the result uniformly.
"""
if isinstance(file_object, str):
try:
return json.loads(file_object)
except (TypeError, ValueError):
return None
return file_object
try:
if isinstance(file_object, str):
return _FILE_OBJECT_DICT_ADAPTER.validate_json(file_object)
return _FILE_OBJECT_DICT_ADAPTER.validate_python(file_object)
except (TypeError, ValueError):
return None
def _empty_list_response() -> dict[str, Any]:
def _empty_list_response() -> PassthroughListResponse:
return {
"object": "list",
"data": [],
@ -819,7 +846,7 @@ def _empty_list_response() -> dict[str, Any]:
}
def _parse_list_limit(query_params: dict[str, Any] | None) -> tuple[int, int]:
def _parse_list_limit(query_params: Mapping[str, str] | None) -> tuple[int, int]:
params = query_params or {}
try:
raw_limit = int(params.get("limit", 20))
@ -830,17 +857,17 @@ def _parse_list_limit(query_params: dict[str, Any] | None) -> tuple[int, int]:
async def _build_list_where_with_cursor(
prisma_client: Any,
resource_kind: str,
prisma_client: PrismaClient,
resource_kind: _ListResourceKind,
provider: str,
owner_filter: dict[str, Any],
query_params: dict[str, Any] | None,
) -> tuple[dict[str, Any], str]:
owner_filter: dict[str, object],
query_params: Mapping[str, str] | None,
) -> tuple[dict[str, object], str]:
"""Return a Prisma ``where`` clause and fetch order for a list query."""
params = query_params or {}
after_id: str | None = params.get("after")
before_id: str | None = params.get("before")
where: dict[str, Any] = dict(owner_filter)
where: dict[str, object] = dict(owner_filter)
fetch_order = "desc"
cursor_id = after_id or before_id
@ -850,14 +877,14 @@ async def _build_list_where_with_cursor(
if not cursor_id or not _managed_id_matches_provider(cursor_id, provider):
return where, fetch_order
cursor_table = (
ManagedFileRepository(prisma_client).table
if resource_kind == "files"
else ManagedObjectRepository(prisma_client).table
)
cursor_field = "unified_file_id" if resource_kind == "files" else "unified_object_id"
try:
cursor_row = await cursor_table.find_first(where={**owner_filter, cursor_field: cursor_id})
cursor_where: dict[str, object] = {**owner_filter, cursor_field: cursor_id}
cursor_row: ManagedFileRowLike | ManagedObjectRowLike | None
if resource_kind == "files":
cursor_row = as_file_row(await file_table_actions(prisma_client).find_first(where=cursor_where))
else:
cursor_row = as_object_row(await object_table_actions(prisma_client).find_first(where=cursor_where))
if cursor_row is not None:
if after_id:
op = "lt"
@ -867,7 +894,7 @@ async def _build_list_where_with_cursor(
# created_at is not unique, so the boundary must also compare the
# unique id (the secondary sort key) to avoid skipping or repeating
# rows that share the cursor row's timestamp across a page boundary.
boundary = {
boundary: dict[str, object] = {
"OR": [
{"created_at": {op: cursor_row.created_at}},
{
@ -884,110 +911,117 @@ async def _build_list_where_with_cursor(
return where, fetch_order
async def _fetch_list_rows(
prisma_client: Any,
resource_kind: str,
where: dict[str, Any],
async def _fetch_file_list_rows(
prisma_client: PrismaClient,
provider: str,
where: dict[str, object],
fetch_order: str,
fetch_limit: int,
) -> list[Any] | None:
# created_at is not unique, so a second sort on the unique id column gives a
# total order, keeping the limit+1 page boundary and cursor deterministic
# across rows that share a created_at timestamp.
) -> list[ManagedFileRowLike] | None:
"""Fetch list rows scoped to *provider* at the DB level: file rows carry
``_passthrough_provider:{provider}`` in ``flat_model_file_ids`` (see
``_mint_or_reuse_file``), since the file table has no provider column.
Pushing the scope into the query means a single DB round-trip serves the
page, with no application-layer scanning that could truncate large pools.
created_at is not unique, so a second sort on the unique id column gives a
total order, keeping the limit+1 page boundary and cursor deterministic
across rows that share a created_at timestamp.
A DB failure returns ``None`` (fail closed) so the caller serves an empty
page and never falls through to the upstream provider.
"""
scoped_where: dict[str, object] = {
**where,
"flat_model_file_ids": {"has": _passthrough_provider_marker(provider)},
}
try:
if resource_kind == "files":
return await ManagedFileRepository(prisma_client).table.find_many(
where=where,
return as_file_rows(
await file_table_actions(prisma_client).find_many(
where=scoped_where,
order=[{"created_at": fetch_order}, {"unified_file_id": fetch_order}],
take=fetch_limit,
)
return await ManagedObjectRepository(prisma_client).table.find_many(
where={**where, "file_purpose": "batch"},
order=[{"created_at": fetch_order}, {"unified_object_id": fetch_order}],
take=fetch_limit,
)
except Exception:
verbose_proxy_logger.warning("managed_id_rewriter: list DB query failed", exc_info=True)
return None
async def _fetch_provider_scoped_list_rows(
prisma_client: Any,
resource_kind: str,
async def _fetch_batch_list_rows(
prisma_client: PrismaClient,
provider: str,
where: dict[str, Any],
where: dict[str, object],
fetch_order: str,
raw_limit: int,
fetch_limit: int,
) -> tuple[list[Any], bool]:
"""Fetch one page of list rows scoped to *provider* at the DB level.
Both resource kinds carry a provider-distinguishing value that the query
filters on directly: object rows namespace ``model_object_id`` as
``passthrough:{provider}:{raw}`` (see ``_mint_or_reuse_object``) and file
rows carry ``_passthrough_provider:{provider}`` in ``flat_model_file_ids``
(see ``_mint_or_reuse_file``), since the file table has no provider column.
Pushing the scope into the query means a single DB round-trip serves the
page, with no application-layer scanning that could truncate large pools.
A DB failure returns an empty page (fail closed) so the caller never falls
through to the upstream provider.
) -> list[ManagedObjectRowLike] | None:
"""Fetch list rows scoped to *provider* at the DB level: object rows
namespace ``model_object_id`` as ``passthrough:{provider}:{raw}`` (see
``_mint_or_reuse_object``), so the query filters on it directly. Ordering
and failure semantics mirror ``_fetch_file_list_rows``.
"""
scoped_where = dict(where)
if resource_kind == "files":
scoped_where["flat_model_file_ids"] = {"has": _passthrough_provider_marker(provider)}
else:
scoped_where["model_object_id"] = {"startswith": f"passthrough:{provider}:"}
scoped_where: dict[str, object] = {
**where,
"model_object_id": {"startswith": f"passthrough:{provider}:"},
"file_purpose": "batch",
}
try:
return as_object_rows(
await object_table_actions(prisma_client).find_many(
where=scoped_where,
order=[{"created_at": fetch_order}, {"unified_object_id": fetch_order}],
take=fetch_limit,
)
)
except Exception:
verbose_proxy_logger.warning("managed_id_rewriter: list DB query failed", exc_info=True)
return None
rows = await _fetch_list_rows(prisma_client, resource_kind, scoped_where, fetch_order, fetch_limit)
_RowT = TypeVar("_RowT")
def _page_from_rows(rows: list[_RowT] | None, raw_limit: int, fetch_order: str) -> tuple[list[_RowT], bool]:
if rows is None:
return [], False
effective_limit = min(raw_limit, 100)
has_more = len(rows) > effective_limit
page = rows[:effective_limit]
if fetch_order == "asc":
page = list(reversed(page))
return list(reversed(page)), has_more
return page, has_more
def _serialize_file_list_item(row: Any) -> dict[str, Any]:
item: dict[str, Any] = {
def _serialize_file_list_item(row: ManagedFileRowLike) -> dict[str, JsonValue]:
item: dict[str, JsonValue] = {
"id": row.unified_file_id,
"object": "file",
"created_at": int(row.created_at.timestamp()) if row.created_at else None,
}
file_object = _parse_file_object(row.file_object)
if isinstance(file_object, dict):
if file_object is not None:
item.update(file_object)
item["id"] = row.unified_file_id # managed ID always wins over stored raw id
return item
def _serialize_batch_list_item(row: Any) -> dict[str, Any]:
item: dict[str, Any] = {}
def _serialize_batch_list_item(row: ManagedObjectRowLike) -> dict[str, JsonValue]:
item: dict[str, JsonValue] = {}
file_object = _parse_file_object(row.file_object)
if isinstance(file_object, dict):
if file_object is not None:
item.update(file_object)
item["id"] = row.unified_object_id # managed ID always wins
item["object"] = "batch"
return item
def _list_boundary_ids(rows: list[Any], resource_kind: str) -> tuple[str | None, str | None]:
if not rows:
return None, None
id_attr = "unified_file_id" if resource_kind == "files" else "unified_object_id"
return getattr(rows[0], id_attr), getattr(rows[-1], id_attr)
async def list_passthrough_ids_from_db(
provider: str,
route: str,
user_api_key_dict: UserAPIKeyAuth,
prisma_client: Any,
query_params: dict[str, Any] | None = None,
) -> dict[str, Any] | None:
prisma_client: PrismaClient | None,
query_params: Mapping[str, str] | None = None,
) -> PassthroughListResponse | None:
"""Query the DB for managed IDs the caller owns and return an OpenAI-style
paginated list response.
@ -1020,21 +1054,25 @@ async def list_passthrough_ids_from_db(
where, fetch_order = await _build_list_where_with_cursor(
prisma_client, resource_kind, provider, owner_filter, query_params
)
page, has_more = await _fetch_provider_scoped_list_rows(
prisma_client,
resource_kind,
provider,
where,
fetch_order,
raw_limit,
fetch_limit,
)
if resource_kind == "files":
data = [_serialize_file_list_item(row) for row in page]
file_page, has_more = _page_from_rows(
await _fetch_file_list_rows(prisma_client, provider, where, fetch_order, fetch_limit),
raw_limit,
fetch_order,
)
data = [_serialize_file_list_item(row) for row in file_page]
first_id = file_page[0].unified_file_id if file_page else None
last_id = file_page[-1].unified_file_id if file_page else None
else:
data = [_serialize_batch_list_item(row) for row in page]
batch_page, has_more = _page_from_rows(
await _fetch_batch_list_rows(prisma_client, provider, where, fetch_order, fetch_limit),
raw_limit,
fetch_order,
)
data = [_serialize_batch_list_item(row) for row in batch_page]
first_id = batch_page[0].unified_object_id if batch_page else None
last_id = batch_page[-1].unified_object_id if batch_page else None
first_id, last_id = _list_boundary_ids(page, resource_kind)
verbose_proxy_logger.debug(
"managed_id_rewriter: list served from DB provider=%s kind=%s count=%d admin=%s",
provider,
@ -1060,8 +1098,8 @@ async def rewrite_path_ids(
path: str,
provider: str,
user_api_key_dict: UserAPIKeyAuth,
prisma_client: Any,
managed_files_hook: Any,
prisma_client: PrismaClient | None,
managed_files_hook: CustomLogger | None,
) -> str:
"""
Walk URL path segments and resolve any passthrough managed IDs to raw
@ -1092,12 +1130,12 @@ async def rewrite_path_ids(
async def rewrite_query_ids(
params: dict[str, Any] | None,
params: dict[str, object] | None,
provider: str,
user_api_key_dict: UserAPIKeyAuth,
prisma_client: Any,
managed_files_hook: Any,
) -> dict[str, Any] | None:
prisma_client: PrismaClient | None,
managed_files_hook: CustomLogger | None,
) -> dict[str, object] | None:
"""
Walk query param values and resolve any passthrough managed IDs.
Returns *params* unchanged (same object) when nothing is resolved.
@ -1123,13 +1161,19 @@ async def rewrite_query_ids(
return mutated if rewritten_keys else params
def _is_litellm_internal_key(key: object) -> bool:
if not isinstance(key, str):
return False
return key.startswith("litellm_")
async def rewrite_body_ids(
body: dict[str, Any] | None,
body: dict[str, JsonValue] | None,
provider: str,
user_api_key_dict: UserAPIKeyAuth,
prisma_client: Any,
managed_files_hook: Any,
) -> dict[str, Any] | None:
prisma_client: PrismaClient | None,
managed_files_hook: CustomLogger | None,
) -> dict[str, JsonValue] | None:
"""
Recursively walk a request body dict/list and resolve any passthrough
managed IDs. Skips litellm internal keys (``litellm_*``).
@ -1140,15 +1184,15 @@ async def rewrite_body_ids(
budget = _RawIdGuardBudget()
async def _walk(node: Any, depth: int) -> Any:
async def _walk(node: JsonValue, depth: int) -> JsonValue:
if depth >= _MAX_BODY_REWRITE_DEPTH:
return node
if isinstance(node, dict):
result: dict[str, Any] = {}
result: dict[str, JsonValue] = {}
changed_inner = False
for k, v in node.items():
# Skip litellm internal injection keys (e.g. litellm_logging_obj)
if isinstance(k, str) and k.startswith("litellm_"):
if _is_litellm_internal_key(k):
result[k] = v
continue
new_v = await _walk(v, depth + 1)
@ -1156,12 +1200,12 @@ async def rewrite_body_ids(
if new_v is not v:
changed_inner = True
return result if changed_inner else node
elif isinstance(node, list):
if isinstance(node, list):
new_list = [await _walk(item, depth + 1) for item in node]
if any(n is not o for n, o in zip(new_list, node)):
return new_list
return node
elif isinstance(node, str):
if isinstance(node, str):
if is_managed(node):
return await _resolve_one(node, provider, user_api_key_dict, prisma_client, managed_files_hook)
await _guard_raw_provider_id(node, provider, user_api_key_dict, prisma_client, budget)
@ -1169,6 +1213,7 @@ async def rewrite_body_ids(
return node
rewritten = await _walk(body, 0)
if rewritten is not body:
if isinstance(rewritten, dict) and rewritten is not body:
verbose_proxy_logger.debug("managed_id_rewriter: body ids rewritten provider=%s", provider)
return rewritten
return rewritten
return body

148
litellm/proxy/pass_through_endpoints/managed_id_rewriter_types.py Normal file
View file

@ -0,0 +1,148 @@
"""
Typed boundary for ``managed_id_rewriter``'s untyped dependencies.
The managed file/object Prisma tables are reached through
``PrismaTableRepository.table`` (untyped) and the enterprise managed-files
hook arrives as a bare ``CustomLogger``. This module narrows both to
structural protocols once, at the boundary, so the rewriter itself works
with fully typed values.
"""
from __future__ import annotations
from collections.abc import Awaitable, Mapping, Sequence
from datetime import datetime
from typing import TYPE_CHECKING, Literal, Protocol, TypedDict, runtime_checkable
from pydantic import JsonValue
from litellm.repositories.table_repositories import (
ManagedFileRepository,
ManagedObjectRepository,
)
if TYPE_CHECKING:
from litellm.integrations.custom_logger import CustomLogger
from litellm.models.managed_files import LiteLLM_ManagedFileTable
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.utils import PrismaClient
from litellm.types.llms.openai import OpenAIFileObject
@runtime_checkable
class ManagedFileRowLike(Protocol):
unified_file_id: str
created_by: str | None
team_id: str | None
created_at: datetime | None
file_object: object
@runtime_checkable
class ManagedObjectRowLike(Protocol):
unified_object_id: str
created_by: str | None
team_id: str | None
created_at: datetime | None
file_object: object
@runtime_checkable
class ManagedTableActions(Protocol):
def find_first(self, where: Mapping[str, object]) -> Awaitable[object]: ...
def find_many(
self,
where: Mapping[str, object],
order: Sequence[Mapping[str, str]] | Mapping[str, str] | None = None,
take: int | None = None,
) -> Awaitable[Sequence[object]]: ...
def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> Awaitable[object]: ...
def upsert(self, where: Mapping[str, object], data: Mapping[str, Mapping[str, object]]) -> Awaitable[object]: ...
@runtime_checkable
class ManagedFileReaderLike(Protocol):
def get_unified_file_id(
self,
file_id: str,
litellm_parent_otel_span: object = None,
) -> Awaitable[LiteLLM_ManagedFileTable | None]: ...
@runtime_checkable
class ManagedFileWriterLike(Protocol):
def store_unified_file_id(
self,
file_id: str,
file_object: OpenAIFileObject | None,
litellm_parent_otel_span: object,
model_mappings: dict[str, str],
user_api_key_dict: UserAPIKeyAuth,
) -> Awaitable[None]: ...
class PassthroughListResponse(TypedDict):
object: Literal["list"]
data: list[dict[str, JsonValue]]
first_id: str | None
last_id: str | None
has_more: bool
class _TableHolder(Protocol):
@property
def table(self) -> object: ...
def _table_actions(holder: _TableHolder) -> ManagedTableActions:
table = holder.table
if isinstance(table, ManagedTableActions):
return table
raise TypeError("Prisma table object does not expose find/update/upsert actions")
def file_table_actions(prisma_client: PrismaClient) -> ManagedTableActions:
return _table_actions(ManagedFileRepository(prisma_client))
def object_table_actions(prisma_client: PrismaClient) -> ManagedTableActions:
return _table_actions(ManagedObjectRepository(prisma_client))
def as_managed_file_reader(hook: CustomLogger) -> ManagedFileReaderLike:
if isinstance(hook, ManagedFileReaderLike):
return hook
raise TypeError("managed_files hook does not expose get_unified_file_id")
def as_managed_file_writer(hook: CustomLogger) -> ManagedFileWriterLike:
if isinstance(hook, ManagedFileWriterLike):
return hook
raise TypeError("managed_files hook does not expose store_unified_file_id")
def as_file_row(record: object) -> ManagedFileRowLike | None:
if isinstance(record, ManagedFileRowLike):
return record
return None
def as_object_row(record: object) -> ManagedObjectRowLike | None:
if isinstance(record, ManagedObjectRowLike):
return record
return None
def as_file_rows(records: Sequence[object] | None) -> list[ManagedFileRowLike]:
if records is None:
return []
return [record for record in records if isinstance(record, ManagedFileRowLike)]
def as_object_rows(records: Sequence[object] | None) -> list[ManagedObjectRowLike]:
if records is None:
return []
return [record for record in records if isinstance(record, ManagedObjectRowLike)]

365
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
View file

@ -5,10 +5,11 @@ import json
import posixpath
import traceback
from base64 import b64encode
from collections.abc import AsyncGenerator, Mapping
from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence
from datetime import datetime
from functools import partial
from itertools import groupby
from typing import Any, cast
from typing import Literal, Protocol, TypedDict, cast, runtime_checkable
from urllib.parse import urlencode, urlparse
import httpx
@ -23,7 +24,9 @@ from fastapi import (
WebSocket,
status,
)
from fastapi import params as fastapi_params
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError
from starlette.datastructures import UploadFile as StarletteUploadFile
from starlette.websockets import WebSocketState
from websockets.asyncio.client import connect
@ -91,13 +94,94 @@ router = APIRouter()
pass_through_endpoint_logging = PassThroughEndpointLogging()
@runtime_checkable
class RouteRegistrableApp(Protocol):
@property
def routes(self) -> Sequence[object]: ...
def add_api_route(
self,
path: str,
endpoint: Callable[..., object],
*,
methods: list[str] | None = None,
dependencies: Sequence[fastapi_params.Depends] | None = None,
) -> None: ...
@runtime_checkable
class TeamTableClient(Protocol):
def find_unique(self, where: dict[str, str]) -> Awaitable[object]: ...
class RegisteredPassthroughParams(TypedDict):
target: str
custom_headers: dict[str, object] | None
forward_headers: bool | None
merge_query_params: bool | None
default_query_params: dict[str, object] | None
dependencies: Sequence[object] | None
cost_per_request: float | None
guardrails: dict[str, object] | None
timeout: float | None
class RegisteredPassthroughRoute(TypedDict):
endpoint_id: str
path: str
type: Literal["exact", "subpath"]
methods: list[str]
auth: bool
passthrough_params: RegisteredPassthroughParams
# Global registry to track registered pass-through routes and prevent memory leaks
_registered_pass_through_routes: dict[str, dict[str, str | bool | list[str] | dict[str, Any]]] = {}
_registered_pass_through_routes: dict[str, RegisteredPassthroughRoute] = {}
_OBJECT_DICT_ADAPTER: TypeAdapter[dict[str, object]] = TypeAdapter(dict[str, object])
_OPTIONAL_OBJECT_DICT_ADAPTER: TypeAdapter[dict[str, object] | None] = TypeAdapter(dict[str, object] | None)
_EMPTY_ROUTE_PARAMS: Mapping[str, object] = {}
_JSON_VALUE_ADAPTER: TypeAdapter[JsonValue] = TypeAdapter(JsonValue)
_JSON_DICT_ADAPTER: TypeAdapter[dict[str, JsonValue]] = TypeAdapter(dict[str, JsonValue])
_CUSTOM_LOGGER_ADAPTER: TypeAdapter[CustomLogger] = TypeAdapter(
CustomLogger, config=ConfigDict(arbitrary_types_allowed=True)
)
_APP_ADAPTER: TypeAdapter[RouteRegistrableApp] = TypeAdapter(
RouteRegistrableApp, config=ConfigDict(arbitrary_types_allowed=True)
)
_TEAM_TABLE_ADAPTER: TypeAdapter[TeamTableClient] = TypeAdapter(
TeamTableClient, config=ConfigDict(arbitrary_types_allowed=True)
)
_TEAM_METADATA_ADAPTER: TypeAdapter[dict[str, JsonValue] | None] = TypeAdapter(dict[str, JsonValue] | None)
_ALLOWED_ROUTES_ADAPTER: TypeAdapter[list[JsonValue] | dict[str, JsonValue] | str | None] = TypeAdapter(
list[JsonValue] | dict[str, JsonValue] | str | None
)
_ENDPOINT_LIST_ADAPTER: TypeAdapter[list[dict[str, object] | PassThroughGenericEndpoint] | None] = TypeAdapter(
list[dict[str, object] | PassThroughGenericEndpoint] | None
)
def get_response_body(response: httpx.Response) -> dict | None:
def _parse_json_value(raw: str | bytes) -> JsonValue:
return _JSON_VALUE_ADAPTER.validate_python(json.loads(raw))
def _json_dict_or_none(value: JsonValue) -> dict[str, JsonValue] | None:
return value if isinstance(value, dict) else None
def _truthiness_or_none(value: JsonValue) -> bool | None:
return None if value is None else bool(value)
def _read_header(headers: Mapping[str, str], key: str) -> str | None:
return headers.get(key)
def get_response_body(response: httpx.Response) -> dict[str, object] | None:
try:
return response.json()
return _OBJECT_DICT_ADAPTER.validate_python(response.json())
except Exception:
return None
@ -174,9 +258,9 @@ async def chat_completion_pass_through_endpoint(
body = await request.body()
body_str = body.decode()
try:
data = ast.literal_eval(body_str)
data = _OBJECT_DICT_ADAPTER.validate_python(ast.literal_eval(body_str))
except Exception:
data = json.loads(body_str)
data = _OBJECT_DICT_ADAPTER.validate_python(json.loads(body_str))
data["adapter_id"] = adapter_id
@ -848,7 +932,7 @@ async def pass_through_request(
# parsed dict (hooks mutate it, breaking the signature / Content-Length).
# Tolerate request objects without `state` (test fixtures) and only honor
# values httpx accepts for `content=`.
_request_state = getattr(request, "state", None)
_request_state = request.state if hasattr(request, "state") else None
state_raw_body: str | bytes | None = (
getattr(_request_state, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, None)
if _request_state is not None
@ -1125,15 +1209,17 @@ async def pass_through_request(
else:
# SigV4-signed callers (Bedrock) supply the exact pre-signed bytes;
# otherwise httpx encodes the parsed JSON dict as before.
body_kwargs: dict[str, Any] = (
{"content": state_raw_body} if state_raw_body is not None else {"json": _parsed_body}
)
req = async_client.build_request(
build_streaming_request = partial(
async_client.build_request,
request.method,
url,
params=requested_query_params,
headers=headers,
**body_kwargs,
)
req = (
build_streaming_request(content=state_raw_body)
if state_raw_body is not None
else build_streaming_request(json=_OPTIONAL_OBJECT_DICT_ADAPTER.validate_python(_parsed_body))
)
response = await async_client.send(req, stream=stream)
@ -1583,7 +1669,12 @@ def _update_metadata_with_tags_in_header(request: Request, metadata: dict) -> di
async def _parse_request_data_by_content_type(
request: Request,
) -> tuple[Any | None, Any | None, Any | None, Any | None]:
) -> tuple[
dict[str, JsonValue] | dict[str, str] | str | StarletteUploadFile | None,
dict[str, JsonValue] | str | StarletteUploadFile | None,
None,
bool | None,
]:
"""
Parse request data based on content type.
@ -1594,18 +1685,18 @@ async def _parse_request_data_by_content_type(
"""
content_type = request.headers.get("content-type", "")
query_params_data = None
custom_body_data = None
query_params_data: dict[str, JsonValue] | dict[str, str] | str | StarletteUploadFile | None = None
custom_body_data: dict[str, JsonValue] | str | StarletteUploadFile | None = None
file_data = None
stream = None
stream: bool | None = None
if "application/json" in content_type:
# ✅ Handle JSON
try:
body = await request.json()
query_params_data = body.get("query_params")
custom_body_data = body.get("custom_body")
stream = body.get("stream")
body = _JSON_DICT_ADAPTER.validate_python(await request.json())
query_params_data = _json_dict_or_none(body.get("query_params"))
custom_body_data = _json_dict_or_none(body.get("custom_body"))
stream = _truthiness_or_none(body.get("stream"))
except json.JSONDecodeError:
# Handle requests with no body (e.g., DELETE requests)
pass
@ -1613,13 +1704,14 @@ async def _parse_request_data_by_content_type(
# ✅ Try to parse as JSON first (handles misconfigured clients sending JSON with multipart content-type)
# If that fails, skip parsing - pass_through_request will handle actual multipart
try:
body = await request.json()
body = _JSON_DICT_ADAPTER.validate_python(await request.json())
# Successfully parsed as JSON - treat as JSON body
query_params_data = body.get("query_params")
custom_body_data = body.get("custom_body")
stream = body.get("stream")
raw_custom_body = body.get("custom_body")
query_params_data = _json_dict_or_none(body.get("query_params"))
custom_body_data = _json_dict_or_none(raw_custom_body)
stream = _truthiness_or_none(body.get("stream"))
# If custom_body is not set, use the entire body
if custom_body_data is None and body:
if raw_custom_body is None and body:
custom_body_data = body
except (json.JSONDecodeError, Exception):
# Not JSON - this is actual multipart data
@ -1643,17 +1735,17 @@ async def _parse_request_data_by_content_type(
def create_pass_through_route(
endpoint,
target: str,
custom_headers: Mapping[str, Any] | None = None,
custom_headers: Mapping[str, object] | None = None,
_forward_headers: bool | None = False,
_merge_query_params: bool | None = False,
dependencies: list | None = None,
dependencies: Sequence[fastapi_params.Depends] | None = None,
include_subpath: bool | None = False,
cost_per_request: float | None = None,
custom_llm_provider: str | None = None,
is_streaming_request: bool | None = False,
query_params: dict | None = None,
default_query_params: dict | None = None,
guardrails: dict[str, Any] | None = None,
guardrails: dict[str, object] | None = None,
config_file_path: str | None = None,
timeout: float | None = None,
):
@ -1665,7 +1757,9 @@ def create_pass_through_route(
if isinstance(target, CustomLogger):
adapter = target
else:
adapter = get_instance_fn(value=target, config_file_path=config_file_path)
adapter = _CUSTOM_LOGGER_ADAPTER.validate_python(
get_instance_fn(value=target, config_file_path=config_file_path)
)
adapter_id = str(uuid.uuid4())
litellm.adapters = [{"id": adapter_id, "adapter": adapter}]
@ -1725,7 +1819,7 @@ def create_pass_through_route(
status_code=status.HTTP_405_METHOD_NOT_ALLOWED,
detail=f"Method {request.method} is not allowed for pass-through endpoint {path}.",
)
target_params = {
target_params: dict[str, object] = {
"target": target,
"custom_headers": custom_headers,
"forward_headers": _forward_headers,
@ -1736,7 +1830,7 @@ def create_pass_through_route(
}
if passthrough_params is not None:
target_params.update(passthrough_params.get("passthrough_params", {}))
target_params.update(passthrough_params.get("passthrough_params", _EMPTY_ROUTE_PARAMS))
# Extract and cast parameters with proper types
param_target = target_params.get("target") or target
@ -1757,12 +1851,18 @@ def create_pass_through_route(
# Ensure custom_headers is a dict. Botocore returns a HeadersDict
# for SigV4-prepared requests, which is a Mapping but not a dict.
headers_dict = dict(param_custom_headers) if isinstance(param_custom_headers, Mapping) else {}
headers_dict = (
_OBJECT_DICT_ADAPTER.validate_python(param_custom_headers)
if isinstance(param_custom_headers, Mapping)
else {}
)
# Ensure query_params and custom_body are dicts or None
final_query_params = query_params_data if isinstance(query_params_data, dict) else {}
final_query_params: dict[str, object] = (
{key: value for key, value in query_params_data.items()} if isinstance(query_params_data, dict) else {}
)
if query_params:
final_query_params.update(query_params)
final_query_params.update(_OBJECT_DICT_ADAPTER.validate_python(query_params))
# Programmatic callers set LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY on
# request.state (see Bedrock proxy). Parsed JSON envelope otherwise.
state_custom_body: dict | None = getattr(
@ -1884,7 +1984,7 @@ async def websocket_passthrough_request(
# Initialize tracking variables
start_time = datetime.now()
websocket_messages: list[dict[str, Any]] = []
websocket_messages: list[JsonValue] = []
litellm_call_id = str(uuid.uuid4())
verbose_proxy_logger.info(f"WebSocket passthrough ({endpoint}): Starting WebSocket connection to {target}")
@ -1977,10 +2077,9 @@ async def websocket_passthrough_request(
)
### CALL HOOKS ### - modify incoming data / reject request before calling the model
websocket_data: dict[str, Any] = {}
websocket_data = await proxy_logging_obj.pre_call_hook(
await proxy_logging_obj.pre_call_hook(
user_api_key_dict=user_api_key_dict,
data=websocket_data,
data={},
call_type="pass_through_endpoint",
)
@ -2007,14 +2106,14 @@ async def websocket_passthrough_request(
text_data = message.get("text")
bytes_data = message.get("bytes")
if text_data is not None:
if isinstance(text_data, str):
# Try to extract model from client setup message for Vertex AI Live
if endpoint and "/vertex_ai/live" in endpoint:
verbose_proxy_logger.debug(
f"WebSocket passthrough ({endpoint}): Processing client message for model extraction"
)
try:
client_message = json.loads(text_data)
client_message = _parse_json_value(text_data)
if isinstance(client_message, dict) and "setup" in client_message:
setup_data = client_message["setup"]
verbose_proxy_logger.debug(
@ -2044,14 +2143,14 @@ async def websocket_passthrough_request(
verbose_proxy_logger.debug(
f"WebSocket passthrough ({endpoint}): Client message does not contain setup data"
)
except (json.JSONDecodeError, KeyError, TypeError) as e:
except (json.JSONDecodeError, ValidationError, KeyError, TypeError) as e:
verbose_proxy_logger.debug(
f"WebSocket passthrough ({endpoint}): Client message is not a valid setup message: {e}"
)
# Not a JSON message or doesn't contain setup data
await upstream_ws.send(text_data)
elif bytes_data is not None:
elif isinstance(bytes_data, bytes):
await upstream_ws.send(bytes_data)
except asyncio.CancelledError:
raise
@ -2069,7 +2168,7 @@ async def websocket_passthrough_request(
# Ensure raw_response is bytes before decoding
if isinstance(raw_response, str):
raw_response = raw_response.encode("ascii")
setup_response = json.loads(raw_response.decode("ascii"))
setup_response = _parse_json_value(raw_response.decode("ascii"))
verbose_proxy_logger.debug(f"Setup response: {setup_response}")
# Extract model and provider from setup response for Vertex AI Live
@ -2106,17 +2205,17 @@ async def websocket_passthrough_request(
await websocket.send_bytes(upstream_message)
# Parse and collect for cost tracking
try:
message_data = json.loads(upstream_message.decode())
message_data = _parse_json_value(upstream_message.decode())
websocket_messages.append(message_data)
except (json.JSONDecodeError, UnicodeDecodeError):
except (json.JSONDecodeError, ValidationError, UnicodeDecodeError):
pass
else:
await websocket.send_text(upstream_message)
# Parse and collect for cost tracking
try:
message_data = json.loads(upstream_message)
message_data = _parse_json_value(upstream_message)
websocket_messages.append(message_data)
except json.JSONDecodeError:
except (json.JSONDecodeError, ValidationError):
pass
except (ConnectionClosedOK, ConnectionClosedError) as e:
@ -2272,7 +2371,7 @@ async def websocket_passthrough_request(
def _is_streaming_response(response: httpx.Response) -> bool:
_content_type = response.headers.get("content-type")
_content_type = _read_header(response.headers, "content-type")
if _content_type is not None and "text/event-stream" in _content_type:
return True
return False
@ -2290,7 +2389,7 @@ def _should_buffer_passthrough_response(response: httpx.Response) -> bool:
"""
if response.status_code >= 400:
return True
media_type = response.headers.get("content-type", "").split(";")[0].strip().lower()
media_type = (_read_header(response.headers, "content-type") or "").split(";")[0].strip().lower()
return media_type in ("", "application/json") or media_type.endswith("+json")
@ -2342,7 +2441,7 @@ async def _relay_passthrough_response_bytes(
)
def _extract_model_from_vertex_ai_setup(setup_response: dict) -> str | None:
def _extract_model_from_vertex_ai_setup(setup_response: JsonValue) -> str | None:
"""
Extract the model name from Vertex AI Live setup response.
@ -2381,7 +2480,7 @@ class SafeRouteAdder:
"""
@staticmethod
def _is_path_registered(app: FastAPI, path: str, methods: list[str]) -> bool:
def _is_path_registered(app: RouteRegistrableApp, path: str, methods: list[str]) -> bool:
"""
Check if a path with any of the specified methods is already registered on the app.
@ -2406,11 +2505,11 @@ class SafeRouteAdder:
@staticmethod
def add_api_route_if_not_exists(
app: FastAPI,
app: RouteRegistrableApp,
path: str,
endpoint: Any,
endpoint: Callable[..., object],
methods: list[str],
dependencies: list | None = None,
dependencies: Sequence[fastapi_params.Depends] | None = None,
) -> bool:
"""
Add an API route to the app only if it doesn't already exist.
@ -2450,13 +2549,13 @@ class SafeRouteAdder:
class InitPassThroughEndpointHelpers:
@staticmethod
def add_exact_path_route(
app: FastAPI,
app: RouteRegistrableApp,
path: str,
target: str,
custom_headers: dict | None,
forward_headers: bool | None,
merge_query_params: bool | None,
dependencies: list | None,
dependencies: Sequence[fastapi_params.Depends] | None,
cost_per_request: float | None,
endpoint_id: str,
guardrails: dict | None = None,
@ -2533,13 +2632,13 @@ class InitPassThroughEndpointHelpers:
@staticmethod
def add_subpath_route(
app: FastAPI,
app: RouteRegistrableApp,
path: str,
target: str,
custom_headers: dict | None,
forward_headers: bool | None,
merge_query_params: bool | None,
dependencies: list | None,
dependencies: Sequence[fastapi_params.Depends] | None,
cost_per_request: float | None,
endpoint_id: str,
guardrails: dict | None = None,
@ -2624,15 +2723,14 @@ class InitPassThroughEndpointHelpers:
]
for key in keys_to_remove:
route_info = _registered_pass_through_routes[key]
path = route_info.get("path")
if isinstance(path, str):
openai_routes = LiteLLMRoutes.openai_routes.value
if path in openai_routes:
openai_routes.remove(path)
if route_info.get("type") == "subpath":
wildcard_path = path.rstrip("/") + "/*"
if wildcard_path in openai_routes:
openai_routes.remove(wildcard_path)
path = route_info["path"]
openai_routes = LiteLLMRoutes.openai_routes.value
if path in openai_routes:
openai_routes.remove(path)
if route_info.get("type") == "subpath":
wildcard_path = path.rstrip("/") + "/*"
if wildcard_path in openai_routes:
openai_routes.remove(wildcard_path)
del _registered_pass_through_routes[key]
verbose_proxy_logger.debug("Removed pass-through route from registry: %s", key)
@ -2699,7 +2797,7 @@ class InitPassThroughEndpointHelpers:
return False
@staticmethod
def get_registered_pass_through_route(route: str, method: str | None = None) -> dict[str, Any] | None:
def get_registered_pass_through_route(route: str, method: str | None = None) -> RegisteredPassthroughRoute | None:
"""Get passthrough params for a given route and optionally filter by HTTP method"""
comparison_route = InitPassThroughEndpointHelpers._route_for_registry_lookup(route)
for key in _registered_pass_through_routes:
@ -2712,7 +2810,7 @@ class InitPassThroughEndpointHelpers:
# but keep supporting test fixtures / older registry entries that
# only encoded methods in the route key.
methods_entry = _registered_pass_through_routes[key].get("methods", [])
route_methods: list[str] = methods_entry if isinstance(methods_entry, list) else []
route_methods: list[str] = methods_entry if methods_entry else []
if not route_methods and len(parts) == 4:
route_methods = parts[3].split(",")
@ -2740,14 +2838,27 @@ def _get_combined_pass_through_endpoints(
return pass_through_endpoints + config_pass_through_endpoints
class _RegisteredEndpointFields(BaseModel):
path: str | None = None
target: str | None = None
headers: dict[str, object] | None = None
forward_headers: bool | None = None
merge_query_params: bool | None = None
default_query_params: dict[str, object] | None = None
guardrails: dict[str, object] | None = None
methods: list[str] | None = None
cost_per_request: float | None = None
timeout: float | None = None
async def _register_pass_through_endpoint(
endpoint: dict[str, Any] | PassThroughGenericEndpoint,
endpoint: dict[str, object] | PassThroughGenericEndpoint,
app: FastAPI,
premium_user: bool,
visited_endpoints: set[str],
config_file_path: str | None = None,
) -> None:
endpoint_data: dict[str, Any]
endpoint_data: dict[str, object]
if isinstance(endpoint, PassThroughGenericEndpoint):
endpoint_data = endpoint.model_dump()
else:
@ -2757,17 +2868,18 @@ async def _register_pass_through_endpoint(
endpoint_data["id"] = str(uuid.uuid4())
endpoint_id = cast(str, endpoint_data["id"])
target = endpoint_data.get("target")
path = endpoint_data.get("path")
endpoint_fields = _RegisteredEndpointFields.model_validate(endpoint_data)
target = endpoint_fields.target
path = endpoint_fields.path
if path is None:
raise ValueError("Path is required for pass-through endpoint")
custom_headers = await set_env_variables_in_header(custom_headers=endpoint_data.get("headers"))
forward_headers = endpoint_data.get("forward_headers")
merge_query_params = endpoint_data.get("merge_query_params")
default_query_params = endpoint_data.get("default_query_params")
custom_headers = await set_env_variables_in_header(custom_headers=endpoint_fields.headers)
forward_headers = endpoint_fields.forward_headers
merge_query_params = endpoint_fields.merge_query_params
default_query_params = endpoint_fields.default_query_params
auth = endpoint_data.get("auth")
dependencies = None
dependencies: list[fastapi_params.Depends] | None = None
auth_enforced = auth is not None and str(auth).lower() == "true"
if auth_enforced:
@ -2782,10 +2894,10 @@ async def _register_pass_through_endpoint(
if target is None:
return
guardrails = endpoint_data.get("guardrails")
methods = endpoint_data.get("methods")
cost_per_request = endpoint_data.get("cost_per_request")
timeout = endpoint_data.get("timeout")
guardrails = endpoint_fields.guardrails
methods = endpoint_fields.methods
cost_per_request = endpoint_fields.cost_per_request
timeout = endpoint_fields.timeout
verbose_proxy_logger.debug("Initializing pass through endpoint: %s (ID: %s)", path, endpoint_id)
InitPassThroughEndpointHelpers.add_exact_path_route(
@ -2925,12 +3037,12 @@ def _get_pass_through_endpoints_from_config() -> list[PassThroughGenericEndpoint
if isinstance(endpoint, dict):
endpoint_dict = dict(endpoint)
endpoint_dict["is_from_config"] = True
returned_endpoints.append(PassThroughGenericEndpoint(**endpoint_dict))
returned_endpoints.append(PassThroughGenericEndpoint.model_validate(endpoint_dict))
elif isinstance(endpoint, PassThroughGenericEndpoint):
# Create a copy with is_from_config=True
endpoint_dict = endpoint.model_dump()
endpoint_dict["is_from_config"] = True
returned_endpoints.append(PassThroughGenericEndpoint(**endpoint_dict))
returned_endpoints.append(PassThroughGenericEndpoint.model_validate(endpoint_dict))
except ValidationError as e:
verbose_proxy_logger.warning(
"Skipping malformed pass-through endpoint from config: %s",
@ -2954,10 +3066,10 @@ async def _get_pass_through_endpoints_from_db(
response: ConfigFieldInfo = await get_config_general_settings(
field_name="pass_through_endpoints", user_api_key_dict=user_api_key_dict
)
pass_through_endpoint_data = _ENDPOINT_LIST_ADAPTER.validate_python(getattr(response, "field_value"))
except Exception:
return []
pass_through_endpoint_data: list | None = response.field_value
if pass_through_endpoint_data is None:
return []
@ -2965,14 +3077,11 @@ async def _get_pass_through_endpoints_from_db(
if endpoint_id is None:
# Return all endpoints from DB, mark as not from config
for endpoint in pass_through_endpoint_data:
if isinstance(endpoint, dict):
endpoint_dict = dict(endpoint)
endpoint_dict["is_from_config"] = False
returned_endpoints.append(PassThroughGenericEndpoint(**endpoint_dict))
elif isinstance(endpoint, PassThroughGenericEndpoint):
endpoint_dict = endpoint.model_dump()
endpoint_dict["is_from_config"] = False
returned_endpoints.append(PassThroughGenericEndpoint(**endpoint_dict))
endpoint_dict = (
endpoint.model_dump() if isinstance(endpoint, PassThroughGenericEndpoint) else dict(endpoint)
)
endpoint_dict["is_from_config"] = False
returned_endpoints.append(PassThroughGenericEndpoint.model_validate(endpoint_dict))
else:
# Find specific endpoint by ID
found_endpoint = _find_endpoint_by_id(pass_through_endpoint_data, endpoint_id)
@ -2983,7 +3092,7 @@ async def _get_pass_through_endpoints_from_db(
else dict(found_endpoint)
)
endpoint_dict["is_from_config"] = False
returned_endpoints.append(PassThroughGenericEndpoint(**endpoint_dict))
returned_endpoints.append(PassThroughGenericEndpoint.model_validate(endpoint_dict))
return returned_endpoints
@ -3008,9 +3117,8 @@ async def _filter_endpoints_by_team_allowed_routes(
HTTPException: If team is not found
"""
# retrieve team from db
team = await TeamRepository(prisma_client).table.find_unique(
where={"team_id": team_id},
)
team_table = _TEAM_TABLE_ADAPTER.validate_python(getattr(TeamRepository(prisma_client), "table"))
team = await team_table.find_unique(where={"team_id": team_id})
if team is None:
raise HTTPException(
status_code=404,
@ -3018,14 +3126,15 @@ async def _filter_endpoints_by_team_allowed_routes(
)
# retrieve team metadata
team_metadata = team.metadata
if team_metadata is not None and team_metadata.get("allowed_passthrough_routes") is not None:
team_metadata = _TEAM_METADATA_ADAPTER.validate_python(getattr(team, "metadata"))
allowed_routes = (
_ALLOWED_ROUTES_ADAPTER.validate_python(team_metadata.get("allowed_passthrough_routes"))
if team_metadata is not None
else None
)
if allowed_routes is not None:
## FILTER pass_through_endpoints by allowed_passthrough_routes
pass_through_endpoints = [
endpoint
for endpoint in pass_through_endpoints
if endpoint.path in team_metadata.get("allowed_passthrough_routes")
]
pass_through_endpoints = [endpoint for endpoint in pass_through_endpoints if endpoint.path in allowed_routes]
return pass_through_endpoints
@ -3115,7 +3224,7 @@ async def update_pass_through_endpoints(
detail={"error": "No pass-through endpoints found"},
)
pass_through_endpoint_data: list | None = response.field_value
pass_through_endpoint_data = _ENDPOINT_LIST_ADAPTER.validate_python(getattr(response, "field_value"))
if pass_through_endpoint_data is None:
raise HTTPException(
status_code=404,
@ -3134,7 +3243,7 @@ async def update_pass_through_endpoints(
# Find the index for updating the list
endpoint_index = None
for idx, endpoint in enumerate(pass_through_endpoint_data):
_endpoint = PassThroughGenericEndpoint(**endpoint) if isinstance(endpoint, dict) else endpoint
_endpoint = PassThroughGenericEndpoint.model_validate(endpoint) if isinstance(endpoint, dict) else endpoint
if _endpoint.id == endpoint_id:
endpoint_index = idx
break
@ -3165,7 +3274,7 @@ async def update_pass_through_endpoints(
endpoint_dict.pop("is_from_config", None)
# Create updated endpoint object
updated_endpoint = PassThroughGenericEndpoint(**endpoint_dict)
updated_endpoint = PassThroughGenericEndpoint.model_validate(endpoint_dict)
# Update the list
pass_through_endpoint_data[endpoint_index] = endpoint_dict
@ -3186,9 +3295,10 @@ async def update_pass_through_endpoints(
_custom_headers: dict | None = updated_endpoint.headers or {}
_custom_headers = await set_env_variables_in_header(custom_headers=_custom_headers)
request_app = _APP_ADAPTER.validate_python(getattr(request, "app"))
if updated_endpoint.include_subpath:
InitPassThroughEndpointHelpers.add_subpath_route(
app=request.app,
app=request_app,
path=updated_endpoint.path,
target=updated_endpoint.target,
custom_headers=_custom_headers,
@ -3205,7 +3315,7 @@ async def update_pass_through_endpoints(
)
else:
InitPassThroughEndpointHelpers.add_exact_path_route(
app=request.app,
app=request_app,
path=updated_endpoint.path,
target=updated_endpoint.target,
custom_headers=_custom_headers,
@ -3257,29 +3367,28 @@ async def create_pass_through_endpoints(
if data_dict.get("id") is None:
data_dict["id"] = str(uuid.uuid4())
if response.field_value is None:
response.field_value = [data_dict]
elif isinstance(response.field_value, list):
response.field_value.append(data_dict)
existing_field_value = _ENDPOINT_LIST_ADAPTER.validate_python(getattr(response, "field_value"))
updated_field_value = [data_dict] if existing_field_value is None else [*existing_field_value, data_dict]
## Update db
updated_data = ConfigFieldUpdate(
field_name="pass_through_endpoints",
field_value=response.field_value,
field_value=updated_field_value,
config_type="general_settings",
)
await update_config_general_settings(data=updated_data, user_api_key_dict=user_api_key_dict)
# Return the created endpoint with the generated ID
created_endpoint = PassThroughGenericEndpoint(**data_dict)
created_endpoint = PassThroughGenericEndpoint.model_validate(data_dict)
# Register the new route
_custom_headers: dict | None = created_endpoint.headers or {}
_custom_headers = await set_env_variables_in_header(custom_headers=_custom_headers)
request_app = _APP_ADAPTER.validate_python(getattr(request, "app"))
if created_endpoint.include_subpath:
InitPassThroughEndpointHelpers.add_subpath_route(
app=request.app,
app=request_app,
path=created_endpoint.path,
target=created_endpoint.target,
custom_headers=_custom_headers,
@ -3296,7 +3405,7 @@ async def create_pass_through_endpoints(
)
else:
InitPassThroughEndpointHelpers.add_exact_path_route(
app=request.app,
app=request_app,
path=created_endpoint.path,
target=created_endpoint.target,
custom_headers=_custom_headers,
@ -3344,8 +3453,8 @@ async def delete_pass_through_endpoints(
response = ConfigFieldInfo(field_name="pass_through_endpoints", field_value=None)
## Update field by removing endpoint
pass_through_endpoint_data: list | None = response.field_value
if response.field_value is None or pass_through_endpoint_data is None:
pass_through_endpoint_data = _ENDPOINT_LIST_ADAPTER.validate_python(getattr(response, "field_value"))
if pass_through_endpoint_data is None:
raise HTTPException(
status_code=400,
detail={"error": "There are no pass-through endpoints setup."},
@ -3363,7 +3472,7 @@ async def delete_pass_through_endpoints(
# Find the index for deleting from the list
endpoint_index = None
for idx, endpoint in enumerate(pass_through_endpoint_data):
_endpoint = PassThroughGenericEndpoint(**endpoint) if isinstance(endpoint, dict) else endpoint
_endpoint = PassThroughGenericEndpoint.model_validate(endpoint) if isinstance(endpoint, dict) else endpoint
if _endpoint.id == endpoint_id:
endpoint_index = idx
break
@ -3393,7 +3502,7 @@ async def delete_pass_through_endpoints(
def _find_endpoint_by_id(
endpoints_data: list,
endpoints_data: list[dict[str, object] | PassThroughGenericEndpoint],
endpoint_id: str,
) -> PassThroughGenericEndpoint | None:
"""
@ -3407,14 +3516,14 @@ def _find_endpoint_by_id(
Found endpoint or None if not found
"""
for endpoint in endpoints_data:
_endpoint: PassThroughGenericEndpoint | None = None
if isinstance(endpoint, dict):
_endpoint = PassThroughGenericEndpoint(**endpoint)
elif isinstance(endpoint, PassThroughGenericEndpoint):
_endpoint = endpoint
_endpoint = (
endpoint
if isinstance(endpoint, PassThroughGenericEndpoint)
else PassThroughGenericEndpoint.model_validate(endpoint)
)
# Only compare IDs to IDs
if _endpoint is not None and _endpoint.id == endpoint_id:
if _endpoint.id == endpoint_id:
return _endpoint
return None

2
litellm/responses/litellm_completion_transformation/session_handler.py
View file

@ -137,7 +137,7 @@ class ResponsesSessionHandler:
model_response = ModelResponse(**_response_output)
for choice in model_response.choices:
if hasattr(choice, "message"):
chat_completion_message_history.append(getattr(choice, "message"))
chat_completion_message_history.append(choice.message)
return chat_completion_message_history
@staticmethod

523
litellm/responses/litellm_completion_transformation/transformation.py
View file

@ -4,12 +4,13 @@ Handles transforming from Responses API -> LiteLLM completion (Chat Completion
import json
import re
from collections.abc import Sequence
from typing import Any, Literal, cast
from collections.abc import Iterable, Mapping, Sequence
from typing import Literal, Protocol, cast
from openai.types.responses import ResponseFunctionToolCall
from openai.types.responses.response_create_params import ResponseInputParam
from openai.types.responses.tool_param import FunctionToolParam
from pydantic import BaseModel, InstanceOf, TypeAdapter
from typing_extensions import TypedDict
from litellm._logging import verbose_logger
@ -27,6 +28,7 @@ from litellm.types.llms.openai import (
ChatCompletionImageUrlObject,
ChatCompletionResponseMessage,
ChatCompletionSystemMessage,
ChatCompletionTextObject,
ChatCompletionToolCallChunk,
ChatCompletionToolCallFunctionChunk,
ChatCompletionToolMessage,
@ -72,6 +74,52 @@ from .custom_tools import (
unwrap_custom_tool_arguments,
)
_STR_OBJECT_DICT_ADAPTER: TypeAdapter[dict[str, object]] = TypeAdapter(dict[str, object])
_STR_OBJECT_DICT_INSTANCE_ADAPTER = TypeAdapter(InstanceOf[dict[str, object]])
_OBJECT_LIST_INSTANCE_ADAPTER = TypeAdapter(InstanceOf[list[object]])
_OBJECT_SEQUENCE_INSTANCE_ADAPTER = TypeAdapter(InstanceOf[Sequence[object]])
_OBJECT_ITERABLE_INSTANCE_ADAPTER = TypeAdapter(InstanceOf[Iterable[object]])
_TEXT_OBJECT_LIST_INSTANCE_ADAPTER = TypeAdapter(InstanceOf[list[ChatCompletionTextObject]])
def _as_str_object_dict(value: object) -> dict[str, object] | None:
return _STR_OBJECT_DICT_INSTANCE_ADAPTER.validate_python(value) if isinstance(value, dict) else None
def _as_object_list(value: object) -> list[object] | None:
return _OBJECT_LIST_INSTANCE_ADAPTER.validate_python(value) if isinstance(value, list) else None
def _as_object_sequence(value: object) -> Sequence[object] | None:
if isinstance(value, (str, bytes)) or not isinstance(value, Sequence):
return None
return _OBJECT_SEQUENCE_INSTANCE_ADAPTER.validate_python(value)
def _as_object_iterable(value: object) -> Iterable[object] | None:
if isinstance(value, (str, bytes)) or not isinstance(value, Iterable):
return None
return _OBJECT_ITERABLE_INSTANCE_ADAPTER.validate_python(value)
def _as_tool_message_content(value: object) -> str | list[ChatCompletionTextObject] | None:
if isinstance(value, str):
return value
return _TEXT_OBJECT_LIST_INSTANCE_ADAPTER.validate_python(value) if isinstance(value, list) else None
def _model_as_dict(model: BaseModel) -> dict[str, object]:
return _STR_OBJECT_DICT_ADAPTER.validate_python(model.model_dump())
class SupportsApplyPatchOperation(Protocol):
@property
def call_id(self) -> str: ...
@property
def operation(self) -> BaseModel: ...
########### Initialize Classes used for Responses API ###########
TOOL_CALLS_CACHE = InMemoryCache()
@ -116,8 +164,8 @@ class LiteLLMCompletionResponsesConfig:
@staticmethod
def _transform_tool_choice(
tool_choice: Any,
) -> str | dict[str, Any] | None:
tool_choice: str | Mapping[str, object] | None,
) -> str | Mapping[str, object] | None:
"""
Transform tool_choice from various formats to OpenAI Chat Completion format.
@ -145,7 +193,8 @@ class LiteLLMCompletionResponsesConfig:
tool_choice_type = tool_choice.get("type")
# If it has a function with name, it's standard OpenAI format - pass through
if tool_choice.get("function") and tool_choice.get("function", {}).get("name"):
function_dict = _as_str_object_dict(tool_choice.get("function"))
if function_dict is not None and function_dict.get("name"):
return tool_choice
# Handle Cursor IDE dict formats without function name
@ -189,7 +238,7 @@ class LiteLLMCompletionResponsesConfig:
responses_api_request: ResponsesAPIOptionalRequestParams,
custom_llm_provider: str | None = None,
stream: bool | None = None,
extra_headers: dict[str, Any] | None = None,
extra_headers: dict[str, str] | None = None,
**kwargs,
) -> dict:
"""
@ -446,7 +495,9 @@ class LiteLLMCompletionResponsesConfig:
if not chat_completion_messages:
continue
deduped_in_place: list[Any] = []
deduped_in_place: list[
AllMessageValues | GenericChatCompletionMessage | ChatCompletionResponseMessage
] = []
for m in chat_completion_messages:
role = ""
if isinstance(m, dict):
@ -456,19 +507,17 @@ class LiteLLMCompletionResponsesConfig:
# Drop assistant tool_calls wrappers if we already have this call_id
if role == "assistant":
tool_calls: Any = (
tool_calls: object = (
m.get("tool_calls") if isinstance(m, dict) else getattr(m, "tool_calls", None)
)
call_id = ""
if (
isinstance(tool_calls, Sequence)
and not isinstance(tool_calls, (str, bytes))
and len(tool_calls) > 0
):
first_call = tool_calls[0]
call_id_raw = (
first_call.get("id")
if isinstance(first_call, dict)
tool_calls_seq = _as_object_sequence(tool_calls)
if tool_calls_seq is not None and len(tool_calls_seq) > 0:
first_call = tool_calls_seq[0]
first_call_dict = _as_str_object_dict(first_call)
call_id_raw: object = (
first_call_dict.get("id")
if first_call_dict is not None
else getattr(first_call, "id", None)
)
if call_id_raw:
@ -518,23 +567,19 @@ class LiteLLMCompletionResponsesConfig:
call_id = ""
if role == "assistant":
tool_calls: Any = None
tool_calls: object = None
if isinstance(tool_call_message, dict):
tool_calls = tool_call_message.get("tool_calls")
else:
tool_calls = getattr(tool_call_message, "tool_calls", None)
if (
isinstance(tool_calls, Sequence)
and not isinstance(tool_calls, (str, bytes))
and len(tool_calls) > 0
):
first_call = tool_calls[0]
call_id_raw = None
if isinstance(first_call, dict):
call_id_raw = first_call.get("id")
else:
call_id_raw = getattr(first_call, "id", None)
tool_calls_seq = _as_object_sequence(tool_calls)
if tool_calls_seq is not None and len(tool_calls_seq) > 0:
first_call = tool_calls_seq[0]
first_call_dict = _as_str_object_dict(first_call)
call_id_raw: object = (
first_call_dict.get("id") if first_call_dict is not None else getattr(first_call, "id", None)
)
if call_id_raw:
call_id = str(call_id_raw)
@ -562,25 +607,40 @@ class LiteLLMCompletionResponsesConfig:
return False
@staticmethod
def _find_previous_assistant_idx(messages: list[Any], current_idx: int) -> int | None:
def _find_previous_assistant_idx(
messages: Sequence[
AllMessageValues
| GenericChatCompletionMessage
| ChatCompletionResponseMessage
| ChatCompletionMessageToolCall
| Message
],
current_idx: int,
) -> int | None:
"""Find the index of the previous assistant message."""
for j in range(current_idx - 1, -1, -1):
if messages[j].get("role") == "assistant":
candidate = messages[j]
candidate_dict = _as_str_object_dict(candidate)
role = candidate_dict.get("role") if candidate_dict is not None else getattr(candidate, "role", None)
if role == "assistant":
return j
return None
@staticmethod
def _recover_tool_call_id_from_assistant(assistant_message: Any, message: Any) -> str:
def _recover_tool_call_id_from_assistant(assistant_message: object, message: object) -> str:
"""Try to recover empty tool_call_id from assistant message's tool_calls."""
tool_calls_raw = (
assistant_message.get("tool_calls")
if isinstance(assistant_message, dict)
assistant_dict = _as_str_object_dict(assistant_message)
tool_calls_raw: object = (
assistant_dict.get("tool_calls")
if assistant_dict is not None
else getattr(assistant_message, "tool_calls", None)
)
if tool_calls_raw and isinstance(tool_calls_raw, list) and len(tool_calls_raw) > 0:
first_tool_call = tool_calls_raw[0]
if isinstance(first_tool_call, dict):
tool_call_id_raw = first_tool_call.get("id", "")
tool_calls_list = _as_object_list(tool_calls_raw)
if tool_calls_list is not None and len(tool_calls_list) > 0:
first_tool_call = tool_calls_list[0]
first_tool_call_dict = _as_str_object_dict(first_tool_call)
if first_tool_call_dict is not None:
tool_call_id_raw: object = first_tool_call_dict.get("id", "")
return str(tool_call_id_raw) if tool_call_id_raw is not None else ""
elif hasattr(first_tool_call, "id"):
tool_call_id_raw = getattr(first_tool_call, "id", None)
@ -588,28 +648,32 @@ class LiteLLMCompletionResponsesConfig:
return ""
@staticmethod
def _get_tool_calls_list(assistant_message: Any) -> list[Any]:
def _get_tool_calls_list(assistant_message: object) -> list[object]:
"""Extract tool_calls as a list from assistant message."""
tool_calls_raw = (
assistant_message.get("tool_calls")
if isinstance(assistant_message, dict)
assistant_dict = _as_str_object_dict(assistant_message)
tool_calls_raw: object = (
assistant_dict.get("tool_calls")
if assistant_dict is not None
else getattr(assistant_message, "tool_calls", None)
)
if tool_calls_raw is None:
return []
if isinstance(tool_calls_raw, list):
return tool_calls_raw
if hasattr(tool_calls_raw, "__iter__") and not isinstance(tool_calls_raw, (str, bytes)):
return list(tool_calls_raw)
tool_calls_list = _as_object_list(tool_calls_raw)
if tool_calls_list is not None:
return tool_calls_list
tool_calls_iterable = _as_object_iterable(tool_calls_raw)
if tool_calls_iterable is not None:
return list(tool_calls_iterable)
return []
@staticmethod
def _check_tool_call_exists(tool_calls: list[Any], tool_call_id: str) -> bool:
def _check_tool_call_exists(tool_calls: Sequence[object], tool_call_id: str) -> bool:
"""Check if a tool_call with the given ID exists in the list."""
for tool_call in tool_calls:
tool_call_id_to_check: str | None = None
if isinstance(tool_call, dict):
tool_call_id_to_check = tool_call.get("id")
tool_call_id_to_check: object = None
tool_call_dict = _as_str_object_dict(tool_call)
if tool_call_dict is not None:
tool_call_id_to_check = tool_call_dict.get("id")
elif hasattr(tool_call, "id"):
tool_call_id_to_check = getattr(tool_call, "id", None)
if tool_call_id_to_check == tool_call_id:
@ -617,52 +681,51 @@ class LiteLLMCompletionResponsesConfig:
return False
@staticmethod
def _reconstruct_tool_call_from_tools(tool_call_id: str, tools: list[Any]) -> dict[str, Any] | None:
def _reconstruct_tool_call_from_tools(tool_call_id: str, tools: Sequence[object]) -> dict[str, object] | None:
"""Reconstruct a minimal tool_call definition from tools list."""
for tool in tools:
if isinstance(tool, dict):
tool_function = tool.get("function") or {}
tool_name = tool_function.get("name") or tool.get("name") or ""
if tool_name:
return {
"id": tool_call_id,
"type": "function",
"function": {
"name": tool_name,
"arguments": "{}", # We don't know the arguments, use empty
},
}
tool_dict = _as_str_object_dict(tool)
if tool_dict is None:
continue
tool_function_dict: Mapping[str, object] = _as_str_object_dict(tool_dict.get("function")) or {}
tool_name = tool_function_dict.get("name") or tool_dict.get("name") or ""
if tool_name:
return {
"id": tool_call_id,
"type": "function",
"function": {
"name": tool_name,
"arguments": "{}", # We don't know the arguments, use empty
},
}
return None
@staticmethod
def _get_mapping_or_attr_value(obj: Any, key: str, default: Any = None) -> Any:
def _get_mapping_or_attr_value(obj: object, key: str, default: object = None) -> object:
"""
Safely read a field from dict-like or attribute-based objects.
"""
if obj is None:
return default
if isinstance(obj, dict):
return obj.get(key, default)
obj_dict = _as_str_object_dict(obj)
if obj_dict is not None:
return obj_dict.get(key, default)
getter = getattr(obj, "get", None)
if callable(getter):
try:
return getter(key, default)
except (TypeError, AttributeError):
pass
if isinstance(obj, BaseModel):
return _model_as_dict(obj).get(key, default)
return getattr(obj, key, default)
@staticmethod
def _create_tool_call_chunk(
tool_use_definition: dict[str, Any], tool_call_id: str, index: int
tool_use_definition: dict[str, object], tool_call_id: str, index: int
) -> ChatCompletionToolCallChunk:
"""Create a ChatCompletionToolCallChunk from tool_use_definition."""
function_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(tool_use_definition, "function")
function_name_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(function_raw, "name")
function_arguments_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(function_raw, "arguments")
function: dict[str, Any] = {
function: dict[str, object] = {
"name": function_name_raw or "",
"arguments": function_arguments_raw or "{}",
}
@ -681,15 +744,16 @@ class LiteLLMCompletionResponsesConfig:
)
@staticmethod
def _normalize_tool_use_definition(tool_use_definition: Any, tool_call_id: str) -> dict[str, Any] | None:
def _normalize_tool_use_definition(tool_use_definition: object, tool_call_id: str) -> dict[str, object] | None:
"""
Normalize cached tool_call definitions to a dict-like shape consumed by _create_tool_call_chunk.
"""
if not tool_use_definition:
return None
if isinstance(tool_use_definition, dict):
normalized_definition: dict[str, Any] = dict(tool_use_definition)
tool_use_definition_dict = _as_str_object_dict(tool_use_definition)
if tool_use_definition_dict is not None:
normalized_definition: dict[str, object] = dict(tool_use_definition_dict)
else:
tool_use_id_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(tool_use_definition, "id")
tool_use_type_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(tool_use_definition, "type")
@ -722,20 +786,22 @@ class LiteLLMCompletionResponsesConfig:
return normalized_definition
@staticmethod
def _add_tool_call_to_assistant(assistant_message: Any, tool_call_chunk: ChatCompletionToolCallChunk) -> None:
def _add_tool_call_to_assistant(assistant_message: object, tool_call_chunk: ChatCompletionToolCallChunk) -> None:
"""Add a tool_call to an assistant message."""
if isinstance(assistant_message, dict):
prev_assistant_dict = cast(dict[str, Any], assistant_message)
if "tool_calls" not in prev_assistant_dict:
prev_assistant_dict["tool_calls"] = []
tool_calls_list = prev_assistant_dict["tool_calls"]
if isinstance(tool_calls_list, list):
assistant_dict = _as_str_object_dict(assistant_message)
if assistant_dict is not None:
if "tool_calls" not in assistant_dict:
assistant_dict["tool_calls"] = []
tool_calls_list = _as_object_list(assistant_dict["tool_calls"])
if tool_calls_list is not None:
tool_calls_list.append(tool_call_chunk)
elif hasattr(assistant_message, "tool_calls"):
if assistant_message.tool_calls is None:
assistant_message.tool_calls = []
if isinstance(assistant_message.tool_calls, list):
assistant_message.tool_calls.append(tool_call_chunk)
if getattr(assistant_message, "tool_calls", None) is None:
setattr(assistant_message, "tool_calls", [])
existing_tool_calls: object = getattr(assistant_message, "tool_calls", None)
existing_tool_calls_list = _as_object_list(existing_tool_calls)
if existing_tool_calls_list is not None:
existing_tool_calls_list.append(tool_call_chunk)
@staticmethod
def _ensure_tool_results_have_corresponding_tool_calls(
@ -746,7 +812,7 @@ class LiteLLMCompletionResponsesConfig:
| ChatCompletionMessageToolCall
| Message
],
tools: list[Any] | None = None,
tools: Sequence[object] | None = None,
) -> list[
AllMessageValues
| GenericChatCompletionMessage
@ -808,9 +874,8 @@ class LiteLLMCompletionResponsesConfig:
)
if tool_call_id:
# Type-safe way to set tool_call_id on tool message
if isinstance(message, dict):
# Cast to dict to allow setting tool_call_id
message_dict = cast(dict[str, Any], message)
message_dict = _as_str_object_dict(message)
if message_dict is not None:
message_dict["tool_call_id"] = tool_call_id
elif hasattr(message, "tool_call_id"):
setattr(message, "tool_call_id", tool_call_id)
@ -862,7 +927,7 @@ class LiteLLMCompletionResponsesConfig:
@staticmethod
def _transform_responses_api_input_item_to_chat_completion_message(
input_item: Any,
input_item: Mapping[str, object],
) -> list[AllMessageValues | GenericChatCompletionMessage | ChatCompletionResponseMessage]:
"""
Transform a Responses API input item into a Chat Completion message
@ -897,9 +962,10 @@ class LiteLLMCompletionResponsesConfig:
# Since guardrails skip None content anyway, we return empty list to exclude it from structured messages
if content is None:
return []
role_value = input_item.get("role")
return [
GenericChatCompletionMessage(
role=input_item.get("role") or "user",
role=role_value if isinstance(role_value, str) and role_value else "user",
content=LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content(
content
),
@ -907,7 +973,7 @@ class LiteLLMCompletionResponsesConfig:
]
@staticmethod
def _is_input_item_tool_call_output(input_item: Any) -> bool:
def _is_input_item_tool_call_output(input_item: Mapping[str, object]) -> bool:
"""
Check if the input item is a tool call output
"""
@ -920,7 +986,7 @@ class LiteLLMCompletionResponsesConfig:
]
@staticmethod
def _is_input_item_function_call(input_item: Any) -> bool:
def _is_input_item_function_call(input_item: Mapping[str, object]) -> bool:
"""
Check if the input item is a function call or custom tool call.
Both need to be reconstructed as assistant tool_calls for Chat
@ -930,7 +996,7 @@ class LiteLLMCompletionResponsesConfig:
@staticmethod
def _transform_responses_api_tool_call_output_to_chat_completion_message(
tool_call_output: dict[str, Any],
tool_call_output: Mapping[str, object],
) -> list[AllMessageValues | GenericChatCompletionMessage | ChatCompletionResponseMessage]:
"""
ChatCompletionToolMessage is used to indicate the output from a tool call
@ -942,8 +1008,8 @@ class LiteLLMCompletionResponsesConfig:
return []
def _normalize_function_call_output_to_tool_content(
output: Any,
) -> Any:
output: object,
) -> str | list[dict[str, object]]:
"""
Normalize Responses API function_call_output.output into a shape that downstream
chat adapters (esp. Gemini) can reliably consume.
@ -964,22 +1030,25 @@ class LiteLLMCompletionResponsesConfig:
return output
# Some adapters represent tool output as a list of "input_*" parts
if isinstance(output, list):
normalized_blocks: list[dict[str, Any]] = []
output_list = _as_object_list(output)
if output_list is not None:
normalized_blocks: list[dict[str, object]] = []
text_acc: list[str] = []
for part in output:
if not isinstance(part, dict):
for part in output_list:
part_dict = _as_str_object_dict(part)
if part_dict is None:
continue
part_type = part.get("type")
part_type = part_dict.get("type")
if part_type in ("input_text", "output_text", "text"):
txt = part.get("text")
txt = part_dict.get("text")
if isinstance(txt, str) and txt:
text_acc.append(txt)
normalized_blocks.append({"type": "text", "text": txt})
elif part_type in ("input_image", "image_url"):
image_url_val = part.get("image_url") or part.get("url")
if isinstance(image_url_val, dict):
url = image_url_val.get("url")
image_url_val = part_dict.get("image_url") or part_dict.get("url")
image_url_dict = _as_str_object_dict(image_url_val)
if image_url_dict is not None:
url = image_url_dict.get("url")
if isinstance(url, str) and url:
normalized_blocks.append({"type": "image_url", "image_url": {"url": url}})
elif isinstance(image_url_val, str) and image_url_val:
@ -1012,9 +1081,11 @@ class LiteLLMCompletionResponsesConfig:
except Exception:
return str(output)
normalized_output = _normalize_function_call_output_to_tool_content(tool_call_output.get("output"))
tool_message_content = _as_tool_message_content(normalized_output)
tool_output_message = ChatCompletionToolMessage(
role="tool",
content=_normalize_function_call_output_to_tool_content(tool_call_output.get("output")),
content=tool_message_content if tool_message_content is not None else "",
tool_call_id=str(call_id),
)
@ -1066,7 +1137,7 @@ class LiteLLMCompletionResponsesConfig:
@staticmethod
def _transform_responses_api_function_call_to_chat_completion_message(
function_call: dict[str, Any],
function_call: Mapping[str, object],
) -> list[AllMessageValues | GenericChatCompletionMessage | ChatCompletionResponseMessage]:
"""
Transform a Responses API function_call into a Chat Completion message with tool calls
@ -1087,15 +1158,17 @@ class LiteLLMCompletionResponsesConfig:
# Create a tool call for the function call. Custom tool calls
# store their payload in "input" (raw string) rather than
# "arguments" (JSON string), so normalize to arguments here.
raw_arguments = function_call.get("arguments")
raw_arguments: object = function_call.get("arguments")
if not raw_arguments and function_call.get("type") == "custom_tool_call":
raw_input = function_call.get("input") or ""
raw_arguments = json.dumps({"content": raw_input}) if raw_input else ""
call_id_value = function_call.get("call_id") or function_call.get("id") or ""
name_value = function_call.get("name") or ""
tool_call = ChatCompletionToolCallChunk(
id=function_call.get("call_id") or function_call.get("id") or "",
id=str(call_id_value),
type="function",
function=ChatCompletionToolCallFunctionChunk(
name=function_call.get("name") or "",
name=str(name_value),
arguments=str(raw_arguments or ""),
),
index=0,
@ -1111,16 +1184,17 @@ class LiteLLMCompletionResponsesConfig:
return [chat_completion_response_message]
@staticmethod
def _resolve_file_id(item: dict[str, Any]) -> str | None:
def _resolve_file_id(item: Mapping[str, object]) -> str | None:
"""
Return the effective file_id for a Responses API input_file item.
Explicit file_id takes precedence; file_url is used as fallback so
downstream providers (Anthropic, Gemini) can handle the URL natively.
"""
return item.get("file_id") or item.get("file_url") or None
resolved = item.get("file_id") or item.get("file_url") or None
return resolved if isinstance(resolved, str) else None
@staticmethod
def _transform_input_file_item_to_file_item(item: dict[str, Any]) -> dict[str, Any]:
def _transform_input_file_item_to_file_item(item: Mapping[str, object]) -> dict[str, object]:
"""
Transform a Responses API input_file item to a Chat Completion file item
@ -1130,35 +1204,38 @@ class LiteLLMCompletionResponsesConfig:
Returns:
Dictionary with transformed file structure for Chat Completion
"""
file_dict: dict[str, Any] = {}
file_dict: dict[str, object] = {}
file_id = LiteLLMCompletionResponsesConfig._resolve_file_id(item)
if file_id:
file_dict["file_id"] = file_id
if item.get("file_data"):
file_dict["file_data"] = item["file_data"]
new_item: dict[str, Any] = {"type": "file", "file": file_dict}
new_item: dict[str, object] = {"type": "file", "file": file_dict}
if "cache_control" in item:
new_item["cache_control"] = item["cache_control"]
return new_item
@staticmethod
def _transform_input_image_item_to_image_item(
item: dict[str, Any],
item: Mapping[str, object],
) -> ChatCompletionImageObject:
"""
Transform a Responses API input_image item to a Chat Completion image item
"""
url_value = item.get("image_url")
detail_value = item.get("detail")
image_url_obj = ChatCompletionImageUrlObject(
url=item.get("image_url") or "", detail=item.get("detail") or "auto"
url=url_value if isinstance(url_value, str) and url_value else "",
detail=detail_value if isinstance(detail_value, str) and detail_value else "auto",
)
return ChatCompletionImageObject(type="image_url", image_url=image_url_obj)
@staticmethod
def _transform_responses_api_content_to_chat_completion_content(
content: Any,
) -> str | list[str | dict[str, Any]]:
content: object,
) -> str | list[str | dict[str, object]]:
"""
Transform a Responses API content into a Chat Completion content
@ -1171,37 +1248,41 @@ class LiteLLMCompletionResponsesConfig:
return ""
elif isinstance(content, str):
return content
elif isinstance(content, list):
content_list: list[str | dict[str, Any]] = []
for item in content:
elif (content_items := _as_object_list(content)) is not None:
content_list: list[str | dict[str, object]] = []
for item in content_items:
if isinstance(item, str):
content_list.append(item)
elif isinstance(item, dict):
if item.get("type") == "input_file":
content_list.append(
LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item(item)
)
elif item.get("type") == "input_image":
image_block = dict(
LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item(item)
)
if "cache_control" in item:
image_block["cache_control"] = item["cache_control"]
content_list.append(image_block)
else:
# Skip text blocks with None text to avoid downstream errors
text_value = item.get("text")
if text_value is None:
continue
content_block: dict[str, Any] = {
"type": LiteLLMCompletionResponsesConfig._get_chat_completion_request_content_type(
item.get("type") or "text"
),
"text": text_value,
}
if "cache_control" in item:
content_block["cache_control"] = item["cache_control"]
content_list.append(content_block)
continue
item_dict = _as_str_object_dict(item)
if item_dict is None:
continue
if item_dict.get("type") == "input_file":
content_list.append(
LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item(item_dict)
)
elif item_dict.get("type") == "input_image":
image_block: dict[str, object] = dict(
LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item(item_dict)
)
if "cache_control" in item_dict:
image_block["cache_control"] = item_dict["cache_control"]
content_list.append(image_block)
else:
# Skip text blocks with None text to avoid downstream errors
text_value = item_dict.get("text")
if text_value is None:
continue
block_type = item_dict.get("type")
content_block: dict[str, object] = {
"type": LiteLLMCompletionResponsesConfig._get_chat_completion_request_content_type(
block_type if isinstance(block_type, str) and block_type else "text"
),
"text": text_value,
}
if "cache_control" in item_dict:
content_block["cache_control"] = item_dict["cache_control"]
content_list.append(content_block)
return content_list
else:
raise ValueError(f"Invalid content type: {type(content)}")
@ -1283,7 +1364,7 @@ class LiteLLMCompletionResponsesConfig:
parameters = dict(typed_tool.get("parameters", {}) or {})
if not parameters or "type" not in parameters:
parameters["type"] = "object"
chat_completion_tool: dict[str, Any] = {
chat_completion_tool: dict[str, object] = {
"type": "function",
"function": {
"name": typed_tool.get("name") or "",
@ -1324,7 +1405,7 @@ class LiteLLMCompletionResponsesConfig:
@staticmethod
def transform_chat_completion_tool_params_to_responses_api_tools(
chat_completion_tools: list[ChatCompletionToolParam | OpenAIMcpServerTool] | None,
) -> list[dict[str, Any]]:
) -> list[dict[str, object]]:
"""
Transform Chat Completion tool params (e.g. from guardrail output) back to
Responses API request tool format. Inverse of
@ -1332,17 +1413,18 @@ class LiteLLMCompletionResponsesConfig:
"""
if chat_completion_tools is None or not chat_completion_tools:
return []
result: list[dict[str, Any]] = []
result: list[dict[str, object]] = []
for tool in chat_completion_tools:
if not isinstance(tool, dict):
result.append(tool) # type: ignore
continue
if tool.get("type") == "function":
fn = cast(dict[str, Any], tool.get("function") or {})
parameters = dict(fn.get("parameters", {}) or {})
fn: Mapping[str, object] = _as_str_object_dict(tool.get("function")) or {}
parameters_dict = _as_str_object_dict(fn.get("parameters"))
parameters: dict[str, object] = dict(parameters_dict) if parameters_dict is not None else {}
if not parameters or "type" not in parameters:
parameters["type"] = "object"
responses_tool: dict[str, Any] = {
responses_tool: dict[str, object] = {
"type": "function",
"name": fn.get("name") or "",
"description": fn.get("description") or "",
@ -1360,9 +1442,17 @@ class LiteLLMCompletionResponsesConfig:
result.append(responses_tool)
else:
# mcp or other: pass through unchanged
result.append(dict(tool))
passthrough_tool: dict[str, object] = dict(tool)
result.append(passthrough_tool)
return result
@staticmethod
def _read_provider_specific_fields(source: BaseModel) -> dict[str, object] | None:
raw_dict = _as_str_object_dict(_model_as_dict(source).get("provider_specific_fields"))
if raw_dict:
return raw_dict
return None
@staticmethod
def transform_chat_completion_tools_to_responses_tools(
chat_completion_response: ModelResponse,
@ -1414,25 +1504,9 @@ class LiteLLMCompletionResponsesConfig:
responses_tools.append(custom_item)
else:
# Build regular function_call output item
provider_specific_fields: dict | None = None
if hasattr(tool, "provider_specific_fields") and getattr(tool, "provider_specific_fields", None):
provider_specific_fields = getattr(tool, "provider_specific_fields")
if not isinstance(provider_specific_fields, dict):
provider_specific_fields = (
dict(provider_specific_fields) # type: ignore
if hasattr(provider_specific_fields, "__dict__")
else {}
)
elif hasattr(function_definition, "provider_specific_fields") and getattr(
function_definition, "provider_specific_fields", None
):
provider_specific_fields = getattr(function_definition, "provider_specific_fields")
if not isinstance(provider_specific_fields, dict):
provider_specific_fields = (
dict(provider_specific_fields) # type: ignore
if hasattr(provider_specific_fields, "__dict__")
else {}
)
provider_specific_fields = LiteLLMCompletionResponsesConfig._read_provider_specific_fields(
tool
) or LiteLLMCompletionResponsesConfig._read_provider_specific_fields(function_definition)
output_tool_call: ResponseFunctionToolCall = ResponseFunctionToolCall(
name=tool_name,
@ -1496,9 +1570,9 @@ class LiteLLMCompletionResponsesConfig:
@staticmethod
def convert_response_function_tool_call_to_chat_completion_tool_call(
tool_call_item: Any,
tool_call_item: object,
index: int = 0,
) -> dict[str, Any]:
) -> dict[str, object]:
"""
Convert ResponseFunctionToolCall to ChatCompletionToolCallChunk format.
@ -1509,37 +1583,36 @@ class LiteLLMCompletionResponsesConfig:
Returns:
Dictionary in ChatCompletionToolCallChunk format
"""
# Extract provider_specific_fields if present
provider_specific_fields = getattr(tool_call_item, "provider_specific_fields", None)
if provider_specific_fields and not isinstance(provider_specific_fields, dict):
provider_specific_fields = (
dict(provider_specific_fields) if hasattr(provider_specific_fields, "__dict__") else {}
)
elif hasattr(tool_call_item, "get") and callable(tool_call_item.get): # type: ignore
provider_fields = tool_call_item.get("provider_specific_fields") # type: ignore
if provider_fields:
provider_specific_fields = (
provider_fields
if isinstance(provider_fields, dict)
else (
dict(provider_fields) # type: ignore
if hasattr(provider_fields, "__dict__")
else {}
)
)
if isinstance(tool_call_item, BaseModel):
dumped = _model_as_dict(tool_call_item)
name_value: object = dumped.get("name")
arguments_value: object = dumped.get("arguments")
item_id_value: object = dumped.get("id")
call_id_value: object = dumped.get("call_id")
psf_value: object = dumped.get("provider_specific_fields")
else:
name_value = getattr(tool_call_item, "name", None)
arguments_value = getattr(tool_call_item, "arguments", None)
item_id_value = getattr(tool_call_item, "id", None)
call_id_value = getattr(tool_call_item, "call_id", None)
psf_value = getattr(tool_call_item, "provider_specific_fields", None)
function_dict: dict[str, Any] = {
"name": tool_call_item.name,
"arguments": tool_call_item.arguments,
# Extract provider_specific_fields if present
psf_dict = _as_str_object_dict(psf_value)
provider_specific_fields = psf_dict if psf_dict else None
function_dict: dict[str, object] = {
"name": name_value,
"arguments": arguments_value,
}
if provider_specific_fields:
function_dict["provider_specific_fields"] = provider_specific_fields
tool_call_dict: dict[str, Any] = {
tool_call_dict: dict[str, object] = {
"id": LiteLLMCompletionResponsesConfig._tool_call_id_from_responses_item(
getattr(tool_call_item, "id", None),
getattr(tool_call_item, "call_id", None),
item_id_value if isinstance(item_id_value, str) else None,
call_id_value if isinstance(call_id_value, str) else None,
),
"function": function_dict,
"type": "function",
@ -1553,9 +1626,9 @@ class LiteLLMCompletionResponsesConfig:
@staticmethod
def convert_apply_patch_tool_call_to_chat_completion_tool_call(
tool_call_item: Any,
tool_call_item: SupportsApplyPatchOperation,
index: int = 0,
) -> dict[str, Any]:
) -> dict[str, object]:
"""
Convert ResponseApplyPatchToolCall to ChatCompletionToolCallChunk format.
@ -1570,10 +1643,8 @@ class LiteLLMCompletionResponsesConfig:
Returns:
Dictionary in ChatCompletionToolCallChunk format
"""
import json
operation_dict = tool_call_item.operation.model_dump()
tool_call_dict: dict[str, Any] = {
operation_dict = _model_as_dict(tool_call_item.operation)
tool_call_dict: dict[str, object] = {
"id": tool_call_item.call_id,
"function": {
"name": "apply_patch",
@ -1711,11 +1782,11 @@ class LiteLLMCompletionResponsesConfig:
"""
output_items: list = []
for choice in chat_completion_response.choices or []:
message = getattr(choice, "message", None)
message: Message | None = getattr(choice, "message", None)
if not message:
continue
psf = getattr(message, "provider_specific_fields", None)
if not psf or not isinstance(psf, dict):
psf = message.provider_specific_fields
if not psf:
continue
results = psf.get("code_interpreter_results")
if results and isinstance(results, list):
@ -1783,13 +1854,16 @@ class LiteLLMCompletionResponsesConfig:
"""
image_generation_items: list[OutputImageGenerationCall] = []
images = getattr(choice.message, "images", [])
images = choice.message.images or []
if not images:
return image_generation_items
for idx, image_item in enumerate(images):
# Extract base64 from data URL
image_url = image_item.get("image_url", {}).get("url", "")
image_url_value: object = image_item.get("image_url")
image_url_dict = _as_str_object_dict(image_url_value)
url_value = image_url_dict.get("url", "") if image_url_dict is not None else ""
image_url = url_value if isinstance(url_value, str) else ""
base64_data = LiteLLMCompletionResponsesConfig._extract_base64_from_data_url(image_url)
if base64_data:
@ -2034,8 +2108,8 @@ class LiteLLMCompletionResponsesConfig:
@staticmethod
def _transform_text_format_to_response_format(
text_param: dict[str, Any] | Any,
) -> dict[str, Any] | None:
text_param: object,
) -> dict[str, object] | None:
"""
Transform Responses API text.format parameter to Chat Completion response_format parameter.
@ -2062,18 +2136,19 @@ class LiteLLMCompletionResponsesConfig:
if not text_param:
return None
if isinstance(text_param, dict):
format_param = text_param.get("format")
if format_param and isinstance(format_param, dict):
format_type = format_param.get("type")
text_param_dict = _as_str_object_dict(text_param)
if text_param_dict is not None:
format_param_dict = _as_str_object_dict(text_param_dict.get("format"))
if format_param_dict:
format_type = format_param_dict.get("type")
if format_type == "json_schema":
return {
"type": "json_schema",
"json_schema": {
"name": format_param.get("name", "response_schema"),
"schema": format_param.get("schema", {}),
"strict": format_param.get("strict", False),
"name": format_param_dict.get("name", "response_schema"),
"schema": format_param_dict.get("schema", {}),
"strict": format_param_dict.get("strict", False),
},
}
elif format_type == "json_object":

6
litellm/responses/mcp/mcp_streaming_iterator.py
View file

@ -64,7 +64,7 @@ async def create_mcp_list_tools_events(
# Convert tools to dict format for the event
mcp_tools_dict = []
for tool in filtered_mcp_tools:
if hasattr(tool, "model_dump") and callable(getattr(tool, "model_dump")):
if hasattr(tool, "model_dump") and callable(tool.model_dump):
# Type cast to help mypy understand this is safe after hasattr check
mcp_tools_dict.append(cast(Any, tool).model_dump())
elif hasattr(tool, "__dict__"):
@ -103,9 +103,9 @@ async def create_mcp_list_tools_events(
# Add input_schema if available
if hasattr(tool, "inputSchema"):
tool_dict["input_schema"] = getattr(tool, "inputSchema")
tool_dict["input_schema"] = tool.inputSchema
elif hasattr(tool, "input_schema"):
tool_dict["input_schema"] = getattr(tool, "input_schema")
tool_dict["input_schema"] = tool.input_schema
formatted_tools.append(tool_dict)

834
litellm/responses/streaming_iterator.py
View file

File diff suppressed because it is too large Load diff

28
ruff-strict-budget.json
View file

@ -1,30 +1,30 @@
{
"ANN001": {
"limit": 3104
"limit": 3080
},
"ANN002": {
"limit": 69
},
"ANN003": {
"limit": 831
"limit": 829
},
"ANN201": {
"limit": 2137
"limit": 2119
},
"ANN202": {
"limit": 941
"limit": 939
},
"ANN204": {
"limit": 724
"limit": 723
},
"ANN205": {
"limit": 127
"limit": 126
},
"ANN206": {
"limit": 130
},
"ANN401": {
"limit": 1851
"limit": 1726
},
"ASYNC230": {
"limit": 14
@ -42,7 +42,7 @@
"limit": 84
},
"B010": {
"limit": 194
"limit": 192
},
"B018": {
"limit": 5
@ -72,7 +72,7 @@
"limit": 22
},
"C408": {
"limit": 14
"limit": 13
},
"C414": {
"limit": 7
@ -81,7 +81,7 @@
"limit": 4
},
"C901": {
"limit": 312
"limit": 311
},
"D419": {
"limit": 9
@ -255,16 +255,16 @@
"limit": 480
},
"S110": {
"limit": 236
"limit": 235
},
"S112": {
"limit": 24
},
"SIM101": {
"limit": 61
"limit": 59
},
"SIM102": {
"limit": 324
"limit": 323
},
"SIM103": {
"limit": 129
@ -321,7 +321,7 @@
"limit": 121
},
"TRY300": {
"limit": 879
"limit": 878
},
"UP006": {
"limit": 0

3
tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py
View file

@ -2649,7 +2649,7 @@ class _TxProxyModelTable:
class _TxPrismaClient:
"""Minimal prisma stub whose ``db.tx()`` yields a transaction and records commit."""
"""Minimal prisma stub whose ``tx()`` yields a transaction and records commit."""
def __init__(self, rows):
self.events: list = []
@ -2668,6 +2668,7 @@ class _TxPrismaClient:
self.db = MagicMock()
self.db.tx = MagicMock(return_value=_TxCM())
self.tx = self.db.tx
class _RecordingRouter:

8
type-discipline-budget.json
View file

@ -1,9 +1,9 @@
{
"LIT001": {
"limit": 23350
"limit": 23263
},
"LIT002": {
"limit": 27256
"limit": 27216
},
"LIT003": {
"limit": 292
@ -15,7 +15,7 @@
"limit": 0
},
"LIT006": {
"limit": 1105
"limit": 1085
},
"LIT007": {
"limit": 0
@ -24,6 +24,6 @@
"limit": 1004
},
"LIT009": {
"limit": 2465
"limit": 2426
}
}