scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-10 03:28:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 85

chore(typing): replace Any seams with real types across responses, proxy, and provider adapters

Browse source 
Replace Any-typed payload dicts, record shapes, and provider request/response
seams with TypedDicts, Protocols, and precise annotations in the ten litellm/
files carrying the highest combined basedpyright reportAny + reportExplicitAny
counts. No behavior changes.

Adds a regression test covering the managed-id list path so a prisma client
missing the managed tables keeps returning a fail-closed empty page.
This commit is contained in:
mateo-berri 2026-08-04 20:05:35 +00:00
parent 956d5177d1
commit 42564e896f
No known key found for this signature in database
18 changed files with 1174 additions and 641 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

28
basedpyright-code-budget.json
View file

@ -1,9 +1,9 @@
{
"reportAny": {
"limit": 29682
"limit": 29082
},
"reportArgumentType": {
"limit": 2645
"limit": 2635
},
"reportAssignmentType": {
"limit": 329
@ -24,7 +24,7 @@
"limit": 42
},
"reportExplicitAny": {
"limit": 9440
"limit": 9198
},
"reportFunctionMemberAccess": {
"limit": 11
@ -54,10 +54,10 @@
"limit": 0
},
"reportMissingParameterType": {
"limit": 5848
"limit": 5843
},
"reportMissingTypeArgument": {
"limit": 15850
"limit": 15834
},
"reportMissingTypeStubs": {
"limit": 41
@ -72,7 +72,7 @@
"limit": 0
},
"reportOptionalMemberAccess": {
"limit": 1079
"limit": 1078
},
"reportOptionalOperand": {
"limit": 0
@ -90,7 +90,7 @@
"limit": 12
},
"reportReturnType": {
"limit": 219
"limit": 218
},
"reportTypedDictNotRequiredAccess": {
"limit": 27
@ -99,22 +99,22 @@
"limit": 0
},
"reportUnknownArgumentType": {
"limit": 45297
"limit": 45277
},
"reportUnknownLambdaType": {
"limit": 113
},
"reportUnknownMemberType": {
"limit": 40411
"limit": 40303
},
"reportUnknownParameterType": {
"limit": 20301
"limit": 20285
},
"reportUnknownVariableType": {
"limit": 31968
"limit": 31883
},
"reportUnnecessaryCast": {
"limit": 177
"limit": 175
},
"reportUnnecessaryComparison": {
"limit": 1021
@ -123,7 +123,7 @@
"limit": 7
},
"reportUnnecessaryIsInstance": {
"limit": 1204
"limit": 1203
},
"reportUntypedBaseClass": {
"limit": 165
@ -138,7 +138,7 @@
"limit": 204
},
"reportUnusedImport": {
"limit": 1003
"limit": 1002
},
"reportUnusedVariable": {
"limit": 1297

184
litellm/a2a_protocol/litellm_completion_bridge/handler.py
View file

@ -10,7 +10,7 @@ A2A Streaming Events (in order):
4. Status update (kind: "status-update") - Final status "completed" with final=true
"""
from collections.abc import AsyncIterator
from collections.abc import AsyncIterator, Mapping
from typing import Any
import litellm
@ -21,6 +21,8 @@ from litellm.a2a_protocol.litellm_completion_bridge.transformation import (
)
from litellm.a2a_protocol.providers.config_manager import A2AProviderConfigManager
from litellm.interactions.agents.utils import merge_agent_headers
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.types.utils import ModelResponse
# litellm_params key carrying the authenticated principal (hashed virtual key) so
# A2A provider configs can scope provider-side state (e.g. LangFlow session memory)
@ -44,6 +46,72 @@ class A2ACompletionBridgeHandler:
Static methods for handling A2A requests via LiteLLM completion.
"""
@staticmethod
def _build_completion_params(
params: dict[str, Any],
litellm_params: Mapping[str, Any],
api_base: str | None,
agent_extra_headers: Mapping[str, str] | None,
*,
stream: bool,
) -> Mapping[str, Any]:
# Extract message from params
message = params.get("message", {})
# Transform A2A message to OpenAI format
openai_messages = A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message)
# Get completion params
custom_llm_provider = litellm_params.get("custom_llm_provider")
model = litellm_params.get("model", "agent")
# Build full model string if provider specified
# Skip prepending if model already starts with the provider prefix
if custom_llm_provider and not model.startswith(f"{custom_llm_provider}/"):
full_model = f"{custom_llm_provider}/{model}"
else:
full_model = model
if stream:
verbose_logger.info("A2A completion bridge streaming: model=%s, api_base=%s", full_model, api_base)
else:
verbose_logger.info("A2A completion bridge: model=%s, api_base=%s", full_model, api_base)
# Build completion params dict
completion_params: dict[str, Any] = {
"model": full_model,
"messages": openai_messages,
"api_base": api_base,
"stream": stream,
}
# Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.)
litellm_params_to_add = {
k: v
for k, v in litellm_params.items()
if k not in ("model", "custom_llm_provider") and k not in _AGENT_ONLY_PARAMS
}
completion_params.update(litellm_params_to_add)
# Apply forward metadata AFTER the litellm_params merge so the helper
# sees any agent-owner-configured ``extra_body.metadata`` and can keep
# those keys authoritative over the client-supplied A2A metadata.
A2ACompletionBridgeTransformation.apply_forward_metadata_to_completion_params(
completion_params=completion_params,
a2a_message=message,
params=params,
)
if agent_extra_headers:
completion_params["extra_headers"] = merge_agent_headers(
dynamic_headers=agent_extra_headers,
static_headers=completion_params.get("extra_headers"),
)
return completion_params
@staticmethod
async def _acompletion(completion_params: Mapping[str, Any]) -> ModelResponse | CustomStreamWrapper:
return await litellm.acompletion(**completion_params)
@staticmethod
async def handle_non_streaming(
request_id: str,
@ -53,7 +121,7 @@ class A2ACompletionBridgeHandler:
agent_extra_headers: dict[str, str] | None = None,
*,
_skip_a2a_provider_routing: bool = False,
) -> dict[str, Any]:
) -> dict[str, object]:
"""
Handle non-streaming A2A request via litellm.acompletion.
@ -86,56 +154,16 @@ class A2ACompletionBridgeHandler:
agent_extra_headers=agent_extra_headers,
)
# Extract message from params
message = params.get("message", {})
# Transform A2A message to OpenAI format
openai_messages = A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message)
# Get completion params
custom_llm_provider = litellm_params.get("custom_llm_provider")
model = litellm_params.get("model", "agent")
# Build full model string if provider specified
# Skip prepending if model already starts with the provider prefix
if custom_llm_provider and not model.startswith(f"{custom_llm_provider}/"):
full_model = f"{custom_llm_provider}/{model}"
else:
full_model = model
verbose_logger.info("A2A completion bridge: model=%s, api_base=%s", full_model, api_base)
# Build completion params dict
completion_params: dict[str, Any] = {
"model": full_model,
"messages": openai_messages,
"api_base": api_base,
"stream": False,
}
# Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.)
litellm_params_to_add = {
k: v
for k, v in litellm_params.items()
if k not in ("model", "custom_llm_provider") and k not in _AGENT_ONLY_PARAMS
}
completion_params.update(litellm_params_to_add)
# Apply forward metadata AFTER the litellm_params merge so the helper
# sees any agent-owner-configured ``extra_body.metadata`` and can keep
# those keys authoritative over the client-supplied A2A metadata.
A2ACompletionBridgeTransformation.apply_forward_metadata_to_completion_params(
completion_params=completion_params,
a2a_message=message,
completion_params = A2ACompletionBridgeHandler._build_completion_params(
params=params,
litellm_params=litellm_params,
api_base=api_base,
agent_extra_headers=agent_extra_headers,
stream=False,
)
if agent_extra_headers:
completion_params["extra_headers"] = merge_agent_headers(
dynamic_headers=agent_extra_headers,
static_headers=completion_params.get("extra_headers"),
)
# Call litellm.acompletion
response = await litellm.acompletion(**completion_params)
response = await A2ACompletionBridgeHandler._acompletion(completion_params)
# Transform response to A2A format
a2a_response = A2ACompletionBridgeTransformation.openai_response_to_a2a_response(
@ -156,7 +184,7 @@ class A2ACompletionBridgeHandler:
agent_extra_headers: dict[str, str] | None = None,
*,
_skip_a2a_provider_routing: bool = False,
) -> AsyncIterator[dict[str, Any]]:
) -> AsyncIterator[dict[str, object]]:
"""
Handle streaming A2A request via litellm.acompletion with stream=True.
@ -198,60 +226,20 @@ class A2ACompletionBridgeHandler:
return
# Extract message from params
message = params.get("message", {})
# Create streaming context
ctx = A2AStreamingContext(
request_id=request_id,
input_message=message,
input_message=params.get("message", {}),
)
# Transform A2A message to OpenAI format
openai_messages = A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message)
# Get completion params
custom_llm_provider = litellm_params.get("custom_llm_provider")
model = litellm_params.get("model", "agent")
# Build full model string if provider specified
# Skip prepending if model already starts with the provider prefix
if custom_llm_provider and not model.startswith(f"{custom_llm_provider}/"):
full_model = f"{custom_llm_provider}/{model}"
else:
full_model = model
verbose_logger.info("A2A completion bridge streaming: model=%s, api_base=%s", full_model, api_base)
# Build completion params dict
completion_params: dict[str, Any] = {
"model": full_model,
"messages": openai_messages,
"api_base": api_base,
"stream": True,
}
# Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.)
litellm_params_to_add = {
k: v
for k, v in litellm_params.items()
if k not in ("model", "custom_llm_provider") and k not in _AGENT_ONLY_PARAMS
}
completion_params.update(litellm_params_to_add)
# Apply forward metadata AFTER the litellm_params merge so the helper
# sees any agent-owner-configured ``extra_body.metadata`` and can keep
# those keys authoritative over the client-supplied A2A metadata.
A2ACompletionBridgeTransformation.apply_forward_metadata_to_completion_params(
completion_params=completion_params,
a2a_message=message,
completion_params = A2ACompletionBridgeHandler._build_completion_params(
params=params,
litellm_params=litellm_params,
api_base=api_base,
agent_extra_headers=agent_extra_headers,
stream=True,
)
if agent_extra_headers:
completion_params["extra_headers"] = merge_agent_headers(
dynamic_headers=agent_extra_headers,
static_headers=completion_params.get("extra_headers"),
)
# 1. Emit initial task event (kind: "task", status: "submitted")
task_event = A2ACompletionBridgeTransformation.create_task_event(ctx)
yield task_event
@ -266,7 +254,7 @@ class A2ACompletionBridgeHandler:
yield working_event
# Call litellm.acompletion with streaming
response = await litellm.acompletion(**completion_params)
response = await A2ACompletionBridgeHandler._acompletion(completion_params)
# 3. Accumulate content and emit artifact update
accumulated_text = ""
@ -312,7 +300,7 @@ async def handle_a2a_completion(
litellm_params: dict[str, Any],
api_base: str | None = None,
agent_extra_headers: dict[str, str] | None = None,
) -> dict[str, Any]:
) -> dict[str, object]:
"""Convenience function for non-streaming A2A completion."""
return await A2ACompletionBridgeHandler.handle_non_streaming(
request_id=request_id,
@ -329,7 +317,7 @@ async def handle_a2a_completion_streaming(
litellm_params: dict[str, Any],
api_base: str | None = None,
agent_extra_headers: dict[str, str] | None = None,
) -> AsyncIterator[dict[str, Any]]:
) -> AsyncIterator[dict[str, object]]:
"""Convenience function for streaming A2A completion."""
async for chunk in A2ACompletionBridgeHandler.handle_streaming(
request_id=request_id,

33
litellm/google_genai/adapters/handler.py
View file

@ -1,7 +1,8 @@
from collections.abc import AsyncIterator, Coroutine
from typing import Any, cast
from collections.abc import AsyncIterator, Coroutine, Mapping
from typing import cast
import litellm
from litellm.types.google_genai.adapters import GenerateContentCompletionKwargs
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import ModelResponse
@ -17,12 +18,12 @@ class GenerateContentToCompletionHandler:
@staticmethod
def _prepare_completion_kwargs(
model: str,
contents: list[dict[str, Any]] | dict[str, Any],
config: dict[str, Any] | None = None,
contents: list[dict[str, object]] | dict[str, object],
config: dict[str, object] | None = None,
stream: bool = False,
litellm_params: GenericLiteLLMParams | None = None,
extra_kwargs: dict[str, Any] | None = None,
) -> dict[str, Any]:
extra_kwargs: Mapping[str, object] | None = None,
) -> GenerateContentCompletionKwargs:
"""Prepare kwargs for litellm.completion/acompletion"""
# Transform generate_content request to completion format
@ -34,7 +35,7 @@ class GenerateContentToCompletionHandler:
**(extra_kwargs or {}),
)
completion_kwargs: dict[str, Any] = dict(completion_request)
completion_kwargs = dict(completion_request)
# Forward extra_kwargs that should be passed to completion call
if extra_kwargs is not None:
@ -48,17 +49,17 @@ class GenerateContentToCompletionHandler:
if stream:
completion_kwargs["stream"] = stream
return completion_kwargs
return GenerateContentCompletionKwargs(**completion_kwargs)
@staticmethod
async def async_generate_content_handler(
model: str,
contents: list[dict[str, Any]] | dict[str, Any],
contents: list[dict[str, object]] | dict[str, object],
litellm_params: GenericLiteLLMParams,
config: dict[str, Any] | None = None,
config: dict[str, object] | None = None,
stream: bool = False,
**kwargs,
) -> dict[str, Any] | AsyncIterator[bytes]:
**kwargs: object,
) -> dict[str, object] | AsyncIterator[bytes]:
"""Handle generate_content call asynchronously using completion adapter"""
completion_kwargs = GenerateContentToCompletionHandler._prepare_completion_kwargs(
@ -103,13 +104,13 @@ class GenerateContentToCompletionHandler:
@staticmethod
def generate_content_handler(
model: str,
contents: list[dict[str, Any]] | dict[str, Any],
contents: list[dict[str, object]] | dict[str, object],
litellm_params: GenericLiteLLMParams,
config: dict[str, Any] | None = None,
config: dict[str, object] | None = None,
stream: bool = False,
_is_async: bool = False,
**kwargs,
) -> dict[str, Any] | AsyncIterator[bytes] | Coroutine[Any, Any, dict[str, Any] | AsyncIterator[bytes]]:
**kwargs: object,
) -> dict[str, object] | AsyncIterator[bytes] | Coroutine[None, None, dict[str, object] | AsyncIterator[bytes]]:
"""Handle generate_content call using completion adapter"""
if _is_async:

306
litellm/integrations/code_interpreter_interception/handler.py
View file

@ -9,13 +9,18 @@ captured stdout back through the typed agentic loop plan.
import json
import time
import uuid
from typing import Any, Literal, TypedDict, cast
from collections.abc import Sequence
from typing import TYPE_CHECKING, Any, Literal, Protocol, TypedDict, runtime_checkable
from pydantic import ValidationError
import litellm
from litellm._logging import verbose_logger
from litellm.integrations.custom_logger import CustomLogger
from litellm.llms.base_llm.sandbox.transformation import (
CodeExecutionResult,
ContainerHandle,
)
from litellm.types.integrations.code_interpreter_interception import (
CodeInterpreterInterceptionConfig,
)
@ -37,6 +42,9 @@ from litellm.types.utils import (
ModelResponse,
)
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
LITELLM_CODE_EXECUTION_TOOL_NAME = "litellm_code_execution"
_INTERCEPTION_ACTIVE_KEY = "_code_interpreter_interception_active"
_SANDBOX_KEY = "_code_interpreter_interception_sandbox_key"
@ -109,26 +117,87 @@ class ChatCompletionFunctionToolChoice(TypedDict):
CodeExecutionFunctionToolChoice = ResponsesFunctionToolChoice | ChatCompletionFunctionToolChoice
def _extract_session_id(kwargs: dict[str, Any]) -> str | None:
class SandboxToolParams(TypedDict):
sandbox_provider: str
api_key: str | None
api_base: str | None
class SandboxConfigProtocol(Protocol):
async def acreate_sandbox(self) -> ContainerHandle: ...
async def arun_code(self, *, container: ContainerHandle, code: str) -> CodeExecutionResult: ...
async def adelete_sandbox(self, *, container: ContainerHandle) -> object: ...
@runtime_checkable
class _SupportsOutput(Protocol):
output: object
_CachedContainer = tuple[ContainerHandle, SandboxToolParams | None, float, str | None]
def _output_item_type(item: object) -> object:
if isinstance(item, dict):
item_mapping: dict[str, object] = item
return item_mapping.get("type")
return getattr(item, "type", None)
def _tool_call_arguments(arguments: object) -> str:
if isinstance(arguments, str):
return arguments
return "" if arguments is None else str(arguments)
def _narrow_tool_call(tool_call: dict[str, object]) -> CodeExecutionToolCall:
tool_call_id = tool_call.get("id")
call_id = tool_call.get("call_id")
return {
"id": tool_call_id if isinstance(tool_call_id, str) else None,
"call_id": call_id if isinstance(call_id, str) else None,
"type": "function",
"name": LITELLM_CODE_EXECUTION_TOOL_NAME,
"arguments": _tool_call_arguments(tool_call.get("arguments")),
}
def _extract_session_id(kwargs: dict[str, object]) -> str | None:
for meta_key in ("metadata", "litellm_metadata"):
meta = kwargs.get(meta_key)
if isinstance(meta, dict):
sid = meta.get("session_id")
metadata: dict[str, object] = meta
sid = metadata.get("session_id")
if sid and isinstance(sid, str):
return sid
return None
def _extract_identity(kwargs: dict[str, Any]) -> str:
return kwargs.get("user_api_key_hash") or ""
def _extract_identity(kwargs: dict[str, object]) -> str:
identity = kwargs.get("user_api_key_hash")
return identity if isinstance(identity, str) else ""
def _resolve_sandbox_tool(sandbox_tool_name: str | None) -> dict[str, Any] | None:
def _resolve_sandbox_tool(sandbox_tool_name: str | None) -> SandboxToolParams | None:
if sandbox_tool_name is None:
return None
try:
from litellm.sandbox.sandbox_tools import resolve_sandbox_tool
except ImportError:
return None
return resolve_sandbox_tool(sandbox_tool_name)
resolved: dict[str, object] | None = resolve_sandbox_tool(sandbox_tool_name)
if resolved is None:
return None
provider = resolved.get("sandbox_provider")
api_key = resolved.get("api_key")
api_base = resolved.get("api_base")
return SandboxToolParams(
sandbox_provider=provider if isinstance(provider, str) else "",
api_key=api_key if isinstance(api_key, str) else None,
api_base=api_base if isinstance(api_base, str) else None,
)
class CodeInterpreterInterceptionLogger(CustomLogger):
@ -149,14 +218,14 @@ class CodeInterpreterInterceptionLogger(CustomLogger):
enabled: bool = True,
enabled_providers: list[str] | None = None,
sandbox_tool_name: str | None = None,
sandbox_config: Any | None = None,
sandbox_config: SandboxConfigProtocol | None = None,
):
super().__init__()
self.enabled = enabled
self.enabled_providers = enabled_providers
self.sandbox_tool_name = sandbox_tool_name
self.sandbox_config = sandbox_config
self._container_cache: dict[str, tuple[Any, dict[str, Any] | None, float, str | None]] = {}
self._container_cache: dict[str, _CachedContainer] = {}
@classmethod
def from_config_yaml(cls, config: CodeInterpreterInterceptionConfig) -> "CodeInterpreterInterceptionLogger":
@ -174,16 +243,13 @@ class CodeInterpreterInterceptionLogger(CustomLogger):
params: CodeInterpreterInterceptionConfig = {}
if "code_interpreter_interception_params" in litellm_settings:
params = litellm_settings["code_interpreter_interception_params"]
elif "code_interpreter_interception" in callback_specific_params and isinstance(
callback_specific_params["code_interpreter_interception"], dict
):
params = cast(
CodeInterpreterInterceptionConfig,
callback_specific_params["code_interpreter_interception"],
)
elif isinstance(callback_specific_params.get("code_interpreter_interception"), dict):
params = callback_specific_params["code_interpreter_interception"]
return CodeInterpreterInterceptionLogger.from_config_yaml(params)
async def async_pre_call_deployment_hook(self, kwargs: dict[str, Any], call_type: CallTypes | None) -> dict | None:
async def async_pre_call_deployment_hook(
self, kwargs: dict[str, object], call_type: CallTypes | None
) -> dict | None:
if not kwargs.get("_agentic_loop_depth"):
kwargs.pop(_INTERCEPTION_ACTIVE_KEY, None)
kwargs.pop(_SANDBOX_KEY, None)
@ -229,13 +295,14 @@ class CodeInterpreterInterceptionLogger(CustomLogger):
return kwargs
@staticmethod
def _strip_interception_metadata(kwargs: dict[str, Any]) -> None:
def _strip_interception_metadata(kwargs: dict[str, object]) -> None:
metadata = kwargs.get(_LITELLM_METADATA_KEY)
if not isinstance(metadata, dict):
return
current_metadata: dict[str, object] = metadata
filtered_metadata = {
key: value
for key, value in metadata.items()
for key, value in current_metadata.items()
if not is_interception_internal_key(key)
and not key.startswith("_agentic_loop")
and key != "max_agentic_loops"
@ -247,9 +314,9 @@ class CodeInterpreterInterceptionLogger(CustomLogger):
kwargs.pop(_LITELLM_METADATA_KEY, None)
@staticmethod
def _write_interception_metadata(kwargs: dict[str, Any]) -> None:
metadata = kwargs.get(_LITELLM_METADATA_KEY)
metadata = dict(metadata) if isinstance(metadata, dict) else {}
def _write_interception_metadata(kwargs: dict[str, object]) -> None:
existing = kwargs.get(_LITELLM_METADATA_KEY)
metadata: dict[str, object] = dict(existing) if isinstance(existing, dict) else {}
for key in (_INTERCEPTION_ACTIVE_KEY, _SANDBOX_KEY, _SESSION_SCOPED_KEY, _CONVERTED_STREAM_KEY):
if key in kwargs:
metadata[key] = kwargs[key]
@ -296,20 +363,21 @@ class CodeInterpreterInterceptionLogger(CustomLogger):
}
@staticmethod
def _tool_choice_targets_code_interpreter(tool_choice: Any) -> bool:
def _tool_choice_targets_code_interpreter(tool_choice: object) -> bool:
if not isinstance(tool_choice, dict):
return False
function = tool_choice.get("function")
choice: dict[str, object] = tool_choice
function = choice.get("function")
return (
tool_choice.get("type") == "code_interpreter"
or tool_choice.get("name") == "code_interpreter"
or tool_choice.get("name") == LITELLM_CODE_EXECUTION_TOOL_NAME
choice.get("type") == "code_interpreter"
or choice.get("name") == "code_interpreter"
or choice.get("name") == LITELLM_CODE_EXECUTION_TOOL_NAME
or (isinstance(function, dict) and function.get("name") == LITELLM_CODE_EXECUTION_TOOL_NAME)
)
def _resolve_provider(self, kwargs: dict[str, Any]) -> str | None:
def _resolve_provider(self, kwargs: dict[str, object]) -> str | None:
provider = kwargs.get("custom_llm_provider")
if provider:
if isinstance(provider, str) and provider:
return provider
model = kwargs.get("model")
if not isinstance(model, str):
@ -321,7 +389,7 @@ class CodeInterpreterInterceptionLogger(CustomLogger):
async def async_should_run_agentic_loop(
self,
response: Any,
response: object,
model: str,
messages: list[dict],
tools: list[dict] | None,
@ -351,12 +419,12 @@ class CodeInterpreterInterceptionLogger(CustomLogger):
tools: dict,
model: str,
messages: list[dict],
response: Any,
anthropic_messages_provider_config: Any,
anthropic_messages_optional_request_params: dict,
logging_obj: Any,
response: object,
anthropic_messages_provider_config: object,
anthropic_messages_optional_request_params: dict[str, object],
logging_obj: "LiteLLMLoggingObj",
stream: bool,
kwargs: dict,
kwargs: dict[str, object],
) -> AgenticLoopPlan:
if kwargs.get("_agentic_loop_api_surface") == CHAT_COMPLETION_AGENTIC_SURFACE:
return await self._build_chat_completion_agentic_loop_plan(
@ -368,14 +436,14 @@ class CodeInterpreterInterceptionLogger(CustomLogger):
)
await self._prune_expired_cache()
tool_calls = cast(list[CodeExecutionToolCall], tools.get("tool_calls", []))
sandbox_key = kwargs.get(_SANDBOX_KEY)
tool_calls = self._agentic_tool_calls(tools)
sandbox_key = self._extract_sandbox_key(kwargs)
is_session = bool(kwargs.get(_SESSION_SCOPED_KEY))
identity = _extract_identity(kwargs) if is_session else None
container, params = await self._get_or_create_container(cache_key=sandbox_key, identity=identity)
try:
container_id = cast(str | None, getattr(container, "id", None))
container_id = self._container_id(container)
input_list = self._normalize_messages(messages)
code_interpreter_calls: list[CodeInterpreterCall] = []
for tool_call in tool_calls:
@ -443,14 +511,14 @@ class CodeInterpreterInterceptionLogger(CustomLogger):
kwargs: dict[str, object],
) -> AgenticLoopPlan:
await self._prune_expired_cache()
tool_calls = cast(list[CodeExecutionToolCall], tools.get("tool_calls", []))
sandbox_key = cast(str | None, kwargs.get(_SANDBOX_KEY))
tool_calls = self._agentic_tool_calls(tools)
sandbox_key = self._extract_sandbox_key(kwargs)
is_session = bool(kwargs.get(_SESSION_SCOPED_KEY))
identity = _extract_identity(cast(dict[str, Any], kwargs)) if is_session else None
identity = _extract_identity(kwargs) if is_session else None
container, params = await self._get_or_create_container(cache_key=sandbox_key, identity=identity)
try:
container_id = cast(str | None, getattr(container, "id", None))
container_id = self._container_id(container)
tool_results = [
await self._build_chat_completion_tool_result(
container=container,
@ -489,10 +557,28 @@ class CodeInterpreterInterceptionLogger(CustomLogger):
},
)
@staticmethod
def _container_id(container: ContainerHandle) -> str | None:
container_id: object = getattr(container, "id", None)
return container_id if isinstance(container_id, str) else None
@staticmethod
def _agentic_tool_calls(tools: dict[str, object]) -> list[CodeExecutionToolCall]:
tool_calls = tools.get("tool_calls")
if not isinstance(tool_calls, list):
return []
items: list[object] = tool_calls
return [_narrow_tool_call(item) for item in items if isinstance(item, dict)]
@staticmethod
def _extract_sandbox_key(kwargs: dict[str, object]) -> str | None:
sandbox_key = kwargs.get(_SANDBOX_KEY)
return sandbox_key if isinstance(sandbox_key, str) else None
async def _build_chat_completion_tool_result(
self,
container: object,
params: dict[str, Any] | None,
container: ContainerHandle,
params: SandboxToolParams | None,
tool_call: CodeExecutionToolCall,
container_id: str | None,
) -> tuple[ChatCompletionToolMessage, CodeInterpreterCall]:
@ -517,10 +603,15 @@ class CodeInterpreterInterceptionLogger(CustomLogger):
)
async def async_agentic_loop_cleanup_hook(self, plan: AgenticLoopPlan, kwargs: dict) -> None:
metadata = plan.metadata or {} if plan else {}
metadata: dict[str, object] = plan.metadata or {} if plan else {}
if metadata.get("is_session_scoped"):
return
await self._delete_container_for_cache_key(metadata.get("sandbox_key"))
await self._delete_container_for_cache_key(self._metadata_sandbox_key(metadata))
@staticmethod
def _metadata_sandbox_key(metadata: dict[str, object]) -> str | None:
sandbox_key = metadata.get("sandbox_key")
return sandbox_key if isinstance(sandbox_key, str) else None
@staticmethod
def _filter_agentic_loop_kwargs(kwargs: dict[str, object]) -> dict[str, object]:
@ -531,12 +622,12 @@ class CodeInterpreterInterceptionLogger(CustomLogger):
and not is_interception_internal_key(k, prefixes=NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES)
}
def _get_followup_tools(self, tools: object, call_type: CallTypes | None) -> list[dict[str, Any]] | None:
def _get_followup_tools(self, tools: object, call_type: CallTypes | None) -> list[dict[str, object]] | None:
if not isinstance(tools, list):
return None
return [
(
self._get_function_tool(call_type=call_type)
dict(self._get_function_tool(call_type=call_type))
if isinstance(tool, dict) and tool.get("type") == "code_interpreter"
else tool
)
@ -549,34 +640,42 @@ class CodeInterpreterInterceptionLogger(CustomLogger):
k: v for k, v in optional_params.items() if k != "tools" and not (k == "tool_choice" and drop_tool_choice)
}
async def async_post_agentic_loop_response_hook(self, response: Any, plan: AgenticLoopPlan, kwargs: dict) -> Any:
metadata = plan.metadata or {} if plan else {}
async def async_post_agentic_loop_response_hook(
self, response: object, plan: AgenticLoopPlan, kwargs: dict
) -> object:
metadata: dict[str, object] = plan.metadata or {} if plan else {}
if not metadata.get("is_session_scoped"):
await self._delete_container_for_cache_key(metadata.get("sandbox_key"))
await self._delete_container_for_cache_key(self._metadata_sandbox_key(metadata))
calls = metadata.get("code_interpreter_calls")
if not calls:
if not calls or not isinstance(calls, list):
return response
is_dict = isinstance(response, dict)
output = response.get("output") if is_dict else getattr(response, "output", None)
if not isinstance(output, list):
if isinstance(response, dict):
response_mapping: dict[str, object] = response
merged = self._merge_code_interpreter_calls(response_mapping.get("output"), calls)
if merged is not None:
response_mapping["output"] = merged
return response
def _item_type(item: Any) -> Any:
return item.get("type") if isinstance(item, dict) else getattr(item, "type", None)
insert_at = next(
(i for i, item in enumerate(output) if _item_type(item) == "message"),
len(output),
)
new_output = output[:insert_at] + list(calls) + output[insert_at:]
if is_dict:
response["output"] = new_output
else:
response.output = new_output
if not isinstance(response, _SupportsOutput):
return response
merged = self._merge_code_interpreter_calls(response.output, calls)
if merged is not None:
response.output = merged
return response
@staticmethod
def _merge_code_interpreter_calls(output: object, calls: Sequence[object]) -> list[object] | None:
if not isinstance(output, list):
return None
items: list[object] = output
insert_at = next(
(i for i, item in enumerate(items) if _output_item_type(item) == "message"),
len(items),
)
return items[:insert_at] + list(calls) + items[insert_at:]
@staticmethod
def _parse_code(arguments: str) -> str:
try:
@ -584,7 +683,7 @@ class CodeInterpreterInterceptionLogger(CustomLogger):
except (json.JSONDecodeError, TypeError, AttributeError):
return ""
async def _run_tool_call(self, container: Any, params: dict[str, Any] | None, arguments: str) -> str:
async def _run_tool_call(self, container: ContainerHandle, params: SandboxToolParams | None, arguments: str) -> str:
try:
code = json.loads(arguments).get("code", "") if arguments else ""
except (json.JSONDecodeError, TypeError):
@ -601,7 +700,7 @@ class CodeInterpreterInterceptionLogger(CustomLogger):
self,
cache_key: str | None,
identity: str | None = None,
) -> tuple[Any, dict[str, Any] | None]:
) -> tuple[ContainerHandle, SandboxToolParams | None]:
if cache_key:
cached = self._container_cache.get(cache_key)
if cached is not None:
@ -623,7 +722,7 @@ class CodeInterpreterInterceptionLogger(CustomLogger):
self._container_cache.pop(lru_key, None)
await self._delete_container(container=lru_entry[0], params=lru_entry[1])
async def _create_container(self) -> tuple[Any, dict[str, Any] | None]:
async def _create_container(self) -> tuple[ContainerHandle, SandboxToolParams | None]:
if self.sandbox_config is not None:
return await self.sandbox_config.acreate_sandbox(), None
@ -641,7 +740,9 @@ class CodeInterpreterInterceptionLogger(CustomLogger):
)
return container, params
async def _run_code(self, container: Any, params: dict[str, Any] | None, code: str) -> Any:
async def _run_code(
self, container: ContainerHandle, params: SandboxToolParams | None, code: str
) -> CodeExecutionResult:
if self.sandbox_config is not None:
return await self.sandbox_config.arun_code(container=container, code=code)
if params is None:
@ -653,7 +754,7 @@ class CodeInterpreterInterceptionLogger(CustomLogger):
api_key=params.get("api_key"),
)
async def _delete_container(self, container: Any, params: dict[str, Any] | None) -> None:
async def _delete_container(self, container: ContainerHandle, params: SandboxToolParams | None) -> None:
try:
if self.sandbox_config is not None:
await self.sandbox_config.adelete_sandbox(container=container)
@ -677,7 +778,7 @@ class CodeInterpreterInterceptionLogger(CustomLogger):
return
await self._delete_container(container=cached[0], params=cached[1])
def _normalize_messages(self, messages: Any) -> list[dict[str, Any]]:
def _normalize_messages(self, messages: object) -> list[dict[str, object]]:
if isinstance(messages, str):
return [{"role": "user", "content": messages}]
if isinstance(messages, list):
@ -686,7 +787,8 @@ class CodeInterpreterInterceptionLogger(CustomLogger):
def _extract_code_execution_tool_calls(self, response: object) -> list[CodeExecutionToolCall]:
if isinstance(response, dict):
output = response.get("output", [])
response_mapping: dict[str, object] = response
output: object = response_mapping.get("output", [])
else:
output = getattr(response, "output", []) or []
if not isinstance(output, list):
@ -702,9 +804,7 @@ class CodeInterpreterInterceptionLogger(CustomLogger):
if self._is_code_execution_call(item)
]
def _extract_chat_completion_code_execution_tool_calls(
self, response: ModelResponse | dict[str, Any]
) -> list[CodeExecutionToolCall]:
def _extract_chat_completion_code_execution_tool_calls(self, response: object) -> list[CodeExecutionToolCall]:
model_response = self._to_model_response(response)
if model_response is None:
return []
@ -743,44 +843,46 @@ class CodeInterpreterInterceptionLogger(CustomLogger):
@staticmethod
def _build_chat_completion_assistant_message(
tool_calls: list[CodeExecutionToolCall],
tool_calls: Sequence[CodeExecutionToolCall],
) -> ChatCompletionAssistantMessage:
assistant_tool_calls: list[ChatCompletionAssistantToolCall] = [
{
"id": tool_call.get("id"),
"type": "function",
"function": {
"name": LITELLM_CODE_EXECUTION_TOOL_NAME,
"arguments": tool_call.get("arguments", ""),
},
}
for tool_call in tool_calls
]
return {
"role": "assistant",
"tool_calls": [
cast(
ChatCompletionAssistantToolCall,
{
"id": tool_call.get("id"),
"type": "function",
"function": {
"name": LITELLM_CODE_EXECUTION_TOOL_NAME,
"arguments": tool_call.get("arguments", ""),
},
},
)
for tool_call in tool_calls
],
"tool_calls": assistant_tool_calls,
}
@staticmethod
def _to_model_response(
response: ModelResponse | dict[str, Any],
) -> ModelResponse | None:
def _to_model_response(response: object) -> ModelResponse | None:
if isinstance(response, ModelResponse):
return response
if not isinstance(response, dict):
return None
response_fields: dict[str, object] = response
try:
return ModelResponse(**response)
return ModelResponse(**response_fields)
except (TypeError, ValidationError):
return None
def _is_code_execution_call(self, item: Any) -> bool:
def _is_code_execution_call(self, item: object) -> bool:
if isinstance(item, dict):
return item.get("type") == "function_call" and item.get("name") == LITELLM_CODE_EXECUTION_TOOL_NAME
return (
getattr(item, "type", None) == "function_call"
and getattr(item, "name", None) == LITELLM_CODE_EXECUTION_TOOL_NAME
)
item_mapping: dict[str, object] = item
return (
item_mapping.get("type") == "function_call"
and item_mapping.get("name") == LITELLM_CODE_EXECUTION_TOOL_NAME
)
item_type: object = getattr(item, "type", None)
item_name: object = getattr(item, "name", None)
return item_type == "function_call" and item_name == LITELLM_CODE_EXECUTION_TOOL_NAME
async def _prune_expired_cache(self) -> None:
now = time.time()

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

