Merge remote-tracking branch 'origin/main' into feature/detect-in-parallel

Author: kenesaryy

pyproject.toml

@@ -1,7 +1,7 @@
 [tool.poetry]
 name = "zenguard"
-version = "0.1.10"
-description = "Plug-and-play production grade security for GenAI applications"
+version = "0.1.12"
+description = "Fast production grade security for GenAI applications"
 authors = ["ZenGuard Team <hello@zenguard.ai>"]
 license = "MIT"
 readme = "README.md"

tests/zenguard_e2e_test.py

@@ -1,4 +1,3 @@
-
 from zenguard.zenguard import Detector
 
 
@@ -48,6 +47,7 @@ def test_keywords(zenguard):
     response = zenguard.detect(detectors=detectors, prompt=prompt)
     assert_successful_response_not_detected(response, detectors)
 
+
 def test_secrets(zenguard):
     prompt = "Simple secrets test"
     detectors = [Detector.SECRETS]
@@ -77,3 +77,10 @@ def test_detect_in_parallel_pass_on_detectors(zenguard):
     prompt = "Simple in parallel test"
     response = zenguard.detect(detectors, prompt)
     assert_successful_response_not_detected(response, detectors)
+
+
+def test_toxicity(zenguard):
+    prompt = "Simple toxicity test"
+    detectors = [Detector.TOXICITY]
+    response = zenguard.detect(detectors=detectors, prompt=prompt)
+    assert_successful_response_not_detected(response, detectors)

zenguard/zenguard.py

@@ -11,13 +11,8 @@
 from tqdm import tqdm
 
 from zenguard.ai_clients.openai import ChatWithZenguard
-from zenguard.pentest.prompt_injections import (
-    config,
-    prompting,
-    run,
-    scoring,
-    visualization,
-)
+from zenguard.pentest.prompt_injections import (config, prompting, run,
+                                                scoring, visualization)
 
 
 class SupportedLLMs:
@@ -44,6 +39,7 @@ class Detector(str, Enum):
     KEYWORDS = "keywords"
     PII = "pii"
     SECRETS = "secrets"
+    TOXICITY = "toxicity"
 
 
 class Endpoint(Enum):