from __future__ import annotations import asyncio import time from typing import TYPE_CHECKING, Any import httpx from yuxi.channels.models import DeliveryResult from yuxi.utils.logging_config import logger if TYPE_CHECKING: from yuxi.channels.adapters.dingding.token import DingDingTokenManager STREAM_UPDATE_MIN_INTERVAL_S = 0.16 STREAM_SIGNIFICANT_INCREMENT = 18 STREAM_NATURAL_BOUNDARIES = ("\n", "。", "！", "？", "；", "：", ",", ". ", "! ", "? ", "; ", ": ") STREAM_PLACEHOLDER_TEXT = "⏳ Thinking..." MAX_CONSECUTIVE_FAILURES = 5 BACKOFF_BASE_S = 0.5 BACKOFF_MAX_S = 8.0 class DingDingStreamingSession: def __init__( self, token_manager: DingDingTokenManager, open_conversation_id: str, robot_code: str, chat_id: str, *, config: dict[str, Any] | None = None, http_client: httpx.AsyncClient | None = None, ): self._token_manager = token_manager self._open_conversation_id = open_conversation_id self._robot_code = robot_code self._chat_id = chat_id self._config = config or {} self._http_client = http_client self._msg_id: str | None = None self._accumulated = "" self._last_update = 0.0 self._closed = False self._consecutive_failures = 0 self._started = False self._chat_type = "group" if chat_id.startswith("group_") else "direct" stream_cfg = self._config.get("streaming", {}) if not isinstance(stream_cfg, dict): stream_cfg = {} self._stream_mode = stream_cfg.get("mode", "block") block_cfg = stream_cfg.get("block", {}) if not isinstance(block_cfg, dict): block_cfg = {} self._update_interval_ms = block_cfg.get("update_interval_ms", 800) self._coalesce_min_chars = block_cfg.get("coalesce_min_chars", 100) async def start(self, initial_text: str = "") -> DeliveryResult: if self._started: return DeliveryResult(success=True, message_id=self._msg_id or "") self._started = True send_text = initial_text or STREAM_PLACEHOLDER_TEXT try: from .send import send_markdown result = await send_markdown( self._token_manager, self._open_conversation_id, self._robot_code, "ForcePilot", send_text, chat_type=self._chat_type, http_client=self._http_client, ) if result.success: self._msg_id = result.message_id self._consecutive_failures = 0 else: self._consecutive_failures += 1 return result except Exception as e: logger.error(f"[DingDingStream] start error: {e}") self._consecutive_failures += 1 return DeliveryResult(success=False, error=str(e)) async def send_chunk(self, chunk: str, finished: bool = False) -> DeliveryResult: if self._closed: return DeliveryResult(success=False, error="Stream closed") if not self._msg_id: return await self.start(chunk) if finished: return await self._finish() self._accumulated += chunk return await self._maybe_update() async def _finish(self) -> DeliveryResult: self._closed = True try: from .send import send_markdown result = await send_markdown( self._token_manager, self._open_conversation_id, self._robot_code, "ForcePilot", self._accumulated, chat_type=self._chat_type, http_client=self._http_client, ) if result.success: self._consecutive_failures = 0 else: self._consecutive_failures += 1 return result except Exception as e: logger.error(f"[DingDingStream] finish error: {e}") self._consecutive_failures += 1 return DeliveryResult(success=False, error=str(e)) async def _maybe_update(self) -> DeliveryResult: now = time.monotonic() elapsed_ms = (now - self._last_update) * 1000 should_send = ( len(self._accumulated) >= self._coalesce_min_chars and elapsed_ms >= self._update_interval_ms ) or any(self._accumulated.endswith(b) for b in STREAM_NATURAL_BOUNDARIES) if not should_send: return DeliveryResult(success=True, message_id=self._msg_id) self._last_update = now try: result = await self._edit_message(self._accumulated) if result.success: self._consecutive_failures = 0 else: self._consecutive_failures += 1 logger.debug(f"[DingDingStream] edit failed (consecutive={self._consecutive_failures}): {result.error}") return result except Exception as e: self._consecutive_failures += 1 logger.debug(f"[DingDingStream] edit error: {e}") return DeliveryResult(success=False, error=str(e)) async def _edit_message(self, content: str) -> DeliveryResult: if not self._token_manager: return DeliveryResult(success=False, error="Not connected") try: token = await self._token_manager.get_token() http_client = self._http_client or httpx.AsyncClient(timeout=httpx.Timeout(15)) url = f"https://api.dingtalk.com/v1.0/robot/messages/{self._msg_id}" headers = { "x-acs-dingtalk-access-token": token, "Content-Type": "application/json", } import json payload = { "msgParam": json.dumps({"content": content}, ensure_ascii=False), "msgKey": "sampleMarkdown", "openConversationId": self._open_conversation_id, "robotCode": self._robot_code, } resp = await http_client.put(url, json=payload, headers=headers) if resp.status_code != 200: return DeliveryResult(success=False, error=f"DingDing API HTTP {resp.status_code}") data = resp.json() return DeliveryResult(success=True, message_id=data.get("processQueryKey", self._msg_id)) except Exception as e: logger.error(f"[DingDingStream] _edit_message failed: {e}") return DeliveryResult(success=False, error=str(e)) @property def message_id(self) -> str | None: return self._msg_id @property def accumulated_text(self) -> str: return self._accumulated @property def is_closed(self) -> bool: return self._closed @property def consecutive_failures(self) -> int: return self._consecutive_failures def should_degrade(self) -> bool: return self._consecutive_failures >= MAX_CONSECUTIVE_FAILURES def backoff_delay(self) -> float: if self._consecutive_failures == 0: return 0.0 return min(BACKOFF_BASE_S * (2 ** (self._consecutive_failures - 1)), BACKOFF_MAX_S) class DingDingStreamManager: def __init__(self, adapter): self._adapter = adapter self._sessions: dict[str, DingDingStreamingSession] = {} self._pending_block_text: dict[str, str] = {} self._lock = asyncio.Lock() async def send_chunk(self, chat_id: str, msg_id: str, chunk: str, finished: bool) -> DeliveryResult: config = getattr(self._adapter, "config", {}) stream_cfg = config.get("streaming", {}) if not isinstance(stream_cfg, dict): stream_cfg = {} stream_mode = stream_cfg.get("mode", "block") if stream_mode == "partial": return await self._send_partial(chat_id, msg_id, chunk, finished) return await self._send_block(chat_id, msg_id, chunk, finished) async def _send_partial(self, chat_id: str, msg_id: str, chunk: str, finished: bool) -> DeliveryResult: session = self._sessions.get(chat_id) if session is None or session.is_closed: adapter_config = getattr(self._adapter, "config", {}) accounts = adapter_config.get("accounts", {}) robot_code = accounts.get("default", {}).get("robot_code", "") open_conv_id = chat_id.replace("group_", "").replace("dm_", "") token_mgr = getattr(self._adapter, "_token_manager", None) http_client = await self._get_http_client() session = DingDingStreamingSession( token_mgr, open_conv_id, robot_code, chat_id, config=adapter_config, http_client=http_client, ) self._sessions[chat_id] = session if session.should_degrade(): logger.warning( f"[DingDingStream] Degrading partial stream for {chat_id}" f" after {session.consecutive_failures} failures" ) return await self._fallback_block_send(chat_id) return await session.send_chunk(chunk, finished) async def _send_block(self, chat_id: str, msg_id: str, chunk: str, finished: bool) -> DeliveryResult: async with self._lock: if not finished: self._pending_block_text[chat_id] = self._pending_block_text.get(chat_id, "") + chunk return DeliveryResult(success=True, message_id=msg_id or "pending") final_text = self._pending_block_text.pop(chat_id, chunk) return await self._fallback_block_send(chat_id, final_text) async def _fallback_block_send(self, chat_id: str, text: str = "") -> DeliveryResult: if not text: return DeliveryResult(success=True) accounts = getattr(self._adapter, "config", {}).get("accounts", {}) robot_code = accounts.get("default", {}).get("robot_code", "") open_conv_id = chat_id.replace("group_", "").replace("dm_", "") chat_type = "group" if chat_id.startswith("group_") else "direct" token_mgr = getattr(self._adapter, "_token_manager", None) http_client = await self._get_http_client() from .send import send_markdown async with self._lock: return await send_markdown( token_mgr, open_conv_id, robot_code, "ForcePilot", text, chat_type=chat_type, http_client=http_client, ) async def _get_http_client(self) -> httpx.AsyncClient: if hasattr(self._adapter, "_get_http_client"): return await self._adapter._get_http_client() return httpx.AsyncClient(timeout=httpx.Timeout(15)) def close_session(self, chat_id: str) -> None: self._sessions.pop(chat_id, None) self._pending_block_text.pop(chat_id, None) def cleanup(self) -> None: self._sessions.clear() self._pending_block_text.clear()