@ -1,5 +1,6 @@
from collections.abc import AsyncIterator, Coroutine, Iterator
from typing import (
TYPE_CHECKING,
Any,
cast,
)
@ -24,6 +25,10 @@ from litellm.types.llms.anthropic_messages.anthropic_response import (
from litellm.types.utils import ModelResponse
from litellm.utils import get_model_info
if TYPE_CHECKING:
from litellm.proxy._types import UserAPIKeyAuth
from litellm.router import Router
# Anthropic-only keys already mapped by the translator; strip on extra_kwargs re-merge.
ANTHROPIC_ONLY_REQUEST_KEYS: frozenset[str] = frozenset({"output_config"})
@ -67,8 +72,8 @@ async def _prepare_context_managed_request(
context_management_spec: Any,
litellm_metadata: dict | None,
additional_drop_params: list[str] | None,
llm_router: Any,
user_api_key_auth: Any = None,
llm_router: "Router | None",
user_api_key_auth: "UserAPIKeyAuth | None" = None,
) -> PolyfillResult | None:
"""Apply client compaction history, then optional context_management polyfill."""
from litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact import (
@ -152,7 +157,7 @@ def _polyfill_will_run(
COMPACT_EDIT_TYPE,
)
return any(isinstance(edit, dict) and edit.get("type") == COMPACT_EDIT_TYPE for edit in edits)
return any(edit.get("type") == COMPACT_EDIT_TYPE for edit in edits)
def _spec_has_non_compact_edits(
@ -178,10 +183,7 @@ def _spec_has_non_compact_edits(
COMPACT_EDIT_TYPE,
)
return any(
isinstance(edit, dict) and isinstance(edit.get("type"), str) and edit.get("type") != COMPACT_EDIT_TYPE
for edit in edits
)
return any(isinstance(edit.get("type"), str) and edit.get("type") != COMPACT_EDIT_TYPE for edit in edits)
def _context_management_explicitly_dropped(additional_drop_params: list[str] | None) -> bool:
@ -231,8 +233,8 @@ async def _run_polyfill_if_enabled(
context_management_spec: Any,
litellm_metadata: dict | None,
additional_drop_params: list[str] | None,
llm_router: Any,
user_api_key_auth: Any = None,
llm_router: "Router | None",
user_api_key_auth: "UserAPIKeyAuth | None" = None,
) -> PolyfillResult | None:
"""Run the async context_management polyfill if a spec is present.
@ -342,7 +344,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:
@ -531,11 +533,11 @@ class LiteLLMMessagesToCompletionTransformationHandler:
top_p: float | None = None,
output_format: dict | None = None,
**kwargs,
) -> AnthropicMessagesResponse | AsyncIterator[Any] | Iterator[bytes]:
) -> 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)
litellm_router = kwargs.pop("litellm_router", None)
litellm_router: Router | None = kwargs.pop("litellm_router", None)
if litellm_router is None:
try:
from litellm.proxy.proxy_server import llm_router as _proxy_router
@ -545,7 +547,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
pass
proxy_litellm_metadata = _extract_proxy_litellm_metadata(kwargs)
user_api_key_auth = (
user_api_key_auth: UserAPIKeyAuth | None = (
proxy_litellm_metadata.get("user_api_key_auth") if proxy_litellm_metadata is not None else None
)
@ -629,8 +631,8 @@ class LiteLLMMessagesToCompletionTransformationHandler:
) -> (
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:
@ -670,7 +672,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
# ``llm_router`` is ``None``, which is safe to call from the bridged
# loop. The async ``async_anthropic_messages_handler`` path is
# unaffected because it ``await``s within the original event loop.
litellm_router = kwargs.pop("litellm_router", None)
litellm_router: Router | None = kwargs.pop("litellm_router", None)
# Skip the async bridge entirely when there is nothing for either the
# polyfill or the client-history slice-only fallback to do. The vast
@ -682,7 +684,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
polyfill_result: PolyfillResult | None = None
else:
proxy_litellm_metadata = _extract_proxy_litellm_metadata(kwargs)
user_api_key_auth = (
user_api_key_auth: UserAPIKeyAuth | None = (
proxy_litellm_metadata.get("user_api_key_auth") if proxy_litellm_metadata is not None else None
)
polyfill_result = run_async_function(

86
litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py
View file

@ -4,11 +4,12 @@ import copy
import json
import traceback
from collections import deque
from collections.abc import AsyncIterator, Iterator
from collections.abc import AsyncIterator, Iterator, Sequence
from typing import (
TYPE_CHECKING,
Any,
Literal,
Protocol,
get_args,
)
@ -19,7 +20,9 @@ from litellm._uuid import uuid
from litellm.types.llms.anthropic import (
AppliedEdit,
CompactionBlock,
ContentBlockDelta,
ContextManagementResponse,
MessageBlockDelta,
StreamingContentBlockDeltaType,
UsageDelta,
UsageIteration,
@ -33,6 +36,25 @@ if TYPE_CHECKING:
_STREAMING_DELTA_TYPES = frozenset(get_args(StreamingContentBlockDeltaType))
class _UsageDeltaWithIterations(UsageDelta, total=False):
iterations: list[UsageIteration]
class _ChunkStream(Protocol):
def __iter__(self) -> "Iterator[ModelResponseStream]": ...
def __aiter__(self) -> "AsyncIterator[ModelResponseStream]": ...
def _optional_attr(obj: object, name: str) -> object:
return getattr(obj, name, None)
def _optional_attr_sequence(obj: object, name: str) -> Sequence[object]:
value = getattr(obj, name, None)
return value if value else ()
def _delta_payload_field(delta_type: StreamingContentBlockDeltaType) -> str:
match delta_type:
case "text_delta":
@ -67,29 +89,29 @@ class _CombinedChunkSplitter:
would advance them out of sync.
"""
def __init__(self, completion_stream: Any):
self._stream = completion_stream
self._sync_iter: Iterator[Any] | None = None
self._async_iter: AsyncIterator[Any] | None = None
self._buffer: deque = deque()
def __init__(self, completion_stream: _ChunkStream):
self._stream: _ChunkStream = completion_stream
self._sync_iter: Iterator[ModelResponseStream] | None = None
self._async_iter: AsyncIterator[ModelResponseStream] | None = None
self._buffer: deque[ModelResponseStream] = deque()
@staticmethod
def _is_combined(chunk: Any) -> bool:
def _is_combined(chunk: "ModelResponseStream") -> bool:
"""True if ``chunk`` carries response content AND a finish_reason."""
choices = getattr(chunk, "choices", None)
choices = _optional_attr_sequence(chunk, "choices")
if not choices:
return False
choice = choices[0]
if getattr(choice, "finish_reason", None) is None:
if _optional_attr(choice, "finish_reason") is None:
return False
delta = getattr(choice, "delta", None)
delta = _optional_attr(choice, "delta")
if delta is None:
return False
return bool(
getattr(delta, "content", None)
or getattr(delta, "tool_calls", None)
or getattr(delta, "reasoning_content", None)
or getattr(delta, "thinking_blocks", None)
_optional_attr(delta, "content")
or _optional_attr(delta, "tool_calls")
or _optional_attr(delta, "reasoning_content")
or _optional_attr(delta, "thinking_blocks")
)
_PAYLOAD_FIELD_GROUPS: "tuple[tuple[str, ...], ...]" = (
@ -124,21 +146,21 @@ class _CombinedChunkSplitter:
normalized to ``reasoning_content`` so the synthesized block start
stays empty and the thinking text is emitted exactly once.
"""
choices = getattr(chunk, "choices", None)
if not choices or len(choices) != 1:
choices = _optional_attr_sequence(chunk, "choices")
if len(choices) != 1:
return (chunk,)
delta = getattr(choices[0], "delta", None)
delta = _optional_attr(choices[0], "delta")
if delta is None:
return (chunk,)
tool_calls = getattr(delta, "tool_calls", None)
tool_calls = _optional_attr_sequence(delta, "tool_calls")
if tool_calls and not any(
getattr(getattr(tool_call, "function", None), "name", None) for tool_call in tool_calls
_optional_attr(_optional_attr(tool_call, "function"), "name") for tool_call in tool_calls
):
return (chunk,)
present_groups = tuple(
group
for group in _CombinedChunkSplitter._PAYLOAD_FIELD_GROUPS
if any(getattr(delta, field, None) for field in group)
if any(_optional_attr(delta, field) for field in group)
)
if len(present_groups) <= 1:
return (chunk,)
@ -177,7 +199,7 @@ class _CombinedChunkSplitter:
return {"reasoning_content": thinking_text}
@staticmethod
def _split(chunk: Any) -> list[Any]:
def _split(chunk: "ModelResponseStream") -> "list[ModelResponseStream]":
"""Return ``[chunk]``, or ``[content_chunk, finish_chunk]`` if combined."""
if not _CombinedChunkSplitter._is_combined(chunk):
return [chunk]
@ -199,10 +221,10 @@ class _CombinedChunkSplitter:
finish_delta.thinking_blocks = None
return [content_chunk, finish_chunk]
def __iter__(self) -> "Iterator[Any]":
def __iter__(self) -> "Iterator[ModelResponseStream]":
return self
def __next__(self) -> Any:
def __next__(self) -> "ModelResponseStream":
if self._buffer:
return self._buffer.popleft()
if self._sync_iter is None:
@ -215,10 +237,10 @@ class _CombinedChunkSplitter:
)
return self._buffer.popleft()
def __aiter__(self) -> "AsyncIterator[Any]":
def __aiter__(self) -> "AsyncIterator[ModelResponseStream]":
return self
async def __anext__(self) -> Any:
async def __anext__(self) -> "ModelResponseStream":
if self._buffer:
return self._buffer.popleft()
if self._async_iter is None:
@ -251,14 +273,14 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
sent_content_block_finish: bool = False
current_content_block_type: Literal["text", "tool_use", "thinking"] = "text"
sent_last_message: bool = False
holding_chunk: Any | None = None
holding_stop_reason_chunk: Any | None = None
holding_chunk: ContentBlockDelta | None = None
holding_stop_reason_chunk: MessageBlockDelta | None = None
queued_usage_chunk: bool = False
current_content_block_index: int = 0
def __init__(
self,
completion_stream: Any,
completion_stream: _ChunkStream,
model: str,
tool_name_mapping: dict[str, str] | None = None,
applied_edits: list[AppliedEdit] | None = None,
@ -299,7 +321,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
text="",
)
def _merge_usage_into_held_stop_reason_chunk(self, chunk: Any) -> dict[str, Any]:
def _merge_usage_into_held_stop_reason_chunk(self, chunk: Any) -> MessageBlockDelta:
"""Merge usage data from ``chunk`` into the held ``message_delta`` chunk.
Shared by both the sync ``__next__`` and async ``__anext__`` paths so
@ -325,7 +347,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
merged_chunk["context_management"] = ContextManagementResponse(applied_edits=list(self.applied_edits))
return self._augment_message_delta_usage(merged_chunk)
def _ensure_context_management_attached(self, message_delta_chunk: dict[str, Any]) -> dict[str, Any]:
def _ensure_context_management_attached(self, message_delta_chunk: MessageBlockDelta) -> MessageBlockDelta:
"""Attach ``context_management`` to a ``message_delta`` chunk if
``self.applied_edits`` is non-empty and the chunk does not already
carry it. Returns the (possibly new) chunk dict.
@ -340,7 +362,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
augmented["context_management"] = ContextManagementResponse(applied_edits=list(self.applied_edits))
return augmented
def _augment_message_delta_usage(self, message_delta_chunk: dict[str, Any]) -> dict[str, Any]:
def _augment_message_delta_usage(self, message_delta_chunk: MessageBlockDelta) -> MessageBlockDelta:
"""Attach polyfill compaction iteration usage to the final message_delta.
Also defensively re-attaches ``context_management`` so the direct
@ -357,7 +379,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
input_tokens = usage.get("input_tokens", 0) or 0
output_tokens = usage.get("output_tokens", 0) or 0
augmented = message_delta_chunk.copy()
augmented_usage = dict(usage)
augmented_usage: _UsageDeltaWithIterations = {**usage}
iterations: list[UsageIteration] = list(self.iterations_usage)
# Only emit a ``message`` iteration when we have real token data.
# Without a separate usage chunk (e.g. provider sent finish_reason

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

@ -12,7 +12,7 @@ MCP Spec Reference:
import typing
from collections.abc import Mapping, Sequence
from typing import Any, NamedTuple, Optional, Protocol, Union
from typing import Any, NamedTuple, Optional, Protocol, Union, runtime_checkable
if typing.TYPE_CHECKING:
from fastapi import Request
@ -24,6 +24,7 @@ if typing.TYPE_CHECKING:
from litellm.proxy.utils import ProxyLogging
from fastapi import HTTPException
from pydantic import TypeAdapter
from litellm._logging import verbose_logger
@ -295,8 +296,14 @@ def _convert_mcp_content_to_openai(
return _convert_single_content(content)
@runtime_checkable
class _TextContentLike(Protocol):
@property
def text(self) -> object: ...
def _convert_single_content(
content: Any,
content: object,
) -> "dict[str, object] | list[dict[str, object]]":
"""Convert a single MCP content item to OpenAI format.
@ -308,12 +315,14 @@ def _convert_single_content(
"""
import json
content_type = getattr(content, "type", None)
content_type: str | None = getattr(content, "type", None)
if content_type == "text":
if not isinstance(content, _TextContentLike):
raise AttributeError(f"{type(content).__name__!r} object has no attribute 'text'")
return {"type": "text", "text": content.text}
elif content_type == "image":
data = getattr(content, "data", "")
mime_type = getattr(content, "mimeType", "image/png")
data: str = getattr(content, "data", "")
mime_type: str = getattr(content, "mimeType", "image/png")
return {
"type": "image_url",
"image_url": {"url": f"data:{mime_type};base64,{data}"},
@ -339,13 +348,16 @@ def _convert_single_content(
# The ``_marker_type`` key lets the message-level converter
# hoist this into the ``tool_calls`` array on the assistant
# message instead of embedding it inline as a content part.
tool_use_id: str = getattr(content, "id", f"call_{id(content)}")
tool_name: str = getattr(content, "name", "")
tool_input: dict[str, object] = getattr(content, "input", {})
return {
"_marker_type": "tool_use",
"id": getattr(content, "id", f"call_{id(content)}"),
"id": tool_use_id,
"type": "function",
"function": {
"name": getattr(content, "name", ""),
"arguments": json.dumps(getattr(content, "input", {}), default=str),
"name": tool_name,
"arguments": json.dumps(tool_input, default=str),
},
}
elif content_type == "tool_result":
@ -581,12 +593,28 @@ def _convert_mcp_tool_choice_to_openai(
return "auto"
class _SamplingToolCallFunction(Protocol):
@property
def name(self) -> str | None: ...
@property
def arguments(self) -> object: ...
class _SamplingToolCall(Protocol):
@property
def id(self) -> str | None: ...
@property
def function(self) -> _SamplingToolCallFunction: ...
class _SamplingResponseMessage(Protocol):
@property
def content(self) -> str | None: ...
@property
def tool_calls(self) -> Sequence[object] | None: ...
def tool_calls(self) -> Sequence[_SamplingToolCall] | None: ...
class _SamplingResponseChoice(Protocol):
@ -605,6 +633,21 @@ class _SamplingCompletionResponse(Protocol):
def model(self) -> str | None: ...
_TOOL_ARGUMENTS_ADAPTER = TypeAdapter(dict[str, object])
def _parse_tool_arguments(arguments: object) -> "dict[str, object]":
"""Decode OpenAI tool-call arguments into the MCP ``input`` mapping."""
import json
if not isinstance(arguments, str):
return _TOOL_ARGUMENTS_ADAPTER.validate_python(arguments)
try:
return _TOOL_ARGUMENTS_ADAPTER.validate_python(json.loads(arguments))
except (json.JSONDecodeError, TypeError):
return {"raw": arguments}
def _convert_openai_response_to_mcp_result(
response: _SamplingCompletionResponse,
model_name: str,
@ -641,7 +684,7 @@ def _convert_openai_response_to_mcp_result(
stop_reason = "endTurn"
actual_model: str = getattr(response, "model", model_name) or model_name
# Check if response has tool calls
tool_calls = getattr(message, "tool_calls", None)
tool_calls = message.tool_calls if hasattr(message, "tool_calls") else None
if tool_calls:
# Build ToolUseContent items
content_parts: list[SamplingMessageContentBlock] = []
@ -650,20 +693,14 @@ def _convert_openai_response_to_mcp_result(
content_parts.append(TextContent(type="text", text=message.content))
# Convert tool calls to MCP ToolUseContent
for tc in tool_calls:
import json
tool_input = tc.function.arguments
if isinstance(tool_input, str):
try:
tool_input = json.loads(tool_input)
except (json.JSONDecodeError, TypeError):
tool_input = {"raw": tool_input}
content_parts.append(
ToolUseContent(
type="tool_use",
id=tc.id,
name=tc.function.name,
input=tool_input,
ToolUseContent.model_validate(
{
"type": "tool_use",
"id": tc.id,
"name": tc.function.name,
"input": _parse_tool_arguments(tc.function.arguments),
}
)
)
return CreateMessageResultWithTools(
@ -1101,7 +1138,7 @@ async def _build_completion_kwargs(
messages=params.messages,
system_prompt=params.systemPrompt,
)
completion_kwargs: dict[str, Any] = {
completion_kwargs: dict[str, object] = {
"model": model,
"messages": openai_messages,
"max_tokens": params.maxTokens,
@ -1116,9 +1153,7 @@ async def _build_completion_kwargs(
openai_tool_choice = _convert_mcp_tool_choice_to_openai(params.toolChoice)
if openai_tool_choice is not None:
completion_kwargs["tool_choice"] = openai_tool_choice
completion_kwargs["metadata"] = {}
if params.metadata:
completion_kwargs["metadata"]["mcp_metadata"] = params.metadata
completion_kwargs["metadata"] = {"mcp_metadata": params.metadata} if params.metadata else {}
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
from litellm.proxy.proxy_server import proxy_config

193
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 Callable, Sequence
from contextvars import ContextVar
from dataclasses import dataclass, field
from datetime import datetime
@ -16,9 +16,9 @@ from typing import (
TYPE_CHECKING,
Any,
Literal,
Protocol,
TypedDict,
Union,
cast,
)
from litellm import DualCache
@ -54,6 +54,7 @@ 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]
@ -300,6 +301,13 @@ _TPM_FLOOR_FRACTION = 4
PARALLEL_REQUEST_SLOT_TTL_SECONDS = 3600
CacheCounterValue = int | float | str | bytes
CacheCounterValues = Sequence[CacheCounterValue | None]
ParallelGaugeCacheValue = dict[str, object] | int | float | str | bytes
class RateLimitDescriptorRateLimitObject(TypedDict, total=False):
requests_per_unit: int | None
tokens_per_unit: int | None
@ -342,6 +350,42 @@ class RateLimitResponseWithDescriptors(TypedDict):
response: RateLimitResponse
class WindowKeyMetadata(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
DescriptorAtomicGroup = tuple[list[str], list[int], list[AtomicCounterMeta]]
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: ...
@dataclass(slots=True)
class RequestRateLimiterStash:
"""
@ -459,7 +503,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 +521,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:
@ -606,12 +650,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
keys: list[str],
now_int: int,
window_size: int,
) -> list[Any]:
) -> CacheCounterValues:
"""
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[CacheCounterValue | None] = []
# Process each window/counter pair
for i in range(0, len(keys), 2):
@ -620,7 +664,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
increment_value = 1
# Get the window start time
window_start = await self.internal_usage_cache.async_get_cache(
window_start: CacheCounterValue | None = await self.internal_usage_cache.async_get_cache(
key=window_key,
litellm_parent_otel_span=None,
local_only=True,
@ -647,7 +691,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
results.append(increment_value) # counter
else:
# Increment the counter
current_counter = await self.internal_usage_cache.async_get_cache(
current_counter: CacheCounterValue | None = await self.internal_usage_cache.async_get_cache(
key=counter_key,
litellm_parent_otel_span=None,
local_only=True,
@ -681,8 +725,8 @@ 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],
cache_values: CacheCounterValues,
key_metadata: dict[str, WindowKeyMetadata],
) -> RateLimitResponse:
"""
Check if the cache values are over the limit.
@ -781,11 +825,36 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
return groups
async def _batch_get_counter_values(
self,
keys: list[str],
parent_otel_span: Span | None,
local_only: bool,
) -> CacheCounterValues | None:
"""Typed view over the DualCache batch read of window/counter keys."""
return await self.internal_usage_cache.async_batch_get_cache(
keys=keys,
parent_otel_span=parent_otel_span,
local_only=local_only,
)
async def _batch_get_gauge_values(
self,
keys: list[str],
parent_otel_span: Span | None,
) -> Sequence[ParallelGaugeCacheValue | None] | None:
"""Typed view over the DualCache batch read of parallel-request gauges."""
return await self.internal_usage_cache.async_batch_get_cache(
keys=keys,
parent_otel_span=parent_otel_span,
local_only=True,
)
async def _execute_redis_batch_rate_limiter_script(
self,
keys_to_fetch: list[str],
now_int: int,
) -> list[Any]:
) -> CacheCounterValues:
"""
Execute Redis operations grouped by hash tag for cluster compatibility.
@ -794,17 +863,17 @@ 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[CacheCounterValue | None] = []
for hash_tag, group_keys in key_groups.items():
try:
group_cache_values = await self.batch_rate_limiter_script(
group_cache_values: CacheCounterValues = await self.batch_rate_limiter_script(
keys=group_keys,
args=[now_int, self.window_size], # Use integer timestamp
)
@ -868,7 +937,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 = await self._batch_get_counter_values(
keys=keys_to_fetch,
parent_otel_span=parent_otel_span,
local_only=True,
@ -882,7 +951,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
## IF under limit in-memory, check Redis
if read_only:
# READ-ONLY MODE: Just read current values without incrementing
cache_values = await self.internal_usage_cache.async_batch_get_cache(
cache_values = await self._batch_get_counter_values(
keys=keys_to_fetch,
parent_otel_span=parent_otel_span,
local_only=False, # Check Redis too
@ -890,9 +959,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
@ -951,14 +1018,14 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
self,
descriptors: list[RateLimitDescriptor],
skip_tpm_check: bool,
) -> tuple[list[str], dict[str, dict[str, Any]], list[ParallelRequestGauge]]:
) -> tuple[list[str], dict[str, WindowKeyMetadata], 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, WindowKeyMetadata] = {}
gauges: list[ParallelRequestGauge] = []
for descriptor in descriptors:
descriptor_key = descriptor["key"]
@ -1014,7 +1081,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: ParallelGaugeCacheValue | None) -> int:
"""
In-flight count from a cached gauge value: a dict of slot_id ->
acquire timestamp when the in-memory registry is authoritative, or
@ -1051,7 +1118,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
if read_only:
if self.parallel_count_script is not None:
try:
raw_counts = await self.parallel_count_script(
raw_counts: list[CacheCounterValue] = await self.parallel_count_script(
keys=gauge_keys,
args=[PARALLEL_REQUEST_SLOT_TTL_SECONDS for _ in gauges],
)
@ -1080,7 +1147,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
if self.parallel_acquire_script is not None:
try:
raw = await self.parallel_acquire_script(
raw: list[CacheCounterValue] = await self.parallel_acquire_script(
keys=gauge_keys,
args=[
arg for gauge in gauges for arg in (gauge["limit"], PARALLEL_REQUEST_SLOT_TTL_SECONDS, slot_id)
@ -1116,10 +1183,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
gauge_keys: list[str],
parent_otel_span: Span | None = None,
) -> list[int]:
values = await self.internal_usage_cache.async_batch_get_cache(
values = await self._batch_get_gauge_values(
keys=gauge_keys,
parent_otel_span=parent_otel_span,
local_only=True,
)
if values is None:
return [0 for _ in gauge_keys]
@ -1145,7 +1211,7 @@ 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: ParallelGaugeCacheValue | None = await self.internal_usage_cache.async_get_cache(
key=gauge["counter_key"],
litellm_parent_otel_span=parent_otel_span,
local_only=True,
@ -1200,7 +1266,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
return
if self.parallel_release_script is not None:
try:
raw = await self.parallel_release_script(
raw: list[CacheCounterValue] = await self.parallel_release_script(
keys=counter_keys,
args=[slot_id for _ in counter_keys],
)
@ -1218,7 +1284,7 @@ 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: ParallelGaugeCacheValue | None = await self.internal_usage_cache.async_get_cache(
key=counter_key,
litellm_parent_otel_span=parent_otel_span,
local_only=True,
@ -1226,7 +1292,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
if isinstance(raw_value, dict):
if slot_id not in raw_value:
continue
new_value: dict[str, float] | int = {key: ts for key, ts in raw_value.items() if key != slot_id}
new_value: dict[str, object] | int = {key: ts for key, ts in raw_value.items() if key != slot_id}
elif raw_value is None:
continue
else:
@ -1277,7 +1343,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[DescriptorAtomicGroup] = []
for descriptor, increment_amounts in zip(descriptors, increments):
keys, args, meta = self._build_descriptor_atomic_payload(
descriptor=descriptor,
@ -1300,7 +1366,7 @@ 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: list[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,
@ -1311,7 +1377,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
self,
descriptor: RateLimitDescriptor,
increment_amounts: dict[Literal["requests", "tokens"], int],
) -> tuple[list[str], list[Any], list[dict[str, Any]]]:
) -> DescriptorAtomicGroup:
"""
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 +1391,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)
rate_limit_types: tuple[Literal["requests", "tokens"], ...] = ("requests", "tokens")
for rlt in rate_limit_types:
if rlt == "requests":
limit_value = rate_limit.get("requests_per_unit")
inc_amount = int(increment_amounts.get("requests", 0) or 0)
@ -1365,7 +1431,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: list[DescriptorAtomicGroup],
parent_otel_span: Span | None = None,
) -> RateLimitResponse:
"""
@ -1374,8 +1440,9 @@ 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[list[AtomicCounterMeta]] = []
statuses: list[RateLimitStatus] = []
raw: list[CacheCounterValue]
for _idx, (keys, args, meta) in enumerate(descriptor_groups):
try:
@ -1396,7 +1463,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
self.window_size,
)
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: list[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,
@ -1414,7 +1481,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
async def _refund_applied_descriptor_groups(
self,
applied: list[list[dict[str, Any]]],
applied: list[list[AtomicCounterMeta]],
) -> None:
"""
Decrement counters for descriptor groups already applied via Lua.
@ -1440,8 +1507,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
def _build_atomic_response(
self,
raw: list[Any],
per_counter_meta: list[dict[str, Any]],
raw: list[CacheCounterValue],
per_counter_meta: list[AtomicCounterMeta],
) -> RateLimitResponse:
"""Convert Lua script return value to RateLimitResponse.
@ -1492,7 +1559,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: list[AtomicCounterMeta],
parent_otel_span: Span | None = None,
) -> RateLimitResponse:
"""In-memory all-or-nothing check-and-increment. Caller holds lock.
@ -1507,27 +1574,25 @@ 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: CacheCounterValue | None = await self.internal_usage_cache.async_get_cache(
key=meta["window_key"],
litellm_parent_otel_span=parent_otel_span,
local_only=True,
)
window_expired = window_start is None or (now_int - int(window_start)) >= window_size
current_counter = (
0
raw_counter: CacheCounterValue | None = (
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.internal_usage_cache.async_get_cache(
key=meta["counter_key"],
litellm_parent_otel_span=parent_otel_span,
local_only=True,
)
)
current_counter = 0 if window_expired else int(raw_counter or 0)
over_limit = (
current_counter + meta["increment"] > meta["current_limit"]
if meta["increment"] > 0
@ -1919,7 +1984,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
@ -2245,7 +2310,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()
@ -2772,15 +2837,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
@staticmethod
def _merge_ratelimit_statuses_into_additional_headers(
additional_headers: dict[str, Any],
additional_headers: dict[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"]
@ -3014,9 +3079,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
async def async_logging_hook(
self,
kwargs: dict,
result: Any,
result: object,
call_type: str,
) -> tuple[dict, Any]:
) -> tuple[dict, object]:
"""
Mirror the pre-call rate-limit snapshot into the SLP so streaming
success callbacks see the same ``x-ratelimit-*`` headers the
@ -3033,8 +3098,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

311
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 Callable, Mapping, Sequence
from typing import TYPE_CHECKING, TypeVar, overload
from urllib.parse import quote, unquote
from fastapi import HTTPException
from pydantic import JsonValue
from litellm._logging import verbose_proxy_logger
from litellm.llms.base_llm.managed_resources.isolation import (
@ -48,9 +50,30 @@ from litellm.repositories.table_repositories import (
ManagedObjectRepository,
)
from litellm.types.llms.openai import OpenAIFileObject
from litellm.types.passthrough_endpoints.managed_id_rewriter import (
ManagedFileIdReader,
ManagedFileIdWriter,
ManagedFileRow,
ManagedFileTable,
ManagedListResponse,
ManagedObjectRow,
ManagedObjectTable,
ManagedResourceRow,
ManagedTable,
PrismaWhere,
PrismaWhereValue,
ResourceKind,
SortOrder,
)
from .managed_id_codec import ManagedIdPayload, decode, is_managed, new_managed_id
if TYPE_CHECKING:
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy.utils import PrismaClient
_RowT = TypeVar("_RowT", bound=ManagedResourceRow)
# ---------------------------------------------------------------------------
# Field map
# ---------------------------------------------------------------------------
@ -172,7 +195,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
@ -197,7 +220,7 @@ class _RawIdGuardBudget:
# ---------------------------------------------------------------------------
# Maps (provider, canonical_path) -> "files" | "batches"
_LIST_ROUTE_TABLE: dict[tuple[str, str], str] = {
_LIST_ROUTE_TABLE: dict[tuple[str, str], ResourceKind] = {
("openai", "/v1/files"): "files",
("openai", "/v1/batches"): "batches",
("azure", "/v1/files"): "files",
@ -259,12 +282,20 @@ def _canonical_path(route: str) -> str:
# ---------------------------------------------------------------------------
def _file_table(prisma_client: PrismaClient) -> ManagedFileTable:
return ManagedFileRepository(prisma_client).table
def _object_table(prisma_client: PrismaClient) -> ManagedObjectTable:
return ManagedObjectRepository(prisma_client).table
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.
@ -305,7 +336,7 @@ async def _resolve_one(
# 2. DB lookup — pick table based on raw ID prefix
if any(raw_id.startswith(p) for p in _FILE_PREFIXES):
# File table — use hook's internal cache for speed when available
if managed_files_hook is not None:
if isinstance(managed_files_hook, ManagedFileIdReader):
try:
file_row = await managed_files_hook.get_unified_file_id(
managed_id,
@ -322,9 +353,7 @@ 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 = await _file_table(prisma_client).find_first(where={"unified_file_id": managed_id})
if db_row is not None:
row_created_by = db_row.created_by
row_team_id = db_row.team_id
@ -338,9 +367,7 @@ 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 = await _object_table(prisma_client).find_first(where={"unified_object_id": managed_id})
if obj_row is not None:
row_created_by = obj_row.created_by
row_team_id = obj_row.team_id
@ -372,7 +399,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,7 +425,7 @@ 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(
candidates = await _file_table(prisma_client).find_many(
where={"flat_model_file_ids": {"has": raw_id}},
)
except Exception:
@ -419,7 +446,7 @@ 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(
existing = await _object_table(prisma_client).find_first(
where={"model_object_id": f"passthrough:{provider}:{raw_id}"}
)
except Exception:
@ -434,7 +461,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, JsonValue] | 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 +469,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 +482,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, JsonValue] | None = None,
is_create_route: bool = True,
) -> str:
"""Return an existing managed file ID or mint + store a new one."""
@ -479,7 +506,7 @@ 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(
candidates: list[ManagedFileRow] = await _file_table(prisma_client).find_many(
where={"flat_model_file_ids": {"has": raw_id}},
order={"created_at": "asc"},
)
@ -524,6 +551,8 @@ async def _mint_or_reuse_file(
raw_id.split("-", 1)[0],
)
if managed_files_hook is not None:
if not isinstance(managed_files_hook, ManagedFileIdWriter):
return raw_id
try:
await managed_files_hook.store_unified_file_id(
file_id=managed_id,
@ -551,9 +580,9 @@ async def _mint_or_reuse_object(
raw_id: str,
provider: str,
file_purpose: str,
body_snapshot: dict,
body_snapshot: Mapping[str, JsonValue],
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."""
@ -569,7 +598,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: ManagedObjectRow, 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 +627,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(prisma_client).update(
where={"unified_object_id": existing.unified_object_id},
data={
"file_object": json.dumps(body_snapshot),
@ -618,9 +647,7 @@ 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 = await _object_table(prisma_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)
existing = None
@ -635,7 +662,7 @@ async def _mint_or_reuse_object(
raw_id.split("_", 1)[0],
)
try:
await ManagedObjectRepository(prisma_client).table.upsert(
await _object_table(prisma_client).upsert(
where={"unified_object_id": managed_id},
data={
"create": {
@ -659,9 +686,7 @@ 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 = await _object_table(prisma_client).find_first(where={"model_object_id": namespaced_model_object_id})
except Exception:
raced = None
if raced is not None:
@ -681,11 +706,11 @@ async def rewrite_response_ids(
provider: str,
method: str,
route: str,
body: dict,
body: dict[str, JsonValue],
user_api_key_dict: UserAPIKeyAuth,
prisma_client: Any,
managed_files_hook: Any,
) -> dict:
prisma_client: PrismaClient | None,
managed_files_hook: CustomLogger | None,
) -> dict[str, JsonValue]:
"""
Mint managed IDs for raw provider values listed in
``BUILTIN_OUTPUT_ID_FIELD_MAP`` and swap them into *body*.
@ -795,7 +820,7 @@ 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:
def _parse_file_object(file_object: JsonValue) -> JsonValue:
"""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
@ -809,7 +834,7 @@ def _parse_file_object(file_object: Any) -> Any:
return file_object
def _empty_list_response() -> dict[str, Any]:
def _empty_list_response() -> ManagedListResponse:
return {
"object": "list",
"data": [],
@ -819,7 +844,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,18 +855,18 @@ 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: ResourceKind,
provider: str,
owner_filter: dict[str, Any],
query_params: dict[str, Any] | None,
) -> tuple[dict[str, Any], str]:
owner_filter: Mapping[str, PrismaWhereValue],
query_params: Mapping[str, str] | None,
) -> tuple[PrismaWhere, SortOrder]:
"""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)
fetch_order = "desc"
where: PrismaWhere = dict(owner_filter)
fetch_order: SortOrder = "desc"
cursor_id = after_id or before_id
# A cursor minted for a different provider would resolve to that provider's
@ -850,10 +875,8 @@ 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_table: ManagedFileTable | ManagedObjectTable = (
_file_table(prisma_client) if resource_kind == "files" else _object_table(prisma_client)
)
cursor_field = "unified_file_id" if resource_kind == "files" else "unified_object_id"
try:
@ -867,7 +890,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: PrismaWhere = {
"OR": [
{"created_at": {op: cursor_row.created_at}},
{
@ -885,25 +908,19 @@ async def _build_list_where_with_cursor(
async def _fetch_list_rows(
prisma_client: Any,
resource_kind: str,
where: dict[str, Any],
fetch_order: str,
open_table: Callable[[], ManagedTable[_RowT]],
where: PrismaWhere,
id_field: str,
fetch_order: SortOrder,
fetch_limit: int,
) -> list[Any] | None:
) -> list[_RowT] | 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.
try:
if resource_kind == "files":
return await ManagedFileRepository(prisma_client).table.find_many(
where=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}],
return await open_table().find_many(
where=where,
order=[{"created_at": fetch_order}, {id_field: fetch_order}],
take=fetch_limit,
)
except Exception:
@ -912,15 +929,15 @@ async def _fetch_list_rows(
async def _fetch_provider_scoped_list_rows(
prisma_client: Any,
resource_kind: str,
provider: str,
where: dict[str, Any],
fetch_order: str,
open_table: Callable[[], ManagedTable[_RowT]],
where: PrismaWhere,
provider_scope: PrismaWhere,
id_field: str,
fetch_order: SortOrder,
raw_limit: int,
fetch_limit: int,
) -> tuple[list[Any], bool]:
"""Fetch one page of list rows scoped to *provider* at the DB level.
) -> tuple[list[_RowT], bool]:
"""Fetch one page of list rows scoped to a 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
@ -931,15 +948,10 @@ async def _fetch_provider_scoped_list_rows(
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.
through to the upstream provider. ``open_table`` is opened inside that
guarded region so a client missing the managed tables fails closed too.
"""
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}:"}
rows = await _fetch_list_rows(prisma_client, resource_kind, scoped_where, fetch_order, fetch_limit)
rows = await _fetch_list_rows(open_table, {**where, **provider_scope}, id_field, fetch_order, fetch_limit)
if rows is None:
return [], False
@ -951,8 +963,8 @@ async def _fetch_provider_scoped_list_rows(
return page, has_more
def _serialize_file_list_item(row: Any) -> dict[str, Any]:
item: dict[str, Any] = {
def _serialize_file_list_item(row: ManagedFileRow) -> 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,
@ -964,8 +976,8 @@ def _serialize_file_list_item(row: Any) -> dict[str, Any]:
return item
def _serialize_batch_list_item(row: Any) -> dict[str, Any]:
item: dict[str, Any] = {}
def _serialize_batch_list_item(row: ManagedObjectRow) -> dict[str, JsonValue]:
item: dict[str, JsonValue] = {}
file_object = _parse_file_object(row.file_object)
if isinstance(file_object, dict):
item.update(file_object)
@ -974,20 +986,19 @@ def _serialize_batch_list_item(row: Any) -> dict[str, Any]:
return item
def _list_boundary_ids(rows: list[Any], resource_kind: str) -> tuple[str | None, str | None]:
def _list_boundary_ids(rows: Sequence[_RowT], get_id: Callable[[_RowT], 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)
return get_id(rows[0]), get_id(rows[-1])
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,
) -> ManagedListResponse | None:
"""Query the DB for managed IDs the caller owns and return an OpenAI-style
paginated list response.
@ -1020,21 +1031,31 @@ 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 = await _fetch_provider_scoped_list_rows(
lambda: _file_table(prisma_client),
where,
{"flat_model_file_ids": {"has": _passthrough_provider_marker(provider)}},
"unified_file_id",
fetch_order,
raw_limit,
fetch_limit,
)
data = [_serialize_file_list_item(row) for row in file_page]
first_id, last_id = _list_boundary_ids(file_page, lambda row: row.unified_file_id)
else:
data = [_serialize_batch_list_item(row) for row in page]
object_page, has_more = await _fetch_provider_scoped_list_rows(
lambda: _object_table(prisma_client),
where,
{"model_object_id": {"startswith": f"passthrough:{provider}:"}, "file_purpose": "batch"},
"unified_object_id",
fetch_order,
raw_limit,
fetch_limit,
)
data = [_serialize_batch_list_item(row) for row in object_page]
first_id, last_id = _list_boundary_ids(object_page, lambda row: row.unified_object_id)
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,
@ -1056,12 +1077,16 @@ async def list_passthrough_ids_from_db(
# ---------------------------------------------------------------------------
def _is_litellm_internal_key(key: object) -> bool:
return isinstance(key, str) and key.startswith("litellm_")
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 +1117,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 +1148,33 @@ async def rewrite_query_ids(
return mutated if rewritten_keys else params
@overload
async def rewrite_body_ids(
body: dict[str, Any] | None,
body: 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: ...
@overload
async def rewrite_body_ids(
body: list[object],
provider: str,
user_api_key_dict: UserAPIKeyAuth,
prisma_client: PrismaClient | None,
managed_files_hook: CustomLogger | None,
) -> list[object]: ...
async def rewrite_body_ids(
body: dict[str, object] | list[object] | None,
provider: str,
user_api_key_dict: UserAPIKeyAuth,
prisma_client: PrismaClient | None,
managed_files_hook: CustomLogger | None,
) -> dict[str, object] | list[object] | None:
"""
Recursively walk a request body dict/list and resolve any passthrough
managed IDs. Skips litellm internal keys (``litellm_*``).
@ -1140,27 +1185,33 @@ async def rewrite_body_ids(
budget = _RawIdGuardBudget()
async def _walk(node: Any, depth: int) -> Any:
async def _walk_mapping(node: dict[str, object], depth: int) -> dict[str, object]:
result: dict[str, object] = {}
changed_inner = False
for k, v in node.items():
# Skip litellm internal injection keys (e.g. litellm_logging_obj)
if _is_litellm_internal_key(k):
result[k] = v
continue
new_v = await _walk(v, depth + 1)
result[k] = new_v
if new_v is not v:
changed_inner = True
return result if changed_inner else node
async def _walk_sequence(node: list[object], depth: int) -> list[object]:
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
async def _walk(node: object, depth: int) -> object:
if depth >= _MAX_BODY_REWRITE_DEPTH:
return node
if isinstance(node, dict):
result: dict[str, Any] = {}
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_"):
result[k] = v
continue
new_v = await _walk(v, depth + 1)
result[k] = new_v
if new_v is not v:
changed_inner = True
return result if changed_inner else node
return await _walk_mapping(node, depth)
elif 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
return await _walk_sequence(node, depth)
elif isinstance(node, str):
if is_managed(node):
return await _resolve_one(node, provider, user_api_key_dict, prisma_client, managed_files_hook)
@ -1168,7 +1219,7 @@ async def rewrite_body_ids(
return node
return node
rewritten = await _walk(body, 0)
rewritten = await _walk_sequence(body, 0) if isinstance(body, list) else await _walk_mapping(body, 0)
if rewritten is not body:
verbose_proxy_logger.debug("managed_id_rewriter: body ids rewritten provider=%s", provider)
return rewritten

2
litellm/repositories/table_repositories.py
View file

@ -27,7 +27,7 @@ class PrismaTableRepository:
return self._prisma_client
@property
def table(self) -> Any:
def table(self) -> Any: # any-ok: Prisma table actions are reached through the untyped client wrapper
return wrap_table_actions_for_config_sync(
actions=getattr(self.prisma_client.db, self.table_name),
table_name=self.table_name,

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

@ -1,3 +1,4 @@
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, cast
from litellm._logging import verbose_logger
@ -23,6 +24,8 @@ from litellm.types.llms.openai import (
if TYPE_CHECKING:
from mcp.types import Tool as MCPTool
from litellm.proxy._types import UserAPIKeyAuth
else:
MCPTool = Any
@ -31,9 +34,9 @@ MAX_MCP_TOOL_CALL_ROUNDS = 5
async def create_mcp_list_tools_events(
mcp_tools_with_litellm_proxy: list[ToolParam],
user_api_key_auth: Any,
user_api_key_auth: "UserAPIKeyAuth | None",
base_item_id: str,
pre_processed_mcp_tools: list[Any],
pre_processed_mcp_tools: list[MCPTool],
) -> list[ResponsesAPIStreamingResponse]:
"""Create MCP discovery events using pre-processed tools from the parent"""
@ -258,8 +261,8 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
base_iterator: Any, # Can be None - will be created internally
mcp_events: list[ResponsesAPIStreamingResponse],
tool_server_map: dict[str, str],
mcp_tools_with_litellm_proxy: list[Any] | None = None,
user_api_key_auth: Any = None,
mcp_tools_with_litellm_proxy: Sequence[Mapping[str, object]] | None = None,
user_api_key_auth: "UserAPIKeyAuth | None" = None,
original_request_params: dict[str, Any] | None = None,
):
# MCP setup

294
litellm/responses/streaming_iterator.py
View file

@ -9,10 +9,11 @@ from collections.abc import Mapping
from datetime import datetime
from functools import lru_cache
from types import MappingProxyType
from typing import Any, Literal
from typing import TYPE_CHECKING, Any, Literal
import httpx
from openai._streaming import SSEDecoder
from typing_extensions import TypeIs
import litellm
from litellm.constants import (
@ -30,10 +31,23 @@ from litellm.litellm_core_utils.llm_response_utils.response_metadata import (
from litellm.litellm_core_utils.thread_pool_executor import executor
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
from litellm.responses.utils import ResponseAPILoggingUtils, ResponsesAPIRequestUtils
from litellm.types.llms.openai import ResponsesAPIStreamEvents
from litellm.types.llms.openai import (
PART_UNION_TYPES,
ResponsesAPIResponse,
ResponsesAPIStreamEvents,
ResponsesAPIStreamingResponse,
)
from litellm.types.utils import CallTypes
from litellm.utils import async_post_call_success_deployment_hook
if TYPE_CHECKING:
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.responses.streaming_websocket import (
PresidioGuardrailCallback,
ResponsesBackendWebSocket,
ResponsesClientWebSocket,
)
@lru_cache(maxsize=1)
def _get_openai_response_types():
@ -42,7 +56,25 @@ def _get_openai_response_types():
return openai_types
def _log_background_task_failure(task: asyncio.Task[Any], *, task_name: str) -> None:
def _is_json_object(value: object) -> TypeIs[dict[str, object]]: # guard-ok: trivial isinstance; JSON keys are str
return isinstance(value, dict)
def _is_json_array(value: object) -> TypeIs[list[object]]: # guard-ok: trivial isinstance narrowing
return isinstance(value, list)
def _is_str_mapping(value: object) -> TypeIs[dict[str, str]]: # guard-ok: verifies every value is str
return _is_json_object(value) and all(isinstance(item, str) for item in value.values())
def _model_id_from_metadata(litellm_metadata: dict[str, object] | None) -> str | None:
model_info = litellm_metadata.get("model_info") if litellm_metadata else None
model_id = model_info.get("id") if _is_json_object(model_info) else None
return model_id if isinstance(model_id, str) else None
def _log_background_task_failure(task: asyncio.Task[object], *, task_name: str) -> None:
if task.cancelled():
return
exception = task.exception()
@ -121,9 +153,9 @@ class BaseResponsesAPIStreamingIterator:
model: str,
responses_api_provider_config: BaseResponsesAPIConfig | None,
logging_obj: LiteLLMLoggingObj,
litellm_metadata: dict[str, Any] | None = None,
litellm_metadata: dict[str, object] | None = None,
custom_llm_provider: str | None = None,
request_data: dict[str, Any] | None = None,
request_data: dict[str, object] | None = None,
call_type: str | None = None,
):
self.response = response
@ -131,7 +163,7 @@ class BaseResponsesAPIStreamingIterator:
self.logging_obj = logging_obj
self.finished = False
self.responses_api_provider_config = responses_api_provider_config
self.completed_response: Any | None = None
self.completed_response: ResponsesAPIStreamingResponse | None = None
self.start_time = getattr(logging_obj, "start_time", datetime.now())
self._failure_handled = False # Track if failure handler has been called
self._yielded_first_chunk = False
@ -145,7 +177,7 @@ class BaseResponsesAPIStreamingIterator:
# track request context for hooks
self.litellm_metadata = litellm_metadata
self.custom_llm_provider = custom_llm_provider
self.request_data: dict[str, Any] = request_data or {}
self.request_data: dict[str, object] = request_data or {}
self.call_type: str | None = call_type
# set hidden params for response headers (e.g., x-litellm-model-id)
@ -154,9 +186,8 @@ class BaseResponsesAPIStreamingIterator:
model=model or "",
optional_params=self.logging_obj.model_call_details.get("litellm_params", {}),
)
_model_info: dict = litellm_metadata.get("model_info", {}) if litellm_metadata else {}
self._hidden_params = {
"model_id": _model_info.get("id", None),
self._hidden_params: dict[str, object] = {
"model_id": _model_id_from_metadata(litellm_metadata),
"api_base": _api_base,
"custom_llm_provider": custom_llm_provider,
}
@ -176,7 +207,7 @@ class BaseResponsesAPIStreamingIterator:
llm_provider=self.custom_llm_provider or "",
)
def _process_chunk(self, chunk) -> Any | None:
def _process_chunk(self, chunk: str) -> ResponsesAPIStreamingResponse | None:
"""Process a single chunk of data from the stream"""
if not chunk:
return None
@ -227,9 +258,7 @@ class BaseResponsesAPIStreamingIterator:
_delta = getattr(openai_responses_api_chunk, "delta", None)
if isinstance(_delta, str):
self._generated_content += _delta
_stream_model_id = (
self.litellm_metadata.get("model_info", {}).get("id") if self.litellm_metadata else None
)
_stream_model_id = _model_id_from_metadata(self.litellm_metadata)
if _event_type in (
ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
@ -277,11 +306,7 @@ class BaseResponsesAPIStreamingIterator:
if item:
encrypted_content = getattr(item, "encrypted_content", None)
if encrypted_content and isinstance(encrypted_content, str):
model_id = (
self.litellm_metadata.get("model_info", {}).get("id")
if self.litellm_metadata
else None
)
model_id = _model_id_from_metadata(self.litellm_metadata)
if model_id:
wrapped_content = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id(
encrypted_content, model_id
@ -401,7 +426,7 @@ class BaseResponsesAPIStreamingIterator:
)
self._handle_failure(exception)
def _record_failed_response_usage(self, response_obj: Any | None) -> None:
def _record_failed_response_usage(self, response_obj: ResponsesAPIResponse | None) -> None:
if response_obj is None or self.logging_obj is None:
return
usage_obj = getattr(response_obj, "usage", None)
@ -451,7 +476,7 @@ class BaseResponsesAPIStreamingIterator:
is_pre_first_chunk=not self._yielded_first_chunk,
)
def _get_completed_response_object(self) -> Any | None:
def _get_completed_response_object(self) -> ResponsesAPIResponse | None:
openai_types = _get_openai_response_types()
completed_response = self.completed_response
if isinstance(completed_response, openai_types.ResponsesAPIResponse):
@ -527,7 +552,9 @@ class BaseResponsesAPIStreamingIterator:
self._completed_response_cached = True
async def _call_post_streaming_deployment_hook(self, chunk):
async def _call_post_streaming_deployment_hook(
self, chunk: ResponsesAPIStreamingResponse
) -> ResponsesAPIStreamingResponse:
"""
Allow callbacks to modify streaming chunks before returning (parity with chat).
"""
@ -564,7 +591,9 @@ class BaseResponsesAPIStreamingIterator:
except Exception:
return chunk
async def call_post_streaming_hooks_for_testing(self, chunk):
async def call_post_streaming_hooks_for_testing(
self, chunk: ResponsesAPIStreamingResponse
) -> ResponsesAPIStreamingResponse:
"""
Helper to invoke streaming deployment hooks explicitly (used in tests).
"""
@ -687,9 +716,9 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
model: str,
responses_api_provider_config: BaseResponsesAPIConfig,
logging_obj: LiteLLMLoggingObj,
litellm_metadata: dict[str, Any] | None = None,
litellm_metadata: dict[str, object] | None = None,
custom_llm_provider: str | None = None,
request_data: dict[str, Any] | None = None,
request_data: dict[str, object] | None = None,
call_type: str | None = None,
):
super().__init__(
@ -707,7 +736,7 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
def __aiter__(self):
return self
async def __anext__(self) -> Any:
async def __anext__(self) -> ResponsesAPIStreamingResponse:
try:
self._check_max_streaming_duration()
while True:
@ -769,9 +798,9 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
model: str,
responses_api_provider_config: BaseResponsesAPIConfig,
logging_obj: LiteLLMLoggingObj,
litellm_metadata: dict[str, Any] | None = None,
litellm_metadata: dict[str, object] | None = None,
custom_llm_provider: str | None = None,
request_data: dict[str, Any] | None = None,
request_data: dict[str, object] | None = None,
call_type: str | None = None,
):
super().__init__(
@ -856,9 +885,9 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
model: str,
responses_api_provider_config: BaseResponsesAPIConfig,
logging_obj: LiteLLMLoggingObj,
litellm_metadata: dict[str, Any] | None = None,
litellm_metadata: dict[str, object] | None = None,
custom_llm_provider: str | None = None,
request_data: dict[str, Any] | None = None,
request_data: dict[str, object] | None = None,
call_type: str | None = None,
):
transformed = responses_api_provider_config.transform_response_api_response(
@ -880,10 +909,10 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
def _set_events_from_response(
self,
transformed: Any,
transformed: ResponsesAPIResponse,
logging_obj: LiteLLMLoggingObj,
) -> None:
self._events = _build_synthetic_response_events(
self._events: list[ResponsesAPIStreamingResponse] = _build_synthetic_response_events(
transformed=transformed,
logging_obj=logging_obj,
chunk_size=self.CHUNK_SIZE,
@ -894,7 +923,7 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
def __aiter__(self):
return self
async def __anext__(self) -> Any:
async def __anext__(self) -> ResponsesAPIStreamingResponse:
if self._idx >= len(self._events):
raise StopAsyncIteration
evt = self._events[self._idx]
@ -908,7 +937,7 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
def __iter__(self):
return self
def __next__(self) -> Any:
def __next__(self) -> ResponsesAPIStreamingResponse:
if self._idx >= len(self._events):
raise StopIteration
evt = self._events[self._idx]
@ -923,9 +952,9 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
def __init__(
self,
response: Any,
response: ResponsesAPIResponse,
logging_obj: LiteLLMLoggingObj,
request_data: dict[str, Any] | None = None,
request_data: dict[str, object] | None = None,
call_type: str | None = None,
):
BaseResponsesAPIStreamingIterator.__init__(
@ -941,13 +970,13 @@ class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
)
self._completed_response_cache_hit = True
self._persist_completed_response_before_logging = False
self._events: list[Any] = []
self._events: list[ResponsesAPIStreamingResponse] = []
self._idx = 0
self._set_events_from_response(transformed=response, logging_obj=logging_obj)
def _set_events_from_response(
self,
transformed: Any,
transformed: ResponsesAPIResponse,
logging_obj: LiteLLMLoggingObj,
) -> None:
self._events = _build_synthetic_response_events(
@ -961,7 +990,7 @@ class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
def __aiter__(self):
return self
async def __anext__(self) -> Any:
async def __anext__(self) -> ResponsesAPIStreamingResponse:
if self._idx >= len(self._events):
raise StopAsyncIteration
evt = self._events[self._idx]
@ -975,7 +1004,7 @@ class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
def __iter__(self):
return self
def __next__(self) -> Any:
def __next__(self) -> ResponsesAPIStreamingResponse:
if self._idx >= len(self._events):
raise StopIteration
evt = self._events[self._idx]
@ -1000,8 +1029,8 @@ def _build_response_status_event(
"response.created",
"response.in_progress",
],
transformed: Any,
) -> Any:
transformed: ResponsesAPIResponse,
) -> ResponsesAPIStreamingResponse:
openai_types = _get_openai_response_types()
in_progress_response = transformed.model_copy(
deep=True,
@ -1018,10 +1047,10 @@ def _build_content_part_done_event(
output_index: int,
content_index: int,
part_payload: dict[str, Any],
) -> Any | None:
) -> ResponsesAPIStreamingResponse | None:
openai_types = _get_openai_response_types()
part_type = part_payload.get("type")
part: Any
part: PART_UNION_TYPES
if part_type == "output_text":
annotations = [
openai_types.BaseLiteLLMOpenAIResponseObject(**annotation)
@ -1057,7 +1086,7 @@ def _build_content_part_done_event(
def _add_text_like_part_events(
*,
events: list[Any],
events: list[ResponsesAPIStreamingResponse],
item_id: str,
output_index: int,
content_index: int,
@ -1123,13 +1152,13 @@ def _add_text_like_part_events(
def _build_synthetic_response_events(
*,
transformed: Any,
transformed: ResponsesAPIResponse,
logging_obj: LiteLLMLoggingObj,
chunk_size: int,
) -> list[Any]:
) -> list[ResponsesAPIStreamingResponse]:
openai_types = _get_openai_response_types()
if litellm.include_cost_in_streaming_usage and logging_obj is not None:
usage_obj: Any | None = getattr(transformed, "usage", None)
usage_obj = transformed.usage if hasattr(transformed, "usage") else None
if usage_obj is not None:
try:
cost: float | None = logging_obj._response_cost_calculator(result=transformed)
@ -1138,7 +1167,7 @@ def _build_synthetic_response_events(
except Exception:
pass
events: list[Any] = [
events: list[ResponsesAPIStreamingResponse] = [
_build_response_status_event(openai_types.ResponsesAPIStreamEvents.RESPONSE_CREATED, transformed),
_build_response_status_event(openai_types.ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS, transformed),
]
@ -1292,34 +1321,34 @@ class ResponsesWebSocketStreaming:
def __init__(
self,
websocket: Any,
backend_ws: Any,
websocket: ResponsesClientWebSocket,
backend_ws: ResponsesBackendWebSocket,
logging_obj: LiteLLMLoggingObj,
user_api_key_dict: Any | None = None,
request_data: dict | None = None,
user_api_key_dict: UserAPIKeyAuth | None = None,
request_data: dict[str, object] | None = None,
first_message: str | None = None,
guardrail_callbacks: list[Any] | None = None,
output_guardrail_callbacks: list[Any] | None = None,
output_guardrail_callbacks: list[PresidioGuardrailCallback] | None = None,
authorized_model: str | None = None,
):
self.websocket = websocket
self.backend_ws = backend_ws
self.logging_obj = logging_obj
self.user_api_key_dict = user_api_key_dict
self.request_data: dict = request_data or {}
self.messages: list[dict] = []
self.input_messages: list[dict[str, str]] = []
self.request_data: dict[str, object] = request_data or {}
self.messages: list[dict[str, object]] = []
self.input_messages: list[dict[str, object]] = []
self.first_message = first_message
self.guardrail_callbacks: list[Any] = guardrail_callbacks or []
self.output_guardrail_callbacks: list[Any] = output_guardrail_callbacks or []
self.output_guardrail_callbacks: list[PresidioGuardrailCallback] = output_guardrail_callbacks or []
# Model name authorized at connection time; enforced on every
# response.create frame to prevent deployment-substitution attacks.
self.authorized_model: str | None = authorized_model
def _should_store_event(self, event_obj: dict) -> bool:
def _should_store_event(self, event_obj: dict[str, object]) -> bool:
return event_obj.get("type") in RESPONSES_WS_LOGGED_EVENT_TYPES
def _store_event(self, event: Any) -> None:
def _store_event(self, event: str | bytes | dict[str, object]) -> None:
if isinstance(event, bytes):
event = event.decode("utf-8")
if isinstance(event, str):
@ -1333,12 +1362,12 @@ class ResponsesWebSocketStreaming:
if self._should_store_event(event_obj):
self.messages.append(event_obj)
def _collect_input_from_client_event(self, message: Any) -> None:
def _collect_input_from_client_event(self, message: object) -> None:
"""Extract user input content from response.create for logging."""
try:
if isinstance(message, str):
msg_obj = json.loads(message)
elif isinstance(message, dict):
elif _is_json_object(message):
msg_obj = message
else:
return
@ -1351,24 +1380,24 @@ class ResponsesWebSocketStreaming:
self.input_messages.append({"role": "user", "content": input_items})
return
if isinstance(input_items, list):
if _is_json_array(input_items):
for item in input_items:
if not isinstance(item, dict):
if not _is_json_object(item):
continue
if item.get("type") == "message" and item.get("role") == "user":
content = item.get("content", [])
if isinstance(content, str):
self.input_messages.append({"role": "user", "content": content})
elif isinstance(content, list):
elif _is_json_array(content):
for c in content:
if isinstance(c, dict) and c.get("type") == "input_text":
if _is_json_object(c) and c.get("type") == "input_text":
text = c.get("text", "")
if text:
self.input_messages.append({"role": "user", "content": text})
except (json.JSONDecodeError, AttributeError, TypeError):
pass
def _store_input(self, message: Any) -> None:
def _store_input(self, message: object) -> None:
self._collect_input_from_client_event(message)
if self.logging_obj:
self.logging_obj.pre_call(input=message, api_key="")
@ -1429,7 +1458,7 @@ class ResponsesWebSocketStreaming:
finally:
await self._log_messages()
def _enforce_authorized_model(self, msg_obj: dict) -> bool:
def _enforce_authorized_model(self, msg_obj: dict[str, object]) -> bool:
"""
Overwrite any ``model`` field in a ``response.create`` frame with the
connection-authorized model to prevent deployment-substitution attacks.
@ -1444,7 +1473,7 @@ class ResponsesWebSocketStreaming:
return False
modified = False
nested = msg_obj.get("response")
if isinstance(nested, dict):
if _is_json_object(nested):
if nested.get("model") != self.authorized_model:
nested["model"] = self.authorized_model
modified = True
@ -1495,8 +1524,9 @@ class ResponsesWebSocketStreaming:
# nested: {"type": "response.create", "response": {"input": ..., "instructions": ...}}
# Mask "input" and "instructions" in both shapes so PII is never
# forwarded unmasked regardless of where the client places it.
nested_response = msg_obj.get("response") if isinstance(msg_obj.get("response"), dict) else None
text_containers: list[tuple[dict, str]] = []
nested_candidate = msg_obj.get("response")
nested_response = nested_candidate if _is_json_object(nested_candidate) else None
text_containers: list[tuple[dict[str, object], str]] = []
for container in (msg_obj, nested_response):
if container is None:
continue
@ -1517,9 +1547,9 @@ class ResponsesWebSocketStreaming:
)
modified = True
elif isinstance(field_value, list):
elif _is_json_array(field_value):
for item in field_value:
if not isinstance(item, dict):
if not _is_json_object(item):
continue
for item_field in ("content", "output"):
value = item.get(item_field)
@ -1531,15 +1561,16 @@ class ResponsesWebSocketStreaming:
request_data=self.request_data,
)
modified = True
elif isinstance(value, list):
elif _is_json_array(value):
for block in value:
if (
isinstance(block, dict)
and block.get("type") in RESPONSES_WS_MASKABLE_TEXT_BLOCK_TYPES
and isinstance(block.get("text"), str)
if not _is_json_object(block):
continue
block_text = block.get("text")
if block.get("type") in RESPONSES_WS_MASKABLE_TEXT_BLOCK_TYPES and isinstance(
block_text, str
):
block["text"] = await cb.check_pii(
text=block["text"],
text=block_text,
output_parse_pii=True,
presidio_config=presidio_config,
request_data=self.request_data,
@ -1590,7 +1621,9 @@ class ResponsesWebSocketStreaming:
if not self.guardrail_callbacks:
return response_str
pii_tokens: dict[str, str] = (self.request_data.get("metadata") or {}).get("pii_tokens", {})
metadata = self.request_data.get("metadata")
raw_pii_tokens = metadata.get("pii_tokens") if _is_json_object(metadata) else None
pii_tokens: dict[str, str] = raw_pii_tokens if _is_str_mapping(raw_pii_tokens) else {}
if not pii_tokens:
return response_str
@ -1604,17 +1637,18 @@ class ResponsesWebSocketStreaming:
if event_type == "response.completed":
modified = False
response_obj = evt_obj.get("response") or {}
if not isinstance(response_obj, dict):
response_obj = evt_obj.get("response")
if not _is_json_object(response_obj):
return response_str
for output_item in response_obj.get("output") or []:
if not isinstance(output_item, dict):
output_items = response_obj.get("output")
for output_item in output_items if _is_json_array(output_items) else []:
if not _is_json_object(output_item):
continue
content = output_item.get("content") or []
if not isinstance(content, list):
content = output_item.get("content")
if not _is_json_array(content):
continue
for content_block in content:
if not isinstance(content_block, dict):
if not _is_json_object(content_block):
continue
text = content_block.get("text")
if isinstance(text, str):
@ -1660,11 +1694,12 @@ class ResponsesWebSocketStreaming:
modified = False
for cb in self.output_guardrail_callbacks:
presidio_config = cb.get_presidio_settings_from_request_data(self.request_data)
response_obj = evt_obj.get("response") or {}
if not isinstance(response_obj, dict):
response_obj = evt_obj.get("response")
if not _is_json_object(response_obj):
continue
for output_item in response_obj.get("output") or []:
if not isinstance(output_item, dict):
output_items = response_obj.get("output")
for output_item in output_items if _is_json_array(output_items) else []:
if not _is_json_object(output_item):
continue
arguments = output_item.get("arguments")
if isinstance(arguments, str):
@ -1677,10 +1712,10 @@ class ResponsesWebSocketStreaming:
if masked_args != arguments:
output_item["arguments"] = masked_args
modified = True
summary = output_item.get("summary") or []
if isinstance(summary, list):
summary = output_item.get("summary")
if _is_json_array(summary):
for summary_block in summary:
if not isinstance(summary_block, dict):
if not _is_json_object(summary_block):
continue
summary_text = summary_block.get("text")
if isinstance(summary_text, str):
@ -1693,11 +1728,11 @@ class ResponsesWebSocketStreaming:
if masked_summary != summary_text:
summary_block["text"] = masked_summary
modified = True
content = output_item.get("content") or []
if not isinstance(content, list):
content = output_item.get("content")
if not _is_json_array(content):
continue
for content_block in content:
if not isinstance(content_block, dict):
if not _is_json_object(content_block):
continue
text = content_block.get("text")
if isinstance(text, str):
@ -1756,12 +1791,12 @@ class ResponsesWebSocketStreaming:
# Managed WebSocket mode (HTTP-backed, provider-agnostic)
# ---------------------------------------------------------------------------
_RESPONSE_CREATE_PARAMS: frozenset = (
_RESPONSE_CREATE_PARAMS: frozenset[str] = (
_get_openai_response_types().ResponsesAPIRequestParams.__required_keys__
| _get_openai_response_types().ResponsesAPIRequestParams.__optional_keys__
)
_MANAGED_WS_SKIP_KWARGS: frozenset = frozenset(
_MANAGED_WS_SKIP_KWARGS: frozenset[str] = frozenset(
{
"litellm_logging_obj",
"litellm_call_id",
@ -1793,17 +1828,17 @@ class ManagedResponsesWebSocketHandler:
def __init__(
self,
websocket: Any,
websocket: ResponsesClientWebSocket,
model: str,
logging_obj: LiteLLMLoggingObj,
user_api_key_dict: Any | None = None,
user_api_key_dict: UserAPIKeyAuth | None = None,
litellm_metadata: dict[str, Any] | None = None,
api_key: str | None = None,
api_base: str | None = None,
timeout: float | None = None,
custom_llm_provider: str | None = None,
first_message: str | None = None,
**kwargs: Any,
**kwargs: object,
) -> None:
self.websocket = websocket
self.model = model
@ -1820,12 +1855,12 @@ class ManagedResponsesWebSocketHandler:
self._connection_provider = self._resolve_provider(model) or custom_llm_provider
self.first_message = first_message
# Carry through safe pass-through kwargs (e.g. extra_headers)
self.extra_kwargs: dict[str, Any] = {k: v for k, v in kwargs.items() if k not in _MANAGED_WS_SKIP_KWARGS}
self.extra_kwargs: dict[str, object] = {k: v for k, v in kwargs.items() if k not in _MANAGED_WS_SKIP_KWARGS}
# In-memory session history: response_id → full accumulated message list.
# Keyed by the DECODED (pre-encoding) response ID from response.completed.
# This avoids the async DB-write race condition where spend logs haven't
# been committed yet when the next response.create arrives.
self._session_history: dict[str, list[dict[str, Any]]] = {}
self._session_history: dict[str, list[dict[str, object]]] = {}
# ------------------------------------------------------------------
# Internal helpers
@ -1854,7 +1889,7 @@ class ManagedResponsesWebSocketHandler:
except Exception:
pass
def _get_history_messages(self, previous_response_id: str) -> list[dict[str, Any]]:
def _get_history_messages(self, previous_response_id: str) -> list[dict[str, object]]:
"""
Return accumulated message history for *previous_response_id*.
@ -1865,7 +1900,7 @@ class ManagedResponsesWebSocketHandler:
raw_id = decoded.get("response_id", previous_response_id)
return list(self._session_history.get(raw_id, []))
def _store_history(self, response_id: str, messages: list[dict[str, Any]]) -> None:
def _store_history(self, response_id: str, messages: list[dict[str, object]]) -> None:
"""
Store the complete accumulated message history for *response_id*.
@ -1875,13 +1910,14 @@ class ManagedResponsesWebSocketHandler:
self._session_history[response_id] = messages
@staticmethod
def _extract_response_id(completed_event: dict[str, Any]) -> str | None:
def _extract_response_id(completed_event: dict[str, object]) -> str | None:
"""
Pull the raw (decoded) response ID out of a ``response.completed`` event.
Returns *None* if the event doesn't contain a usable ID.
"""
resp_obj = completed_event.get("response", {})
encoded_id: str | None = resp_obj.get("id") if isinstance(resp_obj, dict) else None
raw_id = resp_obj.get("id") if _is_json_object(resp_obj) else None
encoded_id: str | None = raw_id if isinstance(raw_id, str) else None
if not encoded_id:
return None
decoded = ResponsesAPIRequestUtils._decode_responses_api_response_id(encoded_id)
@ -1890,7 +1926,7 @@ class ManagedResponsesWebSocketHandler:
@staticmethod
def _extract_output_messages(
completed_event: dict[str, Any],
) -> list[dict[str, Any]]:
) -> list[dict[str, object]]:
"""
Convert the output items in a ``response.completed`` event into
Responses API message dicts suitable for the next turn's ``input``.
@ -1898,7 +1934,7 @@ class ManagedResponsesWebSocketHandler:
resp_obj = completed_event.get("response", {})
if not isinstance(resp_obj, dict):
return []
messages: list[dict[str, Any]] = []
messages: list[dict[str, object]] = []
for item in resp_obj.get("output", []) or []:
if not isinstance(item, dict):
continue
@ -1925,7 +1961,7 @@ class ManagedResponsesWebSocketHandler:
return messages
@staticmethod
def _input_to_messages(input_val: Any) -> list[dict[str, Any]]:
def _input_to_messages(input_val: object) -> list[dict[str, object]]:
"""
Normalise the ``input`` field of a ``response.create`` event to a list
of Responses API message dicts.
@ -1938,15 +1974,15 @@ class ManagedResponsesWebSocketHandler:
"content": [{"type": "input_text", "text": input_val}],
}
]
if isinstance(input_val, list):
return [item for item in input_val if isinstance(item, dict)]
if _is_json_array(input_val):
return [item for item in input_val if _is_json_object(item)]
return []
# ------------------------------------------------------------------
# _process_response_create sub-methods
# ------------------------------------------------------------------
async def _parse_message(self, raw_message: str) -> dict[str, Any] | None:
async def _parse_message(self, raw_message: str) -> dict[str, object] | None:
"""Parse raw WS text; return the message dict or None (JSON error / ignored type)."""
try:
msg_obj = json.loads(raw_message)
@ -1959,10 +1995,10 @@ class ManagedResponsesWebSocketHandler:
return msg_obj
@staticmethod
def _is_warmup_frame(msg_obj: dict[str, Any]) -> bool:
def _is_warmup_frame(msg_obj: dict[str, object]) -> bool:
"""Return True for a response.create whose generate flag is false."""
nested = msg_obj.get("response")
source = nested if isinstance(nested, dict) and nested else msg_obj
source = nested if _is_json_object(nested) and nested else msg_obj
return source.get("generate") is False
@staticmethod
@ -1975,13 +2011,13 @@ class ManagedResponsesWebSocketHandler:
return str(raw_id).startswith(_WARMUP_RESPONSE_ID_PREFIX)
@staticmethod
def _warmup_source_params(msg_obj: dict[str, Any]) -> dict[str, Any]:
def _warmup_source_params(msg_obj: dict[str, object]) -> dict[str, object]:
nested = msg_obj.get("response")
if isinstance(nested, dict) and nested:
if _is_json_object(nested) and nested:
return nested
return {k: v for k, v in msg_obj.items() if k != "type"}
def _build_warmup_response(self, msg_obj: dict[str, Any]) -> dict[str, Any]:
def _build_warmup_response(self, msg_obj: dict[str, object]) -> dict[str, object]:
"""Build a minimal completed Responses API object for a warmup ack."""
source = self._warmup_source_params(msg_obj)
wire_model = source.get("model") or self.model_group or self.model
@ -1999,7 +2035,7 @@ class ManagedResponsesWebSocketHandler:
},
}
async def _send_warmup_ack(self, msg_obj: dict[str, Any]) -> None:
async def _send_warmup_ack(self, msg_obj: dict[str, object]) -> None:
"""
Acknowledge a generate=false prewarm without calling the provider.
@ -2022,7 +2058,7 @@ class ManagedResponsesWebSocketHandler:
await self.websocket.send_text(serialized)
@staticmethod
def _build_base_call_kwargs(msg_obj: dict[str, Any]) -> dict[str, Any]:
def _build_base_call_kwargs(msg_obj: dict[str, object]) -> dict[str, Any]:
"""
Extract Responses API params from the event, handling both wire formats:
Nested: {"type": "response.create", "response": {"input": [...], ...}}
@ -2030,7 +2066,7 @@ class ManagedResponsesWebSocketHandler:
"""
nested = msg_obj.get("response")
response_params: dict[str, Any] = (
nested if isinstance(nested, dict) and nested else {k: v for k, v in msg_obj.items() if k != "type"}
nested if _is_json_object(nested) and nested else {k: v for k, v in msg_obj.items() if k != "type"}
)
return {
param: response_params[param]
@ -2042,8 +2078,8 @@ class ManagedResponsesWebSocketHandler:
self,
call_kwargs: dict[str, Any],
previous_response_id: str | None,
current_messages: list[dict[str, Any]],
prior_history: list[dict[str, Any]],
current_messages: list[dict[str, object]],
prior_history: list[dict[str, object]],
) -> None:
"""Prepend in-memory turn history, or fall back to DB-based reconstruction."""
if not previous_response_id:
@ -2131,7 +2167,7 @@ class ManagedResponsesWebSocketHandler:
call_kwargs.setdefault("litellm_params", {})
call_kwargs["litellm_params"]["proxy_server_request"] = proxy_server_request
async def _stream_and_forward(self, model: str, call_kwargs: dict[str, Any]) -> dict[str, Any] | None:
async def _stream_and_forward(self, model: str, call_kwargs: dict[str, Any]) -> dict[str, object] | None:
"""
Stream ``litellm.aresponses`` and forward every chunk over the WebSocket.
@ -2139,7 +2175,7 @@ class ManagedResponsesWebSocketHandler:
directly (before serialization) to avoid a redundant JSON round-trip on
every chunk. Returns the completed event dict, or ``None``.
"""
completed_event: dict[str, Any] | None = None
completed_event: dict[str, object] | None = None
stream_response = await litellm.aresponses(model=model, **call_kwargs)
async for chunk in stream_response: # type: ignore[union-attr]
if chunk is None:
@ -2163,9 +2199,9 @@ class ManagedResponsesWebSocketHandler:
def _save_turn_history(
self,
completed_event: dict[str, Any] | None,
prior_history: list[dict[str, Any]],
current_messages: list[dict[str, Any]],
completed_event: dict[str, object] | None,
prior_history: list[dict[str, object]],
current_messages: list[dict[str, object]],
) -> None:
"""Store this turn in in-memory history for future previous_response_id lookups."""
if completed_event is None:

21
litellm/types/google_genai/adapters.py Normal file
View file

@ -0,0 +1,21 @@
from typing_extensions import TypedDict
from litellm.types.llms.openai import (
AllMessageValues,
ChatCompletionToolChoiceStringValues,
ChatCompletionToolParam,
)
class GenerateContentCompletionKwargs(TypedDict, total=False):
model: str
messages: list[AllMessageValues]
temperature: float
max_tokens: int
top_p: float
stop: str | list[str]
tools: list[ChatCompletionToolParam]
tool_choice: ChatCompletionToolChoiceStringValues
stream: bool
metadata: dict[str, object]
extra_headers: dict[str, str] | None

123
litellm/types/passthrough_endpoints/managed_id_rewriter.py Normal file
View file

@ -0,0 +1,123 @@
"""
Typed surfaces for the passthrough managed-ID rewriter.
Prisma's generated client is untyped at the ``litellm`` boundary, so the row
shapes, table actions, and query fragments the rewriter touches are declared
here as protocols instead of leaking ``Any`` through every call site.
"""
from __future__ import annotations
from collections.abc import Mapping, Sequence
from datetime import datetime
from typing import (
TYPE_CHECKING,
Literal,
Protocol,
TypeAlias,
TypedDict,
TypeVar,
runtime_checkable,
)
from pydantic import JsonValue
if TYPE_CHECKING:
from litellm.models.managed_files import LiteLLM_ManagedFileTable
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.llms.openai import OpenAIFileObject
SortOrder: TypeAlias = Literal["asc", "desc"]
ResourceKind: TypeAlias = Literal["files", "batches"]
PrismaWhereValue: TypeAlias = (
"str | int | bool | datetime | None | Mapping[str, PrismaWhereValue] | Sequence[PrismaWhereValue]"
)
PrismaWhere: TypeAlias = "Mapping[str, PrismaWhereValue]"
PrismaOrder: TypeAlias = "Mapping[str, SortOrder]"
ManagedRowData: TypeAlias = "Mapping[str, str | None]"
class ManagedResourceRow(Protocol):
"""Columns shared by ``LiteLLM_ManagedFileTable`` and ``LiteLLM_ManagedObjectTable`` rows."""
created_by: str | None
team_id: str | None
created_at: datetime | None
file_object: JsonValue
class ManagedFileRow(ManagedResourceRow, Protocol):
unified_file_id: str
class ManagedObjectRow(ManagedResourceRow, Protocol):
unified_object_id: str
RowT = TypeVar("RowT", bound=ManagedResourceRow)
class ManagedTable(Protocol[RowT]):
"""The Prisma table actions the rewriter reads rows through."""
async def find_first(self, *, where: PrismaWhere) -> RowT | None: ...
async def find_many(
self,
*,
where: PrismaWhere,
order: PrismaOrder | Sequence[PrismaOrder] | None = None,
take: int | None = None,
) -> list[RowT]: ...
class ManagedFileTable(ManagedTable[ManagedFileRow], Protocol): ...
class ManagedObjectTable(ManagedTable[ManagedObjectRow], Protocol):
async def update(self, *, where: PrismaWhere, data: ManagedRowData) -> ManagedObjectRow | None: ...
async def upsert(self, *, where: PrismaWhere, data: Mapping[str, ManagedRowData]) -> ManagedObjectRow: ...
@runtime_checkable
class ManagedFileIdReader(Protocol):
"""Row lookup on the enterprise managed-files hook.
The proxy hook registry is untyped and hands back a bare ``CustomLogger``,
so this protocol is an ``isinstance`` target: the rewriter checks the method
is really there before calling it. It is kept separate from
``ManagedFileIdWriter`` so a hook implementing only one of the two is
narrowed on exactly the capability about to be used.
"""
async def get_unified_file_id(
self,
file_id: str,
litellm_parent_otel_span: object = None,
) -> LiteLLM_ManagedFileTable | None: ...
@runtime_checkable
class ManagedFileIdWriter(Protocol):
"""Row persistence on the enterprise managed-files hook."""
async 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,
) -> None: ...
class ManagedListResponse(TypedDict):
"""OpenAI-style paginated list body served from the managed-resource tables."""
object: Literal["list"]
data: list[dict[str, JsonValue]]
first_id: str | None
last_id: str | None
has_more: bool

41
litellm/types/responses/streaming_websocket.py Normal file
View file

@ -0,0 +1,41 @@
from __future__ import annotations
from typing import Protocol
from litellm.types.guardrails import PresidioPerRequestConfig
class ResponsesClientWebSocket(Protocol):
"""Client-facing websocket surface used by the Responses API websocket handlers."""
async def send_text(self, data: str) -> None: ...
async def receive_text(self) -> str: ...
class ResponsesBackendWebSocket(Protocol):
"""Upstream provider websocket surface used when proxying a native Responses API socket."""
async def recv(self, decode: bool = ...) -> str | bytes: ...
async def send(self, message: str) -> None: ...
async def close(self) -> None: ...
class PresidioGuardrailCallback(Protocol):
"""
Duck-typed PII guardrail surface consumed by the Responses API websocket handlers.
Declared structurally so the SDK does not import from the proxy guardrail package.
"""
def get_presidio_settings_from_request_data(self, data: dict[str, object]) -> PresidioPerRequestConfig | None: ...
async def check_pii(
self,
text: str,
output_parse_pii: bool,
presidio_config: PresidioPerRequestConfig | None,
request_data: dict[str, object],
) -> str: ...

10
ruff-strict-budget.json
View file

@ -1,18 +1,18 @@
{
"ANN001": {
"limit": 3097
"limit": 3094
},
"ANN002": {
"limit": 69
},
"ANN003": {
"limit": 831
"limit": 829
},
"ANN201": {
"limit": 2137
"limit": 2136
},
"ANN202": {
"limit": 941
"limit": 940
},
"ANN204": {
"limit": 724
@ -24,7 +24,7 @@
"limit": 130
},
"ANN401": {
"limit": 1848
"limit": 1762
},
"ASYNC230": {
"limit": 14

43
tests/pass_through_unit_tests/test_passthrough_managed_ids.py
View file

@ -1257,6 +1257,23 @@ class TestRewriteBodyIds:
assert result["files"][0] == "file-nested" # type: ignore[index]
assert result["files"][1] == "raw-string" # type: ignore[index]
@pytest.mark.asyncio
async def test_top_level_list_body_resolved(self):
"""A request body that is a JSON array (not an object) is still walked,
so managed IDs inside it are resolved instead of raising."""
mid = encode("openai", "u", "file-top-level")
hook = _managed_files_hook()
file_row = MagicMock()
file_row.created_by = "user-1"
file_row.team_id = "team-1"
hook.get_unified_file_id = AsyncMock(return_value=file_row)
body = [{"input_file_id": mid}, "raw-string"]
result = await rewrite_body_ids(body, "openai", _user(), None, hook)
assert result is not body
assert result == [{"input_file_id": "file-top-level"}, "raw-string"]
@pytest.mark.asyncio
async def test_forged_managed_id_raises_404(self):
"""An unknown managed ID in the body raises 404 (not passed to upstream)."""
@ -1853,6 +1870,32 @@ class TestListPassthroughIdsFromDb:
assert result["data"] == []
assert result["has_more"] is False
@pytest.mark.asyncio
async def test_list_missing_managed_table_returns_empty_not_error(self):
"""A generated prisma client whose db has no managed tables must fail
closed with an empty list. Opening the table raises AttributeError, and
letting it escape turns an empty 200 into a 500 at the passthrough
endpoint."""
class _DbWithoutManagedTables:
pass
pc = MagicMock()
pc.db = _DbWithoutManagedTables()
for route in ("/openai/v1/files", "/openai/v1/batches"):
result = await list_passthrough_ids_from_db(
provider="openai",
route=route,
user_api_key_dict=_admin_user(),
prisma_client=pc,
)
assert result is not None
assert result["object"] == "list"
assert result["data"] == []
assert result["has_more"] is False
@pytest.mark.asyncio
async def test_list_returns_empty_for_caller_without_identity(self):
"""Caller with neither user_id nor team_id should get an empty list."""

4
type-discipline-budget.json
View file

@ -3,7 +3,7 @@
"limit": 23349
},
"LIT002": {
"limit": 27252
"limit": 27242
},
"LIT003": {
"limit": 292
@ -15,7 +15,7 @@
"limit": 0
},
"LIT006": {
"limit": 1105
"limit": 1096
},
"LIT007": {
"limit": 0