daily_stock_analysis/bot/platforms/dingtalk_stream.py

# -*- coding: utf-8 -*-
"""
===================================
钉钉 Stream 模式适配器
===================================

使用钉钉官方 Stream SDK 接入机器人，无需公网 IP 和 Webhook 配置。

优势：
- 不需要公网 IP 或域名
- 不需要配置 Webhook URL
- 通过 WebSocket 长连接接收消息
- 更简单的接入方式

依赖：
pip install dingtalk-stream

钉钉 Stream SDK：
https://github.com/open-dingtalk/dingtalk-stream-sdk-python
"""

import logging
import asyncio
import threading
from datetime import datetime
from typing import Optional, Callable, Any

logger = logging.getLogger(__name__)

# 尝试导入钉钉 Stream SDK
try:
    import dingtalk_stream
    from dingtalk_stream import AckMessage

    DINGTALK_STREAM_AVAILABLE = True
except ImportError:
    DINGTALK_STREAM_AVAILABLE = False
    logger.warning("[DingTalk Stream] dingtalk-stream SDK 未安装，Stream 模式不可用")
    logger.warning("[DingTalk Stream] 请运行: pip install dingtalk-stream")

from bot.models import BotMessage, BotResponse, ChatType


class DingtalkStreamHandler:
    """
    钉钉 Stream 模式消息处理器
    
    将 Stream SDK 的回调转换为统一的 BotMessage 格式，
    并调用命令分发器处理。
    """

    def __init__(self, on_message: Callable[[BotMessage], BotResponse]):
        """
        Args:
            on_message: 消息处理回调函数，接收 BotMessage 返回 BotResponse
        """
        self._on_message = on_message
        self._logger = logger

    @staticmethod
    def _truncate_log_content(text: str, max_len: int = 200) -> str:
        cleaned = text.replace("\n", " ").strip()
        if len(cleaned) > max_len:
            return f"{cleaned[:max_len]}..."
        return cleaned

    def _log_incoming_message(self, message: BotMessage) -> None:
        content = message.raw_content or message.content or ""
        summary = self._truncate_log_content(content)
        self._logger.info(
            "[DingTalk Stream] Incoming message: msg_id=%s user_id=%s chat_id=%s chat_type=%s content=%s",
            message.message_id,
            message.user_id,
            message.chat_id,
            getattr(message.chat_type, "value", message.chat_type),
            summary,
        )

    if DINGTALK_STREAM_AVAILABLE:
        class _ChatbotHandler(dingtalk_stream.ChatbotHandler):
            """内部消息处理器"""

            def __init__(self, parent: 'DingtalkStreamHandler'):
                super().__init__()
                self._parent = parent
                self.logger = logger

            async def process(self, callback: dingtalk_stream.CallbackMessage):
                """处理收到的消息"""
                try:
                    # 解析消息
                    incoming = dingtalk_stream.ChatbotMessage.from_dict(callback.data)

                    # 转换为统一格式
                    bot_message = self._parent._parse_stream_message(incoming, callback.data)

                    if bot_message:
                        self._parent._log_incoming_message(bot_message)
                        # 调用消息处理回调
                        response = self._parent._on_message(bot_message)

                        # 发送回复
                        if response and response.text:
                            # 构建 @用户 前缀（群聊场景下需要在文本中包含 @用户名）
                            if response.at_user and incoming.sender_nick:
                                if response.markdown:
                                    self.reply_markdown(
                                        title="股票分析助手",
                                        text=f"@{incoming.sender_nick} " + response.text,
                                        incoming_message=incoming
                                    )
                                else:
                                    self.reply_text(response.text, incoming)

                    return AckMessage.STATUS_OK, 'OK'

                except Exception as e:
                    self.logger.error(f"[DingTalk Stream] 处理消息失败: {e}")
                    self.logger.exception(e)
                    return AckMessage.STATUS_SYSTEM_EXCEPTION, str(e)

        def create_handler(self) -> '_ChatbotHandler':
            """创建 SDK 需要的处理器实例"""
            return self._ChatbotHandler(self)

    def _parse_stream_message(self, incoming: Any, raw_data: dict) -> Optional[BotMessage]:
        """
        解析 Stream 消息为统一格式
        
        Args:
            incoming: ChatbotMessage 对象
            raw_data: 原始回调数据
        """
        try:
            raw_data = dict(raw_data or {})

            # 获取消息内容
            raw_content = incoming.text.content if incoming.text else ''

            # 提取命令（去除 @机器人）
            content = self._extract_command(raw_content)

            # 会话类型
            conversation_type = getattr(incoming, 'conversation_type', None)
            if conversation_type == '1':
                chat_type = ChatType.PRIVATE
            elif conversation_type == '2':
                chat_type = ChatType.GROUP
            else:
                chat_type = ChatType.UNKNOWN

            # 是否 @了机器人（Stream 模式下收到的消息一般都是 @机器人的）
            mentioned = True

            # 提取 sessionWebhook，便于异步推送
            session_webhook = (
                    getattr(incoming, 'session_webhook', None)
                    or raw_data.get('sessionWebhook')
                    or raw_data.get('session_webhook')
            )
            if session_webhook:
                raw_data['_session_webhook'] = session_webhook

            return BotMessage(
                platform='dingtalk',
                message_id=getattr(incoming, 'msg_id', '') or '',
                user_id=getattr(incoming, 'sender_id', '') or '',
                user_name=getattr(incoming, 'sender_nick', '') or '',
                chat_id=getattr(incoming, 'conversation_id', '') or '',
                chat_type=chat_type,
                content=content,
                raw_content=raw_content,
                mentioned=mentioned,
                mentions=[],
                timestamp=datetime.now(),
                raw_data=raw_data,
            )

        except Exception as e:
            logger.error(f"[DingTalk Stream] 解析消息失败: {e}")
            return None

    def _extract_command(self, text: str) -> str:
        """提取命令内容（去除 @机器人）"""
        import re
        text = re.sub(r'^@[\S]+\s*', '', text.strip())
        return text.strip()


