AstrBot 4.3.3__py3-none-any.whl → 4.5.0__py3-none-any.whl

astrbot/core/platform/sources/wecom_ai_bot/wecomai_server.py

@@ -0,0 +1,166 @@
+"""
+企业微信智能机器人 HTTP 服务器
+处理企业微信智能机器人的 HTTP 回调请求
+"""
+
+import asyncio
+from typing import Dict, Any, Optional, Callable
+
+import quart
+from astrbot.api import logger
+
+from .wecomai_api import WecomAIBotAPIClient
+from .wecomai_utils import WecomAIBotConstants
+
+
+class WecomAIBotServer:
+    """企业微信智能机器人 HTTP 服务器"""
+
+    def __init__(
+        self,
+        host: str,
+        port: int,
+        api_client: WecomAIBotAPIClient,
+        message_handler: Optional[
+            Callable[[Dict[str, Any], Dict[str, str]], Any]
+        ] = None,
+    ):
+        """初始化服务器
+
+        Args:
+            host: 监听地址
+            port: 监听端口
+            api_client: API客户端实例
+            message_handler: 消息处理回调函数
+        """
+        self.host = host
+        self.port = port
+        self.api_client = api_client
+        self.message_handler = message_handler
+
+        self.app = quart.Quart(__name__)
+        self._setup_routes()
+
+        self.shutdown_event = asyncio.Event()
+
+    def _setup_routes(self):
+        """设置 Quart 路由"""
+
+        # 使用 Quart 的 add_url_rule 方法添加路由
+        self.app.add_url_rule(
+            "/webhook/wecom-ai-bot",
+            view_func=self.verify_url,
+            methods=["GET"],
+        )
+
+        self.app.add_url_rule(
+            "/webhook/wecom-ai-bot",
+            view_func=self.handle_message,
+            methods=["POST"],
+        )
+
+    async def verify_url(self):
+        """验证回调 URL"""
+        args = quart.request.args
+        msg_signature = args.get("msg_signature")
+        timestamp = args.get("timestamp")
+        nonce = args.get("nonce")
+        echostr = args.get("echostr")
+
+        if not all([msg_signature, timestamp, nonce, echostr]):
+            logger.error("URL 验证参数缺失")
+            return "verify fail", 400
+
+        # 类型检查确保不为 None
+        assert msg_signature is not None
+        assert timestamp is not None
+        assert nonce is not None
+        assert echostr is not None
+
+        logger.info("收到企业微信智能机器人 WebHook URL 验证请求。")
+        result = self.api_client.verify_url(msg_signature, timestamp, nonce, echostr)
+        return result, 200, {"Content-Type": "text/plain"}
+
+    async def handle_message(self):
+        """处理消息回调"""
+        args = quart.request.args
+        msg_signature = args.get("msg_signature")
+        timestamp = args.get("timestamp")
+        nonce = args.get("nonce")
+
+        if not all([msg_signature, timestamp, nonce]):
+            logger.error("消息回调参数缺失")
+            return "缺少必要参数", 400
+
+        # 类型检查确保不为 None
+        assert msg_signature is not None
+        assert timestamp is not None
+        assert nonce is not None
+
+        logger.debug(
+            f"收到消息回调，msg_signature={msg_signature}, timestamp={timestamp}, nonce={nonce}"
+        )
+
+        try:
+            # 获取请求体
+            post_data = await quart.request.get_data()
+
+            # 确保 post_data 是 bytes 类型
+            if isinstance(post_data, str):
+                post_data = post_data.encode("utf-8")
+
+            # 解密消息
+            ret_code, message_data = await self.api_client.decrypt_message(
+                post_data, msg_signature, timestamp, nonce
+            )
+
+            if ret_code != WecomAIBotConstants.SUCCESS or not message_data:
+                logger.error("消息解密失败，错误码: %d", ret_code)
+                return "消息解密失败", 400
+
+            # 调用消息处理器
+            response = None
+            if self.message_handler:
+                try:
+                    response = await self.message_handler(
+                        message_data, {"nonce": nonce, "timestamp": timestamp}
+                    )
+                except Exception as e:
+                    logger.error("消息处理器执行异常: %s", e)
+                    return "消息处理异常", 500
+
+            if response:
+                return response, 200, {"Content-Type": "text/plain"}
+            else:
+                return "success", 200, {"Content-Type": "text/plain"}
+
+        except Exception as e:
+            logger.error("处理消息时发生异常: %s", e)
+            return "内部服务器错误", 500
+
+    async def start_server(self):
+        """启动服务器"""
+        logger.info("启动企业微信智能机器人服务器，监听 %s:%d", self.host, self.port)
+
+        try:
+            await self.app.run_task(
+                host=self.host,
+                port=self.port,
+                shutdown_trigger=self.shutdown_trigger,
+            )
+        except Exception as e:
+            logger.error("服务器运行异常: %s", e)
+            raise
+
+    async def shutdown_trigger(self):
+        """关闭触发器"""
+        await self.shutdown_event.wait()
+
+    async def shutdown(self):
+        """关闭服务器"""
+        logger.info("企业微信智能机器人服务器正在关闭...")
+        self.shutdown_event.set()
+
+    def get_app(self):
+        """获取 Quart 应用实例"""
+        return self.app

