Spaces:
Running
Running
| """ | |
| MCP Client for connecting to TraceMind-mcp-server | |
| Uses MCP protocol over HTTP to call remote MCP tools | |
| """ | |
| import os | |
| import asyncio | |
| from typing import Optional, Dict, Any, List | |
| from mcp import ClientSession, StdioServerParameters | |
| from mcp.client.sse import sse_client | |
| import aiohttp | |
| class MCPClient: | |
| """Client for interacting with TraceMind MCP Server""" | |
| def __init__(self, server_url: Optional[str] = None): | |
| """ | |
| Initialize MCP Client | |
| Args: | |
| server_url: URL of the TraceMind-mcp-server endpoint | |
| If None, uses MCP_SERVER_URL from environment | |
| """ | |
| self.server_url = server_url or os.getenv( | |
| 'MCP_SERVER_URL', | |
| 'https://mcp-1st-birthday-tracemind-mcp-server.hf.space/gradio_api/mcp/' | |
| ) | |
| self.session: Optional[ClientSession] = None | |
| self._initialized = False | |
| async def initialize(self): | |
| """Initialize connection to MCP server""" | |
| if self._initialized: | |
| return | |
| try: | |
| # Connect to SSE endpoint | |
| async with sse_client(self.server_url) as (read, write): | |
| async with ClientSession(read, write) as session: | |
| self.session = session | |
| await session.initialize() | |
| self._initialized = True | |
| # List available tools for verification | |
| tools_result = await session.list_tools() | |
| print(f"✅ Connected to TraceMind MCP Server at {self.server_url}") | |
| print(f"📊 Available tools: {len(tools_result.tools)}") | |
| for tool in tools_result.tools: | |
| print(f" - {tool.name}: {tool.description}") | |
| except Exception as e: | |
| print(f"❌ Failed to connect to MCP server: {e}") | |
| raise | |
| async def analyze_leaderboard( | |
| self, | |
| leaderboard_repo: str = "kshitijthakkar/smoltrace-leaderboard", | |
| metric_focus: str = "overall", | |
| time_range: str = "last_week", | |
| top_n: int = 5, | |
| hf_token: Optional[str] = None, | |
| gemini_api_key: Optional[str] = None | |
| ) -> str: | |
| """ | |
| Call the analyze_leaderboard tool on MCP server | |
| Args: | |
| leaderboard_repo: HuggingFace dataset repo for leaderboard | |
| metric_focus: Focus metric (overall, accuracy, cost, latency, co2) | |
| time_range: Time range filter (last_week, last_month, all_time) | |
| top_n: Number of top models to highlight | |
| hf_token: HuggingFace API token (optional if public dataset) | |
| gemini_api_key: Google Gemini API key (optional, server may have it) | |
| Returns: | |
| AI-generated analysis of the leaderboard | |
| """ | |
| if not self._initialized: | |
| await self.initialize() | |
| try: | |
| # Build arguments | |
| args = { | |
| "leaderboard_repo": leaderboard_repo, | |
| "metric_focus": metric_focus, | |
| "time_range": time_range, | |
| "top_n": top_n | |
| } | |
| # Add optional tokens if provided | |
| if hf_token: | |
| args["hf_token"] = hf_token | |
| if gemini_api_key: | |
| args["gemini_api_key"] = gemini_api_key | |
| # Call MCP tool | |
| result = await self.session.call_tool("analyze_leaderboard", arguments=args) | |
| # Extract text from result | |
| if result.content and len(result.content) > 0: | |
| return result.content[0].text | |
| else: | |
| return "No analysis generated" | |
| except Exception as e: | |
| return f"❌ Error calling analyze_leaderboard: {str(e)}" | |
| async def debug_trace( | |
| self, | |
| trace_data: Dict[str, Any], | |
| question: str, | |
| metrics_data: Optional[Dict[str, Any]] = None, | |
| hf_token: Optional[str] = None, | |
| gemini_api_key: Optional[str] = None | |
| ) -> str: | |
| """ | |
| Call the debug_trace tool on MCP server | |
| Args: | |
| trace_data: OpenTelemetry trace data (dict with spans) | |
| question: User question about the trace | |
| metrics_data: Optional GPU metrics data | |
| hf_token: HuggingFace API token | |
| gemini_api_key: Google Gemini API key | |
| Returns: | |
| AI-generated answer to the trace question | |
| """ | |
| if not self._initialized: | |
| await self.initialize() | |
| try: | |
| args = { | |
| "trace_data": trace_data, | |
| "question": question | |
| } | |
| if metrics_data: | |
| args["metrics_data"] = metrics_data | |
| if hf_token: | |
| args["hf_token"] = hf_token | |
| if gemini_api_key: | |
| args["gemini_api_key"] = gemini_api_key | |
| result = await self.session.call_tool("debug_trace", arguments=args) | |
| if result.content and len(result.content) > 0: | |
| return result.content[0].text | |
| else: | |
| return "No answer generated" | |
| except Exception as e: | |
| return f"❌ Error calling debug_trace: {str(e)}" | |
| async def estimate_cost( | |
| self, | |
| model: str, | |
| agent_type: str = "both", | |
| num_tests: int = 100, | |
| hardware: Optional[str] = None, | |
| hf_token: Optional[str] = None, | |
| gemini_api_key: Optional[str] = None | |
| ) -> str: | |
| """ | |
| Call the estimate_cost tool on MCP server | |
| Args: | |
| model: Model name (e.g., 'openai/gpt-4', 'meta-llama/Llama-3.1-8B') | |
| agent_type: Agent type (tool, code, both) | |
| num_tests: Number of tests to run | |
| hardware: Hardware type (cpu, gpu_a10, gpu_h200) | |
| hf_token: HuggingFace API token | |
| gemini_api_key: Google Gemini API key | |
| Returns: | |
| Cost estimation with breakdown | |
| """ | |
| if not self._initialized: | |
| await self.initialize() | |
| try: | |
| args = { | |
| "model": model, | |
| "agent_type": agent_type, | |
| "num_tests": num_tests | |
| } | |
| if hardware: | |
| args["hardware"] = hardware | |
| if hf_token: | |
| args["hf_token"] = hf_token | |
| if gemini_api_key: | |
| args["gemini_api_key"] = gemini_api_key | |
| result = await self.session.call_tool("estimate_cost", arguments=args) | |
| if result.content and len(result.content) > 0: | |
| return result.content[0].text | |
| else: | |
| return "No estimation generated" | |
| except Exception as e: | |
| return f"❌ Error calling estimate_cost: {str(e)}" | |
| async def compare_runs( | |
| self, | |
| run_data_list: List[Dict[str, Any]], | |
| focus_metrics: Optional[List[str]] = None, | |
| hf_token: Optional[str] = None, | |
| gemini_api_key: Optional[str] = None | |
| ) -> str: | |
| """ | |
| Call the compare_runs tool on MCP server | |
| Args: | |
| run_data_list: List of run data dicts from leaderboard | |
| focus_metrics: List of metrics to focus on | |
| hf_token: HuggingFace API token | |
| gemini_api_key: Google Gemini API key | |
| Returns: | |
| AI-generated comparison analysis | |
| """ | |
| if not self._initialized: | |
| await self.initialize() | |
| try: | |
| args = { | |
| "run_data_list": run_data_list | |
| } | |
| if focus_metrics: | |
| args["focus_metrics"] = focus_metrics | |
| if hf_token: | |
| args["hf_token"] = hf_token | |
| if gemini_api_key: | |
| args["gemini_api_key"] = gemini_api_key | |
| result = await self.session.call_tool("compare_runs", arguments=args) | |
| if result.content and len(result.content) > 0: | |
| return result.content[0].text | |
| else: | |
| return "No comparison generated" | |
| except Exception as e: | |
| return f"❌ Error calling compare_runs: {str(e)}" | |
| async def analyze_results( | |
| self, | |
| results_data: List[Dict[str, Any]], | |
| analysis_focus: str = "optimization", | |
| hf_token: Optional[str] = None, | |
| gemini_api_key: Optional[str] = None | |
| ) -> str: | |
| """ | |
| Call the analyze_results tool on MCP server | |
| Args: | |
| results_data: List of test case results | |
| analysis_focus: Focus area (optimization, failures, performance, cost) | |
| hf_token: HuggingFace API token | |
| gemini_api_key: Google Gemini API key | |
| Returns: | |
| AI-generated results analysis with recommendations | |
| """ | |
| if not self._initialized: | |
| await self.initialize() | |
| try: | |
| args = { | |
| "results_data": results_data, | |
| "analysis_focus": analysis_focus | |
| } | |
| if hf_token: | |
| args["hf_token"] = hf_token | |
| if gemini_api_key: | |
| args["gemini_api_key"] = gemini_api_key | |
| result = await self.session.call_tool("analyze_results", arguments=args) | |
| if result.content and len(result.content) > 0: | |
| return result.content[0].text | |
| else: | |
| return "No analysis generated" | |
| except Exception as e: | |
| return f"❌ Error calling analyze_results: {str(e)}" | |
| async def get_dataset_info( | |
| self, | |
| dataset_repo: str, | |
| hf_token: Optional[str] = None, | |
| gemini_api_key: Optional[str] = None | |
| ) -> str: | |
| """ | |
| Call the get_dataset tool on MCP server (resource) | |
| Args: | |
| dataset_repo: HuggingFace dataset repo | |
| hf_token: HuggingFace API token | |
| gemini_api_key: Google Gemini API key | |
| Returns: | |
| Dataset information and structure | |
| """ | |
| if not self._initialized: | |
| await self.initialize() | |
| try: | |
| args = { | |
| "dataset_repo": dataset_repo | |
| } | |
| if hf_token: | |
| args["hf_token"] = hf_token | |
| if gemini_api_key: | |
| args["gemini_api_key"] = gemini_api_key | |
| result = await self.session.call_tool("get_dataset", arguments=args) | |
| if result.content and len(result.content) > 0: | |
| return result.content[0].text | |
| else: | |
| return "No dataset info generated" | |
| except Exception as e: | |
| return f"❌ Error calling get_dataset: {str(e)}" | |
| async def close(self): | |
| """Close the MCP client session""" | |
| if self.session: | |
| # Note: ClientSession doesn't have an explicit close method | |
| # The context manager handles cleanup | |
| self.session = None | |
| self._initialized = False | |
| # Singleton instance for use across the app | |
| _mcp_client_instance: Optional[MCPClient] = None | |
| def get_mcp_client() -> MCPClient: | |
| """Get or create the global MCP client instance""" | |
| global _mcp_client_instance | |
| if _mcp_client_instance is None: | |
| _mcp_client_instance = MCPClient() | |
| return _mcp_client_instance | |