"""
流式端点示例 — @rpc_method 返回异步生成器，框架自动逐块多回复推送

网关在 TOML 中声明 stream = true 时，该端点以 SSE（text/event-stream）输出:
  POST /api/v1/calc/stream_chat  {"prompt": "你好", "tokens": 5}
  响应: data: {"token": "..."} ... data: {"done": true}

约定:
  - 插件侧写法：函数返回生成器（yield），框架逐块 query.reply() 推送
  - {"done": true} 是业务结束标记（插件自定义），不是协议哨兵——
    流式协议以生成器自然结束为流终止信号
  - 中途异常时框架自动发送 {"error": ...} 收尾块
"""
import asyncio
import logging

from chongming_worker import rpc_method

logger = logging.getLogger("plugins.example.stream_chat")


@rpc_method(
    subject="example.stream_chat",
    method="POST",
    summary="模拟 AI 流式对话（SSE，逐 token 输出）",
    stream=True,
)
async def stream_chat(prompt: str = "你好", tokens: int = 10):
    """逐 token 模拟生成：把 return 改成 yield，框架即自动流式推送"""
    logger.info("收到流式请求: prompt=%s, tokens=%d", prompt, tokens)
    for i in range(tokens):
        await asyncio.sleep(0.05)
        yield {"token": f"{prompt}·{i + 1}", "index": i}
    yield {"done": True}