Update app.py

app.py

@@ -74,19 +74,28 @@ def query_model(model_name: str, messages: List[Dict[str, str]]) -> str:
     except Exception as e:
         return f"{model_name} error: {str(e)}"
 
-def respond(message: str, history: List[List[str]], session_id: str) -> Generator[str, None, None]:
-    """Handle sequential model responses with session tracking"""
+def respond(message: str, history: List[List[str]], session_id: str) -> List[List[str]]:
+    """Handle sequential model responses with context preservation"""
     # Load session history
     session = session_manager.load_session(session_id)
-    messages = [{"role": "user", "content": message}]
     
-    # Store user message and update session
+    # Build context from session history
+    messages = []
+    for entry in session["history"]:
+        if entry["type"] == "user":
+            messages.append({"role": "user", "content": entry["content"]})
+        else:
+            messages.append({"role": "assistant", "content": f"{entry['model']}: {entry['content']}"})
+    
+    # Add current message
+    messages.append({"role": "user", "content": message})
     session["history"].append({
         "timestamp": datetime.now().isoformat(),
         "type": "user",
         "content": message
     })
-    session_manager.save_session(session_id, session)
+    
+    responses = []
     
     # Get first model's response
     response1 = query_model("Qwen2.5-Coder-32B-Instruct", messages)
@@ -96,14 +105,8 @@ def respond(message: str, history: List[List[str]], session_id: str) -> Generato
         "model": "Qwen2.5-Coder-32B-Instruct",
         "content": response1
     })
-    session_manager.save_session(session_id, session)
-    yield f"**Qwen2.5-Coder-32B-Instruct**:\n{response1}"
-    
-    # Add first response to context
-    messages.append({
-        "role": "assistant",
-        "content": f"Previous response: {response1}"
-    })
+    messages.append({"role": "assistant", "content": f"Qwen2.5-Coder-32B-Instruct: {response1}"})
+    responses.append(f"**Qwen2.5-Coder-32B-Instruct**:\n{response1}")
     
     # Get second model's response
     response2 = query_model("Qwen2.5-72B-Instruct", messages)
@@ -113,14 +116,8 @@ def respond(message: str, history: List[List[str]], session_id: str) -> Generato
         "model": "Qwen2.5-72B-Instruct",
         "content": response2
     })
-    session_manager.save_session(session_id, session)
-    yield f"**Qwen2.5-72B-Instruct**:\n{response2}"
-    
-    # Add second response to context
-    messages.append({
-        "role": "assistant",
-        "content": f"Previous responses: {response1}\n{response2}"
-    })
+    messages.append({"role": "assistant", "content": f"Qwen2.5-72B-Instruct: {response2}"})
+    responses.append(f"**Qwen2.5-72B-Instruct**:\n{response2}")
     
     # Get final model's response
     response3 = query_model("Llama3.3-70B-Instruct", messages)
@@ -130,8 +127,14 @@ def respond(message: str, history: List[List[str]], session_id: str) -> Generato
         "model": "Llama3.3-70B-Instruct",
         "content": response3
     })
+    messages.append({"role": "assistant", "content": f"Llama3.3-70B-Instruct: {response3}"})
+    responses.append(f"**Llama3.3-70B-Instruct**:\n{response3}")
+    
+    # Save final session state
     session_manager.save_session(session_id, session)
-    yield f"**Llama3.3-70B-Instruct**:\n{response3}"
+    
+    # Return responses in Gradio chat format
+    return [[message, response] for response in responses]
 
 # Create the Gradio interface with session management
 with gr.Blocks(title="Multi-LLM Collaboration Chat") as demo: