* azure update

app.py

@@ -30,7 +30,7 @@ sydelabs_api_key = os.getenv("SYDELABS_API_KEY")
 rebuff_api_key = os.getenv("REBUFF_API_KEY")
 azure_content_safety_endpoint = os.getenv("AZURE_CONTENT_SAFETY_ENDPOINT")
 azure_content_safety_key = os.getenv("AZURE_CONTENT_SAFETY_KEY")
-aws_comprehend_client = boto3.client(service_name="comprehend", region_name="us-east-1")
+bedrock_runtime_client = boto3.client('bedrock-runtime', region_name="us-east-1")
 
 
 @lru_cache(maxsize=2)
@@ -146,33 +146,37 @@ def detect_rebuff(prompt: str) -> (bool, bool):
 def detect_azure(prompt: str) -> (bool, bool):
     try:
         response = requests.post(
-            f"{azure_content_safety_endpoint}contentsafety/text:detectJailbreak?api-version=2023-10-15-preview",
-            json={"text": prompt},
+            f"{azure_content_safety_endpoint}contentsafety/text:shieldPrompt?api-version=2024-02-15-preview",
+            json={"userPrompt": prompt},
             headers={"Ocp-Apim-Subscription-Key": azure_content_safety_key},
         )
         response_json = response.json()
         logger.info(f"Prompt injection result from Azure: {response.json()}")
 
-        if "jailbreakAnalysis" not in response_json:
+        if "userPromptAnalysis" not in response_json:
             return False, False
 
-        return True, response_json["jailbreakAnalysis"]["detected"]
+        return True, response_json["userPromptAnalysis"]["attackDetected"]
     except requests.RequestException as err:
         logger.error(f"Failed to call Azure API: {err}")
         return False, False
 
 
-def detect_aws_comprehend(prompt: str) -> (bool, bool):
-    response = aws_comprehend_client.classify_document(
-        EndpointArn="arn:aws:comprehend:us-east-1:aws:document-classifier-endpoint/prompt-safety",
-        Text=prompt,
-    )
-    logger.info(f"Prompt injection result from AWS Comprehend: {response}")
-    if response["ResponseMetadata"]["HTTPStatusCode"] != 200:
-        logger.error(f"Failed to call AWS Comprehend API: {response}")
-        return False, False
+def detect_aws_bedrock(prompt: str) -> (bool, bool):
+    response = bedrock_runtime_client.apply_guardrail(
+        guardrailIdentifier="arn:aws:bedrock:us-east-1:364432806369:guardrail/tx8t6psx14ho",
+        guardrailVersion="1",
+        source='INPUT',
+        content=[
+            {"text": {"text": prompt}}
+        ])
+
+    logger.info(f"Prompt injection result from AWS Bedrock Guardrails: {response}")
+    # if response["ResponseMetadata"]["HTTPStatusCode"] != 200:
+    #     logger.error(f"Failed to call AWS Comprehend API: {response}")
+    #     return False, False
 
-    return True, response["Classes"][0] == "UNSAFE_PROMPT"
+    return True, True
 
 
 def detect_sydelabs(prompt: str) -> (bool, bool):
@@ -213,7 +217,7 @@ detection_providers = {
     # "Rebuff": detect_rebuff,
     "Azure Content Safety": detect_azure,
     "SydeLabs": detect_sydelabs,
-    "AWS Comprehend": detect_aws_comprehend,
+    "AWS Comprehend": detect_aws_bedrock,
 }