360 lines
19 KiB
Python
360 lines
19 KiB
Python
# -*- 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()
|
||
|
||
async with self.agent.iter(
|
||
conversation_id=conversation_id,
|
||
user_prompt=user_prompt,
|
||
message_history=message_history,
|
||
) as run:
|
||
# 当前运行唯一标识
|
||
current_run_id = run.ctx.state.run_id
|
||
yield Event(
|
||
event_kind=EventKind.RUN_START,
|
||
run_id=current_run_id,
|
||
)
|
||
async for node in run:
|
||
if self.agent.is_model_request_node(
|
||
node=node
|
||
) or self.agent.is_call_tools_node(node=node):
|
||
async with node.stream(run.ctx) as stream:
|
||
# 将原始流转为事件流
|
||
stream = run.ctx.deps.root_capability.wrap_run_event_stream(
|
||
ctx=agent_graph.build_run_context(ctx=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=run.new_messages(),
|
||
)
|