TR updates, first round

Author: lpozo

python-mcp-client/source-code-final/mcp_client/__main__.py

@@ -4,25 +4,27 @@
 from mcp_client.mcp_client import MCPClient
 
 
-def main():
-    """Entry point for the mcp-client CLI app."""
+async def main() -> None:
+    """Run the MCP client with the specified options."""
     args = parse_args()
     if not args.server_path.exists():
         print(f"Error: Server script '{args.server_path}' not found")
         return
 
-    async def run():
-        try:
-            async with MCPClient(args.server_path) as client:
-                if args.members:
-                    await client.list_all_members()
-                elif args.chat:
-                    await client.run_chat()
-        except RuntimeError as e:
-            print(e)
+    try:
+        async with MCPClient(str(args.server_path)) as client:
+            if args.members:
+                await client.list_all_members()
+            elif args.chat:
+                await client.run_chat()
+    except RuntimeError as e:
+        print(e)
+
 
-    asyncio.run(run())
+def cli_main():
+    """Entry point for the mcp-client CLI app."""
+    asyncio.run(main())
 
 
 if __name__ == "__main__":
-    main()
+    asyncio.run(main())

python-mcp-client/source-code-final/mcp_client/handlers.py

@@ -89,7 +89,10 @@ async def _execute_tool(self, tool_call) -> dict:
         tool_args = json.loads(tool_call.function.arguments or "{}")
 
         try:
-            result = await self.client_session.call_tool(tool_name, tool_args)
+            result = await self.client_session.call_tool(
+                tool_name,
+                tool_args,
+            )
             content = result.content[0].text if result.content else ""
             log = f"[Used {tool_name}({tool_args})]"
         except Exception as e:

python-mcp-client/source-code-final/mcp_client/mcp_client.py

@@ -1,6 +1,6 @@
 import sys
 from contextlib import AsyncExitStack
-from typing import Self
+from typing import Any, Awaitable, Callable, Self
 
 from mcp import ClientSession, StdioServerParameters
 from mcp.client.stdio import stdio_client
@@ -17,12 +17,39 @@ class MCPClient:
             # Call client methods here...
     """
 
-    client_session: ClientSession
-
     def __init__(self, server_path: str):
         self.server_path = server_path
         self.exit_stack = AsyncExitStack()
 
+    async def __aenter__(self) -> Self:
+        self.client_session = await self._connect_to_server()
+        return self
+
+    async def __aexit__(self, *_) -> None:
+        await self.exit_stack.aclose()
+
+    async def _connect_to_server(self) -> ClientSession:
+        try:
+            read, write = await self.exit_stack.enter_async_context(
+                stdio_client(
+                    server=StdioServerParameters(
+                        command="sh",
+                        args=[
+                            "-c",
+                            f"{sys.executable} {self.server_path} 2>/dev/null",
+                        ],
+                        env=None,
+                    )
+                )
+            )
+            client_session = await self.exit_stack.enter_async_context(
+                ClientSession(read, write)
+            )
+            await client_session.initialize()
+            return client_session
+        except Exception:
+            raise RuntimeError("Error: Failed to connect to server")
+
     async def list_all_members(self) -> None:
         """List all available tools, prompts, and resources."""
         print("MCP Server Members")
@@ -38,28 +65,13 @@ async def list_all_members(self) -> None:
 
         print("\n" + "=" * 50)
 
-    async def run_chat(self) -> None:
-        """Start interactive chat with MCP server using OpenAI."""
-        try:
-            handler = OpenAIQueryHandler(self.client_session)
-            await chat.run_chat(handler)
-        except RuntimeError as e:
-            print(e)
-
-    async def __aenter__(self) -> Self:
-        self.client_session = await self._connect_to_server()
-        return self
-
-    async def __aexit__(self, *_) -> None:
-        await self.exit_stack.aclose()
-
     async def _list_section(
         self,
         section: str,
-        list_func,
+        list_method: Callable[[], Awaitable[Any]],
     ) -> None:
         try:
-            items = getattr(await list_func(), section)
+            items = getattr(await list_method(), section)
             if items:
                 print(f"\n{section.upper()} ({len(items)}):")
                 print("-" * 30)
@@ -71,24 +83,10 @@ async def _list_section(
         except Exception as e:
             print(f"\n{section.upper()}: Error - {e}")
 
-    async def _connect_to_server(self) -> ClientSession:
+    async def run_chat(self) -> None:
+        """Start interactive chat with MCP server using OpenAI."""
         try:
-            stdio, write = await self.exit_stack.enter_async_context(
-                stdio_client(
-                    server=StdioServerParameters(
-                        command="sh",
-                        args=[
-                            "-c",
-                            f"{sys.executable} {self.server_path} 2>/dev/null",
-                        ],
-                        env=None,
-                    )
-                )
-            )
-            client_session = await self.exit_stack.enter_async_context(
-                ClientSession(stdio, write)
-            )
-            await client_session.initialize()
-            return client_session
-        except Exception:
-            raise RuntimeError("Error: Failed to connect to server")
+            handler = OpenAIQueryHandler(self.client_session)
+            await chat.run_chat(handler)
+        except RuntimeError as e:
+            print(e)

python-mcp-client/source-code-final/mcp_server/mcp_server.py

@@ -10,7 +10,7 @@ dependencies = [
 ]
 
 [project.scripts]
-mcp-client = "mcp_client.__main__:main"
+mcp-client = "mcp_client.__main__:cli_main"
 
 [build-system]
 requires = ["setuptools>=61.0", "wheel"]