From 9fd1d19e935d238a7ca20820e3aac71c6d24d7a2 Mon Sep 17 00:00:00 2001 From: Raven95676 Date: Mon, 7 Apr 2025 11:52:29 +0800 Subject: [PATCH] =?UTF-8?q?=E5=88=86=E7=A6=BB=E6=B5=81=E5=BC=8F=E4=B8=8E?= =?UTF-8?q?=E9=9D=9E=E6=B5=81=E5=BC=8F=E5=93=8D=E5=BA=94=E5=A4=84=E7=90=86?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../process_stage/method/llm_request.py | 249 +++++++++++------- 1 file changed, 150 insertions(+), 99 deletions(-) diff --git a/astrbot/core/pipeline/process_stage/method/llm_request.py b/astrbot/core/pipeline/process_stage/method/llm_request.py index 25cfc722..fc74966a 100644 --- a/astrbot/core/pipeline/process_stage/method/llm_request.py +++ b/astrbot/core/pipeline/process_stage/method/llm_request.py @@ -189,15 +189,28 @@ class LLMRequestSubStage(Stage): ) return - 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 + 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( @@ -210,14 +223,6 @@ class LLMRequestSubStage(Stage): # 保存到历史记录 await self._save_to_history(event, req, final_llm_response) - # 流式输出完成后，将完整的LLM响应设置为事件结果 - if bool(self.streaming_response): - event.clear_result() - async for _ in self._handle_llm_response( - event, req, final_llm_response - ): - pass - except BaseException as e: logger.error(traceback.format_exc()) event.set_result( @@ -243,38 +248,28 @@ class LLMRequestSubStage(Stage): event: AstrMessageEvent, req: ProviderRequest, llm_response: LLMResponse, - ) -> AsyncGenerator[None, None]: - """处理 LLM 响应。 + ) -> AsyncGenerator[Union[None, ProviderRequest], None]: + """处理非流式 LLM 响应。 Returns: - bool: 是否需要继续调用 LLM + AsyncGenerator[Union[None, ProviderRequest], None]: 如果返回 ProviderRequest，表示需要再次调用 LLM Yields: - Iterator[bool]: 将 event 交付给下一个 stage + Iterator[Union[None, ProviderRequest]]: 将 event 交付给下一个 stage 或者返回 ProviderRequest 表示需要再次调用 LLM """ - is_stream = bool(self.streaming_response) - 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 - if is_stream - else ResultContentType.LLM_RESULT - ) + ).set_result_content_type(ResultContentType.LLM_RESULT) ) else: event.set_result( MessageEventResult() .message(llm_response.completion_text) - .set_result_content_type( - ResultContentType.STREAMING_FINISH - if is_stream - else ResultContentType.LLM_RESULT - ) + .set_result_content_type(ResultContentType.LLM_RESULT) ) elif llm_response.role == "err": event.set_result( @@ -283,83 +278,139 @@ class LLMRequestSubStage(Stage): ) ) elif llm_response.role == "tool": - # function calling - tool_call_result: list[ToolCallMessageSegment] = [] - logger.info( - f"触发 {len(llm_response.tools_call_name)} 个函数调用: {llm_response.tools_call_name}" + # 处理函数工具调用 + 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}" + ) ) - 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}" + 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 content的类型可能包括list[TextContent | ImageContent | EmbeddedResource]，这里只处理了TextContent。 + tool_call_result.append( + ToolCallMessageSegment( + role="tool", + tool_call_id=func_tool_id, + content=res.content[0].text, + ) ) - 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 content的类型可能包括list[TextContent | ImageContent | EmbeddedResource]，这里只处理了TextContent。 + else: + 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=res.content[0].text, + content=resp, ) ) - else: - 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: - 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)}", - ) + else: + 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, + 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) ) - 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