Improve interactivity

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

mcp_client.py

@@ -11,12 +11,10 @@ from mcp.client.stdio import stdio_client
 from anthropic import Anthropic
 from anthropic._exceptions import OverloadedError
 from dotenv import load_dotenv
+import functools
 
 load_dotenv()
 
-loop = asyncio.new_event_loop()
-asyncio.set_event_loop(loop)
-
 # SYSTEM_PROMPT = f"""You are a helpful assistant and today is {datetime.now().strftime("%Y-%m-%d")}.
 
 # You do not have any knowledge of the World Development Indicators (WDI) data. However, you can use the tools provided to answer questions.
@@ -76,6 +74,10 @@ LLM_MODEL = "claude-3-5-haiku-20241022"
 # When a tool is needed for any step, ensure to add the token `TOOL_USE`.
 
 
+loop = asyncio.new_event_loop()
+asyncio.set_event_loop(loop)
+
+
 class MCPClientWrapper:
     def __init__(self):
         self.session = None
@@ -83,15 +85,55 @@ class MCPClientWrapper:
         self.anthropic = Anthropic()
         self.tools = []
 
-    def connect(self, server_path: str) -> str:
-        return loop.run_until_complete(self._connect(server_path))
+    # def connect(self, server_path: str) -> str:
+    #     return loop.run_until_complete(self._connect(server_path))
+
+    # async def _connect(self, server_path: str) -> str:
+    #     if self.exit_stack:
+    #         await self.exit_stack.aclose()
+
+    #     self.exit_stack = AsyncExitStack()
+
+    #     is_python = server_path.endswith(".py")
+    #     command = "python" if is_python else "node"
+
+    #     server_params = StdioServerParameters(
+    #         command=command,
+    #         args=[server_path],
+    #         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()
 
-    async def _connect(self, server_path: str) -> str:
+    #     response = await self.session.list_tools()
+    #     self.tools = [
+    #         {
+    #             "name": tool.name,
+    #             "description": tool.description,
+    #             "input_schema": tool.inputSchema,
+    #         }
+    #         for tool in response.tools
+    #     ]
+
+    #     print(self.tools)
+
+    #     tool_names = [tool["name"] for tool in self.tools]
+    #     return f"Connected to MCP server. Available tools: {', '.join(tool_names)}"
+
+    async def connect(self, server_path: str) -> str:
+        # If there's an existing session, close it
         if self.exit_stack:
             await self.exit_stack.aclose()
 
         self.exit_stack = AsyncExitStack()
-
         is_python = server_path.endswith(".py")
         command = "python" if is_python else "node"
 
@@ -101,11 +143,13 @@ class MCPClientWrapper:
             env={"PYTHONIOENCODING": "utf-8", "PYTHONUNBUFFERED": "1"},
         )
 
+        # Launch MCP subprocess and bind streams on the *current* running loop
         stdio_transport = await self.exit_stack.enter_async_context(
             stdio_client(server_params)
         )
         self.stdio, self.write = stdio_transport
 
+        # Create ClientSession on this same loop
         self.session = await self.exit_stack.enter_async_context(
             ClientSession(self.stdio, self.write)
         )
@@ -121,14 +165,13 @@ class MCPClientWrapper:
             for tool in response.tools
         ]
 
-        print(self.tools)
-
+        print("Available tools:", self.tools)
         tool_names = [tool["name"] for tool in self.tools]
         return f"Connected to MCP server. Available tools: {', '.join(tool_names)}"
 
-    def process_message(
+    async def process_message(
         self, message: str, history: List[Union[Dict[str, Any], ChatMessage]]
-    ) -> tuple:
+    ):
         if not self.session:
             messages = history + [
                 {"role": "user", "content": message},
@@ -137,17 +180,25 @@ class MCPClientWrapper:
                     "content": "Please connect to an MCP server first.",
                 },
             ]
+            yield messages, gr.Textbox(value="")
         else:
-            new_messages = loop.run_until_complete(
-                self._process_query(message, history)
-            )
-            messages = history + [{"role": "user", "content": message}] + new_messages
+            # new_messages = loop.run_until_complete(
+            #     self._process_query(message, history)
+            # )
+            # messages = history + [{"role": "user", "content": message}] + new_messages
+
+            messages = history + [{"role": "user", "content": message}]
+
+            yield messages, gr.Textbox(value="")
+
+            async for partial in self._process_query(message, history):
+                messages.extend(partial)
+
+                yield messages, gr.Textbox(value="")
 
         with open("messages.log.jsonl", "a+") as fl:
             fl.write(json.dumps(dict(time=f"{datetime.now()}", messages=messages)))
 
-        return messages, gr.Textbox(value="")
-
     async def _process_query(
         self, message: str, history: List[Union[Dict[str, Any], ChatMessage]]
     ):
@@ -173,7 +224,7 @@ class MCPClientWrapper:
                 tools=self.tools,
             )
         except OverloadedError:
-            return [
+            yield [
                 {
                     "role": "assistant",
                     "content": "The LLM API is overloaded now, try again later...",
@@ -181,6 +232,7 @@ class MCPClientWrapper:
             ]
 
         result_messages = []
+        partial_messages = []
 
         print(response.content)
         contents = response.content
@@ -194,6 +246,9 @@ class MCPClientWrapper:
             if content.type == "text":
                 result_messages.append({"role": "assistant", "content": content.text})
                 claude_messages.append({"role": "assistant", "content": content.text})
+                partial_messages.append(result_messages[-1])
+                yield [result_messages[-1]]
+                partial_messages = []
 
                 # if (
                 #     auto_calls < MAX_CALLS
@@ -242,6 +297,8 @@ class MCPClientWrapper:
                         },
                     }
                 )
+                partial_messages.append(result_messages[-1])
+                yield [result_messages[-1]]
 
                 result_messages.append(
                     {
@@ -256,6 +313,8 @@ class MCPClientWrapper:
                         },
                     }
                 )
+                partial_messages.append(result_messages[-1])
+                yield [result_messages[-1]]
 
                 print(f"Calling tool: {tool_name} with args: {tool_args}")
                 result = await self.session.call_tool(tool_name, tool_args)
@@ -274,6 +333,9 @@ class MCPClientWrapper:
                         },
                     }
                 )
+                partial_messages.append(result_messages[-1])
+                yield [result_messages[-1]]
+                partial_messages = []
 
                 result_content = result.content
                 print(result_content)
@@ -304,6 +366,9 @@ class MCPClientWrapper:
                         },
                     }
                 )
+                partial_messages.append(result_messages[-1])
+                yield [result_messages[-1]]
+                partial_messages = []
 
                 # claude_messages.append(
                 #     {
@@ -337,7 +402,7 @@ class MCPClientWrapper:
                     )
                     auto_calls += 1
                 except OverloadedError:
-                    return [
+                    yield [
                         {
                             "role": "assistant",
                             "content": "The LLM API is overloaded now, try again later...",
@@ -353,7 +418,7 @@ class MCPClientWrapper:
                 #         {"role": "assistant", "content": next_response.content[0].text}
                 #     )
 
-        return result_messages
+        # yield result_messages
 
 
 client = MCPClientWrapper()