mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
Merge 492c8753fc into 6085538b33
This commit is contained in:
Yana Tokareva
GitHub
commit
d7d06da34d
9 changed files with 144 additions and 9 deletions
|
|
@ -253,8 +253,7 @@ class DeepKeepGuardrail(CustomGuardrail):
|
|||
getattr(logging_obj, "litellm_trace_id", None) if logging_obj else None,
|
||||
exc_info=error,
|
||||
)
|
||||
return_inputs: Final[GenericGuardrailAPIInputs] = {}
|
||||
return_inputs.update(inputs)
|
||||
return_inputs: Final[GenericGuardrailAPIInputs] = {**inputs}
|
||||
return return_inputs
|
||||
|
||||
def _handle_guardrail_request_error(
|
||||
|
|
|
|||
|
|
@ -413,8 +413,7 @@ class GenericGuardrailAPI(CustomGuardrail):
|
|||
exc_info=error,
|
||||
)
|
||||
# Keep flow going - treat as action=NONE (no modifications)
|
||||
return_inputs: Final[GenericGuardrailAPIInputs] = {}
|
||||
return_inputs.update(inputs)
|
||||
return_inputs: Final[GenericGuardrailAPIInputs] = {**inputs}
|
||||
return return_inputs
|
||||
|
||||
def _build_request_headers(self) -> Mapping[str, str]:
|
||||
|
|
@ -592,6 +591,7 @@ class GenericGuardrailAPI(CustomGuardrail):
|
|||
tools=([GuardrailToolParam.model_validate(t) for t in tools] if tools else None),
|
||||
structured_messages=None if self.forward_messages_as_sent else structured_messages,
|
||||
tool_calls=tool_calls,
|
||||
mcp_tool_call=inputs.get("mcp_tool_call"),
|
||||
additional_provider_specific_params=additional_params,
|
||||
input_type=input_type,
|
||||
model=model,
|
||||
|
|
|
|||
|
|
@ -1213,8 +1213,7 @@ class StraikerGuardrail(CustomGuardrail):
|
|||
inputs: GenericGuardrailAPIInputs,
|
||||
parsed: StraikerWebhookResponse,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
return_inputs: Final[GenericGuardrailAPIInputs] = {}
|
||||
return_inputs.update(inputs)
|
||||
return_inputs: Final[GenericGuardrailAPIInputs] = {**inputs}
|
||||
if parsed.texts is not None:
|
||||
return_inputs["texts"] = parsed.texts
|
||||
return return_inputs
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@ from typing import TYPE_CHECKING, Any, Final
|
|||
|
||||
from fastapi import HTTPException
|
||||
from mcp.types import Tool as MCPTool
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
|
||||
|
|
@ -43,7 +44,7 @@ from litellm.types.llms.openai import (
|
|||
ChatCompletionToolParam,
|
||||
ChatCompletionToolParamFunctionChunk,
|
||||
)
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs, GenericGuardrailAPIMCPToolCall
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from mcp.types import CallToolResult
|
||||
|
|
@ -52,6 +53,29 @@ if TYPE_CHECKING:
|
|||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
||||
|
||||
_MCP_ARGUMENTS: Final[TypeAdapter[Mapping[str, JsonValue]]] = TypeAdapter(Mapping[str, JsonValue])
|
||||
|
||||
|
||||
def _optional_text(value: object) -> str | None:
|
||||
return value if isinstance(value, str) else None
|
||||
|
||||
|
||||
def _mcp_tool_call(
|
||||
tool_name: str, description: object, arguments: object, data: Mapping[str, object]
|
||||
) -> GenericGuardrailAPIMCPToolCall:
|
||||
return GenericGuardrailAPIMCPToolCall(
|
||||
mcp_server_id=_optional_text(data.get("mcp_server_id")),
|
||||
mcp_server_name=_optional_text(data.get("mcp_server_name")),
|
||||
mcp_server_url=_optional_text(data.get("mcp_server_url")),
|
||||
mcp_server_transport=_optional_text(data.get("mcp_server_transport")),
|
||||
mcp_server_auth_type=_optional_text(data.get("mcp_server_auth_type")),
|
||||
mcp_server_is_byok=is_byok if isinstance(is_byok := data.get("mcp_server_is_byok"), bool) else None,
|
||||
tool_name=tool_name,
|
||||
tool_description=_optional_text(description),
|
||||
arguments=_MCP_ARGUMENTS.validate_python(arguments) if isinstance(arguments, Mapping) else {},
|
||||
)
|
||||
|
||||
|
||||
def _blocked(reason: str) -> HTTPException:
|
||||
return HTTPException(status_code=400, detail={"error": f"Content blocked: {reason}"})
|
||||
|
||||
|
|
@ -188,6 +212,7 @@ class MCPGuardrailTranslationHandler(BaseTranslation):
|
|||
inputs: Final[GenericGuardrailAPIInputs] = GenericGuardrailAPIInputs(
|
||||
tools=[tool_def],
|
||||
texts=list(scanned_texts),
|
||||
mcp_tool_call=_mcp_tool_call(str(mcp_tool_name), mcp_tool_description, mcp_arguments, data),
|
||||
)
|
||||
|
||||
guarded: Final = await guardrail_to_apply.apply_guardrail(
|
||||
|
|
|
|||
|
|
@ -6252,6 +6252,12 @@ class MCPServerManager:
|
|||
"arguments": arguments,
|
||||
"server_name": server_name,
|
||||
"mcp_rate_limit_server_name": server.alias or server.server_name or server.name,
|
||||
"mcp_server_id": server.server_id,
|
||||
"mcp_server_url": redact_mcp_resource_url(server.url),
|
||||
"mcp_server_transport": server.transport,
|
||||
"mcp_server_auth_type": server.auth_type,
|
||||
"mcp_server_is_byok": server.is_byok,
|
||||
"client_ip": client_ip,
|
||||
"user_api_key_auth": checked_auth,
|
||||
"user_api_key_user_id": checked_auth.user_id if checked_auth else None,
|
||||
"user_api_key_team_id": checked_auth.team_id if checked_auth else None,
|
||||
|
|
|
|||
|
|
@ -2236,6 +2236,7 @@ class ProxyLogging:
|
|||
"user_api_key_user_id": kwargs.get("user_api_key_user_id"),
|
||||
"user_api_key_team_id": kwargs.get("user_api_key_team_id"),
|
||||
"user_api_key_end_user_id": kwargs.get("user_api_key_end_user_id"),
|
||||
"requester_ip_address": kwargs.get("client_ip"),
|
||||
}
|
||||
|
||||
# Create synthetic LLM data that guardrails can process
|
||||
|
|
@ -2256,6 +2257,11 @@ class ProxyLogging:
|
|||
# synthetic call_mcp_tool payload (otherwise a key with
|
||||
# mcp_rpm_limit could exceed it via the MCP path).
|
||||
"mcp_server_name": kwargs.get("mcp_rate_limit_server_name"),
|
||||
"mcp_server_id": kwargs.get("mcp_server_id"),
|
||||
"mcp_server_url": kwargs.get("mcp_server_url"),
|
||||
"mcp_server_transport": kwargs.get("mcp_server_transport"),
|
||||
"mcp_server_auth_type": kwargs.get("mcp_server_auth_type"),
|
||||
"mcp_server_is_byok": kwargs.get("mcp_server_is_byok"),
|
||||
# Raw Bearer token from the original HTTP request — allows guardrails
|
||||
# (e.g. MCPJWTSigner) to independently verify the caller's identity
|
||||
# before re-signing an outbound token (FR-5 verify+re-sign).
|
||||
|
|
@ -2264,6 +2270,11 @@ class ProxyLogging:
|
|||
}
|
||||
user_api_key_auth: Final = kwargs.get("user_api_key_auth")
|
||||
if isinstance(user_api_key_auth, UserAPIKeyAuth):
|
||||
synthetic_metadata.update(
|
||||
LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_auth)
|
||||
)
|
||||
synthetic_metadata["actor_agent_id"] = user_api_key_auth.agent_id
|
||||
synthetic_metadata["target_agent_id"] = user_api_key_auth.invoked_agent_id
|
||||
add_guardrails_from_auth_metadata(
|
||||
user_api_key_dict=user_api_key_auth,
|
||||
data=synthetic_data,
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ from collections.abc import Mapping, Sequence
|
|||
from typing import Any, Final, Literal, cast # noqa: TID251 # JSON chat rows have no typed constructor across roles
|
||||
|
||||
from pydantic import ConfigDict, Field, SerializerFunctionWrapHandler, model_serializer
|
||||
from typing_extensions import TypedDict
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm.types.guardrails import GuardrailNumRetries
|
||||
from litellm.types.llms.base import LiteLLMBaseModel
|
||||
|
|
@ -11,7 +11,7 @@ from litellm.types.llms.openai import (
|
|||
ChatCompletionToolCallChunk,
|
||||
)
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
|
||||
from litellm.types.utils import ChatCompletionMessageToolCall
|
||||
from litellm.types.utils import ChatCompletionMessageToolCall, GenericGuardrailAPIMCPToolCall
|
||||
|
||||
|
||||
class GuardrailToolParam(LiteLLMBaseModel):
|
||||
|
|
@ -43,6 +43,9 @@ class GenericGuardrailAPIMetadata(TypedDict, total=False):
|
|||
user_api_key_team_alias: str | None
|
||||
user_api_key_end_user_id: str | None
|
||||
user_api_key_org_id: str | None
|
||||
requester_ip_address: ReadOnly[str | None]
|
||||
actor_agent_id: ReadOnly[str | None]
|
||||
target_agent_id: ReadOnly[str | None]
|
||||
|
||||
|
||||
class GenericGuardrailAPIOptionalParams(LiteLLMBaseModel):
|
||||
|
|
@ -164,6 +167,10 @@ class GenericGuardrailAPIRequest(LiteLLMBaseModel):
|
|||
additional_provider_specific_params: dict[str, object] | None = None
|
||||
tool_calls: list[ChatCompletionToolCallChunk] | list[ChatCompletionMessageToolCall] | None = None
|
||||
model: str | None = None # the model being used for the LLM call
|
||||
mcp_tool_call: GenericGuardrailAPIMCPToolCall | None = Field(
|
||||
default=None,
|
||||
description="The MCP server, tool, and named arguments of the MCP tool call being checked.",
|
||||
)
|
||||
|
||||
|
||||
def coerce_stream_holdback_value(value: Any) -> int:
|
||||
|
|
|
|||
|
|
@ -5362,6 +5362,18 @@ class PriorityReservationSettings(LiteLLMBaseModel):
|
|||
model_config = ConfigDict(protected_namespaces=())
|
||||
|
||||
|
||||
class GenericGuardrailAPIMCPToolCall(TypedDict, total=False):
|
||||
mcp_server_id: ReadOnly[str | None]
|
||||
mcp_server_name: ReadOnly[str | None]
|
||||
mcp_server_url: ReadOnly[str | None]
|
||||
mcp_server_transport: ReadOnly[str | None]
|
||||
mcp_server_auth_type: ReadOnly[str | None]
|
||||
mcp_server_is_byok: ReadOnly[bool | None]
|
||||
tool_name: ReadOnly[str]
|
||||
tool_description: ReadOnly[str | None]
|
||||
arguments: ReadOnly[Mapping[str, JsonValue]]
|
||||
|
||||
|
||||
class GenericGuardrailAPIInputs(TypedDict, total=False):
|
||||
texts: list[str] # extracted text from the LLM response - for basic text guardrails
|
||||
images: list[str] # extracted images from the LLM response - for image guardrails
|
||||
|
|
@ -5371,6 +5383,7 @@ class GenericGuardrailAPIInputs(TypedDict, total=False):
|
|||
AllMessageValues
|
||||
] # structured messages sent to the LLM - indicates if text is from system or user
|
||||
model: str | None # the model being used for the LLM call
|
||||
mcp_tool_call: ReadOnly[GenericGuardrailAPIMCPToolCall]
|
||||
stream_holdback_chars: list[
|
||||
int
|
||||
] # trailing chars to withhold from streaming emission per text (word-boundary safety)
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ from mcp.types import CallToolResult, TextContent
|
|||
import litellm
|
||||
from litellm import ModelResponse
|
||||
from litellm._version import version as litellm_version
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.exceptions import GuardrailRaisedException, Timeout
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, MaskedHTTPStatusError
|
||||
from litellm.proxy._experimental.mcp_server.guardrail_translation.handler import (
|
||||
|
|
@ -34,7 +35,9 @@ from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api.generic_guar
|
|||
_HEADER_PRESENT_PLACEHOLDER,
|
||||
)
|
||||
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import UnifiedLLMGuardrails
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.mcp import MCPAuth, MCPTransport
|
||||
from litellm.types.utils import Choices, GenericGuardrailAPIInputs, Message
|
||||
|
||||
|
||||
|
|
@ -2692,3 +2695,75 @@ async def test_during_mcp_call_blocked_raises() -> None:
|
|||
)
|
||||
|
||||
assert exc_info.value.message == "credential detected in tool input"
|
||||
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_mcp_call_sends_the_server_tool_arguments_and_caller() -> None:
|
||||
sent_bodies: Final[SimpleQueue[bytes]] = SimpleQueue()
|
||||
|
||||
async def answer(request: httpx.Request) -> httpx.Response:
|
||||
sent_bodies.put(request.content)
|
||||
return httpx.Response(200, json={"action": "NONE"}, request=request)
|
||||
|
||||
guardrail: Final = _mcp_guardrail("pre_mcp_call")
|
||||
arguments: Final = {"repo": "octo/hello", "options": {"force": True, "depth": 2}}
|
||||
caller: Final = UserAPIKeyAuth(
|
||||
api_key="hashed-key",
|
||||
key_alias="ci-bot",
|
||||
user_id="u-1",
|
||||
user_email="dev@example.com",
|
||||
team_id="t-1",
|
||||
team_alias="platform",
|
||||
org_id="o-1",
|
||||
end_user_id="e-1",
|
||||
agent_id="agent-1",
|
||||
).model_copy(update={"invoked_agent_id": "agent-2"})
|
||||
pre_hook_kwargs: Final = {
|
||||
"name": "delete_repo",
|
||||
"arguments": arguments,
|
||||
"server_name": "github",
|
||||
"mcp_rate_limit_server_name": "github",
|
||||
"mcp_server_id": "srv-1",
|
||||
"mcp_server_url": "https://mcp.example.com",
|
||||
"mcp_server_transport": MCPTransport.http,
|
||||
"mcp_server_auth_type": MCPAuth.oauth2,
|
||||
"mcp_server_is_byok": True,
|
||||
"client_ip": "203.0.113.7",
|
||||
"tool_description": "Delete a repository",
|
||||
"user_api_key_auth": caller,
|
||||
}
|
||||
proxy_logging_obj: Final = ProxyLogging(user_api_key_cache=DualCache())
|
||||
data: Final = proxy_logging_obj._convert_mcp_to_llm_format(
|
||||
proxy_logging_obj._create_mcp_request_object_from_kwargs(pre_hook_kwargs), pre_hook_kwargs
|
||||
)
|
||||
|
||||
with patch.object(guardrail, "async_handler", AsyncHTTPHandler(transport=httpx.MockTransport(answer))):
|
||||
await MCPGuardrailTranslationHandler().process_input_messages(data=data, guardrail_to_apply=guardrail)
|
||||
|
||||
body: Final = json.loads(sent_bodies.get_nowait())
|
||||
assert body["mcp_tool_call"] == {
|
||||
"mcp_server_id": "srv-1",
|
||||
"mcp_server_name": "github",
|
||||
"mcp_server_url": "https://mcp.example.com",
|
||||
"mcp_server_transport": "http",
|
||||
"mcp_server_auth_type": "oauth2",
|
||||
"mcp_server_is_byok": True,
|
||||
"tool_name": "delete_repo",
|
||||
"tool_description": "Delete a repository",
|
||||
"arguments": arguments,
|
||||
}
|
||||
assert body["request_data"] == {
|
||||
"user_api_key_hash": "hashed-key",
|
||||
"user_api_key_alias": "ci-bot",
|
||||
"user_api_key_user_id": "u-1",
|
||||
"user_api_key_user_email": "dev@example.com",
|
||||
"user_api_key_team_id": "t-1",
|
||||
"user_api_key_team_alias": "platform",
|
||||
"user_api_key_end_user_id": "e-1",
|
||||
"user_api_key_org_id": "o-1",
|
||||
"requester_ip_address": "203.0.113.7",
|
||||
"actor_agent_id": "agent-1",
|
||||
"target_agent_id": "agent-2",
|
||||
}
|
||||
assert body["texts"] == ["Delete a repository", "octo/hello"]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue