AstrBot/astrbot/core/provider/sources/openai_source.py

import asyncio
import base64
import inspect
import json
import os
import random
import re
from collections.abc import AsyncGenerator

from openai import AsyncAzureOpenAI, AsyncOpenAI
from openai._exceptions import NotFoundError
from openai.lib.streaming.chat._completions import ChatCompletionStreamState
from openai.types.chat.chat_completion import ChatCompletion
from openai.types.chat.chat_completion_chunk import ChatCompletionChunk

import astrbot.core.message.components as Comp
from astrbot import logger
from astrbot.api.provider import Provider
from astrbot.core.agent.message import Message
from astrbot.core.agent.tool import ToolSet
from astrbot.core.message.message_event_result import MessageChain
from astrbot.core.provider.entities import LLMResponse, ToolCallsResult
from astrbot.core.utils.io import download_image_by_url

from ..register import register_provider_adapter


@register_provider_adapter(
    "openai_chat_completion",
    "OpenAI API Chat Completion 提供商适配器",
)
class ProviderOpenAIOfficial(Provider):
    def __init__(self, provider_config, provider_settings) -> None:
        super().__init__(provider_config, provider_settings)
        self.chosen_api_key = None
        self.api_keys: list = super().get_keys()
        self.chosen_api_key = self.api_keys[0] if len(self.api_keys) > 0 else None
        self.timeout = provider_config.get("timeout", 120)
        self.custom_headers = provider_config.get("custom_headers", {})
        if isinstance(self.timeout, str):
            self.timeout = int(self.timeout)

        if not isinstance(self.custom_headers, dict) or not self.custom_headers:
            self.custom_headers = None
        else:
            for key in self.custom_headers:
                self.custom_headers[key] = str(self.custom_headers[key])

        if "api_version" in provider_config:
            # Using Azure OpenAI API
            self.client = AsyncAzureOpenAI(
                api_key=self.chosen_api_key,
                api_version=provider_config.get("api_version", None),
                default_headers=self.custom_headers,
                base_url=provider_config.get("api_base", ""),
                timeout=self.timeout,
            )
        else:
            # Using OpenAI Official API
            self.client = AsyncOpenAI(
                api_key=self.chosen_api_key,
                base_url=provider_config.get("api_base", None),
                default_headers=self.custom_headers,
                timeout=self.timeout,
            )

        self.default_params = inspect.signature(
            self.client.chat.completions.create,
        ).parameters.keys()

        model_config = provider_config.get("model_config", {})
        model = model_config.get("model", "unknown")
        self.set_model(model)

        self.reasoning_key = "reasoning_content"

    def _maybe_inject_xai_search(self, payloads: dict, **kwargs):
        """当开启 xAI 原生搜索时，向请求体注入 Live Search 参数。

        - 仅在 provider_config.xai_native_search 为 True 时生效
        - 默认注入 {"mode": "auto"}
        - 允许通过 kwargs 使用 xai_search_mode 覆盖（on/auto/off）
        """
        if not bool(self.provider_config.get("xai_native_search", False)):
            return

        mode = kwargs.get("xai_search_mode", "auto")
        mode = str(mode).lower()
        if mode not in ("auto", "on", "off"):
            mode = "auto"

        # off 时不注入，保持与未开启一致
        if mode == "off":
            return

        # OpenAI SDK 不识别的字段会在 _query/_query_stream 中放入 extra_body
        payloads["search_parameters"] = {"mode": mode}

    async def get_models(self):
        try:
            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
        except NotFoundError as e:
            raise Exception(f"获取模型列表失败：{e}")

    async def _query(self, payloads: dict, tools: ToolSet | None) -> LLMResponse:
        if tools:
            model = payloads.get("model", "").lower()
            omit_empty_param_field = "gemini" in model
            tool_list = tools.get_func_desc_openai_style(
                omit_empty_parameter_field=omit_empty_param_field,
            )
            if tool_list:
                payloads["tools"] = tool_list

        # 不在默认参数中的参数放在 extra_body 中
        extra_body = {}
        to_del = []
        for key in payloads:
            if key not in self.default_params:
                extra_body[key] = payloads[key]
                to_del.append(key)
        for key in to_del:
            del payloads[key]

        # 读取并合并 custom_extra_body 配置
        custom_extra_body = self.provider_config.get("custom_extra_body", {})
        if isinstance(custom_extra_body, dict):
            extra_body.update(custom_extra_body)

        model = payloads.get("model", "").lower()

        # 针对 deepseek 模型的特殊处理：deepseek-reasoner调用必须移除 tools ，否则将被切换至 deepseek-chat
        if model == "deepseek-reasoner" and "tools" in payloads:
            del payloads["tools"]

        completion = await self.client.chat.completions.create(
            **payloads,
            stream=False,
            extra_body=extra_body,
        )

        if not isinstance(completion, ChatCompletion):
            raise Exception(
                f"API 返回的 completion 类型错误：{type(completion)}: {completion}。",
            )

        logger.debug(f"completion: {completion}")

        llm_response = await self._parse_openai_completion(completion, tools)

        return llm_response

    async def _query_stream(
        self,
        payloads: dict,
        tools: ToolSet | None,
    ) -> AsyncGenerator[LLMResponse, None]:
        """流式查询API，逐步返回结果"""
        if tools:
            model = payloads.get("model", "").lower()
            omit_empty_param_field = "gemini" in model
            tool_list = tools.get_func_desc_openai_style(
                omit_empty_parameter_field=omit_empty_param_field,
            )
            if tool_list:
                payloads["tools"] = tool_list

        # 不在默认参数中的参数放在 extra_body 中
        extra_body = {}

        # 读取并合并 custom_extra_body 配置
        custom_extra_body = self.provider_config.get("custom_extra_body", {})
        if isinstance(custom_extra_body, dict):
            extra_body.update(custom_extra_body)

        to_del = []
        for key in payloads:
            if key not in self.default_params:
                extra_body[key] = payloads[key]
                to_del.append(key)
        for key in to_del:
            del payloads[key]

        stream = await self.client.chat.completions.create(
            **payloads,
            stream=True,
            extra_body=extra_body,
        )

        llm_response = LLMResponse("assistant", is_chunk=True)

        state = ChatCompletionStreamState()

        async for chunk in stream:
            try:
                state.handle_chunk(chunk)
            except Exception as e:
                logger.warning("Saving chunk state error: " + str(e))
            if len(chunk.choices) == 0:
                continue
            delta = chunk.choices[0].delta
            # logger.debug(f"chunk delta: {delta}")
            # handle the content delta
            reasoning = self._extract_reasoning_content(chunk)
            _y = False
            if reasoning:
                llm_response.reasoning_content = reasoning
                _y = True
            if delta.content:
                completion_text = delta.content
                llm_response.result_chain = MessageChain(
                    chain=[Comp.Plain(completion_text)],
                )
                _y = True
            if _y:
                yield llm_response

        final_completion = state.get_final_completion()
        llm_response = await self._parse_openai_completion(final_completion, tools)

        yield llm_response

    def _extract_reasoning_content(
        self,
        completion: ChatCompletion | ChatCompletionChunk,
    ) -> str:
        """Extract reasoning content from OpenAI ChatCompletion if available."""
        reasoning_text = ""
        if len(completion.choices) == 0:
            return reasoning_text
        if isinstance(completion, ChatCompletion):
            choice = completion.choices[0]
            reasoning_attr = getattr(choice.message, self.reasoning_key, None)
            if reasoning_attr:
                reasoning_text = str(reasoning_attr)
        elif isinstance(completion, ChatCompletionChunk):
            delta = completion.choices[0].delta
            reasoning_attr = getattr(delta, self.reasoning_key, None)
            if reasoning_attr:
                reasoning_text = str(reasoning_attr)
        return reasoning_text

    async def _parse_openai_completion(
        self, completion: ChatCompletion, tools: ToolSet | None
    ) -> LLMResponse:
        """Parse OpenAI ChatCompletion into LLMResponse"""
        llm_response = LLMResponse("assistant")

        if len(completion.choices) == 0:
            raise Exception("API 返回的 completion 为空。")
        choice = completion.choices[0]

        # parse the text completion
        if choice.message.content is not None:
            # text completion
            completion_text = str(choice.message.content).strip()
            # specially, some providers may set <think> tags around reasoning content in the completion text,
            # we use regex to remove them, and store then in reasoning_content field
            reasoning_pattern = re.compile(r"<think>(.*?)</think>", re.DOTALL)
            matches = reasoning_pattern.findall(completion_text)
            if matches:
                llm_response.reasoning_content = "\n".join(
                    [match.strip() for match in matches],
                )
                completion_text = reasoning_pattern.sub("", completion_text).strip()
            llm_response.result_chain = MessageChain().message(completion_text)

        # parse the reasoning content if any
        # the priority is higher than the <think> tag extraction
        llm_response.reasoning_content = self._extract_reasoning_content(completion)

        # parse tool calls if any
        if choice.message.tool_calls and tools is not None:
            args_ls = []
            func_name_ls = []
            tool_call_ids = []
            tool_call_extra_content_dict = {}
            for tool_call in choice.message.tool_calls:
                if isinstance(tool_call, str):
                    # workaround for #1359
                    tool_call = json.loads(tool_call)
                if tools is None:
                    # 工具集未提供
                    # Should be unreachable
                    raise Exception("工具集未提供")
                for tool in tools.func_list:
                    if (
                        tool_call.type == "function"
                        and tool.name == tool_call.function.name
                    ):
                        # workaround for #1454
                        if isinstance(tool_call.function.arguments, str):
                            args = json.loads(tool_call.function.arguments)
                        else:
                            args = tool_call.function.arguments
                        args_ls.append(args)
                        func_name_ls.append(tool_call.function.name)
                        tool_call_ids.append(tool_call.id)

                        # gemini-2.5 / gemini-3 series extra_content handling
                        extra_content = getattr(tool_call, "extra_content", None)
                        if extra_content is not None:
                            tool_call_extra_content_dict[tool_call.id] = extra_content
            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_call_ids
            llm_response.tools_call_extra_content = tool_call_extra_content_dict
        # specially handle finish reason
        if choice.finish_reason == "content_filter":
            raise Exception(
                "API 返回的 completion 由于内容安全过滤被拒绝(非 AstrBot)。",
            )
        if llm_response.completion_text is None and not llm_response.tools_call_args:
            logger.error(f"API 返回的 completion 无法解析：{completion}。")
            raise Exception(f"API 返回的 completion 无法解析：{completion}。")

        llm_response.raw_completion = completion

        return llm_response

    async def _prepare_chat_payload(
        self,
        prompt: str | None,
        image_urls: list[str] | None = None,
        contexts: list[dict] | list[Message] | None = None,
        system_prompt: str | None = None,
        tool_calls_result: ToolCallsResult | list[ToolCallsResult] | None = None,
        model: str | None = None,
        **kwargs,
    ) -> tuple:
        """准备聊天所需的有效载荷和上下文"""
        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"]

        # tool calls result
        if tool_calls_result:
            if isinstance(tool_calls_result, ToolCallsResult):
                context_query.extend(tool_calls_result.to_openai_messages())
            else:
                for tcr in tool_calls_result:
                    context_query.extend(tcr.to_openai_messages())

        model_config = self.provider_config.get("model_config", {})
        model_config["model"] = model or self.get_model()

        payloads = {"messages": context_query, **model_config}

        # xAI origin search tool inject
        self._maybe_inject_xai_search(payloads, **kwargs)

        return payloads, context_query

    async def _handle_api_error(
        self,
        e: Exception,
        payloads: dict,
        context_query: list,
        func_tool: ToolSet | None,
        chosen_key: str,
        available_api_keys: list[str],
        retry_cnt: int,
        max_retries: int,
    ) -> tuple:
        """处理API错误并尝试恢复"""
        if "429" in str(e):
            logger.warning(
                f"API 调用过于频繁，尝试使用其他 Key 重试。当前 Key: {chosen_key[:12]}",
            )
            # 最后一次不等待
            if retry_cnt < max_retries - 1:
                await asyncio.sleep(1)
            available_api_keys.remove(chosen_key)
            if len(available_api_keys) > 0:
                chosen_key = random.choice(available_api_keys)
                return (
                    False,
                    chosen_key,
                    available_api_keys,
                    payloads,
                    context_query,
                    func_tool,
                )
            raise e
        if "maximum context length" in str(e):
            logger.warning(
                f"上下文长度超过限制。尝试弹出最早的记录然后重试。当前记录条数: {len(context_query)}",
            )
            await self.pop_record(context_query)
            payloads["messages"] = context_query
            return (
                False,
                chosen_key,
                available_api_keys,
                payloads,
                context_query,
                func_tool,
            )
        if "The model is not a VLM" in str(e):  # siliconcloud
            # 尝试删除所有 image
            new_contexts = await self._remove_image_from_context(context_query)
            payloads["messages"] = new_contexts
            context_query = new_contexts
            return (
                False,
                chosen_key,
                available_api_keys,
                payloads,
                context_query,
                func_tool,
            )
        if (
            "Function calling is not enabled" in str(e)
            or ("tool" in str(e).lower() and "support" in str(e).lower())
            or ("function" in str(e).lower() and "support" in str(e).lower())
        ):
            # openai, ollama, gemini openai, siliconcloud 的错误提示与 code 不统一，只能通过字符串匹配
            logger.info(
                f"{self.get_model()} 不支持函数工具调用，已自动去除，不影响使用。",
            )
            payloads.pop("tools", None)
            return False, chosen_key, available_api_keys, payloads, context_query, None
        # logger.error(f"发生了错误。Provider 配置如下: {self.provider_config}")

        if "tool" in str(e).lower() and "support" in str(e).lower():
            logger.error("疑似该模型不支持函数调用工具调用。请输入 /tool off_all")

        if "Connection error." in str(e):
            proxy = os.environ.get("http_proxy", None)
            if proxy:
                logger.error(
                    f"可能为代理原因，请检查代理是否正常。当前代理: {proxy}",
                )

        raise e

    async def text_chat(
        self,
        prompt=None,
        session_id=None,
        image_urls=None,
        func_tool=None,
        contexts=None,
        system_prompt=None,
        tool_calls_result=None,
        model=None,
        **kwargs,
    ) -> LLMResponse:
        payloads, context_query = await self._prepare_chat_payload(
            prompt,
            image_urls,
            contexts,
            system_prompt,
            tool_calls_result,
            model=model,
            **kwargs,
        )

        llm_response = None
        max_retries = 10
        available_api_keys = self.api_keys.copy()
        chosen_key = random.choice(available_api_keys)

        last_exception = None
        retry_cnt = 0
        for retry_cnt in range(max_retries):
            try:
                self.client.api_key = chosen_key
                llm_response = await self._query(payloads, func_tool)
                break
            except Exception as e:
                last_exception = e
                (
                    success,
                    chosen_key,
                    available_api_keys,
                    payloads,
                    context_query,
                    func_tool,
                ) = await self._handle_api_error(
                    e,
                    payloads,
                    context_query,
                    func_tool,
                    chosen_key,
                    available_api_keys,
                    retry_cnt,
                    max_retries,
                )
                if success:
                    break

        if retry_cnt == max_retries - 1 or llm_response is None:
            logger.error(f"API 调用失败，重试 {max_retries} 次仍然失败。")
            if last_exception is None:
                raise Exception("未知错误")
            raise last_exception
        return llm_response

    async def text_chat_stream(
        self,
        prompt=None,
        session_id=None,
        image_urls=None,
        func_tool=None,
        contexts=None,
        system_prompt=None,
        tool_calls_result=None,
        model=None,
        **kwargs,
    ) -> AsyncGenerator[LLMResponse, None]:
        """流式对话，与服务商交互并逐步返回结果"""
        payloads, context_query = await self._prepare_chat_payload(
            prompt,
            image_urls,
            contexts,
            system_prompt,
            tool_calls_result,
            model=model,
            **kwargs,
        )

        max_retries = 10
        available_api_keys = self.api_keys.copy()
        chosen_key = random.choice(available_api_keys)

        last_exception = None
        retry_cnt = 0
        for retry_cnt in range(max_retries):
            try:
                self.client.api_key = chosen_key
                async for response in self._query_stream(payloads, func_tool):
                    yield response
                break
            except Exception as e:
                last_exception = e
                (
                    success,
                    chosen_key,
                    available_api_keys,
                    payloads,
                    context_query,
                    func_tool,
                ) = await self._handle_api_error(
                    e,
                    payloads,
                    context_query,
                    func_tool,
                    chosen_key,
                    available_api_keys,
                    retry_cnt,
                    max_retries,
                )
                if success:
                    break

        if retry_cnt == max_retries - 1:
            logger.error(f"API 调用失败，重试 {max_retries} 次仍然失败。")
            if last_exception is None:
                raise Exception("未知错误")
            raise last_exception

    async def _remove_image_from_context(self, contexts: list):
        """从上下文中删除所有带有 image 的记录"""
        new_contexts = []

        for context in contexts:
            if "content" in context and isinstance(context["content"], list):
                # continue
                new_content = []
                for item in context["content"]:
                    if isinstance(item, dict) and "image_url" in item:
                        continue
                    new_content.append(item)
                if not new_content:
                    # 用户只发了图片
                    new_content = [{"type": "text", "text": "[图片]"}]
                context["content"] = new_content
            new_contexts.append(context)
        return new_contexts

    def get_current_key(self) -> str:
        return self.client.api_key

    def get_keys(self) -> list[str]:
        return self.api_keys

    def set_key(self, key):
        self.client.api_key = key

    async def assemble_context(
        self,
        text: str,
        image_urls: list[str] | None = None,
    ) -> dict:
        """组装成符合 OpenAI 格式的 role 为 user 的消息段"""
        if image_urls:
            user_content = {
                "role": "user",
                "content": [{"type": "text", "text": text if text else "[图片]"}],
            }
            for image_url in image_urls:
                if image_url.startswith("http"):
                    image_path = await download_image_by_url(image_url)
                    image_data = await self.encode_image_bs64(image_path)
                elif image_url.startswith("file:///"):
                    image_path = image_url.replace("file:///", "")
                    image_data = await self.encode_image_bs64(image_path)
                else:
                    image_data = await self.encode_image_bs64(image_url)
                if not image_data:
                    logger.warning(f"图片 {image_url} 得到的结果为空，将忽略。")
                    continue
                user_content["content"].append(
                    {
                        "type": "image_url",
                        "image_url": {"url": image_data},
                    },
                )
            return user_content
        return {"role": "user", "content": text}

    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