diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py index cb1b223fb47..e94f7329885 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py @@ -22,6 +22,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.llms.anthropic.common_utils import anthropic_server_tool_call_params from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -163,11 +164,11 @@ class lakeraAI_Moderation(CustomGuardrail): for message in data["messages"]: role = message.get("role") if role in stringified_roles: - if "tool_calls" in message: - tool_call_messages = [ - *tool_call_messages, - *message["tool_calls"], - ] + tool_call_messages = [ + *tool_call_messages, + *anthropic_server_tool_call_params(message), + *message.get("tool_calls", ()), + ] if role == Role.SYSTEM.value: # we need this for later system_message = message continue @@ -185,10 +186,11 @@ class lakeraAI_Moderation(CustomGuardrail): if system_message is not None: if not litellm.add_function_to_prompt: content = system_message.get("content") - function_input: Final = [] - for tool_call in tool_call_messages: - if "function" in tool_call: - function_input.append(tool_call["function"]["arguments"]) + function_input: Final[tuple[str, ...]] = tuple( + tool_call["function"]["arguments"] + for tool_call in tool_call_messages + if "function" in tool_call + ) if len(function_input) > 0: content += " Function Input: " + " ".join(function_input) diff --git a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py index 5902a07dea2..0530cf22f62 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py @@ -45,6 +45,7 @@ from litellm.llms.anthropic.common_utils import ( anthropic_server_tool_calls, with_server_tool_call_arguments, write_anthropic_server_tool_calls, + write_request_server_tool_calls, ) from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, @@ -571,6 +572,15 @@ class LassoGuardrail(CustomGuardrail): elif role == "assistant" and not content and orig_msg.get("tool_calls"): msg["tool_calls"] = self._update_tool_calls_from_masked(orig_msg["tool_calls"], masked_tool_use) + if role == "assistant": + write_request_server_tool_calls( + msg, + tuple( + self._masked_server_tool_call(server_tool_call, masked_tool_use) + for server_tool_call in anthropic_server_tool_call_params(orig_msg) + ), + ) + result.append(msg) return result @@ -711,7 +721,7 @@ class LassoGuardrail(CustomGuardrail): expanded.append({"role": role, "content": content}) if role == "assistant": - for call in iter_items(msg.get("tool_calls") or []): + for call in (*anthropic_server_tool_call_params(msg), *iter_items(msg.get("tool_calls") or [])): call_id, name, input_data = self._extract_tool_call_fields(call) if not call_id or not name: verbose_proxy_logger.warning( diff --git a/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py b/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py index 6150fccac0c..c76e8e9c084 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py +++ b/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py @@ -18,6 +18,7 @@ from litellm.integrations.custom_guardrail import ( log_guardrail_information, ) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.anthropic.common_utils import anthropic_server_tool_call_params from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -180,8 +181,8 @@ class QualifireGuardrail(CustomGuardrail): api_message["tool_call_id"] = tool_call_id # Handle tool calls if present - tool_calls = msg.get("tool_calls") - if tool_calls and isinstance(tool_calls, list): + tool_calls = [*anthropic_server_tool_call_params(msg), *(msg.get("tool_calls") or [])] + if tool_calls: api_tool_calls = [] for tc in tool_calls: if isinstance(tc, dict): diff --git a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index a7fa752f654..c90a9d8d229 100644 --- a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -6,10 +6,11 @@ with guardrail transformations, including tool calls. """ import json -from collections.abc import Mapping +from collections.abc import Iterable, Mapping from typing import Any, Final, Literal, Optional import pytest +from openai.types.responses import ResponseFunctionWebSearch from litellm.integrations.custom_guardrail import CustomGuardrail @@ -149,14 +150,14 @@ def _server_tool_call(call_id: str, name: str, arguments: Mapping[str, object]) return {"id": call_id, "type": "function", "function": {"name": name, "arguments": json.dumps(arguments)}} -def _tool_call_ids_names_and_arguments(tool_calls) -> list[tuple[str, str, str]]: +def _tool_call_ids_names_and_arguments(tool_calls: Iterable[Mapping[str, Any]]) -> list[tuple[str, str, str]]: return [ (tool_call["id"], tool_call["function"]["name"], tool_call["function"]["arguments"]) for tool_call in tool_calls ] -def _hosted_web_call_inputs(web_search_calls) -> list[tuple[str, str, Optional[str]]]: +def _hosted_web_call_inputs(web_search_calls: Iterable[ResponseFunctionWebSearch]) -> list[tuple[str, str, Optional[str]]]: return [ (call.id, call.status, call.action.query if call.action.type == "search" else call.action.url) for call in web_search_calls @@ -177,9 +178,9 @@ class QueryRedactingGuardrail(CustomGuardrail): input_type: Literal["request", "response"], logging_obj: Optional[Any] = None, ) -> GenericGuardrailAPIInputs: - def redacted(tool_call): - arguments = json.loads(tool_call["function"]["arguments"]) - redacted_arguments = { + def redacted(tool_call: Mapping[str, Any]) -> Mapping[str, Any]: + arguments: Final[dict[str, object]] = json.loads(tool_call["function"]["arguments"]) + redacted_arguments: Final = { key: "[REDACTED]" if key in ("query", "url") else value for key, value in arguments.items() } if redacted_arguments == arguments: diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai.py new file mode 100644 index 00000000000..6880d0ac491 --- /dev/null +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai.py @@ -0,0 +1,57 @@ +import json +from typing import Final + +import httpx +import pytest +import respx + +import litellm +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.guardrail_hooks.lakera_ai import lakeraAI_Moderation + + +@pytest.mark.asyncio +async def test_prompt_check_appends_replayed_server_tool_call_arguments_to_system_message( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "add_function_to_prompt", False) + route: Final = respx_mock.post("https://api.lakera.ai/v1/prompt_injection").mock( + return_value=httpx.Response(200, json={"results": [{"flagged": False}]}) + ) + moderation: Final = lakeraAI_Moderation(api_key="test-key", moderation_check="pre_call") + + await moderation._check( + data={ + "messages": [ + {"role": "system", "content": "You are a shopping assistant."}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_01", + "type": "function", + "function": {"name": "lookup_order", "arguments": '{"order_id": "A-17"}'}, + } + ], + "provider_specific_fields": { + "server_tool_calls": [ + { + "id": "srvtoolu_01", + "type": "function", + "function": {"name": "web_fetch", "arguments": '{"url": "https://example.com"}'}, + } + ] + }, + }, + ] + }, + user_api_key_dict=UserAPIKeyAuth(), + call_type="completion", + ) + + assert json.loads(route.calls.last.request.content)["input"][0] == { + "role": "system", + "content": 'You are a shopping assistant. Function Input: {"url": "https://example.com"} {"order_id": "A-17"}', + } diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_lasso.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_lasso.py index 6b6849f4f49..634e511fb91 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_lasso.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_lasso.py @@ -1691,3 +1691,113 @@ async def test_post_call_masking_rewrites_server_tool_call_arguments( } ] } + + +def _replayed_web_fetch_request(url: str) -> dict[str, object]: + return { + "messages": [ + {"role": "user", "content": "open the order page"}, + { + "role": "assistant", + "content": None, + "provider_specific_fields": { + "server_tool_calls": [ + { + "id": "srvtoolu_01", + "type": "function", + "function": {"name": "web_fetch", "arguments": json.dumps({"url": url})}, + } + ] + }, + }, + {"role": "user", "content": "summarize it"}, + ] + } + + +@pytest.mark.asyncio +async def test_pre_call_scan_sends_replayed_server_tool_calls_as_tool_use_entries( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + lasso_guardrail: Final = LassoGuardrail( + lasso_api_key="test-key", + user_id="test-user", + conversation_id="test-conversation", + guardrail_name="test-guard", + event_hook="pre_call", + default_on=True, + ) + route: Final = respx_mock.post("https://server.lasso.security/gateway/v3/classify").mock( + return_value=httpx.Response(200, json={"violations_detected": False, "deputies": {}, "findings": {}}) + ) + + await lasso_guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data=_replayed_web_fetch_request("https://example.com/?email=jane@example.com"), + call_type="completion", + ) + + assert json.loads(route.calls.last.request.content)["messages"] == [ + {"role": "user", "content": "open the order page"}, + { + "role": "model", + "content": { + "type": "tool_use", + "id": "srvtoolu_01", + "name": "web_fetch", + "input": {"url": "https://example.com/?email=jane@example.com"}, + }, + }, + {"role": "user", "content": "summarize it"}, + ] + + +@pytest.mark.asyncio +async def test_pre_call_masking_rewrites_replayed_server_tool_call_arguments( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + lasso_guardrail: Final = LassoGuardrail( + lasso_api_key="test-key", + user_id="test-user", + conversation_id="test-conversation", + mask=True, + guardrail_name="test-guard", + event_hook="pre_call", + default_on=True, + ) + respx_mock.post("https://server.lasso.security/gateway/v3/classifix").mock( + return_value=httpx.Response( + 200, + json={ + "violations_detected": True, + "deputies": {"pattern-detection": True}, + "findings": {"pattern-detection": [{"action": "AUTO_MASKING", "severity": "HIGH"}]}, + "messages": [ + {"role": "user", "content": "open the order page"}, + { + "role": "model", + "content": { + "type": "tool_use", + "id": "srvtoolu_01", + "name": "web_fetch", + "input": {"url": "https://example.com/?email="}, + }, + }, + {"role": "user", "content": "summarize it"}, + ], + }, + ) + ) + + result: Final = await lasso_guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data=_replayed_web_fetch_request("https://example.com/?email=jane@example.com"), + call_type="completion", + ) + + assert isinstance(result, dict) + assert result["messages"] == _replayed_web_fetch_request("https://example.com/?email=")["messages"] diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_qualifire.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_qualifire.py index dfd54cff730..ef97a335562 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_qualifire.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_qualifire.py @@ -2,6 +2,7 @@ Unit tests for Qualifire guardrail integration. """ +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -735,3 +736,41 @@ class TestQualifireGuardrailRegistry: assert "qualifire" in guardrail_class_registry assert guardrail_class_registry["qualifire"] == QualifireGuardrail + + +def test_convert_messages_sends_replayed_server_tool_calls_with_client_tool_calls() -> None: + from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( + QualifireGuardrail, + ) + + guardrail: Final = QualifireGuardrail(api_key="test_key", guardrail_name="test_guardrail") + + result: Final = guardrail._convert_messages_to_api_format( + [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_123", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"location": "NYC"}'}, + } + ], + "provider_specific_fields": { + "server_tool_calls": [ + { + "id": "srvtoolu_01", + "type": "function", + "function": {"name": "web_fetch", "arguments": '{"url": "https://example.com"}'}, + } + ] + }, + }, + ] + ) + + assert result[0]["tool_calls"] == [ + {"id": "srvtoolu_01", "name": "web_fetch", "arguments": {"url": "https://example.com"}}, + {"id": "call_123", "name": "get_weather", "arguments": {"location": "NYC"}}, + ]