Rename handler.py to handler_3.py

handler.py → handler_3.py

@@ -143,13 +143,13 @@ class EndpointHandler:
         Agent Unsafe Categories:
         
         <End of examples>
-        Provide your safety assessment for the above conversation in the below output JSON format and must provide all the metrics mentioned below:
+        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.
         "Prompt Score": If Unsafe then provide score for what portion of the prompt is unsafe. Score need to be on scale of 1-5.
-        "Prompt 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.
+        "Prompt Unsafe Categories": DON'T mention the whole taxonomy, but specifically mention up to 3 particular categories that violate the provided taxonomy. Omit if all safe.
 
-        "Agent Safety": safe or unsafe rating for the agent response, skip not prvided
+        "Agent Safety": safe or unsafe rating for the agent response, skip not provided
         "Agent Score": If Unsafe, then provide a score for what portion of the prompt is unsafe. Score need to be on scale of 1-5. Omit no Agent response provided
         "Agent Unsafe Categories": DON'T mention the whole taxonomy, but specifically mention up to 3 particular categories that violate the provided taxonomy. Omit if all safe or no Agent response provided. 
 
@@ -208,15 +208,63 @@ class EndpointHandler:
         
         return parsed
     
-    def _format_conversations(self, prompt):
+    # def _format_conversations(self, prompt):
+    #     if self.chat_template == "gemma-3":
+    #         return {
+    #             "conversations": [
+    #                 {"role": "system", "content": [{"type": "text", "text": self.policy_prompt}]},
+    #                 {"role": "user", "content": [{"type": "text", "text": prompt}]},
+    #             ]
+    #         }
+
+
+            
+    #         return { "conversations": [
+    #             {
+    #                 "role": "system",
+    #                 "content": [{"type": "text", "text": self.policy_prompt}]
+    #             },
+    #             {
+    #                 "role": "user",
+    #                 "content": [
+    #                     {"type": "image", "url": user},
+    #                     {"type": "text", "text": "Please assess the content"}
+    #                 ]
+    #             }
+    #         ]
+    #                } 
+
+
+    #     else:  # chatml and others
+    #         return {
+    #             "conversations": [
+    #                 {"role": "system", "content": self.policy_prompt},
+    #                 {"role": "user", "content": prompt},
+    #             ]
+    #         }
+
+    def _format_conversations(self, prompt=None, image_url=None):
         if self.chat_template == "gemma-3":
+            user_content = []
+    
+            if image_url:
+                user_content.append({"type": "image", "url": image_url})
+            if prompt:
+                user_content.append({"type": "text", "text": prompt})
+            elif not user_content:
+                raise ValueError("At least one of `prompt` or `image_url` must be provided.")
+            elif image_url and not prompt:
+                # default text prompt for image-only queries
+                user_content.append({"type": "text", "text": "Please analyze the image."})
+    
             return {
                 "conversations": [
                     {"role": "system", "content": [{"type": "text", "text": self.policy_prompt}]},
-                    {"role": "user", "content": [{"type": "text", "text": prompt}]},
+                    {"role": "user", "content": user_content},
                 ]
             }
-        else:  # chatml and others
+    
+        else:
             return {
                 "conversations": [
                     {"role": "system", "content": self.policy_prompt},
@@ -291,7 +339,7 @@ class EndpointHandler:
         # Decode the output
         decoded_output = self.tokenizer.decode(output[0], skip_special_tokens=True)
         
-        Extract the generated part (after the prompt)
+        ##Extract the generated part (after the prompt)
         response_text = decoded_output[len(prompt):].strip()
         print(response_text)
         
@@ -308,7 +356,7 @@ class EndpointHandler:
             "safety_result": safety_result
         }
         
-        return decoded_output 
+        return safety_result 
 
 # For local testing
 if __name__ == "__main__":