AstrBot 3.5.6__py3-none-any.whl → 4.7.0__py3-none-any.whl

astrbot/core/provider/provider.py

@@ -1,161 +1,170 @@
 import abc
-from typing import List
-from astrbot.core.db import BaseDatabase
-from typing import TypedDict, AsyncGenerator
-from astrbot.core.provider.func_tool_manager import FuncCall
-from astrbot.core.provider.entities import LLMResponse, ToolCallsResult
-from dataclasses import dataclass
+import asyncio
+from collections.abc import AsyncGenerator
 
-
-class Personality(TypedDict):
-    prompt: str = ""
-    name: str = ""
-    begin_dialogs: List[str] = []
-    mood_imitation_dialogs: List[str] = []
-
-    # cache
-    _begin_dialogs_processed: List[dict] = []
-    _mood_imitation_dialogs_processed: str = ""
-
-
-@dataclass
-class ProviderMeta:
-    id: str
-    model: str
-    type: str
+from astrbot.core.agent.message import Message
+from astrbot.core.agent.tool import ToolSet
+from astrbot.core.provider.entities import (
+    LLMResponse,
+    ProviderMeta,
+    RerankResult,
+    ToolCallsResult,
+)
+from astrbot.core.provider.register import provider_cls_map
 
 
 class AbstractProvider(abc.ABC):
+    """Provider Abstract Class"""
+
     def __init__(self, provider_config: dict) -> None:
         super().__init__()
         self.model_name = ""
         self.provider_config = provider_config
 
     def set_model(self, model_name: str):
-        """设置当前使用的模型名称"""
+        """Set the current model name"""
         self.model_name = model_name
 
     def get_model(self) -> str:
-        """获得当前使用的模型名称"""
+        """Get the current model name"""
         return self.model_name
 
     def meta(self) -> ProviderMeta:
-        """获取 Provider 的元数据"""
-        return ProviderMeta(
-            id=self.provider_config["id"],
+        """Get the provider metadata"""
+        provider_type_name = self.provider_config["type"]
+        meta_data = provider_cls_map.get(provider_type_name)
+        if not meta_data:
+            raise ValueError(f"Provider type {provider_type_name} not registered")
+        meta = ProviderMeta(
+            id=self.provider_config.get("id", "default"),
             model=self.get_model(),
-            type=self.provider_config["type"],
+            type=provider_type_name,
+            provider_type=meta_data.provider_type,
         )
+        return meta
 
 
 class Provider(AbstractProvider):
+    """Chat Provider"""
+
     def __init__(
         self,
         provider_config: dict,
         provider_settings: dict,
-        persistant_history: bool = True,
-        db_helper: BaseDatabase = None,
-        default_persona: Personality = None,
     ) -> None:
         super().__init__(provider_config)
-
         self.provider_settings = provider_settings
 
-        self.curr_personality: Personality = default_persona
-        """维护了当前的使用的 persona，即人格。可能为 None"""
-
     @abc.abstractmethod
     def get_current_key(self) -> str:
-        raise NotImplementedError()
+        raise NotImplementedError
 
-    def get_keys(self) -> List[str]:
+    def get_keys(self) -> list[str]:
         """获得提供商 Key"""
-        return self.provider_config.get("key", [])
+        keys = self.provider_config.get("key", [""])
+        return keys or [""]
 
     @abc.abstractmethod
     def set_key(self, key: str):
-        raise NotImplementedError()
+        raise NotImplementedError
 
     @abc.abstractmethod
-    def get_models(self) -> List[str]:
+    async def get_models(self) -> list[str]:
         """获得支持的模型列表"""
-        raise NotImplementedError()
+        raise NotImplementedError
 
     @abc.abstractmethod
     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,
-        tool_calls_result: ToolCallsResult = None,
+        prompt: str | None = None,
+        session_id: str | None = None,
+        image_urls: list[str] | None = None,
+        func_tool: ToolSet | None = None,
+        contexts: list[Message] | list[dict] | None = None,
+        system_prompt: str | None = None,
+        tool_calls_result: ToolCallsResult | list[ToolCallsResult] | None = None,
+        model: str | None = None,
         **kwargs,
     ) -> LLMResponse:
         """获得 LLM 的文本对话结果。会使用当前的模型进行对话。
 
         Args:
-            prompt: 提示词
+            prompt: 提示词，和 contexts 二选一使用，如果都指定，则会将 prompt（以及可能的 image_urls） 作为最新的一条记录添加到 contexts 中
             session_id: 会话 ID(此属性已经被废弃)
             image_urls: 图片 URL 列表
-            tools: Function-calling 工具
-            contexts: 上下文
+            tools: tool set
+            contexts: 上下文，和 prompt 二选一使用
             tool_calls_result: 回传给 LLM 的工具调用结果。参考: https://platform.openai.com/docs/guides/function-calling
             kwargs: 其他参数
 
         Notes:
             - 如果传入了 image_urls，将会在对话时附上图片。如果模型不支持图片输入，将会抛出错误。
             - 如果传入了 tools，将会使用 tools 进行 Function-calling。如果模型不支持 Function-calling，将会抛出错误。
+
         """
         ...
 
     async def text_chat_stream(
         self,
-        prompt: str,
-        session_id: str = None,
-        image_urls: List[str] = None,
-        func_tool: FuncCall = None,
-        contexts: List = None,
-        system_prompt: str = None,
-        tool_calls_result: ToolCallsResult = None,
+        prompt: str | None = None,
+        session_id: str | None = None,
+        image_urls: list[str] | None = None,
+        func_tool: ToolSet | None = None,
+        contexts: list[Message] | list[dict] | None = None,
+        system_prompt: str | None = None,
+        tool_calls_result: ToolCallsResult | list[ToolCallsResult] | None = None,
+        model: str | None = None,
         **kwargs,
     ) -> AsyncGenerator[LLMResponse, None]:
         """获得 LLM 的流式文本对话结果。会使用当前的模型进行对话。在生成的最后会返回一次完整的结果。
 
         Args:
-            prompt: 提示词
+            prompt: 提示词，和 contexts 二选一使用，如果都指定，则会将 prompt（以及可能的 image_urls） 作为最新的一条记录添加到 contexts 中
             session_id: 会话 ID(此属性已经被废弃)
             image_urls: 图片 URL 列表
-            tools: Function-calling 工具
-            contexts: 上下文
+            tools: tool set
+            contexts: 上下文，和 prompt 二选一使用
             tool_calls_result: 回传给 LLM 的工具调用结果。参考: https://platform.openai.com/docs/guides/function-calling
             kwargs: 其他参数
 
         Notes:
             - 如果传入了 image_urls，将会在对话时附上图片。如果模型不支持图片输入，将会抛出错误。
             - 如果传入了 tools，将会使用 tools 进行 Function-calling。如果模型不支持 Function-calling，将会抛出错误。
+
         """
         ...
 
-    async def pop_record(self, context: List):
-        """
-        弹出 context 第一条非系统提示词对话记录
-        """
+    async def pop_record(self, context: list):
+        """弹出 context 第一条非系统提示词对话记录"""
         poped = 0
         indexs_to_pop = []
         for idx, record in enumerate(context):
             if record["role"] == "system":
                 continue
-            else:
-                indexs_to_pop.append(idx)
-                poped += 1
-                if poped == 2:
-                    break
+            indexs_to_pop.append(idx)
+            poped += 1
+            if poped == 2:
+                break
 
         for idx in reversed(indexs_to_pop):
             context.pop(idx)
 
