Update app.py

app.py

@@ -3,9 +3,11 @@ import os
 import requests
 import threading
 from datetime import datetime
-from typing import List, Dict, Any
+from typing import List, Dict, Any, Generator
+from session_manager import SessionManager
 
-# Get the Hugging Face API key from Spaces secrets
+# Initialize session manager and get HF API key
+session_manager = SessionManager()
 HF_API_KEY = os.getenv("HF_API_KEY")
 
 # Model endpoints configuration
@@ -72,21 +74,29 @@ 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) -> str:
+def respond(message: str, history: List[List[str]], session_id: str) -> Generator[str, None, None]:
     """Handle sequential model responses with session tracking"""
     # Load session history
     session = session_manager.load_session(session_id)
     messages = [{"role": "user", "content": message}]
     
-    # Store user message in session
+    # Store user message and update session
     session["history"].append({
         "timestamp": datetime.now().isoformat(),
         "type": "user",
         "content": message
     })
+    session_manager.save_session(session_id, session)
     
     # Get first model's response
     response1 = query_model("Qwen2.5-Coder-32B-Instruct", messages)
+    session["history"].append({
+        "timestamp": datetime.now().isoformat(),
+        "type": "assistant",
+        "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
@@ -97,6 +107,13 @@ def respond(message: str, history: List[List[str]], session_id: str) -> str:
     
     # Get second model's response
     response2 = query_model("Qwen2.5-72B-Instruct", messages)
+    session["history"].append({
+        "timestamp": datetime.now().isoformat(),
+        "type": "assistant",
+        "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
@@ -107,6 +124,13 @@ def respond(message: str, history: List[List[str]], session_id: str) -> str:
     
     # Get final model's response
     response3 = query_model("Llama3.3-70B-Instruct", messages)
+    session["history"].append({
+        "timestamp": datetime.now().isoformat(),
+        "type": "assistant",
+        "model": "Llama3.3-70B-Instruct",
+        "content": response3
+    })
+    session_manager.save_session(session_id, session)
     yield f"**Llama3.3-70B-Instruct**:\n{response3}"
 
 # Create the Gradio interface with session management
@@ -137,4 +161,4 @@ with gr.Blocks(title="Multi-LLM Collaboration Chat") as demo:
     )
 
 if __name__ == "__main__":
-    chat_interface.launch(share=True)
+    demo.launch(share=True)