liubiren
/
Python
1
1
Fork 0
Code Issues Pull Requests Packages Projects Releases Activity

1

This commit is contained in:
liubiren 2026-09-11 20:24:02 +08:00
parent 31828a7be3
commit 1ac121e9cb
11 changed files with 205 additions and 192 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

36
agent/alembic/versions/3a74f94a7ad4_.py
View File

@ -1,36 +0,0 @@
"""empty message
Revision ID: 3a74f94a7ad4
Revises: d52f1ee0f9c4
Create Date: 2026-09-03 12:30:13.538072
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
import sqlmodel
# revision identifiers, used by Alembic.
revision: str = '3a74f94a7ad4'
down_revision: Union[str, Sequence[str], None] = 'd52f1ee0f9c4'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
"""Upgrade schema."""
# ### commands auto generated by Alembic - please adjust! ###
with op.batch_alter_table('conversationrecord', schema=None) as batch_op:
batch_op.add_column(sa.Column('workflow', sqlmodel.sql.sqltypes.AutoString(), server_default=sa.text("('')"), nullable=False))
# ### end Alembic commands ###
def downgrade() -> None:
"""Downgrade schema."""
# ### commands auto generated by Alembic - please adjust! ###
with op.batch_alter_table('conversationrecord', schema=None) as batch_op:
batch_op.drop_column('workflow')
# ### end Alembic commands ###

16
agent/alembic/versions/d52f1ee0f9c4_.py → agent/alembic/versions/3e4ccabe44c3_.py
View File

@ -1,8 +1,8 @@
"""empty message
Revision ID: d52f1ee0f9c4
Revision ID: 3e4ccabe44c3
Revises:
Create Date: 2026-08-11 14:59:37.246745
Create Date: 2026-09-11 14:22:48.429299
"""
from typing import Sequence, Union
@ -12,7 +12,7 @@ import sqlalchemy as sa
import sqlmodel
# revision identifiers, used by Alembic.
revision: str = 'd52f1ee0f9c4'
revision: str = '3e4ccabe44c3'
down_revision: Union[str, Sequence[str], None] = None
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
@ -37,9 +37,10 @@ def upgrade() -> None:
sa.Column('id', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
sa.Column('user_id', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
sa.Column('description', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
sa.Column('usage', sa.JSON(), nullable=False),
sa.Column('is_deleted', sa.Boolean(), nullable=False),
sa.Column('created_at', sa.DateTime(), nullable=False),
sa.Column('is_deleted', sa.Boolean(), nullable=False),
sa.Column('usage', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
sa.Column('workflow', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
sa.PrimaryKeyConstraint('id')
)
with op.batch_alter_table('conversationrecord', schema=None) as batch_op:
@ -49,9 +50,10 @@ def upgrade() -> None:
op.create_table('messagerecord',
sa.Column('id', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
sa.Column('conversation_id', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
sa.Column('type', sa.Enum('USER_PROMPT', 'THINKING', 'TOOL_CALL', 'RESULT_OUTPUT', name='messagetype'), nullable=False),
sa.Column('type', sa.Enum('USER_PROMPT', 'THINKING', 'TOOL_CALL', 'TEXT', name='messagetype'), nullable=False),
sa.Column('title', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
sa.Column('content', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
sa.Column('approval', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
sa.PrimaryKeyConstraint('id')
)
with op.batch_alter_table('messagerecord', schema=None) as batch_op:
@ -60,7 +62,7 @@ def upgrade() -> None:
op.create_table('runrecord',
sa.Column('id', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
sa.Column('conversation_id', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
sa.Column('new_messages', sa.JSON(), nullable=False),
sa.Column('new_messages', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
sa.PrimaryKeyConstraint('id')
)
with op.batch_alter_table('runrecord', schema=None) as batch_op:

10
agent/application/pages/conversation.py
View File

@ -475,9 +475,8 @@ def tool_call(item: MessageHistoryItem) -> rx.Component:
rx.button(
"发送",
# 点击事件：审批
on_click=lambda: ConversationState.approve(item.id),
on_click=lambda: ConversationState.approve(item.id, True),
# 是否禁用绑定验证码发送禁用状态
disabled=AuthState.is_captcha_code_sending_button_disabled,
margin="0 8px",
background_color="var(--prismui-background-color-1)",
line_height="16px",
@ -487,11 +486,6 @@ def tool_call(item: MessageHistoryItem) -> rx.Component:
"var(--prismui-color-4)",
"var(--prismui-color-5)",
),
cursor=rx.cond(
AuthState.is_captcha_code_sending_button_disabled,
"not-allowed",
"pointer",
),
),
width="100%",
height="45px",
@ -669,7 +663,7 @@ def message_history() -> rx.Component:
align_items="center",
width="100%",
padding="0 16px",
margin_bottom="8px",
margin="8px 0",
overflow="hidden",
),
rx.fragment(),

6
agent/application/states/auth.py
View File

@ -208,8 +208,9 @@ class AuthState(rx.State):
"""
恢复会话状态
"""
# 恢复会话状态
# 获取会话状态
conversation_state = await self.get_state(ConversationState)
# 恢复会话状态
await conversation_state.resume_conversation_state(user_id=self.user_id)
@rx.event
@ -253,8 +254,9 @@ class AuthState(rx.State):
# 创建用户记录
user_id = await db_state.create_user_record(email=self.email)
# 恢复当前用户会话状态
# 获取会话状态
conversation_state = await self.get_state(ConversationState)
# 恢复当前用户的会话状态
await conversation_state.resume_conversation_state(user_id=user_id)
self.user_id = user_id

171
agent/application/states/conversation.py
View File

@ -4,18 +4,15 @@
"""
from datetime import datetime
from typing import AsyncGenerator, cast
from pydantic_ai import (
DeferredToolRequests,
DeferredToolResults,
ToolApproved,
ToolDenied,
)
from pydantic_ai.messages import (
FunctionToolCallEvent,
FunctionToolResultEvent,
LoadCapabilityCallPart,
ToolReturnPart,
PartDeltaEvent,
PartEndEvent,
PartStartEvent,
@ -24,13 +21,14 @@ from pydantic_ai.messages import (
ThinkingPart,
ThinkingPartDelta,
ToolCallPart,
ToolSearchCallPart,
ToolReturnPart,
)
from pydantic_ai.run import AgentRunResultEvent, AgentRunResult
from pydantic_ai.run import AgentRunResultEvent
import reflex as rx
from application.states.database import DatabaseState
from application.states.models import (
Approval,
Conversation,
ConversationHistoryItem,
Message,
@ -38,11 +36,8 @@ from application.states.models import (
MessageType,
Workflow,
WorkflowType,
Approval,
workflow_dump_json,
workflow_validate_json,
usage_validate_json,
)
from application.workshop.events import AgentPromptEvent
def format_conversation_created_at(created_at: datetime) -> str:
@ -280,34 +275,7 @@ class ConversationState(rx.State):
# 获取数据库状态
db_state = await self.get_db_state()
# 无审批实例
if not approval:
# 若用户提示词为空则直接返回
if not conversation.user_prompt:
return
# 将正在运行设置为是
conversation.is_running = True
# 将等待流式输出设置为是
conversation.awaiting_stream = True
yield # 通知前端更新渲染
# 先创建消息记录再添加消息实例
conversation.messages.update(
await db_state.create_message_record(
conversation_id=conversation.id,
message=Message(
type=MessageType.USER_PROMPT,
content=(user_prompt := conversation.user_prompt),
),
)
)
# 清空用户提示词
conversation.user_prompt = ""
yield # 通知前端更新渲染
# 有审批实例
else:
if isinstance(approval, Approval):
if not (user_prompt := approval.user_prompt):
return
@ -315,14 +283,50 @@ class ConversationState(rx.State):
conversation.awaiting_stream = True
yield # 通知前端更新渲染
else:
# 若用户提示词为空则直接返回
if not (user_prompt := conversation.user_prompt):
return
# 清空用户提示词
conversation.user_prompt = ""
# 将正在运行设置为是
conversation.is_running = True
# 将等待流式输出设置为是
conversation.awaiting_stream = True
# 先创建消息记录再添加消息实例
conversation.messages.update(
await db_state.create_message_record(
conversation_id=conversation.id,
message=Message(
type=MessageType.USER_PROMPT,
content=user_prompt,
),
)
)
yield # 通知前端更新渲染
# 初始化片段索引映射为消息实例唯一标识字典
index_map_to_message_id: dict[int, str] = {}
# 初始化工具调用片段工具名称映射为消息实例唯一标识字典
tool_name_map_to_message_id: dict[str, str] = {}
# 查询消息历史列表
message_history = await db_state.read_message_history(
conversation_id=self.actived_conversation_id
)
# 初始化工作流
if conversation.workflow:
match conversation.workflow.type:
case WorkflowType.TALK:
from application.workshop.talk import (
run_stream_events,
)
# 运行并流式输出事件
stream_events = run_stream_events(
user_prompt=user_prompt,
message_history=message_history,
usage=conversation.usage,
)
case WorkflowType.BOOK_FLIGHT:
from application.workshop.book_flight import (
run_stream_events,
@ -332,25 +336,10 @@ class ConversationState(rx.State):
stream_events = run_stream_events(
deps=conversation.workflow.deps,
user_prompt=user_prompt,
message_history=await db_state.read_message_history(
conversation_id=self.actived_conversation_id
),
message_history=message_history,
usage=conversation.usage,
deferred_tool_results=deferred_tool_results,
)
else:
from application.workshop.talk import (
run_stream_events,
)
# 运行并流式输出事件
stream_events = run_stream_events(
user_prompt=user_prompt,
message_history=await db_state.read_message_history(
conversation_id=self.actived_conversation_id
),
usage=conversation.usage,
)
# 消息列表
messages: list[Message] = []
@ -382,7 +371,9 @@ class ConversationState(rx.State):
yield # 通知前端更新渲染
# 工具调用分片
case ToolCallPart(tool_name=tool_name):
case ToolCallPart(
tool_call_id=tool_call_id, tool_name=tool_name
):
# 构建消息实例
message = Message(
type=MessageType.TOOL_CALL,
@ -391,8 +382,10 @@ class ConversationState(rx.State):
)
# 添加至消息字典
conversation.messages[message.id] = message
# 将消息实例唯一标识与工具调用片段工具名称映射
tool_name_map_to_message_id[tool_name] = message.id
# 将消息实例唯一标识与工具调用唯一标识映射
conversation.workflow.tool_call_id_map_to_message_id[
tool_call_id
] = message.id
yield # 通知前端更新渲染
# 文本分片
@ -436,7 +429,7 @@ class ConversationState(rx.State):
):
match part:
# 思考分片
case ThinkingPart(content=content):
case ThinkingPart():
# 获取消息实例
message = conversation.messages[
index_map_to_message_id[index]
@ -447,7 +440,7 @@ class ConversationState(rx.State):
messages.append(message)
# 文本分片
case TextPart(content=content):
case TextPart():
# 获取消息实例
message = conversation.messages[
index_map_to_message_id[index]
@ -458,13 +451,18 @@ class ConversationState(rx.State):
case FunctionToolResultEvent(part=part):
match part:
# 工具返回分片
case ToolReturnPart(tool_name=tool_name, content=content):
case ToolReturnPart(
tool_call_id=tool_call_id,
tool_name=tool_name,
content=content,
):
# 获取消息实例
message = conversation.messages[
tool_name_map_to_message_id[tool_name]
conversation.workflow.tool_call_id_map_to_message_id[
tool_call_id
]
]
message.is_running = False
yield # 通知前端更新渲染
message.title = f"{tool_name}已完成"
message.content = cast(
str, content
@ -472,26 +470,39 @@ class ConversationState(rx.State):
yield # 通知前端更新渲染
messages.append(message)
# ========== 智能体提示事件 ==========
case AgentPromptEvent(content=content):
# 构建消息实例
message = Message(type=MessageType.TEXT, content=content)
# 添加至消息字典
conversation.messages[message.id] = message
yield # 通知前端更新渲染
messages.append(message)
# ========== 智能体运行结果事件 ==========
case AgentRunResultEvent(result=result):
if isinstance(output := result.output, DeferredToolRequests):
for approval_ in output.approvals:
for appr in output.approvals:
# 获取消息实例
message = conversation.messages[
tool_name_map_to_message_id[approval_.tool_name]
conversation.workflow.tool_call_id_map_to_message_id[
tool_call_id := appr.tool_call_id
]
message.content = output.metadata[
tool_call_id := approval_.tool_call_id
]["content"]
]
message.content = output.metadata[tool_call_id]["content"]
message.approval = Approval(
tool_call_id=tool_call_id,
)
yield # 通知前端更新渲染
messages.append(message)
else:
# 将工作流设置为空
conversation.workflow = None
# 若工作流类型非对话则设置为对话
if conversation.workflow.type != WorkflowType.TALK:
conversation.workflow = Workflow()
# 将正在运行设置为否
conversation.is_running = False
yield # 通知前端更新渲染
# 更新会话记录
await db_state.update_conversation_record(
@ -499,15 +510,12 @@ class ConversationState(rx.State):
usage=result.usage,
workflow=conversation.workflow,
)
# 创建运行记录
await db_state.create_run_record(
conversation_id=conversation.id,
new_messages=result.new_messages(),
)
self.conversations[self.actived_conversation_id] = conversation
yield # 通知前端更新渲染
# 批量创建消息记录
await db_state.create_message_records(
conversation_id=conversation.id,
@ -578,7 +586,6 @@ class ConversationState(rx.State):
if not conversation:
return
# 获取审批实例
if not (message := conversation.messages.get(item_id)) or not isinstance(
approval := message.approval, Approval
):
@ -587,14 +594,15 @@ class ConversationState(rx.State):
# 更新审批结果
approval.result = result
# 运行
await self.run(
async for _ in self.run(
approval,
DeferredToolResults(
approvals={
approval.tool_call_id: ToolApproved() if result else ToolDenied()
}
),
)
):
...
@rx.event
async def set_workflow(self, type: WorkflowType) -> None:
@ -617,6 +625,11 @@ class ConversationState(rx.State):
f"帮我查询并预定在 {workflow.deps.date.strftime('%Y-%m-%d')} 从 {workflow.deps.origin_airport_code} 到 {workflow.deps.destination_airport_code} 的航班"
)
case _:
# 初始化对话工作流
workflow = Workflow()
# 设置工作流
conversation.workflow = workflow
# 获取数据库状态
db_state = await self.get_db_state()

21
agent/application/states/database.py
View File

@ -2,9 +2,10 @@
"""
数据库状态
"""
from dataclasses import asdict
from datetime import datetime, timedelta
from random import choices
from dataclasses import asdict
from pydantic_ai import ModelMessage, ModelMessagesTypeAdapter, RunUsage
from pydantic_ai._uuid import uuid7
import reflex as rx
@ -15,12 +16,12 @@ from application.states.models import (
Message,
MessageType,
Workflow,
approval_dump_json,
approval_validate_json,
usage_dump_json,
usage_validate_json,
workflow_dump_json,
workflow_validate_json,
approval_dump_json,
approval_validate_json,
)
@ -346,7 +347,17 @@ class DatabaseState(rx.State):
"""
async with rx.asession() as session:
for message in messages:
record = MessageRecord(
result = await session.exec(
select(MessageRecord).where(MessageRecord.id == message.id)
)
record = result.first()
if record:
record.title = message.title
record.content = message.content
record.approval = approval_dump_json(message.approval)
else:
session.add(
MessageRecord(
id=message.id,
conversation_id=conversation_id,
type=message.type,
@ -354,7 +365,7 @@ class DatabaseState(rx.State):
content=message.content,
approval=approval_dump_json(message.approval),
)
session.add(record)
)
await session.commit()
async def read_message_history(self, conversation_id: str) -> list[ModelMessage]:

14
agent/application/states/models.py
View File

@ -66,6 +66,7 @@ class WorkflowType(StrEnum):
工作类型枚举
"""
TALK = "talk"
BOOK_FLIGHT = "预定航班"
@ -74,8 +75,11 @@ class Workflow(BaseModel):
工作流类
"""
type: WorkflowType = Field(..., description="工作类型")
type: WorkflowType = Field(default=WorkflowType.TALK, description="工作类型")
deps: Any = Field(default=None, description="工作依赖项")
tool_call_id_map_to_message_id: dict[str, str] = Field(
default_factory=dict, description="工具调用唯一标识映射为消息实例唯一标识字典"
)
class Conversation(BaseModel):
@ -88,7 +92,7 @@ class Conversation(BaseModel):
created_at: datetime = Field(..., description="会话创建日期时间")
usage: RunUsage = Field(..., description="会话使用量")
messages: dict[str, Message] = Field(default_factory=dict, description="消息字典")
workflow: Workflow | None = Field(default=None, description="工作流")
workflow: Workflow = Field(default_factory=Workflow, description="工作，默认为对话")
user_prompt: str = Field(default="", description="用户提示词")
is_running: bool = Field(
default=False, description="正在运行，True 表示正在运行， False 表示运行结束"
@ -119,7 +123,7 @@ class MessageHistoryItem(Message):
# 工作流适配器
WorkflowAdapter = TypeAdapter(Any)
WorkflowAdapter = TypeAdapter(Workflow)
def workflow_validate_json(workflow: str) -> Any:
@ -127,11 +131,11 @@ def workflow_validate_json(workflow: str) -> Any:
将工作流反序列化
"""
if not workflow:
return None
return Workflow()
return WorkflowAdapter.validate_json(workflow)
def workflow_dump_json(workflow: Any) -> str:
def workflow_dump_json(workflow: Workflow) -> str:
"""
将工作流序列化
:param workflow: 工作流

47
agent/application/workshop/book_flight.py
View File

@ -16,24 +16,25 @@ from pydantic_ai import (
RunContext,
RunUsage,
UsageLimits,
AgentRunResult,
)
from pydantic_ai.messages import (
PartDeltaEvent,
TextPart,
AgentStreamEvent,
FunctionToolResultEvent,
PartStartEvent,
ToolCallPart,
ToolReturnPart,
)
from pydantic_ai.run import AgentRunResult, AgentRunResultEvent
from pydantic_ai.run import AgentRunResultEvent
from application.states.models import (
NoResult,
Approval,
)
from application.states.models import NoResult
from application.workshop.models import (
DEEPSEEK_V4_FLASH_MODEL,
DEEPSEEK_V4_FLASH_MODEL_SETTINGS,
)
from application.workshop.events import AgentPromptEvent
class Flight(BaseModel):
@ -174,6 +175,9 @@ async def search_flights(ctx: RunContext[Deps]) -> list[Flight]:
**设计工具需先设计业务流程**
本示例按照 查询航班 -> 推荐航班 -> 预定航班
"""
if ctx.deps.searched_flights:
return ctx.deps.searched_flights
# 提取到的所有航班信息
extracted_flights = (
await extraction_flights_agent.run(
@ -204,13 +208,14 @@ async def search_flights(ctx: RunContext[Deps]) -> list[Flight]:
return searched_flights
@agent.tool(
name="预定航班",
)
@agent.tool(name="预定航班")
async def book_flight(ctx: RunContext[Deps]) -> Flight | NoResult:
"""
预定航班
"""
if ctx.deps.booked_flight:
return ctx.deps.booked_flight
if (searched_flights := ctx.deps.searched_flights) is None:
raise ModelRetry("必须先使用 search_flights 查询航班")
@ -289,7 +294,6 @@ async def run_stream_events(
"""
运行并流式输出事件
"""
tool_names = set()
async with agent.run_stream_events(
user_prompt=user_prompt,
deps=deps,
@ -308,8 +312,6 @@ async def run_stream_events(
case ToolCallPart(tool_name=tool_name):
match tool_name:
case "查询航班" | "预定航班":
if tool_name not in tool_names:
tool_names.add(tool_name)
yield event
case FunctionToolResultEvent(part):
@ -327,14 +329,19 @@ async def run_stream_events(
event.part.content += f"{searched_flight.number} {searched_flight.airfare}$ 于 {searched_flight.date.strftime('%Y-%m-%d')} 从 {searched_flight.origin_airport_code} 到 {searched_flight.destination_airport_code}\n\n"
yield event
case "预定航班":
if isinstance(content, NoResult):
event.part.content = (
"未提取到航班号，请检查后重试"
)
if isinstance(content, Flight):
event.part.content = f"已预定航班：\n{content.number} {content.airfare}$ 于 {content.date.strftime('%Y-%m-%d')} 从 {content.origin_airport_code} 到 {content.destination_airport_code}\n"
yield event
# 预定航班无需业务定义
case AgentRunResultEvent:
case AgentRunResultEvent(result=result):
match result:
case AgentRunResult(output=output):
match output:
case Flight():
yield AgentPromptEvent(
content=f"已预定航班：\n\n{output.number} {output.airfare}$ 于 {output.date.strftime('%Y-%m-%d')} 从 {output.origin_airport_code} 到 {output.destination_airport_code}\n\n"
)
case NoResult():
yield AgentPromptEvent(
content="未提取到航班号，请检查后重试"
)
case DeferredToolRequests():
yield event

16
agent/application/workshop/events.py Normal file
View File

@ -0,0 +1,16 @@
# -*- coding: utf-8 -*-
"""
智能体事件模块
"""
from dataclasses import dataclass
from pydantic_ai.messages import CustomEvent
@dataclass(kw_only=True)
class AgentPromptEvent(CustomEvent):
"""
智能体提示事件
"""
content: str

2
agent/application/workshop/models.py
View File

@ -1,6 +1,6 @@
# -*- coding: utf-8 -*-
"""
智能体相关模块
智能体模型模块
"""
from pydantic_ai import ModelSettings
from pydantic_ai.models.openai import OpenAIChatModel

BIN
agent/database.db
View File

Binary file not shown.