diff --git a/README.md b/README.md index 61fe990..f6fcc9d 100644 --- a/README.md +++ b/README.md @@ -118,6 +118,7 @@ async def add(request: AddGetRequest) -> AddGetResponse: | `subject` | `str` | Zenoh 主题（必需），如 `"example.add"` | | `method` | `str` | HTTP 方法：`"GET"` / `"POST"` / `"MQTT"` | | `response_model` | `type[BaseModel]` | Pydantic 响应模型类 | +| `stream` | `bool` | 流式端点：函数返回生成器，框架逐块多回复推送（默认 `false`） | **装饰器行为**： @@ -127,6 +128,31 @@ async def add(request: AddGetRequest) -> AddGetResponse: | `POST` | JSON body | 自动绑定到 Pydantic 模型 | | `MQTT` | MQTT payload JSON | 同 POST | +### 流式端点（SSE 多回复） + +`@rpc_method(..., stream=True)` 的函数返回**生成器**（async / sync 均可），框架自动将每个 +`yield` 的 dict 作为独立 Zenoh 回复块推送（Query/Reply 原生多回复），网关以 SSE +（`text/event-stream`）逐块转发给 HTTP 客户端。无需端口、进程或新依赖： + +```python +@rpc_method(subject="example.stream_chat", method="POST", stream=True) +async def stream_chat(prompt: str = "你好", tokens: int = 10): + for i in range(tokens): + await asyncio.sleep(0.05) # 模拟逐 token 生成 + yield {"token": f"{prompt}·{i + 1}", "index": i} + yield {"done": True} +``` + +**协议约定**（框架自动处理，插件无需关心底层实现）： + +- 生成器正常结束 = 流结束。**没有协议层结束哨兵**；`{"done": True}` 这类字段只是业务标记 +- 生成器中途异常 → 框架发送 `{"error": ...}` 收尾块后停流 +- 客户端断开 / 查询被网关放弃 → 逐块 reply 失败时框架立即停流，避免 Worker 空转 +- 每个块必须是可 JSON 序列化的对象，单块建议小于 64KB（zenoh 消息上限内） + +**网关侧配合**：对应 TOML 方法节声明 `stream = true` 并配大 `timeout`（流总寿命上限）， +详见网关仓库 README「流式端点（SSE）」。非流式端点（不声明 `stream`）行为完全不变。 + ### 请求信封协议 Worker 运行时收到的请求负载统一为信封结构（网关 HTTP 路由 / MQTT Bridge / `chongming-request` Client 自动组装）： diff --git a/src/chongming_worker/decorators.py b/src/chongming_worker/decorators.py index c096fa0..612ecc5 100644 --- a/src/chongming_worker/decorators.py +++ b/src/chongming_worker/decorators.py @@ -108,12 +108,15 @@ class RPCMethod: timeout: float = 3.0, request_model: Optional[Type[BaseModel]] = None, response_model: Optional[Type[BaseModel]] = None, + stream: bool = False, ): self.func = func self.subject = subject self.method = method.upper() self.summary = summary or func.__doc__ or "" self.timeout = timeout + # 流式端点：函数返回生成器，运行时逐块多回复推送（网关以 SSE 输出） + self.stream = bool(stream) self.name = func.__name__ # 绑定的 Pydantic 模型（由 gen_models 生成，或手动传入） @@ -302,6 +305,7 @@ def rpc_method( timeout: float = 3.0, request_model: Optional[Type[BaseModel]] = None, response_model: Optional[Type[BaseModel]] = None, + stream: bool = False, ): """ 装饰器：将函数注册为 RPC 端点 @@ -322,6 +326,20 @@ def rpc_method( def add(a: float, b: float): return a + b + # 流式端点（stream=True）：函数返回生成器，框架逐块多回复推送 + # 网关在 TOML 路由配置中同时声明 stream = true 时以 SSE 输出， + # 客户端逐块收到 data: {"token": ...} 事件（详见 runtime.stream_replies_to_query） + @rpc_method( + subject="chat.stream", + method="POST", + summary="Streaming chat", + stream=True, + ) + async def stream_chat(prompt: str): + async for token in llm.stream(prompt): + yield {"token": token} + yield {"done": True} + 参数: subject: 接口名称，格式为 "worker_type.operation" method: HTTP 方法 (GET/POST/PUT/DELETE) @@ -329,6 +347,7 @@ def rpc_method( timeout: 超时时间（秒） request_model: 可选的 Pydantic 请求模型（用于验证请求参数） response_model: 可选的 Pydantic 响应模型（用于格式化响应） + stream: 是否为流式端点（函数返回生成器，逐块回复；网关侧配 stream = true） """ def decorator(func: Callable) -> Callable: rpc = RPCMethod( @@ -339,6 +358,7 @@ def rpc_method( timeout=timeout, request_model=request_model, response_model=response_model, + stream=stream, ) # 提取 worker_type（subject 的第一部分） worker_type = subject.split(".")[0] if "." in subject else "default" @@ -535,7 +555,7 @@ def generate_route_configs(worker_type: str) -> List[Dict[str, Any]]: for _key, rpc in methods.items(): # 从 subject 提取操作名（最后一部分） op_name = rpc.subject.split(".")[-1] if "." in rpc.subject else rpc.subject - configs.append({ + config = { "subject": rpc.subject, "path": f"/{op_name}", "method": rpc.method, @@ -543,7 +563,11 @@ def generate_route_configs(worker_type: str) -> List[Dict[str, Any]]: "response_model": _build_response_spec(rpc), "summary": rpc.summary, "timeout": rpc.timeout, - }) + } + # 流式端点标记（网关据此生成 SSE 路由；老网关忽略此字段） + if rpc.stream: + config["stream"] = True + configs.append(config) return configs diff --git a/src/chongming_worker/runtime.py b/src/chongming_worker/runtime.py index 525d536..0bf3893 100644 --- a/src/chongming_worker/runtime.py +++ b/src/chongming_worker/runtime.py @@ -5,6 +5,7 @@ Worker 运行时 — 通用的生命周期管理 """ import os import asyncio +import inspect import json import logging import orjson @@ -22,6 +23,48 @@ logger = logging.getLogger("chongming_worker.runtime") # 是否已提示过"非信封载荷"警告（每个进程仅提示一次，避免刷屏） _warned_legacy_payload = False + +async def stream_replies_to_query(query: zenoh.Query, result) -> None: + """ + 将生成器结果逐块回复给查询方（Zenoh 多回复流式）。 + + 支持异步生成器与同步生成器，每个 yield 的 dict 作为独立回复块发送。 + 约定（与网关 call_stream / SSE 端点配合）: + - 生成器正常结束 = 流结束，无结束哨兵 + - 生成器中途异常 -> 发送 {"error": ...} 收尾块后停流 + - query.reply 失败（查询已终止，如客户端断开/网关放弃）-> 立即停流， + 避免 Worker 对已无接收方的查询继续生成 + + Args: + query: Zenoh Query 对象（仅使用 key_expr 与 reply） + result: 异步或同步生成器，逐块 yield 可 JSON 序列化的对象 + """ + async def _reply(chunk) -> bool: + """单块回复；返回 False 表示查询已终止，应停流""" + try: + query.reply( + key_expr=query.key_expr, + payload=orjson.dumps(chunk), + ) + return True + except Exception as e: + logger.debug("[Worker] 流式回复失败（查询可能已终止）: %s", e) + return False + + try: + if inspect.isasyncgen(result): + async for chunk in result: + if not await _reply(chunk): + break + else: + for chunk in result: + if not await _reply(chunk): + break + except Exception as e: + # 生成器中途异常：以 {"error": ...} 收尾块结束流 + logger.error("[Worker] 流式生成器异常，发送错误块后停止: %s", e) + await _reply({"error": str(e)}) + async def create_worker( session: zenoh.Session, worker_type: str, @@ -72,6 +115,7 @@ async def create_worker( """创建处理特定 subject 的回调函数""" loop = asyncio.get_running_loop() async def _handle_query(query: zenoh.Query): + matched = None # 方法匹配结果（异常分支判断流式错误格式用） try: # 解析请求负载（信封协议: {"meta": {...}, "data": {...}}） payload_bytes = bytes(query.payload) if query.payload else b"" @@ -104,7 +148,7 @@ async def create_worker( ) return - # 调用 RPC 方法（支持同步和异步） + # 调用 RPC 方法（支持同步和异步；返回生成器时进入流式模式） # user_info 从 meta 信封注入（框架参数，不进业务 payload） result = await matched.call( data, @@ -113,23 +157,34 @@ async def create_worker( user_info=meta.get("user_info"), ) - # 回复 - query.reply( - key_expr=query.key_expr, - payload=orjson.dumps(result) - ) + if inspect.isasyncgen(result) or inspect.isgenerator(result): + # 流式模式：@rpc_method 函数返回生成器，逐块多回复推送 + await stream_replies_to_query(query, result) + else: + # 一次性回复（原有逻辑） + query.reply( + key_expr=query.key_expr, + payload=orjson.dumps(result) + ) logger.debug("[Worker] %s 回复 %s %s", worker_id, matched.method, subject) except Exception as e: - error_response = { - "result": 0.0, - "operation": "error", - "error": str(e), - "worker": worker_id - } - query.reply( - key_expr=query.key_expr, - payload=json.dumps(error_response).encode("utf-8") - ) + if matched is not None and matched.stream: + # 流式端点错误：以 {"error": ...} 收尾块结束流 + query.reply( + key_expr=query.key_expr, + payload=orjson.dumps({"error": str(e)}) + ) + else: + error_response = { + "result": 0.0, + "operation": "error", + "error": str(e), + "worker": worker_id + } + query.reply( + key_expr=query.key_expr, + payload=json.dumps(error_response).encode("utf-8") + ) logger.error("[Worker] %s 异常 %s: %s", worker_id, subject, e) def callback(query: zenoh.Query):