96 lines
3.0 KiB
Python
96 lines
3.0 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import logging
|
||
|
|
import time
|
||
|
|
|
||
|
|
from yuxi.channel.extensions.mattermost.client import MattermostClient
|
||
|
|
from yuxi.channel.extensions.mattermost.errors import MattermostError
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
class DraftStreamManager:
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
client: MattermostClient,
|
||
|
|
channel_id: str,
|
||
|
|
max_chars: int = 4000,
|
||
|
|
throttle_ms: int = 1200,
|
||
|
|
root_id: str | None = None,
|
||
|
|
):
|
||
|
|
self.client = client
|
||
|
|
self.channel_id = channel_id
|
||
|
|
self.max_chars = max_chars
|
||
|
|
self.throttle_ms = max(throttle_ms, 250)
|
||
|
|
self.root_id = root_id
|
||
|
|
self.stream_post_id: str | None = None
|
||
|
|
self.last_update_at: float = 0.0
|
||
|
|
self._dirty = False
|
||
|
|
self._full_text = ""
|
||
|
|
|
||
|
|
async def update(self, text: str) -> None:
|
||
|
|
self._full_text = text
|
||
|
|
now = time.monotonic()
|
||
|
|
elapsed_ms = (now - self.last_update_at) * 1000
|
||
|
|
|
||
|
|
if self.last_update_at > 0 and elapsed_ms < self.throttle_ms:
|
||
|
|
self._dirty = True
|
||
|
|
return
|
||
|
|
|
||
|
|
await self._flush(text)
|
||
|
|
|
||
|
|
async def finalize(self, text: str | None = None) -> str | None:
|
||
|
|
if text is not None:
|
||
|
|
self._full_text = text
|
||
|
|
if self.stream_post_id:
|
||
|
|
try:
|
||
|
|
truncated = self._full_text[: self.max_chars]
|
||
|
|
await self.client.update_post(
|
||
|
|
self.stream_post_id,
|
||
|
|
{"message": truncated},
|
||
|
|
)
|
||
|
|
return self.stream_post_id
|
||
|
|
except MattermostError:
|
||
|
|
pass
|
||
|
|
return None
|
||
|
|
|
||
|
|
async def discard(self) -> None:
|
||
|
|
if self.stream_post_id:
|
||
|
|
try:
|
||
|
|
await self.client.delete_post(self.stream_post_id)
|
||
|
|
except MattermostError:
|
||
|
|
pass
|
||
|
|
self.stream_post_id = None
|
||
|
|
|
||
|
|
async def _flush(self, text: str) -> None:
|
||
|
|
truncated = text[: self.max_chars]
|
||
|
|
if len(text) > self.max_chars:
|
||
|
|
truncated += "..."
|
||
|
|
try:
|
||
|
|
if self.stream_post_id:
|
||
|
|
await self.client.update_post(self.stream_post_id, {"message": truncated})
|
||
|
|
else:
|
||
|
|
payload: dict = {
|
||
|
|
"channel_id": self.channel_id,
|
||
|
|
"message": truncated + "\n\n_Thinking…_",
|
||
|
|
}
|
||
|
|
if self.root_id:
|
||
|
|
payload["root_id"] = self.root_id
|
||
|
|
post = await self.client.create_post(payload)
|
||
|
|
self.stream_post_id = post["id"]
|
||
|
|
self.last_update_at = time.monotonic()
|
||
|
|
self._dirty = False
|
||
|
|
except MattermostError:
|
||
|
|
pass
|
||
|
|
|
||
|
|
async def maybe_flush(self) -> None:
|
||
|
|
if self._dirty and self._full_text:
|
||
|
|
await self._flush(self._full_text)
|
||
|
|
|
||
|
|
def can_finalize_in_place(self, target_channel_id: str, target_root_id: str | None = None) -> bool:
|
||
|
|
if self.channel_id != target_channel_id:
|
||
|
|
return False
|
||
|
|
if self.root_id != target_root_id:
|
||
|
|
return False
|
||
|
|
return self.stream_post_id is not None
|