Fix session consistency, move Lasso API version away from source code (#17316)

* store and fetch lasso-conversation id from cache

* include gateway/v# in the baseUrl to allow simpler version migrations in the future

* add tests for cached conversation ID

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
 

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,

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