updated the handler.py

handler.py

@@ -3,6 +3,7 @@ import torch
 import logging
 import time
 import traceback
+import json
 from typing import Dict, List, Any, Union, Generator
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 from threading import Thread
@@ -117,13 +118,14 @@ class EndpointHandler:
             Dict[str, Any]: Model response
         """
         start_time = time.time()
-        logger.info(f"Starting request processing: {data}")
+        logger.info(f"Starting request processing")
         
         try:
             # Extract input parameters with defaults
             if "inputs" not in data:
                 logger.warning("No 'inputs' field in request data")
-                return {"generated_text": "Error: Missing 'inputs' field in request"}
+                error_msg = "Missing 'inputs' field in request"
+                return self._format_error_response(error_msg)
             
             # Track user and system messages
             user_message = ""
@@ -158,10 +160,12 @@ class EndpointHandler:
                         user_message = msg.get("content", "")
             else:
                 logger.warning(f"Unsupported input format: {type(data['inputs'])}")
-                return {"generated_text": "Error: Unsupported input format. Expected string or messages object."}
+                error_msg = "Unsupported input format. Expected string or messages object."
+                return self._format_error_response(error_msg)
             
-            logger.info(f"Extracted user message: '{user_message}'")
-            logger.info(f"Extracted system message: '{system_message}'")
+            logger.info(f"Extracted user message length: {len(user_message)} characters")
+            if system_message:
+                logger.info(f"Extracted system message length: {len(system_message)} characters")
             
             # Format the prompt with system and user messages
             prompt = self.format_prompt_with_system(user_message, system_message)
@@ -199,18 +203,25 @@ class EndpointHandler:
                 
                 logger.info(f"Response generation completed, text length: {len(response_text) if isinstance(response_text, str) else 'N/A'}")
                 
-                # Return response - keeping it simple for debug purposes
-                return {"generated_text": response_text}
+                # Format and return response in OpenAI format
+                if isinstance(response_text, str):
+                    return self._format_openai_response(
+                        response_text, 
+                        input_ids.shape[1], 
+                        len(self.tokenizer.encode(response_text)) if response_text else 0
+                    )
+                else:
+                    return self._format_error_response(f"Error during generation: {response_text}")
                 
             except RuntimeError as e:
                 logger.error(f"Runtime Error during generation: {str(e)}")
                 logger.error(traceback.format_exc())
-                return {"generated_text": f"Error during generation: {str(e)}"}
+                return self._format_error_response(f"Error during generation: {str(e)}")
                 
         except Exception as e:
             logger.error(f"Unexpected error during request processing: {str(e)}")
             logger.error(traceback.format_exc())
-            return {"generated_text": f"Unexpected error: {str(e)}"}
+            return self._format_error_response(f"Unexpected error: {str(e)}")
         finally:
             duration = time.time() - start_time
             logger.info(f"Request processing completed in {duration:.2f} seconds")
@@ -298,6 +309,87 @@ class EndpointHandler:
             logger.error(f"Error in _safe_generate: {str(e)}")
             logger.error(traceback.format_exc())
             return f"Generation error: {str(e)}. Please try a simpler input."
+    
+    def _format_openai_response(self, response_text, prompt_tokens, completion_tokens):
+        """Format the response in OpenAI-style format"""
+        try:
+            # Create a response ID
+            response_id = f"phi4-{int(time.time())}"
+            
+            # Build OpenAI-compatible response
+            openai_response = {
+                "id": response_id,
+                "object": "chat.completion",
+                "created": int(time.time()),
+                "model": "phi-4-mini",
+                "choices": [
+                    {
+                        "index": 0,
+                        "message": {
+                            "role": "assistant",
+                            "content": response_text
+                        },
+                        "finish_reason": "stop"
+                    }
+                ],
+                "usage": {
+                    "prompt_tokens": prompt_tokens,
+                    "completion_tokens": completion_tokens,
+                    "total_tokens": prompt_tokens + completion_tokens
+                }
+            }
+            
+            # For compatibility with Hugging Face UI, include the generated_text field
+            openai_response["generated_text"] = response_text
+            
+            logger.info(f"Formatted OpenAI-style response: {len(json.dumps(openai_response))} bytes")
+            return openai_response
+            
+        except Exception as e:
+            logger.error(f"Error formatting OpenAI response: {str(e)}")
+            # Fall back to simple response
+            return {"generated_text": response_text}
+    
+    def _format_error_response(self, error_message):
+        """Format an error response in OpenAI-style format"""
+        try:
+            error_response = {
+                "id": f"phi4-error-{int(time.time())}",
+                "object": "chat.completion",
+                "created": int(time.time()),
+                "model": "phi-4-mini",
+                "choices": [
+                    {
+                        "index": 0,
+                        "message": {
+                            "role": "assistant",
+                            "content": f"Error: {error_message}"
+                        },
+                        "finish_reason": "error"
+                    }
+                ],
+                "usage": {
+                    "prompt_tokens": 0,
+                    "completion_tokens": 0,
+                    "total_tokens": 0
+                },
+                "error": {
+                    "message": error_message,
+                    "type": "invalid_request_error",
+                    "code": "error"
+                }
+            }
+            
+            # For compatibility with Hugging Face UI, include the generated_text field
+            error_response["generated_text"] = f"Error: {error_message}"
+            
+            logger.info(f"Formatted error response: {len(json.dumps(error_response))} bytes")
+            return error_response
+            
+        except Exception as e:
+            logger.error(f"Error formatting error response: {str(e)}")
+            # Fall back to simple error response
+            return {"generated_text": f"Error: {error_message}"}
 
 # For local testing
 if __name__ == "__main__":
@@ -316,4 +408,4 @@ if __name__ == "__main__":
     
     # Run the test
     result = handler(test_with_messages)
-    print(result)
+    print(json.dumps(result, indent=2))