+    def _ensure_message_to_dicts(
+        self,
+        messages: list[dict] | list[Message] | None,
+    ) -> list[dict]:
+        """Convert a list of Message objects to a list of dictionaries."""
+        if not messages:
+            return []
+        dicts: list[dict] = []
+        for message in messages:
+            if isinstance(message, Message):
+                dicts.append(message.model_dump())
+            else:
+                dicts.append(message)
+
+        return dicts
+
 
 class STTProvider(AbstractProvider):
     def __init__(self, provider_config: dict, provider_settings: dict) -> None:
@@ -166,7 +175,7 @@ class STTProvider(AbstractProvider):
     @abc.abstractmethod
     async def get_text(self, audio_url: str) -> str:
         """获取音频的文本"""
-        raise NotImplementedError()
+        raise NotImplementedError
 
 
 class TTSProvider(AbstractProvider):
@@ -178,4 +187,110 @@ class TTSProvider(AbstractProvider):
     @abc.abstractmethod
     async def get_audio(self, text: str) -> str:
         """获取文本的音频，返回音频文件路径"""
-        raise NotImplementedError()
+        raise NotImplementedError
+
+
+class EmbeddingProvider(AbstractProvider):
+    def __init__(self, provider_config: dict, provider_settings: dict) -> None:
+        super().__init__(provider_config)
+        self.provider_config = provider_config
+        self.provider_settings = provider_settings
+
+    @abc.abstractmethod
+    async def get_embedding(self, text: str) -> list[float]:
+        """获取文本的向量"""
+        ...
+
+    @abc.abstractmethod
+    async def get_embeddings(self, text: list[str]) -> list[list[float]]:
+        """批量获取文本的向量"""
+        ...
+
+    @abc.abstractmethod
+    def get_dim(self) -> int:
+        """获取向量的维度"""
+        ...
+
+    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} 次: {e!s}",
+                            )
+                        # 等待一段时间后重试，使用指数退避
+                        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:
+        super().__init__(provider_config)
+        self.provider_config = provider_config
+        self.provider_settings = provider_settings
+
+    @abc.abstractmethod
+    async def rerank(
+        self,
+        query: str,
+        documents: list[str],
+        top_n: int | None = None,
+    ) -> list[RerankResult]:
+        """获取查询和文档的重排序分数"""
+        ...

astrbot/core/provider/register.py

@@ -1,11 +1,11 @@
-from typing import List, Dict
-from .entities import ProviderMetaData, ProviderType
 from astrbot.core import logger
+
+from .entities import ProviderMetaData, ProviderType
 from .func_tool_manager import FuncCall
 
-provider_registry: List[ProviderMetaData] = []
+provider_registry: list[ProviderMetaData] = []
 """维护了通过装饰器注册的 Provider"""
-provider_cls_map: Dict[str, ProviderMetaData] = {}
+provider_cls_map: dict[str, ProviderMetaData] = {}
 """维护了 Provider 类型名称和 ProviderMetadata 的映射"""
 
 llm_tools = FuncCall()
@@ -15,15 +15,15 @@ def register_provider_adapter(
     provider_type_name: str,
     desc: str,
     provider_type: ProviderType = ProviderType.CHAT_COMPLETION,
-    default_config_tmpl: dict = None,
-    provider_display_name: str = None,
+    default_config_tmpl: dict | None = None,
+    provider_display_name: str | None = None,
 ):
     """用于注册平台适配器的带参装饰器"""
 
     def decorator(cls):
         if provider_type_name in provider_cls_map:
             raise ValueError(
-                f"检测到大模型提供商适配器 {provider_type_name} 已经注册，可能发生了大模型提供商适配器类型命名冲突。"
+                f"检测到大模型提供商适配器 {provider_type_name} 已经注册，可能发生了大模型提供商适配器类型命名冲突。",
             )
 
         # 添加必备选项
@@ -36,6 +36,8 @@ def register_provider_adapter(
                 default_config_tmpl["id"] = provider_type_name
 
         pm = ProviderMetaData(
+            id="default",  # will be replaced when instantiated
+            model=None,
             type=provider_type_name,
             desc=desc,
             provider_type=provider_type,