diff --git a/docs/my-website/docs/proxy/guardrails/lasso_security.md b/docs/my-website/docs/proxy/guardrails/lasso_security.md index 21528790af..113e3f8974 100644 --- a/docs/my-website/docs/proxy/guardrails/lasso_security.md +++ b/docs/my-website/docs/proxy/guardrails/lasso_security.md @@ -35,7 +35,7 @@ guardrails: guardrail: lasso mode: "pre_call" api_key: os.environ/LASSO_API_KEY - api_base: "https://server.lasso.security" + api_base: "https://server.lasso.security/gateway/v3" - guardrail_name: "lasso-post-guard" litellm_params: guardrail: lasso @@ -228,7 +228,7 @@ Expected response: ## PII Masking with Lasso -Lasso supports automatic PII detection and masking using the `/gateway/v1/classifix` endpoint. When enabled, sensitive information like emails, phone numbers, and other PII will be automatically masked with appropriate placeholders. +Lasso supports automatic PII detection and masking using the `/classifix` endpoint. When enabled, sensitive information like emails, phone numbers, and other PII will be automatically masked with appropriate placeholders. ### Enabling PII Masking diff --git a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py index 99d2b82400..ea8f1b0a97 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py @@ -33,6 +33,8 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.integrations.custom_guardrail import dc as global_cache + from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -100,7 +102,7 @@ class LassoGuardrail(CustomGuardrail): ) self.api_base = ( - api_base or os.getenv("LASSO_API_BASE") or "https://server.lasso.security" + api_base or os.getenv("LASSO_API_BASE") or "https://server.lasso.security/gateway/v3" ) verbose_proxy_logger.debug( @@ -125,7 +127,7 @@ class LassoGuardrail(CustomGuardrail): async def async_pre_call_hook( self, user_api_key_dict: UserAPIKeyAuth, - cache: DualCache, + cache: DualCache, # Deprecated, use global_cache instead (kept to align with CustomGuardrail interface) data: dict, call_type: Literal[ "completion", @@ -150,10 +152,10 @@ class LassoGuardrail(CustomGuardrail): return data # Get or generate conversation_id and store it in data for post-call consistency - conversation_id = self._get_or_generate_conversation_id(data, cache) - data.setdefault("_lasso_internal", {})["conversation_id"] = conversation_id + # The conversation_id is being stored in the cache so it can be used by the post_call hook + self._get_or_generate_conversation_id(data, global_cache) - return await self._run_lasso_guardrail(data, cache, message_type="PROMPT") + return await self._run_lasso_guardrail(data, global_cache, message_type="PROMPT") @log_guardrail_information async def async_moderation_hook( @@ -213,17 +215,12 @@ class LassoGuardrail(CustomGuardrail): "litellm_call_id": data.get("litellm_call_id"), } - # Copy stored conversation_id from pre-call hook - if data.get("_lasso_internal", {}).get("conversation_id") and isinstance(response_data, dict): - response_data.setdefault("_lasso_internal", {})["conversation_id"] = data["_lasso_internal"][ - "conversation_id" - ] # Handle masking for post-call if self.mask: - headers = self._prepare_headers(response_data) - payload = self._prepare_payload(response_messages, "COMPLETION", response_data) - api_url = f"{self.api_base}/gateway/v3/classifix" + headers = self._prepare_headers(response_data, global_cache) + payload = self._prepare_payload(response_messages, response_data, global_cache, "COMPLETION") + api_url = f"{self.api_base}/classifix" try: lasso_response = await self._call_lasso_api(headers=headers, payload=payload, api_url=api_url) @@ -241,7 +238,7 @@ class LassoGuardrail(CustomGuardrail): raise LassoGuardrailAPIError(f"Failed to apply post-call masking: {str(e)}") else: # Use the same data for conversation_id consistency (no cache access needed) - await self._run_lasso_guardrail(response_data, cache=None, message_type="COMPLETION") + await self._run_lasso_guardrail(response_data, cache=global_cache, message_type="COMPLETION") verbose_proxy_logger.debug("Post-call Lasso validation completed") else: verbose_proxy_logger.warning("No response messages found to validate") @@ -306,7 +303,7 @@ class LassoGuardrail(CustomGuardrail): async def _run_lasso_guardrail( self, data: dict, - cache: Optional[DualCache], + cache: DualCache, message_type: Literal["PROMPT", "COMPLETION"] = "PROMPT", ): """ @@ -345,14 +342,14 @@ class LassoGuardrail(CustomGuardrail): async def _handle_classification( self, data: dict, - cache: Optional[DualCache], + cache: DualCache, message_type: Literal["PROMPT", "COMPLETION"], messages: List[Dict[str, str]], ) -> dict: """Handle classification without masking.""" try: headers = self._prepare_headers(data, cache) - payload = self._prepare_payload(messages, message_type, data, cache) + payload = self._prepare_payload(messages, data, cache, message_type) response = await self._call_lasso_api(headers=headers, payload=payload) self._process_lasso_response(response) return data @@ -363,15 +360,15 @@ class LassoGuardrail(CustomGuardrail): async def _handle_masking( self, data: dict, - cache: Optional[DualCache], + cache: DualCache, message_type: Literal["PROMPT", "COMPLETION"], messages: List[Dict[str, str]], ) -> dict: """Handle masking with classifix endpoint.""" try: headers = self._prepare_headers(data, cache) - payload = self._prepare_payload(messages, message_type, data, cache) - api_url = f"{self.api_base}/gateway/v3/classifix" + payload = self._prepare_payload(messages, data, cache, message_type) + api_url = f"{self.api_base}/classifix" response = await self._call_lasso_api(headers=headers, payload=payload, api_url=api_url) self._process_lasso_response(response) @@ -437,7 +434,7 @@ class LassoGuardrail(CustomGuardrail): }, ) - def _prepare_headers(self, data: dict, cache: Optional[DualCache] = None) -> Dict[str, str]: + def _prepare_headers(self, data: dict, cache: DualCache) -> Dict[str, str]: """Prepare headers for the Lasso API request.""" if not self.lasso_api_key: raise LassoGuardrailMissingSecrets( @@ -455,13 +452,7 @@ class LassoGuardrail(CustomGuardrail): headers["lasso-user-id"] = self.user_id # Always include conversation_id (generated or provided) - if cache is not None: - conversation_id = self._get_or_generate_conversation_id(data, cache) - else: - # For post-call hook, use stored conversation_id or generate a new one - conversation_id = ( - data.get("_lasso_internal", {}).get("conversation_id") or self.conversation_id or self._generate_ulid() - ) + conversation_id = self._get_or_generate_conversation_id(data, cache) headers["lasso-conversation-id"] = conversation_id @@ -470,9 +461,9 @@ class LassoGuardrail(CustomGuardrail): def _prepare_payload( self, messages: List[Dict[str, str]], + data: dict, + cache: DualCache, message_type: Literal["PROMPT", "COMPLETION"] = "PROMPT", - data: Optional[dict] = None, - cache: Optional[DualCache] = None, ) -> Dict[str, Any]: """ Prepare the payload for the Lasso API request. @@ -490,20 +481,9 @@ class LassoGuardrail(CustomGuardrail): payload["userId"] = self.user_id # Always include sessionId (conversation_id - generated or provided) - if data is not None: - if cache is not None: - conversation_id = self._get_or_generate_conversation_id(data, cache) - else: - # For post-call hook, use stored conversation_id or fallback - conversation_id = ( - data.get("_lasso_internal", {}).get("conversation_id") - or self.conversation_id - or self._generate_ulid() - ) + conversation_id = self._get_or_generate_conversation_id(data, cache) - payload["sessionId"] = conversation_id - elif self.conversation_id: - payload["sessionId"] = self.conversation_id + payload["sessionId"] = conversation_id return payload @@ -514,7 +494,7 @@ class LassoGuardrail(CustomGuardrail): api_url: Optional[str] = None, ) -> LassoResponse: """Call the Lasso API and return the response.""" - url = api_url or f"{self.api_base}/gateway/v3/classify" + url = api_url or f"{self.api_base}/classify" verbose_proxy_logger.debug(f"Calling Lasso API with messageType: {payload.get('messageType')}") response = await self.async_handler.post( url=url, 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 c63974ac3f..87542c974a 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lasso.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lasso.py @@ -1,6 +1,7 @@ import os import sys import pytest +import uuid from unittest.mock import patch, MagicMock from httpx import Response, Request from fastapi import HTTPException @@ -77,10 +78,11 @@ class TestLassoGuardrail: assert guardrail.lasso_api_key == "test-api-key" assert guardrail.user_id == "test-user" assert guardrail.conversation_id == "test-conversation" - assert guardrail.api_base == "https://server.lasso.security" + assert guardrail.api_base == "https://server.lasso.security/gateway/v3" @pytest.mark.asyncio async def test_pre_call_no_violations(self): + from litellm.integrations.custom_guardrail import dc as global_cache """Test pre-call hook with no violations detected.""" # Setup guardrail guardrail = LassoGuardrail( @@ -90,12 +92,16 @@ class TestLassoGuardrail: default_on=True ) + test_call_id = str(uuid.uuid4()) + assert global_cache.get_cache(f"lasso_conversation_id:{test_call_id}") is None + # Test data data = { "messages": [ {"role": "user", "content": "Hello, how are you?"} ], - "metadata": {} + "metadata": {}, + "litellm_call_id": test_call_id } # Mock successful API response with no violations @@ -118,13 +124,14 @@ class TestLassoGuardrail: request=Request(method="POST", url="https://server.lasso.security/gateway/v3/classify"), ) + local_cache = DualCache() with patch( "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", return_value=mock_response ): result = await guardrail.async_pre_call_hook( user_api_key_dict=UserAPIKeyAuth(), - cache=DualCache(), + cache=local_cache, data=data, call_type="completion" ) @@ -132,6 +139,11 @@ class TestLassoGuardrail: # Should return original data when no violations detected assert result == data + # Verify that the conversation_id is stored in the global cache but not the local cache + cache_key = f"lasso_conversation_id:{test_call_id}" + assert global_cache.get_cache(cache_key) is not None + assert local_cache.get_cache(cache_key) is None + @pytest.mark.asyncio async def test_pre_call_with_violations(self): """Test pre-call hook with violations detected.""" @@ -466,9 +478,10 @@ class TestLassoGuardrail: ) messages = [{"role": "user", "content": "Test message"}] + cache = DualCache() # Test PROMPT payload - prompt_payload = guardrail._prepare_payload(messages, "PROMPT") + prompt_payload = guardrail._prepare_payload(messages, {}, cache, "PROMPT") assert prompt_payload["messageType"] == "PROMPT" assert prompt_payload["messages"] == messages assert prompt_payload["userId"] == "test-user" @@ -476,7 +489,7 @@ class TestLassoGuardrail: # Test COMPLETION payload completion_messages = [{"role": "assistant", "content": "Test response"}] - completion_payload = guardrail._prepare_payload(completion_messages, "COMPLETION") + completion_payload = guardrail._prepare_payload(completion_messages, {}, cache, "COMPLETION") assert completion_payload["messageType"] == "COMPLETION" assert completion_payload["messages"] == completion_messages assert completion_payload["userId"] == "test-user" @@ -489,9 +502,9 @@ class TestLassoGuardrail: user_id="test-user", conversation_id="test-conversation" ) - + cache = DualCache() data = {"litellm_call_id": "test-call-id"} - headers = guardrail._prepare_headers(data) + headers = guardrail._prepare_headers(data, cache) assert headers["lasso-api-key"] == "test-api-key" assert headers["Content-Type"] == "application/json" assert headers["lasso-user-id"] == "test-user" @@ -499,7 +512,7 @@ class TestLassoGuardrail: # Test without optional fields guardrail_minimal = LassoGuardrail(lasso_api_key="test-api-key") - headers_minimal = guardrail_minimal._prepare_headers(data) + headers_minimal = guardrail_minimal._prepare_headers(data, cache) assert headers_minimal["lasso-api-key"] == "test-api-key" assert headers_minimal["Content-Type"] == "application/json" assert "lasso-user-id" not in headers_minimal