Update helper.py

helper.py

@@ -11,16 +11,14 @@ def get_fn(model_name: str, **model_kwargs) -> Callable:
     try:
         client = OpenAI(
             base_url="http://192.222.58.60:8000/v1", 
-            api_key="tela",  
+            api_key="tela",  # Replace with your actual API key or use environment variables
         )
     except Exception as e:
         print(f"The API or base URL were not defined: {str(e)}")
-        raise e  # It's better to raise the exception to prevent the app from running without a client
+        raise e  # Prevent the app from running without a client
 
     def predict(
-        message: str,
-        history: list,
-        system_prompt: str,
+        messages: list,  # Preprocessed messages from preprocess function
         temperature: float,
         max_tokens: int,
         top_k: int,
@@ -28,20 +26,6 @@ def get_fn(model_name: str, **model_kwargs) -> Callable:
         top_p: float
     ) -> Generator[str, None, None]:
         try:
-            # Initialize the messages list with the system prompt
-            messages = [
-                {"role": "system", "content": system_prompt}
-            ]
-            
-            # Append the conversation history
-            for user_msg, assistant_msg in history:
-                messages.append({"role": "user", "content": user_msg})
-                if assistant_msg:
-                    messages.append({"role": "assistant", "content": assistant_msg})
-            
-            # Append the latest user message
-            messages.append({"role": "user", "content": message})
-            
             # Call the OpenAI API with the formatted messages
             response = client.chat.completions.create(
                 model=model_name,  
@@ -55,7 +39,7 @@ def get_fn(model_name: str, **model_kwargs) -> Callable:
                 # Ensure response_format is set correctly; typically it's a string like 'text'
                 response_format="text",
             )
-    
+
             response_text = ""
             # Iterate over the streaming response
             for chunk in response:
@@ -68,46 +52,33 @@ def get_fn(model_name: str, **model_kwargs) -> Callable:
             
             if not response_text.strip():
                 yield "I apologize, but I was unable to generate a response. Please try again."
-    
+
         except Exception as e:
             print(f"Error during generation: {str(e)}")
             yield f"An error occurred: {str(e)}"
-    
-    return predict
-
 
+    return predict
 
-def get_image_base64(url: str, ext: str):
+def get_image_base64(url: str, ext: str) -> str:
     with open(url, "rb") as image_file:
         encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
-    return "data:image/" + ext + ";base64," + encoded_string
-
+    return f"data:image/{ext};base64,{encoded_string}"
 
-def handle_user_msg(message: str):
-    if type(message) is str:
+def handle_user_msg(message: str) -> str:
+    if isinstance(message, str):
         return message
-    elif type(message) is dict:
-        if message["files"] is not None and len(message["files"]) > 0:
-            ext = os.path.splitext(message["files"][-1])[1].strip(".")
-            if ext.lower() in ["png", "jpg", "jpeg", "gif", "pdf"]:
+    elif isinstance(message, dict):
+        if message.get("files"):
+            ext = os.path.splitext(message["files"][-1])[1].strip(".").lower()
+            if ext in ["png", "jpg", "jpeg", "gif", "pdf"]:
                 encoded_str = get_image_base64(message["files"][-1], ext)
+                return f"{message.get('text', '')}\n![Image]({encoded_str})"
             else:
-                raise NotImplementedError(f"Not supported file type {ext}")
-            content = [
-                    {"type": "text", "text": message["text"]},
-                    {
-                        "type": "image_url",
-                        "image_url": {
-                            "url": encoded_str,
-                        }
-                    },
-                ]
+                raise NotImplementedError(f"Unsupported file type: {ext}")
         else:
-            content = message["text"]
-        return content
+            return message.get("text", "")
     else:
-        raise NotImplementedError
-
+        raise NotImplementedError("Unsupported message type")
 
 def get_interface_args(pipeline: str):
     if pipeline == "chat":
@@ -138,7 +109,6 @@ def get_interface_args(pipeline: str):
         raise ValueError(f"Unsupported pipeline type: {pipeline}")
     return inputs, outputs, preprocess, postprocess
 
-
 def registry(name: str = None, **kwargs) -> gr.ChatInterface:
     """Create a Gradio Interface with similar styling and parameters."""
     
@@ -187,7 +157,6 @@ def registry(name: str = None, **kwargs) -> gr.ChatInterface:
             gr.Slider(0.0, 2.0, value=1.1, label="Repetition penalty"),
             gr.Slider(0.0, 1.0, value=0.95, label="Top P sampling"),
         ],
-        # Optionally, you can customize other ChatInterface parameters here
     )
     
     return interface