270 lines
8.6 KiB
Python
270 lines
8.6 KiB
Python
|
import asyncio
|
||
|
|
import logging
|
||
|
|
import random
|
||
|
|
import time
|
||
|
|
from collections.abc import Callable, Coroutine
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
from yuxi.channel.message.models import ReplyPayload, ReplyStage, StreamingChunk, UnifiedMessage
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
_DEFAULT_MIN_CHARS = 80
|
||
|
|
_DEFAULT_DEBOUNCE_MS = 150.0
|
||
|
|
_DEFAULT_TIMEOUT_MS = 10_000.0
|
||
|
|
|
||
|
|
OnIdleCallback = Callable[[], None]
|
||
|
|
OnErrorCallback = Callable[[Exception], None]
|
||
|
|
BeforeDeliverCallback = Callable[[str], Coroutine[Any, Any, str | None]]
|
||
|
|
|
||
|
|
|
||
|
|
class ReplyDispatcher:
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
send_fn: Callable[..., Coroutine[Any, Any, str | None]],
|
||
|
|
msg: UnifiedMessage,
|
||
|
|
*,
|
||
|
|
response_prefix: str = "",
|
||
|
|
human_delay: tuple[float, float] | None = None,
|
||
|
|
timeout_ms: float = _DEFAULT_TIMEOUT_MS,
|
||
|
|
on_idle: OnIdleCallback | None = None,
|
||
|
|
on_error: OnErrorCallback | None = None,
|
||
|
|
before_deliver: BeforeDeliverCallback | None = None,
|
||
|
|
):
|
||
|
|
self._send_fn = send_fn
|
||
|
|
self._msg = msg
|
||
|
|
self._response_prefix = response_prefix
|
||
|
|
self._human_delay = human_delay
|
||
|
|
self._timeout_ms = timeout_ms
|
||
|
|
self._on_idle = on_idle
|
||
|
|
self._on_error = on_error
|
||
|
|
self._before_deliver = before_deliver
|
||
|
|
|
||
|
|
self._send_chain: asyncio.Lock = asyncio.Lock()
|
||
|
|
self._flush_lock: asyncio.Lock = asyncio.Lock()
|
||
|
|
self._pending = 0
|
||
|
|
self._complete = asyncio.Event()
|
||
|
|
self._done = False
|
||
|
|
|
||
|
|
self._blocks: list[str] = []
|
||
|
|
self._block_deadline: float = 0.0
|
||
|
|
self._block_timer: asyncio.Task | None = None
|
||
|
|
self._sent_payload_keys: set[str] = set()
|
||
|
|
self._sent_content_keys: set[str] = set()
|
||
|
|
|
||
|
|
self._tool_count = 0
|
||
|
|
self._block_count = 0
|
||
|
|
self._final_count = 0
|
||
|
|
self._failed_count = 0
|
||
|
|
self._cancelled_count = 0
|
||
|
|
self._started_at: float | None = None
|
||
|
|
|
||
|
|
self._abort_event = asyncio.Event()
|
||
|
|
|
||
|
|
async def enqueue(self, chunk: StreamingChunk) -> bool:
|
||
|
|
if self._abort_event.is_set():
|
||
|
|
self._cancelled_count += 1
|
||
|
|
return False
|
||
|
|
|
||
|
|
if chunk.stage == ReplyStage.TOOL:
|
||
|
|
return await self._dispatch_tool(chunk)
|
||
|
|
elif chunk.stage == ReplyStage.BLOCK:
|
||
|
|
return await self._dispatch_block(chunk)
|
||
|
|
elif chunk.stage == ReplyStage.FINAL:
|
||
|
|
return await self._dispatch_final(chunk)
|
||
|
|
return False
|
||
|
|
|
||
|
|
async def _dispatch_tool(self, chunk: StreamingChunk) -> bool:
|
||
|
|
self._pending += 1
|
||
|
|
self._tool_count += 1
|
||
|
|
try:
|
||
|
|
prefix = self._response_prefix if self._response_prefix else ""
|
||
|
|
content = f"{prefix}{chunk.content}" if chunk.content else ""
|
||
|
|
if not content:
|
||
|
|
return False
|
||
|
|
|
||
|
|
await self._send_ordered(content)
|
||
|
|
self._block_count += 1
|
||
|
|
return True
|
||
|
|
except Exception as e:
|
||
|
|
logger.exception("Tool dispatch failed")
|
||
|
|
self._failed_count += 1
|
||
|
|
if self._on_error:
|
||
|
|
self._on_error(e)
|
||
|
|
raise
|
||
|
|
finally:
|
||
|
|
self._pending -= 1
|
||
|
|
self._maybe_fire_idle()
|
||
|
|
|
||
|
|
async def _dispatch_block(self, chunk: StreamingChunk) -> bool:
|
||
|
|
if not chunk.content:
|
||
|
|
return False
|
||
|
|
|
||
|
|
async with self._flush_lock:
|
||
|
|
self._blocks.append(chunk.content)
|
||
|
|
now = time.monotonic()
|
||
|
|
|
||
|
|
if self._block_deadline == 0.0:
|
||
|
|
self._block_deadline = now + _DEFAULT_DEBOUNCE_MS / 1000.0
|
||
|
|
self._block_timer = asyncio.create_task(self._flush_after_debounce())
|
||
|
|
|
||
|
|
if sum(len(b) for b in self._blocks) >= _DEFAULT_MIN_CHARS:
|
||
|
|
await self._flush_blocks()
|
||
|
|
return True
|
||
|
|
|
||
|
|
async def _flush_after_debounce(self) -> None:
|
||
|
|
try:
|
||
|
|
await asyncio.sleep(_DEFAULT_DEBOUNCE_MS / 1000.0)
|
||
|
|
except asyncio.CancelledError:
|
||
|
|
return
|
||
|
|
async with self._flush_lock:
|
||
|
|
await self._flush_blocks()
|
||
|
|
|
||
|
|
async def _flush_blocks(self) -> None:
|
||
|
|
if not self._blocks:
|
||
|
|
return
|
||
|
|
combined = "".join(self._blocks)
|
||
|
|
self._blocks.clear()
|
||
|
|
self._block_deadline = 0.0
|
||
|
|
if self._block_timer and not self._block_timer.done():
|
||
|
|
self._block_timer.cancel()
|
||
|
|
self._block_timer = None
|
||
|
|
|
||
|
|
payload = ReplyPayload(
|
||
|
|
target_id=self._target_id,
|
||
|
|
content=combined,
|
||
|
|
reply_to_id=self._msg.msg_id if self._block_count == 0 else None,
|
||
|
|
)
|
||
|
|
if payload.content_key in self._sent_content_keys:
|
||
|
|
return
|
||
|
|
self._sent_content_keys.add(payload.content_key)
|
||
|
|
|
||
|
|
self._pending += 1
|
||
|
|
self._block_count += 1
|
||
|
|
try:
|
||
|
|
await self._send_ordered(combined)
|
||
|
|
except Exception as e:
|
||
|
|
logger.exception("Block dispatch failed")
|
||
|
|
self._failed_count += 1
|
||
|
|
if self._on_error:
|
||
|
|
self._on_error(e)
|
||
|
|
raise
|
||
|
|
finally:
|
||
|
|
self._pending -= 1
|
||
|
|
self._maybe_fire_idle()
|
||
|
|
|
||
|
|
async def _dispatch_final(self, chunk: StreamingChunk) -> bool:
|
||
|
|
async with self._flush_lock:
|
||
|
|
await self._flush_blocks()
|
||
|
|
|
||
|
|
if chunk.content and chunk.content != "finished":
|
||
|
|
self._pending += 1
|
||
|
|
self._final_count += 1
|
||
|
|
try:
|
||
|
|
await self._send_ordered(chunk.content)
|
||
|
|
except Exception as e:
|
||
|
|
logger.exception("Final dispatch failed")
|
||
|
|
self._failed_count += 1
|
||
|
|
if self._on_error:
|
||
|
|
self._on_error(e)
|
||
|
|
raise
|
||
|
|
finally:
|
||
|
|
self._pending -= 1
|
||
|
|
self._maybe_fire_idle()
|
||
|
|
self._done = True
|
||
|
|
self._complete.set()
|
||
|
|
return True
|
||
|
|
|
||
|
|
self._done = True
|
||
|
|
self._complete.set()
|
||
|
|
self._maybe_fire_idle()
|
||
|
|
return True
|
||
|
|
|
||
|
|
async def _send_ordered(self, content: str) -> None:
|
||
|
|
async with self._send_chain:
|
||
|
|
if self._abort_event.is_set():
|
||
|
|
return
|
||
|
|
|
||
|
|
if self._human_delay and self._block_count > 0:
|
||
|
|
delay = self._human_delay[0] + (self._human_delay[1] - self._human_delay[0]) * random.random()
|
||
|
|
await asyncio.sleep(delay / 1000.0)
|
||
|
|
|
||
|
|
if self._response_prefix and self._block_count == 0:
|
||
|
|
prefix = self._response_prefix.format(
|
||
|
|
agent_name=self._msg.channel_type,
|
||
|
|
model="",
|
||
|
|
)
|
||
|
|
content = f"{prefix}{content}"
|
||
|
|
|
||
|
|
if self._before_deliver:
|
||
|
|
try:
|
||
|
|
content = await self._before_deliver(content)
|
||
|
|
except Exception as e:
|
||
|
|
logger.exception("before_deliver hook failed")
|
||
|
|
if self._on_error:
|
||
|
|
self._on_error(e)
|
||
|
|
return
|
||
|
|
if content is None:
|
||
|
|
self._cancelled_count += 1
|
||
|
|
return
|
||
|
|
|
||
|
|
await self._send_fn(content)
|
||
|
|
|
||
|
|
@property
|
||
|
|
def _target_id(self) -> str:
|
||
|
|
return self._msg.group.id if self._msg.group and self._msg.group.id else self._msg.sender.id
|
||
|
|
|
||
|
|
async def wait_idle(self, timeout_ms: float | None = None) -> None:
|
||
|
|
timeout = timeout_ms or self._timeout_ms
|
||
|
|
try:
|
||
|
|
await asyncio.wait_for(self._complete.wait(), timeout=timeout / 1000.0)
|
||
|
|
except TimeoutError:
|
||
|
|
self._abort_event.set()
|
||
|
|
self._done = True
|
||
|
|
self._complete.set()
|
||
|
|
logger.warning("ReplyDispatcher idle timeout after %.0fms, aborting", timeout)
|
||
|
|
|
||
|
|
def mark_complete(self) -> None:
|
||
|
|
self._done = True
|
||
|
|
self._complete.set()
|
||
|
|
self._maybe_fire_idle()
|
||
|
|
|
||
|
|
async def abort(self) -> None:
|
||
|
|
async with self._flush_lock:
|
||
|
|
if self._blocks:
|
||
|
|
await self._flush_blocks()
|
||
|
|
self._abort_event.set()
|
||
|
|
self._complete.set()
|
||
|
|
|
||
|
|
def _maybe_fire_idle(self) -> None:
|
||
|
|
if self._done and self._pending == 0 and self._on_idle:
|
||
|
|
self._on_idle()
|
||
|
|
|
||
|
|
@property
|
||
|
|
def tool_count(self) -> int:
|
||
|
|
return self._tool_count
|
||
|
|
|
||
|
|
@property
|
||
|
|
def block_count(self) -> int:
|
||
|
|
return self._block_count
|
||
|
|
|
||
|
|
@property
|
||
|
|
def final_count(self) -> int:
|
||
|
|
return self._final_count
|
||
|
|
|
||
|
|
@property
|
||
|
|
def failed_count(self) -> int:
|
||
|
|
return self._failed_count
|
||
|
|
|
||
|
|
@property
|
||
|
|
def cancelled_count(self) -> int:
|
||
|
|
return self._cancelled_count
|
||
|
|
|
||
|
|
@property
|
||
|
|
def queued_counts(self) -> dict[str, int]:
|
||
|
|
return {"tool": self._tool_count, "block": self._block_count, "final": self._final_count}
|
||
|
|
|
||
|
|
@property
|
||
|
|
def is_idle(self) -> bool:
|
||
|
|
return self._done and self._pending == 0
|