# -*- coding: utf-8 -*- """ Pydantic AI 智能体和相关模块 """ # 列举导入模块 from typing import AsyncGenerator, List, Optional from pydantic_ai import _agent_graph as agent_graph from pydantic_ai import Agent, ThinkingPartDelta from pydantic_ai.capabilities import AgentCapability from pydantic_ai.messages import ( FunctionToolCallEvent, FunctionToolResultEvent, LoadCapabilityCallPart, ModelMessage, PartDeltaEvent, PartEndEvent, PartStartEvent, TextPart, TextPartDelta, ThinkingPart, ToolCallPart, ToolSearchCallPart, ) from pydantic_ai.models.openai import OpenAIChatModel from pydantic_ai.output import OutputSpec from pydantic_ai.providers.openai import OpenAIProvider from application.models import Event, EventKind, PartKind DEFAULT_INSTRUCTIONS: str = """ # 角色 专业友好AI助手，结构化解答各类问题。 # 输出硬性规则 1. 全文强制标准Markdown，禁止纯文本；不要额外说明排版格式，直接输出内容； 2. 层级使用 `#/##/###`，列表用 `-` 无序列表或数字有序列表； 3. 代码块用 ```语言名``` 包裹； 4. 重点内容标注 **粗体**/*斜体*； 5. 思考、工具日志仅输出文本，适配前端折叠面板，禁止输出HTML标签； 6. 内容分点拆分，排版整洁适配前端Markdown渲染。 # 行文要求 语言通俗，逻辑完整简洁，无多余废话。 """ class AIAgent: """ 基于 Pydantic AI 封装的智能体 """ def __init__( self, instructions: Optional[str] = None, capabilities: Optional[List[AgentCapability]] = None, output_type: OutputSpec = str, retries: int = 1, ): """ 初始化 :param instructions: 指令 :param capabilities: 技能列表，默认为不使用技能 :param output_type: 输出类型 :param retries: 重试次数，默认为1次 :return: 智能体实例 """ # 一次会话（conversation）包含若干次运行（run），每次运行包含用户提示词（user_prompt）和输出（output）。其中，输出包含若干片分片（Part） # 本次运行工具调用唯一标识与片段索引映射表 self.tool_call_ids: dict[str, int] = {} # 若指令为空则使用默认指令 if not instructions: instructions = DEFAULT_INSTRUCTIONS # 初始化智能体 self.agent = Agent( model=OpenAIChatModel( model_name="deepseek-v4-flash", provider=OpenAIProvider( base_url="https://tokenhub.tencentmaas.com/v1", api_key="sk-D9Y1mCe8VlvNqLuSC4mAjqEwxJ2nW4C0h8a7EPn8kg9RLsHq", ), ), instructions=instructions, capabilities=capabilities, output_type=output_type, retries=retries, ) async def run( self, conversation_id: str, user_prompt: str | List[str], message_history: Optional[List[ModelMessage]] = None, ) -> AsyncGenerator[Event]: """ 运行 :param conversation_id: 会话唯一标识 :param user_prompt: 用户提示词（用户提示词） :param message_history: 消息历史，默认为None :yield: AsyncGenerator[Event] """ # 重置工具调用唯一标识与片段索引映射表 self.tool_call_ids.clear() self.agent.run() async with self.agent.iter( conversation_id=conversation_id, user_prompt=user_prompt, message_history=message_history, ) as agent_run: # 当前运行唯一标识 current_run_id = agent_run.ctx.state.run_id yield Event( event_kind=EventKind.RUN_START, run_id=current_run_id, ) async for node in agent_run: if self.agent.is_model_request_node( node=node ) or self.agent.is_call_tools_node(node=node): async with node.stream(agent_run.ctx) as stream: # 将原始流转为事件流 stream = ( agent_run.ctx.deps.root_capability.wrap_run_event_stream( ctx=agent_graph.build_run_context(ctx=agent_run.ctx), stream=stream, ) ) async for event in stream: match event: # ========== 分片开始事件 ========== case PartStartEvent( event_kind=event_kind, index=index, part=part, previous_part_kind=previous_part_kind, ): match part: # 思考分片开始事件 case ThinkingPart( part_kind=part_kind, content=content ): yield Event( event_kind=EventKind(event_kind), event_content=content, part_index=index, previous_part_kind=PartKind( previous_part_kind ), part_kind=PartKind(part_kind), run_id=current_run_id, ) # 检索工具分片开始事件 case ToolSearchCallPart( tool_kind=tool_kind, tool_call_id=tool_call_id, tool_name=tool_name, ): # 记录工具调用唯一标识与片段索引映射 self.tool_call_ids[tool_call_id] = index yield Event( event_kind=EventKind(event_kind), part_index=index, previous_part_kind=PartKind( previous_part_kind ), part_kind=PartKind(tool_kind), tool_name=tool_name, run_id=current_run_id, ) # 加载能力分片开始事件 case LoadCapabilityCallPart( tool_kind=tool_kind, tool_call_id=tool_call_id, tool_name=tool_name, ): # 记录工具调用唯一标识与片段索引映射 self.tool_call_ids[tool_call_id] = index yield Event( event_kind=EventKind(event_kind), part_index=index, previous_part_kind=PartKind( previous_part_kind ), part_kind=PartKind(tool_kind), tool_name=tool_name, run_id=current_run_id, ) # 调用工具分片开始事件 case ToolCallPart( part_kind=part_kind, tool_call_id=tool_call_id, tool_name=tool_name, ): # 记录工具调用唯一标识与片段索引映射 self.tool_call_ids[tool_call_id] = index yield Event( event_kind=EventKind(event_kind), part_index=index, previous_part_kind=PartKind( previous_part_kind ), part_kind=PartKind(part_kind), tool_name=tool_name, run_id=current_run_id, ) # 文本分片开始事件 case TextPart( part_kind=part_kind, content=content ): yield Event( event_kind=EventKind(event_kind), event_content=content, part_index=index, previous_part_kind=PartKind( previous_part_kind ), part_kind=PartKind(part_kind), run_id=current_run_id, ) # ========== 分片增量事件（前端仅就文本片段实现打字机效果） ========== case PartDeltaEvent( event_kind=event_kind, index=index, delta=delta ): match delta: # 思考分片增量事件 case ThinkingPartDelta( part_delta_kind=part_delta_kind, content_delta=content_delta, ): yield Event( event_kind=EventKind(event_kind), event_content=content_delta or "", part_index=index, part_kind=PartKind(part_delta_kind), run_id=current_run_id, ) # 文本分片增量事件 case TextPartDelta( part_delta_kind=part_delta_kind, content_delta=content_delta, ): yield Event( event_kind=EventKind(event_kind), event_content=content_delta, part_index=index, part_kind=PartKind(part_delta_kind), run_id=current_run_id, ) # ========== 分片结束事件 ========== case PartEndEvent( event_kind=event_kind, index=index, part=part, next_part_kind=next_part_kind, ): match part: # 思考分片结束事件 case ThinkingPart( part_kind=part_kind, content=content ): yield Event( event_kind=EventKind(event_kind), event_content=content, part_index=index, part_kind=PartKind(part_kind), next_part_kind=PartKind(next_part_kind), run_id=current_run_id, ) # 检索工具分片结束事件 case ToolSearchCallPart( tool_kind=tool_kind, tool_name=tool_name ): yield Event( event_kind=EventKind(event_kind), part_index=index, part_kind=PartKind(tool_kind), next_part_kind=PartKind(next_part_kind), tool_name=tool_name, run_id=current_run_id, ) # 加载能力分片结束事件 case LoadCapabilityCallPart( tool_kind=tool_kind, tool_name=tool_name ): yield Event( event_kind=EventKind(event_kind), part_index=index, part_kind=PartKind(tool_kind), next_part_kind=PartKind(next_part_kind), tool_name=tool_name, run_id=current_run_id, ) # 调用工具分片结束事件 case ToolCallPart( part_kind=part_kind, tool_name=tool_name ): yield Event( event_kind=EventKind(event_kind), part_index=index, part_kind=PartKind(part_kind), next_part_kind=PartKind(next_part_kind), tool_name=tool_name, run_id=current_run_id, ) # 文本分片结束事件 case TextPart( part_kind=part_kind, content=content ): yield Event( event_kind=EventKind(event_kind), event_content=content, part_index=index, part_kind=PartKind(part_kind), next_part_kind=PartKind(next_part_kind), run_id=current_run_id, ) # ========== 函数工具调用事件 ========== case FunctionToolCallEvent( event_kind=event_kind, part=part, ): yield Event( event_kind=EventKind(event_kind), part_index=self.tool_call_ids.get( part.tool_call_id ), part_kind=PartKind(part.part_kind), tool_name=part.tool_name, run_id=current_run_id, ) # ========== 函数工具回调事件 ========== case FunctionToolResultEvent( event_kind=event_kind, content=content, part=part, ): yield Event( event_kind=EventKind(event_kind), event_content=( content if isinstance(content, str) else "" ), # 暂仅处理文本类型，后续处理多模态和缓存 part_index=self.tool_call_ids.get( part.tool_call_id ), part_kind=PartKind(part.part_kind), tool_name=part.tool_name, run_id=current_run_id, ) yield Event( event_kind=EventKind.RUN_END, run_id=current_run_id, run_new_messages=agent_run.new_messages(), )