Merge pull request #27 from ZenGuard-AI/feature/add-toxicity-api

enhancement - add toxicity api to client

Author: Baur

tests/zenguard_e2e_test.py

@@ -1,4 +1,3 @@
-
 from zenguard.zenguard import Detector
 
 
@@ -42,8 +41,16 @@ def test_keywords(zenguard):
     response = zenguard.detect(detectors=detectors, prompt=prompt)
     assert_successful_response_not_detected(response)
 
+
 def test_secrets(zenguard):
     prompt = "Simple secrets test"
     detectors = [Detector.SECRETS]
     response = zenguard.detect(detectors=detectors, prompt=prompt)
     assert_successful_response_not_detected(response)
+
+
+def test_toxicity(zenguard):
+    prompt = "Simple toxicity test"
+    detectors = [Detector.TOXICITY]
+    response = zenguard.detect(detectors=detectors, prompt=prompt)
+    assert_successful_response_not_detected(response)

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(Enum):
     BANNED_TOPICS = "v1/detect/topics/banned"
     KEYWORDS = "v1/detect/keywords"
     SECRETS = "v1/detect/secrets"
+    TOXICITY = "v1/detect/toxicity"
 
 
 class Endpoint(Enum):