Kris
1a7690aaa0
新增了完整的流式响应处理模块，包含代码围栏解析、分块策略、指令处理、流式草稿循环等核心能力，支持按段落/句子等规则切割响应块，支持进度更新与指令回调，完善流式交互的全流程支持
421 lines
15 KiB
Python
421 lines
15 KiB
Python
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import logging
|
||
import time
|
||
from collections.abc import Awaitable, Callable
|
||
|
||
from yuxi.channel.runtime.backoff import BackoffConfig, ErrorBackoff
|
||
from yuxi.channel.config.defaults import TIMEOUT
|
||
from yuxi.channel.streaming.directives import (
|
||
DirectiveHandler,
|
||
StreamingDirective,
|
||
StreamingDirectiveParser,
|
||
)
|
||
from yuxi.channel.streaming.models import DraftStreamSession, DraftStreamState
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
SendOrEditFn = Callable[[str, str | None], Awaitable[str | None]]
|
||
ProgressCallback = Callable[[str], Awaitable[None]]
|
||
|
||
_MAX_FLUSH_ITERATIONS = 100
|
||
_TOOL_STATUS_PREFIX = "__tool__:"
|
||
|
||
_TOOL_STATUS_LABELS: dict[str, str] = {
|
||
"thinking": "🤔 思考中...",
|
||
"searching": "🔍 搜索中...",
|
||
"reading": "📖 阅读中...",
|
||
"generating": "✍️ 生成中...",
|
||
"executing": "⚙️ 执行中...",
|
||
"summarizing": "📝 总结中...",
|
||
}
|
||
|
||
|
||
def get_tool_status_label(tool_status: str) -> str:
|
||
return _TOOL_STATUS_LABELS.get(tool_status, f"... {tool_status} ...")
|
||
|
||
|
||
class DraftStreamLoop:
|
||
"""流式草稿消息循环引擎。
|
||
|
||
``update()`` 方法要求调用方每次传入**完整的累积文本**
|
||
（而非增量 chunk），因为内部 ``StreamingDirectiveParser``
|
||
需要对全文做指令扫描。传入增量文本会导致未发送内容静默丢失。
|
||
"""
|
||
|
||
def __init__(self, send_or_edit: SendOrEditFn, throttle_ms: int = 1000, max_retries: int = 3):
|
||
self._send_or_edit = send_or_edit
|
||
self._throttle_ms = throttle_ms
|
||
self._active_session: DraftStreamSession | None = None
|
||
self._update_queue: asyncio.Queue[str | None] | None = None
|
||
self._loop_task: asyncio.Task[None] | None = None
|
||
self._lock = asyncio.Lock()
|
||
|
||
self._last_sent_at: float = 0.0
|
||
self._pending_text: str = ""
|
||
self._in_flight: asyncio.Task[str | None] | None = None
|
||
self._timer: asyncio.Task[None] | None = None
|
||
self._flush_task: asyncio.Task[None] | None = None
|
||
|
||
self._directive_parser = StreamingDirectiveParser()
|
||
self._on_directive: DirectiveHandler | None = None
|
||
|
||
self._on_progress: ProgressCallback | None = None
|
||
|
||
self._send_backoff = ErrorBackoff(
|
||
config=BackoffConfig(
|
||
base_delay=0.5,
|
||
max_delay=5.0,
|
||
exponent=2.0,
|
||
jitter=True,
|
||
jitter_factor=0.1,
|
||
max_retries=max_retries,
|
||
),
|
||
)
|
||
|
||
def set_directive_handler(self, handler: DirectiveHandler) -> None:
|
||
self._on_directive = handler
|
||
|
||
def set_progress_callback(self, callback: ProgressCallback) -> None:
|
||
self._on_progress = callback
|
||
|
||
async def start(self, session: DraftStreamSession) -> None:
|
||
async with self._lock:
|
||
if self._active_session is not None:
|
||
await self._abort_current()
|
||
self._active_session = session
|
||
session.state = DraftStreamState.STREAMING
|
||
self._update_queue = asyncio.Queue()
|
||
self._directive_parser.reset()
|
||
self._loop_task = asyncio.create_task(self._run_loop(session))
|
||
|
||
async def update(self, full_text: str, tool_status: str | None = None) -> None:
|
||
"""将**完整累积文本**送入流式处理循环。
|
||
|
||
调用方必须传入当前已累积的全部文本内容，而非增量 chunk。
|
||
内部 ``StreamingDirectiveParser`` 会对全文做指令解析与去重。
|
||
"""
|
||
if self._update_queue is None:
|
||
return
|
||
await self._update_queue.put(full_text)
|
||
if tool_status is not None:
|
||
await self._update_queue.put(f"{_TOOL_STATUS_PREFIX}{tool_status}")
|
||
|
||
async def finalize(self) -> None:
|
||
if self._update_queue is not None:
|
||
await self._update_queue.put(None)
|
||
if self._loop_task is not None:
|
||
try:
|
||
await asyncio.wait_for(self._loop_task, timeout=TIMEOUT.stream_finalize)
|
||
except TimeoutError:
|
||
logger.warning("DraftStreamLoop finalize timed out after 30s, forcing cancel")
|
||
self._loop_task.cancel()
|
||
try:
|
||
await self._loop_task
|
||
except asyncio.CancelledError:
|
||
pass
|
||
finally:
|
||
self._loop_task = None
|
||
if self._active_session is not None and self._active_session.state == DraftStreamState.FINISHED:
|
||
self._active_session = None
|
||
|
||
async def abort(self) -> None:
|
||
async with self._lock:
|
||
await self._abort_current()
|
||
|
||
async def flush(self) -> None:
|
||
self._cancel_timer()
|
||
iterations = 0
|
||
while not self._is_stopped and iterations < _MAX_FLUSH_ITERATIONS:
|
||
iterations += 1
|
||
if self._in_flight and not self._in_flight.done():
|
||
await self._in_flight
|
||
continue
|
||
|
||
text = self._pending_text
|
||
if not text.strip():
|
||
self._pending_text = ""
|
||
return
|
||
|
||
self._pending_text = ""
|
||
msg_id = self._active_session.message_id if self._active_session else None
|
||
key = f"flush:{id(self)}"
|
||
|
||
async def _do_flush() -> str | None:
|
||
return await self._send_or_edit(text, msg_id)
|
||
|
||
async def _retry_flush() -> str | None:
|
||
return await self._send_backoff.execute(key, _do_flush)
|
||
|
||
current = asyncio.create_task(_retry_flush())
|
||
|
||
def _clear_in_flight(t: asyncio.Task[str | None]) -> None:
|
||
if self._in_flight is t:
|
||
self._in_flight = None
|
||
|
||
current.add_done_callback(_clear_in_flight)
|
||
self._in_flight = current
|
||
|
||
try:
|
||
sent = await current
|
||
if sent is None:
|
||
self._pending_text = text
|
||
return
|
||
self._last_sent_at = time.monotonic()
|
||
if not self._pending_text:
|
||
return
|
||
except Exception:
|
||
logger.exception(
|
||
"DraftStreamLoop flush failed after %d retries",
|
||
self._send_backoff.config.max_retries,
|
||
)
|
||
self._pending_text = text
|
||
return
|
||
|
||
if iterations >= _MAX_FLUSH_ITERATIONS:
|
||
logger.warning("DraftStreamLoop.flush() exceeded max iterations (%s), breaking", _MAX_FLUSH_ITERATIONS)
|
||
|
||
async def update_text(self, text: str) -> None:
|
||
if self._is_stopped:
|
||
return
|
||
async with self._lock:
|
||
self._pending_text = text
|
||
if self._in_flight and not self._in_flight.done():
|
||
self._schedule()
|
||
return
|
||
if self._timer is None and (time.monotonic() - self._last_sent_at) >= (self._throttle_ms / 1000):
|
||
try:
|
||
self._flush_task = asyncio.create_task(self.flush())
|
||
except RuntimeError:
|
||
return
|
||
self._flush_task.add_done_callback(self._on_flush_done)
|
||
return
|
||
self._schedule()
|
||
|
||
def stop(self) -> None:
|
||
self._pending_text = ""
|
||
self._cancel_timer()
|
||
|
||
def reset_pending(self) -> None:
|
||
self._pending_text = ""
|
||
|
||
def reset_throttle_window(self) -> None:
|
||
self._last_sent_at = 0.0
|
||
self._cancel_timer()
|
||
|
||
async def wait_for_in_flight(self) -> None:
|
||
if self._in_flight and not self._in_flight.done():
|
||
await self._in_flight
|
||
|
||
@property
|
||
def _is_stopped(self) -> bool:
|
||
if self._active_session is None:
|
||
return True
|
||
return self._active_session.state not in (DraftStreamState.STREAMING, DraftStreamState.FINALIZING)
|
||
|
||
async def _abort_current(self) -> None:
|
||
if self._loop_task and not self._loop_task.done():
|
||
self._loop_task.cancel()
|
||
try:
|
||
await self._loop_task
|
||
except asyncio.CancelledError:
|
||
pass
|
||
self._loop_task = None
|
||
if self._active_session:
|
||
self._active_session.state = DraftStreamState.ABORTED
|
||
self._active_session = None
|
||
self._update_queue = None
|
||
self._cancel_timer()
|
||
self._pending_text = ""
|
||
self._directive_parser.reset()
|
||
|
||
async def _run_loop(self, session: DraftStreamSession) -> None:
|
||
last_update = time.monotonic()
|
||
pending_text = ""
|
||
|
||
try:
|
||
while session.state == DraftStreamState.STREAMING:
|
||
try:
|
||
item = await asyncio.wait_for(
|
||
self._update_queue.get(),
|
||
timeout=TIMEOUT.stream_queue_poll,
|
||
)
|
||
except TimeoutError:
|
||
now = time.monotonic()
|
||
if pending_text and (now - last_update) >= (self._throttle_ms / 1000):
|
||
last_update = now
|
||
await self._send_update(session, pending_text)
|
||
pending_text = ""
|
||
session.last_update_ts = time.monotonic()
|
||
continue
|
||
|
||
if item is None:
|
||
if pending_text:
|
||
await self._send_update(session, pending_text)
|
||
session.state = DraftStreamState.FINISHED
|
||
return
|
||
|
||
if item.startswith(_TOOL_STATUS_PREFIX):
|
||
status = item[len(_TOOL_STATUS_PREFIX):]
|
||
session.tool_status = status
|
||
if self._on_progress is not None:
|
||
try:
|
||
await self._on_progress(status)
|
||
except Exception:
|
||
logger.exception("Progress callback failed for status %s", status)
|
||
if pending_text:
|
||
await self._send_update(session, pending_text)
|
||
pending_text = ""
|
||
continue
|
||
|
||
clean_text, directives = self._directive_parser.feed(item)
|
||
await self._dispatch_directives(directives)
|
||
pending_text = clean_text
|
||
session.accumulated_text = clean_text
|
||
|
||
now = time.monotonic()
|
||
elapsed_ms = int((now - last_update) * 1000)
|
||
delta = len(pending_text) - len(session.last_sent_text)
|
||
|
||
is_natural_boundary = bool(pending_text and pending_text[-1] in "\n。！？!?；;：:")
|
||
|
||
should_send = (
|
||
delta >= session.significant_delta_chars
|
||
or is_natural_boundary
|
||
or (delta > 0 and elapsed_ms >= session.throttle_ms)
|
||
)
|
||
|
||
if should_send and pending_text:
|
||
last_update = now
|
||
await self._send_update(session, pending_text)
|
||
pending_text = ""
|
||
session.last_update_ts = time.monotonic()
|
||
except asyncio.CancelledError:
|
||
session.state = DraftStreamState.ABORTED
|
||
raise
|
||
|
||
async def _send_update(self, session: DraftStreamSession, text: str) -> None:
|
||
key = f"send_update:{id(self)}"
|
||
|
||
async def _do_send() -> None:
|
||
new_id = await self._send_or_edit(text, session.message_id)
|
||
if new_id:
|
||
session.message_id = new_id
|
||
session.last_sent_text = text
|
||
|
||
try:
|
||
await self._send_backoff.execute(key, _do_send)
|
||
except Exception:
|
||
logger.exception(
|
||
"DraftStreamLoop send_or_edit failed after %d retries",
|
||
self._send_backoff.config.max_retries,
|
||
)
|
||
|
||
async def _dispatch_directives(self, directives: list[StreamingDirective]) -> None:
|
||
if not directives or self._on_directive is None:
|
||
return
|
||
for directive in directives:
|
||
try:
|
||
await self._on_directive(directive)
|
||
except Exception:
|
||
logger.exception("Directive handler failed for %s", directive.type)
|
||
|
||
def _schedule(self) -> None:
|
||
if self._timer is not None:
|
||
return
|
||
delay = max(0, (self._throttle_ms / 1000) - (time.monotonic() - self._last_sent_at))
|
||
|
||
async def _delayed_flush():
|
||
await asyncio.sleep(delay)
|
||
await self.flush()
|
||
|
||
try:
|
||
self._timer = asyncio.create_task(_delayed_flush())
|
||
except RuntimeError:
|
||
return
|
||
self._timer.add_done_callback(self._on_timer_done)
|
||
|
||
def _cancel_timer(self) -> None:
|
||
if self._timer is not None and not self._timer.done():
|
||
self._timer.cancel()
|
||
self._timer = None
|
||
|
||
def _on_flush_done(self, task: asyncio.Task) -> None:
|
||
if self._flush_task is task:
|
||
self._flush_task = None
|
||
try:
|
||
task.result()
|
||
except Exception:
|
||
logger.exception("DraftStreamLoop background flush failed")
|
||
|
||
def _on_timer_done(self, task: asyncio.Task) -> None:
|
||
if self._timer is task:
|
||
self._timer = None
|
||
try:
|
||
task.result()
|
||
except Exception:
|
||
logger.exception("DraftStreamLoop timer flush failed")
|
||
|
||
|
||
class FinalizableDraftStreamControls:
|
||
def __init__(
|
||
self,
|
||
stream_loop: DraftStreamLoop,
|
||
session: DraftStreamSession,
|
||
target_id: str,
|
||
account: dict | None = None,
|
||
):
|
||
self._stream_loop = stream_loop
|
||
self._session = session
|
||
self._target_id = target_id
|
||
self._account = account or {}
|
||
self._started = False
|
||
|
||
async def start(self) -> None:
|
||
if self._started:
|
||
return
|
||
await self._stream_loop.start(self._session)
|
||
self._started = True
|
||
|
||
async def update(self, full_text: str, tool_status: str | None = None) -> None:
|
||
if not self._started:
|
||
await self.start()
|
||
await self._stream_loop.update(full_text, tool_status)
|
||
|
||
async def finalize(self) -> None:
|
||
if self._started:
|
||
await self._stream_loop.finalize()
|
||
self._started = False
|
||
|
||
async def abort(self) -> None:
|
||
if self._started:
|
||
await self._stream_loop.abort()
|
||
self._started = False
|
||
|
||
def set_progress_callback(self, callback: ProgressCallback) -> None:
|
||
self._stream_loop.set_progress_callback(callback)
|
||
|
||
@property
|
||
def session(self) -> DraftStreamSession:
|
||
return self._session
|
||
|
||
async def update_text(self, text: str) -> None:
|
||
if self._session.stopped or self._session.final:
|
||
return
|
||
await self.start()
|
||
await self._stream_loop.update_text(text)
|
||
|
||
async def stop(self) -> None:
|
||
self._session.final = True
|
||
await self._stream_loop.flush()
|
||
|
||
async def seal(self) -> None:
|
||
self._session.final = True
|
||
self._stream_loop.stop()
|
||
await self._stream_loop.wait_for_in_flight()
|
||
|
||
def discard_pending(self) -> None:
|
||
self._session.stopped = True
|
||
self._stream_loop.stop()
|