ForcePilot/backend/package/yuxi/channels/adapters/feishu/stream.py

from __future__ import annotations

import asyncio
import json
import time
from dataclasses import dataclass
from typing import Literal

from yuxi.utils.logging_config import logger

try:
    import lark_oapi

    HAS_LARK_SDK = True
except ImportError:
    HAS_LARK_SDK = False
    lark_oapi = None  # type: ignore

STREAM_UPDATE_MIN_INTERVAL_S = 0.16
STREAM_SIGNIFICANT_INCREMENT = 18
STREAM_NATURAL_BOUNDARIES = ("\n", "。", "！", "？", "；", "：", ",", ". ", "! ", "? ", "; ", ": ")
STREAM_PLACEHOLDER_TEXT = "⏳ Thinking..."
STREAM_START_BACKOFF_S = 60

CHAR_PRINT_INTERVAL_S = 0.05
CHAR_PRINT_STEP = 1

StreamSendMode = Literal["create", "reply", "root_create"]


def merge_streaming_text(existing: str, incoming: str) -> str:
    if not existing:
        return incoming
    if not incoming:
        return existing
    if incoming.startswith(existing):
        return incoming
    if existing.startswith(incoming):
        return existing

    for i in range(len(incoming) - 1, 0, -1):
        prefix = incoming[:i]
        if existing.endswith(prefix):
            return existing + incoming[i:]

    overlap_start = max(0, len(existing) - len(incoming))
    for j in range(1, len(incoming)):
        suffix = incoming[j:]
        if existing.startswith(suffix):
            return incoming[:j] + existing

    return existing + incoming


@dataclass
class FeishuStreamingSession:
    """管理单个流式卡片的完整生命周期"""

    _client: object
    _chat_id: str
    _msg_id: str | None = None
    _reply_to_msg_id: str | None = None
    _root_id: str | None = None
    _send_mode: StreamSendMode = "create"

    def __post_init__(self):
        self._card_id: str | None = None
        self._message_id: str | None = self._msg_id
        self._accumulated = ""
        self._last_update = 0.0
        self._closed = False

    async def start(self) -> str:
        if not HAS_LARK_SDK:
            return ""

        try:
            placeholder_card = {
                "header": {"title": {"tag": "plain_text", "content": STREAM_PLACEHOLDER_TEXT}, "template": "blue"},
                "elements": [{"tag": "markdown", "content": ""}],
            }

            create_card_req = (
                lark_oapi.api.cardkit.v1.CreateCardRequest.builder()
                .request_body(lark_oapi.api.cardkit.v1.CreateCardRequestBody.builder().card(placeholder_card).build())
                .build()
            )
            create_resp = await self._client.cardkit.v1.card.create(create_card_req)
            if not create_resp.success():
                logger.warning(f"[FeishuStream] Failed to create card: {create_resp.msg}")
                return ""

            self._card_id = create_resp.data.get("card_id", "")

            msg_content = json.dumps(
                {"config": {"streaming_mode": True}, "card_id": self._card_id},
                ensure_ascii=False,
            )

            chat_type = "private" if self._chat_id.startswith("dm_") else "group"

            if self._send_mode == "reply" and self._reply_to_msg_id:
                msg_req = _make_reply_message(self._reply_to_msg_id, msg_content)
            elif self._send_mode == "root_create":
                msg_req = _make_interactive_message(self._chat_id, chat_type, msg_content)
                if self._reply_to_msg_id:
                    msg_req.root_id = self._reply_to_msg_id
            else:
                msg_req = _make_interactive_message(self._chat_id, chat_type, msg_content)
                if self._root_id:
                    msg_req.root_id = self._root_id

            msg_resp = await self._client.im.message.create(msg_req)
            if msg_resp.success():
                self._message_id = msg_resp.data.get("message_id", "")
                logger.info(f"[FeishuStream] Card sent, card_id={self._card_id}, message_id={self._message_id}")
                return self._card_id
            logger.warning(f"[FeishuStream] Failed to send card message: {msg_resp.msg}")
            return ""
        except Exception as e:
            logger.error(f"[FeishuStream] start error: {e}")
            return ""

    async def update(self, full_text: str) -> None:
        if not self._card_id or self._closed:
            return

        now = time.monotonic()
        merged = merge_streaming_text(self._accumulated, full_text)
        elapsed = now - self._last_update
        increment = len(merged) - len(self._accumulated)

        skip = (
            elapsed < STREAM_UPDATE_MIN_INTERVAL_S
            and increment < STREAM_SIGNIFICANT_INCREMENT
            and not any(full_text.endswith(b) for b in STREAM_NATURAL_BOUNDARIES)
        )
        if skip:
            self._accumulated = merged
            return

        self._accumulated = merged
        self._last_update = now

        try:
            update_req = (
                lark_oapi.api.cardkit.v1.UpdateCardContentRequest.builder()
                .card_id(self._card_id)
                .request_body(
                    lark_oapi.api.cardkit.v1.UpdateCardContentRequestBody.builder()
                    .elements([{"tag": "markdown", "content": merged}])
                    .build()
                )
                .build()
            )
            resp = await self._client.cardkit.v1.card.update_content(update_req)
            if not resp.success():
                logger.debug(f"[FeishuStream] Update failed: {resp.msg}")
        except Exception as e:
            logger.debug(f"[FeishuStream] Update error: {e}")

    async def close(self) -> None:
        if self._closed or not self._card_id:
            return
        self._closed = True

        try:
            close_req = (
                lark_oapi.api.cardkit.v1.UpdateCardSettingsRequest.builder()
                .card_id(self._card_id)
                .request_body(
                    lark_oapi.api.cardkit.v1.UpdateCardSettingsRequestBody.builder().streaming_mode(False).build()
                )
                .build()
            )
            resp = await self._client.cardkit.v1.card.update_settings(close_req)
            if resp.success():
                logger.info(f"[FeishuStream] Card closed, card_id={self._card_id}")
        except Exception as e:
            logger.warning(f"[FeishuStream] Close error: {e}")


def _make_interactive_message(receive_id: str, chat_type: str, content_json: str) -> Any:
    receive_id_type = "open_id" if chat_type == "private" else "chat_id"

    return (
        lark_oapi.api.im.v1.CreateMessageRequest.builder()
        .receive_id_type(receive_id_type)
        .request_body(
            lark_oapi.api.im.v1.CreateMessageRequestBody.builder()
            .receive_id(receive_id)
            .msg_type("interactive")
            .content(content_json)
            .build()
        )
        .build()
    )


def _make_reply_message(reply_to_msg_id: str, content_json: str) -> Any:
    return (
        lark_oapi.api.im.v1.ReplyMessageRequest.builder()
        .message_id(reply_to_msg_id)
        .request_body(
            lark_oapi.api.im.v1.ReplyMessageRequestBody.builder().msg_type("interactive").content(content_json).build()
        )
        .build()
    )


class CharacterStreamingSession:
    def __init__(
        self,
        session: FeishuStreamingSession,
        *,
        interval_s: float = CHAR_PRINT_INTERVAL_S,
        step: int = CHAR_PRINT_STEP,
    ):
        self._session = session
        self._interval_s = interval_s
        self._step = step
        self._displayed_len = 0

    async def update(self, full_text: str) -> None:
        while self._displayed_len < len(full_text):
            next_len = min(self._displayed_len + self._step, len(full_text))
            display_text = full_text[:next_len]
            self._displayed_len = next_len
            await self._session.update(display_text)
            if next_len < len(full_text):
                await asyncio.sleep(self._interval_s)

    async def close(self) -> None:
        await self._session.close()