fix async error

app.py

@@ -1,10 +1,12 @@
 import asyncio
 import os
 import json
-from typing import List, Dict, Any, Union
+from typing import List, Dict, Any, Union, Optional
 from contextlib import AsyncExitStack
 import mimetypes
 import tempfile
+import threading
+from concurrent.futures import ThreadPoolExecutor
 
 import gradio as gr
 from gradio.components.chatbot import ChatMessage
@@ -17,15 +19,13 @@ from dotenv import load_dotenv
 
 load_dotenv()
 
-loop = asyncio.new_event_loop()
-asyncio.set_event_loop(loop)
-
 class MCPClientWrapper:
     def __init__(self):
-        self.session = None
-        self.exit_stack = None
+        self.session: Optional[ClientSession] = None
+        self.exit_stack: Optional[AsyncExitStack] = None
         self.tools = []
         self.connected = False
+        self._connection_lock = threading.Lock()
         
         # Initialize all LLM clients
         self.anthropic_client = None
@@ -61,10 +61,9 @@ class MCPClientWrapper:
         
         try:
             if os.getenv("LLAMAINDEX_API_KEY"):
-                # Using OpenAI-compatible endpoint for Llama
                 self.llama_client = OpenAI(
                     api_key=os.getenv("LLAMAINDEX_API_KEY"),
-                    base_url="https://api.llamaindex.ai/v1"  # Adjust based on your provider
+                    base_url="https://api.llamaindex.ai/v1"
                 )
         except Exception as e:
             print(f"⚠️ Failed to initialize Llama client: {e}")
@@ -118,28 +117,43 @@ class MCPClientWrapper:
         self.current_model = model
         return f"✅ Switched to {provider}: {model}"
     
-    def connect(self) -> str:
-        return loop.run_until_complete(self._connect())
-    
-    async def _connect(self) -> str:
+    async def _cleanup_connection(self):
+        """Safely cleanup existing connection."""
         if self.exit_stack:
-            await self.exit_stack.aclose()
-        
-        self.exit_stack = AsyncExitStack()
-        
-        server_path = "gradio_mcp_server.py"
-        
-        server_params = StdioServerParameters(
-            command="python",
-            args=[server_path],
-            env={"PYTHONIOENCODING": "utf-8", "PYTHONUNBUFFERED": "1"}
-        )
-        
+            try:
+                await self.exit_stack.aclose()
+            except Exception as e:
+                print(f"Warning: Error during cleanup: {e}")
+            finally:
+                self.exit_stack = None
+                self.session = None
+                self.connected = False
+    
+    async def _establish_connection(self) -> str:
+        """Establish MCP connection in proper async context."""
         try:
-            stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params))
-            self.stdio, self.write = stdio_transport
+            # Clean up any existing connection
+            await self._cleanup_connection()
+            
+            self.exit_stack = AsyncExitStack()
+            
+            server_path = "gradio_mcp_server.py"
+            server_params = StdioServerParameters(
+                command="python",
+                args=[server_path],
+                env={"PYTHONIOENCODING": "utf-8", "PYTHONUNBUFFERED": "1"}
+            )
+            
+            # Enter the async context managers
+            stdio_transport = await self.exit_stack.enter_async_context(
+                stdio_client(server_params)
+            )
+            stdio, write = stdio_transport
+            
+            self.session = await self.exit_stack.enter_async_context(
+                ClientSession(stdio, write)
+            )
             
-            self.session = await self.exit_stack.enter_async_context(ClientSession(self.stdio, self.write))
             await self.session.initialize()
             
             response = await self.session.list_tools()
@@ -152,26 +166,56 @@ class MCPClientWrapper:
             self.connected = True
             tool_names = [tool["name"] for tool in self.tools]
             return f"✅ Connected to MCP Weather Server. Available tools: {', '.join(tool_names)}"
+            
         except Exception as e:
             self.connected = False
+            await self._cleanup_connection()
             return f"❌ Failed to connect to MCP server: {str(e)}"
     
