diff --git a/reme4/components/client/base_client.py b/reme4/components/client/base_client.py index a5b05b6c..350c51eb 100644 --- a/reme4/components/client/base_client.py +++ b/reme4/components/client/base_client.py @@ -24,18 +24,18 @@ class BaseClient(BaseComponent): """Close the client and release resources.""" @abstractmethod - def _execute(self) -> AsyncGenerator[str, None]: + def _execute(self, action: str, payload: dict) -> AsyncGenerator[str, None]: """Backend-specific execution; yield text chunks (single yield for non-streaming backends).""" @abstractmethod async def list_actions(self) -> list[dict]: """Discover available actions on the server; each dict is the raw backend descriptor.""" - async def __call__(self) -> AsyncGenerator[str, None]: + async def __call__(self, action: str, **kwargs) -> AsyncGenerator[str, None]: """Dispatch: action='list' returns the action catalog; otherwise delegate to _execute().""" - if getattr(self, "action", None) == "list": + if action == "list": actions = await self.list_actions() yield json.dumps(actions, indent=2, ensure_ascii=False) return - async for chunk in self._execute(): + async for chunk in self._execute(action, kwargs): yield chunk diff --git a/reme4/components/client/http_client.py b/reme4/components/client/http_client.py index efec1ff7..871c0d3a 100644 --- a/reme4/components/client/http_client.py +++ b/reme4/components/client/http_client.py @@ -19,7 +19,6 @@ class HttpClient(BaseClient): def __init__( self, - action: str, host: str | None = None, port: int | None = None, timeout: float = 30.0, @@ -40,7 +39,6 @@ class HttpClient(BaseClient): else: host, port = REME_DEFAULT_HOST, REME_DEFAULT_PORT - self.action = action self.base_url = f"http://{host}:{port}" self.timeout = timeout @@ -49,7 +47,7 @@ class HttpClient(BaseClient): if self.client is None: self.client = httpx.AsyncClient(base_url=self.base_url, timeout=self.timeout) - async def _iter_stream_chunks(self) -> AsyncGenerator[StreamChunk, None]: + async def _iter_stream_chunks(self, action: str, payload: dict) -> AsyncGenerator[StreamChunk, None]: """Send request and yield raw StreamChunks; auto-detects JSON vs SSE via Content-Type. For JSON responses: yields a single CONTENT chunk with the raw response body. @@ -58,7 +56,7 @@ class HttpClient(BaseClient): if self.client is None: raise RuntimeError("Client not initialized. Call _start() first.") - async with self.client.stream("POST", f"/{self.action}", json=self.kwargs) as resp: + async with self.client.stream("POST", f"/{action}", json=payload) as resp: resp.raise_for_status() ctype = resp.headers.get("content-type", "") @@ -66,11 +64,11 @@ class HttpClient(BaseClient): async for line in resp.aiter_lines(): if not line.startswith("data:"): continue - payload = line[len("data:") :] - if payload.strip() == "[DONE]": + data_str = line[len("data:") :] + if data_str.strip() == "[DONE]": return try: - data = json.loads(payload) + data = json.loads(data_str) except json.JSONDecodeError: continue chunk = StreamChunk(**data) @@ -85,9 +83,9 @@ class HttpClient(BaseClient): body = await resp.aread() yield StreamChunk(chunk_type=ChunkEnum.CONTENT, chunk=body.decode()) - async def stream_chunks(self) -> AsyncGenerator[StreamChunk, None]: + async def stream_chunks(self, action: str, **kwargs) -> AsyncGenerator[StreamChunk, None]: """HTTP-specific richer access: yield raw StreamChunk objects (no display formatting).""" - async for chunk in self._iter_stream_chunks(): + async for chunk in self._iter_stream_chunks(action, kwargs): yield chunk async def list_actions(self) -> list[dict]: @@ -129,11 +127,11 @@ class HttpClient(BaseClient): return "\n".join(parts) # pylint: disable=invalid-overridden-method - async def _execute(self) -> AsyncGenerator[str, None]: + async def _execute(self, action: str, payload: dict) -> AsyncGenerator[str, None]: """Yield text chunks for CLI display; JSON responses are pretty-formatted.""" - async for chunk in self._iter_stream_chunks(): - payload = chunk.chunk - text = payload if isinstance(payload, str) else json.dumps(payload, ensure_ascii=False) + async for chunk in self._iter_stream_chunks(action, payload): + chunk_payload = chunk.chunk + text = chunk_payload if isinstance(chunk_payload, str) else json.dumps(chunk_payload, ensure_ascii=False) yield self._format_for_display(text) async def _close(self) -> None: diff --git a/reme4/components/client/mcp_client.py b/reme4/components/client/mcp_client.py index b8b17a9b..6cb390cb 100644 --- a/reme4/components/client/mcp_client.py +++ b/reme4/components/client/mcp_client.py @@ -26,25 +26,24 @@ class MCPClient(BaseClient): Usage: # SSE (default) - client = MCPClient(action="my_tool", host="localhost", port=8000, query="hello") + client = MCPClient(host="localhost", port=8000) async with client: - async for text in client(): + async for text in client(action="my_tool", query="hello"): print(text) # Streamable HTTP - client = MCPClient(action="my_tool", transport="streamable-http", host="localhost", port=8000) + client = MCPClient(transport="streamable-http", host="localhost", port=8000) # Stdio - client = MCPClient(action="my_tool", transport="stdio", command="python", args=["server.py"]) + client = MCPClient(transport="stdio", command="python", args=["server.py"]) # Custom transport object from fastmcp.client import SSETransport - client = MCPClient(action="my_tool", transport=SSETransport(url="http://host:port/sse")) + client = MCPClient(transport=SSETransport(url="http://host:port/sse")) """ def __init__( self, - action: str, transport: str | Any = "sse", host: str | None = None, port: int | None = None, @@ -71,7 +70,6 @@ class MCPClient(BaseClient): self.host = host self.port = port - self.action = action self.transport = transport self.timeout = timeout @@ -97,11 +95,11 @@ class MCPClient(BaseClient): await self.client.__aenter__() # pylint: disable=invalid-overridden-method - async def _execute(self) -> AsyncGenerator[str, None]: + async def _execute(self, action: str, payload: dict) -> AsyncGenerator[str, None]: if self.client is None: raise RuntimeError("Client not initialized. Call _start() first.") - result: CallToolResult = await self.client.call_tool(self.action, self.kwargs) + result: CallToolResult = await self.client.call_tool(action, payload) yield self._extract_text(result) async def list_actions(self) -> list[dict]: diff --git a/reme4/reme.py b/reme4/reme.py index fe74bb7a..c8ad95a0 100644 --- a/reme4/reme.py +++ b/reme4/reme.py @@ -18,8 +18,8 @@ async def call_server(action: str, **kwargs): """Call the appropriate server component.""" backend: str = kwargs.pop("backend", "http") client_cls = R.get(ComponentEnum.CLIENT, backend) - async with client_cls(action=action, **kwargs) as client: - async for chunk in client(): + async with client_cls() as client: + async for chunk in client(action=action, **kwargs): print(chunk, end="", flush=True) print() diff --git a/reme4/utils/common_utils.py b/reme4/utils/common_utils.py index 7df0dbef..8c83535b 100644 --- a/reme4/utils/common_utils.py +++ b/reme4/utils/common_utils.py @@ -209,8 +209,8 @@ async def call_action( from ..components.client.http_client import HttpClient pieces: list[str] = [] - async with HttpClient(action=action, host=host, port=port, timeout=timeout, **kwargs) as client: - async for chunk in client.stream_chunks(): + async with HttpClient(host=host, port=port, timeout=timeout) as client: + async for chunk in client.stream_chunks(action, **kwargs): payload = chunk.chunk pieces.append(payload if isinstance(payload, str) else json.dumps(payload, ensure_ascii=False)) raw = "".join(pieces) diff --git a/reme4/utils/service_utils.py b/reme4/utils/service_utils.py index d1a39b9d..aebabc2a 100644 --- a/reme4/utils/service_utils.py +++ b/reme4/utils/service_utils.py @@ -13,8 +13,8 @@ async def find_reme(host: str, port: int) -> str: from ..components.client.http_client import HttpClient try: - async with HttpClient(action="health_check", host=host, port=port, timeout=2.0) as client: - async for _ in client(): + async with HttpClient(host=host, port=port, timeout=2.0) as client: + async for _ in client(action="health_check"): break return "reme" except Exception: