Update app.py

app.py

@@ -71,49 +71,54 @@ def query_model(model_name: str, messages: List[Dict[str, str]]) -> str:
     except Exception as e:
         return f"{model_name} error: {str(e)}"
 
+class ConversationState:
+    def __init__(self):
+        self.messages = []
+        
+    def add_user_message(self, message: str):
+        self.messages.append({"role": "user", "content": message})
+        
+    def add_assistant_message(self, model_name: str, message: str):
+        self.messages.append({
+            "role": "assistant",
+            "model": model_name,
+            "content": message
+        })
+        
+    def get_context(self) -> List[Dict[str, str]]:
+        return [
+            {
+                "role": msg["role"],
+                "content": f"{msg.get('model', '')}: {msg['content']}" if msg["role"] == "assistant" else msg["content"]
+            }
+            for msg in self.messages
+        ]
+
+conversation_state = ConversationState()
+
 def respond(message: str, history: List[List[str]]) -> str:
-    """Handle sequential model responses with continuous context"""
-    # Build full conversation history
-    messages = []
-    for user_msg, assistant_msg in history:
-        messages.append({"role": "user", "content": user_msg})
-        if assistant_msg:
-            # Split the assistant message to get individual model responses
-            model_responses = assistant_msg.split("\n\n")
-            for response in model_responses:
-                if "**" in response:  # Only add valid model responses
-                    messages.append({"role": "assistant", "content": response})
+    """Handle sequential model responses with continuous conversation context"""
     
-    # Add current message
-    messages.append({"role": "user", "content": message})
-    current_output = ""
+    # Add current message to conversation state
+    conversation_state.add_user_message(message)
+    current_context = conversation_state.get_context()
     
     # Get first model's response
-    response1 = query_model("Qwen2.5-Coder-32B-Instruct", messages)
-    current_output = f"**Qwen2.5-Coder-32B-Instruct**:\n{response1}"
-    yield current_output
-    
-    # Add first response to context
-    messages.append({
-        "role": "assistant",
-        "content": f"Previous response: {response1}"
-    })
-    
-    # Get second model's response
-    response2 = query_model("Qwen2.5-72B-Instruct", messages)
-    current_output = f"{current_output}\n\n**Qwen2.5-72B-Instruct**:\n{response2}"
-    yield current_output
+    response1 = query_model("Qwen2.5-Coder-32B-Instruct", current_context)
+    conversation_state.add_assistant_message("Qwen2.5-Coder-32B-Instruct", response1)
+    yield f"**Qwen2.5-Coder-32B-Instruct**:\n{response1}"
     
-    # Add second response to context
-    messages.append({
-        "role": "assistant",
-        "content": f"Previous responses: {response1}\n{response2}"
-    })
+    # Get second model's response with updated context
+    current_context = conversation_state.get_context()
+    response2 = query_model("Qwen2.5-72B-Instruct", current_context)
+    conversation_state.add_assistant_message("Qwen2.5-72B-Instruct", response2)
+    yield f"**Qwen2.5-72B-Instruct**:\n{response2}"
     
-    # Get final model's response
-    response3 = query_model("Llama3.3-70B-Instruct", messages)
-    current_output = f"{current_output}\n\n**Llama3.3-70B-Instruct**:\n{response3}"
-    yield current_output
+    # Get final model's response with complete context
+    current_context = conversation_state.get_context()
+    response3 = query_model("Llama3.3-70B-Instruct", current_context)
+    conversation_state.add_assistant_message("Llama3.3-70B-Instruct", response3)
+    yield f"**Llama3.3-70B-Instruct**:\n{response3}"
 
 # Create the Gradio interface
 chat_interface = gr.ChatInterface(