Merge pull request #39 from ZenGuard-AI/baur/branch

[feature] add async detect and reporting for prompt_injection detector

Author: Baur

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "zenguard"
-version = "0.1.15"
+version = "0.1.16"
 description = "Fast production grade security for GenAI applications"
 authors = ["ZenGuard Team <hello@zenguard.ai>"]
 license = "MIT"

tests/zenguard_e2e_test.py

@@ -1,4 +1,9 @@
-from zenguard.zenguard import Detector
+from unittest.mock import Mock, patch
+
+import httpx
+import pytest
+
+from zenguard.zenguard import API_REPORT_PROMPT_INJECTIONS, Detector
 
 
 def assert_successful_response_not_detected(response):
@@ -74,14 +79,14 @@ def test_update_detectors(zenguard):
     assert response is None
 
 
-def test_detect_in_parallel(zenguard):
+def test_detect_in_parallel_error_no_detectors(zenguard):
     detectors = [Detector.SECRETS, Detector.ALLOWED_TOPICS]
     response = zenguard.update_detectors(detectors=detectors)
     assert response is None
 
     prompt = "Simple in parallel test"
-    response = zenguard.detect([], prompt)
-    assert "No detectors" in response["error"]
+    with pytest.raises(ValueError):
+        response = zenguard.detect([], prompt)
 
 
 def test_detect_in_parallel_pass_on_detectors(zenguard):
@@ -91,3 +96,50 @@ def test_detect_in_parallel_pass_on_detectors(zenguard):
     response = zenguard.detect(detectors, prompt)
     assert_detectors_response(response, detectors)
     assert "error" not in response
+
+
+def test_prompt_injection_async(zenguard):
+    prompt = "Simple prompt injection test"
+    detectors = [Detector.PROMPT_INJECTION]
+    zenguard.detect_async(detectors=detectors, prompt=prompt)
+
+
+def test_detect_error_no_detectors(zenguard):
+    prompt = "Simple prompt injection test"
+    with pytest.raises(ValueError):
+        zenguard.detect_async([], prompt)
+
+
+def test_report_with_valid_detector_and_days(zenguard):
+    with patch("httpx.post") as mock_post:
+        mock_response = Mock()
+        # TODO(baur): Update this to the actual response
+        mock_response.json.return_value = {"prompt_injections": 10}
+        mock_post.return_value = mock_response
+
+        result = zenguard.report(detector=Detector.PROMPT_INJECTION, days=7)
+
+        assert result == {"prompt_injections": 10}
+        mock_post_args, mock_post_kwargs = mock_post.call_args
+
+        # Assert only the relevant parts of the API call
+        assert API_REPORT_PROMPT_INJECTIONS in mock_post_args[0]
+        assert mock_post_kwargs["json"] == {"days": 7}
+
+
+def test_report_with_invalid_detector(zenguard):
+    with pytest.raises(ValueError):
+        zenguard.report(detector=Detector.PII, days=7)
+
+
+def test_report_with_request_error(zenguard):
+    with patch("httpx.post") as mock_post:
+        mock_post.side_effect = httpx.RequestError("Connection error")
+
+        with pytest.raises(RuntimeError) as exc_info:
+            zenguard.report(detector=Detector.PROMPT_INJECTION)
+
+        assert (
+            str(exc_info.value)
+            == "An error occurred while making the request: Connection error"
+        )

zenguard/zenguard.py

@@ -19,6 +19,8 @@
     visualization,
 )
 
+API_REPORT_PROMPT_INJECTIONS = "v1/report/prompt_injections"
+
 
 class SupportedLLMs:
     CHATGPT = "chatgpt"
@@ -99,8 +101,11 @@ def __init__(self, config: ZenGuardConfig):
             raise ValueError(f"LLM {config.llm} is not supported")
 
     def detect(self, detectors: list[Detector], prompt: str):
+        """
+        Uses detectors to evaluate the prompt and return the results.
+        """
         if len(detectors) == 0:
-            return {"error": "No detectors were provided"}
+            raise ValueError("No detectors were provided")
 
         url = self._backend
         if len(detectors) == 1:
@@ -118,15 +123,48 @@ def detect(self, detectors: list[Detector], prompt: str):
                 timeout=20,
             )
         except httpx.RequestError as e:
-            print(e)
-            return {"error": str(e)}
-
-        if response.status_code != 200:
-            print(response.json())
-            return {"error": str(response.json())}
+            raise RuntimeError(
+                f"An error occurred while making the request: {str(e)}"
+            ) from e
+        except httpx.HTTPStatusError as e:
+            raise RuntimeError(
+                f"Received an unexpected status code: {response.status_code}\nResponse content: {response.json()}"
+            ) from e
 
         return response.json()
 
+    def detect_async(self, detectors: list[Detector], prompt: str):
+        """
+        Same as detect function but asynchroneous.
+        """
+        if len(detectors) == 0:
+            raise ValueError("No detectors were provided")
+
+        if detectors[0] != Detector.PROMPT_INJECTION:
+            raise ValueError(
+                "Only Prompt Injection detector is supported for async detection"
+            )
+
+        url = self._backend + convert_detector_to_api(detectors[0]) + "_async"
+        json = {"messages": [prompt]}
+
+        try:
+            response = httpx.post(
+                url,
+                json=json,
+                headers={"x-api-key": self._api_key},
+                timeout=20,
+            )
+            response.raise_for_status()
+        except httpx.RequestError as e:
+            raise RuntimeError(
+                f"An error occurred while making the request: {str(e)}"
+            ) from e
+        except httpx.HTTPStatusError as e:
+            raise RuntimeError(
+                f"Received an unexpected status code: {response.status_code}\nResponse content: {response.json()}"
+            ) from e
+
     def _attack_zenguard(self, detector: Detector, attacks: list[str]):
         attacks = tqdm(attacks)
         for attack in attacks:
@@ -171,3 +209,40 @@ def update_detectors(self, detectors: list[Detector]):
 
         if response.status_code != 200:
             return {"error": str(response.json())}
+
+    def report(self, detector: Detector, days: int = None):
+        """
+        Get a report of the detections made by the detector in the last days.
+        Days is optional and if not provided, it will return all the detections.
+        Days is int and will give back the number of detections made in the last days.
+        """
+
+        if detector != Detector.PROMPT_INJECTION:
+            raise ValueError(
+                "Only Prompt Injection detector is currently supported for reports"
+            )
+
+        json = {}
+        if days:
+            json = {"days": days}
+
+        url = self._backend + API_REPORT_PROMPT_INJECTIONS
+
+        try:
+            response = httpx.post(
+                url,
+                json=json,
+                headers={"x-api-key": self._api_key},
+                timeout=20,
+            )
+            response.raise_for_status()
+        except httpx.RequestError as e:
+            raise RuntimeError(
+                f"An error occurred while making the request: {str(e)}"
+            ) from e
+        except httpx.HTTPStatusError as e:
+            raise RuntimeError(
+                f"Received an unexpected status code: {response.status_code}\nResponse content: {response.text}"
+            ) from e
+
+        return response.json()