+    def connect(self) -> str:
+        """Thread-safe connection method for Gradio."""
+        with self._connection_lock:
+            try:
+                # Create new event loop for this operation
+                try:
+                    loop = asyncio.get_event_loop()
+                except RuntimeError:
+                    loop = asyncio.new_event_loop()
+                    asyncio.set_event_loop(loop)
+                
+                if loop.is_running():
+                    # If loop is already running, we need to run in a thread
+                    import concurrent.futures
+                    with concurrent.futures.ThreadPoolExecutor() as executor:
+                        future = executor.submit(self._run_connection_in_new_loop)
+                        return future.result()
+                else:
+                    return loop.run_until_complete(self._establish_connection())
+            except Exception as e:
+                return f"❌ Connection error: {str(e)}"
+    
+    def _run_connection_in_new_loop(self) -> str:
+        """Run connection in a new event loop (for thread safety)."""
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        try:
+            return loop.run_until_complete(self._establish_connection())
+        finally:
+            loop.close()
+    
     def read_uploaded_file(self, file_path: str) -> str:
         """Read and process uploaded file content."""
         if not file_path or not os.path.exists(file_path):
             return ""
         
         try:
-            # Get file info
             file_size = os.path.getsize(file_path)
             file_name = os.path.basename(file_path)
             mime_type, _ = mimetypes.guess_type(file_path)
             
-            # Check file size (limit to 10MB)
             if file_size > 10 * 1024 * 1024:
-                return f"\n\n🔍 **File Upload Error**: {file_name} is too large (>10MB). Please upload a smaller file."
+                return f"\n\n📄 **File Upload Error**: {file_name} is too large (>10MB). Please upload a smaller file."
             
-            # Try to read as text
             encodings_to_try = ['utf-8', 'utf-16', 'latin-1', 'cp1252']
             
             for encoding in encodings_to_try:
@@ -179,12 +223,11 @@ class MCPClientWrapper:
                     with open(file_path, 'r', encoding=encoding) as f:
                         content = f.read()
                     
-                    # If content is too long, truncate it
-                    max_chars = 50000  # Roughly 50k characters
+                    max_chars = 50000
                     if len(content) > max_chars:
                         content = content[:max_chars] + f"\n\n[Content truncated - showing first {max_chars} characters of {len(content)} total]"
                     
-                    file_info = f"\n\n🔍 **Uploaded File**: {file_name}"
+                    file_info = f"\n\n📄 **Uploaded File**: {file_name}"
                     if mime_type:
                         file_info += f" ({mime_type})"
                     file_info += f" - {file_size:,} bytes\n\n```\n{content}\n```"
@@ -194,18 +237,16 @@ class MCPClientWrapper:
                 except UnicodeDecodeError:
                     continue
             
-            # If all text encodings fail, it's likely a binary file
-            return f"\n\n🔍 **File Upload**: {file_name} appears to be a binary file and cannot be displayed as text."
+            return f"\n\n📄 **File Upload**: {file_name} appears to be a binary file and cannot be displayed as text."
             
         except Exception as e:
-            return f"\n\n🔍 **File Upload Error**: Could not read {file_name}: {str(e)}"
+            return f"\n\n📄 **File Upload Error**: Could not read {file_name}: {str(e)}"
     
     def _convert_tools_for_provider(self, provider: str):
         """Convert MCP tools format to provider-specific format."""
         if provider == "claude":
             return self.tools
         elif provider in ["openai", "llama"]:
-            # Convert to OpenAI tools format
             openai_tools = []
             for tool in self.tools:
                 openai_tools.append({
@@ -218,7 +259,6 @@ class MCPClientWrapper:
                 })
             return openai_tools
         elif provider == "mistral":
