diff --git a/reme4/components/client/base_client.py b/reme4/components/client/base_client.py index 4ca93298..a5b05b6c 100644 --- a/reme4/components/client/base_client.py +++ b/reme4/components/client/base_client.py @@ -1,6 +1,8 @@ """Base client abstraction.""" +import json from abc import abstractmethod +from collections.abc import AsyncGenerator from ..base_component import BaseComponent from ...enumeration import ComponentEnum @@ -22,5 +24,18 @@ class BaseClient(BaseComponent): """Close the client and release resources.""" @abstractmethod - async def __call__(self) -> str: - """Execute the configured action""" + def _execute(self) -> 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]: + """Dispatch: action='list' returns the action catalog; otherwise delegate to _execute().""" + if getattr(self, "action", None) == "list": + actions = await self.list_actions() + yield json.dumps(actions, indent=2, ensure_ascii=False) + return + async for chunk in self._execute(): + yield chunk diff --git a/reme4/components/client/http_client.py b/reme4/components/client/http_client.py index 96958c2c..1ccb7620 100644 --- a/reme4/components/client/http_client.py +++ b/reme4/components/client/http_client.py @@ -2,17 +2,20 @@ import json import os +from collections.abc import AsyncGenerator import httpx from .base_client import BaseClient from ..component_registry import R from ...constants import REME_SERVICE_INFO, REME_DEFAULT_HOST, REME_DEFAULT_PORT +from ...enumeration import ChunkEnum +from ...schema import StreamChunk @R.register("http") class HttpClient(BaseClient): - """HTTP client that communicates with ReMe service via REST API.""" + """HTTP client that auto-adapts to JSON or SSE endpoints via Content-Type.""" def __init__( self, @@ -46,15 +49,68 @@ class HttpClient(BaseClient): if self.client is None: self.client = httpx.AsyncClient(base_url=self.base_url, timeout=self.timeout) - async def __call__(self) -> str: - """Send POST request to the configured action endpoint.""" + async def _iter_stream_chunks(self) -> AsyncGenerator[StreamChunk, None]: + """Send request and yield StreamChunks; auto-detects JSON vs SSE via Content-Type.""" if self.client is None: raise RuntimeError("Client not initialized. Call _start() first.") - response = await self.client.post(f"/{self.action}", json=self.kwargs) - response.raise_for_status() - result = response.json() - return json.dumps(result, indent=2, ensure_ascii=False) + async with self.client.stream("POST", f"/{self.action}", json=self.kwargs) as resp: + resp.raise_for_status() + ctype = resp.headers.get("content-type", "") + + if ctype.startswith("text/event-stream"): + async for line in resp.aiter_lines(): + if not line.startswith("data:"): + continue + payload = line[len("data:") :] + if payload.strip() == "[DONE]": + return + try: + data = json.loads(payload) + except json.JSONDecodeError: + continue + chunk = StreamChunk(**data) + if chunk.chunk_type == ChunkEnum.ERROR: + # Surface server-side errors as exceptions so callers don't + # mistake error chunks for valid content. + raise RuntimeError(str(chunk.chunk)) + if chunk.done: + return + yield chunk + else: + body = await resp.aread() + text = body.decode() + try: + data = json.loads(text) + pretty = json.dumps(data, indent=2, ensure_ascii=False) + except json.JSONDecodeError: + pretty = text + yield StreamChunk(chunk_type=ChunkEnum.CONTENT, chunk=pretty) + + async def stream_chunks(self) -> AsyncGenerator[StreamChunk, None]: + """HTTP-specific richer access: yield full StreamChunk objects with chunk_type/metadata.""" + async for chunk in self._iter_stream_chunks(): + yield chunk + + async def list_actions(self) -> list[dict]: + """Return raw OpenAPI operations; each dict gets an `action` key (path without leading '/').""" + if self.client is None: + raise RuntimeError("Client not initialized. Call _start() first.") + resp = await self.client.get("/openapi.json") + resp.raise_for_status() + spec = resp.json() + actions: list[dict] = [] + for path, methods in spec.get("paths", {}).items(): + for method, op in methods.items(): + actions.append({"action": path.lstrip("/"), "method": method.upper(), **op}) + return actions + + # pylint: disable=invalid-overridden-method + async def _execute(self) -> AsyncGenerator[str, None]: + """Yield text chunks; one yield for JSON endpoints, many for SSE.""" + async for chunk in self._iter_stream_chunks(): + payload = chunk.chunk + yield payload if isinstance(payload, str) else json.dumps(payload, ensure_ascii=False) async def _close(self) -> None: """Close the HTTP client.""" diff --git a/reme4/components/client/mcp_client.py b/reme4/components/client/mcp_client.py index ce74d690..b8b17a9b 100644 --- a/reme4/components/client/mcp_client.py +++ b/reme4/components/client/mcp_client.py @@ -2,6 +2,7 @@ import json import os +from collections.abc import AsyncGenerator from typing import Any from fastmcp import Client @@ -27,7 +28,8 @@ class MCPClient(BaseClient): # SSE (default) client = MCPClient(action="my_tool", host="localhost", port=8000, query="hello") async with client: - result = await client() + async for text in client(): + print(text) # Streamable HTTP client = MCPClient(action="my_tool", transport="streamable-http", host="localhost", port=8000) @@ -94,12 +96,20 @@ class MCPClient(BaseClient): self.client = Client(self._build_transport(), timeout=self.timeout) await self.client.__aenter__() - async def __call__(self) -> str: + # pylint: disable=invalid-overridden-method + async def _execute(self) -> 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) - return self._extract_text(result) + yield self._extract_text(result) + + async def list_actions(self) -> list[dict]: + """Return raw MCP Tool dumps; each dict gets an `action` key (the tool name).""" + if self.client is None: + raise RuntimeError("Client not initialized. Call _start() first.") + tools = await self.client.list_tools() + return [tool.model_dump() for tool in tools] # pylint: disable=unnecessary-dunder-call async def _close(self) -> None: diff --git a/reme4/config/default.yaml b/reme4/config/default.yaml index b054b620..f06d6c5a 100644 --- a/reme4/config/default.yaml +++ b/reme4/config/default.yaml @@ -40,6 +40,29 @@ jobs: steps: - backend: health_check_step + - backend: stream + name: stream_demo + description: "stream demo job: repeat query 10x and stream char-by-char" + parameters: + type: object + properties: + query: + type: string + description: "query to echo" + repeat: + type: integer + description: "number of times to repeat the query" + default: 10 + interval: + type: number + description: "seconds between chunks" + default: 0.1 + required: + - query + steps: + - backend: stream_demo_step1 + - backend: stream_demo_step2 + components: # 1. tokenizer — no dependencies tokenizer: diff --git a/reme4/reme.py b/reme4/reme.py index 98562d18..fe74bb7a 100644 --- a/reme4/reme.py +++ b/reme4/reme.py @@ -19,7 +19,9 @@ async def call_server(action: str, **kwargs): backend: str = kwargs.pop("backend", "http") client_cls = R.get(ComponentEnum.CLIENT, backend) async with client_cls(action=action, **kwargs) as client: - print(await client()) + async for chunk in client(): + print(chunk, end="", flush=True) + print() def main(): diff --git a/reme4/steps/base_step.py b/reme4/steps/base_step.py index 7cd45b55..a710f457 100644 --- a/reme4/steps/base_step.py +++ b/reme4/steps/base_step.py @@ -11,6 +11,7 @@ from agentscope.token import TokenCounterBase from ..components.embedding import BaseEmbeddingModel from ..components.file_parser import BaseFileParser from ..components.file_store import BaseFileStore +from ..components.file_watcher import BaseFileWatcher from ..components.prompt_handler import PromptHandler from ..components.runtime_context import RuntimeContext from ..enumeration import ComponentEnum @@ -121,6 +122,11 @@ class BaseStep(ABC): """Return the embedding model component.""" return self._resolve("embedding", BaseEmbeddingModel, ComponentEnum.EMBEDDING_MODEL) + @property + def file_watcher(self) -> BaseFileWatcher: + """Return the file watcher component.""" + return self._resolve("file_watcher", BaseFileWatcher, ComponentEnum.FILE_WATCHER) + def prompt_format(self, prompt_name: str, **kwargs) -> str: """Format a named prompt template with the given kwargs.""" return self.prompt.prompt_format(prompt_name=prompt_name, **kwargs) diff --git a/reme4/steps/common/__init__.py b/reme4/steps/common/__init__.py index dc57b38c..fb67df6a 100644 --- a/reme4/steps/common/__init__.py +++ b/reme4/steps/common/__init__.py @@ -2,11 +2,14 @@ from .demo import DemoEchoStep1, DemoEchoStep2 from .health_check import HealthCheckStep +from .stream_demo import StreamDemoStep1, StreamDemoStep2 from .version import VersionStep __all__ = [ "DemoEchoStep1", "DemoEchoStep2", "HealthCheckStep", + "StreamDemoStep1", + "StreamDemoStep2", "VersionStep", ] diff --git a/reme4/steps/common/stream_demo.py b/reme4/steps/common/stream_demo.py new file mode 100644 index 00000000..450077aa --- /dev/null +++ b/reme4/steps/common/stream_demo.py @@ -0,0 +1,42 @@ +"""Streaming demo steps: step1 prepares text, step2 streams it char-by-char.""" + +import asyncio + +from ..base_step import BaseStep +from ...components import R +from ...enumeration import ChunkEnum + + +@R.register("stream_demo_step1") +class StreamDemoStep1(BaseStep): + """Read query from context, repeat it 10x, write back for Step2 to stream.""" + + async def execute(self): + assert self.context is not None + query = self.context.get("query", "") + repeat = int(self.context.get("repeat", 10)) + + stream_text = (query * repeat) if query else "" + + self.logger.info(f"[{self.name}] query={query!r}, repeat={repeat}, len={len(stream_text)}") + + self.context["stream_text"] = stream_text + return self.context.response + + +@R.register("stream_demo_step2") +class StreamDemoStep2(BaseStep): + """Stream stream_text char-by-char as CONTENT chunks with 0.1s pacing.""" + + async def execute(self): + assert self.context is not None + stream_text: str = self.context.get("stream_text", "") + interval = float(self.context.get("interval", 0.1)) + + self.logger.info(f"[{self.name}] streaming {len(stream_text)} chars, interval={interval}s") + + for ch in stream_text: + await self.context.add_stream_string(ch, ChunkEnum.CONTENT) + await asyncio.sleep(interval) + + return self.context.response diff --git a/reme4/utils/service_utils.py b/reme4/utils/service_utils.py index f7468893..d1a39b9d 100644 --- a/reme4/utils/service_utils.py +++ b/reme4/utils/service_utils.py @@ -5,15 +5,17 @@ import socket import subprocess import sys -from ..components.client.http_client import HttpClient from ..constants import REME_DEFAULT_HOST, REME_DEFAULT_PORT async def find_reme(host: str, port: int) -> str: """Probe host:port. Returns 'reme', 'occupied', or 'free'.""" + from ..components.client.http_client import HttpClient + try: - async with HttpClient(action="health_check", host=host, port=port, timeout=2.0) as c: - await c() + async with HttpClient(action="health_check", host=host, port=port, timeout=2.0) as client: + async for _ in client(): + break return "reme" except Exception: pass @@ -26,8 +28,8 @@ async def find_reme(host: str, port: int) -> str: return "occupied" -def _run(cmd: list[str]) -> str: - """Run cmd; return stdout, or empty on failure.""" +def _sh(cmd: list[str]) -> str: + """Run cmd; return stdout, or '' on failure.""" try: return subprocess.check_output(cmd, stderr=subprocess.DEVNULL, text=True) except (subprocess.CalledProcessError, FileNotFoundError): @@ -35,41 +37,32 @@ def _run(cmd: list[str]) -> str: def _pid_on_port(port: int) -> int | None: - """Return PID listening on TCP port, or None.""" - out = _run(["lsof", "-nP", f"-iTCP:{port}", "-sTCP:LISTEN", "-t"]).strip() + """PID listening on TCP port, or None.""" + out = _sh(["lsof", "-nP", f"-iTCP:{port}", "-sTCP:LISTEN", "-t"]).strip() return int(out.splitlines()[0]) if out else None def _scan_reme_procs() -> list[tuple[int, str, int]]: - """Find 'reme ... start' processes. Returns [(pid, host, port), ...].""" - procs = [] - for line in _run(["pgrep", "-af", "reme.* start"]).splitlines(): + """List running 'reme ... start' processes as (pid, host, port).""" + procs: list[tuple[int, str, int]] = [] + for line in _sh(["pgrep", "-af", "reme.* start"]).splitlines(): parts = line.split() - if not parts: - continue - try: - pid = int(parts[0]) - except ValueError: + if not parts or not parts[0].isdigit(): continue host, port = REME_DEFAULT_HOST, REME_DEFAULT_PORT for t in parts[1:]: if t.startswith("service.host="): host = t.split("=", 1)[1] - elif t.startswith("service.port="): - try: - port = int(t.split("=", 1)[1]) - except ValueError: - pass - procs.append((pid, host, port)) + elif t.startswith("service.port=") and t.split("=", 1)[1].isdigit(): + port = int(t.split("=", 1)[1]) + procs.append((int(parts[0]), host, port)) return procs async def locate_reme() -> tuple[str, int, int | None] | None: - """Locate a running reme. Returns (host, port, pid) or None.""" - # 1. Try default host:port + """Find a running reme: try default port, then scanned processes.""" if await find_reme(REME_DEFAULT_HOST, REME_DEFAULT_PORT) == "reme": return REME_DEFAULT_HOST, REME_DEFAULT_PORT, _pid_on_port(REME_DEFAULT_PORT) - # 2. Scan reme processes and probe each for pid, host, port in _scan_reme_procs(): if await find_reme(host, port) == "reme": return host, port, pid @@ -77,11 +70,7 @@ async def locate_reme() -> tuple[str, int, int | None] | None: def precheck_start(svc_config: dict | None) -> bool: - """Pre-flight check before `start`. Returns True if caller should proceed. - - Prints a message and returns False if reme is already running. - Exits with code 1 if the port is occupied by a non-reme process. - """ + """Pre-flight check for `start`: False if reme is up, exits 1 on port conflict.""" host = (svc_config or {}).get("host") or REME_DEFAULT_HOST port = (svc_config or {}).get("port") or REME_DEFAULT_PORT status = asyncio.run(find_reme(host, port)) @@ -90,8 +79,7 @@ def precheck_start(svc_config: dict | None) -> bool: return False if status == "occupied": print( - f"port {port} is occupied by another process. " - f"Start with a different port: reme4 start service.port=", + f"port {port} occupied. Start on another port: reme4 start service.port=", file=sys.stderr, ) sys.exit(1) @@ -99,11 +87,10 @@ def precheck_start(svc_config: dict | None) -> bool: def cli_find_reme() -> None: - """Handle the `find_reme` CLI action: print HOST/PORT/PID or a hint.""" + """Handle `reme find_reme`: print HOST/PORT/PID or a hint to start reme.""" found = asyncio.run(locate_reme()) - if found: - host, port, pid = found - print(f"HOST={host} PORT={port} PID={pid if pid is not None else 'unknown'}") - else: + if not found: print("reme not started. Try: reme start", file=sys.stderr) sys.exit(1) + host, port, pid = found + print(f"HOST={host} PORT={port} PID={pid or 'unknown'}")