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

astrbot/core/provider/sources/anthropic_source.py

@@ -1,60 +1,115 @@
-from typing import List
+import base64
+import json
+from collections.abc import AsyncGenerator
 from mimetypes import guess_type
 
+import anthropic
 from anthropic import AsyncAnthropic
 from anthropic.types import Message
 
-from astrbot.core.utils.io import download_image_by_url
-from astrbot.core.db import BaseDatabase
-from astrbot.api.provider import Provider, Personality
 from astrbot import logger
-from astrbot.core.provider.func_tool_manager import FuncCall
+from astrbot.api.provider import Provider
+from astrbot.core.provider.entities import LLMResponse
+from astrbot.core.provider.func_tool_manager import ToolSet
+from astrbot.core.utils.io import download_image_by_url
+
 from ..register import register_provider_adapter
-from astrbot.core.message.message_event_result import MessageChain
-from astrbot.core.provider.entities import LLMResponse, ToolCallsResult
-from .openai_source import ProviderOpenAIOfficial
 
 
 @register_provider_adapter(
-    "anthropic_chat_completion", "Anthropic Claude API 提供商适配器"
+    "anthropic_chat_completion",
+    "Anthropic Claude API 提供商适配器",
 )
-class ProviderAnthropic(ProviderOpenAIOfficial):
+class ProviderAnthropic(Provider):
     def __init__(
         self,
-        provider_config: dict,
-        provider_settings: dict,
-        db_helper: BaseDatabase,
-        persistant_history=True,
-        default_persona: Personality = None,
+        provider_config,
+        provider_settings,
     ) -> None:
