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

astrbot/core/provider/entities.py

@@ -1,74 +1,60 @@
-import enum
 import base64
+import enum
 import json
-from astrbot.core.utils.io import download_image_by_url
-from astrbot import logger
 from dataclasses import dataclass, field
-from typing import List, Dict, Type
-from .func_tool_manager import FuncCall
+from typing import Any
+
+from anthropic.types import Message as AnthropicMessage
+from google.genai.types import GenerateContentResponse
 from openai.types.chat.chat_completion import ChatCompletion
-from openai.types.chat.chat_completion_message_tool_call import (
-    ChatCompletionMessageToolCall,
+
+import astrbot.core.message.components as Comp
+from astrbot import logger
+from astrbot.core.agent.message import (
+    AssistantMessageSegment,
+    ToolCall,
+    ToolCallMessageSegment,
 )
+from astrbot.core.agent.tool import ToolSet
 from astrbot.core.db.po import Conversation
 from astrbot.core.message.message_event_result import MessageChain
-import astrbot.core.message.components as Comp
+from astrbot.core.utils.io import download_image_by_url
 
 
 class ProviderType(enum.Enum):
     CHAT_COMPLETION = "chat_completion"
     SPEECH_TO_TEXT = "speech_to_text"
     TEXT_TO_SPEECH = "text_to_speech"
+    EMBEDDING = "embedding"
+    RERANK = "rerank"
 
 
 @dataclass
-class ProviderMetaData:
+class ProviderMeta:
+    """The basic metadata of a provider instance."""
+
+    id: str
+    """the unique id of the provider instance that user configured"""
+    model: str | None
+    """the model name of the provider instance currently used"""
     type: str
-    """提供商适配器名称，如 openai, ollama"""
-    desc: str = ""
-    """提供商适配器描述."""
+    """the name of the provider adapter, such as openai, ollama"""
     provider_type: ProviderType = ProviderType.CHAT_COMPLETION
-    cls_type: Type = None
-
-    default_config_tmpl: dict = None
-    """平台的默认配置模板"""
-    provider_display_name: str = None
-    """显示在 WebUI 配置页中的提供商名称，如空则是 type"""
+    """the capability type of the provider adapter"""
 
 
 @dataclass
-class ToolCallMessageSegment:
-    """OpenAI 格式的上下文中 role 为 tool 的消息段。参考: https://platform.openai.com/docs/guides/function-calling"""
-
-    tool_call_id: str
-    content: str
-    role: str = "tool"
+class ProviderMetaData(ProviderMeta):
+    """The metadata of a provider adapter for registration."""
 
-    def to_dict(self):
-        return {
-            "tool_call_id": self.tool_call_id,
-            "content": self.content,
-            "role": self.role,
-        }
-
-
-@dataclass
-class AssistantMessageSegment:
-    """OpenAI 格式的上下文中 role 为 assistant 的消息段。参考: https://platform.openai.com/docs/guides/function-calling"""
-
-    content: str = None
-    tool_calls: List[ChatCompletionMessageToolCall | Dict] = None
-    role: str = "assistant"
-
-    def to_dict(self):
-        ret = {
-            "role": self.role,
-        }
-        if self.content:
-            ret["content"] = self.content
-        elif self.tool_calls:
-            ret["tool_calls"] = self.tool_calls
-        return ret
+    desc: str = ""
+    """the short description of the provider adapter"""
+    cls_type: Any = None
+    """the class type of the provider adapter"""
+    default_config_tmpl: dict | None = None
+    """the default configuration template of the provider adapter"""
+    provider_display_name: str | None = None
+    """the display name of the provider shown in the WebUI configuration page; if empty, the type is used"""
 
 
 @dataclass
@@ -77,44 +63,70 @@ class ToolCallsResult:
 
     tool_calls_info: AssistantMessageSegment
     """函数调用的信息"""
-    tool_calls_result: List[ToolCallMessageSegment]
+    tool_calls_result: list[ToolCallMessageSegment]
     """函数调用的结果"""
 
-    def to_openai_messages(self) -> List[Dict]:
+    def to_openai_messages(self) -> list[dict]:
         ret = [
-            self.tool_calls_info.to_dict(),
-            *[item.to_dict() for item in self.tool_calls_result],
+            self.tool_calls_info.model_dump(),
+            *[item.model_dump() for item in self.tool_calls_result],
         ]
         return ret
 
+    def to_openai_messages_model(
+        self,
+    ) -> list[AssistantMessageSegment | ToolCallMessageSegment]:
+        return [
+            self.tool_calls_info,
+            *self.tool_calls_result,
+        ]
+
 
 @dataclass
 class ProviderRequest:
-    prompt: str
+    prompt: str | None = None
     """提示词"""
-    session_id: str = ""
+    session_id: str | None = ""
     """会话 ID"""
-    image_urls: List[str] = None
+    image_urls: list[str] = field(default_factory=list)
     """图片 URL 列表"""
-    func_tool: FuncCall = None
+    func_tool: ToolSet | None = None
     """可用的函数工具"""
-    contexts: List = None
-    """上下文。格式与 openai 的上下文格式一致：
+    contexts: list[dict] = field(default_factory=list)
+    """
+    OpenAI 格式上下文列表。
     参考 https://platform.openai.com/docs/api-reference/chat/create#chat-create-messages
     """
     system_prompt: str = ""
     """系统提示词"""
-    conversation: Conversation = None
-
-    tool_calls_result: ToolCallsResult = None
+    conversation: Conversation | None = None
+    """关联的对话对象"""
+    tool_calls_result: list[ToolCallsResult] | ToolCallsResult | None = None
     """附加的上次请求后工具调用的结果。参考: https://platform.openai.com/docs/guides/function-calling#handling-function-calls"""
+    model: str | None = None
+    """模型名称，为 None 时使用提供商的默认模型"""
 
     def __repr__(self):
-        return f"ProviderRequest(prompt={self.prompt}, session_id={self.session_id}, image_urls={self.image_urls}, func_tool={self.func_tool}, contexts={self._print_friendly_context()}, system_prompt={self.system_prompt.strip()}, tool_calls_result={self.tool_calls_result})"
+        return (
+            f"ProviderRequest(prompt={self.prompt}, session_id={self.session_id}, "
+            f"image_count={len(self.image_urls or [])}, "
+            f"func_tool={self.func_tool}, "
+            f"contexts={self._print_friendly_context()}, "
+            f"system_prompt={self.system_prompt}, "
+            f"conversation_id={self.conversation.cid if self.conversation else 'N/A'}, "
+        )
 
     def __str__(self):
         return self.__repr__()
 
+    def append_tool_calls_result(self, tool_calls_result: ToolCallsResult):
+        """添加工具调用结果到请求中"""
+        if not self.tool_calls_result:
+            self.tool_calls_result = []
+        if isinstance(self.tool_calls_result, ToolCallsResult):
+            self.tool_calls_result = [self.tool_calls_result]
+        self.tool_calls_result.append(tool_calls_result)
+
     def _print_friendly_context(self):
         """打印友好的消息上下文。将 image_url 的值替换为 <Image>"""
         if not self.contexts:
@@ -150,12 +162,14 @@ class ProviderRequest:
 
         return result_parts
 
-    async def assemble_context(self) -> Dict:
+    async def assemble_context(self) -> dict:
         """将请求(prompt 和 image_urls)包装成 OpenAI 的消息格式。"""
         if self.image_urls:
             user_content = {
                 "role": "user",
-                "content": [{"type": "text", "text": self.prompt if self.prompt else "[图片]"}],
+                "content": [
+                    {"type": "text", "text": self.prompt if self.prompt else "[图片]"},
+                ],
             }
             for image_url in self.image_urls:
                 if image_url.startswith("http"):
@@ -170,11 +184,10 @@ class ProviderRequest:
                     logger.warning(f"图片 {image_url} 得到的结果为空，将忽略。")
                     continue
                 user_content["content"].append(
-                    {"type": "image_url", "image_url": {"url": image_data}}
+                    {"type": "image_url", "image_url": {"url": image_data}},
                 )
             return user_content
-        else:
-            return {"role": "user", "content": self.prompt}
+        return {"role": "user", "content": self.prompt}
 
     async def _encode_image_bs64(self, image_url: str) -> str:
         """将图片转换为 base64"""
@@ -189,34 +202,44 @@ class ProviderRequest:
 @dataclass
 class LLMResponse:
     role: str
-    """角色, assistant, tool, err"""
-    result_chain: MessageChain = None
-    """返回的消息链"""
-    tools_call_args: List[Dict[str, any]] = field(default_factory=list)
-    """工具调用参数"""
-    tools_call_name: List[str] = field(default_factory=list)
-    """工具调用名称"""
-    tools_call_ids: List[str] = field(default_factory=list)
-    """工具调用 ID"""
-
-    raw_completion: ChatCompletion = None
-    _new_record: Dict[str, any] = None
+    """The role of the message, e.g., assistant, tool, err"""
+    result_chain: MessageChain | None = None
+    """A chain of message components representing the text completion from LLM."""
+    tools_call_args: list[dict[str, Any]] = field(default_factory=list)
+    """Tool call arguments."""
+    tools_call_name: list[str] = field(default_factory=list)
+    """Tool call names."""
+    tools_call_ids: list[str] = field(default_factory=list)
+    """Tool call IDs."""
+    tools_call_extra_content: dict[str, dict[str, Any]] = field(default_factory=dict)
+    """Tool call extra content. tool_call_id -> extra_content dict"""
+    reasoning_content: str = ""
+    """The reasoning content extracted from the LLM, if any."""
+
+    raw_completion: (
+        ChatCompletion | GenerateContentResponse | AnthropicMessage | None
+    ) = None
+    """The raw completion response from the LLM provider."""
 
     _completion_text: str = ""
+    """The plain text of the completion."""
 
     is_chunk: bool = False
-    """是否是流式输出的单个 Chunk"""
+    """Indicates if the response is a chunked response."""
 
     def __init__(
         self,
         role: str,
         completion_text: str = "",
-        result_chain: MessageChain = None,
-        tools_call_args: List[Dict[str, any]] = None,
-        tools_call_name: List[str] = None,
-        tools_call_ids: List[str] = None,
-        raw_completion: ChatCompletion = None,
-        _new_record: Dict[str, any] = None,
+        result_chain: MessageChain | None = None,
+        tools_call_args: list[dict[str, Any]] | None = None,
+        tools_call_name: list[str] | None = None,
+        tools_call_ids: list[str] | None = None,
+        tools_call_extra_content: dict[str, dict[str, Any]] | None = None,
+        raw_completion: ChatCompletion
+        | GenerateContentResponse
+        | AnthropicMessage
+        | None = None,
         is_chunk: bool = False,
     ):
         """初始化 LLMResponse
@@ -228,6 +251,7 @@ class LLMResponse:
             tools_call_args (List[Dict[str, any]], optional): 工具调用参数. Defaults to None.
             tools_call_name (List[str], optional): 工具调用名称. Defaults to None.
             raw_completion (ChatCompletion, optional): 原始响应, OpenAI 格式. Defaults to None.
+
         """
         if tools_call_args is None:
             tools_call_args = []
@@ -235,6 +259,8 @@ class LLMResponse:
             tools_call_name = []
         if tools_call_ids is None:
             tools_call_ids = []
+        if tools_call_extra_content is None:
+            tools_call_extra_content = {}
 
         self.role = role
         self.completion_text = completion_text
@@ -242,8 +268,8 @@ class LLMResponse:
         self.tools_call_args = tools_call_args
         self.tools_call_name = tools_call_name
         self.tools_call_ids = tools_call_ids
+        self.tools_call_extra_content = tools_call_extra_content
         self.raw_completion = raw_completion
-        self._new_record = _new_record
         self.is_chunk = is_chunk
 
     @property
@@ -264,18 +290,48 @@ class LLMResponse:
         else:
             self._completion_text = value
 
-    def to_openai_tool_calls(self) -> List[Dict]:
-        """将工具调用信息转换为 OpenAI 格式"""
+    def to_openai_tool_calls(self) -> list[dict]:
+        """Convert to OpenAI tool calls format. Deprecated, use to_openai_to_calls_model instead."""
+        ret = []
+        for idx, tool_call_arg in enumerate(self.tools_call_args):
+            payload = {
+                "id": self.tools_call_ids[idx],
+                "function": {
+                    "name": self.tools_call_name[idx],
+                    "arguments": json.dumps(tool_call_arg),
+                },
+                "type": "function",
+            }
+            if self.tools_call_extra_content.get(self.tools_call_ids[idx]):
+                payload["extra_content"] = self.tools_call_extra_content[
+                    self.tools_call_ids[idx]
+                ]
+            ret.append(payload)
+        return ret
+
+    def to_openai_to_calls_model(self) -> list[ToolCall]:
+        """The same as to_openai_tool_calls but return pydantic model."""
         ret = []
         for idx, tool_call_arg in enumerate(self.tools_call_args):
             ret.append(
-                {
-                    "id": self.tools_call_ids[idx],
-                    "function": {
-                        "name": self.tools_call_name[idx],
-                        "arguments": json.dumps(tool_call_arg),
-                    },
-                    "type": "function",
-                }
+                ToolCall(
+                    id=self.tools_call_ids[idx],
+                    function=ToolCall.FunctionBody(
+                        name=self.tools_call_name[idx],
+                        arguments=json.dumps(tool_call_arg),
+                    ),
+                    # the extra_content will not serialize if it's None when calling ToolCall.model_dump()
+                    extra_content=self.tools_call_extra_content.get(
+                        self.tools_call_ids[idx]
+                    ),
+                ),
             )
         return ret
+
+
+@dataclass
+class RerankResult:
+    index: int
+    """在候选列表中的索引位置"""
+    relevance_score: float
+    """相关性分数"""