scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 152

chore: merge litellm_internal_staging into typing cleanup branch

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Devin AI 2026-08-10 16:06:44 +00:00
commit 901e9e0022
23 changed files with 865 additions and 616 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

28
basedpyright-code-budget.json
View file

@ -1,18 +1,18 @@
{
"reportAny": {
"limit": 26575
"limit": 26391
},
"reportArgumentType": {
"limit": 2626
"limit": 2614
},
"reportAssignmentType": {
"limit": 329
"limit": 327
},
"reportAttributeAccessIssue": {
"limit": 514
},
"reportCallIssue": {
"limit": 116
"limit": 114
},
"reportConstantRedefinition": {
"limit": 40
@ -24,7 +24,7 @@
"limit": 19
},
"reportExplicitAny": {
"limit": 8704
"limit": 8319
},
"reportFunctionMemberAccess": {
"limit": 7
@ -54,10 +54,10 @@
"limit": 0
},
"reportMissingParameterType": {
"limit": 5835
"limit": 5825
},
"reportMissingTypeArgument": {
"limit": 15746
"limit": 15695
},
"reportMissingTypeStubs": {
"limit": 40
@ -90,7 +90,7 @@
"limit": 8
},
"reportReturnType": {
"limit": 217
"limit": 213
},
"reportTypedDictNotRequiredAccess": {
"limit": 26
@ -99,22 +99,22 @@
"limit": 0
},
"reportUnknownArgumentType": {
"limit": 45026
"limit": 45004
},
"reportUnknownLambdaType": {
"limit": 113
},
"reportUnknownMemberType": {
"limit": 39703
"limit": 39649
},
"reportUnknownParameterType": {
"limit": 20203
"limit": 20132
},
"reportUnknownVariableType": {
"limit": 31170
"limit": 31156
},
"reportUnnecessaryCast": {
"limit": 122
"limit": 118
},
"reportUnnecessaryComparison": {
"limit": 701
@ -123,7 +123,7 @@
"limit": 5
},
"reportUnnecessaryIsInstance": {
"limit": 862
"limit": 857
},
"reportUntypedBaseClass": {
"limit": 0

150
litellm/evals/main.py
View file

@ -7,7 +7,7 @@ import asyncio
import contextvars
from collections.abc import Coroutine
from functools import partial
from typing import Any, Final
from typing import Final
import httpx
@ -21,8 +21,10 @@ from litellm.types.llms.openai_evals import (
CancelRunResponse,
CreateEvalRequest,
CreateRunRequest,
DataSourceConfig,
DeleteEvalResponse,
Eval,
GraderConfig,
ListEvalsParams,
ListEvalsResponse,
ListRunsParams,
@ -41,13 +43,13 @@ DEFAULT_OPENAI_API_BASE: Final = "https://api.openai.com"
@client
async def acreate_eval(
data_source_config: dict[str, Any],
testing_criteria: list[dict[str, Any]],
data_source_config: DataSourceConfig,
testing_criteria: list[GraderConfig],
name: str | None = None,
metadata: dict[str, Any] | None = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
metadata: dict[str, object] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
extra_body: dict[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
@ -110,17 +112,17 @@ async def acreate_eval(
@client
def create_eval(
data_source_config: dict[str, Any],
testing_criteria: list[dict[str, Any]],
data_source_config: DataSourceConfig,
testing_criteria: list[GraderConfig],
name: str | None = None,
metadata: dict[str, Any] | None = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
metadata: dict[str, object] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
extra_body: dict[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> Eval | Coroutine[Any, Any, Eval]:
) -> Eval | Coroutine[object, object, Eval]:
"""
Create a new evaluation
@ -231,8 +233,8 @@ async def alist_evals(
before: str | None = None,
order: str | None = None,
order_by: str | None = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
@ -300,12 +302,12 @@ def list_evals(
before: str | None = None,
order: str | None = None,
order_by: str | None = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> ListEvalsResponse | Coroutine[Any, Any, ListEvalsResponse]:
) -> ListEvalsResponse | Coroutine[object, object, ListEvalsResponse]:
"""
List all evaluations
@ -413,8 +415,8 @@ def list_evals(
@client
async def aget_eval(
eval_id: str,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
@ -470,12 +472,12 @@ async def aget_eval(
@client
def get_eval(
eval_id: str,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> Eval | Coroutine[Any, Any, Eval]:
) -> Eval | Coroutine[object, object, Eval]:
"""
Get an evaluation by ID
@ -564,10 +566,10 @@ def get_eval(
async def aupdate_eval(
eval_id: str,
name: str | None = None,
metadata: dict[str, Any] | None = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
metadata: dict[str, object] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
extra_body: dict[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
@ -630,14 +632,14 @@ async def aupdate_eval(
def update_eval(
eval_id: str,
name: str | None = None,
metadata: dict[str, Any] | None = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
metadata: dict[str, object] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
extra_body: dict[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> Eval | Coroutine[Any, Any, Eval]:
) -> Eval | Coroutine[object, object, Eval]:
"""
Update an evaluation
@ -783,8 +785,8 @@ def update_eval(
@client
async def adelete_eval(
eval_id: str,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
@ -840,12 +842,12 @@ async def adelete_eval(
@client
def delete_eval(
eval_id: str,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> DeleteEvalResponse | Coroutine[Any, Any, DeleteEvalResponse]:
) -> DeleteEvalResponse | Coroutine[object, object, DeleteEvalResponse]:
"""
Delete an evaluation
@ -933,8 +935,8 @@ def delete_eval(
@client
async def acancel_eval(
eval_id: str,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
@ -990,12 +992,12 @@ async def acancel_eval(
@client
def cancel_eval(
eval_id: str,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> CancelEvalResponse | Coroutine[Any, Any, CancelEvalResponse]:
) -> CancelEvalResponse | Coroutine[object, object, CancelEvalResponse]:
"""
Cancel a running evaluation
@ -1092,12 +1094,12 @@ def cancel_eval(
@client
async def acreate_run(
eval_id: str,
data_source: dict[str, Any],
data_source: dict[str, object],
name: str | None = None,
metadata: dict[str, Any] | None = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
metadata: dict[str, object] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
extra_body: dict[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
@ -1161,16 +1163,16 @@ async def acreate_run(
@client
def create_run(
eval_id: str,
data_source: dict[str, Any],
data_source: dict[str, object],
name: str | None = None,
metadata: dict[str, Any] | None = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
metadata: dict[str, object] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
extra_body: dict[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> Run | Coroutine[Any, Any, Run]:
) -> Run | Coroutine[object, object, Run]:
"""
Create a new run for an evaluation
@ -1280,8 +1282,8 @@ async def alist_runs(
after: str | None = None,
before: str | None = None,
order: str | None = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
@ -1349,12 +1351,12 @@ def list_runs(
after: str | None = None,
before: str | None = None,
order: str | None = None,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> ListRunsResponse | Coroutine[Any, Any, ListRunsResponse]:
) -> ListRunsResponse | Coroutine[object, object, ListRunsResponse]:
"""
List all runs for an evaluation
@ -1462,8 +1464,8 @@ def list_runs(
async def aget_run(
eval_id: str,
run_id: str,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
@ -1522,12 +1524,12 @@ async def aget_run(
def get_run(
eval_id: str,
run_id: str,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> Run | Coroutine[Any, Any, Run]:
) -> Run | Coroutine[object, object, Run]:
"""
Get a specific run
@ -1618,8 +1620,8 @@ def get_run(
async def acancel_run(
eval_id: str,
run_id: str,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
@ -1678,12 +1680,12 @@ async def acancel_run(
def cancel_run(
eval_id: str,
run_id: str,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> CancelRunResponse | Coroutine[Any, Any, CancelRunResponse]:
) -> CancelRunResponse | Coroutine[object, object, CancelRunResponse]:
"""
Cancel a running run
@ -1783,8 +1785,8 @@ def cancel_run(
async def adelete_run(
eval_id: str,
run_id: str,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
@ -1843,12 +1845,12 @@ async def adelete_run(
def delete_run(
eval_id: str,
run_id: str,
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
**kwargs,
) -> RunDeleteResponse | Coroutine[Any, Any, RunDeleteResponse]:
) -> RunDeleteResponse | Coroutine[object, object, RunDeleteResponse]:
"""
Delete a run

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

@ -10,10 +10,7 @@ import asyncio
import math
import uuid
from collections.abc import AsyncIterator, Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, TypeVar, cast
from pydantic import TypeAdapter, ValidationError
from typing_extensions import TypeIs
from typing import TYPE_CHECKING, Any, Final, TypedDict, cast
import litellm
from litellm._logging import verbose_logger
@ -75,42 +72,17 @@ WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY: Final = "_websearch_interception_emit_native_b
# ``web_search_tool_result`` blocks to inject into the final response.
WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY: Final = "websearch_native_blocks"
_ResponseT = TypeVar("_ResponseT")
_CONTENT_ATTR: Final = "content"
_OBJECT_ITEMS_ADAPTER: Final[TypeAdapter[tuple[object, ...]]] = TypeAdapter(tuple[object, ...])
_NATIVE_BLOCKS_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, object], ...]]] = TypeAdapter(
tuple[Mapping[str, object], ...]
)
_WEBSEARCH_CONFIG_ADAPTER: Final[TypeAdapter[WebSearchInterceptionConfig]] = TypeAdapter(WebSearchInterceptionConfig)
class _PlanMetadataView(TypedDict):
websearch_native_blocks: Sequence[Mapping[str, object]] | None
def _is_json_object(value: object) -> TypeIs[dict[str, object]]: # guard-ok: trivial isinstance; JSON keys are str
return isinstance(value, dict)
class _AgenticLoopParamsView(TypedDict):
agentic_loop_params: AgenticLoopParams
def _parse_object_items(value: object) -> tuple[object, ...]:
try:
return _OBJECT_ITEMS_ADAPTER.validate_python(value)
except ValidationError:
return ()
def _parse_native_blocks(value: object) -> tuple[Mapping[str, object], ...]:
try:
return _NATIVE_BLOCKS_ADAPTER.validate_python(value)
except ValidationError:
return ()
def _parse_websearch_config(value: object) -> WebSearchInterceptionConfig:
try:
return _WEBSEARCH_CONFIG_ADAPTER.validate_python(value)
except ValidationError:
return {}
class _WebSearchSettingsView(TypedDict):
websearch_interception_params: WebSearchInterceptionConfig
class WebSearchInterceptionLogger(CustomLogger):
@ -434,7 +406,7 @@ class WebSearchInterceptionLogger(CustomLogger):
return tool.get("name")
@classmethod
def _sync_forced_tool_choice(cls, tool_choice: Any, converted_tools: list[dict[str, object]]) -> object:
def _sync_forced_tool_choice(cls, tool_choice: object, converted_tools: Sequence[Mapping[str, object]]) -> object:
"""Repoint a forced ``tool_choice`` at ``litellm_web_search`` when it
names a web-search tool that was just converted away.
@ -502,7 +474,7 @@ class WebSearchInterceptionLogger(CustomLogger):
kwargs[WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY] = True
# Convert native web search tools to LiteLLM standard
converted_tools: Final = []
converted_tools: Final[list[dict[str, object]]] = []
for tool in tools:
if is_web_search_tool(tool):
standard_tool = get_litellm_web_search_tool()
@ -873,7 +845,10 @@ class WebSearchInterceptionLogger(CustomLogger):
Anthropic-native clients (Claude Desktop, the Anthropic SDK) can
render citations / sources alongside the model's textual reply.
"""
native_blocks: Final = _parse_native_blocks(plan.metadata.get(WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY))
metadata_view: Final[_PlanMetadataView] = {
"websearch_native_blocks": plan.metadata.get(WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY)
}
native_blocks: Final = metadata_view["websearch_native_blocks"]
if not native_blocks:
return response
return self._inject_native_blocks(response, native_blocks)
@ -923,20 +898,17 @@ class WebSearchInterceptionLogger(CustomLogger):
)
@staticmethod
def _inject_native_blocks(
response: _ResponseT,
native_blocks: Sequence[Mapping[str, object]],
) -> _ResponseT:
def _inject_native_blocks(response: Any, native_blocks: Sequence[Mapping[str, object]]) -> Any:
"""Prepend native blocks to response content, dict or object form."""
if not native_blocks:
return response
if _is_json_object(response):
existing_items: Final = _parse_object_items(response.get(_CONTENT_ATTR))
response[_CONTENT_ATTR] = [*native_blocks, *existing_items]
if isinstance(response, dict):
existing = response.get("content") or []
response["content"] = list(native_blocks) + list(existing)
return response
existing_attr_items: Final = _parse_object_items(getattr(response, _CONTENT_ATTR, None))
existing = getattr(response, "content", None) or []
try:
setattr(response, _CONTENT_ATTR, [*native_blocks, *existing_attr_items])
response.content = list(native_blocks) + list(existing)
except (AttributeError, TypeError):
# Object refused write — fall through and leave the response
# untouched rather than crash the request.
@ -1321,8 +1293,10 @@ class WebSearchInterceptionLogger(CustomLogger):
kwargs_for_followup: Final = self._prepare_followup_kwargs(kwargs)
if logging_obj is not None:
agentic_params: Final[AgenticLoopParams] = logging_obj.model_call_details.get("agentic_loop_params", {})
full_model_name = agentic_params.get("model", model)
agentic_view: Final[_AgenticLoopParamsView] = {
"agentic_loop_params": logging_obj.model_call_details.get("agentic_loop_params", {})
}
full_model_name = agentic_view["agentic_loop_params"].get("model", model)
verbose_logger.debug(
"WebSearchInterception: Built anthropic request patch [call_id=%s model=%s messages=%d searches=%d]",
_call_id,
@ -1718,7 +1692,7 @@ class WebSearchInterceptionLogger(CustomLogger):
@staticmethod
def initialize_from_proxy_config(
litellm_settings: Mapping[str, object],
litellm_settings: dict[str, Any],
callback_specific_params: Mapping[str, object],
) -> "WebSearchInterceptionLogger":
"""
@ -1741,11 +1715,19 @@ class WebSearchInterceptionLogger(CustomLogger):
)
"""
# Get websearch_interception_params from litellm_settings or callback_specific_params
raw_params: Final = (
litellm_settings["websearch_interception_params"]
if "websearch_interception_params" in litellm_settings
else callback_specific_params.get("websearch_interception")
)
websearch_params: WebSearchInterceptionConfig = {}
if "websearch_interception_params" in litellm_settings:
settings_view: Final[_WebSearchSettingsView] = {
"websearch_interception_params": litellm_settings["websearch_interception_params"]
}
websearch_params = settings_view["websearch_interception_params"]
elif "websearch_interception" in callback_specific_params and isinstance(
callback_specific_params["websearch_interception"], dict
):
websearch_params = cast(
WebSearchInterceptionConfig,
callback_specific_params["websearch_interception"],
)
# Use classmethod to initialize from config
return WebSearchInterceptionLogger.from_config_yaml(_parse_websearch_config(raw_params))
return WebSearchInterceptionLogger.from_config_yaml(websearch_params)

3
litellm/litellm_core_utils/cli_token_utils.py
View file

@ -8,6 +8,7 @@ This module has no dependencies on proxy code and can be safely imported at the
import json
import os
import time
from collections.abc import Mapping
from pathlib import Path
from typing import Final
@ -71,7 +72,7 @@ def get_litellm_gateway_api_key(
return token_data["key"]
def is_cli_token_fresh(token_data: dict, buffer_hours: float = 0.1) -> bool:
def is_cli_token_fresh(token_data: Mapping[str, object], buffer_hours: float = 0.1) -> bool:
"""Check whether a cached CLI token (as stored in token.json) is still
within its expiration window. Used by `lite auth print-token` to fail
fast, without a network round trip, once the cached token is past

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

@ -1,12 +1,13 @@
from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping
from typing import (
TYPE_CHECKING,
Any,
Final,
TypeAlias,
cast,
)
from typing_extensions import TypedDict
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.asyncify import run_async_function
@ -39,6 +40,11 @@ _AnthropicSystem: TypeAlias = "str | list[dict[str, object]] | None"
_ContextManagementSpec: TypeAlias = "dict[str, object] | list[dict[str, object]] | None"
class _CompletionKwargs(TypedDict, total=False, extra_items=object):
model: str
custom_llm_provider: str
def _messages_have_compaction_block(messages: _AnthropicMessages) -> bool:
"""Return True when any message carries a ``compaction`` content block."""
for msg in messages:
@ -312,7 +318,7 @@ ANTHROPIC_ADAPTER: Final = AnthropicAdapter()
class LiteLLMMessagesToCompletionTransformationHandler:
@staticmethod
def _route_openai_thinking_to_responses_api_if_needed(
completion_kwargs: dict[str, Any],
completion_kwargs: _CompletionKwargs,
*,
thinking: Mapping[str, object] | None,
) -> None:
@ -377,7 +383,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
@staticmethod
def _normalize_reasoning_effort(
completion_kwargs: dict[str, Any],
completion_kwargs: _CompletionKwargs,
) -> None:
"""
Normalize reasoning_effort values based on target model capabilities.
@ -393,7 +399,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
if reasoning_effort is None:
return
model: Final = cast(str, completion_kwargs.get("model", ""))
model: Final = completion_kwargs.get("model", "")
custom_llm_provider: Final = completion_kwargs.get("custom_llm_provider")
if isinstance(reasoning_effort, str):
@ -417,19 +423,19 @@ class LiteLLMMessagesToCompletionTransformationHandler:
max_tokens: int,
messages: _AnthropicMessages,
model: str,
metadata: dict | None = None,
metadata: dict[str, object] | None = None,
stop_sequences: list[str] | None = None,
stream: bool | None = False,
system: _AnthropicSystem = None,
temperature: float | None = None,
thinking: dict | None = None,
tool_choice: dict | None = None,
tools: list[dict] | None = None,
thinking: dict[str, object] | None = None,
tool_choice: dict[str, object] | None = None,
tools: list[dict[str, object]] | None = None,
top_k: int | None = None,
top_p: float | None = None,
output_format: dict | None = None,
output_format: dict[str, object] | None = None,
extra_kwargs: Mapping[str, object] | None = None,
) -> tuple[dict[str, Any], dict[str, str]]:
) -> tuple[_CompletionKwargs, dict[str, str]]:
"""Prepare kwargs for litellm.completion/acompletion.
Returns:
@ -486,7 +492,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
if openai_request is None:
raise ValueError("Failed to translate request to OpenAI format")
completion_kwargs: Final[dict[str, Any]] = dict(openai_request)
completion_kwargs: Final[_CompletionKwargs] = {**openai_request}
if stream:
completion_kwargs["stream"] = stream
@ -538,17 +544,17 @@ class LiteLLMMessagesToCompletionTransformationHandler:
max_tokens: int,
messages: _AnthropicMessages,
model: str,
metadata: dict | None = None,
metadata: dict[str, object] | None = None,
stop_sequences: list[str] | None = None,
stream: bool | None = False,
system: str | None = None,
temperature: float | None = None,
thinking: dict | None = None,
tool_choice: dict | None = None,
thinking: dict[str, object] | None = None,
tool_choice: dict[str, object] | None = None,
tools: list[dict[str, object]] | None = None,
top_k: int | None = None,
top_p: float | None = None,
output_format: dict | None = None,
output_format: dict[str, object] | None = None,
**kwargs,
) -> AnthropicMessagesResponse | AsyncIterator[bytes] | Iterator[bytes]:
"""Handle non-Anthropic models asynchronously using the adapter"""
@ -625,17 +631,17 @@ class LiteLLMMessagesToCompletionTransformationHandler:
max_tokens: int,
messages: _AnthropicMessages,
model: str,
metadata: dict | None = None,
metadata: dict[str, object] | None = None,
stop_sequences: list[str] | None = None,
stream: bool | None = False,
system: str | None = None,
temperature: float | None = None,
thinking: dict | None = None,
tool_choice: dict | None = None,
thinking: dict[str, object] | None = None,
tool_choice: dict[str, object] | None = None,
tools: list[dict[str, object]] | None = None,
top_k: int | None = None,
top_p: float | None = None,
output_format: dict | None = None,
output_format: dict[str, object] | None = None,
_is_async: bool = False,
**kwargs,
) -> (

36
litellm/llms/anthropic/experimental_pass_through/messages/mcp_handler.py
View file

@ -7,8 +7,8 @@ tool through a ``tool_use`` content block, and results are fed back as
``tool_result`` blocks in a user message.
"""
from collections.abc import AsyncIterator, Mapping, Sequence
from typing import Any, Final
from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Mapping, Sequence
from typing import Any, Final, NamedTuple
from litellm._logging import verbose_logger
from litellm.responses.mcp.request_context import MCPRequestContext
@ -24,14 +24,18 @@ from litellm.types.llms.anthropic_messages.anthropic_response import (
MAX_MCP_TOOL_USE_ITERATIONS: Final = 10
def _get_response_content(response: AnthropicMessagesResponse) -> Sequence[Mapping[str, Any]]:
class _AnthropicMessagesCall(NamedTuple):
fn: Callable[..., Awaitable[AnthropicMessagesResponse | Iterator[bytes] | AsyncIterator[object]]]
def _get_response_content(response: AnthropicMessagesResponse) -> Sequence[Mapping[str, object]]:
content: Final = response.get("content")
if not isinstance(content, list):
return ()
return tuple(block for block in content if isinstance(block, dict))
def _extract_tool_use_blocks(response: AnthropicMessagesResponse) -> Sequence[Mapping[str, Any]]:
def _extract_tool_use_blocks(response: AnthropicMessagesResponse) -> Sequence[Mapping[str, object]]:
"""Return the ``tool_use`` content blocks the model emitted."""
return tuple(block for block in _get_response_content(response) if block.get("type") == "tool_use")
@ -41,7 +45,7 @@ def _get_stop_reason(response: AnthropicMessagesResponse) -> str | None:
return stop_reason if isinstance(stop_reason, str) else None
def _build_tool_result_message(tool_results: Sequence[Mapping[str, Any]]) -> AnthropicMessagesUserMessageParam:
def _build_tool_result_message(tool_results: Sequence[Mapping[str, object]]) -> AnthropicMessagesUserMessageParam:
"""Turn executed tool results into the user message Anthropic expects."""
return AnthropicMessagesUserMessageParam(
role="user",
@ -58,11 +62,11 @@ def _build_tool_result_message(tool_results: Sequence[Mapping[str, Any]]) -> Ant
async def anthropic_messages_with_mcp(
max_tokens: int,
messages: Sequence[Mapping[str, Any]],
messages: Sequence[Mapping[str, object]],
model: str,
tools: Sequence[Mapping[str, Any]] | None = None,
tools: Sequence[Mapping[str, object]] | None = None,
**kwargs: Any, # kwargs-ok: forwarded verbatim to litellm.anthropic_messages, which owns the param contract
) -> AnthropicMessagesResponse | AsyncIterator[Any]:
) -> AnthropicMessagesResponse | Iterator[bytes] | AsyncIterator[object]:
"""
Expand litellm_proxy MCP references for `/v1/messages` and run the tool loop.
@ -81,7 +85,7 @@ async def anthropic_messages_with_mcp(
mcp_references, other_tools = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools)
if not mcp_references:
return await litellm.anthropic_messages(
return await _AnthropicMessagesCall(fn=litellm.anthropic_messages).fn(
max_tokens=max_tokens,
messages=list(messages),
model=model,
@ -114,7 +118,7 @@ async def anthropic_messages_with_mcp(
)
stream: Final = bool(kwargs.pop("stream", False))
base_call_args: Final[Mapping[str, Any]] = {
base_call_args: Final[Mapping[str, object]] = {
"max_tokens": max_tokens,
"model": model,
"tools": all_tools or None,
@ -123,10 +127,12 @@ async def anthropic_messages_with_mcp(
}
if not should_auto_execute:
return await litellm.anthropic_messages(messages=list(messages), stream=stream, **base_call_args)
return await _AnthropicMessagesCall(fn=litellm.anthropic_messages).fn(
messages=list(messages), stream=stream, **base_call_args
)
working_messages: Sequence[Mapping[str, Any]] = tuple(messages)
response: AnthropicMessagesResponse = await litellm.anthropic_messages(
working_messages: Sequence[Mapping[str, object]] = tuple(messages)
response: AnthropicMessagesResponse = await _AnthropicMessagesCall(fn=litellm.anthropic_messages).fn(
messages=list(working_messages), stream=False, **base_call_args
)
@ -161,7 +167,9 @@ async def anthropic_messages_with_mcp(
{"role": "assistant", "content": list(_get_response_content(response))},
_build_tool_result_message(tool_results),
)
response = await litellm.anthropic_messages(messages=list(working_messages), stream=False, **base_call_args)
response = await _AnthropicMessagesCall(fn=litellm.anthropic_messages).fn(
messages=list(working_messages), stream=False, **base_call_args
)
else:
verbose_logger.warning(
"MCP tool loop hit its %s iteration cap for model %s; returning the last response",

33
litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py
View file

@ -8,7 +8,12 @@ from collections.abc import AsyncIterator, Coroutine
from typing import Any, Final
import litellm
from litellm.types.llms.anthropic import AnthropicMessagesRequest
from litellm.types.llms.anthropic import (
AllAnthropicToolsValues,
AnthropicMessagesRequest,
AnthropicOutputConfig,
AnthropicOutputSchema,
)
from litellm.types.llms.anthropic_messages.anthropic_response import (
AnthropicMessagesResponse,
)
@ -27,24 +32,24 @@ def _build_responses_kwargs(
model: str,
context_management: dict | None = None,
metadata: dict | None = None,
output_config: dict | None = None,
output_config: AnthropicOutputConfig | None = None,
stop_sequences: list[str] | None = None,
stream: bool | None = False,
system: str | None = None,
temperature: float | None = None,
thinking: dict | None = None,
tool_choice: dict | None = None,
tools: list[dict] | None = None,
tools: list[AllAnthropicToolsValues | dict] | None = None,
top_k: int | None = None,
top_p: float | None = None,
output_format: dict | None = None,
output_format: AnthropicOutputSchema | None = None,
extra_kwargs: dict[str, Any] | None = None,
) -> dict[str, Any]:
"""
Build the kwargs dict to pass directly to litellm.responses() / litellm.aresponses().
"""
# Build a typed AnthropicMessagesRequest for the adapter
request_data: Final[dict[str, Any]] = {
request_data: Final[AnthropicMessagesRequest] = {
"model": model,
"messages": messages,
"max_tokens": max_tokens,
@ -128,19 +133,19 @@ class LiteLLMMessagesToResponsesAPIHandler:
model: str,
context_management: dict | None = None,
metadata: dict | None = None,
output_config: dict | None = None,
output_config: AnthropicOutputConfig | None = None,
stop_sequences: list[str] | None = None,
stream: bool | None = False,
system: str | None = None,
temperature: float | None = None,
thinking: dict | None = None,
tool_choice: dict | None = None,
tools: list[dict] | None = None,
tools: list[AllAnthropicToolsValues | dict] | None = None,
top_k: int | None = None,
top_p: float | None = None,
output_format: dict | None = None,
output_format: AnthropicOutputSchema | None = None,
**kwargs,
) -> AnthropicMessagesResponse | AsyncIterator:
) -> AnthropicMessagesResponse | AsyncIterator[bytes]:
responses_kwargs: Final = _build_responses_kwargs(
max_tokens=max_tokens,
messages=messages,
@ -179,23 +184,23 @@ class LiteLLMMessagesToResponsesAPIHandler:
model: str,
context_management: dict | None = None,
metadata: dict | None = None,
output_config: dict | None = None,
output_config: AnthropicOutputConfig | None = None,
stop_sequences: list[str] | None = None,
stream: bool | None = False,
system: str | None = None,
temperature: float | None = None,
thinking: dict | None = None,
tool_choice: dict | None = None,
tools: list[dict] | None = None,
tools: list[AllAnthropicToolsValues | dict] | None = None,
top_k: int | None = None,
top_p: float | None = None,
output_format: dict | None = None,
output_format: AnthropicOutputSchema | None = None,
_is_async: bool = False,
**kwargs,
) -> (
AnthropicMessagesResponse
| AsyncIterator[Any]
| Coroutine[Any, Any, AnthropicMessagesResponse | AsyncIterator[Any]]
| AsyncIterator[bytes]
| Coroutine[None, None, AnthropicMessagesResponse | AsyncIterator[bytes]]
):
if _is_async:
return LiteLLMMessagesToResponsesAPIHandler.async_anthropic_messages_handler(

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

@ -15,6 +15,8 @@ from collections.abc import Mapping, Sequence
from typing import Any, Final, NamedTuple, Optional, Protocol, Union, runtime_checkable
if typing.TYPE_CHECKING:
from collections.abc import Awaitable, Callable
from fastapi import Request
from mcp.client.session import ClientSession
from mcp.shared.context import RequestContext
@ -28,8 +30,9 @@ if typing.TYPE_CHECKING:
ToolUseContent,
)
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.utils import ProxyLogging
from litellm.types.utils import ModelResponse
from fastapi import HTTPException
from pydantic import TypeAdapter
@ -1016,7 +1019,7 @@ async def _run_budget_checks(
general_settings=general_settings or {},
route="/chat/completions",
llm_router=_llm_router,
proxy_logging_obj=typing.cast("ProxyLogging", _proxy_logging_obj),
proxy_logging_obj=_proxy_logging_obj,
valid_token=user_api_key_auth,
request=dummy_request,
)
@ -1176,15 +1179,19 @@ async def _build_completion_kwargs(
)
class _AcompletionCall(NamedTuple):
fn: "Callable[..., Awaitable[ModelResponse | CustomStreamWrapper]]"
async def _run_guardrails_and_call_llm(
completion_kwargs: dict[str, Any],
completion_kwargs: dict[str, object],
user_api_key_auth: "UserAPIKeyAuth",
) -> Any:
try:
from litellm.proxy.proxy_server import proxy_logging_obj as _plo
if _plo is not None:
completion_kwargs = await typing.cast("ProxyLogging", _plo).pre_call_hook(
completion_kwargs = await _plo.pre_call_hook(
user_api_key_dict=user_api_key_auth,
data=completion_kwargs,
call_type="acompletion",
@ -1204,10 +1211,10 @@ async def _run_guardrails_and_call_llm(
from litellm.proxy.proxy_server import llm_router
if llm_router is not None:
return await llm_router.acompletion(**completion_kwargs)
return await litellm.acompletion(**completion_kwargs)
return await _AcompletionCall(fn=llm_router.acompletion).fn(**completion_kwargs)
return await _AcompletionCall(fn=litellm.acompletion).fn(**completion_kwargs)
except ImportError:
return await litellm.acompletion(**completion_kwargs)
return await _AcompletionCall(fn=litellm.acompletion).fn(**completion_kwargs)
async def handle_sampling_create_message(

68
litellm/proxy/agent_endpoints/a2a_endpoints.py
View file

@ -11,7 +11,7 @@ The A2A SDK can point to LiteLLM's URL and invoke agents registered with LiteLLM
"""
import json
from collections.abc import AsyncGenerator
from collections.abc import AsyncGenerator, Mapping
from copy import deepcopy
from typing import TYPE_CHECKING, Any, Final
from urllib.parse import urlparse
@ -36,7 +36,7 @@ from litellm.proxy.agent_endpoints.databricks_oauth import (
)
from litellm.proxy.agent_endpoints.utils import merge_agent_headers
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.utils import get_custom_url
from litellm.proxy.utils import ProxyLogging, get_custom_url
from litellm.types.utils import all_litellm_params
if TYPE_CHECKING:
@ -46,7 +46,7 @@ if TYPE_CHECKING:
router: Final = APIRouter()
_PASCAL_TO_WIRE: Final[dict[str, str]] = {
_PASCAL_TO_WIRE: Final[Mapping[str, str]] = {
"SendMessage": "message/send",
"SendStreamingMessage": "message/stream",
"GetTask": "tasks/get",
@ -118,9 +118,9 @@ def _caller_identity_headers(user_api_key_dict: UserAPIKeyAuth) -> dict[str, str
def _forwarding_headers(
user_api_key_dict: UserAPIKeyAuth,
request_data: dict[str, Any],
agent_extra_headers: dict[str, str] | None,
) -> dict[str, str] | None:
request_data: Mapping[str, object],
agent_extra_headers: Mapping[str, str] | None,
) -> Mapping[str, str] | None:
sanitized: Final = (
{k: v for k, v in agent_extra_headers.items() if not k.lower().startswith("x-litellm-")}
if agent_extra_headers
@ -136,7 +136,7 @@ def _forwarding_headers(
def _jsonrpc_error(
request_id: Any | None,
request_id: object,
code: int,
message: str,
status_code: int = 400,
@ -162,7 +162,7 @@ def _get_agent(agent_id: str):
return agent
def _enforce_inbound_trace_id(agent: Any, request: Request) -> None:
def _enforce_inbound_trace_id(agent: "AgentResponse", request: Request) -> None:
"""Raise 400 if agent requires x-litellm-trace-id on inbound calls and it is missing."""
agent_litellm_params: Final = agent.litellm_params or {}
if not agent_litellm_params.get("require_trace_id_on_calls_to_agent"):
@ -181,8 +181,8 @@ def _enforce_inbound_trace_id(agent: Any, request: Request) -> None:
async def _forward_jsonrpc(
agent_url: str,
body: dict[str, Any],
extra_headers: dict[str, str] | None = None,
body: dict[str, object],
extra_headers: Mapping[str, str] | None = None,
) -> dict[str, Any]:
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.types.llms.custom_http import httpxSpecialProvider
@ -205,11 +205,11 @@ async def _forward_jsonrpc(
async def _a2a_sse_event_source(
agent_url: str,
body: dict[str, Any],
request_id: Any | None = None,
extra_headers: dict[str, str] | None = None,
body: Mapping[str, object],
request_id: str | int | None = None,
extra_headers: Mapping[str, str] | None = None,
served_version: A2AVersion = "0.3",
) -> AsyncGenerator[dict, None]:
) -> AsyncGenerator[Mapping[str, object], None]:
"""Stream an upstream A2A SSE response as parsed JSON-RPC event dicts.
Upstream HTTP/JSON-RPC errors are surfaced as a single JSON-RPC error event
@ -234,7 +234,7 @@ async def _a2a_sse_event_source(
try:
if not resp.is_success:
error_body: Final = await resp.aread()
error_event: dict[str, Any] | None = None
error_event: Mapping[str, object] | None = None
try:
parsed: Final = json.loads(error_body)
if isinstance(parsed, dict) and "error" in parsed:
@ -267,12 +267,12 @@ async def _a2a_sse_event_source(
async def _forward_jsonrpc_sse(
agent_url: str,
body: dict[str, Any],
request_id: Any | None = None,
extra_headers: dict[str, str] | None = None,
proxy_logging_obj: Any | None = None,
user_api_key_dict: Any | None = None,
request_data: dict[str, Any] | None = None,
body: Mapping[str, object],
request_id: str | int | None = None,
extra_headers: Mapping[str, str] | None = None,
proxy_logging_obj: ProxyLogging | None = None,
user_api_key_dict: UserAPIKeyAuth | None = None,
request_data: dict[str, object] | None = None,
served_version: A2AVersion = "0.3",
) -> StreamingResponse:
event_source: Final = _a2a_sse_event_source(
@ -283,10 +283,10 @@ async def _forward_jsonrpc_sse(
served_version=served_version,
)
def _serialize_chunk(chunk: Any) -> str:
def _serialize_chunk(chunk: object) -> str:
return f"data: {json.dumps(chunk)}\n\n"
def _serialize_error(proxy_exc: Any) -> str:
def _serialize_error(proxy_exc: object) -> str:
return (
"data: "
+ json.dumps(
@ -331,17 +331,17 @@ async def _forward_jsonrpc_sse(
async def _handle_stream_message(
api_base: str | None,
request_id: Any,
params: dict[str, Any],
litellm_params: dict[str, Any] | None = None,
request_id: str | int,
params: dict[str, object],
litellm_params: dict[str, object] | None = None,
agent_id: str | None = None,
metadata: dict[str, Any] | None = None,
proxy_server_request: dict[str, Any] | None = None,
metadata: dict[str, object] | None = None,
proxy_server_request: dict[str, object] | None = None,
*,
agent_extra_headers: dict[str, str] | None = None,
user_api_key_dict: UserAPIKeyAuth | None = None,
request_data: dict[str, Any] | None = None,
proxy_logging_obj: Any | None = None,
request_data: dict[str, object] | None = None,
proxy_logging_obj: ProxyLogging | None = None,
served_version: A2AVersion = "0.3",
) -> StreamingResponse:
"""Handle message/stream method via SDK functions.
@ -430,7 +430,7 @@ async def _handle_stream_message(
obj = normalize_stream_event(obj, served_version, request_id=request_id)
return json.dumps(obj) + "\n"
def _ndjson_error(proxy_exc: Any) -> str:
def _ndjson_error(proxy_exc: object) -> str:
return (
json.dumps(
{
@ -669,7 +669,7 @@ async def invoke_agent_a2a(
agent_name: Final = agent_card_params.get("name", agent_id)
# Get litellm_params (may include custom_llm_provider for completion bridge)
litellm_params = agent.litellm_params or {}
litellm_params: dict[str, object] = agent.litellm_params or {}
custom_llm_provider: Final = litellm_params.get("custom_llm_provider")
# Hand the authenticated key hash to the completion bridge so provider
@ -725,7 +725,7 @@ async def invoke_agent_a2a(
request_data = data
# Build merged headers for the backend agent
static_headers: Final[dict[str, str]] = dict(agent.static_headers or {})
static_headers: Final[Mapping[str, str]] = dict(agent.static_headers or {})
raw_headers: Final = dict(request.headers)
normalized: Final = {k.lower(): v for k, v in raw_headers.items()}
@ -893,7 +893,7 @@ async def invoke_agent_a2a(
detail="Push notification URL must be a string",
)
_validate_push_notification_url(callback_url)
forward_body = {
forward_body: dict[str, object] = {
"jsonrpc": "2.0",
"id": request_id,
"method": method,

102
litellm/proxy/client/cli/commands/auth.py
View file

@ -11,6 +11,7 @@ import click
import requests
from rich.console import Console
from rich.table import Table
from typing_extensions import NotRequired, TypedDict
from litellm.constants import CLI_JWT_EXPIRATION_HOURS
from litellm.litellm_core_utils.cli_token_utils import is_cli_token_fresh
@ -18,6 +19,57 @@ from litellm.litellm_core_utils.cli_token_utils import is_cli_token_fresh
from .private_json import write_private_json
class CliTokenData(TypedDict):
base_url: str
key: str
user_id: str
user_email: str
user_role: str
auth_header_name: str
jwt_token: str
timestamp: float
class CliTeam(TypedDict, total=False):
team_id: str | None
team_alias: str | None
models: list[str]
max_budget: float | None
class CliContextObj(TypedDict):
base_url: str
base_url_explicit: NotRequired[bool]
class CliPollData(TypedDict, total=False):
status: str
key: str
user_id: str
teams: list[str]
team_details: object
requires_team_selection: bool
team_id: str
class CliPollRequestKwargs(TypedDict, total=False):
timeout: int
headers: dict[str, str]
class CliSsoStartData(TypedDict):
login_id: str
poll_secret: str
user_code: str
class CliAuthResult(TypedDict):
api_key: str
user_id: str | None
teams: list[str]
team_id: str | None
# Token storage utilities
def get_token_file_path() -> str:
"""Get the path to store the authentication token"""
@ -27,12 +79,12 @@ def get_token_file_path() -> str:
return str(config_dir / "token.json")
def save_token(token_data: dict[str, Any]) -> None:
def save_token(token_data: CliTokenData) -> None:
"""Save token data to file"""
write_private_json(get_token_file_path(), token_data)
def load_token() -> dict[str, Any] | None:
def load_token() -> CliTokenData | None:
"""Load token data from file"""
token_file: Final = get_token_file_path()
if not os.path.exists(token_file):
@ -65,7 +117,7 @@ def get_stored_api_key(expected_base_url: str | None = None) -> str | None:
# Team selection utilities
def display_teams_table(teams: list[dict[str, Any]]) -> None:
def display_teams_table(teams: list[CliTeam]) -> None:
"""Display teams in a formatted table"""
console: Final = Console()
@ -165,7 +217,7 @@ def display_interactive_team_selection(teams: list[dict[str, Any]], selected_ind
for i, team in enumerate(teams):
team_alias = team.get("team_alias") or "N/A"
team_id = team.get("team_id", "N/A")
models = team.get("models", [])
models: list[str] = team.get("models", [])
max_budget = team.get("max_budget")
# Format models list
@ -249,10 +301,11 @@ def prompt_team_selection_fallback(
while True:
try:
choice = click.prompt(
prompt_response: str = click.prompt(
"\nSelect a team by entering the index number (or 'skip' to continue without a team)",
type=str,
).strip()
)
choice = prompt_response.strip()
if choice.lower() == "skip":
return None
@ -275,7 +328,7 @@ def prompt_team_selection_fallback(
def _response_error_detail(response: requests.Response) -> str | None:
try:
body: Final = response.json()
body: Final[dict[str, object] | list[object] | str | int | float | bool | None] = response.json()
except ValueError:
return None
detail: Final = body.get("detail") if isinstance(body, dict) else None
@ -309,15 +362,15 @@ def _poll_for_ready_data(
other_status_log_every: int = 10,
http_error_log_every: int = 10,
connection_error_log_every: int = 10,
) -> dict[str, Any] | None:
) -> CliPollData | None:
for attempt in range(total_timeout // poll_interval):
try:
request_kwargs: dict[str, Any] = {"timeout": request_timeout}
request_kwargs: CliPollRequestKwargs = {"timeout": request_timeout}
if headers is not None:
request_kwargs["headers"] = headers
response = requests.get(url, **request_kwargs)
if response.status_code == 200:
data = response.json()
data: CliPollData = response.json()
status = data.get("status")
if status == "ready":
return data
@ -341,7 +394,7 @@ def _poll_for_ready_data(
return None
def _normalize_teams(teams, team_details):
def _normalize_teams(teams: object, team_details: object) -> list[CliTeam]:
"""If team_details are a
Args:
@ -365,7 +418,7 @@ def _normalize_teams(teams, team_details):
return []
def _start_cli_sso_flow(base_url: str) -> dict[str, Any]:
def _start_cli_sso_flow(base_url: str) -> CliSsoStartData:
start_url: Final = f"{base_url}/sso/cli/start"
try:
response: Final = requests.post(start_url, timeout=10)
@ -389,7 +442,7 @@ def _start_cli_sso_flow(base_url: str) -> dict[str, Any]:
)
try:
data: Final = response.json()
data: Final[CliSsoStartData] = response.json()
except ValueError:
content_type: Final = response.headers.get("content-type", "unknown")
raise ValueError(
@ -398,7 +451,7 @@ def _start_cli_sso_flow(base_url: str) -> dict[str, Any]:
f"Response starts with: {response.text[:200]!r}"
)
required_fields: Final = ("login_id", "poll_secret", "user_code")
required_fields: Final[tuple[str, ...]] = ("login_id", "poll_secret", "user_code")
missing_fields: Final = tuple(field for field in required_fields if not isinstance(data.get(field), str))
if missing_fields:
raise ValueError(
@ -412,7 +465,7 @@ def _get_cli_sso_poll_headers(poll_secret: str) -> dict[str, str]:
return {"x-litellm-cli-poll-secret": poll_secret}
def _poll_for_authentication(base_url: str, key_id: str, poll_secret: str) -> dict | None:
def _poll_for_authentication(base_url: str, key_id: str, poll_secret: str) -> CliAuthResult | None:
"""
Poll the server for authentication completion and handle team selection.
@ -431,7 +484,7 @@ def _poll_for_authentication(base_url: str, key_id: str, poll_secret: str) -> di
teams = data.get("teams", [])
team_details: Final = data.get("team_details")
user_id = data.get("user_id")
normalized_teams: Final[list[dict[str, Any]]] = _normalize_teams(teams, team_details)
normalized_teams: Final[list[CliTeam]] = _normalize_teams(teams, team_details)
if not normalized_teams:
click.echo("Warning: No teams available for selection.")
return None
@ -478,7 +531,7 @@ def _poll_for_authentication(base_url: str, key_id: str, poll_secret: str) -> di
def _handle_team_selection_during_polling(
base_url: str, key_id: str, poll_secret: str, teams: list[dict[str, Any]]
base_url: str, key_id: str, poll_secret: str, teams: list[CliTeam]
) -> str | None:
"""
Handle team selection and re-poll with selected team_id.
@ -522,7 +575,7 @@ def _handle_team_selection_during_polling(
return None
def _render_and_prompt_for_team_selection(teams: list[dict[str, Any]]) -> str | None:
def _render_and_prompt_for_team_selection(teams: list[CliTeam]) -> str | None:
"""Render teams table and prompt user for a team selection.
Returns the selected team_id as a string, or None if selection was
@ -546,10 +599,11 @@ def _render_and_prompt_for_team_selection(teams: list[dict[str, Any]]) -> str |
# Simple selection
while True:
try:
choice = click.prompt(
prompt_response: str = click.prompt(
"\nSelect a team by entering the index number (or 'skip' to use first team)",
type=str,
).strip()
)
choice = prompt_response.strip()
if choice.lower() == "skip":
# Default to the first team's ID if the user skips an
@ -582,7 +636,8 @@ def login(ctx: click.Context):
from litellm.constants import LITELLM_CLI_SOURCE_IDENTIFIER
from litellm.proxy.client.cli.interface import show_commands
base_url: Final = ctx.obj["base_url"]
ctx_obj: Final[CliContextObj] = ctx.obj
base_url: Final = ctx_obj["base_url"]
try:
cli_sso_flow: Final = _start_cli_sso_flow(base_url=base_url)
@ -675,8 +730,9 @@ def print_token(ctx: click.Context):
# explicitly pointed us at a server, trust whichever one `lite login`
# actually issued this token for -- that's the whole point of not
# needing a wrapper command.
if ctx.obj.get("base_url_explicit"):
base_url: Final = ctx.obj["base_url"]
ctx_obj: Final[CliContextObj] = ctx.obj
if ctx_obj.get("base_url_explicit"):
base_url: Final = ctx_obj["base_url"]
if token_data.get("base_url") != base_url.rstrip("/"):
click.echo("Not authenticated for this server. Run 'lite login'.", err=True)
sys.exit(1)

9
litellm/proxy/db/exception_handler.py
View file

@ -1,5 +1,5 @@
from collections.abc import Awaitable, Callable
from typing import Any, Final
from typing import Any, Final, TypeVar
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import (
@ -311,14 +311,17 @@ def _coerce_timeout(value: Any, fallback: float) -> float:
return fallback
_ReadResultT: Final = TypeVar("_ReadResultT")
async def call_with_db_reconnect_retry(
prisma_client: Any,
coro_factory: Callable[[], Awaitable[Any]],
coro_factory: Callable[[], Awaitable[_ReadResultT]],
*,
reason: str,
timeout_seconds: float | None = None,
lock_timeout_seconds: float | None = None,
) -> Any:
) -> _ReadResultT:
"""Run a Prisma read coroutine with one transport-reconnect-and-retry.
The canonical "self-heal a transient DB transport blip" wrapper used by

169
litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py
View file

@ -9,10 +9,10 @@ import asyncio
import json
import os
import re
from collections.abc import AsyncGenerator
from collections.abc import AsyncGenerator, Coroutine, Mapping, Sequence
from datetime import datetime
from re import Pattern
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict, cast
import yaml
from fastapi import HTTPException
@ -28,6 +28,7 @@ from litellm.types.utils import (
GenericGuardrailAPIInputs,
GuardrailStatus,
GuardrailTracingDetail,
ModelResponse,
ModelResponseStream,
)
@ -83,6 +84,46 @@ WORD_NUMBER_SEQUENCE_PATTERN: Final = re.compile(
WORD_NUMBER_TOKEN_FINDER: Final = re.compile(rf"(?:{WORD_NUMBER_TOKEN_REGEX})", re.IGNORECASE)
class ConditionalCategoryConfig(TypedDict):
identifier_words: Sequence[str]
block_words: Sequence[str]
action: ContentFilterAction
severity: str
class CompiledPatternEntry(TypedDict):
regex: Pattern[str]
pattern_name: str
action: ContentFilterAction
keyword_regex: Pattern[str] | None
allow_word_numbers: bool
class _PatternExtraLookup(TypedDict):
keyword_pattern: str | None
allow_word_numbers: bool
class _CategoryConfigView(TypedDict):
category: object
enabled: object
action: object
category_file: str | None
class CategoryFileData(TypedDict, total=False):
category_name: str
description: str
default_action: str
keywords: Sequence[Mapping[str, str]]
exceptions: Sequence[str]
identifier_words: Sequence[str]
always_block_keywords: Sequence[Mapping[str, str]]
inherit_from: str
additional_block_words: Sequence[str]
phrase_patterns: Sequence[str]
# Helper data structure for category-based detection
class CategoryConfig:
"""Configuration for a content category."""
@ -92,13 +133,13 @@ class CategoryConfig:
category_name: str,
description: str,
default_action: ContentFilterAction,
keywords: list[dict[str, str]],
exceptions: list[str],
identifier_words: list[str] | None = None,
always_block_keywords: list[dict[str, str]] | None = None,
keywords: Sequence[Mapping[str, str]],
exceptions: Sequence[str],
identifier_words: Sequence[str] | None = None,
always_block_keywords: Sequence[Mapping[str, str]] | None = None,
inherit_from: str | None = None,
additional_block_words: list[str] | None = None,
phrase_patterns: list[str] | None = None,
additional_block_words: Sequence[str] | None = None,
phrase_patterns: Sequence[str] | None = None,
):
self.category_name = category_name
self.description = description
@ -151,7 +192,7 @@ class ContentFilterGuardrail(CustomGuardrail):
severity_threshold: str = "medium",
llm_router: Router | None = None,
image_model: str | None = None,
competitor_intent_config: dict[str, Any] | None = None,
competitor_intent_config: dict[str, object] | None = None,
**kwargs,
):
"""
@ -194,9 +235,7 @@ class ContentFilterGuardrail(CustomGuardrail):
# Always-block keywords are checked after exceptions (exceptions take precedence)
self.always_block_category_keywords: dict[str, tuple[str, str, ContentFilterAction]] = {}
# Store conditional categories (identifier_words + block_words)
self.conditional_categories: dict[
str, dict[str, Any]
] = {} # category_name -> {identifier_words, block_words, action, severity}
self.conditional_categories: dict[str, ConditionalCategoryConfig] = {}
# Competitor intent checker (optional; airline uses major_airlines.json, generic requires competitors)
self._competitor_intent_checker: BaseCompetitorIntentChecker | None = None
@ -212,7 +251,7 @@ class ContentFilterGuardrail(CustomGuardrail):
normalized_blocked_words: Final = self._normalize_blocked_words(blocked_words)
# Compile regex patterns
self.compiled_patterns: list[dict[str, Any]] = []
self.compiled_patterns: list[CompiledPatternEntry] = []
for pattern_config in normalized_patterns:
self._add_pattern(pattern_config)
@ -250,7 +289,7 @@ class ContentFilterGuardrail(CustomGuardrail):
"Loaded %s categories with %s keywords", len(self.loaded_categories), len(self.category_keywords)
)
def _init_competitor_intent_checker(self, competitor_intent_config: dict[str, Any]) -> None:
def _init_competitor_intent_checker(self, competitor_intent_config: dict[str, object]) -> None:
try:
competitor_intent_type: Final = competitor_intent_config.get("competitor_intent_type", "airline")
if competitor_intent_type == "generic":
@ -293,6 +332,15 @@ class ContentFilterGuardrail(CustomGuardrail):
result.append(word)
return result
@staticmethod
def _category_config_view(cat_config: ContentFilterCategoryConfig) -> _CategoryConfigView:
return {
"category": cat_config.get("category"),
"enabled": cat_config.get("enabled", True),
"action": cat_config.get("action"),
"category_file": cat_config.get("category_file"),
}
@staticmethod
def _assert_within_categories_dir(path: str, categories_dir: str) -> None:
"""Raise ValueError if path escapes the categories directory."""
@ -395,7 +443,8 @@ class ContentFilterGuardrail(CustomGuardrail):
categories_dir: Final = os.path.join(os.path.dirname(__file__), "categories")
for cat_config in categories:
category_name = cat_config.get("category")
view = self._category_config_view(cat_config)
category_name = view["category"]
if not category_name or not isinstance(category_name, str):
verbose_proxy_logger.warning("Category name missing or invalid in config, skipping")
continue
@ -405,12 +454,12 @@ class ContentFilterGuardrail(CustomGuardrail):
verbose_proxy_logger.warning("Category name '%s' contains invalid characters, skipping", category_name)
continue
enabled = cat_config.get("enabled", True)
action = cat_config.get("action")
enabled = view["enabled"]
action = view["action"]
severity_threshold = (
cat_config.get("severity_threshold", self.severity_threshold) or self.severity_threshold
)
custom_file = cat_config.get("category_file")
custom_file = view["category_file"]
if not enabled:
verbose_proxy_logger.debug("Category %s is disabled, skipping", category_name)
@ -514,7 +563,7 @@ class ContentFilterGuardrail(CustomGuardrail):
categories_dir: Directory containing category files
"""
try:
block_words: Final = []
block_words: Final[list[str]] = []
inherit_from = category_config_obj.inherit_from
# Load inherited block words if specified
@ -605,11 +654,7 @@ class ContentFilterGuardrail(CustomGuardrail):
"""
if file_path.lower().endswith(".json"):
return self._load_category_file_json(file_path)
with open(file_path, "r") as f:
data: Final = yaml.safe_load(f)
# Handle always_block_keywords if present
always_block: Final = data.get("always_block_keywords", [])
data: Final = self._read_category_yaml(file_path)
return CategoryConfig(
category_name=data.get("category_name", "unknown"),
@ -618,12 +663,17 @@ class ContentFilterGuardrail(CustomGuardrail):
keywords=data.get("keywords", []),
exceptions=data.get("exceptions", []),
identifier_words=data.get("identifier_words"),
always_block_keywords=always_block,
always_block_keywords=data.get("always_block_keywords", []),
inherit_from=data.get("inherit_from"),
additional_block_words=data.get("additional_block_words"),
phrase_patterns=data.get("phrase_patterns"),
)
@staticmethod
def _read_category_yaml(file_path: str) -> CategoryFileData:
with open(file_path, "r") as f:
return yaml.safe_load(f)
def _load_category_file_json(self, file_path: str) -> CategoryConfig:
"""
Load a category from the harm_toxic_abuse-style JSON format.
@ -682,13 +732,13 @@ class ContentFilterGuardrail(CustomGuardrail):
pattern_config: ContentFilterPattern configuration
"""
try:
extra_config: dict[str, Any] = {}
extra_config: _PatternExtraLookup = {"keyword_pattern": None, "allow_word_numbers": False}
if pattern_config.pattern_type == "prebuilt":
if not pattern_config.pattern_name:
raise ValueError("pattern_name is required for prebuilt patterns")
compiled = get_compiled_pattern(pattern_config.pattern_name)
pattern_name = pattern_config.pattern_name
extra_config = PATTERN_EXTRA_CONFIG.get(pattern_name, {}) or {}
extra_config = self._lookup_pattern_extra(pattern_name)
elif pattern_config.pattern_type == "regex":
if not pattern_config.pattern:
raise ValueError("pattern is required for regex patterns")
@ -697,9 +747,8 @@ class ContentFilterGuardrail(CustomGuardrail):
else:
raise ValueError(f"Unknown pattern_type: {pattern_config.pattern_type}")
keyword_regex: Pattern | None = None
if extra_config.get("keyword_pattern"):
keyword_regex = re.compile(extra_config["keyword_pattern"], re.IGNORECASE)
keyword_pattern: Final = extra_config["keyword_pattern"]
keyword_regex: Final = re.compile(keyword_pattern, re.IGNORECASE) if keyword_pattern else None
self.compiled_patterns.append(
{
@ -707,7 +756,7 @@ class ContentFilterGuardrail(CustomGuardrail):
"pattern_name": pattern_name,
"action": pattern_config.action,
"keyword_regex": keyword_regex,
"allow_word_numbers": bool(extra_config.get("allow_word_numbers")),
"allow_word_numbers": extra_config["allow_word_numbers"],
}
)
verbose_proxy_logger.debug("Added pattern: %s with action %s", pattern_name, pattern_config.action)
@ -715,6 +764,14 @@ class ContentFilterGuardrail(CustomGuardrail):
verbose_proxy_logger.error("Error adding pattern %s: %s", pattern_config, e)
raise
@staticmethod
def _lookup_pattern_extra(pattern_name: str) -> _PatternExtraLookup:
extra: Final = PATTERN_EXTRA_CONFIG.get(pattern_name)
return {
"keyword_pattern": extra.get("keyword_pattern") if extra is not None else None,
"allow_word_numbers": bool(extra.get("allow_word_numbers")) if extra is not None else False,
}
def _load_blocked_words_file(self, file_path: str) -> None:
"""
Load blocked words from a YAML file.
@ -754,18 +811,16 @@ class ContentFilterGuardrail(CustomGuardrail):
except Exception as e:
raise Exception(f"Error loading blocked words file {file_path}: {e}")
def _find_pattern_spans(self, text: str, pattern_entry: dict[str, Any]) -> list[tuple[int, int]]:
def _find_pattern_spans(self, text: str, pattern_entry: CompiledPatternEntry) -> list[tuple[int, int]]:
"""Return all match spans for a pattern, applying contextual rules if required."""
regex: Final[Pattern] = pattern_entry["regex"]
keyword_regex: Final[Pattern | None] = pattern_entry.get("keyword_regex")
regex: Final[Pattern[str]] = pattern_entry["regex"]
keyword_regex: Final[Pattern[str] | None] = pattern_entry.get("keyword_regex")
allow_word_numbers: Final[bool] = pattern_entry.get("allow_word_numbers", False)
keyword_matches: list[re.Match] | None = None
if keyword_regex is not None:
keyword_matches = list(keyword_regex.finditer(text))
if not keyword_matches:
return []
keyword_matches: Final = list(keyword_regex.finditer(text)) if keyword_regex is not None else None
if keyword_matches is not None and not keyword_matches:
return []
match_spans: Final[list[tuple[int, int]]] = []
@ -795,7 +850,7 @@ class ContentFilterGuardrail(CustomGuardrail):
self,
value_start: int,
value_end: int,
keyword_matches: list[re.Match],
keyword_matches: Sequence[re.Match[str]],
text: str,
) -> bool:
"""Check if a value is separated from a keyword by an allowed gap."""
@ -861,7 +916,7 @@ class ContentFilterGuardrail(CustomGuardrail):
def _convert_word_number_sequence(self, sequence: str) -> str | None:
"""Convert a spelled-out digit sequence (e.g., 'One-Two') into digits."""
tokens: Final = WORD_NUMBER_TOKEN_FINDER.findall(sequence)
tokens: Final[list[str]] = WORD_NUMBER_TOKEN_FINDER.findall(sequence)
if not tokens:
return None
@ -1328,7 +1383,7 @@ class ContentFilterGuardrail(CustomGuardrail):
HTTPException: If sensitive content is detected and action is BLOCK
"""
# Collect all exceptions from loaded categories
all_exceptions: Final = []
all_exceptions: Final[list[str]] = []
for category in self.loaded_categories.values():
all_exceptions.extend(category.exceptions)
@ -1404,7 +1459,7 @@ class ContentFilterGuardrail(CustomGuardrail):
if not (images and self.image_model and self.llm_router):
return
tasks: Final = []
tasks: Final[list[Coroutine[object, object, ModelResponse]]] = []
for image in images:
task = self.llm_router.acompletion(
model=self.image_model,
@ -1425,12 +1480,10 @@ class ContentFilterGuardrail(CustomGuardrail):
tasks.append(task)
responses: Final = await asyncio.gather(*tasks)
descriptions: Final = []
descriptions: Final[list[str]] = []
for response in responses:
choice = response.choices[0]
message = getattr(choice, "message", None)
if message and getattr(message, "content", None):
image_description = message.content
image_description = self._describe_image_response_content(response)
if image_description:
verbose_proxy_logger.debug("Image description: %s", image_description)
descriptions.append(image_description)
else:
@ -1447,7 +1500,7 @@ class ContentFilterGuardrail(CustomGuardrail):
except HTTPException as e:
# e.detail can be a string or dict
if isinstance(e.detail, dict) and "error" in e.detail:
detail_dict = cast(dict[str, Any], e.detail)
detail_dict = cast(dict[str, str], e.detail)
detail_dict["error"] = detail_dict["error"] + " (Image description): " + description
elif isinstance(e.detail, str):
e.detail = e.detail + " (Image description): " + description
@ -1455,6 +1508,14 @@ class ContentFilterGuardrail(CustomGuardrail):
e.detail = "Content blocked: Image description detected" + description
raise e
@staticmethod
def _describe_image_response_content(response: ModelResponse) -> str | None:
choice = response.choices[0]
message = getattr(choice, "message", None)
if message and getattr(message, "content", None):
return message.content
return None
def _count_masked_entities(
self,
detections: list[ContentFilterDetection],
@ -1484,12 +1545,12 @@ class ContentFilterGuardrail(CustomGuardrail):
category = category_detection["category"]
masked_entity_count[category] = masked_entity_count.get(category, 0) + 1
def _build_match_details(self, detections: list[ContentFilterDetection]) -> list[dict]:
def _build_match_details(self, detections: list[ContentFilterDetection]) -> list[dict[str, object]]:
"""Build match_details list from content filter detections."""
match_details: Final[list[dict]] = []
match_details: Final[list[dict[str, object]]] = []
for detection in detections:
action_taken = detection.get("action", detection.get("action_hint", ""))
detail: dict = {"type": detection["type"], "action_taken": action_taken}
detail: dict[str, object] = {"type": detection["type"], "action_taken": action_taken}
if detection["type"] == "pattern":
detail["detection_method"] = "regex"
detail["snippet"] = cast(PatternDetection, detection).get("pattern_name", "")
@ -1510,7 +1571,7 @@ class ContentFilterGuardrail(CustomGuardrail):
def _get_detection_methods(self, detections: list[ContentFilterDetection]) -> str:
"""Get comma-separated detection methods used."""
methods: Final[set] = set()
methods: Final[set[str]] = set()
for detection in detections:
if detection["type"] == "pattern":
methods.add("regex")
@ -1659,7 +1720,7 @@ class ContentFilterGuardrail(CustomGuardrail):
guardrail_json_response = exception_str if exception_str else [dict(detection) for detection in detections]
# Competitor intent: add confidence and classification to tracing if present
tracing_kw: Final[dict[str, Any]] = {
tracing_kw: Final[GuardrailTracingDetail] = {
"guardrail_id": self.config_guardrail_id or self.guardrail_name,
"policy_template": self.config_policy_template or self._get_policy_templates(),
"detection_method": (self._get_detection_methods(detections) if detections else None),

2
litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.py
View file

@ -55,7 +55,7 @@ for pattern_data in _PATTERNS_DATA["patterns"]:
PATTERN_EXTRA_CONFIG[pattern_data["name"]] = extra_config
def get_compiled_pattern(pattern_name: str) -> Pattern:
def get_compiled_pattern(pattern_name: str) -> Pattern[str]:
"""
Get a compiled regex pattern by name.

109
litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py
View file

@ -73,12 +73,14 @@ import hashlib
import os
import re
import time
from typing import Any, Final, Optional
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, Optional
import jwt
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey
from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey, RSAPublicKey
from typing_extensions import NotRequired, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.caching import DualCache
@ -90,13 +92,28 @@ from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.utils import CallTypesLiteral
if TYPE_CHECKING:
from jwt.types import Options
class _OIDCDiscoveryDocument(TypedDict, total=False):
jwks_uri: str
class _JWTDecodeKwargs(TypedDict):
algorithms: Sequence[str]
options: "Options"
audience: NotRequired[str]
issuer: NotRequired[str]
# Module-level singleton for the JWKS discovery endpoint to access.
_mcp_jwt_signer_instance: Optional["MCPJWTSigner"] = None
_MCP_JWT_CALL_TYPES: Final = frozenset({"call_mcp_tool", "list_mcp_tools"})
# Simple in-memory JWKS cache: keyed by JWKS URI → (keys_list, fetched_at).
_jwks_cache: Final[dict[str, tuple]] = {}
_jwks_cache: Final[dict[str, tuple[Sequence[Mapping[str, object]], float]]] = {}
_JWKS_CACHE_TTL: Final = 3600 # 1 hour
@ -133,7 +150,7 @@ def _int_to_base64url(n: int) -> str:
return base64.urlsafe_b64encode(n.to_bytes(byte_length, byteorder="big")).rstrip(b"=").decode("ascii")
def _compute_kid(public_key: Any) -> str:
def _compute_kid(public_key: RSAPublicKey) -> str:
"""Derive a key ID from the public key's DER encoding (SHA-256, first 16 hex chars)."""
der_bytes: Final = public_key.public_bytes(
encoding=serialization.Encoding.DER,
@ -142,7 +159,7 @@ def _compute_kid(public_key: Any) -> str:
return hashlib.sha256(der_bytes).hexdigest()[:16]
async def _fetch_jwks(jwks_uri: str) -> list[dict[str, Any]]:
async def _fetch_jwks(jwks_uri: str) -> Sequence[Mapping[str, object]]:
"""
Fetch and cache a JWKS from the given URI.
@ -163,12 +180,13 @@ async def _fetch_jwks(jwks_uri: str) -> list[dict[str, Any]]:
client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check)
resp: Final = await client.get(jwks_uri, headers={"Accept": "application/json"})
resp.raise_for_status()
keys = resp.json().get("keys", [])
_jwks_cache[jwks_uri] = (keys, now)
return keys
jwks_body: Final[Mapping[str, Sequence[Mapping[str, object]]]] = resp.json()
fetched_keys: Final = jwks_body.get("keys", [])
_jwks_cache[jwks_uri] = (fetched_keys, now)
return fetched_keys
async def _fetch_oidc_discovery(discovery_uri: str) -> dict[str, Any]:
async def _fetch_oidc_discovery(discovery_uri: str) -> _OIDCDiscoveryDocument:
"""Fetch an OIDC discovery document and return its parsed JSON."""
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
@ -178,7 +196,8 @@ async def _fetch_oidc_discovery(discovery_uri: str) -> dict[str, Any]:
client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check)
resp: Final = await client.get(discovery_uri, headers={"Accept": "application/json"})
resp.raise_for_status()
return resp.json()
document: Final[_OIDCDiscoveryDocument] = resp.json()
return document
class MCPJWTSigner(CustomGuardrail):
@ -230,8 +249,8 @@ class MCPJWTSigner(CustomGuardrail):
# FR-12: End-user identity mapping
end_user_claim_sources: list[str] | None = None,
# FR-13: Claim operations
add_claims: dict[str, Any] | None = None,
set_claims: dict[str, Any] | None = None,
add_claims: Mapping[str, object] | None = None,
set_claims: Mapping[str, object] | None = None,
remove_claims: list[str] | None = None,
# FR-14: Two-token model
channel_token_audience: str | None = None,
@ -283,7 +302,7 @@ class MCPJWTSigner(CustomGuardrail):
self.verify_issuer: str | None = verify_issuer
self.verify_audience: str | None = verify_audience
# Cached OIDC discovery document (fetched lazily, TTL = 24 h)
self._oidc_discovery_doc: dict[str, Any] | None = None
self._oidc_discovery_doc: _OIDCDiscoveryDocument | None = None
self._oidc_discovery_fetched_at: float = 0.0
# --- FR-12: End-user identity mapping ---
@ -294,8 +313,8 @@ class MCPJWTSigner(CustomGuardrail):
]
# --- FR-13: Claim operations ---
self.add_claims: dict[str, Any] = add_claims or {}
self.set_claims: dict[str, Any] = set_claims or {}
self.add_claims: Mapping[str, object] = add_claims or {}
self.set_claims: Mapping[str, object] = set_claims or {}
self.remove_claims: list[str] = remove_claims or []
# --- FR-14: Two-token model ---
@ -347,7 +366,7 @@ class MCPJWTSigner(CustomGuardrail):
"""
return 3600 if self._persistent_key else 300
def get_jwks(self) -> dict[str, Any]:
def get_jwks(self) -> Mapping[str, Sequence[Mapping[str, str]]]:
"""
Return the JWKS for the RSA public key.
Used by GET /.well-known/jwks.json so MCP servers can verify tokens.
@ -374,7 +393,7 @@ class MCPJWTSigner(CustomGuardrail):
# the IdP, short enough to pick up jwks_uri changes after key rotation.
_OIDC_DISCOVERY_TTL = 86400
async def _get_oidc_discovery(self) -> dict[str, Any]:
async def _get_oidc_discovery(self) -> _OIDCDiscoveryDocument:
"""Fetch and cache the OIDC discovery document with a 24-hour TTL.
Only caches when the doc contains a 'jwks_uri' so that a transient or
@ -391,7 +410,7 @@ class MCPJWTSigner(CustomGuardrail):
return doc
return self._oidc_discovery_doc or {}
async def _verify_incoming_jwt(self, raw_token: str) -> dict[str, Any]:
async def _verify_incoming_jwt(self, raw_token: str) -> dict[str, object]:
"""
Verify an incoming Bearer JWT against the configured IdP's JWKS.
@ -438,8 +457,8 @@ class MCPJWTSigner(CustomGuardrail):
# it infers from the key type (RSAPublicKey → RS256).
alg: Final = getattr(signing_jwk, "algorithm_name", None) or "RS256"
decode_options: Final[dict[str, Any]] = {"verify_exp": True}
decode_kwargs: Final[dict[str, Any]] = {
decode_options: Final[Options] = {"verify_exp": True}
decode_kwargs: Final[_JWTDecodeKwargs] = {
"algorithms": [alg],
"options": decode_options,
}
@ -451,10 +470,10 @@ class MCPJWTSigner(CustomGuardrail):
if self.verify_issuer:
decode_kwargs["issuer"] = self.verify_issuer
payload: Final[dict[str, Any]] = jwt.decode(raw_token, signing_jwk.key, **decode_kwargs)
payload: Final[dict[str, object]] = jwt.decode(raw_token, signing_jwk.key, **decode_kwargs)
return payload
async def _introspect_opaque_token(self, token: str) -> dict[str, Any]:
async def _introspect_opaque_token(self, token: str) -> dict[str, object]:
"""
Perform RFC 7662 token introspection for opaque (non-JWT) tokens.
@ -479,7 +498,7 @@ class MCPJWTSigner(CustomGuardrail):
headers={"Accept": "application/json"},
)
resp.raise_for_status()
result: Final[dict[str, Any]] = resp.json()
result: Final[dict[str, object]] = resp.json()
if not result.get("active", False):
raise jwt.exceptions.ExpiredSignatureError(
"MCPJWTSigner: incoming token is inactive (introspection returned active=false)"
@ -492,7 +511,7 @@ class MCPJWTSigner(CustomGuardrail):
def _validate_required_claims(
self,
jwt_claims: dict[str, Any] | None,
jwt_claims: Mapping[str, object] | None,
) -> None:
"""
Raise HTTP 403 if any required_claims are absent from the verified
@ -522,7 +541,7 @@ class MCPJWTSigner(CustomGuardrail):
def _resolve_end_user_identity(
self,
user_api_key_dict: UserAPIKeyAuth,
jwt_claims: dict[str, Any] | None,
jwt_claims: Mapping[str, object] | None,
) -> str:
"""
Resolve the outbound JWT 'sub' using the ordered end_user_claim_sources list.
@ -545,19 +564,19 @@ class MCPJWTSigner(CustomGuardrail):
value = str(raw) if raw else None
elif source == "litellm:user_id":
uid = getattr(user_api_key_dict, "user_id", None)
uid = user_api_key_dict.user_id
value = str(uid) if uid else None
elif source == "litellm:email":
email = getattr(user_api_key_dict, "user_email", None)
email = user_api_key_dict.user_email
value = str(email) if email else None
elif source == "litellm:end_user_id":
eid = getattr(user_api_key_dict, "end_user_id", None)
eid = user_api_key_dict.end_user_id
value = str(eid) if eid else None
elif source == "litellm:team_id":
tid = getattr(user_api_key_dict, "team_id", None)
tid = user_api_key_dict.team_id
value = str(tid) if tid else None
else:
@ -568,7 +587,7 @@ class MCPJWTSigner(CustomGuardrail):
return value
# Final fallback for service accounts with no user identity
token: Final = getattr(user_api_key_dict, "token", None) or getattr(user_api_key_dict, "api_key", None)
token: Final = user_api_key_dict.token or user_api_key_dict.api_key
if token:
return "apikey:" + hashlib.sha256(str(token).encode()).hexdigest()[:16]
return "litellm-proxy"
@ -615,7 +634,7 @@ class MCPJWTSigner(CustomGuardrail):
# FR-13: Claim operations
# ------------------------------------------------------------------
def _apply_claim_operations(self, claims: dict[str, Any]) -> dict[str, Any]:
def _apply_claim_operations(self, claims: dict[str, object]) -> dict[str, object]:
"""Apply add_claims, set_claims, and remove_claims to the claim dict."""
# add_claims: insert only when key is absent
for k, v in self.add_claims.items():
@ -637,9 +656,9 @@ class MCPJWTSigner(CustomGuardrail):
def _passthrough_optional_claims(
self,
claims: dict[str, Any],
jwt_claims: dict[str, Any] | None,
) -> dict[str, Any]:
claims: dict[str, object],
jwt_claims: Mapping[str, object] | None,
) -> dict[str, object]:
"""Forward optional_claims from verified incoming token into the outbound JWT."""
if not self.optional_claims or not jwt_claims:
return claims
@ -656,7 +675,7 @@ class MCPJWTSigner(CustomGuardrail):
self,
user_api_key_dict: UserAPIKeyAuth,
data: dict,
jwt_claims: dict[str, Any] | None = None,
jwt_claims: Mapping[str, object] | None = None,
call_type: CallTypesLiteral | None = None,
) -> dict[str, Any]:
"""
@ -669,7 +688,7 @@ class MCPJWTSigner(CustomGuardrail):
jwt_claims if available. None for pure API-key requests.
"""
now: Final = int(time.time())
claims: dict[str, Any] = {
claims: dict[str, object] = {
"iss": self.issuer,
"aud": self.audience,
"iat": now,
@ -681,18 +700,18 @@ class MCPJWTSigner(CustomGuardrail):
claims["sub"] = self._resolve_end_user_identity(user_api_key_dict, jwt_claims)
# email passthrough when available from LiteLLM context
user_email: Final = getattr(user_api_key_dict, "user_email", None)
user_email: Final = user_api_key_dict.user_email
if user_email:
claims["email"] = user_email
# act — RFC 8693 delegation claim (team/org context)
team_id: Final = getattr(user_api_key_dict, "team_id", None)
org_id: Final = getattr(user_api_key_dict, "org_id", None)
team_id: Final = user_api_key_dict.team_id
org_id: Final = user_api_key_dict.org_id
act_sub: Final = team_id or org_id or "litellm-proxy"
claims["act"] = {"sub": act_sub}
# end_user_id when set separately from user_id
end_user_id: Final = getattr(user_api_key_dict, "end_user_id", None)
end_user_id: Final = user_api_key_dict.end_user_id
if end_user_id:
claims["end_user_id"] = end_user_id
@ -710,8 +729,8 @@ class MCPJWTSigner(CustomGuardrail):
def _build_channel_token_claims(
self,
base_claims: dict[str, Any],
) -> dict[str, Any]:
base_claims: Mapping[str, object],
) -> dict[str, object]:
"""
Build claims for the channel token (FR-14 two-token model).
@ -776,7 +795,7 @@ class MCPJWTSigner(CustomGuardrail):
# ------------------------------------------------------------------
# FR-5: Verify incoming token before re-signing
# ------------------------------------------------------------------
jwt_claims: dict[str, Any] | None = None
jwt_claims: dict[str, object] | None = None
raw_token: Final[str | None] = hook_data.get("incoming_bearer_token")
if self.access_token_discovery_uri and raw_token:
@ -810,7 +829,7 @@ class MCPJWTSigner(CustomGuardrail):
# Fall back to LiteLLM-decoded JWT claims (available when proxy uses JWT auth).
if jwt_claims is None:
jwt_claims = getattr(user_api_key_dict, "jwt_claims", None)
jwt_claims = user_api_key_dict.jwt_claims
# ------------------------------------------------------------------
# FR-15: Validate required claims
@ -896,7 +915,7 @@ async def inject_mcp_jwt_headers_for_upstream(
if auth_hdr.lower().startswith("bearer "):
incoming_bearer_token = auth_hdr[len("bearer ") :]
hook_data: Final[dict[str, Any]] = {
hook_data: Final = {
"mcp_tool_name": "" if for_list_tools else mcp_tool_name,
"incoming_bearer_token": incoming_bearer_token,
"extra_headers": merged,

91
litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py
View file

@ -8,6 +8,7 @@ Provides real-time threat detection, DLP, URL filtering, content masking, and po
import json
import os
import re
from collections.abc import AsyncIterable, Mapping, Sequence
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final, Literal, Optional
from urllib.parse import urlparse
@ -166,7 +167,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
GuardrailEventHooks.during_mcp_call: GuardrailEventHooks.during_call,
}
def should_run_guardrail(self, data: Any, event_type: GuardrailEventHooks) -> bool:
def should_run_guardrail(self, data: Mapping[str, object], event_type: GuardrailEventHooks) -> bool:
if super().should_run_guardrail(data, event_type):
return True
compat: Final = self._MCP_COMPAT_MAP.get(event_type)
@ -175,7 +176,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
return True
return False
def _extract_text_from_messages(self, messages: list[dict[str, Any]]) -> str:
def _extract_text_from_messages(self, messages: Sequence[Mapping[str, object]]) -> str:
"""Extract text content from messages array."""
if not isinstance(messages, list) or not messages:
return ""
@ -242,10 +243,10 @@ class PanwPrismaAirsHandler(CustomGuardrail):
self,
content: str = "",
is_response: bool = False,
metadata: dict[str, Any] | None = None,
call_id: str | None = None,
metadata: Mapping[str, object] | None = None,
call_id: object = None,
tool_event: dict[str, Any] | None = None,
) -> dict[str, Any]:
) -> dict[str, object]:
"""Call PANW Prisma AIRS API to scan content or a tool_event."""
if tool_event is None and not content.strip():
@ -275,7 +276,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
else:
app_name_value = self.app_name # Defaults to "LiteLLM"
panw_metadata: Final = {
panw_metadata: Final[dict[str, object]] = {
"app_user": (
(metadata.get("app_user") or metadata.get("user") or "litellm_user") if metadata else "litellm_user"
),
@ -295,13 +296,13 @@ class PanwPrismaAirsHandler(CustomGuardrail):
panw_metadata["litellm_trace_id"] = metadata["litellm_trace_id"]
# Build contents: tool_event takes priority, else prompt/response text
contents: list[dict[str, Any]]
contents: Sequence[Mapping[str, object]]
if tool_event is not None:
contents = [{"tool_event": tool_event}]
else:
contents = [{"response" if is_response else "prompt": content}]
payload: Final = {
payload: Final[dict[str, object]] = {
"metadata": panw_metadata,
"contents": contents,
}
@ -325,7 +326,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
# If neither profile_name nor profile_id is provided, PANW API will use the
# profile linked to the API key (if configured in Strata Cloud Manager)
if profile_name or profile_id:
ai_profile: Final = {}
ai_profile: Final[dict[str, object]] = {}
if profile_id:
ai_profile["profile_id"] = profile_id
if profile_name:
@ -333,7 +334,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
payload["ai_profile"] = ai_profile
if is_response and tool_event is None:
payload["metadata"]["is_response"] = True
panw_metadata["is_response"] = True
headers: Final = {
"Content-Type": "application/json",
@ -355,7 +356,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
)
response.raise_for_status()
result: Final = response.json()
result: Final[dict[str, object]] = response.json()
# Validate response format
if "action" not in result:
@ -489,7 +490,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
)
return "unknown"
def _get_masked_text(self, scan_result: dict[str, Any], is_response: bool = False) -> str | None:
def _get_masked_text(self, scan_result: Mapping[str, object], is_response: bool = False) -> str | None:
"""Extract masked text from PANW scan result."""
masked_key: Final = "response_masked_data" if is_response else "prompt_masked_data"
masked_data: Final = scan_result.get(masked_key)
@ -511,7 +512,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
@staticmethod
def _apply_mcp_masking(
request_data: dict,
original_args: Any,
original_args: object,
masked_text: str,
*,
is_blocked: bool = True,
@ -544,7 +545,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
# If the original args were structured, preserve the type.
if isinstance(original_args, (dict, list)):
try:
parsed: Final = json.loads(masked_text)
parsed: Final[object] = json.loads(masked_text)
except (json.JSONDecodeError, TypeError):
raise HTTPException(
status_code=400,
@ -556,7 +557,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
}
},
)
masked_value: Any = parsed
masked_value: object = parsed
else:
masked_value = masked_text
@ -572,7 +573,9 @@ class PanwPrismaAirsHandler(CustomGuardrail):
else:
verbose_proxy_logger.info("PANW Prisma AIRS: MCP request allowed with PII masking applied")
def _apply_masking_to_messages(self, messages: list[dict[str, Any]], masked_text: str) -> list[dict[str, Any]]:
def _apply_masking_to_messages(
self, messages: list[dict[str, object]], masked_text: str
) -> Sequence[Mapping[str, object]]:
"""Apply masked text to the last user message."""
if not messages:
return messages
@ -622,7 +625,9 @@ class PanwPrismaAirsHandler(CustomGuardrail):
if hasattr(choice.message.function_call, "arguments"):
choice.message.function_call.arguments = masked_text
def _build_error_detail(self, scan_result: dict[str, Any], is_response: bool = False) -> dict[str, Any]:
def _build_error_detail(
self, scan_result: Mapping[str, object], is_response: bool = False
) -> Mapping[str, Mapping[str, object]]:
"""Build enhanced error detail with scan information."""
action_type: Final = "Response" if is_response else "Prompt"
code_suffix: Final = "_response_blocked" if is_response else "_blocked"
@ -642,7 +647,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
},
)
error_detail: Final = {
error_detail: Final[dict[str, dict[str, object]]] = {
"error": {
"message": error_msg,
"type": "guardrail_violation",
@ -672,12 +677,12 @@ class PanwPrismaAirsHandler(CustomGuardrail):
def _handle_api_error_with_logging(
self,
scan_result: dict[str, Any],
data: dict[str, Any],
scan_result: dict[str, object],
data: dict[str, object],
start_time: datetime,
event_type: GuardrailEventHooks,
is_response: bool = False,
) -> dict[str, Any] | None:
) -> None:
"""Handle API errors with fail-open/fail-closed logic."""
end_time: Final = datetime.now()
duration: Final = (end_time - start_time).total_seconds()
@ -722,7 +727,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
add_guardrail_to_applied_guardrails_header(
request_data=data, guardrail_name=f"{self.guardrail_name}:unscanned"
)
return None
return
raise HTTPException(
status_code=500,
@ -783,7 +788,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
return metadata
@staticmethod
def _extract_text_from_sse_bytes(chunks: list[bytes]) -> str:
def _extract_text_from_sse_bytes(chunks: Sequence[bytes]) -> str:
"""Extract text from Anthropic SSE byte chunks (content_block_delta → text_delta)."""
texts: Final[list[str]] = []
raw: Final = b"".join(chunks).decode("utf-8", errors="replace")
@ -804,7 +809,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
return "".join(texts)
@staticmethod
def _extract_text_from_streaming_events(chunks: list) -> str:
def _extract_text_from_streaming_events(chunks: Sequence[object]) -> str:
"""Extract text from /v1/responses streaming events (object or dict)."""
def _attr(c, key):
@ -960,7 +965,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
cache: DualCache,
data: dict[str, Any],
call_type: CallTypesLiteral,
) -> dict[str, Any] | None:
) -> dict[str, object] | None:
"""
Pre-call hook to scan user prompts before sending to LLM.
@ -1075,10 +1080,10 @@ class PanwPrismaAirsHandler(CustomGuardrail):
@log_guardrail_information
async def async_post_call_success_hook(
self,
data: dict[str, Any],
data: dict[str, object],
user_api_key_dict: UserAPIKeyAuth,
response: Any,
) -> Any:
response: object,
) -> object:
"""
Post-call hook to scan LLM responses before returning to user.
@ -1193,7 +1198,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
assembled_model_response: ModelResponse,
request_data: dict,
start_time: datetime,
) -> tuple[bool, ModelResponse, dict[str, Any]]:
) -> tuple[bool, ModelResponse, dict[str, object]]:
"""
Scan assembled streaming response and apply masking if needed.
Returns (content_was_modified, response, scan_result).
@ -1255,8 +1260,8 @@ class PanwPrismaAirsHandler(CustomGuardrail):
async def async_post_call_streaming_iterator_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
response: Any,
request_data: dict,
response: AsyncIterable[object],
request_data: dict[str, object],
):
"""
Process streaming response chunks and scan the assembled response.
@ -1367,7 +1372,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
# returns a proper JSON error response with the correct status code.
# (Raising from a generator hits create_response's generic except → 500.)
detail: Final = e.detail if isinstance(e.detail, dict) else {"message": str(e.detail)}
error_obj: Final[dict[str, Any]] = dict(detail.get("error", detail))
error_obj: Final[dict[str, object]] = dict(detail.get("error", detail))
error_obj["code"] = e.status_code
yield f"data: {json.dumps({'error': error_obj})}\n\n"
except Exception as e:
@ -1378,8 +1383,8 @@ class PanwPrismaAirsHandler(CustomGuardrail):
self,
tool_calls: list,
is_response: bool,
metadata: dict[str, Any],
call_id: str,
metadata: Mapping[str, object],
call_id: object,
request_data: dict,
start_time: datetime,
) -> None:
@ -1416,7 +1421,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
tool_name = func.get("name")
# --- build tool_event payload (canonical PANW schema) -----------
tool_event: dict[str, Any] = {
tool_event: dict[str, object] = {
"metadata": {
"ecosystem": "openai",
"method": "tools/call",
@ -1472,7 +1477,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
@staticmethod
def _is_anthropic_request(
request_data: dict,
request_data: Mapping[str, object],
logging_obj: Optional["LiteLLMLoggingObj"] = None,
) -> bool:
"""Detect if the current request is an Anthropic /v1/messages call."""
@ -1497,7 +1502,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
def _use_latest_user_only(
self,
request_data: dict,
request_data: Mapping[str, object],
logging_obj: Optional["LiteLLMLoggingObj"] = None,
) -> bool:
"""Resolve whether to scan only the latest user message.
@ -1515,8 +1520,8 @@ class PanwPrismaAirsHandler(CustomGuardrail):
@staticmethod
def _get_latest_user_text_indices(
texts: list[str],
messages: list,
texts: Sequence[str],
messages: Sequence[object],
) -> set | None:
"""Return text indices belonging to only the latest scannable human-authored (user or developer) message.
@ -1569,8 +1574,8 @@ class PanwPrismaAirsHandler(CustomGuardrail):
@staticmethod
def _get_scannable_text_indices(
texts: list[str],
structured_messages: list,
texts: Sequence[str],
structured_messages: Sequence[object],
) -> set | None:
"""Derive which ``texts`` indices originate from user/system messages.
@ -1627,7 +1632,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
request_data: dict[str, object],
input_type: Literal["request", "response"],
logging_obj: Optional["LiteLLMLoggingObj"] = None,
) -> GenericGuardrailAPIInputs:
@ -1798,7 +1803,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
# "mcp_tool_name"/"mcp_arguments". Check canonical first, then fallback.
mcp_tool_name: Final = request_data.get("mcp_tool_name") or self._mcp_name_fallback(request_data)
if mcp_tool_name and input_type == "request":
mcp_tool_event: Final[dict[str, Any]] = {
mcp_tool_event: Final[dict[str, object]] = {
"metadata": {
"ecosystem": "mcp",
"method": "tools/call",

3
litellm/proxy/hooks/mcp_semantic_filter/hook.py
View file

@ -5,6 +5,7 @@ Pre-call hook that filters MCP tools semantically before LLM inference.
Reduces context window size and improves tool selection accuracy.
"""
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, Optional
from fastapi import HTTPException
@ -164,7 +165,7 @@ class SemanticToolFilterHook(CustomLogger):
return [name for name in names if name]
@staticmethod
def _narrow_mcp_references(tools: list[Any], selected_tool_names: list[str]) -> list[Any]:
def _narrow_mcp_references(tools: Sequence[Mapping[str, object]], selected_tool_names: list[str]) -> list[object]:
"""
Restrict each litellm_proxy MCP reference to the semantically selected tools.

36
litellm/proxy/management_endpoints/cache_settings_endpoints.py
View file

@ -12,7 +12,7 @@ import asyncio
import json
from collections.abc import Mapping
from datetime import datetime, timezone
from typing import Any, Final
from typing import TYPE_CHECKING, Any, Final, Protocol
from fastapi import APIRouter, Depends, Header, HTTPException
from pydantic import BaseModel, Field
@ -37,8 +37,26 @@ from litellm.types.management_endpoints import (
CacheSettingsField,
)
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
router: Final = APIRouter()
class _CacheConfigRow(Protocol):
cache_settings: str | Mapping[str, object] | None
class _CacheConfigTable(Protocol):
async def find_unique(self, where: Mapping[str, str]) -> _CacheConfigRow | None: ...
async def upsert(self, where: Mapping[str, str], data: Mapping[str, Mapping[str, str]]) -> _CacheConfigRow: ...
def _cache_config_table(prisma_client: "PrismaClient") -> _CacheConfigTable:
return CacheConfigRepository(prisma_client).table
# Cache fields holding credentials. Masked on read so plaintext Redis /
# Sentinel passwords never leave the server in a GET response. `url` is here
# because a Redis/Valkey URL can embed a password inline
@ -197,7 +215,7 @@ def _saved_secret_is_reusable(incoming: Mapping[str, object], saved: Mapping[str
return True
def _merge_over_saved(incoming: Mapping[str, object], saved: Mapping[str, object]) -> dict[str, Any]:
def _merge_over_saved(incoming: Mapping[str, object], saved: Mapping[str, object]) -> Mapping[str, object]:
"""Keep the stored secret behind any credential the caller echoed back redacted or omitted.
GET returns credentials as the marker and the form never re-prefills a
@ -339,7 +357,7 @@ class CacheSettingsManager:
return normalized1 == normalized2
@staticmethod
async def init_cache_settings_in_db(prisma_client, proxy_config):
async def init_cache_settings_in_db(prisma_client: "PrismaClient", proxy_config):
"""
Initialize cache settings from database into the router on startup.
Only reinitializes if cache params have changed.
@ -349,7 +367,7 @@ class CacheSettingsManager:
try:
cache_config: Final = await call_with_db_reconnect_retry(
prisma_client,
lambda: CacheConfigRepository(prisma_client).table.find_unique(where={"id": "cache_config"}),
lambda: _cache_config_table(prisma_client).find_unique(where={"id": "cache_config"}),
reason="init_cache_settings_in_db_lookup_failure",
)
if cache_config is not None and cache_config.cache_settings:
@ -444,7 +462,7 @@ async def get_cache_settings(
# Read the stored settings (decrypted); an env-only cache has none.
stored: dict[str, object] = {}
if prisma_client is not None:
cache_config = await CacheConfigRepository(prisma_client).table.find_unique(where={"id": "cache_config"})
cache_config = await _cache_config_table(prisma_client).find_unique(where={"id": "cache_config"})
if cache_config is not None and cache_config.cache_settings:
stored = proxy_config._decrypt_db_variables(
variables_dict=_parse_stored_settings(cache_config.cache_settings)
@ -511,9 +529,7 @@ async def test_cache_connection(
saved_settings: dict[str, object] = {}
if prisma_client is not None:
try:
existing_row: Final = await CacheConfigRepository(prisma_client).table.find_unique(
where={"id": "cache_config"}
)
existing_row: Final = await _cache_config_table(prisma_client).find_unique(where={"id": "cache_config"})
if existing_row is not None and existing_row.cache_settings:
saved_settings = proxy_config._decrypt_db_variables(
variables_dict=_parse_stored_settings(existing_row.cache_settings)
@ -590,7 +606,7 @@ async def update_cache_settings(
try:
# Read the stored row first: its decrypted values back any credential the
# caller echoed back redacted, and its key set drives the audit diff.
existing_row: Final = await CacheConfigRepository(prisma_client).table.find_unique(where={"id": "cache_config"})
existing_row: Final = await _cache_config_table(prisma_client).find_unique(where={"id": "cache_config"})
before_settings: dict[str, object] | None = None
saved_settings: dict[str, object] = {}
if existing_row is not None and existing_row.cache_settings:
@ -606,7 +622,7 @@ async def update_cache_settings(
encrypted_settings: Final = proxy_config._encrypt_env_variables(environment_variables=cache_settings)
# Save to database
await CacheConfigRepository(prisma_client).table.upsert(
await _cache_config_table(prisma_client).upsert(
where={"id": "cache_config"},
data={
"create": {

94
litellm/proxy/memory/memory_endpoints.py
View file

@ -18,20 +18,21 @@ Scoping:
"""
import json
from collections.abc import Mapping
from typing import TYPE_CHECKING, Final
from collections.abc import Mapping, Sequence
from datetime import datetime
from typing import TYPE_CHECKING, Final, Protocol
from fastapi import APIRouter, Depends, HTTPException, Query
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import (
CommonProxyErrors,
LiteLLM_TeamTable,
LitellmUserRoles,
UserAPIKeyAuth,
user_api_key_has_admin_view,
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.repositories.prisma_protocols import MemoryRecord
from litellm.repositories.table_repositories import MemoryRepository
from litellm.repositories.team_repository import TeamRepository
from litellm.types.memory_management import (
@ -48,11 +49,54 @@ if TYPE_CHECKING:
router: Final = APIRouter()
class _MemoryRecord(Protocol):
memory_id: str
key: str
value: str
metadata: object
user_id: str | None
team_id: str | None
created_at: datetime | None
created_by: str | None
updated_at: datetime | None
updated_by: str | None
class _MemoryTableActions(Protocol):
async def create(self, data: Mapping[str, object]) -> _MemoryRecord: ...
async def find_many(
self,
where: Mapping[str, object] | None = ...,
order: Mapping[str, str] | None = ...,
skip: int = ...,
take: int = ...,
) -> Sequence[_MemoryRecord]: ...
async def count(self, where: Mapping[str, object] | None = ...) -> int: ...
async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> _MemoryRecord: ...
async def delete(self, where: Mapping[str, object]) -> _MemoryRecord | None: ...
def _memory_table(prisma_client: "PrismaClient") -> _MemoryTableActions:
return MemoryRepository(prisma_client).table
class _TeamTableActions(Protocol):
async def find_unique(self, where: Mapping[str, str]) -> LiteLLM_TeamTable | None: ...
def _team_table(prisma_client: "PrismaClient") -> _TeamTableActions:
return TeamRepository(prisma_client).table
def _serialize_metadata_for_prisma(metadata: object) -> str:
"""
Encode a `metadata` payload for the `Json?` column.
`metadata` is typed `object | None`, so callers may send dicts, lists,
`metadata` is typed `Optional[Any]`, so callers may send dicts, lists,
or JSON scalars (including plain Python strings like `"hello"`).
prisma-client-python rejects raw Python values on `Json?` columns
(`MissingRequiredValueError` / `DataError`), and Postgres `jsonb`
@ -74,18 +118,18 @@ def _visibility_filter(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, object
"""
if user_api_key_has_admin_view(user_api_key_dict):
return None
ors: Final[list[dict[str, object]]] = []
if user_api_key_dict.user_id:
ors.append({"user_id": user_api_key_dict.user_id})
if user_api_key_dict.team_id:
ors.append({"team_id": user_api_key_dict.team_id})
ors: Final = [
{field: value}
for field, value in (("user_id", user_api_key_dict.user_id), ("team_id", user_api_key_dict.team_id))
if value
]
if not ors:
# Caller has neither user_id nor team_id — match nothing.
return {"memory_id": "__no_match__"}
return {"OR": ors}
def _row_to_model(row: MemoryRecord) -> LiteLLM_MemoryRow:
def _row_to_model(row: _MemoryRecord) -> LiteLLM_MemoryRow:
return LiteLLM_MemoryRow(
memory_id=row.memory_id,
key=row.key,
@ -119,7 +163,7 @@ def _internal_error(log_message: str, exc: Exception, default_detail: str) -> HT
async def _assert_write_access(
prisma_client: "PrismaClient", row: MemoryRecord, user_api_key_dict: UserAPIKeyAuth
prisma_client: "PrismaClient", row: _MemoryRecord, user_api_key_dict: UserAPIKeyAuth
) -> None:
"""
Enforce ownership for mutations (PUT/DELETE).
@ -142,8 +186,8 @@ async def _assert_write_access(
"""
if _is_admin(user_api_key_dict):
return
row_user_id: Final = row.user_id
row_team_id: Final = row.team_id
row_user_id: Final = getattr(row, "user_id", None)
row_team_id: Final = getattr(row, "team_id", None)
# Personal ownership.
if row_user_id and row_user_id == user_api_key_dict.user_id:
@ -175,7 +219,7 @@ async def _is_team_admin_for(prisma_client: "PrismaClient", user_api_key_dict: U
)
try:
team_obj: Final = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id})
team_obj: Final = await _team_table(prisma_client).find_unique(where={"team_id": team_id})
except Exception as e:
verbose_proxy_logger.exception("Error loading team for write-auth check (team_id=%s): %s", team_id, e)
return False
@ -288,7 +332,7 @@ async def create_memory(
create_data["metadata"] = _serialize_metadata_for_prisma(body.metadata)
try:
row: Final = await MemoryRepository(prisma_client).rows.create(data=create_data)
row: Final = await _memory_table(prisma_client).create(data=create_data)
except Exception as e:
# Key is globally unique. Any duplicate → 409.
if _is_unique_violation(e):
@ -348,8 +392,8 @@ async def list_memory(
where = {"AND": [key_filter, vis]}
try:
total: Final = await MemoryRepository(prisma_client).rows.count(where=where)
rows: Final = await MemoryRepository(prisma_client).rows.find_many(
total: Final = await _memory_table(prisma_client).count(where=where)
rows: Final = await _memory_table(prisma_client).find_many(
where=where,
order={"updated_at": "desc"},
skip=(page - 1) * page_size,
@ -363,14 +407,12 @@ async def list_memory(
async def _find_memory_for_caller(
prisma_client: "PrismaClient", key: str, user_api_key_dict: UserAPIKeyAuth
) -> MemoryRecord:
) -> _MemoryRecord:
"""Look up a memory row by key, scoped to the caller's visibility."""
key_filter: Final[Mapping[str, object]] = {"key": key}
vis: Final = _visibility_filter(user_api_key_dict)
where: Final[Mapping[str, object]] = key_filter if vis is None else {"AND": [key_filter, vis]}
rows: Final = await MemoryRepository(prisma_client).rows.find_many(
where=where, take=1, order={"updated_at": "desc"}
)
rows = await _memory_table(prisma_client).find_many(where=where, take=1, order={"updated_at": "desc"})
if not rows:
raise HTTPException(status_code=404, detail=f"Memory with key '{key}' not found")
return rows[0]
@ -438,7 +480,7 @@ async def upsert_memory(
)
data["updated_by"] = user_api_key_dict.user_id
async def _find_existing() -> MemoryRecord | None:
async def _find_existing() -> _MemoryRecord | None:
"""Return the caller-visible row for `key`, or None."""
try:
return await _find_memory_for_caller(prisma_client, key, user_api_key_dict)
@ -455,7 +497,7 @@ async def upsert_memory(
# their team) — otherwise a teammate could overwrite a personal
# entry through the OR-based visibility filter.
await _assert_write_access(prisma_client, existing, user_api_key_dict)
row = await MemoryRepository(prisma_client).rows.update(
row = await _memory_table(prisma_client).update(
where={"memory_id": existing.memory_id},
data=data,
)
@ -481,7 +523,7 @@ async def upsert_memory(
if body.metadata is not None:
create_data["metadata"] = _serialize_metadata_for_prisma(body.metadata)
try:
row = await MemoryRepository(prisma_client).rows.create(data=create_data)
row = await _memory_table(prisma_client).create(data=create_data)
except Exception as e:
# Race: a concurrent PUT/POST created the row after our check.
# Re-read and fall back to an update so the PUT stays idempotent
@ -498,7 +540,7 @@ async def upsert_memory(
)
# Same write-authorization check as the non-race path.
await _assert_write_access(prisma_client, existing_after_race, user_api_key_dict)
row = await MemoryRepository(prisma_client).rows.update(
row = await _memory_table(prisma_client).update(
where={"memory_id": existing_after_race.memory_id},
data=data,
)
@ -526,7 +568,7 @@ async def delete_memory(
# Visibility != write authority — see the upsert handler for the rationale.
await _assert_write_access(prisma_client, row, user_api_key_dict)
try:
await MemoryRepository(prisma_client).rows.delete(where={"memory_id": row.memory_id})
await _memory_table(prisma_client).delete(where={"memory_id": row.memory_id})
except Exception as e:
raise _internal_error("Error deleting memory: %s", e, "Internal error deleting memory entry.")

103
litellm/rag/ingestion/s3_vectors_ingestion.py
View file

@ -17,7 +17,8 @@ from __future__ import annotations
import hashlib
import uuid
from typing import TYPE_CHECKING, Any, Final
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, TypedDict
import litellm
from litellm._logging import verbose_logger
@ -35,10 +36,32 @@ from litellm.llms.custom_httpx.http_handler import (
from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion
if TYPE_CHECKING:
import httpx
from litellm import Router
from litellm.types.rag import RAGIngestOptions
class S3VectorDataPayload(TypedDict):
float32: Sequence[float]
class S3VectorEntry(TypedDict):
key: str
data: S3VectorDataPayload
metadata: Mapping[str, str]
class S3VectorsQueryMatch(TypedDict, total=False):
key: str
distance: float
metadata: Mapping[str, str]
class S3VectorsQueryResponse(TypedDict, total=False):
vectors: Sequence[S3VectorsQueryMatch]
class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM):
"""
S3 Vectors RAG ingestion using httpx + AWS SigV4 signing.
@ -66,10 +89,10 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM):
BaseAWSLLM.__init__(self)
# Extract config
self.vector_bucket_name = self.vector_store_config["vector_bucket_name"]
self.index_name = self.vector_store_config.get("index_name")
self.distance_metric = self.vector_store_config.get("distance_metric", S3_VECTORS_DEFAULT_DISTANCE_METRIC)
self.non_filterable_metadata_keys = self.vector_store_config.get(
self.vector_bucket_name: str = self.vector_store_config["vector_bucket_name"]
self.index_name: str | None = self.vector_store_config.get("index_name")
self.distance_metric: str = self.vector_store_config.get("distance_metric", S3_VECTORS_DEFAULT_DISTANCE_METRIC)
self.non_filterable_metadata_keys: Sequence[str] = self.vector_store_config.get(
"non_filterable_metadata_keys",
S3_VECTORS_DEFAULT_NON_FILTERABLE_METADATA_KEYS,
)
@ -78,7 +101,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM):
self.dimension = self._get_dimension_from_config()
# Get AWS region using BaseAWSLLM method
_aws_region: Final = self.vector_store_config.get("aws_region_name")
_aws_region: Final[str | None] = self.vector_store_config.get("aws_region_name")
self.aws_region_name = self.get_aws_region_name_for_non_llm_api_calls(
aws_region_name=str(_aws_region) if _aws_region else None
)
@ -135,7 +158,8 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM):
Returns None if dimension should be auto-detected.
"""
if "dimension" in self.vector_store_config:
return int(self.vector_store_config["dimension"])
configured_dimension: Final[int] = self.vector_store_config["dimension"]
return int(configured_dimension)
return None
async def _ensure_config_initialized(self):
@ -258,7 +282,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM):
get_body: Final = safe_dumps({"vectorBucketName": self.vector_bucket_name})
try:
response = await self._sign_and_execute_request("POST", get_url, data=get_body)
response: httpx.Response = await self._sign_and_execute_request("POST", get_url, data=get_body)
if response.status_code == 200:
verbose_logger.debug("Vector bucket %s exists", self.vector_bucket_name)
return
@ -294,7 +318,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM):
get_body: Final = safe_dumps({"vectorBucketName": self.vector_bucket_name, "indexName": self.index_name})
try:
response = await self._sign_and_execute_request("POST", get_url, data=get_body)
response: httpx.Response = await self._sign_and_execute_request("POST", get_url, data=get_body)
if response.status_code == 200:
verbose_logger.debug("Vector index %s exists", self.index_name)
return
@ -311,7 +335,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM):
)
# Prepare index configuration per AWS API docs
index_config: Final = {
index_config: Final[dict[str, object]] = {
"vectorBucketName": self.vector_bucket_name,
"indexName": self.index_name,
"dataType": "float32",
@ -336,7 +360,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM):
verbose_logger.exception("Error creating vector index: %s", e)
raise
async def _put_vectors(self, vectors: list[dict[str, Any]]):
async def _put_vectors(self, vectors: Sequence[S3VectorEntry]):
"""
Call PutVectors API to store vectors in S3 Vectors.
@ -355,7 +379,9 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM):
}
try:
response: Final = await self._sign_and_execute_request("POST", url, data=safe_dumps(request_body))
response: Final[httpx.Response] = await self._sign_and_execute_request(
"POST", url, data=safe_dumps(request_body)
)
if response.status_code in (200, 201):
verbose_logger.info("Successfully stored %s vectors in index %s", len(vectors), self.index_name)
@ -442,24 +468,18 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM):
raise ValueError(error_msg)
# Prepare vectors for PutVectors API
vectors: Final = []
for i, (chunk, embedding) in enumerate(zip(chunks, embeddings)):
# Build metadata dict
metadata: dict[str, str] = {
"source_text": chunk, # Non-filterable (for reference)
"chunk_index": str(i), # Filterable
}
if filename:
metadata["filename"] = filename # Filterable
vector_obj = {
"key": f"{filename}_{i}" if filename else f"chunk_{i}",
"data": {"float32": embedding},
"metadata": metadata,
}
vectors.append(vector_obj)
vectors: Final = [
S3VectorEntry(
key=f"{filename}_{i}" if filename else f"chunk_{i}",
data=S3VectorDataPayload(float32=embedding),
metadata=(
{"source_text": chunk, "chunk_index": str(i), "filename": filename}
if filename
else {"source_text": chunk, "chunk_index": str(i)}
),
)
for i, (chunk, embedding) in enumerate(zip(chunks, embeddings))
]
# Call PutVectors API
await self._put_vectors(vectors)
@ -468,7 +488,9 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM):
vector_store_id: Final = f"{self.vector_bucket_name}:{self.index_name}"
return vector_store_id, filename
async def query_vector_store(self, vector_store_id: str, query: str, top_k: int = 5) -> dict[str, Any] | None:
async def query_vector_store(
self, vector_store_id: str, query: str, top_k: int = 5
) -> S3VectorsQueryResponse | None:
"""
Query S3 Vectors using QueryVectors API.
@ -489,7 +511,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM):
embedding_model: Final = self.embedding_config.get("model", "text-embedding-3-small")
response = await litellm.aembedding(model=embedding_model, input=[query])
query_embedding: Final = response.data[0]["embedding"]
query_embedding: Final[Sequence[float]] = response.data[0]["embedding"]
# Call QueryVectors API
url: Final = f"https://s3vectors.{self.aws_region_name}.api.aws/QueryVectors"
@ -504,15 +526,18 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM):
}
try:
response = await self._sign_and_execute_request("POST", url, data=safe_dumps(request_body))
query_response: Final[httpx.Response] = await self._sign_and_execute_request(
"POST", url, data=safe_dumps(request_body)
)
if response.status_code == 200:
results: Final = response.json()
if query_response.status_code == 200:
results: Final[S3VectorsQueryResponse] = query_response.json()
matches: Final = results.get("vectors")
verbose_logger.debug("Query returned %s results", len(results.get("vectors", [])))
# Check if query terms appear in results
if results.get("vectors"):
for result in results["vectors"]:
if matches:
for result in matches:
metadata = result.get("metadata", {})
source_text = metadata.get("source_text", "")
if query.lower() in source_text.lower():
@ -521,7 +546,9 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM):
# Return results even if exact match not found
return results
else:
verbose_logger.error("QueryVectors failed with status %s: %s", response.status_code, response.text)
verbose_logger.error(
"QueryVectors failed with status %s: %s", query_response.status_code, query_response.text
)
return None
except Exception as e:
verbose_logger.exception("Error querying vectors: %s", e)

194
litellm/responses/mcp/litellm_proxy_mcp_handler.py
View file

@ -2,7 +2,10 @@ import re
import traceback
from collections.abc import Iterable, Mapping, Sequence
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final, Literal, Optional
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypedDict, overload
from openai.types.chat import ChatCompletionToolParam
from openai.types.responses.function_tool_param import FunctionToolParam
from litellm._logging import verbose_logger
from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG
@ -18,6 +21,7 @@ from litellm.types.llms.openai import (
ResponsesAPIResponse,
ResponsesAPIStreamingResponse,
)
from litellm.types.llms.openai import ToolParam as ResponsesToolParam
from litellm.types.utils import (
CallTypes,
Choices,
@ -36,10 +40,14 @@ else:
MCPTool = Any
# NOTE: We intentionally keep ToolParam as a broad type here to avoid tight coupling
# to optional OpenAI SDK typing symbols in environments that may not have them available.
# `Any` is used to keep mypy compatible with the broader OpenAI tool union types
# passed around in Responses API while still allowing dict-style access at runtime.
ToolParam = Any
ToolParam: TypeAlias = Mapping[str, object]
class MCPToolResult(TypedDict):
tool_call_id: str | None
result: str
name: str | None
LITELLM_PROXY_MCP_SERVER_URL: Final = "litellm_proxy"
LITELLM_PROXY_MCP_SERVER_URL_PREFIX: Final = f"{LITELLM_PROXY_MCP_SERVER_URL}/mcp/"
@ -199,13 +207,12 @@ class LiteLLM_Proxy_MCP_Handler:
_get_tools_from_mcp_servers,
)
mcp_servers: Final[list[str]] = []
if mcp_tools_with_litellm_proxy:
for _tool in mcp_tools_with_litellm_proxy:
# if user specifies servers as server_url: litellm_proxy/mcp/zapier,github then return zapier,github
server_url = _tool.get("server_url", "") if isinstance(_tool, dict) else ""
if isinstance(server_url, str) and server_url.startswith(LITELLM_PROXY_MCP_SERVER_URL_PREFIX):
mcp_servers.append(server_url.split("/")[-1])
mcp_servers: Final = [
server_url.split("/")[-1]
for _tool in (mcp_tools_with_litellm_proxy or ())
for server_url in (_tool.get("server_url", "") if isinstance(_tool, dict) else "",)
if isinstance(server_url, str) and server_url.startswith(LITELLM_PROXY_MCP_SERVER_URL_PREFIX)
]
# Resolve toolset names: collect all toolset IDs first, then apply their
# combined permissions in a single pass so multiple toolsets are unioned
@ -279,15 +286,15 @@ class LiteLLM_Proxy_MCP_Handler:
allowed_mcp_servers=allowed_mcp_servers,
)
server_names: Final[list[str]] = []
for server in allowed_mcp_servers:
if server is None:
continue
server_name = (
getattr(server, "server_name", None) or getattr(server, "alias", None) or getattr(server, "name", None)
server_names: Final = [
server_name
for server in allowed_mcp_servers
if server is not None
for server_name in (
getattr(server, "server_name", None) or getattr(server, "alias", None) or getattr(server, "name", None),
)
if isinstance(server_name, str):
server_names.append(server_name)
if isinstance(server_name, str)
]
return tools, server_names
@ -305,8 +312,8 @@ class LiteLLM_Proxy_MCP_Handler:
List of deduplicated MCP tools
The returned dictionary maps each tool_name to the server_name
"""
seen_names: Final = set()
deduplicated_tools: Final = []
seen_names: Final[set[str]] = set()
deduplicated_tools: Final[list[MCPTool]] = []
tool_server_map: Final[dict[str, str]] = {}
for tool in mcp_tools:
@ -331,7 +338,7 @@ class LiteLLM_Proxy_MCP_Handler:
) -> list[MCPTool]:
"""Filter MCP tools based on allowed_tools parameter from the original tool configs."""
# Collect all allowed tool names from all MCP tool configs
allowed_tool_names: Final = set()
allowed_tool_names: Final[set[str]] = set()
for tool_config in mcp_tools_with_litellm_proxy:
if isinstance(tool_config, dict) and "allowed_tools" in tool_config:
allowed_tools = tool_config.get("allowed_tools", [])
@ -343,23 +350,13 @@ class LiteLLM_Proxy_MCP_Handler:
return mcp_tools
# Filter tools based on allowed names
filtered_tools: Final = []
for mcp_tool in mcp_tools:
if isinstance(mcp_tool, dict):
tool_name = mcp_tool.get("name")
else:
tool_name = getattr(mcp_tool, "name", None)
if not tool_name:
continue
if tool_name in allowed_tool_names:
filtered_tools.append(mcp_tool)
continue
unprefixed_name, _ = split_server_prefix_from_name(tool_name)
if unprefixed_name in allowed_tool_names:
filtered_tools.append(mcp_tool)
filtered_tools: Final = [
mcp_tool
for mcp_tool in mcp_tools
for tool_name in (mcp_tool.get("name") if isinstance(mcp_tool, dict) else getattr(mcp_tool, "name", None),)
if tool_name
and (tool_name in allowed_tool_names or split_server_prefix_from_name(tool_name)[0] in allowed_tool_names)
]
return filtered_tools
@ -448,24 +445,37 @@ class LiteLLM_Proxy_MCP_Handler:
return deduplicated_mcp_tools, tool_server_map
@overload
@staticmethod
def _transform_mcp_tools_to_openai(
mcp_tools: Sequence[MCPTool],
target_format: Literal["responses"] = ...,
) -> list[FunctionToolParam]: ...
@overload
@staticmethod
def _transform_mcp_tools_to_openai(
mcp_tools: Sequence[MCPTool],
target_format: Literal["chat"],
) -> list[ChatCompletionToolParam]: ...
@staticmethod
def _transform_mcp_tools_to_openai(
mcp_tools: Sequence[MCPTool],
target_format: Literal["responses", "chat"] = "responses",
) -> list[Any]:
) -> Sequence[FunctionToolParam | ChatCompletionToolParam]:
"""Transform MCP tools to OpenAI-compatible format."""
from litellm.experimental_mcp_client.tools import (
transform_mcp_tool_to_openai_responses_api_tool,
transform_mcp_tool_to_openai_tool,
)
openai_tools: Final[list[Any]] = []
for mcp_tool in mcp_tools:
if target_format == "chat":
openai_tool = transform_mcp_tool_to_openai_tool(mcp_tool)
else:
openai_tool = transform_mcp_tool_to_openai_responses_api_tool(mcp_tool)
openai_tools.append(openai_tool)
openai_tools: Final = [
transform_mcp_tool_to_openai_tool(mcp_tool)
if target_format == "chat"
else transform_mcp_tool_to_openai_responses_api_tool(mcp_tool)
for mcp_tool in mcp_tools
]
return openai_tools
@ -496,9 +506,9 @@ class LiteLLM_Proxy_MCP_Handler:
return True
@staticmethod
def _extract_tool_calls_from_response(response: ResponsesAPIResponse) -> list[Any]:
def _extract_tool_calls_from_response(response: ResponsesAPIResponse) -> list[object]:
"""Extract tool calls from the response output."""
tool_calls: Final[list[Any]] = []
tool_calls: Final[list[object]] = []
for output_item in response.output:
# Check if this is a function call output item
if isinstance(output_item, dict) and output_item.get("type") == "function_call":
@ -533,7 +543,7 @@ class LiteLLM_Proxy_MCP_Handler:
@staticmethod
def _extract_tool_call_details(
tool_call,
tool_call: object,
) -> tuple[str | None, str | None, str | None]:
"""Extract tool name, arguments, and call_id from a tool call."""
if isinstance(tool_call, dict):
@ -566,7 +576,7 @@ class LiteLLM_Proxy_MCP_Handler:
return tool_name, tool_arguments, tool_call_id
@staticmethod
def _parse_tool_arguments(tool_arguments: Any) -> dict[str, Any]:
def _parse_tool_arguments(tool_arguments: str | None) -> dict[str, object]:
"""Parse tool arguments, handling both string and dict formats."""
import json
@ -591,23 +601,18 @@ class LiteLLM_Proxy_MCP_Handler:
# Fallback to generic handling if MCP types not available
return "Tool executed successfully"
text_parts: Final = []
other_content_types: Final = []
for content_item in result.content:
if isinstance(content_item, TextContent):
# Text content - extract the text
text_parts.append(str(content_item.text))
elif isinstance(content_item, ImageContent):
# Image content
other_content_types.append("Image")
elif isinstance(content_item, EmbeddedResource):
# Embedded resource
other_content_types.append("EmbeddedResource")
else:
# Other unknown content types
content_type = type(content_item).__name__
other_content_types.append(content_type)
text_parts: Final = [
str(content_item.text) for content_item in result.content if isinstance(content_item, TextContent)
]
other_content_types: Final = [
"Image"
if isinstance(content_item, ImageContent)
else "EmbeddedResource"
if isinstance(content_item, EmbeddedResource)
else type(content_item).__name__
for content_item in result.content
if not isinstance(content_item, TextContent)
]
# Combine text parts if any
result_text = " ".join(text_parts) if text_parts else ""
@ -631,7 +636,7 @@ class LiteLLM_Proxy_MCP_Handler:
litellm_call_id: str | None = None,
litellm_trace_id: str | None = None,
request_tags: list[str] | None = None,
) -> list[dict[str, Any]]:
) -> list[MCPToolResult]:
"""Execute tool calls and return results."""
from fastapi import HTTPException
@ -645,11 +650,11 @@ class LiteLLM_Proxy_MCP_Handler:
)
from litellm.proxy.proxy_server import proxy_logging_obj
tool_results: Final = []
tool_results: Final[list[MCPToolResult]] = []
tool_call_id: str | None = None
rules_obj: Final = Rules()
for tool_call in tool_calls:
logging_request_data: dict[str, Any] = {}
logging_request_data: dict[str, object] = {}
tool_name: str | None = None
try:
(
@ -678,7 +683,7 @@ class LiteLLM_Proxy_MCP_Handler:
sanitized_tool_name = strip_known_server_prefix(resolved_tool_name, mcp_server)
start_time = datetime.now()
logging_input = [
logging_input: Sequence[Mapping[str, object]] = [
{
"role": "tool",
"content": {
@ -688,13 +693,14 @@ class LiteLLM_Proxy_MCP_Handler:
}
]
tool_logging_call_id = litellm_call_id or str(uuid.uuid4())
logging_metadata: dict[str, object] = {
"tool_call_id": tool_call_id,
"tool_name": sanitized_tool_name,
"server_name": server_name,
}
logging_request_data = {
"model": f"MCP: {tool_name}",
"metadata": {
"tool_call_id": tool_call_id,
"tool_name": sanitized_tool_name,
"server_name": server_name,
},
"metadata": logging_metadata,
"input": logging_input,
"call_type": CallTypes.call_mcp_tool.value,
"litellm_call_id": tool_logging_call_id,
@ -712,7 +718,7 @@ class LiteLLM_Proxy_MCP_Handler:
if litellm_trace_id:
logging_request_data["litellm_trace_id"] = litellm_trace_id
if request_tags:
logging_request_data["metadata"]["tags"] = request_tags
logging_metadata["tags"] = request_tags
if user_api_key_auth is not None:
from litellm.proxy.litellm_pre_call_utils import (
LiteLLMProxyRequestSetup,
@ -902,16 +908,16 @@ class LiteLLM_Proxy_MCP_Handler:
@staticmethod
def _create_follow_up_messages_for_chat(
original_messages: list[Any],
original_messages: list[object],
response: ModelResponse,
tool_results: Sequence[Mapping[str, object]],
) -> list[Any]:
) -> Sequence[Mapping[str, object]]:
"""Create follow-up chat messages that include tool execution results."""
from copy import deepcopy
from litellm.utils import convert_list_message_to_dict
follow_up_messages: list[Any] = convert_list_message_to_dict(deepcopy(original_messages))
follow_up_messages: list[dict[str, object]] = convert_list_message_to_dict(deepcopy(original_messages))
if not follow_up_messages:
follow_up_messages = []
@ -950,9 +956,9 @@ class LiteLLM_Proxy_MCP_Handler:
response: ResponsesAPIResponse,
tool_results: Sequence[Mapping[str, object]],
original_input: str | ResponseInputParam | None = None,
) -> list[Any]:
) -> list[object]:
"""Create follow-up input with tool results in proper format."""
follow_up_input: Final[list[Any]] = []
follow_up_input: Final[list[object]] = []
# Add original user input if available to maintain conversation context
if original_input:
@ -964,8 +970,8 @@ class LiteLLM_Proxy_MCP_Handler:
follow_up_input.append(original_input)
# Add the assistant message with function calls
assistant_message_content: Final[list[Any]] = []
function_calls: Final[list[dict[str, Any]]] = []
assistant_message_content: Final[list[object]] = []
function_calls: Final[list[dict[str, object]]] = []
for output_item in response.output:
if not isinstance(output_item, dict) and hasattr(output_item, "model_dump"):
@ -1027,7 +1033,7 @@ class LiteLLM_Proxy_MCP_Handler:
async def _make_follow_up_call(
follow_up_input: list[Any],
model: str,
all_tools: list[Any] | None,
all_tools: Sequence[ResponsesToolParam] | None,
response_id: str,
**call_params: Any,
) -> ResponsesAPIResponse | BaseResponsesAPIStreamingIterator:
@ -1044,7 +1050,7 @@ class LiteLLM_Proxy_MCP_Handler:
async def _log_mcp_tool_failure(
*,
proxy_logging_obj: Optional["ProxyLogging"],
user_api_key_auth: Any,
user_api_key_auth: "UserAPIKeyAuth | None",
request_data: dict[str, object],
error: Exception,
) -> None:
@ -1072,7 +1078,7 @@ class LiteLLM_Proxy_MCP_Handler:
all_tools: Sequence[object] | None,
mcp_tools_with_litellm_proxy: list[Mapping[str, object]],
mcp_discovery_events: list[ResponsesAPIStreamingResponse],
call_params: dict[str, Any],
call_params: Mapping[str, object],
previous_response_id: str | None,
tool_server_map: dict[str, str],
**kwargs,
@ -1115,10 +1121,10 @@ class LiteLLM_Proxy_MCP_Handler:
input: str | ResponseInputParam,
model: str,
all_tools: Sequence[object] | None,
call_params: dict[str, Any],
call_params: Mapping[str, object],
previous_response_id: str | None,
**kwargs,
) -> dict[str, Any]:
**kwargs: object,
) -> dict[str, object]:
"""
Build a clean request parameters dictionary for MCP streaming.
@ -1126,7 +1132,7 @@ class LiteLLM_Proxy_MCP_Handler:
in a clean, maintainable way.
"""
# Start with the core required parameters
request_params: Final = {
request_params: Final[dict[str, object]] = {
"input": input,
"model": model,
"tools": all_tools,
@ -1146,7 +1152,7 @@ class LiteLLM_Proxy_MCP_Handler:
@staticmethod
def _create_tool_execution_events(
tool_calls: Sequence[object], tool_results: list[dict[str, Any]]
tool_calls: Sequence[object], tool_results: Sequence[MCPToolResult]
) -> list[ResponsesAPIStreamingResponse]:
"""
Create MCP tool execution events for streaming.

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

@ -19,13 +19,13 @@ from litellm.types.llms.openai import (
ResponsesAPIResponse,
ResponsesAPIStreamEvents,
ResponsesAPIStreamingResponse,
ToolParam,
)
if TYPE_CHECKING:
from mcp.types import Tool as MCPTool
from litellm.proxy._types import UserAPIKeyAuth
from litellm.responses.mcp.litellm_proxy_mcp_handler import MCPToolResult
else:
MCPTool = Any
@ -33,7 +33,7 @@ MAX_MCP_TOOL_CALL_ROUNDS: Final = 5
async def create_mcp_list_tools_events(
mcp_tools_with_litellm_proxy: list[ToolParam],
mcp_tools_with_litellm_proxy: Sequence[Mapping[str, object]],
user_api_key_auth: "UserAPIKeyAuth | None",
base_item_id: str,
pre_processed_mcp_tools: list[MCPTool],
@ -44,13 +44,14 @@ async def create_mcp_list_tools_events(
try:
# Extract MCP server names
mcp_servers: Final = []
for tool in mcp_tools_with_litellm_proxy:
if isinstance(tool, dict) and "server_url" in tool:
server_url = tool.get("server_url")
if isinstance(server_url, str) and server_url.startswith("litellm_proxy/mcp/"):
server_name = server_url.split("/")[-1]
mcp_servers.append(server_name)
_mcp_servers: Final = [
server_url.split("/")[-1]
for tool in mcp_tools_with_litellm_proxy
if isinstance(tool, dict)
and "server_url" in tool
and isinstance(server_url := tool.get("server_url"), str)
and server_url.startswith("litellm_proxy/mcp/")
]
# Emit list tools in progress event
in_progress_event: Final = MCPListToolsInProgressEvent(
@ -65,15 +66,14 @@ async def create_mcp_list_tools_events(
filtered_mcp_tools: Final = pre_processed_mcp_tools
# Convert tools to dict format for the event
mcp_tools_dict: Final = []
for tool in filtered_mcp_tools:
if hasattr(tool, "model_dump") and callable(getattr(tool, "model_dump")):
# Type cast to help mypy understand this is safe after hasattr check
mcp_tools_dict.append(cast(Any, tool).model_dump())
elif hasattr(tool, "__dict__"):
mcp_tools_dict.append(tool.__dict__)
else:
mcp_tools_dict.append({"name": getattr(tool, "name", str(tool))})
_mcp_tools_dict: Final = [
tool.model_dump()
if hasattr(tool, "model_dump") and callable(getattr(tool, "model_dump"))
else tool.__dict__
if hasattr(tool, "__dict__")
else {"name": getattr(tool, "name", str(tool))}
for tool in filtered_mcp_tools
]
# Emit list tools completed event
completed_event: Final = MCPListToolsCompletedEvent(
@ -96,21 +96,18 @@ async def create_mcp_list_tools_events(
server_label = str(server_label_value) if server_label_value is not None else ""
# Format tools for OpenAI output_item.done format
formatted_tools: Final = []
for tool in filtered_mcp_tools:
tool_dict = {
formatted_tools: Final = [
{
"name": getattr(tool, "name", "unknown"),
"description": getattr(tool, "description", ""),
"annotations": {"read_only": False},
**dict.fromkeys(
("input_schema",) if hasattr(tool, "inputSchema") or hasattr(tool, "input_schema") else (),
getattr(tool, "inputSchema", getattr(tool, "input_schema", None)),
),
}
# Add input_schema if available
if hasattr(tool, "inputSchema"):
tool_dict["input_schema"] = getattr(tool, "inputSchema")
elif hasattr(tool, "input_schema"):
tool_dict["input_schema"] = getattr(tool, "input_schema")
formatted_tools.append(tool_dict)
for tool in filtered_mcp_tools
]
# Create the output_item.done event with MCP tools list
output_item_done_event = OutputItemDoneEvent(
@ -166,7 +163,7 @@ async def create_mcp_list_tools_events(
def create_mcp_call_events(
tool_name: str,
tool_call_id: str,
tool_call_id: str | None,
arguments: str,
result: str | None = None,
base_item_id: str | None = None,
@ -256,9 +253,12 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
4. Emits tool execution events in the stream
"""
model: str
tool_results: "Sequence[MCPToolResult]"
def __init__(
self,
base_iterator: Any, # Can be None - will be created internally
base_iterator: "BaseResponsesAPIStreamingIterator | ResponsesAPIResponse | None", # created internally when None
mcp_events: list[ResponsesAPIStreamingResponse],
tool_server_map: dict[str, str],
mcp_tools_with_litellm_proxy: Sequence[Mapping[str, object]] | None = None,
@ -285,7 +285,9 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
self.tool_server_map = tool_server_map
# Iterator references
self.base_iterator: Any | ResponsesAPIResponse | None = base_iterator # Will be created when needed
self.base_iterator: BaseResponsesAPIStreamingIterator | ResponsesAPIResponse | None = (
base_iterator # Will be created when needed
)
# Response collection for tool execution
self.collected_response: ResponsesAPIResponse | None = None

16
ruff-strict-budget.json
View file

@ -1,12 +1,12 @@
{
"ANN001": {
"limit": 3114
"limit": 3106
},
"ANN002": {
"limit": 71
},
"ANN003": {
"limit": 834
"limit": 832
},
"ANN201": {
"limit": 2031
@ -24,7 +24,7 @@
"limit": 133
},
"ANN401": {
"limit": 1499
"limit": 1495
},
"ASYNC230": {
"limit": 11
@ -39,7 +39,7 @@
"limit": 505
},
"B009": {
"limit": 81
"limit": 79
},
"B010": {
"limit": 190
@ -57,7 +57,7 @@
"limit": 3
},
"BLE001": {
"limit": 2921
"limit": 2924
},
"C401": {
"limit": 8
@ -78,7 +78,7 @@
"limit": 1
},
"C901": {
"limit": 314
"limit": 313
},
"D419": {
"limit": 6
@ -192,7 +192,7 @@
"limit": 0
},
"S110": {
"limit": 215
"limit": 218
},
"S112": {
"limit": 22
@ -234,7 +234,7 @@
"limit": 5
},
"TID251": {
"limit": 1229
"limit": 1226
},
"TRY002": {
"limit": 528

10
type-discipline-budget.json
View file

@ -1,9 +1,9 @@
{
"LIT001": {
"limit": 23121
"limit": 23064
},
"LIT002": {
"limit": 27090
"limit": 27166
},
"LIT003": {
"limit": 269
@ -15,7 +15,7 @@
"limit": 0
},
"LIT006": {
"limit": 1083
"limit": 1078
},
"LIT007": {
"limit": 0
@ -27,9 +27,9 @@
"limit": 0
},
"LIT010": {
"limit": 16669
"limit": 16753
},
"LIT011": {
"limit": 5592
"limit": 5598
}
}