-        # Skip OpenAI's __init__ and call Provider's __init__ directly
-        Provider.__init__(
-            self,
+        super().__init__(
             provider_config,
             provider_settings,
-            persistant_history,
-            db_helper,
-            default_persona,
         )
 
-        self.chosen_api_key = None
-        self.api_keys: List = provider_config.get("key", [])
-        self.chosen_api_key = self.api_keys[0] if len(self.api_keys) > 0 else None
+        self.chosen_api_key: str = ""
+        self.api_keys: list = super().get_keys()
+        self.chosen_api_key = self.api_keys[0] if len(self.api_keys) > 0 else ""
         self.base_url = provider_config.get("api_base", "https://api.anthropic.com")
         self.timeout = provider_config.get("timeout", 120)
         if isinstance(self.timeout, str):
             self.timeout = int(self.timeout)
 
         self.client = AsyncAnthropic(
-            api_key=self.chosen_api_key, timeout=self.timeout, base_url=self.base_url
+            api_key=self.chosen_api_key,
+            timeout=self.timeout,
+            base_url=self.base_url,
         )
 
         self.set_model(provider_config["model_config"]["model"])
 
-    async def _query(self, payloads: dict, tools: FuncCall) -> LLMResponse:
+    def _prepare_payload(self, messages: list[dict]):
+        """准备 Anthropic API 的请求 payload
+
+        Args:
+            messages: OpenAI 格式的消息列表，包含用户输入和系统提示等信息
+        Returns:
+            system_prompt: 系统提示内容
+            new_messages: 处理后的消息列表，去除系统提示
+
+        """
+        system_prompt = ""
+        new_messages = []
+        for message in messages:
+            if message["role"] == "system":
+                system_prompt = message["content"]
+            elif message["role"] == "assistant":
+                blocks = []
+                if isinstance(message["content"], str):
+                    blocks.append({"type": "text", "text": message["content"]})
+                if "tool_calls" in message:
+                    for tool_call in message["tool_calls"]:
+                        blocks.append(  # noqa: PERF401
+                            {
+                                "type": "tool_use",
+                                "name": tool_call["function"]["name"],
+                                "input": (
+                                    json.loads(tool_call["function"]["arguments"])
+                                    if isinstance(
+                                        tool_call["function"]["arguments"],
+                                        str,
+                                    )
+                                    else tool_call["function"]["arguments"]
+                                ),
+                                "id": tool_call["id"],
+                            },
+                        )
+                new_messages.append(
+                    {
+                        "role": "assistant",
+                        "content": blocks,
+                    },
+                )
+            elif message["role"] == "tool":
+                new_messages.append(
+                    {
+                        "role": "user",
+                        "content": [
+                            {
+                                "type": "tool_result",
+                                "tool_use_id": message["tool_call_id"],
+                                "content": message["content"],
+                            },
+                        ],
+                    },
+                )
+            else:
+                new_messages.append(message)
+
+        return system_prompt, new_messages
+
+    async def _query(self, payloads: dict, tools: ToolSet | None) -> LLMResponse:
         if tools:
-            tool_list = tools.get_func_desc_anthropic_style()
-            if tool_list:
+            if tool_list := tools.get_func_desc_anthropic_style():
                 payloads["tools"] = tool_list
 
         completion = await self.client.messages.create(**payloads, stream=False)
@@ -64,68 +119,169 @@ class ProviderAnthropic(ProviderOpenAIOfficial):
 
         if len(completion.content) == 0:
             raise Exception("API 返回的 completion 为空。")
-        # TODO: 如果进行函数调用，思维链被截断，用户可能需要思维链的内容
-        # 选最后一条消息，如果要进行函数调用，anthropic会先返回文本消息的思维链，然后再返回函数调用请求
-        content = completion.content[-1]
-
-        llm_response = LLMResponse("assistant")
-
-        if content.type == "text":
-            # text completion
-            completion_text = str(content.text).strip()
-            # llm_response.completion_text = completion_text
-            llm_response.result_chain = MessageChain().message(completion_text)
-
-        # Anthropic每次只返回一个函数调用
-        if completion.stop_reason == "tool_use":
-            # tools call (function calling)
-            args_ls = []
-            func_name_ls = []
-            tool_use_ids = []
-            func_name_ls.append(content.name)
-            args_ls.append(content.input)
-            tool_use_ids.append(content.id)
-            llm_response.role = "tool"
-            llm_response.tools_call_args = args_ls
-            llm_response.tools_call_name = func_name_ls
-            llm_response.tools_call_ids = tool_use_ids
 
-        if not llm_response.completion_text and not llm_response.tools_call_args:
-            logger.error(f"API 返回的 completion 无法解析：{completion}。")
-            raise Exception(f"API 返回的 completion 无法解析：{completion}。")
+        llm_response = LLMResponse(role="assistant")
 
-        llm_response.raw_completion = completion
+        for content_block in completion.content:
+            if content_block.type == "text":
+                completion_text = str(content_block.text).strip()
+                llm_response.completion_text = completion_text
+
+            if content_block.type == "tool_use":
+                llm_response.tools_call_args.append(content_block.input)
+                llm_response.tools_call_name.append(content_block.name)
+                llm_response.tools_call_ids.append(content_block.id)
+        # TODO(Soulter): 处理 end_turn 情况
+        if not llm_response.completion_text and not llm_response.tools_call_args:
+            raise Exception(f"Anthropic API 返回的 completion 无法解析：{completion}。")
 
         return llm_response
 
+    async def _query_stream(
+        self,
+        payloads: dict,
+        tools: ToolSet | None,
+    ) -> AsyncGenerator[LLMResponse, None]:
+        if tools:
+            if tool_list := tools.get_func_desc_anthropic_style():
+                payloads["tools"] = tool_list
+
+        # 用于累积工具调用信息
+        tool_use_buffer = {}
+        # 用于累积最终结果
+        final_text = ""
+        final_tool_calls = []
+
+        async with self.client.messages.stream(**payloads) as stream:
+            assert isinstance(stream, anthropic.AsyncMessageStream)
+            async for event in stream:
+                if event.type == "content_block_start":
+                    if event.content_block.type == "text":
+                        # 文本块开始
+                        yield LLMResponse(
+                            role="assistant",
+                            completion_text="",
+                            is_chunk=True,
+                        )
+                    elif event.content_block.type == "tool_use":
+                        # 工具使用块开始，初始化缓冲区
+                        tool_use_buffer[event.index] = {
+                            "id": event.content_block.id,
+                            "name": event.content_block.name,
+                            "input": {},
+                        }
+
+                elif event.type == "content_block_delta":
+                    if event.delta.type == "text_delta":
+                        # 文本增量
+                        final_text += event.delta.text
+                        yield LLMResponse(
+                            role="assistant",
+                            completion_text=event.delta.text,
+                            is_chunk=True,
+                        )
+                    elif event.delta.type == "input_json_delta":
+                        # 工具调用参数增量
+                        if event.index in tool_use_buffer:
+                            # 累积 JSON 输入
+                            if "input_json" not in tool_use_buffer[event.index]:
+                                tool_use_buffer[event.index]["input_json"] = ""
+                            tool_use_buffer[event.index]["input_json"] += (
+                                event.delta.partial_json
+                            )
+
+                elif event.type == "content_block_stop":
+                    # 内容块结束
+                    if event.index in tool_use_buffer:
+                        # 解析完整的工具调用
+                        tool_info = tool_use_buffer[event.index]
+                        try:
+                            if "input_json" in tool_info:
+                                tool_info["input"] = json.loads(tool_info["input_json"])
+
+                            # 添加到最终结果
+                            final_tool_calls.append(
+                                {
+                                    "id": tool_info["id"],
+                                    "name": tool_info["name"],
+                                    "input": tool_info["input"],
+                                },
+                            )
+
+                            yield LLMResponse(
+                                role="tool",
+                                completion_text="",
+                                tools_call_args=[tool_info["input"]],
+                                tools_call_name=[tool_info["name"]],
+                                tools_call_ids=[tool_info["id"]],
+                                is_chunk=True,
+                            )
+                        except json.JSONDecodeError:
+                            # JSON 解析失败，跳过这个工具调用
+                            logger.warning(f"工具调用参数 JSON 解析失败: {tool_info}")
+
+                        # 清理缓冲区
+                        del tool_use_buffer[event.index]
+
+        # 返回最终的完整结果
+        final_response = LLMResponse(
+            role="assistant",
+            completion_text=final_text,
+            is_chunk=False,
+        )
+
+        if final_tool_calls:
+            final_response.tools_call_args = [
+                call["input"] for call in final_tool_calls
+            ]
+            final_response.tools_call_name = [call["name"] for call in final_tool_calls]
+            final_response.tools_call_ids = [call["id"] for call in final_tool_calls]
+
+        yield final_response
+
     async def text_chat(
         self,
-        prompt: str,
-        session_id: str = None,
-        image_urls: List[str] = [],
-        func_tool: FuncCall = None,
-        contexts=[],
+        prompt=None,
+        session_id=None,
+        image_urls=None,
+        func_tool=None,
+        contexts=None,
         system_prompt=None,
-        tool_calls_result: ToolCallsResult = None,
+        tool_calls_result=None,
+        model=None,
         **kwargs,
     ) -> LLMResponse:
-        if not prompt:
-            prompt = "<image>"
+        if contexts is None:
+            contexts = []
+        new_record = None
+        if prompt is not None:
+            new_record = await self.assemble_context(prompt, image_urls)
+        context_query = self._ensure_message_to_dicts(contexts)
+        if new_record:
+            context_query.append(new_record)
 
-        new_record = await self.assemble_context(prompt, image_urls)
-        context_query = [*contexts, new_record]
+        if system_prompt:
+            context_query.insert(0, {"role": "system", "content": system_prompt})
 
         for part in context_query:
             if "_no_save" in part:
                 del part["_no_save"]
 
+        # tool calls result
         if tool_calls_result:
-            # 暂时这样写。
-            prompt += f"Here are the related results via using tools: {str(tool_calls_result.tool_calls_result)}"
+            if not isinstance(tool_calls_result, list):
+                context_query.extend(tool_calls_result.to_openai_messages())
+            else:
+                for tcr in tool_calls_result:
+                    context_query.extend(tcr.to_openai_messages())
+
+        system_prompt, new_messages = self._prepare_payload(context_query)
 
         model_config = self.provider_config.get("model_config", {})
+        model_config["model"] = model or self.get_model()
+
+        payloads = {"messages": new_messages, **model_config}
 
-        payloads = {"messages": context_query, **model_config}
         # Anthropic has a different way of handling system prompts
         if system_prompt:
             payloads["system"] = system_prompt
@@ -133,32 +289,9 @@ class ProviderAnthropic(ProviderOpenAIOfficial):
         llm_response = None
         try:
             llm_response = await self._query(payloads, func_tool)
-
         except Exception as e:
-            if "maximum context length" in str(e):
-                retry_cnt = 20
-                while retry_cnt > 0:
-                    logger.warning(
-                        f"上下文长度超过限制。尝试弹出最早的记录然后重试。当前记录条数: {len(context_query)}"
-                    )
-                    try:
-                        await self.pop_record(context_query)
-                        response = await self.client.messages.create(
-                            messages=context_query, **model_config
-                        )
-                        llm_response = LLMResponse("assistant")
-                        llm_response.result_chain = MessageChain().message(response.content[0].text)
-                        llm_response.raw_completion = response
-                        return llm_response
-                    except Exception as e:
-                        if "maximum context length" in str(e):
-                            retry_cnt -= 1
-                        else:
-                            raise e
-                return LLMResponse("err", "err: 请尝试 /reset 清除会话记录。")
-            else:
-                logger.error(f"发生了错误。Provider 配置如下: {model_config}")
-                raise e
+            logger.error(f"发生了错误。Provider 配置如下: {model_config}")
+            raise e
 
         return llm_response
 
@@ -171,25 +304,47 @@ class ProviderAnthropic(ProviderOpenAIOfficial):
         contexts=...,
         system_prompt=None,
         tool_calls_result=None,
+        model=None,
         **kwargs,
     ):