astrbot/core/platform/sources/wecom_ai_bot/wecomai_utils.py

@@ -0,0 +1,199 @@
+"""
+企业微信智能机器人工具模块
+提供常量定义、工具函数和辅助方法
+"""
+
+import string
+import random
+import hashlib
+import base64
+import aiohttp
+import asyncio
+from Crypto.Cipher import AES
+from typing import Any, Tuple
+from astrbot.api import logger
+
+
+# 常量定义
+class WecomAIBotConstants:
+    """企业微信智能机器人常量"""
+
+    # 消息类型
+    MSG_TYPE_TEXT = "text"
+    MSG_TYPE_IMAGE = "image"
+    MSG_TYPE_MIXED = "mixed"
+    MSG_TYPE_STREAM = "stream"
+    MSG_TYPE_EVENT = "event"
+
+    # 流消息状态
+    STREAM_CONTINUE = False
+    STREAM_FINISH = True
+
+    # 错误码
+    SUCCESS = 0
+    DECRYPT_ERROR = -40001
+    VALIDATE_SIGNATURE_ERROR = -40002
+    PARSE_XML_ERROR = -40003
+    COMPUTE_SIGNATURE_ERROR = -40004
+    ILLEGAL_AES_KEY = -40005
+    VALIDATE_APPID_ERROR = -40006
+    ENCRYPT_AES_ERROR = -40007
+    ILLEGAL_BUFFER = -40008
+
+
+def generate_random_string(length: int = 10) -> str:
+    """生成随机字符串
+
+    Args:
+        length: 字符串长度，默认为 10
+
+    Returns:
+        随机字符串
+    """
+    letters = string.ascii_letters + string.digits
+    return "".join(random.choice(letters) for _ in range(length))
+
+
+def calculate_image_md5(image_data: bytes) -> str:
+    """计算图片数据的 MD5 值
+
+    Args:
+        image_data: 图片二进制数据
+
+    Returns:
+        MD5 哈希值（十六进制字符串）
+    """
+    return hashlib.md5(image_data).hexdigest()
+
+
+def encode_image_base64(image_data: bytes) -> str:
+    """将图片数据编码为 Base64
+
+    Args:
+        image_data: 图片二进制数据
+
+    Returns:
+        Base64 编码的字符串
+    """
+    return base64.b64encode(image_data).decode("utf-8")
+
+
+def format_session_id(session_type: str, session_id: str) -> str:
+    """格式化会话 ID
+
+    Args:
+        session_type: 会话类型 ("user", "group")
+        session_id: 原始会话 ID
+
+    Returns:
+        格式化后的会话 ID
+    """
+    return f"wecom_ai_bot_{session_type}_{session_id}"
+
+
+def parse_session_id(formatted_session_id: str) -> Tuple[str, str]:
+    """解析格式化的会话 ID
+
+    Args:
+        formatted_session_id: 格式化的会话 ID
+
+    Returns:
+        (会话类型, 原始会话ID)
+    """
+    parts = formatted_session_id.split("_", 3)
+    if (
+        len(parts) >= 4
+        and parts[0] == "wecom"
+        and parts[1] == "ai"
+        and parts[2] == "bot"
+    ):
+        return parts[3], "_".join(parts[4:]) if len(parts) > 4 else ""
+    return "user", formatted_session_id
+
+
+def safe_json_loads(json_str: str, default: Any = None) -> Any:
+    """安全地解析 JSON 字符串
+
+    Args:
+        json_str: JSON 字符串
+        default: 解析失败时的默认值
+
+    Returns:
+        解析结果或默认值
+    """
+    import json
+
+    try:
+        return json.loads(json_str)
+    except (json.JSONDecodeError, TypeError) as e:
+        logger.warning(f"JSON 解析失败: {e}, 原始字符串: {json_str}")
+        return default
+
+
+def format_error_response(error_code: int, error_msg: str) -> str:
+    """格式化错误响应
+
+    Args:
+        error_code: 错误码
+        error_msg: 错误信息
+
+    Returns:
+        格式化的错误响应字符串
+    """
+    return f"Error {error_code}: {error_msg}"
+
+
+async def process_encrypted_image(
+    image_url: str, aes_key_base64: str
+) -> Tuple[bool, str]:
+    """下载并解密加密图片
+
+    Args:
+        image_url: 加密图片的URL
+        aes_key_base64: Base64编码的AES密钥(与回调加解密相同)
+
+    Returns:
+        Tuple[bool, str]: status 为 True 时 data 是解密后的图片数据的 base64 编码，
+            status 为 False 时 data 是错误信息
+    """
+    # 1. 下载加密图片
+    logger.info("开始下载加密图片: %s", image_url)
+    try:
+        async with aiohttp.ClientSession() as session:
+            async with session.get(image_url, timeout=15) as response:
+                response.raise_for_status()
+                encrypted_data = await response.read()
+        logger.info("图片下载成功，大小: %d 字节", len(encrypted_data))
+    except (aiohttp.ClientError, asyncio.TimeoutError) as e:
+        error_msg = f"下载图片失败: {str(e)}"
+        logger.error(error_msg)
+        return False, error_msg
+
+    # 2. 准备AES密钥和IV
+    if not aes_key_base64:
+        raise ValueError("AES密钥不能为空")
+
+    # Base64解码密钥 (自动处理填充)
+    aes_key = base64.b64decode(aes_key_base64 + "=" * (-len(aes_key_base64) % 4))
+    if len(aes_key) != 32:
+        raise ValueError("无效的AES密钥长度: 应为32字节")
+
+    iv = aes_key[:16]  # 初始向量为密钥前16字节
+
+    # 3. 解密图片数据
+    cipher = AES.new(aes_key, AES.MODE_CBC, iv)
+    decrypted_data = cipher.decrypt(encrypted_data)
+
+    # 4. 去除PKCS#7填充 (Python 3兼容写法)
+    pad_len = decrypted_data[-1]  # 直接获取最后一个字节的整数值
+    if pad_len > 32:  # AES-256块大小为32字节
+        raise ValueError("无效的填充长度 (大于32字节)")
+
+    decrypted_data = decrypted_data[:-pad_len]
+    logger.info("图片解密成功，解密后大小: %d 字节", len(decrypted_data))
+
+    # 5. 转换为base64编码
+    base64_data = base64.b64encode(decrypted_data).decode("utf-8")
+    logger.info("图片已转换为base64编码，编码后长度: %d", len(base64_data))
+
+    return True, base64_data

