avsolatorio commited on
Commit
11aa784
·
1 Parent(s): e6e4f17

Add disconnect to client

Browse files

Signed-off-by: Aivin V. Solatorio <avsolatorio@gmail.com>

Files changed (1) hide show
  1. mcp_client.py +13 -3
mcp_client.py CHANGED
@@ -126,6 +126,12 @@ class MCPClientWrapper:
126
  tool_names = [tool["name"] for tool in self.tools]
127
  return f"Connected to MCP server. Available tools: {', '.join(tool_names)}"
128
 
 
 
 
 
 
 
129
  async def process_message(
130
  self, message: str, history: List[Union[Dict[str, Any], ChatMessage]]
131
  ):
@@ -353,7 +359,7 @@ def gradio_interface():
353
 
354
  chatbot = gr.Chatbot(
355
  value=[],
356
- height=600,
357
  type="messages",
358
  show_copy_button=True,
359
  avatar_images=("img/small-user.png", "img/small-robot.png"),
@@ -368,14 +374,18 @@ def gradio_interface():
368
  )
369
  clear_btn = gr.Button("Clear Chat", scale=1)
370
 
371
- connect_btn.click(client.connect, inputs=server_path, outputs=status)
372
 
373
  # Automatically call client.connect(...) as soon as the interface loads
374
- demo.load(fn=client.connect, inputs=server_path, outputs=status)
 
 
375
 
376
  msg.submit(client.process_message, [msg, chatbot], [chatbot, msg])
377
  clear_btn.click(lambda: [], None, chatbot)
378
 
 
 
379
  return demo
380
 
381
 
 
126
  tool_names = [tool["name"] for tool in self.tools]
127
  return f"Connected to MCP server. Available tools: {', '.join(tool_names)}"
128
 
129
+ async def disconnect(self):
130
+ if self.exit_stack:
131
+ await self.exit_stack.aclose()
132
+ self.exit_stack = None
133
+ self.session = None
134
+
135
  async def process_message(
136
  self, message: str, history: List[Union[Dict[str, Any], ChatMessage]]
137
  ):
 
359
 
360
  chatbot = gr.Chatbot(
361
  value=[],
362
+ height="70vh",
363
  type="messages",
364
  show_copy_button=True,
365
  avatar_images=("img/small-user.png", "img/small-robot.png"),
 
374
  )
375
  clear_btn = gr.Button("Clear Chat", scale=1)
376
 
377
+ # connect_btn.click(client.connect, inputs=server_path, outputs=status)
378
 
379
  # Automatically call client.connect(...) as soon as the interface loads
380
+ demo.load(
381
+ fn=client.connect, inputs=server_path, outputs=status, show_progress="full"
382
+ )
383
 
384
  msg.submit(client.process_message, [msg, chatbot], [chatbot, msg])
385
  clear_btn.click(lambda: [], None, chatbot)
386
 
387
+ demo.unload(client.disconnect)
388
+
389
  return demo
390
 
391