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

astrbot/core/pipeline/process_stage/method/llm_request.py

@@ -1,606 +0,0 @@
-"""
-本地 Agent 模式的 LLM 调用 Stage
-"""
-
-import traceback
-import asyncio
-import json
-from typing import Union, AsyncGenerator
-from ...context import PipelineContext
-from ..stage import Stage
-from astrbot.core.platform.astr_message_event import AstrMessageEvent
-from astrbot.core.message.message_event_result import (
-    MessageEventResult,
-    ResultContentType,
-    MessageChain,
-)
-from astrbot.core.message.components import Image
-from astrbot.core import logger
-from astrbot.core.utils.metrics import Metric
-from astrbot.core.provider.entities import (
-    ProviderRequest,
-    LLMResponse,
-    ToolCallMessageSegment,
-    AssistantMessageSegment,
-    ToolCallsResult,
-)
-from astrbot.core.star.star_handler import star_handlers_registry, EventType
-from astrbot.core.star.star import star_map
-from mcp.types import (
-    TextContent,
-    ImageContent,
-    EmbeddedResource,
-    TextResourceContents,
-    BlobResourceContents,
-)
-
-
-class LLMRequestSubStage(Stage):
-    async def initialize(self, ctx: PipelineContext) -> None:
-        self.ctx = ctx
-        self.bot_wake_prefixs = ctx.astrbot_config["wake_prefix"]  # list
-        self.provider_wake_prefix = ctx.astrbot_config["provider_settings"][
-            "wake_prefix"
-        ]  # str
-        self.max_context_length = ctx.astrbot_config["provider_settings"][
-            "max_context_length"
-        ]  # int
-        self.dequeue_context_length = min(
-            max(1, ctx.astrbot_config["provider_settings"]["dequeue_context_length"]),
-            self.max_context_length - 1,
-        )  # int
-        self.streaming_response = ctx.astrbot_config["provider_settings"][
-            "streaming_response"
-        ]  # bool
-
-        for bwp in self.bot_wake_prefixs:
-            if self.provider_wake_prefix.startswith(bwp):
-                logger.info(
-                    f"识别 LLM 聊天额外唤醒前缀 {self.provider_wake_prefix} 以机器人唤醒前缀 {bwp} 开头，已自动去除。"
-                )
-                self.provider_wake_prefix = self.provider_wake_prefix[len(bwp) :]
-
-        self.conv_manager = ctx.plugin_manager.context.conversation_manager
-
-    async def process(
-        self, event: AstrMessageEvent, _nested: bool = False
-    ) -> Union[None, AsyncGenerator[None, None]]:
-        req: ProviderRequest = None
-
-        provider = self.ctx.plugin_manager.context.get_using_provider()
-        if provider is None:
-            return
-
-        if event.get_extra("provider_request"):
-            req = event.get_extra("provider_request")
-            assert isinstance(req, ProviderRequest), (
-                "provider_request 必须是 ProviderRequest 类型。"
-            )
-
-            if req.conversation:
-                all_contexts = json.loads(req.conversation.history)
-                req.contexts = self._process_tool_message_pairs(
-                    all_contexts, remove_tags=True
-                )
-
-        else:
-            req = ProviderRequest(prompt="", image_urls=[])
-            if self.provider_wake_prefix:
-                if not event.message_str.startswith(self.provider_wake_prefix):
-                    return
-            req.prompt = event.message_str[len(self.provider_wake_prefix) :]
-            req.func_tool = self.ctx.plugin_manager.context.get_llm_tool_manager()
-            for comp in event.message_obj.message:
-                if isinstance(comp, Image):
-                    image_path = await comp.convert_to_file_path()
-                    req.image_urls.append(image_path)
-
-            # 获取对话上下文
-            conversation_id = await self.conv_manager.get_curr_conversation_id(
-                event.unified_msg_origin
-            )
-            if not conversation_id:
-                conversation_id = await self.conv_manager.new_conversation(
-                    event.unified_msg_origin
-                )
-            conversation = await self.conv_manager.get_conversation(
-                event.unified_msg_origin, conversation_id
-            )
-            if not conversation:
-                conversation_id = await self.conv_manager.new_conversation(
-                    event.unified_msg_origin
-                )
-                conversation = await self.conv_manager.get_conversation(
-                    event.unified_msg_origin, conversation_id
-                )
-            req.conversation = conversation
-            req.contexts = json.loads(conversation.history)
-
-            event.set_extra("provider_request", req)
-
-        if not req.prompt and not req.image_urls:
-            return
-
-        # 执行请求 LLM 前事件钩子。
-        # 装饰 system_prompt 等功能
-        # 获取当前平台ID
-        platform_id = event.get_platform_id()
-        handlers = star_handlers_registry.get_handlers_by_event_type(
-            EventType.OnLLMRequestEvent, platform_id=platform_id
-        )
-        for handler in handlers:
-            try:
-                logger.debug(
-                    f"hook(on_llm_request) -> {star_map[handler.handler_module_path].name} - {handler.handler_name}"
-                )
-                await handler.handler(event, req)
-            except BaseException:
-                logger.error(traceback.format_exc())
-
-            if event.is_stopped():
-                logger.info(
-                    f"{star_map[handler.handler_module_path].name} - {handler.handler_name} 终止了事件传播。"
-                )
-                return
-
-        if isinstance(req.contexts, str):
-            req.contexts = json.loads(req.contexts)
-
-        # max context length
-        if (
-            self.max_context_length != -1  # -1 为不限制
-            and len(req.contexts) // 2 > self.max_context_length
-        ):
-            logger.debug("上下文长度超过限制，将截断。")
-            req.contexts = req.contexts[
-                -(self.max_context_length - self.dequeue_context_length + 1) * 2 :
-            ]
-            # 找到第一个role 为 user 的索引，确保上下文格式正确
-            index = next(
-                (
-                    i
-                    for i, item in enumerate(req.contexts)
-                    if item.get("role") == "user"
-                ),
-                None,
-            )
-            if index is not None and index > 0:
-                req.contexts = req.contexts[index:]
-
-        # session_id
-        if not req.session_id:
-            req.session_id = event.unified_msg_origin
-
-        async def requesting(req: ProviderRequest):
-            try:
-                need_loop = True
-                while need_loop:
-                    need_loop = False
-                    logger.debug(f"提供商请求 Payload: {req}")
-
-                    final_llm_response = None
-
-                    if self.streaming_response:
-                        stream = provider.text_chat_stream(**req.__dict__)
-                        async for llm_response in stream:
-                            if llm_response.is_chunk:
-                                if llm_response.result_chain:
-                                    yield llm_response.result_chain  # MessageChain
-                                else:
-                                    yield MessageChain().message(
-                                        llm_response.completion_text
-                                    )
-                            else:
-                                final_llm_response = llm_response
-                    else:
-                        final_llm_response = await provider.text_chat(
-                            **req.__dict__
-                        )  # 请求 LLM
-
-                    if not final_llm_response:
-                        raise Exception("LLM response is None.")
-
-                    # 执行 LLM 响应后的事件钩子。
-                    handlers = star_handlers_registry.get_handlers_by_event_type(
-                        EventType.OnLLMResponseEvent
-                    )
-                    for handler in handlers:
-                        try:
-                            logger.debug(
-                                f"hook(on_llm_response) -> {star_map[handler.handler_module_path].name} - {handler.handler_name}"
-                            )
-                            await handler.handler(event, final_llm_response)
-                        except BaseException:
-                            logger.error(traceback.format_exc())
-
-                        if event.is_stopped():
-                            logger.info(
-                                f"{star_map[handler.handler_module_path].name} - {handler.handler_name} 终止了事件传播。"
-                            )
-                            return
-
-                    if self.streaming_response:
-                        # 流式输出的处理
-                        async for result in self._handle_llm_stream_response(
-                            event, req, final_llm_response
-                        ):
-                            if isinstance(result, ProviderRequest):
-                                # 有函数工具调用并且返回了结果，我们需要再次请求 LLM
-                                req = result
-                                need_loop = True
-                            else:
-                                yield
-                    else:
-                        # 非流式输出的处理
-                        async for result in self._handle_llm_response(
-                            event, req, final_llm_response
-                        ):
-                            if isinstance(result, ProviderRequest):
-                                # 有函数工具调用并且返回了结果，我们需要再次请求 LLM
-                                req = result
-                                need_loop = True
-                            else:
-                                yield
-
-                asyncio.create_task(
-                    Metric.upload(
-                        llm_tick=1,
-                        model_name=provider.get_model(),
-                        provider_type=provider.meta().type,
-                    )
-                )
-
-                # 保存到历史记录
-                await self._save_to_history(event, req, final_llm_response)
-
-            except BaseException as e:
-                logger.error(traceback.format_exc())
-                event.set_result(
-                    MessageEventResult().message(
-                        f"AstrBot 请求失败。\n错误类型: {type(e).__name__}\n错误信息: {str(e)}"
-                    )
-                )
-
-        if not self.streaming_response:
-            event.set_extra("tool_call_result", None)
-            async for _ in requesting(req):
-                yield
-        else:
-            event.set_result(
-                MessageEventResult()
-                .set_result_content_type(ResultContentType.STREAMING_RESULT)
-                .set_async_stream(requesting(req))
-            )
-            # 这里使用yield来暂停当前阶段，等待流式输出完成后继续处理
-            yield
-
-            if event.get_extra("tool_call_result"):
-                event.set_result(event.get_extra("tool_call_result"))
-                event.set_extra("tool_call_result", None)
-                yield
-
-        # 暂时直接发出去
-        if img_b64 := event.get_extra("tool_call_img_respond"):
-            await event.send(MessageChain(chain=[Image.fromBase64(img_b64)]))
-            event.set_extra("tool_call_img_respond", None)
-            yield
-
-    async def _handle_llm_response(
-        self,
-        event: AstrMessageEvent,
-        req: ProviderRequest,
-        llm_response: LLMResponse,
-    ) -> AsyncGenerator[Union[None, ProviderRequest], None]:
-        """处理非流式 LLM 响应。
-
-        Returns:
-            AsyncGenerator[Union[None, ProviderRequest], None]: 如果返回 ProviderRequest，表示需要再次调用 LLM
-
-        Yields:
-            Iterator[Union[None, ProviderRequest]]: 将 event 交付给下一个 stage 或者返回 ProviderRequest 表示需要再次调用 LLM
-        """
-        if llm_response.role == "assistant":
-            # text completion
-            if llm_response.result_chain:
-                event.set_result(
-                    MessageEventResult(
-                        chain=llm_response.result_chain.chain
-                    ).set_result_content_type(ResultContentType.LLM_RESULT)
-                )
-            else:
-                event.set_result(
-                    MessageEventResult()
-                    .message(llm_response.completion_text)
-                    .set_result_content_type(ResultContentType.LLM_RESULT)
-                )
-        elif llm_response.role == "err":
-            event.set_result(
-                MessageEventResult().message(
-                    f"AstrBot 请求失败。\n错误信息: {llm_response.completion_text}"
-                )
-            )
-        elif llm_response.role == "tool":
-            # 处理函数工具调用
-            async for result in self._handle_function_tools(event, req, llm_response):
-                yield result
-
-    async def _handle_llm_stream_response(
-        self,
-        event: AstrMessageEvent,
-        req: ProviderRequest,
-        llm_response: LLMResponse,
-    ) -> AsyncGenerator[Union[None, ProviderRequest], None]:
-        """处理流式 LLM 响应。
-
-        专门用于处理流式输出完成后的响应，与非流式响应处理分离。
-
-        Returns:
-            AsyncGenerator[Union[None, ProviderRequest], None]: 如果返回 ProviderRequest，表示需要再次调用 LLM
-
-        Yields:
-            Iterator[Union[None, ProviderRequest]]: 将 event 交付给下一个 stage 或者返回 ProviderRequest 表示需要再次调用 LLM
-        """
-        if llm_response.role == "assistant":
-            # text completion
-            if llm_response.result_chain:
-                event.set_result(
-                    MessageEventResult(
-                        chain=llm_response.result_chain.chain
-                    ).set_result_content_type(ResultContentType.STREAMING_FINISH)
-                )
-            else:
-                event.set_result(
-                    MessageEventResult()
-                    .message(llm_response.completion_text)
-                    .set_result_content_type(ResultContentType.STREAMING_FINISH)
-                )
-        elif llm_response.role == "err":
-            event.set_result(
-                MessageEventResult().message(
-                    f"AstrBot 请求失败。\n错误信息: {llm_response.completion_text}"
-                )
-            )
-        elif llm_response.role == "tool":
-            # 处理函数工具调用
-            async for result in self._handle_function_tools(event, req, llm_response):
-                yield result
-
-    async def _handle_function_tools(
-        self,
-        event: AstrMessageEvent,
-        req: ProviderRequest,
-        llm_response: LLMResponse,
-    ) -> AsyncGenerator[Union[None, ProviderRequest], None]:
-        """处理函数工具调用。
-
-        Returns:
-            AsyncGenerator[Union[None, ProviderRequest], None]: 如果返回 ProviderRequest，表示需要再次调用 LLM
-        """
-        # function calling
-        tool_call_result: list[ToolCallMessageSegment] = []
-        logger.info(
-            f"触发 {len(llm_response.tools_call_name)} 个函数调用: {llm_response.tools_call_name}"
-        )
-        for func_tool_name, func_tool_args, func_tool_id in zip(
-            llm_response.tools_call_name,
-            llm_response.tools_call_args,
-            llm_response.tools_call_ids,
-        ):
-            try:
-                func_tool = req.func_tool.get_func(func_tool_name)
-                if func_tool.origin == "mcp":
-                    logger.info(
-                        f"从 MCP 服务 {func_tool.mcp_server_name} 调用工具函数：{func_tool.name}，参数：{func_tool_args}"
-                    )
-                    client = req.func_tool.mcp_client_dict[func_tool.mcp_server_name]
-                    res = await client.session.call_tool(func_tool.name, func_tool_args)
-                    if res:
-                        # TODO 仅对ImageContent | EmbeddedResource进行了简单的Fallback
-                        if isinstance(res.content[0], TextContent):
-                            tool_call_result.append(
-                                ToolCallMessageSegment(
-                                    role="tool",
-                                    tool_call_id=func_tool_id,
-                                    content=res.content[0].text,
-                                )
-                            )
-                        elif isinstance(res.content[0], ImageContent):
-                            tool_call_result.append(
-                                ToolCallMessageSegment(
-                                    role="tool",
-                                    tool_call_id=func_tool_id,
-                                    content="返回了图片(已直接发送给用户)",
-                                )
-                            )
-                            event.set_extra(
-                                "tool_call_img_respond",
-                                res.content[0].data,
-                            )
-                        elif isinstance(res.content[0], EmbeddedResource):
-                            resource = res.content[0].resource
-                            if isinstance(resource, TextResourceContents):
-                                tool_call_result.append(
-                                    ToolCallMessageSegment(
-                                        role="tool",
-                                        tool_call_id=func_tool_id,
-                                        content=resource.text,
-                                    )
-                                )
-                            elif (
-                                isinstance(resource, BlobResourceContents)
-                                and resource.mimeType
-                                and resource.mimeType.startswith("image/")
-                            ):
-                                tool_call_result.append(
-                                    ToolCallMessageSegment(
-                                        role="tool",
-                                        tool_call_id=func_tool_id,
-                                        content="返回了图片(已直接发送给用户)",
-                                    )
-                                )
-                                event.set_extra(
-                                    "tool_call_img_respond",
-                                    res.content[0].data,
-                                )
-                            else:
-                                tool_call_result.append(
-                                    ToolCallMessageSegment(
-                                        role="tool",
-                                        tool_call_id=func_tool_id,
-                                        content="返回的数据类型不受支持",
-                                    )
-                                )
-                else:
-                    # 获取处理器，过滤掉平台不兼容的处理器
-                    platform_id = event.get_platform_id()
-                    star_md = star_map.get(func_tool.handler_module_path)
-                    if (
-                        star_md
-                        and platform_id in star_md.supported_platforms
-                        and not star_md.supported_platforms[platform_id]
-                    ):
-                        logger.debug(
-                            f"处理器 {func_tool_name}({star_md.name}) 在当前平台不兼容或者被禁用，跳过执行"
-                        )
-                        # 直接跳过，不添加任何消息到tool_call_result
-                        continue
-
-                    logger.info(
-                        f"调用工具函数：{func_tool_name}，参数：{func_tool_args}"
-                    )
-                    # 尝试调用工具函数
-                    wrapper = self._call_handler(
-                        self.ctx, event, func_tool.handler, **func_tool_args
-                    )
-                    async for resp in wrapper:
-                        if resp is not None:  # 有 return 返回
-                            tool_call_result.append(
-                                ToolCallMessageSegment(
-                                    role="tool",
-                                    tool_call_id=func_tool_id,
-                                    content=resp,
-                                )
-                            )
-                        else:
-                            res = event.get_result()
-                            if res and res.chain:
-                                event.set_extra("tool_call_result", res)
-                            yield  # 有生成器返回
-                event.clear_result()  # 清除上一个 handler 的结果
-            except BaseException as e:
-                logger.warning(traceback.format_exc())
-                tool_call_result.append(
-                    ToolCallMessageSegment(
-                        role="tool",
-                        tool_call_id=func_tool_id,
-                        content=f"error: {str(e)}",
-                    )
-                )
-        if tool_call_result:
-            # 函数调用结果
-            req.func_tool = None  # 暂时不支持递归工具调用
-            assistant_msg_seg = AssistantMessageSegment(
-                role="assistant", tool_calls=llm_response.to_openai_tool_calls()
-            )
-            # 在多轮 Tool 调用的情况下，这里始终保持最新的 Tool 调用结果，减少上下文长度。
-            req.tool_calls_result = ToolCallsResult(
-                tool_calls_info=assistant_msg_seg,
-                tool_calls_result=tool_call_result,
-            )
-            yield req  # 再次执行 LLM 请求
-        else:
-            if llm_response.completion_text:
-                event.set_result(
-                    MessageEventResult().message(llm_response.completion_text)
-                )
-
-    async def _save_to_history(
-        self, event: AstrMessageEvent, req: ProviderRequest, llm_response: LLMResponse
-    ):
-        if not req or not req.conversation or not llm_response:
-            return
-
-        if llm_response.role == "assistant":
-            # 文本回复
-            contexts = req.contexts.copy()
-            contexts.append(await req.assemble_context())
-
-            # 记录并标记函数调用结果
-            if req.tool_calls_result:
-                tool_calls_messages = req.tool_calls_result.to_openai_messages()
-
-                # 添加标记
-                for message in tool_calls_messages:
-                    message["_tool_call_history"] = True
-
-                processed_tool_messages = self._process_tool_message_pairs(
-                    tool_calls_messages, remove_tags=False
-                )
-
-                contexts.extend(processed_tool_messages)
-
-            contexts.append(
-                {"role": "assistant", "content": llm_response.completion_text}
-            )
-            contexts_to_save = list(
-                filter(lambda item: "_no_save" not in item, contexts)
-            )
-            await self.conv_manager.update_conversation(
-                event.unified_msg_origin, req.conversation.cid, history=contexts_to_save
-            )
-
-    def _process_tool_message_pairs(self, messages, remove_tags=True):
-        """处理工具调用消息，确保assistant和tool消息成对出现
-
-        Args:
-            messages (list): 消息列表
-            remove_tags (bool): 是否移除_tool_call_history标记
-
-        Returns:
-            list: 处理后的消息列表，保证了assistant和对应tool消息的成对出现
-        """
-        result = []
-        i = 0
-
-        while i < len(messages):
-            current_msg = messages[i]
-
-            # 普通消息直接添加
-            if "_tool_call_history" not in current_msg:
-                result.append(current_msg.copy() if remove_tags else current_msg)
-                i += 1
-                continue
-
-            # 工具调用消息成对处理
-            if current_msg.get("role") == "assistant" and "tool_calls" in current_msg:
-                assistant_msg = current_msg.copy()
-
-                if remove_tags and "_tool_call_history" in assistant_msg:
-                    del assistant_msg["_tool_call_history"]
-
-                related_tools = []
-                j = i + 1
-                while (
-                    j < len(messages)
-                    and messages[j].get("role") == "tool"
-                    and "_tool_call_history" in messages[j]
-                ):
-                    tool_msg = messages[j].copy()
-
-                    if remove_tags:
-                        del tool_msg["_tool_call_history"]
-
-                    related_tools.append(tool_msg)
-                    j += 1
-
-                # 成对的时候添加到结果
-                if related_tools:
-                    result.append(assistant_msg)
-                    result.extend(related_tools)
-
-                i = j  # 跳过已处理
-            else:
-                # 单独的tool消息
-                i += 1
-
-        return result