astrbot/core/provider/manager.py

@@ -1,6 +1,5 @@
 import asyncio
 import traceback
-from typing import List
 
 from astrbot.core import logger, sp
 from astrbot.core.astrbot_config_mgr import AstrBotConfigManager
@@ -28,7 +27,7 @@ class ProviderManager:
         self.persona_mgr = persona_mgr
         self.acm = acm
         config = acm.confs["default"]
-        self.providers_config: List = config["provider"]
+        self.providers_config: list = config["provider"]
         self.provider_settings: dict = config["provider_settings"]
         self.provider_stt_settings: dict = config.get("provider_stt_settings", {})
         self.provider_tts_settings: dict = config.get("provider_tts_settings", {})
@@ -36,15 +35,15 @@ class ProviderManager:
         # 人格相关属性，v4.0.0 版本后被废弃，推荐使用 PersonaManager
         self.default_persona_name = persona_mgr.default_persona
 
-        self.provider_insts: List[Provider] = []
+        self.provider_insts: list[Provider] = []
         """加载的 Provider 的实例"""
-        self.stt_provider_insts: List[STTProvider] = []
+        self.stt_provider_insts: list[STTProvider] = []
         """加载的 Speech To Text Provider 的实例"""
-        self.tts_provider_insts: List[TTSProvider] = []
+        self.tts_provider_insts: list[TTSProvider] = []
         """加载的 Text To Speech Provider 的实例"""