class DingtalkStreamClient:
    """
    钉钉 Stream 模式客户端
    
    封装 dingtalk-stream SDK，提供简单的启动接口。
    
    使用方式：
        client = DingtalkStreamClient()
        client.start()  # 阻塞运行
        
        # 或者在后台运行
        client.start_background()
    """

    def __init__(
            self,
            client_id: Optional[str] = None,
            client_secret: Optional[str] = None
    ):
        """
        Args:
            client_id: 应用 AppKey（不传则从配置读取）
            client_secret: 应用 AppSecret（不传则从配置读取）
        """
        if not DINGTALK_STREAM_AVAILABLE:
            raise ImportError(
                "dingtalk-stream SDK 未安装。\n"
                "请运行: pip install dingtalk-stream"
            )

        from src.config import get_config
        config = get_config()

        self._client_id = client_id or getattr(config, 'dingtalk_app_key', None)
        self._client_secret = client_secret or getattr(config, 'dingtalk_app_secret', None)

        if not self._client_id or not self._client_secret:
            raise ValueError(
                "钉钉 Stream 模式需要配置 DINGTALK_APP_KEY 和 DINGTALK_APP_SECRET"
            )

        self._client: Optional[dingtalk_stream.DingTalkStreamClient] = None
        self._background_thread: Optional[threading.Thread] = None
        self._running = False

    def _create_message_handler(self) -> Callable[[BotMessage], BotResponse]:
        """创建消息处理函数"""

        def handle_message(message: BotMessage) -> BotResponse:
            from bot.dispatcher import get_dispatcher
            dispatcher = get_dispatcher()
            return dispatcher.dispatch(message)

        return handle_message

    def start(self) -> None:
        """
        启动 Stream 客户端（阻塞）
        
        此方法会阻塞当前线程，直到客户端停止。
        """
        logger.info("[DingTalk Stream] 正在启动...")

        # 创建凭证
        credential = dingtalk_stream.Credential(
            self._client_id,
            self._client_secret
        )

        # 创建客户端
        self._client = dingtalk_stream.DingTalkStreamClient(credential)

        # 注册消息处理器
        handler = DingtalkStreamHandler(self._create_message_handler())
        self._client.register_callback_handler(
            dingtalk_stream.chatbot.ChatbotMessage.TOPIC,
            handler.create_handler()
        )

        self._running = True
        logger.info("[DingTalk Stream] 客户端已启动，等待消息...")

        # 启动（阻塞）
        self._client.start_forever()

    def start_background(self) -> None:
        """
        在后台线程启动 Stream 客户端（非阻塞）
        
        适用于与其他服务（如 WebUI）同时运行的场景。
        """
        if self._background_thread and self._background_thread.is_alive():
            logger.warning("[DingTalk Stream] 客户端已在运行")
            return

        self._running = True
        self._background_thread = threading.Thread(
            target=self._run_in_background,
            daemon=True,
            name="DingtalkStreamClient"
        )
        self._background_thread.start()
        logger.info("[DingTalk Stream] 后台客户端已启动")

    def _run_in_background(self) -> None:
        """后台运行（处理异常和重连）"""
        while self._running:
            try:
                self.start()
            except Exception as e:
                logger.error(f"[DingTalk Stream] 运行异常: {e}")
                if self._running:
                    logger.info("[DingTalk Stream] 5 秒后重连...")
                    import time
                    time.sleep(5)

    def stop(self) -> None:
        """停止客户端"""
        self._running = False
        logger.info("[DingTalk Stream] 客户端已停止")

    @property
    def is_running(self) -> bool:
        """是否正在运行"""
        return self._running


# 全局客户端实例
_stream_client: Optional[DingtalkStreamClient] = None


def get_dingtalk_stream_client() -> Optional[DingtalkStreamClient]:
    """获取全局 Stream 客户端实例"""
    global _stream_client

    if _stream_client is None and DINGTALK_STREAM_AVAILABLE:
        try:
            _stream_client = DingtalkStreamClient()
        except (ImportError, ValueError) as e:
            logger.warning(f"[DingTalk Stream] 无法创建客户端: {e}")
            return None

    return _stream_client


def start_dingtalk_stream_background() -> bool:
    """
    在后台启动钉钉 Stream 客户端
    
    Returns:
        是否成功启动
    """
    client = get_dingtalk_stream_client()
    if client:
        client.start_background()
        return True
    return False