diff --git a/agent/application/states/conversation.py b/agent/application/states/conversation.py index 40142c1..e27c980 100644 --- a/agent/application/states/conversation.py +++ b/agent/application/states/conversation.py @@ -30,6 +30,7 @@ from application.states.models import ( Work, WorkType, TaskStatus, + AgentRunWarmUpEvent, ConversationHistoryItem, usage_validate_python, MessageHistoryItem, @@ -308,15 +309,15 @@ class ConversationState(rx.State): # 对话 if not conversation.work: from application.workshop.talk import ( - run_stream_events, - ) + run_stream_events, + ) else: # 匹配工作类型 match conversation.work.type: # 预定航班 case WorkType.BOOK_FLIGHT: from application.workshop.book_flight import run_stream_events - + # 消息列表 messages: list[Message] = [] # 运行并流式输出事件 @@ -357,8 +358,7 @@ class ConversationState(rx.State): # 文本分片开始事件 case TextPart(content=content): # 构建消息实例 - message = Message(type=MessageType.RESULT_OUTPUT) - message.content = content + message = Message(type=MessageType.TEXT, content=content) # 添加至消息字典 conversation.messages[message.id] = message # 将消息实例唯一标识与片段索引映射 @@ -413,6 +413,18 @@ class ConversationState(rx.State): ] messages.append(message) + # ========== 智能体运行预热事件 ========== + case AgentRunWarmUpEvent(content=content): + # 构建消息实例 + message = Message( + type=MessageType.TEXT, + content=content, + ) + # 添加至消息字典 + conversation.messages[message.id] = message + yield # 通知前端更新渲染 + messages.append(message) + # ========== 智能体运行结果事件 ========== case AgentRunResultEvent(result=result): # 更新使用量 diff --git a/agent/application/states/models.py b/agent/application/states/models.py index 4d4b179..5818eab 100644 --- a/agent/application/states/models.py +++ b/agent/application/states/models.py @@ -45,6 +45,7 @@ class MessageType(StrEnum): THINKING = "thinking" WORK = "work" RESULT_OUTPUT = "result_output" + TEXT = "text" class Message(BaseModel): @@ -65,6 +66,17 @@ class Message(BaseModel): ) +class AgentRunWarmUpEvent(BaseModel): + """ + 智能体运行预热事件类 + """ + + content: str = Field( + default="", + description="预热内容", + ) + + class WorkType(StrEnum): """ 工作类型枚举 diff --git a/agent/application/workshop/book_flight.py b/agent/application/workshop/book_flight.py index 1293442..e8fe701 100644 --- a/agent/application/workshop/book_flight.py +++ b/agent/application/workshop/book_flight.py @@ -25,6 +25,7 @@ from pydantic_ai.run import AgentRunResultEvent, AgentRunResult from enum import StrEnum from pydantic_ai._uuid import uuid7 from pydantic_ai.messages import ( + AgentStreamEvent, FunctionToolCallEvent, FunctionToolResultEvent, LoadCapabilityCallPart, @@ -73,6 +74,18 @@ class MessageType(StrEnum): THINKING = "thinking" WORK_OUTPUT = "work_output" RESULT_OUTPUT = "result_output" + TEXT = "text" + + +class AgentRunWarmUpEvent(BaseModel): + """ + 智能体运行预热事件类 + """ + + content: str = Field( + default="", + description="预热内容", + ) class Message(BaseModel): @@ -132,11 +145,13 @@ class Flight(BaseModel): return airport_code -class NoSelectedFlight(BaseModel): +class NoResult(BaseModel): """ - 未选择到航班 + 无结果类 """ + ... + class Deps(BaseModel): """ @@ -149,28 +164,28 @@ class Deps(BaseModel): searched_flights: list[Flight] | None = Field( default=None, description="查询到的航班" ) - recommended_flight: Flight | NoSelectedFlight | None = Field( - default=None, description="推荐航班" + extracted_flight_number: str | None = Field( + default=None, description="提取到的航班号" ) + booked_flight: Flight | None = Field(default=None, description="预定到的航班") # 主智能体 -agent = Agent[Deps, Flight | NoSelectedFlight | DeferredToolRequests]( +agent = Agent[Deps, Flight | NoResult | DeferredToolRequests]( model=DEEPSEEK_V4_FLASH_MODEL, model_settings=MODEL_SETTINGS, deps_type=Deps, - output_type=Flight | NoSelectedFlight | DeferredToolRequests, + output_type=Flight | NoResult | DeferredToolRequests, system_prompt=( "你的任务是帮助用户查询并预定航班，**必须按照下述步骤执行**：", - "1. 使用 match_flight 匹配满足用户需求的航班；", - "2. 使用 select_flight 选择航班。若未选择到航班则返回 NoSelectedFlight；", - "3. 若选择到航班则使用 book_flight 预定航班。由用户确认后预定。", + "1. 使用 search_flights 查询满足用户需求的航班；", + "2. 使用 book_flight 预定航班。若未选择到航班则返回 NoSelectedFlight；", ), ) -# 提取智能体 -extraction_agent = Agent[Deps, list[Flight]]( +# 提取所有航班信息智能体 +extraction_flights_agent = Agent[Deps, list[Flight]]( model=DEEPSEEK_V4_FLASH_MODEL, model_settings=MODEL_SETTINGS, deps_type=Deps, @@ -183,6 +198,20 @@ extraction_agent = Agent[Deps, list[Flight]]( ), ) +# 提取航班号智能体 +extraction_flight_number_agent = Agent[Deps, str | None]( + model=DEEPSEEK_V4_FLASH_MODEL, + model_settings=MODEL_SETTINGS, + deps_type=Deps, + output_type=str | None, + system_prompt=( + "你的任务是从用户提示词中提取航班号且须在查询到的航班中。", + "其中，航班号须由英文、数字和短横线组成、大写，例如 SFO-AK123。", + "若用户提示词中没有航班号则返回 None。", + "禁止编造。", + ), +) + source_material = """ 1. Flight SFO-AK123 @@ -228,7 +257,7 @@ source_material = """ """ -@agent.tool +@agent.tool(name="查询航班") async def search_flights(ctx: RunContext[Deps]) -> list[Flight]: """ 查询航班 @@ -241,76 +270,76 @@ async def search_flights(ctx: RunContext[Deps]) -> list[Flight]: **设计工具需先设计业务流程** 本示例按照 查询航班 -> 推荐航班 -> 预定航班 """ - 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}" + # 提取到的所有航班信息 + extracted_flights = ( + await extraction_flights_agent.run( + f"来源资料：\n{source_material}", + deps=ctx.deps, + usage=ctx.usage, + usage_limits=UsageLimits(request_limit=10), ) + ).output + if not extracted_flights: # 模拟业务规则约束：提取到的所有航班信息不可能为空 + raise ModelRetry(f"提取到的所有航班信息不可能为空") + # 查询到的航班 - ctx.deps.searched_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 - ] + ctx.deps.searched_flights = sorted( + [ + 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 + ], + key=lambda flight: flight.airfare, + ) return ctx.deps.searched_flights -@agent.tool -async def recommend_flight(ctx: RunContext[Deps]) -> Flight | NoSelectedFlight: - """ - 推荐航班 - **如何提高提取准确率** - 1. 系统提示词规则约束 - 2. 降低模型温度 - 3. 输出类型约定和输出模型校验 - 4. 业务规则约束（例如，输出不可为空列表） - """ - if (searched_flights := ctx.deps.searched_flights) is None: - raise ModelRetry("必须先使用 search_flights 查询航班") - ctx.deps.recommended_flight = ( - min(searched_flights, key=lambda flight: flight.airfare) - if searched_flights - else NoSelectedFlight() - ) - return ctx.deps.recommended_flight - - -@agent.tool -async def book_flight(ctx: RunContext[Deps]) -> Flight | NoSelectedFlight: +@agent.tool(name="预定航班") +async def book_flight(ctx: RunContext[Deps]) -> Flight | NoResult: """ 预定航班 """ - if (recommended_flight := ctx.deps.recommended_flight) is None: - raise ModelRetry("必须先使用 recommend_flight 推荐航班") + if (searched_flights := ctx.deps.searched_flights) is None: + raise ModelRetry("必须先使用 search_flights 查询航班") - if isinstance(recommended_flight, NoSelectedFlight): - return recommended_flight - else: - if not ctx.tool_call_approved: - raise ApprovalRequired() - return recommended_flight + if not ctx.tool_call_approved: + raise ApprovalRequired() + + # 提取到的航班号 + ctx.deps.extracted_flight_number = ( + await extraction_flight_number_agent.run( + f"用户提示词：\n{ctx.prompt}\n查询到的航班：\n{(flight_numbers:= [flight.number for flight in searched_flights])}", + deps=ctx.deps, + usage=ctx.usage, + usage_limits=UsageLimits(request_limit=10), + ) + ).output + if not ctx.deps.extracted_flight_number: + return NoResult() + if ctx.deps.extracted_flight_number not in flight_numbers: + raise ModelRetry( + f"提取到的航班号 {ctx.deps.extracted_flight_number} 不在查询到的航班中" + ) + # 预定到的航班 + ctx.deps.booked_flight = next( + flight + for flight in searched_flights + if flight.number == ctx.deps.extracted_flight_number + ) + return ctx.deps.booked_flight @agent.output_validator async def validate_output( - ctx: RunContext[Deps], output: Flight | NoSelectedFlight | DeferredToolRequests -) -> Flight | NoSelectedFlight | DeferredToolRequests: + ctx: RunContext[Deps], + output: Flight | NoResult | DeferredToolRequests, +) -> Flight | NoResult | DeferredToolRequests: """ 输出校验 """ - # 不校验未选择到航班或延迟工具请求 - if isinstance(output, NoSelectedFlight | DeferredToolRequests): + if isinstance(output, (NoResult, DeferredToolRequests)): return output errors = [] @@ -335,14 +364,20 @@ async def run_stream_events( message_history: list[ModelMessage], usage: RunUsage, deferred_tool_results: DeferredToolResults | None = None, -) -> AsyncGenerator: - # 构建消息 - message = Message( - type=MessageType.WORK_OUTPUT, - title="预定航班", - is_running=True, +) -> AsyncGenerator[ + AgentStreamEvent + | AgentRunResultEvent + | AgentRunWarmUpEvent + | Flight + | NoResult + | DeferredToolRequests, + None, +]: + # 构建智能体运行预热事件 + yield AgentRunWarmUpEvent( + content=f"正在预定 {deps.date.strftime('%Y-%m-%d')} 从 {deps.origin_airport_code} 到 {deps.destination_airport_code} 的航班" ) - yield message + tool_names = set() async with agent.run_stream_events( user_prompt=user_prompt, deps=deps, @@ -357,73 +392,35 @@ async def run_stream_events( part=part, ): match part: - case ToolCallPart( - tool_name=tool_name, tool_call_id=tool_call_id - ): + case ToolCallPart(tool_name=tool_name): match tool_name: - case "search_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 "recommend_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 "book_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 "查询航班" | "预定航班": + if tool_name not in tool_names: + tool_names.add(tool_name) + yield event case FunctionToolResultEvent(part): match part: case ToolReturnPart( tool_name=tool_name, - tool_call_id=tool_call_id, content=content, ): match tool_name: - case "search_flights": - task = message.tasks[tool_name] + case "查询航班": + # 查询到的航班 searched_flights = cast(list[Flight], content) - task.title = ( - f"已查询到 {len(searched_flights)} 班航班" - ) - task.content = "\n".join( - [ - f"{searched_flight.number} - {searched_flight.date.strftime('%Y-%m-%d')} 从 {searched_flight.origin_airport_code} 飞往 {searched_flight.destination_airport_code} - {searched_flight.airfare}" - for searched_flight in searched_flights - ] - ) - task.status = TaskStatus.DONE - yield message + event.part.content = f"已查询到 **{len(searched_flights)} 班航班**：\n" + for searched_flight in searched_flights: + 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" + yield event - case "recommend_flight": - task = message.tasks[tool_name] - recommended_flight = cast( - Flight | NoSelectedFlight, content - ) - if isinstance(recommended_flight, Flight): - task.title = "为您推荐" - task.content = f"{recommended_flight.number} - {recommended_flight.date.strftime('%Y-%m-%d')} 从 {recommended_flight.origin_airport_code} 飞往 {recommended_flight.destination_airport_code} - {recommended_flight.airfare}" - else: - task.title = "为您推荐" + case "预定航班": + if isinstance(content, NoResult): + event.part.content = "未查询到符合条件的航班" - task.status = TaskStatus.DONE - yield message + 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( result=result, @@ -433,14 +430,8 @@ async def run_stream_events( match output: case DeferredToolRequests(approvals=approvals): for approval in approvals: - match approval: - case ToolCallPart(tool_name=tool_name): - task = message.tasks[tool_name] - task.content = "请确认是否预定 " - task.status = ( - TaskStatus.PENDING_APPROVAL - ) - yield message + output = approvals + yield event async def main():