update handler.py for text generation

handler.py

@@ -27,21 +27,25 @@ class EndpointHandler:
     def __call__(self, data: Dict):
         """
         Args:
-            data: Dictionary with "inputs" and optional "parameters"
+            data: Dictionary with either string input or structured messages
         Returns:
             Generated text
         """
-        # Extract messages from input
-        messages = data.get("inputs", {}).get("messages", [])
-        if not messages:
-            return {"error": "No messages provided"}
-
-        # Format input text
-        input_text = ""
-        for msg in messages:
-            role = msg.get("role", "")
-            content = msg.get("content", "")
-            input_text += f"{role}: {content}\n"
+        # Handle input
+        if isinstance(data.get("inputs"), str):
+            input_text = data["inputs"]
+        else:
+            # Extract messages from input
+            messages = data.get("inputs", {}).get("messages", [])
+            if not messages:
+                return {"error": "No messages provided"}
+                
+            # Format input text
+            input_text = ""
+            for msg in messages:
+                role = msg.get("role", "")
+                content = msg.get("content", "")
+                input_text += f"{role}: {content}\n"
 
         # Get generation parameters
         params = {**self.default_params}