-            # Convert to Mistral tools format
             mistral_tools = []
             for tool in self.tools:
                 mistral_tools.append({
@@ -270,6 +310,7 @@ class MCPClientWrapper:
             raise Exception(f"Error calling {provider}: {str(e)}")
     
     def process_message(self, message: str, history: List[Union[Dict[str, Any], ChatMessage]], uploaded_file) -> tuple:
+        """Process message in thread-safe manner."""
         if not self.session or not self.connected:
             return history + [
                 {"role": "user", "content": message}, 
@@ -284,10 +325,44 @@ class MCPClientWrapper:
         # Combine message with file content
         full_message = message + file_content
         
-        new_messages = loop.run_until_complete(self._process_query(full_message, history))
-        return history + [{"role": "user", "content": full_message}] + new_messages, gr.Textbox(value=""), gr.File(value=None)
+        try:
+            # Run async processing in new event loop
+            new_messages = self._run_async_processing(full_message, history)
+            return history + [{"role": "user", "content": full_message}] + new_messages, gr.Textbox(value=""), gr.File(value=None)
+        except Exception as e:
+            return history + [
+                {"role": "user", "content": full_message},
+                {"role": "assistant", "content": f"❌ Error processing message: {str(e)}"}
+            ], gr.Textbox(value=""), gr.File(value=None)
+    
+    def _run_async_processing(self, message: str, history: List[Union[Dict[str, Any], ChatMessage]]):
+        """Run async message processing in new event loop."""
+        try:
+            loop = asyncio.get_event_loop()
+        except RuntimeError:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+        
+        if loop.is_running():
+            # Run in thread if event loop is already running
+            import concurrent.futures
+            with concurrent.futures.ThreadPoolExecutor() as executor:
+                future = executor.submit(self._process_in_new_loop, message, history)
+                return future.result()
+        else:
+            return loop.run_until_complete(self._process_query(message, history))
+    
+    def _process_in_new_loop(self, message: str, history: List[Union[Dict[str, Any], ChatMessage]]):
+        """Process query in a completely new event loop."""
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        try:
+            return loop.run_until_complete(self._process_query(message, history))
+        finally:
+            loop.close()
     
     async def _process_query(self, message: str, history: List[Union[Dict[str, Any], ChatMessage]]):
+        """Process the actual query with LLM and tools."""
         claude_messages = []
         for msg in history:
             if isinstance(msg, ChatMessage):
@@ -305,8 +380,6 @@ class MCPClientWrapper:
         except Exception as e:
             return [{"role": "assistant", "content": f"❌ Error with {self.current_provider}: {str(e)}"}]
 
-        result_messages = []
-        
         # Handle different response formats based on provider
         if self.current_provider == "claude":
             return await self._process_claude_response(response, claude_messages)
@@ -315,7 +388,7 @@ class MCPClientWrapper:
         elif self.current_provider == "mistral":
             return await self._process_mistral_response(response, claude_messages)
         
-        return result_messages
+        return []
     
     async def _process_claude_response(self, response, claude_messages):
         """Process Claude API response."""
@@ -348,11 +421,9 @@ class MCPClientWrapper:
                 if isinstance(result_content, list):
                     result_content = "\n".join(str(item) for item in result_content)
                 
-                # Format the response
                 formatted_response = self._format_weather_response(result_content, tool_name)
                 result_messages.append(formatted_response)
                 
-                # Let the LLM analyze and respond
                 claude_messages.append({"role": "user", "content": f"Tool result for {tool_name}: {result_content}"})
                 next_response = await self._call_llm(claude_messages, self.current_provider, self.current_model)
                 
@@ -500,6 +571,7 @@ class MCPClientWrapper:
             }
         }
 
+# Initialize client
 client = MCPClientWrapper()
 
 def gradio_interface():
@@ -547,7 +619,7 @@ def gradio_interface():
         status = gr.Textbox(
             label="🔌 Connection Status", 
             interactive=False, 
-            value="🔄 Connecting to weather server..."
+            value="🔄 Ready to connect..."
         )
         
         # Main chat interface
@@ -556,11 +628,10 @@ def gradio_interface():
             height=600,
             type="messages",
             show_copy_button=True,
-            avatar_images=("👤", "🤖"),
-            bubble_full_width=False
+            avatar_images=("👤", "🤖")
         )
         
-        # File upload component (already exists in your code!)
+        # File upload component
         file_upload = gr.File(
             label="📎 Upload File (optional)",
             file_count="single",
@@ -611,7 +682,7 @@ def gradio_interface():
                 return f"{provider}: {model}", status_msg
             return current_model_display.value, "❌ Please select both provider and model"
         
-        # Auto-connect when the interface loads
+        # Auto-connect function
         def auto_connect():
             return client.connect()