Spaces:
Running
Running
| import json | |
| import subprocess | |
| import sys | |
| from typing import Any, Dict, List, Optional | |
| class MCPClientError(RuntimeError): | |
| pass | |
| class MCPCalculationClient: | |
| def __init__( | |
| self, | |
| run_dir: str, | |
| python_executable: Optional[str] = None, | |
| server_script: Optional[str] = None, | |
| ): | |
| self.run_dir = str(run_dir) | |
| self.python_executable = python_executable or sys.executable | |
| if server_script: | |
| self.server_script = str(server_script) | |
| else: | |
| self.server_script = None | |
| self._proc: Optional[subprocess.Popen] = None | |
| self._id = 0 | |
| def start(self) -> None: | |
| if self._proc is not None: | |
| return | |
| if self.server_script: | |
| cmd = [ | |
| self.python_executable, | |
| self.server_script, | |
| "--stdio", | |
| "--run-dir", | |
| self.run_dir, | |
| ] | |
| else: | |
| cmd = [ | |
| self.python_executable, | |
| "-m", | |
| "mcp_tox_calc.server", | |
| "--stdio", | |
| "--run-dir", | |
| self.run_dir, | |
| ] | |
| self._proc = subprocess.Popen( | |
| cmd, | |
| stdin=subprocess.PIPE, | |
| stdout=subprocess.PIPE, | |
| stderr=subprocess.PIPE, | |
| text=True, | |
| bufsize=1, | |
| ) | |
| self._request("initialize", {"protocolVersion": "2024-11-05", "clientInfo": {"name": "toxra-app", "version": "0.1.0"}}) | |
| def stop(self) -> None: | |
| if self._proc is None: | |
| return | |
| try: | |
| if self._proc.stdin: | |
| self._proc.stdin.close() | |
| self._proc.terminate() | |
| self._proc.wait(timeout=3) | |
| except Exception: | |
| try: | |
| self._proc.kill() | |
| except Exception: | |
| pass | |
| finally: | |
| self._proc = None | |
| def __enter__(self): | |
| self.start() | |
| return self | |
| def __exit__(self, exc_type, exc, tb): | |
| self.stop() | |
| def _request(self, method: str, params: Dict[str, Any]) -> Dict[str, Any]: | |
| if self._proc is None: | |
| raise MCPClientError("MCP server not started.") | |
| if self._proc.stdin is None or self._proc.stdout is None: | |
| raise MCPClientError("MCP server pipes unavailable.") | |
| self._id += 1 | |
| req_id = self._id | |
| request = { | |
| "jsonrpc": "2.0", | |
| "id": req_id, | |
| "method": method, | |
| "params": params, | |
| } | |
| self._proc.stdin.write(json.dumps(request) + "\n") | |
| self._proc.stdin.flush() | |
| while True: | |
| line = self._proc.stdout.readline() | |
| if line == "": | |
| err = "" | |
| if self._proc.stderr is not None: | |
| try: | |
| err = self._proc.stderr.read()[-1500:] | |
| except Exception: | |
| err = "" | |
| raise MCPClientError(f"No response from MCP server. stderr={err}") | |
| line = line.strip() | |
| if not line: | |
| continue | |
| try: | |
| resp = json.loads(line) | |
| except Exception: | |
| continue | |
| if resp.get("id") != req_id: | |
| continue | |
| if "error" in resp: | |
| raise MCPClientError(str(resp["error"])) | |
| result = resp.get("result", {}) | |
| if not isinstance(result, dict): | |
| return {"result": result} | |
| return result | |
| def list_tools(self) -> List[Dict[str, Any]]: | |
| result = self._request("tools/list", {}) | |
| tools = result.get("tools", []) | |
| return tools if isinstance(tools, list) else [] | |
| def call_tool(self, name: str, arguments: Dict[str, Any]) -> Dict[str, Any]: | |
| result = self._request("tools/call", {"name": name, "arguments": arguments}) | |
| content = result.get("content", []) if isinstance(result, dict) else [] | |
| if isinstance(content, list) and content: | |
| first = content[0] | |
| if isinstance(first, dict) and first.get("type") == "json": | |
| data = first.get("json", {}) | |
| return data if isinstance(data, dict) else {"value": data} | |
| return result if isinstance(result, dict) else {"value": result} | |
| def run_batch_cancer_risk(rows: List[Dict[str, Any]], run_dir: str) -> Dict[str, Any]: | |
| with MCPCalculationClient(run_dir=run_dir) as client: | |
| return client.call_tool("run_batch_cancer_risk", {"rows": rows}) | |