diff --git a/docs/my-website/docs/proxy/guardrails/pillar_security.md b/docs/my-website/docs/proxy/guardrails/pillar_security.md index 5ab9f9bf8c..9632376768 100644 --- a/docs/my-website/docs/proxy/guardrails/pillar_security.md +++ b/docs/my-website/docs/proxy/guardrails/pillar_security.md @@ -60,6 +60,8 @@ litellm_settings: set_verbose: true # Enable detailed logging ``` +**Note:** Virtual key context is **automatically passed** as headers - no additional configuration needed! + ### 3. Start the Proxy ```bash @@ -210,7 +212,7 @@ export PILLAR_API_KEY="your_api_key_here" export PILLAR_API_BASE="https://api.pillar.security" export PILLAR_ON_FLAGGED_ACTION="monitor" export PILLAR_FALLBACK_ON_ERROR="allow" -export PILLAR_TIMEOUT="30.0" +export PILLAR_TIMEOUT="5.0" ``` ### Session Tracking diff --git a/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py b/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py index 812fc6a576..74903cc52a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py +++ b/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py @@ -90,6 +90,10 @@ class PillarGuardrail(CustomGuardrail): fallback_on_error: Action when API errors occur ('allow' or 'block') timeout: Timeout for API calls in seconds **kwargs: Additional arguments passed to parent class + + Note: + LiteLLM virtual key context (user_id, team_id, key_alias, etc.) is always + automatically passed as X-LiteLLM-* headers to enable application/user tracking. """ self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) self.api_key = api_key or os.environ.get("PILLAR_API_KEY") @@ -222,7 +226,7 @@ class PillarGuardrail(CustomGuardrail): return data verbose_proxy_logger.debug("Pillar Guardrail: Pre-call hook") - result = await self.run_pillar_guardrail(data) + result = await self.run_pillar_guardrail(data, user_api_key_dict) # Add guardrail name to response headers add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) @@ -265,7 +269,7 @@ class PillarGuardrail(CustomGuardrail): return data verbose_proxy_logger.debug("Pillar Guardrail: During-call moderation hook") - result = await self.run_pillar_guardrail(data) + result = await self.run_pillar_guardrail(data, user_api_key_dict) # Add guardrail name to response headers add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) @@ -315,7 +319,7 @@ class PillarGuardrail(CustomGuardrail): post_call_data["messages"] = data.get("messages", []) + response_messages # Reuse the existing guardrail logic - zero duplication! - await self.run_pillar_guardrail(post_call_data) + await self.run_pillar_guardrail(post_call_data, user_api_key_dict) # Add guardrail name to response headers add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) @@ -326,12 +330,13 @@ class PillarGuardrail(CustomGuardrail): # CORE LOGIC METHOD # ========================================================================= - async def run_pillar_guardrail(self, data: dict) -> dict: + async def run_pillar_guardrail(self, data: dict, user_api_key_dict: UserAPIKeyAuth) -> dict: """ Core method to run the Pillar guardrail scan. Args: data: Request data containing messages and metadata + user_api_key_dict: User API key authentication info containing key context Returns: Original data if safe or in monitor mode @@ -345,7 +350,7 @@ class PillarGuardrail(CustomGuardrail): return data try: - headers = self._prepare_headers() + headers = self._prepare_headers(user_api_key_dict) payload = self._prepare_payload(data) response = await self._call_pillar_api( @@ -403,8 +408,16 @@ class PillarGuardrail(CustomGuardrail): }, ) - def _prepare_headers(self) -> Dict[str, str]: - """Prepare headers for the Pillar API request.""" + def _prepare_headers(self, user_api_key_dict: UserAPIKeyAuth) -> Dict[str, str]: + """ + Prepare headers for the Pillar API request. + + Args: + user_api_key_dict: User API key authentication info containing key context + + Returns: + Dictionary of headers to send to Pillar API + """ if not self.api_key: msg = ( "Couldn't get Pillar API key, either set the `PILLAR_API_KEY` in the environment or " @@ -415,7 +428,7 @@ class PillarGuardrail(CustomGuardrail): headers: Dict[str, str] = { "Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json", - } + } # Add Pillar-specific headers based on configuration self._set_bool_header(headers, "plr_scanners", self.include_scanners) @@ -423,6 +436,20 @@ class PillarGuardrail(CustomGuardrail): self._set_bool_header(headers, "plr_async", self.async_mode) self._set_bool_header(headers, "plr_persist", self.persist_session) + # Always add LiteLLM virtual key context headers (metadata excluded for security) + context_mapping = { + "X-LiteLLM-Key-Name": user_api_key_dict.key_name, + "X-LiteLLM-Key-Alias": user_api_key_dict.key_alias, + "X-LiteLLM-User-Id": user_api_key_dict.user_id, + "X-LiteLLM-User-Email": user_api_key_dict.user_email, + "X-LiteLLM-Team-Id": user_api_key_dict.team_id, + "X-LiteLLM-Team-Name": user_api_key_dict.team_alias, + "X-LiteLLM-Org-Id": user_api_key_dict.org_id, + } + for header_name, value in context_mapping.items(): + if value: + headers[header_name] = str(value) + return headers def _set_bool_header(self, headers: Dict[str, str], header_name: str, value: Optional[bool]) -> None: @@ -517,6 +544,14 @@ class PillarGuardrail(CustomGuardrail): """ Prepare the payload for the Pillar API request following the /api/v1/protect contract. + This method supports multi-modal content (images, files, audio, video, etc.) as messages + are passed through without modification. The messages array can contain any OpenAI-compatible + message structure including: + - Text content (string) + - Multi-modal content blocks (image_url, image_file, audio, video, document, file) + - Attachments + - Tool calls + Args: data: Request data diff --git a/tests/test_litellm/proxy/guardrails/test_pillar_guardrails.py b/tests/test_litellm/proxy/guardrails/test_pillar_guardrails.py index 9aecd01dc5..99a51d20a7 100644 --- a/tests/test_litellm/proxy/guardrails/test_pillar_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_pillar_guardrails.py @@ -234,6 +234,22 @@ def pillar_async_response(): ) +@pytest.fixture +def user_api_key_dict_with_context(): + """Fixture providing UserAPIKeyAuth with complete context.""" + return UserAPIKeyAuth( + token="hashed-test-token", + key_name="production-api-key", + key_alias="prod-key", + user_id="user-123", + user_email="test@example.com", + team_id="team-456", + team_alias="engineering-team", + org_id="org-789", + metadata={"environment": "production", "region": "us-east-1"}, + ) + + @pytest.fixture def mock_llm_response_with_tools(): """Fixture providing a mock LLM response with tool calls.""" @@ -502,6 +518,217 @@ async def test_pre_call_hook_custom_header_overrides( assert captured_headers.get("plr_evidence") == "false" +# ========================================================================= +# LITELLM KEY CONTEXT HEADER TESTS +# ========================================================================= + + +@pytest.mark.asyncio +async def test_litellm_context_headers_automatically_added( + sample_request_data, + user_api_key_dict_with_context, + dual_cache, + pillar_clean_response, +): + """Test that LiteLLM context headers are automatically added (always enabled).""" + guardrail = PillarGuardrail( + guardrail_name="pillar-context-enabled", + api_key="test-pillar-key", + api_base="https://api.pillar.security", + ) + + captured_headers: Dict[str, str] = {} + + async def _mock_post(*args, **kwargs): + captured_headers.update(kwargs.get("headers", {})) + return pillar_clean_response + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new=_mock_post, + ): + await guardrail.async_pre_call_hook( + data=sample_request_data, + cache=dual_cache, + user_api_key_dict=user_api_key_dict_with_context, + call_type="completion", + ) + + # Verify LiteLLM context headers are present + assert "X-LiteLLM-Key-Name" in captured_headers + assert captured_headers["X-LiteLLM-Key-Name"] == "production-api-key" + assert "X-LiteLLM-Key-Alias" in captured_headers + assert captured_headers["X-LiteLLM-Key-Alias"] == "prod-key" + assert "X-LiteLLM-User-Id" in captured_headers + assert captured_headers["X-LiteLLM-User-Id"] == "user-123" + assert "X-LiteLLM-User-Email" in captured_headers + assert captured_headers["X-LiteLLM-User-Email"] == "test@example.com" + assert "X-LiteLLM-Team-Id" in captured_headers + assert captured_headers["X-LiteLLM-Team-Id"] == "team-456" + assert "X-LiteLLM-Team-Name" in captured_headers + assert captured_headers["X-LiteLLM-Team-Name"] == "engineering-team" + assert "X-LiteLLM-Org-Id" in captured_headers + assert captured_headers["X-LiteLLM-Org-Id"] == "org-789" + + # Metadata is NOT sent (may contain sensitive information) + assert "X-LiteLLM-Metadata" not in captured_headers + + +@pytest.mark.asyncio +async def test_litellm_context_with_partial_fields( + sample_request_data, + dual_cache, + pillar_clean_response, +): + """Test that partial LiteLLM context (only some fields present) is handled correctly.""" + # Create UserAPIKeyAuth with only some fields populated + partial_context = UserAPIKeyAuth( + user_id="user-only", + team_id="team-only", + ) + + guardrail = PillarGuardrail( + guardrail_name="pillar-partial-context", + api_key="test-pillar-key", + api_base="https://api.pillar.security", + pass_litellm_key_header=True, + ) + + captured_headers: Dict[str, str] = {} + + async def _mock_post(*args, **kwargs): + captured_headers.update(kwargs.get("headers", {})) + return pillar_clean_response + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new=_mock_post, + ): + await guardrail.async_pre_call_hook( + data=sample_request_data, + cache=dual_cache, + user_api_key_dict=partial_context, + call_type="completion", + ) + + # Verify only populated fields are present + assert "X-LiteLLM-User-Id" in captured_headers + assert captured_headers["X-LiteLLM-User-Id"] == "user-only" + assert "X-LiteLLM-Team-Id" in captured_headers + assert captured_headers["X-LiteLLM-Team-Id"] == "team-only" + + # Verify empty fields are not present + assert "X-LiteLLM-Key-Name" not in captured_headers + assert "X-LiteLLM-User-Email" not in captured_headers + + +# ========================================================================= +# MULTI-MODAL CONTENT TESTS +# ========================================================================= + + +@pytest.mark.asyncio +async def test_multimodal_image_url_support( + user_api_key_dict, + dual_cache, + pillar_clean_response, +): + """Test that messages with image URLs are properly handled.""" + multimodal_data = { + "model": "gpt-4-vision-preview", + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What's in this image?"}, + { + "type": "image_url", + "image_url": { + "url": "https://example.com/image.jpg", + "detail": "high", + }, + }, + ], + } + ], + } + + guardrail = PillarGuardrail( + guardrail_name="pillar-multimodal", + api_key="test-pillar-key", + api_base="https://api.pillar.security", + ) + + captured_payload: Dict[str, Any] = {} + + async def _mock_post(*args, **kwargs): + captured_payload.update(kwargs.get("json", {})) + return pillar_clean_response + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new=_mock_post, + ): + result = await guardrail.async_pre_call_hook( + data=multimodal_data, + cache=dual_cache, + user_api_key_dict=user_api_key_dict, + call_type="completion", + ) + + # Verify multimodal message structure is preserved + assert result == multimodal_data + assert "messages" in captured_payload + assert len(captured_payload["messages"]) == 1 + assert isinstance(captured_payload["messages"][0]["content"], list) + assert captured_payload["messages"][0]["content"][1]["type"] == "image_url" + + +@pytest.mark.asyncio +async def test_multimodal_with_attachments( + user_api_key_dict, + dual_cache, + pillar_clean_response, +): + """Test that messages with file attachments are properly handled.""" + multimodal_data = { + "model": "gpt-4", + "messages": [ + { + "role": "user", + "content": "Analyze this document", + "attachments": [ + { + "file_id": "file-abc123", + "tools": [{"type": "code_interpreter"}], + } + ], + } + ], + } + + guardrail = PillarGuardrail( + guardrail_name="pillar-attachments", + api_key="test-pillar-key", + api_base="https://api.pillar.security", + ) + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + return_value=pillar_clean_response, + ): + result = await guardrail.async_pre_call_hook( + data=multimodal_data, + cache=dual_cache, + user_api_key_dict=user_api_key_dict, + call_type="completion", + ) + + # Verify attachment structure is preserved + assert result == multimodal_data + assert result["messages"][0]["attachments"] is not None + + # ============================================================================ # EDGE CASE TESTS # ============================================================================