-        self.embedding_provider_insts: List[EmbeddingProvider] = []
+        self.embedding_provider_insts: list[EmbeddingProvider] = []
         """加载的 Embedding Provider 的实例"""
-        self.rerank_provider_insts: List[RerankProvider] = []
+        self.rerank_provider_insts: list[RerankProvider] = []
         """加载的 Rerank Provider 的实例"""
         self.inst_map: dict[
             str,
@@ -175,7 +174,11 @@ class ProviderManager:
     async def initialize(self):
         # 逐个初始化提供商
         for provider_config in self.providers_config:
-            await self.load_provider(provider_config)
+            try:
+                await self.load_provider(provider_config)
+            except Exception as e:
+                logger.error(traceback.format_exc())
+                logger.error(e)
 
         # 设置默认提供商
         selected_provider_id = sp.get(
@@ -404,10 +407,12 @@ class ProviderManager:
 
             self.inst_map[provider_config["id"]] = inst
         except Exception as e:
-            logger.error(traceback.format_exc())
             logger.error(
                 f"实例化 {provider_config['type']}({provider_config['id']}) 提供商适配器失败：{e}"
             )
+            raise Exception(
+                f"实例化 {provider_config['type']}({provider_config['id']}) 提供商适配器失败：{e}"
+            )
 
     async def reload(self, provider_config: dict):
         await self.terminate_provider(provider_config["id"])

astrbot/core/provider/provider.py

@@ -1,4 +1,5 @@
 import abc
+import asyncio
 from typing import List
 from typing import AsyncGenerator
 from astrbot.core.agent.tool import ToolSet
@@ -203,6 +204,72 @@ class EmbeddingProvider(AbstractProvider):
         """获取向量的维度"""
         ...
 
+    async def get_embeddings_batch(
+        self,
+        texts: list[str],
+        batch_size: int = 16,
+        tasks_limit: int = 3,
+        max_retries: int = 3,
+        progress_callback=None,
+    ) -> list[list[float]]:
+        """批量获取文本的向量，分批处理以节省内存
+
+        Args:
+            texts: 文本列表
+            batch_size: 每批处理的文本数量
+            tasks_limit: 并发任务数量限制
+            max_retries: 失败时的最大重试次数
+            progress_callback: 进度回调函数，接收参数 (current, total)
+
+        Returns:
+            向量列表
+        """
+        semaphore = asyncio.Semaphore(tasks_limit)
+        all_embeddings: list[list[float]] = []
+        failed_batches: list[tuple[int, list[str]]] = []
+        completed_count = 0
+        total_count = len(texts)
+
+        async def process_batch(batch_idx: int, batch_texts: list[str]):
+            nonlocal completed_count
+            async with semaphore:
+                for attempt in range(max_retries):
+                    try:
+                        batch_embeddings = await self.get_embeddings(batch_texts)
+                        all_embeddings.extend(batch_embeddings)
+                        completed_count += len(batch_texts)
+                        if progress_callback:
+                            await progress_callback(completed_count, total_count)
+                        return
+                    except Exception as e:
+                        if attempt == max_retries - 1:
+                            # 最后一次重试失败，记录失败的批次
+                            failed_batches.append((batch_idx, batch_texts))
+                            raise Exception(
+                                f"批次 {batch_idx} 处理失败，已重试 {max_retries} 次: {str(e)}"
+                            )
+                        # 等待一段时间后重试，使用指数退避
+                        await asyncio.sleep(2**attempt)
+
+        tasks = []
+        for i in range(0, len(texts), batch_size):
+            batch_texts = texts[i : i + batch_size]
+            batch_idx = i // batch_size
+            tasks.append(process_batch(batch_idx, batch_texts))
+
+        # 收集所有任务的结果，包括失败的任务
+        results = await asyncio.gather(*tasks, return_exceptions=True)
+
+        # 检查是否有失败的任务
+        errors = [r for r in results if isinstance(r, Exception)]
+        if errors:
+            error_msg = (
+                f"有 {len(errors)} 个批次处理失败: {'; '.join(str(e) for e in errors)}"
+            )
+            raise Exception(error_msg)
+
+        return all_embeddings
+
 
 class RerankProvider(AbstractProvider):
     def __init__(self, provider_config: dict, provider_settings: dict) -> None:

astrbot/core/provider/sources/anthropic_source.py

@@ -10,7 +10,7 @@ from anthropic.types import Message
 from astrbot.core.utils.io import download_image_by_url
 from astrbot.api.provider import Provider
 from astrbot import logger
-from astrbot.core.provider.func_tool_manager import FuncCall
+from astrbot.core.provider.func_tool_manager import ToolSet
 from ..register import register_provider_adapter
 from astrbot.core.provider.entities import LLMResponse
 from typing import AsyncGenerator
@@ -104,7 +104,7 @@ class ProviderAnthropic(Provider):
 
         return system_prompt, new_messages
 
-    async def _query(self, payloads: dict, tools: FuncCall) -> LLMResponse:
+    async def _query(self, payloads: dict, tools: ToolSet | None) -> LLMResponse:
         if tools:
             if tool_list := tools.get_func_desc_anthropic_style():
                 payloads["tools"] = tool_list
@@ -135,7 +135,7 @@ class ProviderAnthropic(Provider):
         return llm_response
 
     async def _query_stream(
-        self, payloads: dict, tools: FuncCall
+        self, payloads: dict, tools: ToolSet | None
     ) -> AsyncGenerator[LLMResponse, None]:
         if tools:
             if tool_list := tools.get_func_desc_anthropic_style():
@@ -326,7 +326,7 @@ class ProviderAnthropic(Provider):
         async for llm_response in self._query_stream(payloads, func_tool):
             yield llm_response
 
-    async def assemble_context(self, text: str, image_urls: List[str] = None):
+    async def assemble_context(self, text: str, image_urls: List[str] | None = None):
         """组装上下文，支持文本和图片"""
         if not image_urls:
             return {"role": "user", "content": text}

astrbot/core/provider/sources/dashscope_source.py

@@ -1,15 +1,14 @@
 import re
 import asyncio
 import functools
-from typing import List
 from .. import Provider, Personality
 from ..entities import LLMResponse
-from ..func_tool_manager import FuncCall
 from ..register import register_provider_adapter
 from astrbot.core.message.message_event_result import MessageChain
 from .openai_source import ProviderOpenAIOfficial
 from astrbot.core import logger, sp
 from dashscope import Application
+from dashscope.app.application_response import ApplicationResponse
 
 
 @register_provider_adapter("dashscope", "Dashscope APP 适配器。")
@@ -62,11 +61,11 @@ class ProviderDashscope(ProviderOpenAIOfficial):
     async def text_chat(
         self,
         prompt: str,
-        session_id: str = None,
-        image_urls: List[str] = [],
-        func_tool: FuncCall = None,
-        contexts: List = None,
-        system_prompt: str = None,
+        session_id=None,
+        image_urls=[],
+        func_tool=None,
+        contexts=None,
+        system_prompt=None,
         model=None,
         **kwargs,
     ) -> LLMResponse:
@@ -122,6 +121,8 @@ class ProviderDashscope(ProviderOpenAIOfficial):
             )
             response = await asyncio.get_event_loop().run_in_executor(None, partial)
 
+        assert isinstance(response, ApplicationResponse)
+
         logger.debug(f"dashscope resp: {response}")
 
         if response.status_code != 200:
@@ -135,12 +136,12 @@ class ProviderDashscope(ProviderOpenAIOfficial):
                 ),
             )
 
