Update app.py

app.py

@@ -53,7 +53,6 @@ if not API_KEY:
 
 # Pydantic models for request/response validation
 class ChatConfig(BaseModel):
-    # Removed api_key field; only model and temperature are received
     model: str = "mistralai/mistral-small-3.1-24b-instruct:free"
     temperature: Optional[float] = Field(default=0.7, ge=0.0, le=1.0)
 
@@ -82,7 +81,6 @@ class InitResponse(BaseModel):
     session_id: str
     status: str
 
-# Simple HTML interface
 @app.get("/", response_class=HTMLResponse)
 async def root():
     """
@@ -404,9 +402,8 @@ async def root():
     </body>
     </html>
     """
-    return html_content
+    return HTMLResponse(content=html_content)
 
-# Health check endpoint
 @app.get("/health")
 async def health_check():
     """Health check endpoint"""
@@ -419,11 +416,9 @@ async def initialize_chat(config: ChatConfig):
         # Generate a session ID
         session_id = f"session_{datetime.now().strftime('%Y%m%d%H%M%S')}_{uuid.uuid4().hex[:8]}"
         
-        # If the environment variable is missing, raise an error
         if not API_KEY:
             raise HTTPException(status_code=400, detail="The OPENROUTE_API environment variable is not set.")
         
-        # Initialize the chat instance
         chat = EnhancedRecursiveThinkingChat(
             api_key=API_KEY, 
             model=config.model,
@@ -434,7 +429,6 @@ async def initialize_chat(config: ChatConfig):
             "created_at": datetime.now().isoformat(),
             "model": config.model
         }
-        
         return {"session_id": session_id, "status": "initialized"}
     except Exception as e:
         logger.error(f"Error initializing chat: {str(e)}")
@@ -453,8 +447,21 @@ async def send_message_original(request: MessageRequest):
         # Make a direct call to the LLM without recursion logic
         messages = [{"role": "user", "content": request.message}]
         response_data = chat._call_api(messages, temperature=chat.temperature, stream=False)
-        # Extract the text from the response
-        original_text = response_data["choices"][0]["message"]["content"]
+
+        # The structure of response_data depends on the underlying LLM. 
+        # We'll try to handle both "message" and "text" keys as possible.
+        if isinstance(response_data, dict) and "choices" in response_data:
+            first_choice = response_data["choices"][0]
+            if "message" in first_choice and "content" in first_choice["message"]:
+                original_text = first_choice["message"]["content"]
+            elif "text" in first_choice:
+                original_text = first_choice["text"]
+            else:
+                original_text = str(first_choice)
+        else:
+            # If for some reason the response is not in the expected format, just convert to string
+            original_text = str(response_data)
+
         return {"response": original_text.strip()}
     except Exception as e:
         logger.error(f"Error getting original response: {str(e)}")
@@ -462,20 +469,21 @@ async def send_message_original(request: MessageRequest):
 
 @app.post("/api/send_message")
 async def send_message(request: MessageRequest):
-    """Send a message and get a response with the chain-of-thought process (HTTP-based, not streaming)."""
+    """
+    Send a message and get a response with the chain-of-thought process (HTTP-based, not streaming).
+    Primarily left here for completeness, but the user-facing code calls the WebSocket for streaming.
+    """
     try:
         if request.session_id not in chat_instances:
             raise HTTPException(status_code=404, detail="Session not found")
         
         chat = chat_instances[request.session_id]["chat"]
         
-        # Override class parameters if provided
         original_thinking_fn = chat._determine_thinking_rounds
         original_alternatives_fn = chat._generate_alternatives
         original_temperature = getattr(chat, "temperature", 0.7)
         
         if request.thinking_rounds is not None:
-            # Override the thinking rounds determination
             chat._determine_thinking_rounds = lambda _: request.thinking_rounds
         
         if request.alternatives_per_round is not None:
@@ -483,18 +491,16 @@ async def send_message(request: MessageRequest):
                 return original_alternatives_fn(base_response, prompt, request.alternatives_per_round)
             chat._generate_alternatives = modified_generate_alternatives
         
-        # Override temperature if provided
         if request.temperature is not None:
             setattr(chat, "temperature", request.temperature)
         
-        # Process the message
         logger.info(f"Processing message for session {request.session_id}")
         start_time = datetime.now()
         result = chat.think_and_respond(request.message, verbose=True)
         processing_time = (datetime.now() - start_time).total_seconds()
         logger.info(f"Message processed in {processing_time:.2f} seconds")
         
-        # Restore original functions and parameters
+        # Restore original
         chat._determine_thinking_rounds = original_thinking_fn
         chat._generate_alternatives = original_alternatives_fn
         if request.temperature is not None:
@@ -513,21 +519,19 @@ async def send_message(request: MessageRequest):
 
 @app.post("/api/save")
 async def save_conversation(request: SaveRequest):
-    """Save the conversation or the full thinking log"""
+    """Save the conversation or the full thinking log."""
     try:
         if request.session_id not in chat_instances:
             raise HTTPException(status_code=404, detail="Session not found")
         
         chat = chat_instances[request.session_id]["chat"]
         
-        # Generate default filename if not provided
         filename = request.filename
         if filename is None:
             timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
             log_type = "full_log" if request.full_log else "conversation"
             filename = f"recthink_{log_type}_{timestamp}.json"
         
