Update handler.py

handler.py

@@ -117,13 +117,8 @@ class EndpointHandler:
         Provide your safety assessment for the above conversation in the below output JSON format and must provide all the metrics mentioned below:
         
         "Prompt Safety": safe or unsafe rating for the user message.
-        "Response Safety": safe or unsafe rating for the agent response. Omit if no agent response present.
-        
         "Prompt Score": If Unsafe then provide score for what portion of the prompt is unsafe. Score need to be on scale of 1-5.
-        "Response Score": If Unsafe then provide score mentioned which portion of the response is unsafe. Score need to be on scale of 1-5.
-        
         "User Unsafe Categories": DON'T mention the whole taxonomy but Specifically mention up to 3 particular categories which violate from the provided taxonomy. Omit if all safe.
-        "Agent Unsafe Categories": DON'T mention the whole taxonomy but Specifically mention up to 3 particular categories which violate from the provided taxonomy. Omit if all safe. Omit if no agent response present.
         
         Do not include anything other than the output JSON in your response.
         Output JSON:
@@ -145,55 +140,78 @@ class EndpointHandler:
                 ]
             }
 
-    def extract_and_parse_json(self, response):
-        # First try to extract content inside triple backticks
-        match = re.search(r'```(?:json)?\s*(.*?)\s*```', response, re.DOTALL)
-        content = match.group(1).strip() if match else response.strip()
+    # def extract_and_parse_json(self, response):
+    #     # First try to extract content inside triple backticks
+    #     match = re.search(r'```(?:json)?\s*(.*?)\s*```', response, re.DOTALL)
+    #     content = match.group(1).strip() if match else response.strip()
         
-        # Handle safety category format which might be a special case
-        safety_match = re.search(r'"Prompt Safety":\s*"([^"]+)",\s*"Response Safety":\s*"([^"]+)",\s*"Prompt Score":\s*"?([^",]+)"?,\s*"Response Score":\s*"?([^",]+)"?,\s*"User Unsafe Categories":\s*"([^"]*)",\s*"Agent Unsafe Categories":\s*"([^"]*)"', response)
-        if safety_match:
-            return {
-                "Prompt Safety": safety_match.group(1),
-                "Response Safety": safety_match.group(2),
-                "Prompt Score": safety_match.group(3),
-                "Response Score": safety_match.group(4),
-                "User Unsafe Categories": safety_match.group(5),
-                "Agent Unsafe Categories": safety_match.group(6)
-            }
+    #     # Handle safety category format which might be a special case
+    #     safety_match = re.search(r'"Prompt Safety":\s*"([^"]+)",\s*"Response Safety":\s*"([^"]+)",\s*"Prompt Score":\s*"?([^",]+)"?,\s*"Response Score":\s*"?([^",]+)"?,\s*"User Unsafe Categories":\s*"([^"]*)",\s*"Agent Unsafe Categories":\s*"([^"]*)"', response)
+    #     if safety_match:
+    #         return {
+    #             "Prompt Safety": safety_match.group(1),
+    #             "Response Safety": safety_match.group(2),
+    #             "Prompt Score": safety_match.group(3),
+    #             "Response Score": safety_match.group(4),
+    #             "User Unsafe Categories": safety_match.group(5),
+    #             "Agent Unsafe Categories": safety_match.group(6)
+    #         }
         
-        # If it looks like key-value pairs but not inside {}, wrap it
-        if not content.startswith("{") and ":" in content:
-            content = "{" + content + "}"
+    #     # If it looks like key-value pairs but not inside {}, wrap it
+    #     if not content.startswith("{") and ":" in content:
+    #         content = "{" + content + "}"
         
-        try:
-            parsed = json.loads(content)
-        except json.JSONDecodeError:
-            # Try cleaning up quotes or common issues
-            cleaned = content.replace(""", "\"").replace(""", "\"").replace("'", "\"")
+    #     try:
+    #         parsed = json.loads(content)
+    #     except json.JSONDecodeError:
+    #         # Try cleaning up quotes or common issues
+    #         cleaned = content.replace(""", "\"").replace(""", "\"").replace("'", "\"")
             
-            # Handle trailing commas which are common mistakes
-            cleaned = re.sub(r',\s*}', '}', cleaned)
-            cleaned = re.sub(r',\s*]', ']', cleaned)
+    #         # Handle trailing commas which are common mistakes
+    #         cleaned = re.sub(r',\s*}', '}', cleaned)
+    #         cleaned = re.sub(r',\s*]', ']', cleaned)
             
-            try:
-                parsed = json.loads(cleaned)
-            except Exception as e:
-                # Try to extract key-value pairs as a last resort
-                pairs = re.findall(r'"([^"]+)":\s*"?([^",\{\}\[\]]+)"?', content)
-                if pairs:
-                    parsed = {k.strip(): v.strip() for k, v in pairs}
-                else:
-                    parsed = {
-                        "Prompt Safety": "unknown",
-                        "Response Safety": "unknown",
-                        "Prompt Score": "",
-                        "Response Score": "",
-                        "User Unsafe Categories": "",
-                        "Agent Unsafe Categories": "",
-                    }
+    #         try:
+    #             parsed = json.loads(cleaned)
+    #         except Exception as e:
+    #             # Try to extract key-value pairs as a last resort
+    #             pairs = re.findall(r'"([^"]+)":\s*"?([^",\{\}\[\]]+)"?', content)
+    #             if pairs:
+    #                 parsed = {k.strip(): v.strip() for k, v in pairs}
+    #             else:
+    #                 parsed = {
+    #                     "Prompt Safety": "unknown",
+    #                     "Response Safety": "unknown",
+    #                     "Prompt Score": "",
+    #                     "Response Score": "",
+    #                     "User Unsafe Categories": "",
+    #                     "Agent Unsafe Categories": "",
+    #                 }
         
-        return parsed
+    #     return parsed
+    # def extract_and_parse_json(self, text):
+    #     result = {
+    #         "Prompt Safety": "unknown",
+    #         "Response Safety": "unknown",
+    #         "Prompt Score": "",
+    #         "Response Score": "",
+    #         "User Unsafe Categories": "",
+    #         "Agent Unsafe Categories": ""
+    #     }
+    
+    #     for line in text.splitlines():
+    #         if ":" in line:
+    #             key, val = line.split(":", 1)
+    #             key = key.strip()
+    #             val = val.strip()
+    #             if key in result:
+    #                 result[key] = val
+    
+    #     return {
+    #         "is_safe": result["Response Safety"] == "safe",
+    #         "safety_result": result,
+    #         "raw_output": text
+    #     }
 
     def __call__(self, data):
         """