89 lines
2.8 KiB
Python
89 lines
2.8 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import logging
|
||
|
|
import time
|
||
|
|
from typing import TYPE_CHECKING
|
||
|
|
|
||
|
|
if TYPE_CHECKING:
|
||
|
|
from yuxi.channel.extensions.rocketchat.outbound import RocketChatOutbound
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
class RocketChatEditStream:
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
outbound: RocketChatOutbound,
|
||
|
|
room_id: str,
|
||
|
|
thread_id: str | None = None,
|
||
|
|
max_chars: int = 4000,
|
||
|
|
throttle_ms: int = 1200,
|
||
|
|
):
|
||
|
|
self.outbound = outbound
|
||
|
|
self.room_id = room_id
|
||
|
|
self.thread_id = thread_id
|
||
|
|
self.max_chars = max_chars
|
||
|
|
self.throttle_ms = max(throttle_ms, 250)
|
||
|
|
self.stream_msg_id: str | None = None
|
||
|
|
self.last_update_at: float = 0
|
||
|
|
self._dirty: bool = False
|
||
|
|
|
||
|
|
async def update(self, text: str) -> None:
|
||
|
|
now = time.monotonic()
|
||
|
|
elapsed = now - self.last_update_at
|
||
|
|
if elapsed < self.throttle_ms / 1000:
|
||
|
|
self._dirty = True
|
||
|
|
return
|
||
|
|
|
||
|
|
if self._dirty or elapsed >= self.throttle_ms / 1000:
|
||
|
|
await self._flush(text)
|
||
|
|
|
||
|
|
async def _flush(self, text: str) -> None:
|
||
|
|
truncated = text[: self.max_chars]
|
||
|
|
try:
|
||
|
|
if self.stream_msg_id:
|
||
|
|
await self.outbound.edit_message(
|
||
|
|
self.room_id,
|
||
|
|
self.stream_msg_id,
|
||
|
|
truncated,
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
resp = await self.outbound.client.post_message(
|
||
|
|
self.room_id,
|
||
|
|
truncated + "\n\n_Thinking…_",
|
||
|
|
thread_id=self.thread_id,
|
||
|
|
)
|
||
|
|
self.stream_msg_id = resp.get("message", {}).get("_id", "")
|
||
|
|
except Exception as e:
|
||
|
|
logger.debug("Stream flush error: %s", e)
|
||
|
|
self.last_update_at = time.monotonic()
|
||
|
|
self._dirty = False
|
||
|
|
|
||
|
|
async def finalize(self, final_text: str) -> None:
|
||
|
|
if self.stream_msg_id:
|
||
|
|
try:
|
||
|
|
await self.outbound.edit_message(
|
||
|
|
self.room_id,
|
||
|
|
self.stream_msg_id,
|
||
|
|
final_text[: self.max_chars],
|
||
|
|
)
|
||
|
|
except Exception as e:
|
||
|
|
logger.debug("Stream finalize error: %s", e)
|
||
|
|
else:
|
||
|
|
try:
|
||
|
|
await self.outbound.send_text(
|
||
|
|
self.room_id,
|
||
|
|
final_text,
|
||
|
|
thread_id=self.thread_id,
|
||
|
|
)
|
||
|
|
except Exception as e:
|
||
|
|
logger.debug("Stream finalize send error: %s", e)
|
||
|
|
|
||
|
|
async def cancel(self) -> None:
|
||
|
|
if self.stream_msg_id:
|
||
|
|
try:
|
||
|
|
await self.outbound.delete_message(self.room_id, self.stream_msg_id)
|
||
|
|
except Exception:
|
||
|
|
pass
|
||
|
|
self.stream_msg_id = None
|