diff --git a/agent/application/pages/conversation.py b/agent/application/pages/conversation.py index f65697f..790f83d 100644 --- a/agent/application/pages/conversation.py +++ b/agent/application/pages/conversation.py @@ -9,7 +9,7 @@ from application.states.models import ( ConversationHistoryItem, MessageHistoryItem, MessageType, - WorkFlowType, + WorkType, ) @@ -278,8 +278,8 @@ def use_guidance() -> rx.Component: rx.hstack( rx.box( "预定航班", - on_click=lambda: ConversationState.init_work_flow( - WorkFlowType.BOOK_FLIGHT + on_click=lambda: ConversationState.init_work( + WorkType.BOOK_FLIGHT ), padding="10px 16px", background_color="var(--prismui-background-color-6)", @@ -570,8 +570,7 @@ def message_history() -> rx.Component: rx.auto_scroll( rx.vstack( rx.foreach( - ConversationState.message_history, - message_history_item, + ConversationState.message_history, message_history_item ), # 若等待流式输出则显示加载动效 rx.cond( diff --git a/agent/application/states/conversation.py b/agent/application/states/conversation.py index aed9306..5abbf26 100644 --- a/agent/application/states/conversation.py +++ b/agent/application/states/conversation.py @@ -27,10 +27,11 @@ from application.states.models import ( Conversation, Message, MessageType, - WorkFlow, - WorkFlowType, + Work, + WorkType, ConversationHistoryItem, usage_validate_python, + MessageHistoryItem, ) @@ -173,7 +174,7 @@ class ConversationState(rx.State): self.actived_conversation_id = conversation_id @rx.var - def message_history(self) -> list[Message]: + def message_history(self) -> list[MessageHistoryItem]: """ 获取当前会话的消息历史 :return: 当前会话的消息历史 @@ -182,7 +183,10 @@ class ConversationState(rx.State): conversation = self.conversations.get(self.actived_conversation_id) if not conversation: return [] - return list(conversation.messages.values()) + return [ + MessageHistoryItem(**message.model_dump()) + for message in list(conversation.messages.values()) + ] @rx.event async def create_conversation(self) -> None: @@ -302,17 +306,18 @@ class ConversationState(rx.State): # 初始化工具调用唯一标识集合 tool_call_ids: set[str] = set() - # 匹配工作流类型 - match conversation.work_flow: - # 预定航班 - case WorkFlowType.BOOK_FLIGHT: - from application.workshop.book_flight import run_stream_events - - # 非结构化对话 - case _: - from application.workshop.unstructured_dialogue import ( + # 非结构化对话 + if not conversation.work: + from application.workshop.unstructured_dialogue import ( run_stream_events, ) + else: + # 匹配工作类型 + match conversation.work.type: + # 预定航班 + case WorkType.BOOK_FLIGHT: + from application.workshop.book_flight import run_stream_events + # 消息列表 messages: list[Message] = [] @@ -440,22 +445,22 @@ class ConversationState(rx.State): ) @rx.event - async def init_work_flow(self, work_flow_type: WorkFlowType) -> None: + async def init_work(self, work_type: WorkType) -> None: """ - 初始化工作流 + 初始化工作 """ # 当前会话 conversation = self.conversations.get(self.actived_conversation_id) if not conversation: return - match work_flow_type: - # 预定航班工作流 - case WorkFlowType.BOOK_FLIGHT: - from application.workshop.book_flight import init_work_flow + match work_type: + # 预定航班工作 + case WorkType.BOOK_FLIGHT: + from application.workshop.book_flight import init_work # 初始化预定航班任务 - conversation.work_flow = (work_flow := init_work_flow()) + conversation.work = (work := init_work()) @rx.event def toggle_message_history_item_shown(self, message_id: str) -> None: diff --git a/agent/application/states/models.py b/agent/application/states/models.py index 4edc7ec..807433c 100644 --- a/agent/application/states/models.py +++ b/agent/application/states/models.py @@ -12,6 +12,29 @@ from pydantic_ai._uuid import uuid7 from pydantic_ai.usage import UsageLimits +class TaskStatus(StrEnum): + """ + 任务状态枚举 + """ + + RUNNING = "running" + DONE = "done" + ERROR = "error" + + +class Task(BaseModel): + """ + 任务类 + """ + + id: str = Field(default_factory=lambda: str(uuid7()), description="任务唯一标识") + tool_name: str = Field(..., description="工具名称") + tool_call_id: str = Field(..., description="工具调用唯一标识") + status: TaskStatus = Field(default=TaskStatus.RUNNING, description="任务状态") + title: str = Field(default="", description="任务标题") + content: str = Field(default="", description="任务内容") + + class MessageType(StrEnum): """ 消息类型枚举 @@ -19,7 +42,7 @@ class MessageType(StrEnum): USER_PROMPT = "user_prompt" THINKING = "thinking" - TOOL_CALL = "tool_call" + WORK = "work" RESULT_OUTPUT = "result_output" @@ -32,6 +55,7 @@ class Message(BaseModel): type: MessageType = Field(..., description="消息类型") title: str = Field(default="", description="消息标题") content: str = Field(default="", description="消息内容") + tasks: list[Task] = Field(default_factory=list, description="任务列表") is_running: bool = Field( default=False, description="正在运行,True 表示正在运行,False 表示运行完成" ) @@ -40,33 +64,25 @@ class Message(BaseModel): ) -class WorkFlowType(StrEnum): +class WorkType(StrEnum): """ - 工作流类型枚举 + 工作类型枚举 """ BOOK_FLIGHT = "预定航班" -class Deps(BaseModel): - """ - 依赖项类 - """ - - pass - - -class WorkFlow(BaseModel): +class Work(BaseModel): """ 工作流类 """ - type: WorkFlowType = Field(..., description="工作流类型") - deps: Any = Field(default=None, description="工作流依赖项") - usage: RunUsage = Field(default=RunUsage(), description="工作流使用量") + type: WorkType = Field(..., description="工作类型") + deps: Any = Field(default=None, description="工作依赖项") + usage: RunUsage = Field(default=RunUsage(), description="工作使用量") usage_limits: UsageLimits | None = Field( default=None, - description="工作流使用量限制", + description="工作使用量限制", ) @@ -77,7 +93,7 @@ class Conversation(BaseModel): id: str = Field(..., description="会话唯一标识") description: str = Field(..., description="会话描述") - work_flow: WorkFlow | None = Field(default=None, description="工作流") + work: Work | None = Field(default=None, description="工作") user_prompt: str = Field(default="", description="用户提示词") usage: dict[str, Any] = Field(..., description="会话使用量") messages: dict[str, Message] = Field(..., description="消息字典") @@ -101,22 +117,11 @@ class ConversationHistoryItem(BaseModel): created_at: str = Field(..., description="会话创建日期时间") -class MessageHistoryItem(BaseModel): +class MessageHistoryItem(Message): """ 消息历史项类 """ - id: str = Field(..., description="消息唯一标识") - type: MessageType = Field(..., description="消息类型") - title: str = Field(default="", description="消息标题") - content: str = Field(default="", description="消息内容") - is_running: bool = Field( - default=False, description="正在运行,True 表示正在运行,False 表示运行完成" - ) - is_shown: bool = Field( - default=False, description="展示组件,True 表示展示,False 表示隐藏" - ) - # Usage 适配器 UsageAdapter = TypeAdapter(RunUsage) diff --git a/agent/application/workshop/book_flight.py b/agent/application/workshop/book_flight.py index 023e21e..8d45c47 100644 --- a/agent/application/workshop/book_flight.py +++ b/agent/application/workshop/book_flight.py @@ -4,9 +4,8 @@ """ import asyncio import datetime -from typing import AsyncGenerator +from typing import AsyncGenerator, cast -from logfire_api.variables import ValueDoesNotEqual from pydantic import BaseModel, Field, field_validator, TypeAdapter from pydantic_ai import ( Agent, @@ -22,8 +21,77 @@ from pydantic_ai import ( ) from pydantic_ai.models.openai import OpenAIChatModel from pydantic_ai.providers.openai import OpenAIProvider -from sqlalchemy.sql.dml import ReturningDelete from pydantic_ai.run import AgentRunResultEvent +from enum import StrEnum +from pydantic_ai._uuid import uuid7 +from pydantic_ai.messages import ( + FunctionToolCallEvent, + FunctionToolResultEvent, + LoadCapabilityCallPart, + PartDeltaEvent, + PartEndEvent, + PartStartEvent, + TextPart, + TextPartDelta, + ThinkingPart, + ThinkingPartDelta, + ToolCallPart, + ToolSearchCallPart, + ToolReturnPart, +) + + +class TaskStatus(StrEnum): + """ + 任务状态枚举 + """ + + RUNNING = "running" + DONE = "done" + ERROR = "error" + + +class Task(BaseModel): + """ + 任务类 + """ + + tool_name: str = Field(..., description="工具名称") + tool_call_id: str = Field(..., description="工具调用唯一标识") + status: TaskStatus = Field(default=TaskStatus.RUNNING, description="任务状态") + title: str = Field(default="", description="任务标题") + content: str = Field(default="", description="任务内容") + + +class MessageType(StrEnum): + """ + 消息类型枚举 + """ + + USER_PROMPT = "user_prompt" + THINKING = "thinking" + WORK_OUTPUT = "work_output" + RESULT_OUTPUT = "result_output" + + +class Message(BaseModel): + """ + 消息类 + """ + + id: str = Field(default_factory=lambda: str(uuid7()), description="消息唯一标识") + type: MessageType = Field(..., description="消息类型") + title: str = Field(default="", description="消息标题") + content: str = Field(default="", description="消息内容") + tasks: dict[str, Task] = Field( + default_factory=dict, description="任务字典,键为工具调用唯一标识" + ) + is_running: bool = Field( + default=False, description="正在运行,True 表示正在运行,False 表示运行完成" + ) + is_shown: bool = Field( + default=False, description="展示组件,True 表示展示,False 表示隐藏" + ) DEEPSEEK_V4_FLASH_MODEL = OpenAIChatModel( @@ -63,9 +131,9 @@ class Flight(BaseModel): return airport_code -class FlightNoFound(BaseModel): +class NoSelectedFlight(BaseModel): """ - 未查询到航班类 + 未选择到航班 """ @@ -77,36 +145,37 @@ class Deps(BaseModel): date: datetime.date = Field(..., description="日期") origin_airport_code: str = Field(..., description="出发机场代码") destination_airport_code: str = Field(..., description="到达机场代码") - source_material: str = Field(..., description="来源资料") - extracted_flights: list[Flight] | None = Field( - default=None, description="提取到的所有航班信息" + matched_flights: list[Flight] | None = Field( + default=None, description="匹配到的航班" ) - selected_flight: Flight | FlightNoFound | None = Field( - default=None, description="查询到的航班" + selected_flight: Flight | NoSelectedFlight | None = Field( + default=None, description="选择到的航班" ) -# 航班查询与预定智能体 -agent = Agent[Deps, Flight | FlightNoFound | DeferredToolRequests]( +# 主智能体 +agent = Agent[Deps, Flight | NoSelectedFlight | DeferredToolRequests]( model=DEEPSEEK_V4_FLASH_MODEL, model_settings=MODEL_SETTINGS, deps_type=Deps, - output_type=Flight | FlightNoFound | DeferredToolRequests, + output_type=Flight | NoSelectedFlight | DeferredToolRequests, system_prompt=( "你的任务是帮助用户查询并预定航班,**必须按照下述步骤执行**:", - "1. 使用 extract_flights 提取所有航班信息;", - "2. 使用 select_flight 查询航班;", - "3. 若查询到航班则使用 book_flight 预定航班;若未查询到航班则返回 FlightNoFound。", + "1. 使用 match_flight 匹配满足用户需求的航班;", + "2. 使用 select_flight 选择航班。若未选择到航班则返回 NoSelectedFlight;", + "3. 若选择到航班则使用 book_flight 预定航班。由用户确认后预定。", ), ) -# 航班信息提取智能体 -extraction_agent = Agent( + +# 提取智能体 +extraction_agent = Agent[Deps, list[Flight]]( model=DEEPSEEK_V4_FLASH_MODEL, model_settings=MODEL_SETTINGS, + deps_type=Deps, output_type=list[Flight], system_prompt=( - "你的任务是根据提供的来源资料提取出所有航班信息,包括航班号、日期、出发机场代码、到达机场代码和机票价格。", + "你的任务是在来源资料中提取所有航班信息,包括航班号、日期、出发机场代码、到达机场代码和机票价格。", "其中,出发机场代码和到达机场代码须为英文、大写、3位字符串,例如 San Francisco International Airport (SFO) 的机场代码为 SFO。", "若某航班信息不全则跳过该航班,若没有任何航班信息则返回空列表。", "禁止编造。", @@ -114,111 +183,6 @@ extraction_agent = Agent( ) -@agent.tool -async def extract_flights(ctx: RunContext[Deps]) -> list[Flight]: - """ - 提取所有航班信息 - **如何提高提取准确率** - 1. 系统提示词规则约束 - 2. 降低模型温度 - 3. 输出类型约定和输出模型校验 - 4. 业务规则约束(例如,输出不可为空列表) - """ - result = await extraction_agent.run( - ctx.deps.source_material, - usage=ctx.usage, - usage_limits=UsageLimits(request_limit=10), - ) - - if (flight_counts := len(output := result.output)) != 8: # 模拟业务规则约束 - raise ModelRetry(f"所有航班信息数量应为 8 ,不是 {flight_counts}") - # 更新提取到的航班信息 - ctx.deps.extracted_flights = output - return output - - -# 航班查询智能体 -selection_agent = Agent[Deps, Flight | FlightNoFound]( - model=DEEPSEEK_V4_FLASH_MODEL, - model_settings=MODEL_SETTINGS, - deps_type=Deps, - output_type=Flight | FlightNoFound, - system_prompt=( - "你的任务是在提取到的所有航班信息中查询匹配用户提供的日期、出发机场代码和到达机场代码的机票价格最低的航班。", - "若没有符合要求的航班则返回 FlightNoFound。", - ), -) - - -@agent.tool -async def select_flight(ctx: RunContext[Deps]) -> Flight | FlightNoFound: - """ - 查询航班 - """ - extracted_flights = ctx.deps.extracted_flights - if not extracted_flights: - raise ModelRetry("须先使用 extract_flights 提取所有航班信息") - - result = await selection_agent.run( - ( - f"{ctx.prompt}", - f"所有航班信息:\n{TypeAdapter(list[Flight]).dump_json(extracted_flights).decode('utf-8')}", - ), - deps=ctx.deps, - usage=ctx.usage, - usage_limits=UsageLimits(request_limit=10), - ) - ctx.deps.selected_flight = (output := result.output) - return output - - -@selection_agent.output_validator -async def validate_output( - ctx: RunContext[Deps], output: Flight | FlightNoFound -) -> Flight | FlightNoFound: - """ - 航班查询智能体输出校验 - """ - # 不校验未查询到航班 - if isinstance(output, FlightNoFound): - return output - - errors = [] - if output.date != ctx.deps.date: - errors.append(f"航班日期应为 {ctx.deps.date}, 不是 {output.date}") - if output.origin_airport_code != ctx.deps.origin_airport_code: - errors.append( - f"航班出发机场应为 {ctx.deps.origin_airport_code}, 不是 {output.origin_airport_code}" - ) - if output.destination_airport_code != ctx.deps.destination_airport_code: - errors.append( - f"航班到达机场应为 {ctx.deps.destination_airport_code}, 不是 {output.destination_airport_code}" - ) - if errors: - raise ModelRetry("\n".join(errors)) - return output - - -@agent.tool -async def book_flight(ctx: RunContext[Deps]) -> Flight | FlightNoFound: - """ - 预定航班 - """ - if not (selected_flight := ctx.deps.selected_flight): - raise ModelRetry("须先使用 select_flight 查询航班") - - if isinstance(selected_flight, FlightNoFound): - return selected_flight - else: - if not ctx.tool_call_approved: - raise ApprovalRequired( - metadata={ - "reason": f"需要用户确认是否预定 {selected_flight.number} 在 {selected_flight.date.strftime('%Y-%m-%d')} 从 {selected_flight.origin_airport_code} 到 {selected_flight.destination_airport_code} 的航班" - } - ) - return selected_flight - - source_material = """ 1. Flight SFO-AK123 - Price: $350 @@ -263,6 +227,111 @@ source_material = """ """ +@agent.tool +async def match_flights(ctx: RunContext[Deps]) -> list[Flight]: + """ + 匹配满足用户需求的航班 + **如何提高提取准确率** + 1. 系统提示词规则约束 + 2. 降低模型温度 + 3. 输出类型约定和输出模型校验 + 4. 业务规则约束(例如,输出不可为空列表) + + **设计工具需先设计业务流程** + 本示例按照 匹配满足用户需求的航班 -> 选择航班 -> 预定航班 + """ + extraction_agent_result = await extraction_agent.run( + f"来源资料:\n{source_material}", + deps=ctx.deps, + usage=ctx.usage, + usage_limits=UsageLimits(request_limit=10), + ) + if ( + extracted_flight_counts := len( + extracted_flights := extraction_agent_result.output + ) + ) != 8: # 模拟业务规则约束 + raise ModelRetry( + f"提取到的所有航班信息数量应为 8 ,不是 {extracted_flight_counts}" + ) + # 匹配用户需求的航班 + ctx.deps.matched_flights = [ + flight + for flight in extracted_flights + if flight.date == ctx.deps.date + and flight.origin_airport_code == ctx.deps.origin_airport_code + and flight.destination_airport_code == ctx.deps.destination_airport_code + ] + return ctx.deps.matched_flights + + +@agent.tool +async def select_flight(ctx: RunContext[Deps]) -> Flight | NoSelectedFlight: + """ + 选择航班 + **如何提高提取准确率** + 1. 系统提示词规则约束 + 2. 降低模型温度 + 3. 输出类型约定和输出模型校验 + 4. 业务规则约束(例如,输出不可为空列表) + """ + if (matched_flights := ctx.deps.matched_flights) is None: + raise ModelRetry("必须先使用 match_flights 匹配满足用户需求的航班") + ctx.deps.selected_flight = ( + min(matched_flights, key=lambda flight: flight.airfare) + if matched_flights + else NoSelectedFlight() + ) + return ctx.deps.selected_flight + + +@agent.tool +async def book_flight(ctx: RunContext[Deps]) -> Flight | NoSelectedFlight: + """ + 预定航班 + """ + if (selected_flight := ctx.deps.selected_flight) is None: + raise ModelRetry("必须先使用 select_flight 选择航班") + + if isinstance(selected_flight, NoSelectedFlight): + return selected_flight + else: + if not ctx.tool_call_approved: + raise ApprovalRequired( + metadata={ + "reason": f"需要用户确认是否预定 {selected_flight.number} 在 {selected_flight.date.strftime('%Y-%m-%d')} 从 {selected_flight.origin_airport_code} 到 {selected_flight.destination_airport_code} 的航班" + } + ) + return selected_flight + + +@agent.output_validator +async def validate_output( + ctx: RunContext[Deps], output: Flight | NoSelectedFlight | DeferredToolRequests +) -> Flight | NoSelectedFlight | DeferredToolRequests: + """ + 输出校验 + """ + # 不校验未选择到航班或延迟工具请求 + if isinstance(output, NoSelectedFlight | DeferredToolRequests): + return output + + errors = [] + if output.date != ctx.deps.date: + errors.append(f"航班日期应为 {ctx.deps.date}, 不是 {output.date}") + if output.origin_airport_code != ctx.deps.origin_airport_code: + errors.append( + f"航班出发机场应为 {ctx.deps.origin_airport_code}, 不是 {output.origin_airport_code}" + ) + if output.destination_airport_code != ctx.deps.destination_airport_code: + errors.append( + f"航班到达机场应为 {ctx.deps.destination_airport_code}, 不是 {output.destination_airport_code}" + ) + if errors: + raise ModelRetry("\n".join(errors)) + return output + + async def run_stream_events( deps: Deps, user_prompt: str, @@ -270,6 +339,13 @@ async def run_stream_events( usage: RunUsage, deferred_tool_results: DeferredToolResults | None = None, ) -> AsyncGenerator: + # 构建工作输出消息 + message = Message( + type=MessageType.WORK_OUTPUT, + title="正在预定航班", + is_running=True, + ) + yield message async with agent.run_stream_events( user_prompt=user_prompt, deps=deps, @@ -278,13 +354,67 @@ async def run_stream_events( deferred_tool_results=deferred_tool_results, ) as events: async for event in events: - yield event + match event: + case PartStartEvent( + part=part, + ): + match part: + case ToolCallPart( + tool_name=tool_name, tool_call_id=tool_call_id + ): + match tool_name: + case "match_flights": + if tool_name not in message.tasks: + message.tasks[tool_name] = Task( + tool_name=tool_name, + tool_call_id=tool_call_id, + title="正在查询航班", + ) + yield message + + case "select_flight": + if tool_name not in message.tasks: + message.tasks[tool_name] = Task( + tool_name=tool_name, + tool_call_id=tool_call_id, + title="正在选择航班", + ) + yield message + + + case FunctionToolResultEvent(part): + match part: + case ToolReturnPart( + tool_name=tool_name, + tool_call_id=tool_call_id, + content=content, + ): + match tool_name: + case "match_flights": + matched_flights = cast(list[Flight], content) + task = message.tasks[tool_name] + task.title = f"已查询到 {len(matched_flights)} 班航班" + task.content = "\n".join( + [ + f"{matched_flight.number} - {matched_flight.date.strftime('%Y-%m-%d')} - {matched_flight.origin_airport_code} ~ {matched_flight.destination_airport_code} ${matched_flight.airfare}" + for matched_flight in matched_flights + ] + ) + task.status = TaskStatus.DONE + yield message + + case "select_flight": + selected_flight = cast(Flight, content) + task = message.tasks[tool_name] + task.title = f"已选择航班" + task.content = f"{selected_flight.number} - {selected_flight.date.strftime('%Y-%m-%d')} - {selected_flight.origin_airport_code} ~ {selected_flight.destination_airport_code} ${selected_flight.airfare}" + task.status = TaskStatus.DONE + yield message async def main(): # 实例化依赖 deps = Deps( - source_material=source_material, date=datetime.date(2025, 1, 10), origin_airport_code="SFO", destination_airport_code="ANC",