Merge pull request trustyai-explainability#54 from trustyai-explainability/AddHeaderPassthrough

Feat: Add header extraction from shield params

Author: RobGeada

llama_stack_provider_trustyai_fms/detectors/base.py

@@ -325,7 +325,7 @@ def _extract_detector_params(self) -> dict[str, Any]:
 
         return detector_params
 
-    def _prepare_headers(self) -> Headers:
+    def _prepare_headers(self, params: dict[str, Any] = None) -> Headers:
         """Prepare request headers based on configuration"""
         headers: Headers = {
             "accept": "application/json",
@@ -342,6 +342,11 @@ def _prepare_headers(self) -> Headers:
         elif self.config.auth_token:
             headers["Authorization"] = f"Bearer {self.config.auth_token}"
 
+        if params is not None:
+            for k, v in params.items():
+                if k.lower() == "headers" and isinstance(v, dict):
+                    headers.update(v)
+
         return headers
 
     def _prepare_request_payload(

llama_stack_provider_trustyai_fms/detectors/chat.py

@@ -199,7 +199,7 @@ async def _call_detector_api(
         """Call chat detector API with proper endpoint selection"""
         try:
             request = self._prepare_chat_request(messages, params)
-            headers = self._prepare_headers()
+            headers = self._prepare_headers(params)
 
             logger.info("Making detector API request")
             logger.debug(f"Request headers: {headers}")

llama_stack_provider_trustyai_fms/detectors/content.py

@@ -121,7 +121,7 @@ async def _call_detector_api(
         """Call detector API with proper endpoint selection"""
         try:
             request = self._prepare_content_request(content, params)
-            headers = self._prepare_headers()
+            headers = self._prepare_headers(params)
 
             logger.info("Making detector API request")
             logger.debug(f"Request headers: {headers}")

tests/unit/test_detectors.py

@@ -285,6 +285,32 @@ async def test_http_request_success(self, initialized_detector, httpx_mock):
             result.violation is not None
         )  # Should detect violation because score > threshold
 
+    @pytest.mark.asyncio
+    async def test_headers_extracted_from_params(self, initialized_detector):
+        detector = initialized_detector
+
+        # Patch _make_request to capture headers
+        called = {}
+
+        async def fake_make_request(request, headers, timeout=None):
+            called["headers"] = headers
+            # Return a valid detection so the rest of the code works
+            return [{"score": 0.9, "label": "label", "detection_type": "content"}]
+
+        detector._make_request = fake_make_request
+
+        messages = [UserMessage(content="Test content", role="user")]
+        params = {"headers": {"X-Test-Header": "header-value", "another": "foo"}}
+        await detector.run_shield(
+            shield_id=detector.config.detector_id, messages=messages, params=params
+        )
+
+        # Check that headers were extracted and passed
+        assert "X-Test-Header" in called["headers"]
+        assert called["headers"]["X-Test-Header"] == "header-value"
+        assert "another" in called["headers"]
+        assert called["headers"]["another"] == "foo"
+
     @pytest.mark.asyncio
     async def test_http_request_no_violation(self, initialized_detector, httpx_mock):
         detector = initialized_detector