modify: 流式输出

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
 
 

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):