Claude Data Buddy

# client.py """ MCP Client for COMS 6998 - Supports listing tools - Calling server tools - Retrieving resources - Interactive mode - Automated tests """ import asyncio import json from pathlib import Path from mcp.client.session import ClientSession from mcp.client.stdio import stdio_client, StdioServerParameters class MCPFileAnalyzerClient: """Enhanced MCP client with connection management and error handling.""" def __init__(self, server_script_path: str = "main.py"): self.server_script_path = server_script_path self.server_params = StdioServerParameters( command="python", args=[str(Path(server_script_path).resolve())], env=None ) self._stdio_transport = None self.session = None async def connect(self): """Connect to MCP server with proper error handling.""" try: self._stdio_transport = stdio_client(self.server_params) read, write = await self._stdio_transport.__aenter__() self.session = ClientSession(read, write) await self.session.__aenter__() await self.session.initialize() return True except Exception as e: print(f"❌ Failed to connect to MCP server: {e}") return False async def disconnect(self): """Disconnect from server with proper cleanup.""" if self.session: try: await self.session.__aexit__(None, None, None) except Exception: pass self.session = None if self._stdio_transport: try: await self._stdio_transport.__aexit__(None, None, None) except Exception: pass self._stdio_transport = None async def list_tools(self): """List all available tools from the server.""" if not self.session: return None try: return await self.session.list_tools() except Exception as e: print(f"❌ Error listing tools: {e}") return None async def list_resources(self): """List all available resources from the server.""" if not self.session: return None try: return await self.session.list_resources() except Exception as e: print(f"❌ Error listing resources: {e}") return None async def call_tool(self, tool_name: str, arguments: dict = None): """Call a specific tool with arguments and error handling.""" if not self.session: print("❌ Not connected to server") return None if arguments is None: arguments = {} try: result = await self.session.call_tool(tool_name, arguments) return result except Exception as e: return f"❌ Error calling tool '{tool_name}': {e}" async def get_resource(self, uri: str): """Retrieve a resource by URI.""" if not self.session: print("❌ Not connected to server") return None try: from mcp import types result = await self.session.read_resource(types.AnyUrl(uri)) if result.contents: return result.contents[0].text if hasattr(result.contents[0], 'text') else result.contents[0] return None except Exception as e: return f"❌ Error getting resource '{uri}': {e}" async def run_demo(): """Simple non-interactive demo.""" client = MCPFileAnalyzerClient() if not await client.connect(): return try: print("📌 Tools available:") tools = await client.list_tools() if tools: for tool in tools.tools: print(f" - {tool.name}: {tool.description}") print("\n📌 Resources available:") resources = await client.list_resources() if resources: for resource in resources.resources: print(f" - {resource.uri}: {resource.name}") else: print(" (No resources available)") print("\n📌 Creating sample data:") result = await client.call_tool("create_sample", {}) print(result) print("\n📌 Listing data files:") result = await client.call_tool("list_data_files", {}) print(result) print("\n📌 Summarizing sample.csv:") result = await client.call_tool("summarize_csv", {"file_name": "sample.csv"}) print(result) finally: await client.disconnect() async def interactive(): """Interactive mode with command parsing.""" client = MCPFileAnalyzerClient() if not await client.connect(): return try: print("🟢 MCP Client Interactive Mode") print("Type 'exit' or 'quit' to quit.") print("Type 'help' for available commands.") print("\nAvailable tools:") tools = await client.list_tools() if tools: for tool in tools.tools: print(f" - {tool.name}") while True: try: cmd = input("\nEnter command: ").strip() if cmd in ["exit", "quit"]: break elif cmd == "help": print("\nCommands:") print(" list_tools - List all available tools") print(" list_resources - List all available resources") print(" list_files - List data files") print(" summarize <file> - Summarize a CSV file") print(" analyze <file> <operation> - Analyze CSV (operations: describe, head, info, columns)") print(" comprehensive <file> - Comprehensive analysis") print(" compare <file1> <file2> - Compare two files") print(" create <rows> <filename> - Create custom dataset") print(" exit/quit - Exit") elif cmd == "list_tools": tools = await client.list_tools() if tools: for tool in tools.tools: print(f" {tool.name}: {tool.description}") elif cmd == "list_resources": resources = await client.list_resources() if resources and resources.resources: for resource in resources.resources: print(f" {resource.uri}: {resource.name}") else: print(" No resources available") elif cmd == "list_files": result = await client.call_tool("list_data_files", {}) print(result) elif cmd.startswith("summarize "): filename = cmd.split(" ", 1)[1] result = await client.call_tool("summarize_csv", {"file_name": filename}) print(result) elif cmd.startswith("analyze "): parts = cmd.split(" ", 2) if len(parts) >= 3: filename = parts[1] operation = parts[2] result = await client.call_tool("analyze_csv", { "file_name": filename, "operation": operation }) print(result) else: print("Usage: analyze <filename> <operation>") elif cmd.startswith("comprehensive "): filename = cmd.split(" ", 1)[1] result = await client.call_tool("comprehensive_analysis", {"file_name": filename}) print(result) elif cmd.startswith("compare "): parts = cmd.split(" ", 2) if len(parts) >= 3: file1 = parts[1] file2 = parts[2] result = await client.call_tool("compare_files", {"file1": file1, "file2": file2}) print(result) else: print("Usage: compare <file1> <file2>") elif cmd.startswith("create "): parts = cmd.split(" ", 2) if len(parts) >= 3: rows = int(parts[1]) filename = parts[2] result = await client.call_tool("create_custom_dataset", { "rows": rows, "file_name": filename }) print(result) else: print("Usage: create <rows> <filename>") else: # Try as direct tool call try: params_str = input("Enter parameters as JSON dict (or press Enter for {}): ") params = json.loads(params_str) if params_str else {} result = await client.call_tool(cmd, params) print(f"Result: {result}") except json.JSONDecodeError: print("❌ Invalid JSON. Use format: {\"key\": \"value\"}") except Exception as e: print(f"❌ Error: {e}") except KeyboardInterrupt: print("\n🛑 Interrupted by user") break except Exception as e: print(f"❌ Error: {e}") finally: await client.disconnect() # Run default demo if __name__ == "__main__": import sys if len(sys.argv) > 1 and sys.argv[1] == "interactive": asyncio.run(interactive()) else: asyncio.run(run_demo())