Refactor add_text and get_history functions to improve message handling and formatting

app.py

@@ -141,7 +141,6 @@ def clear_history(request: gr.Request):
 
 
 def add_text(state, message, system_prompt, request: gr.Request):
-    print(f"state: {state}")
     if not state:
         state = init_state()
     images = message.get("files", [])
@@ -162,11 +161,22 @@ def add_text(state, message, system_prompt, request: gr.Request):
             return (state, state.to_gradio_chatbot(), textbox) + (no_change_btn,) * 5
     images = [Image.open(path).convert("RGB") for path in images]
 
+    # Init again if send the second image
     if len(images) > 0 and len(state.get_images(source=state.USER)) > 0:
         state = init_state(state)
-    
+
+    # Upload the first image
     if len(images) > 0 and len(state.get_images(source=state.USER)) == 0:
-        text = DEFAULT_IMAGE_TOKEN + "\n" + text
+        if len(state.messages) == 0: ## In case the first message is an image
+            text = DEFAULT_IMAGE_TOKEN + "\n" + system_prompt + "\n" + text
+        else: ## In case the image is uploaded after some text messages
+           first_user_message = state.messages[0]['content']
+           state.update_message(Conversation.USER, DEFAULT_IMAGE_TOKEN + "\n" + first_user_message, None, 0)
+
+    # If the first message is text
+    if len(images) == 0 and len(state.get_images(source=state.USER)) == 0 and len(state.messages) == 0:
+        text = system_prompt + "\n" + text
+
 
     state.set_system_message(system_prompt)
     state.append_message(Conversation.USER, text, images)
@@ -215,6 +225,10 @@ def predict(state,
                     state.update_message(state.USER, DEFAULT_IMAGE_TOKEN + "\n" + first_user_message, None, index)
         
 
+        idx_last_user_mess, message = state.get_user_message(source=state.USER, position='last')
+        if idx_last_user_mess == 0 and DEFAULT_IMAGE_TOKEN not in first_user_message:
+            message = state.system_message + "\n" + message
+
         history = state.get_history()
         logger.info(f"====  History ====\n{history}")
         message = history[-1][0] if len(history) > 0 else ""

conversation.py

@@ -137,35 +137,28 @@ class Conversation:
         return send_messages
     
     def get_history(self):
-        results = []
+        history = []
         system_message = self.system_message
-        messages = self.messages[:-2]
-
+        messages = self.messages[:-1]
         logger.info(f"=== Raw messages ===\n{messages}")
 
+        if len(messages) < 2:
+            return None
+
         for i in range(len(messages)):
             if messages[i]['role'] == 'user':
                 # Create the question by combining system and user messages
                 user_content = messages[i]['content']
 
-                # Check if it's the first user message and contains <image>
-                if i == 0 and '<image>' in user_content:
-                    user_content = user_content.replace('<image>', '')
-                    question = f"<image>\n{system_message}\n{user_content}"
-                elif i == 0 and '<image>' not in user_content:
-                    question = f"{system_message}\n{user_content}"
-                else:
-                    question = f"{user_content}"
-
                 # Check for the corresponding assistant response
                 answer = ""
                 if i + 1 < len(messages) and messages[i + 1]['role'] == 'assistant':
                     answer = messages[i + 1]['content']
 
                 # Add the question-answer pair to results
-                results.append((question, answer))
+                history.append((user_content, answer))
 
-        return results
+        return history
         
 
     def append_message(