-        # raise NotImplementedError("This method is not implemented yet.")
-        # 调用 text_chat 模拟流式
-        llm_response = await self.text_chat(
-            prompt=prompt,
-            session_id=session_id,
-            image_urls=image_urls,
-            func_tool=func_tool,
-            contexts=contexts,
-            system_prompt=system_prompt,
-            tool_calls_result=tool_calls_result,
-        )
-        llm_response.is_chunk = True
-        yield llm_response
-        llm_response.is_chunk = False
-        yield llm_response
+        if contexts is None:
+            contexts = []
+        new_record = None
+        if prompt is not None:
+            new_record = await self.assemble_context(prompt, image_urls)
+        context_query = self._ensure_message_to_dicts(contexts)
+        if new_record:
+            context_query.append(new_record)
+        if system_prompt:
+            context_query.insert(0, {"role": "system", "content": system_prompt})
+
+        for part in context_query:
+            if "_no_save" in part:
+                del part["_no_save"]
 
-    async def assemble_context(self, text: str, image_urls: List[str] = None):
+        # tool calls result
+        if tool_calls_result:
+            if not isinstance(tool_calls_result, list):
+                context_query.extend(tool_calls_result.to_openai_messages())
+            else:
+                for tcr in tool_calls_result:
+                    context_query.extend(tcr.to_openai_messages())
+
+        system_prompt, new_messages = self._prepare_payload(context_query)
+
+        model_config = self.provider_config.get("model_config", {})
+        model_config["model"] = model or self.get_model()
+
+        payloads = {"messages": new_messages, **model_config}
+
+        # Anthropic has a different way of handling system prompts
+        if system_prompt:
+            payloads["system"] = system_prompt
+
+        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 = None):
         """组装上下文，支持文本和图片"""
         if not image_urls:
             return {"role": "user", "content": text}
