diff --git a/utils/agent.py b/utils/agent.py index 0ac8a62..ead6312 100644 --- a/utils/agent.py +++ b/utils/agent.py @@ -6,12 +6,35 @@ # 列举导入模块 from pathlib import Path import time -from typing import AsyncGenerator, List, Optional +from typing import AsyncGenerator, List, Optional, Dict, Literal from uuid import uuid4 from pydantic_ai.messages import ModelMessage from pydantic_ai import Agent as PydanticAIAgent, ModelMessage from pydantic_ai.capabilities import AgentCapability -from pydantic_ai.messages import ModelMessagesTypeAdapter, ModelResponse +from asyncio import Task, create_task +from pydantic import Field + +from pydantic import BaseModel + + +from pydantic_ai.messages import ( + ModelMessagesTypeAdapter, + ModelResponse, + PartStartEvent, + PartDeltaEvent, + PartEndEvent, + TextPart, + ToolSearchCallPart, + NativeToolSearchCallPart, + LoadCapabilityCallPart, + ThinkingPart, + ToolCallPart, + NativeToolCallPart, + TextPartDelta, + ThinkingPartDelta, + ModelResponseStreamEvent, + ToolReturnPart, +) from pydantic_ai.models.openai import OpenAIChatModel from pydantic_ai.output import OutputSpec from pydantic_ai.providers.openai import OpenAIProvider @@ -80,11 +103,11 @@ class AgentMemory(SQLite): except Exception as exception: raise RuntimeError(f"新增对话消息发生异常：{str(exception)}") from exception - def read_message_history(self, session_id: str) -> List[ModelMessage]: + def get_message_history(self, session_id: str) -> List[ModelMessage]: """ - 查询对话消息历史 + 获取指定会话的消息历史 :param session_id: 会话唯一标识 - :return: 对话消息历史 + :return: 消息历史 """ try: with self: @@ -109,6 +132,91 @@ class AgentMemory(SQLite): ) from exception +class Part(BaseModel): + """ + 片段模型 + """ + + type: Literal["thinking", "text"] = Field(..., description="片段类型") + content: str = Field(..., description="片段内容") + timer: Optional[Task] = Field(default=None, description="片段绑定的延时异步任务") + + +class StreamEventsDebouncer: + """ + 流式事件防抖器 + """ + + def __init__(self, debounce_by: float = 0.2): + """ + 初始化 + :param debounce_by: 防抖间隔（秒） + """ + # 片段缓存 + self.parts: Dict[int, Part] = {} + + async def handle_event(self, event): + """ + 处理流式事件 + :param event: 流式事件 + :yield: 处理后流式事件 + """ + # 若为片段开始事件且片段类型为思考或纯文本则创建片段缓存，否则直接返回 + if isinstance(event, PartStartEvent): + index, part = event.index, event.part + match part: + case ThinkingPart(): + part_type = "thinking" + case TextPart(): + part_type = "text" + case _: + yield event + return + + # 缓存片段 + self.parts[index] = Part( + type=part_type, + content=part.content if part.content else "", + ) + # 片段首次触发直接推送原始事件，保证首条内容立刻展示 + yield event + return + + # 若为片段增量事件且增量类型为思考或纯文本则在片段缓存的片段内容追加 + if isinstance(event, PartDeltaEvent): + index, delta = event.index, event.delta + part = self.parts.get(index) + if isinstance(delta, (ThinkingPartDelta, TextPartDelta)): + # 增量内容追加至缓冲区 + part.content += delta.content_delta + # 重置防抖计时器，实现滑动窗口效果 + self._cancel_timer(index) + part.timer = create_task(self._delay_flush(index)) + else: + # 非文本类增量事件直接透传 + yield event + return + + # 命中强制刷新规则：先推送所有缓存聚合内容，再透传当前触发事件 + if isinstance(event, FLUSH_TRIGGER_TYPES): + async for flush_event in self.flush_all(): + yield flush_event + yield event + return + + # 其余未匹配的所有事件原样透传 + yield event + + def _cancel_timer(self, index: int): + """ + 取消指定索引片段的计时器 + :param index: 片段索引 + """ + part = self.parts.get(index) + if part and part.timer is not None and not part.timer.done(): + part.timer.cancel() + + class Agent: """ 智能体，支持： @@ -154,55 +262,47 @@ class Agent: self.agent.to_web() - async def output_message_streamed( + async def stream_messages_events( self, user_prompt: str | List[str] ) -> AsyncGenerator[str, None]: """ - 智能体流式输出消息 + 流式输出消息事件 :param user_prompt: 用户提示词（用户输入消息） - :return: 流式消息 + :return: 消息事件 """ """定义：一次会话（session）包含若干论对话（turn），每一轮对话由用户输入消息（message）和智能体输出消息组成""" - # 查询该会话历史对话消息列表 - message_history = self.agent_memory.read_message_history( + # 获取指定会话的消息历史 + message_history = self.agent_memory.get_message_history( session_id=self.session_id ) async with self.agent.run_stream_events( user_prompt=user_prompt, message_history=message_history, - ) as result: - async for response in result.stream_response(): # 完整结构化相应对象 - if not isinstance(response, ModelResponse): - continue - - if not hasattr(response, "parts") or not response.parts: - continue - - for part in response.parts: - if isinstance(part, str): - content = part.strip() - if content: + ) as events: + async for event in events: + # 处理片段开始事件 + if isinstance(event, PartStartEvent): + part = event.part + if isinstance(part, ThinkingPart): # 处理思考片段 + if content := part.content: yield f"00:{content}" - continue - - match part.part_kind: - case "text": - content = part.content.strip() - if content: - yield f"00:{content}" - case "thinking": - content = part.content.strip() - if content: - yield f"01:{content}" - case "tool-call" | "builtin-tool-call": - yield f"02:技能名称：{part.tool_name}，调用参数：{part.args}" - case "tool-return" | "builtin-tool-return": - yield f"03:技能结果：{part.content}" - case _: - continue + elif isinstance(part, TextPart): # 处理文本片段 + if content := part.content: + yield f"01:{content}" + elif isinstance( + part, + ( + NativeToolSearchCallPart, + NativeToolCallPart, + ToolSearchCallPart, + ToolCallPart, + LoadCapabilityCallPart, + ), + ): # 处理原生搜索、原生工具、自封搜索、自封工具和加载能力调用片段 + yield f"02:技能名称：{part.tool_name}" self.agent_memory.create_new_messages( session_id=self.session_id, - new_messages=result.new_messages(), + new_messages=events.new_messages(), )