diff --git a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py index a72f3e4c3ff0..fa4759f9b787 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py @@ -5,6 +5,7 @@ # # +-------------------------------------------------------------+ +import json import os import uuid from typing import ( @@ -219,11 +220,58 @@ async def async_post_call_success_hook( # Extract messages from the response for validation if isinstance(response, litellm.ModelResponse): - response_messages = [] + response_messages: List[Dict[str, Any]] = [] for choice in response.choices: - if hasattr(choice, "message") and choice.message.content: + if not hasattr(choice, "message"): + continue + msg = choice.message + if msg.content: response_messages.append( - {"role": "assistant", "content": choice.message.content} + {"role": "assistant", "content": msg.content} + ) + for call in getattr(msg, "tool_calls", None) or []: + call_id = ( + call.get("id") + if isinstance(call, dict) + else getattr(call, "id", None) + ) + func = ( + call.get("function") + if isinstance(call, dict) + else getattr(call, "function", None) + ) + name = ( + func.get("name") + if isinstance(func, dict) + else getattr(func, "name", None) + ) + args_str = ( + func.get("arguments") + if isinstance(func, dict) + else getattr(func, "arguments", None) + ) + if not call_id or not name: + continue + input_data = None + if args_str: + try: + parsed = json.loads(args_str) + if isinstance(parsed, dict): + input_data = parsed + except (json.JSONDecodeError, TypeError): + verbose_proxy_logger.debug( + "Failed to parse tool_call arguments as JSON, sending null input" + ) + response_messages.append( + { + "role": "model", + "content": { + "type": "tool_use", + "id": call_id, + "name": name, + "input": input_data, + }, + } ) if response_messages: @@ -366,7 +414,10 @@ async def _run_lasso_guardrail( LassoGuardrailAPIError: If the Lasso API call fails HTTPException: If blocking violations are detected """ - messages: List[Dict[str, str]] = data.get("messages", []) + messages: List[Dict[str, Any]] = data.get("messages", []) + if not messages: + return data + messages = self._expand_messages_for_classification(messages) if not messages: return data @@ -382,7 +433,7 @@ async def _handle_classification( data: dict, cache: DualCache, message_type: Literal["PROMPT", "COMPLETION"], - messages: List[Dict[str, str]], + messages: List[Dict[str, Any]], ) -> dict: """Handle classification without masking.""" try: @@ -400,7 +451,7 @@ async def _handle_masking( data: dict, cache: DualCache, message_type: Literal["PROMPT", "COMPLETION"], - messages: List[Dict[str, str]], + messages: List[Dict[str, Any]], ) -> dict: """Handle masking with classifix endpoint.""" try: @@ -412,9 +463,13 @@ async def _handle_masking( ) self._process_lasso_response(response) - # Apply masking to messages if violations detected and masked messages are available + # Apply masking to messages if violations detected and masked messages are available. + # Map masked content back onto the original OpenAI-format messages so the + # downstream provider receives a compatible payload. if response.get("violations_detected") and response.get("messages"): - data["messages"] = response["messages"] + data["messages"] = self._map_masked_messages_back( + data["messages"], response["messages"] + ) self._log_masking_applied(message_type, dict(response)) return data @@ -422,6 +477,95 @@ async def _handle_masking( await self._handle_api_error(e, message_type) return data # This line won't be reached due to exception, but satisfies type checker + def _map_masked_messages_back( + self, + original_messages: List[Dict[str, Any]], + masked_messages: List[Dict[str, Any]], + ) -> List[Dict[str, Any]]: + """Map Lasso-format masked messages back onto the original OpenAI-format messages. + + Lasso receives expanded messages (tool_use / tool_result blocks) and returns them + in the same Lasso-internal format with sensitive values replaced. Writing those + blocks straight into data["messages"] would corrupt the OpenAI-compatible schema + the downstream provider expects. This helper re-applies only the masked content + while preserving the original structure. + """ + # Index masked content by type so we can look up by id without caring about order. + masked_tool_use: Dict[str, Dict[str, Any]] = {} + masked_tool_result: Dict[str, str] = {} + masked_text: List[str] = [] + + for msg in masked_messages: + content = msg.get("content") + if isinstance(content, dict): + if content.get("type") == "tool_use": + call_id = content.get("id") + if call_id: + masked_tool_use[call_id] = content + elif content.get("type") == "tool_result": + tool_use_id = content.get("tool_use_id") + if tool_use_id: + masked_tool_result[tool_use_id] = content.get("content", "") + elif isinstance(content, str): + masked_text.append(content) + + result: List[Dict[str, Any]] = [] + text_cursor = 0 + + for orig_msg in original_messages: + msg = dict(orig_msg) + role = msg.get("role") + content = msg.get("content") + + if role == "tool": + tool_call_id = msg.get("tool_call_id") + if tool_call_id and tool_call_id in masked_tool_result: + msg["content"] = masked_tool_result[tool_call_id] + + elif isinstance(content, str) and content: + if text_cursor < len(masked_text): + msg["content"] = masked_text[text_cursor] + text_cursor += 1 + if role == "assistant" and orig_msg.get("tool_calls"): + msg["tool_calls"] = self._update_tool_calls_from_masked( + orig_msg["tool_calls"], masked_tool_use + ) + + 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 + ) + + result.append(msg) + + return result + + def _update_tool_calls_from_masked( + self, + tool_calls: List[Any], + masked_tool_use: Dict[str, Dict[str, Any]], + ) -> List[Any]: + """Replace tool_call arguments with masked values returned by Lasso.""" + updated = [] + for call in tool_calls: + call_id = ( + call.get("id") if isinstance(call, dict) else getattr(call, "id", None) + ) + if call_id and call_id in masked_tool_use: + masked_input = masked_tool_use[call_id].get("input") + if masked_input is not None: + if isinstance(call, dict): + call = dict(call) + func = dict(call.get("function", {})) + func["arguments"] = json.dumps(masked_input) + call["function"] = func + else: + func = getattr(call, "function", None) + if func: + func.arguments = json.dumps(masked_input) + updated.append(call) + return updated + async def _handle_api_error( self, error: Exception, @@ -478,6 +622,97 @@ def _log_masking_applied( }, ) + def _expand_messages_for_classification( + self, messages: List[Dict[str, Any]] + ) -> List[Dict[str, Any]]: + """ + Convert raw OpenAI-format messages to Lasso API format with content blocks. + + - assistant messages with `tool_calls` → assistant message per tool_use block + - role=tool messages → developer role + tool_result block + - plain text messages pass through unchanged + """ + expanded: List[Dict[str, Any]] = [] + for msg in messages: + role = msg.get("role", "") + content = msg.get("content") + + if role == "tool": + tool_call_id = msg.get("tool_call_id") + if not tool_call_id: + verbose_proxy_logger.warning( + "Skipping tool message without tool_call_id" + ) + continue + expanded.append( + { + "role": "developer", + "content": { + "type": "tool_result", + "tool_use_id": tool_call_id, + "content": content or "", + }, + } + ) + continue + + if content: + expanded.append({"role": role, "content": content}) + + if role == "assistant": + for call in msg.get("tool_calls") or []: + call_id = ( + call.get("id") + if isinstance(call, dict) + else getattr(call, "id", None) + ) + func = ( + call.get("function") + if isinstance(call, dict) + else getattr(call, "function", None) + ) + if not func: + continue + name = ( + func.get("name") + if isinstance(func, dict) + else getattr(func, "name", None) + ) + args_str = ( + func.get("arguments") + if isinstance(func, dict) + else getattr(func, "arguments", None) + ) + if not call_id or not name: + verbose_proxy_logger.warning( + "Skipping malformed tool_call", + extra={"call_id": call_id, "name": name}, + ) + continue + input_data = None + if args_str: + try: + parsed = json.loads(args_str) + if isinstance(parsed, dict): + input_data = parsed + except (json.JSONDecodeError, TypeError): + verbose_proxy_logger.warning( + "Failed to parse tool_call arguments, dropping input" + ) + expanded.append( + { + "role": "model", + "content": { + "type": "tool_use", + "id": call_id, + "name": name, + "input": input_data, + }, + } + ) + + return expanded + def _prepare_headers(self, data: dict, cache: DualCache) -> Dict[str, str]: """Prepare headers for the Lasso API request.""" if not self.lasso_api_key: @@ -504,7 +739,7 @@ def _prepare_headers(self, data: dict, cache: DualCache) -> Dict[str, str]: def _prepare_payload( self, - messages: List[Dict[str, str]], + messages: List[Dict[str, Any]], data: dict, cache: DualCache, message_type: Literal["PROMPT", "COMPLETION"] = "PROMPT", @@ -513,9 +748,9 @@ def _prepare_payload( Prepare the payload for the Lasso API request. Args: - messages: List of message objects + messages: List of message objects (may contain tool_use/tool_result content blocks) message_type: Type of message - "PROMPT" for input, "COMPLETION" for output - data: Request data (used for conversation_id generation) + data: Request data (used for conversation_id generation and tools extraction) cache: Cache instance for storing conversation_id (optional for post-call) """ payload: Dict[str, Any] = {"messages": messages, "messageType": message_type} @@ -526,9 +761,46 @@ def _prepare_payload( # Always include sessionId (conversation_id - generated or provided) conversation_id = self._get_or_generate_conversation_id(data, cache) - payload["sessionId"] = conversation_id + # Map OpenAI ChatCompletionToolParam array → ToolDefinition array + tools_data: List[Dict[str, Any]] = data.get("tools") or [] + if tools_data: + tool_definitions = [] + for tool in tools_data: + func = ( + tool.get("function") + if isinstance(tool, dict) + else getattr(tool, "function", None) + ) + if not func: + continue + name = ( + func.get("name") + if isinstance(func, dict) + else getattr(func, "name", None) + ) + if not name: + continue + td: Dict[str, Any] = {"name": name} + description = ( + func.get("description") + if isinstance(func, dict) + else getattr(func, "description", None) + ) + if description: + td["description"] = description + parameters = ( + func.get("parameters") + if isinstance(func, dict) + else getattr(func, "parameters", None) + ) + if parameters: + td["parameters"] = parameters + tool_definitions.append(td) + if tool_definitions: + payload["tools"] = tool_definitions + return payload async def _call_lasso_api( @@ -652,23 +924,55 @@ def _parse_violated_deputies(self, response: LassoResponse) -> List[str]: def _apply_masking_to_model_response( self, model_response: litellm.ModelResponse, - masked_messages: List[Dict[str, str]], + masked_messages: List[Dict[str, Any]], ) -> None: """Apply masking to the actual model response when mask=True and masked content is available.""" - masked_index = 0 + # Index masked tool_use blocks by id for O(1) lookup. + masked_tool_use: Dict[str, Dict[str, Any]] = {} + masked_text: List[str] = [] + for msg in masked_messages: + content = msg.get("content") + if isinstance(content, dict) and content.get("type") == "tool_use": + call_id = content.get("id") + if call_id: + masked_tool_use[call_id] = content + elif isinstance(content, str): + masked_text.append(content) + + text_cursor = 0 for choice in model_response.choices: - if ( - hasattr(choice, "message") - and choice.message.content - and masked_index < len(masked_messages) - ): - # Replace the content with the masked version from Lasso - choice.message.content = masked_messages[masked_index]["content"] - masked_index += 1 + if not hasattr(choice, "message"): + continue + msg = choice.message + + if msg.content and text_cursor < len(masked_text): + msg.content = masked_text[text_cursor] + text_cursor += 1 verbose_proxy_logger.debug( - f"Applied masked content to choice {masked_index}" + f"Applied masked text content to choice {text_cursor}" ) + for call in getattr(msg, "tool_calls", None) or []: + call_id = ( + call.get("id") + if isinstance(call, dict) + else getattr(call, "id", None) + ) + if call_id and call_id in masked_tool_use: + masked_input = masked_tool_use[call_id].get("input") + if masked_input is not None: + if isinstance(call, dict): + func = call.get("function", {}) + if isinstance(func, dict): + func["arguments"] = json.dumps(masked_input) + else: + func = getattr(call, "function", None) + if func: + func.arguments = json.dumps(masked_input) + verbose_proxy_logger.debug( + f"Applied masked tool_call arguments for call_id={call_id}" + ) + @staticmethod def get_config_model() -> Optional[Type["GuardrailConfigModel"]]: from litellm.types.proxy.guardrails.guardrail_hooks.lasso import ( diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lasso.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lasso.py index 6286d4ea4091..034b8cf10671 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lasso.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lasso.py @@ -767,3 +767,268 @@ def test_check_for_blocking_actions(self): empty_response = {} blocking_violations = guardrail._check_for_blocking_actions(empty_response) assert len(blocking_violations) == 0 + + # ------------------------------------------------------------------ + # Tool-calling tests + # ------------------------------------------------------------------ + + def test_payload_preparation_with_tools(self): + """_prepare_payload maps OpenAI ChatCompletionToolParam to ToolDefinition shape.""" + guardrail = LassoGuardrail( + lasso_api_key="test-api-key", + conversation_id="test-conversation", + ) + data = { + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get current weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}, + }, + } + ] + } + payload = guardrail._prepare_payload([], data, DualCache(), "PROMPT") + assert "tools" in payload + assert payload["tools"] == [ + { + "name": "get_weather", + "description": "Get current weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}, + } + ] + + def test_payload_preparation_no_tools(self): + """_prepare_payload omits tools key when no tools provided (regression).""" + guardrail = LassoGuardrail( + lasso_api_key="test-api-key", + conversation_id="test-conversation", + ) + messages = [{"role": "user", "content": "Hello"}] + payload = guardrail._prepare_payload(messages, {}, DualCache(), "PROMPT") + assert "tools" not in payload + assert payload["messages"] == messages + + def test_expand_messages_assistant_tool_calls(self): + """Pre-call: assistant tool_calls expand into tool_use content blocks.""" + guardrail = LassoGuardrail(lasso_api_key="test-api-key") + messages = [ + {"role": "user", "content": "What's the weather in NY?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_abc", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city":"NY"}'}, + } + ], + }, + ] + expanded = guardrail._expand_messages_for_classification(messages) + assert len(expanded) == 2 + assert expanded[0] == {"role": "user", "content": "What's the weather in NY?"} + assert expanded[1] == { + "role": "model", + "content": {"type": "tool_use", "id": "call_abc", "name": "get_weather", "input": {"city": "NY"}}, + } + + def test_expand_messages_tool_role(self): + """Pre-call: role=tool messages become developer + tool_result block.""" + guardrail = LassoGuardrail(lasso_api_key="test-api-key") + messages = [ + {"role": "tool", "tool_call_id": "call_abc", "content": "72°F, sunny"}, + ] + expanded = guardrail._expand_messages_for_classification(messages) + assert len(expanded) == 1 + assert expanded[0] == { + "role": "developer", + "content": {"type": "tool_result", "tool_use_id": "call_abc", "content": "72°F, sunny"}, + } + + def test_expand_messages_tool_role_missing_tool_call_id(self): + """Pre-call: tool message without tool_call_id is skipped with a warning.""" + guardrail = LassoGuardrail(lasso_api_key="test-api-key") + messages = [{"role": "tool", "content": "some result"}] + expanded = guardrail._expand_messages_for_classification(messages) + assert expanded == [] + + def test_expand_messages_assistant_with_text_and_tool_calls(self): + """Pre-call: assistant with both text and tool_calls produces text msg + tool_use msg.""" + guardrail = LassoGuardrail(lasso_api_key="test-api-key") + messages = [ + { + "role": "assistant", + "content": "Let me check that for you.", + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}} + ], + } + ] + expanded = guardrail._expand_messages_for_classification(messages) + assert len(expanded) == 2 + assert expanded[0] == {"role": "assistant", "content": "Let me check that for you."} + assert expanded[1]["content"]["type"] == "tool_use" + assert expanded[1]["content"]["name"] == "lookup" + + def test_expand_messages_plain_text_unchanged(self): + """Pre-call: plain text messages pass through without modification (regression).""" + guardrail = LassoGuardrail(lasso_api_key="test-api-key") + messages = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi there!"}, + ] + expanded = guardrail._expand_messages_for_classification(messages) + assert expanded == messages + + @pytest.mark.asyncio + async def test_post_call_with_tool_calls(self): + """Post-call: tool_calls in model response are extracted as tool_use blocks.""" + guardrail = LassoGuardrail( + lasso_api_key="test-api-key", + guardrail_name="test-guard", + event_hook="post_call", + default_on=True, + ) + data = {"messages": [{"role": "user", "content": "run the tool"}]} + + mock_model_response = MagicMock(spec=litellm.ModelResponse) + mock_choice = MagicMock() + mock_choice.message.content = None + tool_call = MagicMock() + tool_call.id = "call_xyz" + tool_call.function.name = "my_tool" + tool_call.function.arguments = '{"param": "value"}' + mock_choice.message.tool_calls = [tool_call] + mock_model_response.choices = [mock_choice] + + captured_payload = {} + + async def capture_post(url, headers, json, timeout): + captured_payload.update(json) + return Response( + status_code=200, + json={"deputies": {}, "findings": {}, "violations_detected": False}, + request=Request(method="POST", url=url), + ) + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + side_effect=capture_post, + ): + result = await guardrail.async_post_call_success_hook( + data=data, + user_api_key_dict=UserAPIKeyAuth(), + response=mock_model_response, + ) + + assert result == mock_model_response + assert len(captured_payload["messages"]) == 1 + assert captured_payload["messages"][0]["content"] == { + "type": "tool_use", + "id": "call_xyz", + "name": "my_tool", + "input": {"param": "value"}, + } + + @pytest.mark.asyncio + async def test_post_call_text_only_regression(self): + """Post-call: text-only response still classified correctly (regression).""" + guardrail = LassoGuardrail( + lasso_api_key="test-api-key", + guardrail_name="test-guard", + event_hook="post_call", + default_on=True, + ) + data = {"messages": [{"role": "user", "content": "Hello"}]} + + mock_model_response = MagicMock(spec=litellm.ModelResponse) + mock_choice = MagicMock() + mock_choice.message.content = "Hi! How can I help?" + mock_choice.message.tool_calls = None + mock_model_response.choices = [mock_choice] + + mock_api_response = Response( + status_code=200, + json={"deputies": {}, "findings": {}, "violations_detected": False}, + request=Request(method="POST", url="https://server.lasso.security/gateway/v3/classify"), + ) + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + return_value=mock_api_response, + ): + result = await guardrail.async_post_call_success_hook( + data=data, + user_api_key_dict=UserAPIKeyAuth(), + response=mock_model_response, + ) + + assert result == mock_model_response + + # ------------------------------------------------------------------ + # _map_masked_messages_back round-trip tests + # ------------------------------------------------------------------ + + def test_map_masked_messages_back_text(self): + """Plain text content is replaced with masked version.""" + guardrail = LassoGuardrail(lasso_api_key="test-api-key") + original = [{"role": "user", "content": "My email is john@example.com"}] + masked = [{"role": "user", "content": "My email is "}] + result = guardrail._map_masked_messages_back(original, masked) + assert result == [{"role": "user", "content": "My email is "}] + + def test_map_masked_messages_back_tool_result(self): + """Tool result content is replaced with masked version.""" + guardrail = LassoGuardrail(lasso_api_key="test-api-key") + original = [{"role": "tool", "tool_call_id": "call_abc", "content": "secret: abc123"}] + masked = [ + { + "role": "developer", + "content": {"type": "tool_result", "tool_use_id": "call_abc", "content": "secret: "}, + } + ] + result = guardrail._map_masked_messages_back(original, masked) + assert result[0]["content"] == "secret: " + + def test_map_masked_messages_back_tool_use_arguments(self): + """Assistant tool_call arguments are replaced with masked values.""" + import json as _json + guardrail = LassoGuardrail(lasso_api_key="test-api-key") + original = [ + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "send_email", "arguments": '{"to":"john@example.com"}'}} + ], + } + ] + masked = [ + { + "role": "model", + "content": {"type": "tool_use", "id": "call_1", "name": "send_email", "input": {"to": ""}}, + } + ] + result = guardrail._map_masked_messages_back(original, masked) + updated_args = _json.loads(result[0]["tool_calls"][0]["function"]["arguments"]) + assert updated_args == {"to": ""} + + def test_map_masked_messages_back_preserves_unmasked(self): + """Messages without sensitive content pass through unchanged.""" + guardrail = LassoGuardrail(lasso_api_key="test-api-key") + original = [ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "My ssn is 123-45-6789"}, + ] + masked = [ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "My ssn is "}, + ] + result = guardrail._map_masked_messages_back(original, masked) + assert result[0]["content"] == "You are helpful." + assert result[1]["content"] == "My ssn is "