Spaces:
Sleeping
Sleeping
File size: 2,378 Bytes
881aff7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 | """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)
|