Update app.py

app.py

@@ -9,9 +9,6 @@ from gradio.components.chatbot import ChatMessage
 from mcp import ClientSession, StdioServerParameters
 from mcp.client.stdio import stdio_client
 from anthropic import Anthropic
-from dotenv import load_dotenv
-
-load_dotenv()
 
 loop = asyncio.new_event_loop()
 asyncio.set_event_loop(loop)
@@ -20,10 +17,26 @@ class MCPClientWrapper:
     def __init__(self):
         self.session = None
         self.exit_stack = None
-        self.anthropic = Anthropic()
+        self.anthropic = None
         self.tools = []
+        self.api_key = None
+    
+    def set_api_key(self, api_key: str) -> str:
+        """Set the Anthropic API key and initialize the client"""
+        if not api_key or not api_key.strip():
+            return "❌ Please provide a valid Anthropic API key"
+        
+        try:
+            self.api_key = api_key.strip()
+            self.anthropic = Anthropic(api_key=self.api_key)
+            return "✅ Anthropic API key set successfully"
+        except Exception as e:
+            return f"❌ Error setting API key: {str(e)}"
     
     def connect(self, server_path: str) -> str:
+        if not self.anthropic:
+            return "❌ Please set your Anthropic API key first"
+        
         return loop.run_until_complete(self._connect(server_path))
     
     async def _connect(self, server_path: str) -> str:
@@ -41,31 +54,46 @@ class MCPClientWrapper:
             env={"PYTHONIOENCODING": "utf-8", "PYTHONUNBUFFERED": "1"}
         )
         
-        stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params))
-        self.stdio, self.write = stdio_transport
-        
-        self.session = await self.exit_stack.enter_async_context(ClientSession(self.stdio, self.write))
-        await self.session.initialize()
-        
-        response = await self.session.list_tools()
-        self.tools = [{ 
-            "name": tool.name,
-            "description": tool.description,
-            "input_schema": tool.inputSchema
-        } for tool in response.tools]
-        
-        tool_names = [tool["name"] for tool in self.tools]
-        return f"Connected to MCP server. Available tools: {', '.join(tool_names)}"
+        try:
+            stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params))
+            self.stdio, self.write = stdio_transport
+            
+            self.session = await self.exit_stack.enter_async_context(ClientSession(self.stdio, self.write))
+            await self.session.initialize()
+            
+            response = await self.session.list_tools()
+            self.tools = [{ 
+                "name": tool.name,
+                "description": tool.description,
+                "input_schema": tool.inputSchema
+            } for tool in response.tools]
+            
+            tool_names = [tool["name"] for tool in self.tools]
+            return f"✅ Connected to MCP server. Available tools: {', '.join(tool_names)}"
+        except Exception as e:
+            return f"❌ Failed to connect to MCP server: {str(e)}"
     
     def process_message(self, message: str, history: List[Union[Dict[str, Any], ChatMessage]]) -> tuple:
+        if not self.anthropic:
+            return history + [
+                {"role": "user", "content": message}, 
+                {"role": "assistant", "content": "❌ Please set your Anthropic API key first."}
+            ], gr.Textbox(value="")
+        
         if not self.session:
             return history + [
                 {"role": "user", "content": message}, 
-                {"role": "assistant", "content": "Please connect to an MCP server first."}
+                {"role": "assistant", "content": "❌ Please connect to an MCP server first."}
             ], gr.Textbox(value="")
         
-        new_messages = loop.run_until_complete(self._process_query(message, history))
-        return history + [{"role": "user", "content": message}] + new_messages, gr.Textbox(value="")
+        try:
+            new_messages = loop.run_until_complete(self._process_query(message, history))
+            return history + [{"role": "user", "content": message}] + new_messages, gr.Textbox(value="")
+        except Exception as e:
+            return history + [
+                {"role": "user", "content": message}, 
+                {"role": "assistant", "content": f"❌ Error processing message: {str(e)}"}
+            ], gr.Textbox(value="")
     
     async def _process_query(self, message: str, history: List[Union[Dict[str, Any], ChatMessage]]):
         claude_messages = []
@@ -192,10 +220,25 @@ class MCPClientWrapper:
 client = MCPClientWrapper()
 
 def gradio_interface():
-    with gr.Blocks(title="MCP Weather Client") as demo:
-        gr.Markdown("# MCP Weather Assistant")
-        gr.Markdown("Connect to your MCP weather server and chat with the assistant")
+    with gr.Blocks(title="MCP client") as demo:
+        gr.Markdown("# MCP test")
+        gr.Markdown("Connect to mcp server and chat with the assistant")
+        
+        # API Key Section
+        with gr.Row(equal_height=True):
+            with gr.Column(scale=4):
+                api_key_input = gr.Textbox(
+                    label="Anthropic API Key",
+                    placeholder="Enter your Anthropic API key (sk-ant-...)",
+                    type="password",
+                    value=""
+                )
+            with gr.Column(scale=1):
+                set_key_btn = gr.Button("Set API Key")
+        
+        api_key_status = gr.Textbox(label="API Key Status", interactive=False)
         
+        # Server Connection Section
         with gr.Row(equal_height=True):
             with gr.Column(scale=4):
                 server_path = gr.Textbox(
@@ -206,8 +249,9 @@ def gradio_interface():
             with gr.Column(scale=1):
                 connect_btn = gr.Button("Connect")
         
-        status = gr.Textbox(label="Connection Status", interactive=False)
+        connection_status = gr.Textbox(label="Connection Status", interactive=False)
         
+        # Chat Section
         chatbot = gr.Chatbot(
             value=[], 
             height=500,
@@ -224,15 +268,14 @@ def gradio_interface():
             )
             clear_btn = gr.Button("Clear Chat", scale=1)
         
-        connect_btn.click(client.connect, inputs=server_path, outputs=status)
+        # Event handlers
+        set_key_btn.click(client.set_api_key, inputs=api_key_input, outputs=api_key_status)
+        connect_btn.click(client.connect, inputs=server_path, outputs=connection_status)
         msg.submit(client.process_message, [msg, chatbot], [chatbot, msg])
         clear_btn.click(lambda: [], None, chatbot)
         
     return demo
 
 if __name__ == "__main__":
-    if not os.getenv("ANTHROPIC_API_KEY"):
-        print("Warning: ANTHROPIC_API_KEY not found in environment. Please set it in your .env file.")
-    
     interface = gradio_interface()
     interface.launch(debug=True)