369 lines
12 KiB
Python
369 lines
12 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
智能体模块
|
||
"""
|
||
|
||
# 列举导入模块
|
||
from pathlib import Path
|
||
import time
|
||
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 asyncio import Task, create_task, sleep, Queue, QueueEmpty
|
||
from pydantic import Field
|
||
|
||
from pydantic import BaseModel
|
||
|
||
|
||
from pydantic_ai.messages import (
|
||
ModelMessagesTypeAdapter,
|
||
ModelResponse,
|
||
ModelMessageEvent,
|
||
PartStartEvent,
|
||
PartDeltaEvent,
|
||
PartEndEvent,
|
||
TextPart,
|
||
ToolSearchCallPart,
|
||
NativeToolSearchCallPart,
|
||
LoadCapabilityCallPart,
|
||
ThinkingPart,
|
||
ToolCallPart,
|
||
NativeToolCallPart,
|
||
TextPartDelta,
|
||
ThinkingPartDelta,
|
||
ModelResponseStreamEvent,
|
||
ToolReturnPart,
|
||
AgentRunResultEvent,
|
||
NativeToolReturnPart,
|
||
)
|
||
from pydantic_ai.models.openai import OpenAIChatModel
|
||
from pydantic_ai.output import OutputSpec
|
||
from pydantic_ai.providers.openai import OpenAIProvider
|
||
|
||
from .sqlite import SQLite
|
||
|
||
|
||
class AgentMemory(SQLite):
|
||
"""
|
||
智能体记忆体，支持：
|
||
create：新增对话消息
|
||
read：查询会话历史消息
|
||
"""
|
||
|
||
def __init__(self):
|
||
"""
|
||
初始化
|
||
"""
|
||
# 构建数据库路径
|
||
super().__init__(database=Path(__file__).parent.resolve() / "agent_memory.db")
|
||
|
||
try:
|
||
with self:
|
||
self.execute(
|
||
sql="""
|
||
CREATE TABLE IF NOT EXISTS new_messages
|
||
(
|
||
--唯一标识
|
||
id TEXT PRIMARY KEY,
|
||
--会话唯一标识
|
||
session_id TEXT NOT NULL,
|
||
--新对话消息
|
||
new_messages TEXT NOT NULL,
|
||
--时间戳（毫秒）
|
||
timestamp INTEGER NOT NULL
|
||
)
|
||
"""
|
||
)
|
||
except Exception as exception:
|
||
raise RuntimeError(f"初始化数据库发生异常：{str(exception)}") from exception
|
||
|
||
def create_new_messages(
|
||
self, session_id: str, new_messages: List[ModelMessage]
|
||
) -> bool:
|
||
"""
|
||
新增新对话消息
|
||
:param session_id: 会话唯一标识
|
||
:param new_messages: 新对话消息
|
||
:return: 新增是否成功
|
||
"""
|
||
try:
|
||
with self:
|
||
return self.execute(
|
||
sql="""
|
||
INSERT INTO new_messages (id, session_id, new_messages, timestamp) VALUES (?, ?, ?, ?)
|
||
""",
|
||
parameters=(
|
||
uuid4().hex.lower(),
|
||
session_id,
|
||
ModelMessagesTypeAdapter.dump_json(new_messages),
|
||
int(time.time() * 1000),
|
||
),
|
||
)
|
||
except Exception as exception:
|
||
raise RuntimeError(f"新增对话消息发生异常：{str(exception)}") from exception
|
||
|
||
def get_message_history(self, session_id: str) -> List[ModelMessage]:
|
||
"""
|
||
获取指定会话的消息历史
|
||
:param session_id: 会话唯一标识
|
||
:return: 消息历史
|
||
"""
|
||
try:
|
||
with self:
|
||
result = self.query_all(
|
||
sql="""
|
||
SELECT new_messages
|
||
FROM new_messages
|
||
WHERE session_id = ?
|
||
ORDER BY timestamp ASC
|
||
""",
|
||
parameters=(session_id,),
|
||
)
|
||
message_history = []
|
||
for row in result:
|
||
message_history.extend(
|
||
ModelMessagesTypeAdapter.validate_json(row["new_messages"])
|
||
)
|
||
return message_history
|
||
except Exception as exception:
|
||
raise RuntimeError(
|
||
f"查询会话历史消息发生异常：{str(exception)}"
|
||
) from exception
|
||
|
||
|
||
class PendingFlushPart(BaseModel):
|
||
"""
|
||
待刷新模型消息片段
|
||
"""
|
||
|
||
type_: Literal["thinking", "text"] = Field(..., description="片段类型")
|
||
content_delta: str = Field(default="", description="片段内容增量")
|
||
task: Optional[Task] = Field(default=None, description="异步协程任务")
|
||
|
||
|
||
class EventsDebouncer:
|
||
"""
|
||
模型消息事件防抖器
|
||
使用 Pydantic-AI 中 run_stream_events 方法时需就其返回的模型消息事件处理，延迟刷新片段类型为思考或文本的片段增量事件以实现防抖（打字机效果）
|
||
在 Pydantic-AI 中模型流式输出包含若干片段，例如思考片段、文本片段等。每个片段包含若干模型消息事件，例如片段开始事件、片段增量事件、片段结束事件等
|
||
"""
|
||
|
||
def __init__(self, flush_delay: float = 0.2):
|
||
"""
|
||
初始化
|
||
:param flush_delay: 刷新延迟（单位为秒）
|
||
"""
|
||
self.flush_delay = flush_delay
|
||
|
||
# 待刷新模型响应片段
|
||
self.pending_flush_parts: Dict[int, PendingFlushPart] = {}
|
||
# 待刷新列队
|
||
self.pending_flush_queue = Queue()
|
||
|
||
async def handle_event(
|
||
self, event: ModelMessageEvent
|
||
) -> AsyncGenerator[ModelMessageEvent, None]:
|
||
"""
|
||
处理模型消息事件
|
||
:param event: 模型消息事件
|
||
:yield: AsyncGenerator
|
||
"""
|
||
# 返回待刷新列队中模型消息事件
|
||
while True:
|
||
try:
|
||
yield self.pending_flush_queue.get_nowait()
|
||
except QueueEmpty:
|
||
break
|
||
|
||
# 若为思考或文本的片段开始事件则先缓存片段再返回，其它事件直接返回
|
||
if isinstance(event, PartStartEvent):
|
||
index, part = event.index, event.part # 片段索引、片段
|
||
match part:
|
||
case ThinkingPart():
|
||
type_ = "thinking"
|
||
case TextPart():
|
||
type_ = "text"
|
||
case _:
|
||
yield event
|
||
return
|
||
# 新增片段缓存
|
||
self.pending_flush_parts[index] = PendingFlushPart(
|
||
type_=type_,
|
||
)
|
||
yield event
|
||
return
|
||
|
||
# 若为思考或文本的片段增量事件则更新缓存，其它事件直接返回
|
||
if isinstance(event, PartDeltaEvent):
|
||
index, delta = event.index, event.delta
|
||
pending_flush_part = self.pending_flush_parts.get(index)
|
||
if pending_flush_part and isinstance(
|
||
delta, (ThinkingPartDelta, TextPartDelta)
|
||
):
|
||
pending_flush_part.content_delta += delta.content_delta or ""
|
||
|
||
# 原理：若持续输出片段增量事件则拦截并重新创建异步协程任务，至输出间隔超过刷新延迟则返回
|
||
# 取消异步协程任务
|
||
self._cancel_task(index=index)
|
||
# 创建异步协程任务
|
||
pending_flush_part.task = create_task(
|
||
coro=self._delay_flush(index=index)
|
||
)
|
||
else:
|
||
yield event
|
||
return
|
||
|
||
# 若为其它事件则先刷新所有片段缓存再直接返回
|
||
for index in list(self.pending_flush_parts.keys()):
|
||
flush_event = await self._flush(index=index)
|
||
if flush_event:
|
||
yield flush_event
|
||
|
||
# 若为片段结束事件则删除该片段缓存
|
||
if isinstance(event, PartEndEvent):
|
||
self.part_caches.pop(event.index, None)
|
||
return
|
||
|
||
def _cancel_task(self, index: int) -> None:
|
||
"""
|
||
取消异步协程任务
|
||
:param index: 片段索引
|
||
:yield: None
|
||
"""
|
||
pending_flush_part = self.pending_flush_parts.get(index)
|
||
if not pending_flush_part or not pending_flush_part.task:
|
||
return
|
||
|
||
# 若异步协程任务未完成则取消
|
||
if not pending_flush_part.task.done():
|
||
pending_flush_part.task.cancel()
|
||
|
||
pending_flush_part.task = None
|
||
|
||
async def _delay_flush(self, index: int):
|
||
"""
|
||
延迟刷新
|
||
:param index: 片段索引
|
||
:return: None
|
||
"""
|
||
await sleep(delay=self.flush_delay)
|
||
flush_event = await self._flush(index=index)
|
||
if flush_event:
|
||
await self.pending_flush_queue.put(item=flush_event)
|
||
|
||
async def _flush(self, index: int) -> Optional[PartDeltaEvent]:
|
||
"""
|
||
刷新
|
||
:param index: 片段索引
|
||
:return: Optional[PartDeltaEvent]
|
||
"""
|
||
pending_flush_part = self.pending_flush_parts.get(index)
|
||
if not pending_flush_part or not pending_flush_part.content_delta:
|
||
return None
|
||
|
||
match pending_flush_part.type_:
|
||
case "thinking":
|
||
delta = ThinkingPartDelta(content_delta=pending_flush_part.content_delta)
|
||
case "text":
|
||
delta = TextPartDelta(content_delta=pending_flush_part.content_delta)
|
||
|
||
pending_flush_part.content_delta = ""
|
||
pending_flush_part.task = None
|
||
return PartDeltaEvent(index=index, delta=delta)
|
||
|
||
|
||
|
||
|
||
|
||
class Agent:
|
||
"""
|
||
智能体，支持：
|
||
1 实例智能体
|
||
2 异步运行
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
session_id: str,
|
||
instructions: str,
|
||
output_type: OutputSpec = str,
|
||
capabilities: Optional[List[AgentCapability]] = None,
|
||
):
|
||
"""
|
||
初始化智能体
|
||
:param session_id: 会话唯一标识
|
||
:param instructions: 指令
|
||
:param skills: 智能体技能列表，默认为不使用技能
|
||
:param output_type: 输出类型
|
||
:return: 智能体实例
|
||
"""
|
||
# 会话唯一标识
|
||
self.session_id = session_id
|
||
|
||
# 实例智能体的记忆体
|
||
self.agent_memory = AgentMemory()
|
||
|
||
# 实例智能体
|
||
self.agent = PydanticAIAgent(
|
||
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=1,
|
||
)
|
||
|
||
self.agent.to_web()
|
||
|
||
async def stream_messages_events(
|
||
self, user_prompt: str | List[str]
|
||
) -> AsyncGenerator[str, None]:
|
||
"""
|
||
流式输出消息事件
|
||
:param user_prompt: 用户提示词（用户输入消息）
|
||
:return: 消息事件
|
||
"""
|
||
"""定义：一次会话（session）包含若干论对话（turn），每一轮对话由用户输入消息（message）和智能体输出消息组成"""
|
||
# 获取指定会话的消息历史
|
||
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 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}"
|
||
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=events.new_messages(),
|
||
)
|