from __future__ import annotations import logging from dataclasses import dataclass from yuxi.channel.config.defaults import ( BLOCK_STREAMING_BREAK, BLOCK_STREAMING_CHUNK_BREAK_PREFERENCE, BLOCK_STREAMING_CHUNK_MAX_CHARS, BLOCK_STREAMING_CHUNK_MIN_CHARS, STREAMING_PREVIEW_MIN_INITIAL_CHARS, STREAMING_PREVIEW_THROTTLE_MS, ) from .utils import get_nio logger = logging.getLogger(__name__) @dataclass class MatrixDraftSession: target_id: str client: object replace_event_id: str | None = None streaming_mode: str = "partial" delivered: bool = False stopped: bool = False total_text: str = "" async def start(self, initial_text: str) -> str | None: content = {"msgtype": "m.text", "body": initial_text} resp = await self.client.room_send(self.target_id, "m.room.message", content) self.replace_event_id = resp.event_id self.delivered = True self.total_text = initial_text return resp.event_id async def append(self, text: str) -> None: if self.stopped or not self.replace_event_id: return self.total_text = text edit_content = { "msgtype": "m.text", "body": text, "m.new_content": {"msgtype": "m.text", "body": text}, "m.relates_to": { "rel_type": "m.replace", "event_id": self.replace_event_id, }, } try: await self.client.room_send(self.target_id, "m.room.message", edit_content) except Exception: logger.debug("Matrix draft append failed, falling back") self.stopped = True async def finalize(self, final_text: str = "") -> str: self.stopped = True if not self.replace_event_id: return "" text = final_text or self.total_text content = {"msgtype": "m.text", "body": text} edit_content = { **content, "m.new_content": content, "m.relates_to": { "rel_type": "m.replace", "event_id": self.replace_event_id, }, } if self.streaming_mode == "quiet": edit_content["com.forcepilot.finalized_preview"] = True try: await self.client.room_send(self.target_id, "m.room.message", edit_content) except Exception: logger.warning("Matrix draft finalize failed for %s", self.replace_event_id) return self.replace_event_id async def abort(self) -> None: self.stopped = True if self.replace_event_id and self.total_text: await self.finalize(self.total_text) streaming_mode = "partial" preview_stream_throttle_ms = STREAMING_PREVIEW_THROTTLE_MS preview_min_initial_chars = STREAMING_PREVIEW_MIN_INITIAL_CHARS block_streaming_enabled = False block_streaming_break = BLOCK_STREAMING_BREAK block_streaming_chunk_min_chars = BLOCK_STREAMING_CHUNK_MIN_CHARS block_streaming_chunk_max_chars = BLOCK_STREAMING_CHUNK_MAX_CHARS block_streaming_chunk_break_preference = BLOCK_STREAMING_CHUNK_BREAK_PREFERENCE block_streaming_coalesce_defaults = None def create_draft_stream_session( target_id: str, config: dict = None, account_id: str = None, client: object = None ) -> object: if client is not None: return MatrixDraftSession( target_id=target_id, client=client, streaming_mode=_resolve_streaming_mode(config, account_id), ) from .config import _apply_env_overrides, _dict_to_account cfg = config or {} aid = account_id or "default" account_data = cfg.get("accounts", {}).get(aid, {}) account = _dict_to_account(account_data) account = _apply_env_overrides(account) if not account.homeserver or not account.access_token: logger.warning("Matrix draft stream session: account not configured") return None nio = get_nio() new_client = nio.AsyncClient( homeserver=account.homeserver, user=account.user_id, device_id=account.device_id, ) new_client.access_token = account.access_token return MatrixDraftSession( target_id=target_id, client=new_client, streaming_mode=account.streaming or streaming_mode, ) def _resolve_streaming_mode(config: dict | None, account_id: str | None) -> str: cfg = config or {} aid = account_id or "default" account_data = cfg.get("accounts", {}).get(aid, {}) return account_data.get("streaming", streaming_mode) def create_block_chunker() -> object: from yuxi.channel.sdk import BlockChunker return BlockChunker()