-        output_text = response.output.get("text", "")
+        output_text = response.output.get("text", "") or ""
         # RAG 引用脚标格式化
         output_text = re.sub(r"<ref>\[(\d+)\]</ref>", r"[\1]", output_text)
         if self.output_reference and response.output.get("doc_references", None):
             ref_str = ""
-            for ref in response.output.get("doc_references", []):
+            for ref in response.output.get("doc_references", []) or []:
                 ref_title = (
                     ref.get("title", "")
                     if ref.get("title")

astrbot/core/provider/sources/dify_source.py

@@ -1,9 +1,7 @@
 import astrbot.core.message.components as Comp
 import os
-from typing import List
 from .. import Provider
 from ..entities import LLMResponse
-from ..func_tool_manager import FuncCall
 from ..register import register_provider_adapter
 from astrbot.core.utils.dify_api_client import DifyAPIClient
 from astrbot.core.utils.io import download_image_by_url, download_file
@@ -55,11 +53,11 @@ class ProviderDify(Provider):
     async def text_chat(
         self,
         prompt: str,
-        session_id: str = None,
-        image_urls: List[str] = None,
-        func_tool: FuncCall = None,
-        contexts: List = None,
-        system_prompt: str = None,
+        session_id=None,
+        image_urls=None,
+        func_tool=None,
+        contexts=None,
+        system_prompt=None,
         tool_calls_result=None,
         model=None,
         **kwargs,
@@ -223,7 +221,7 @@ class ProviderDify(Provider):
             # Chat
             return MessageChain(chain=[Comp.Plain(chunk)])
 
-        async def parse_file(item: dict) -> Comp:
+        async def parse_file(item: dict):
             match item["type"]:
                 case "image":
                     return Comp.Image(file=item["url"], url=item["url"])

astrbot/core/provider/sources/gemini_embedding_source.py

@@ -32,7 +32,6 @@ class GeminiEmbeddingProvider(EmbeddingProvider):
         self.model = provider_config.get(
             "embedding_model", "gemini-embedding-exp-03-07"
         )
-        self.dimension = provider_config.get("embedding_dimensions", 768)
 
     async def get_embedding(self, text: str) -> list[float]:
         """
@@ -60,4 +59,4 @@ class GeminiEmbeddingProvider(EmbeddingProvider):
 
     def get_dim(self) -> int:
         """获取向量的维度"""
-        return self.dimension
+        return self.provider_config.get("embedding_dimensions", 768)

astrbot/core/provider/sources/openai_embedding_source.py

@@ -22,7 +22,6 @@ class OpenAIEmbeddingProvider(EmbeddingProvider):
             timeout=int(provider_config.get("timeout", 20)),
         )
         self.model = provider_config.get("embedding_model", "text-embedding-3-small")
-        self.dimension = provider_config.get("embedding_dimensions", 1024)
 
     async def get_embedding(self, text: str) -> list[float]:
         """
@@ -40,4 +39,4 @@ class OpenAIEmbeddingProvider(EmbeddingProvider):
 
     def get_dim(self) -> int:
         """获取向量的维度"""
-        return self.dimension
+        return self.provider_config.get("embedding_dimensions", 1024)