Spaces:
Sleeping
Sleeping
| """Persistent MCP client for the local paper retrieval server.""" | |
| from __future__ import annotations | |
| import json | |
| import os | |
| import sys | |
| from pathlib import Path | |
| from typing import Any | |
| from mcp import ClientSession, StdioServerParameters | |
| from mcp.client.stdio import stdio_client | |
| PROJECT_ROOT = Path(__file__).resolve().parents[1] | |
| class PaperMCPClient: | |
| """Keeps one MCP session open for retrieval tool calls.""" | |
| def __init__(self, project_root: Path | None = None): | |
| self.project_root = project_root or PROJECT_ROOT | |
| self._stdio_ctx = None | |
| self._session_ctx = None | |
| self._session: ClientSession | None = None | |
| async def connect(self) -> None: | |
| if self._session is not None: | |
| return | |
| server_script = os.getenv("MCP_SERVER_SCRIPT", "waterscope_retrieval_server.py") | |
| server_params = StdioServerParameters( | |
| command=sys.executable, | |
| args=[str(self.project_root / server_script)], | |
| cwd=str(self.project_root), | |
| env={**os.environ}, | |
| ) | |
| self._stdio_ctx = stdio_client(server_params) | |
| read, write = await self._stdio_ctx.__aenter__() | |
| self._session_ctx = ClientSession(read, write) | |
| self._session = await self._session_ctx.__aenter__() | |
| await self._session.initialize() | |
| async def close(self) -> None: | |
| if self._session_ctx is not None: | |
| await self._session_ctx.__aexit__(None, None, None) | |
| self._session_ctx = None | |
| self._session = None | |
| if self._stdio_ctx is not None: | |
| await self._stdio_ctx.__aexit__(None, None, None) | |
| self._stdio_ctx = None | |
| async def call_tool(self, name: str, arguments: dict[str, Any] | None = None) -> Any: | |
| if self._session is None: | |
| raise RuntimeError("MCP client is not connected.") | |
| payload = {key: value for key, value in (arguments or {}).items() if value is not None} | |
| result = await self._session.call_tool(name, payload) | |
| if result.isError: | |
| message = result.content[0].text if result.content else "Unknown MCP tool error" | |
| raise RuntimeError(message) | |
| if not result.content: | |
| return {} | |
| text = result.content[0].text | |
| return json.loads(text) | |