-        # Make sure the output directory exists
         os.makedirs("logs", exist_ok=True)
         file_path = os.path.join("logs", filename)
         
@@ -543,31 +547,28 @@ async def save_conversation(request: SaveRequest):
 
 @app.get("/api/sessions", response_model=SessionResponse)
 async def list_sessions():
-    """List all active chat sessions"""
+    """List all active chat sessions."""
     sessions = []
     for session_id, session_data in chat_instances.items():
         chat = session_data["chat"]
-        message_count = len(chat.conversation_history) // 2  # Each message-response pair counts as 2
-        
+        message_count = len(chat.conversation_history) // 2
         sessions.append(SessionInfo(
             session_id=session_id,
             message_count=message_count,
             created_at=session_data["created_at"],
             model=session_data["model"]
         ))
-    
     return {"sessions": sessions}
 
 @app.get("/api/sessions/{session_id}")
 async def get_session(session_id: str):
-    """Get details for a specific chat session"""
+    """Get details for a specific chat session."""
     if session_id not in chat_instances:
         raise HTTPException(status_code=404, detail="Session not found")
     
     session_data = chat_instances[session_id]
     chat = session_data["chat"]
     
-    # Extract conversation history
     conversation = []
     for i in range(0, len(chat.conversation_history), 2):
         if i+1 < len(chat.conversation_history):
@@ -586,14 +587,12 @@ async def get_session(session_id: str):
 
 @app.delete("/api/sessions/{session_id}")
 async def delete_session(session_id: str):
-    """Delete a chat session"""
+    """Delete a chat session."""
     if session_id not in chat_instances:
         raise HTTPException(status_code=404, detail="Session not found")
-    
     del chat_instances[session_id]
     return {"status": "deleted", "session_id": session_id}
 
-# WebSocket connection manager
 class ConnectionManager:
     def __init__(self):
         self.active_connections: Dict[str, WebSocket] = {}
@@ -612,7 +611,6 @@ class ConnectionManager:
 
 manager = ConnectionManager()
 
-# WebSocket for streaming the thinking process
 @app.websocket("/ws/{session_id}")
 async def websocket_endpoint(websocket: WebSocket, session_id: str):
     try:
@@ -624,40 +622,31 @@ async def websocket_endpoint(websocket: WebSocket, session_id: str):
             return
         
         chat = chat_instances[session_id]["chat"]
-        
-        # Set up a custom callback to stream the thinking process
         original_call_api = chat._call_api
         
         async def stream_callback(chunk):
             await manager.send_json(session_id, {"type": "chunk", "content": chunk})
         
-        # Override the _call_api method to also send updates via WebSocket
         def ws_call_api(messages, temperature=0.7, stream=True):
             result = original_call_api(messages, temperature, stream)
-            # Send the chunk via WebSocket if we're streaming
             if stream:
                 asyncio.create_task(stream_callback(result))
             return result
         
-        # Replace the method temporarily
         chat._call_api = ws_call_api
         
-        # Wait for messages from the client
         while True:
             data = await websocket.receive_text()
             message_data = json.loads(data)
             
             if message_data["type"] == "message":
-                # Process the message
                 start_time = datetime.now()
                 
                 try:
-                    # Get parameters if they exist
                     thinking_rounds = message_data.get("thinking_rounds", None)
                     alternatives_per_round = message_data.get("alternatives_per_round", None)
                     temperature = message_data.get("temperature", None)
                     
-                    # Override if needed
                     original_thinking_fn = chat._determine_thinking_rounds
                     original_alternatives_fn = chat._generate_alternatives
                     original_temperature = getattr(chat, "temperature", 0.7)
@@ -673,24 +662,20 @@ async def websocket_endpoint(websocket: WebSocket, session_id: str):
                     if temperature is not None:
                         setattr(chat, "temperature", temperature)
                     
-                    # Send a status message that we've started processing
                     await manager.send_json(session_id, {
                         "type": "status",
                         "status": "processing",
                         "message": "Starting recursive thinking process..."
                     })
                     
-                    # Process the message with chain-of-thought
                     result = chat.think_and_respond(message_data["content"], verbose=True)
                     processing_time = (datetime.now() - start_time).total_seconds()
                     
-                    # Restore original functions
                     chat._determine_thinking_rounds = original_thinking_fn
                     chat._generate_alternatives = original_alternatives_fn
                     if temperature is not None:
                         setattr(chat, "temperature", original_temperature)
                     
-                    # Send the final result
                     await manager.send_json(session_id, {
                         "type": "final",
                         "response": result["response"],
@@ -706,7 +691,6 @@ async def websocket_endpoint(websocket: WebSocket, session_id: str):
                         "type": "error",
                         "error": error_msg
                     })
-    
     except WebSocketDisconnect:
         logger.info(f"WebSocket disconnected: {session_id}")
         manager.disconnect(session_id)
@@ -718,14 +702,11 @@ async def websocket_endpoint(websocket: WebSocket, session_id: str):
         except:
             pass
     finally:
-        # Restore the original method if needed
         if 'chat' in locals() and 'original_call_api' in locals():
             chat._call_api = original_call_api
         
-        # Make sure to disconnect
         manager.disconnect(session_id)
 
-# Use port 7860 for Hugging Face Spaces
 if __name__ == "__main__":
     port = 7860
     print(f"Starting server on port {port}")