@@ -222,11 +377,36 @@ class ProviderAnthropic(ProviderOpenAIOfficial):
                     "source": {
                         "type": "base64",
                         "media_type": mime_type,
-                        "data": image_data.split("base64,")[1]
-                        if "base64," in image_data
-                        else image_data,
+                        "data": (
+                            image_data.split("base64,")[1]
+                            if "base64," in image_data
+                            else image_data
+                        ),
                     },
-                }
+                },
             )
 
         return {"role": "user", "content": content}
+
+    async def encode_image_bs64(self, image_url: str) -> str:
+        """将图片转换为 base64"""
+        if image_url.startswith("base64://"):
+            return image_url.replace("base64://", "data:image/jpeg;base64,")
+        with open(image_url, "rb") as f:
+            image_bs64 = base64.b64encode(f.read()).decode("utf-8")
+            return "data:image/jpeg;base64," + image_bs64
+        return ""
+
+    def get_current_key(self) -> str:
+        return self.chosen_api_key
+
+    async def get_models(self) -> list[str]:
+        models_str = []
+        models = await self.client.models.list()
+        models = sorted(models.data, key=lambda x: x.id)
+        for model in models:
+            models_str.append(model.id)
+        return models_str
+
+    def set_key(self, key: str):
+        self.chosen_api_key = key