# -*- coding: utf-8 -*- """ 生成产品需求文档智能体 """ import datetime from typing import AsyncGenerator, Literal from pydantic import BaseModel, Field from pydantic_ai import ( Agent, ModelMessage, ModelRetry, RunContext, UsageLimits, ModelMessage, ) from pydantic_ai.run import AgentRunResultEvent from application.tasks.models import DEEPSEEK_V4_FLASH_MODEL from application.states.models import ( TaskType, Task, TaskNode, usage_to_dict, usage_to_object, TaskResultEvent, usage_limits_to_dict, usage_limits_to_object, ) class Deps(BaseModel): flights: str = Field(..., description="航班信息") date: datetime.date = Field(..., description="航班日期") origin: str = Field(..., description="出发机场") destination: str = Field(..., description="到达机场") class FlightDetails(BaseModel): """ 航班详情 """ number: str = Field(description="航班号") date: datetime.date = Field(description="航班日期") origin: str = Field(description="出发机场") destination: str = Field(description="到达机场") price: int = Field(description="机票价格") class NoFlightFound(BaseModel): """ 未查询到航班 """ # 主控智能体 master_agent = Agent[Deps, FlightDetails | NoFlightFound]( model=DEEPSEEK_V4_FLASH_MODEL, deps_type=Deps, output_type=FlightDetails | NoFlightFound, retries=2, system_prompt=("你的工作是在给定日期为用户找到最便宜的航班"), ) # 航班查询智能体 search_agent = Agent( model=DEEPSEEK_V4_FLASH_MODEL, output_type=list[FlightDetails], system_prompt="从给定文本中提取所有航班详细信息", ) @master_agent.tool async def search_flights(ctx: RunContext[Deps]) -> list[FlightDetails]: """ 查询并返回航班详情列表 """ result = await search_agent.run(ctx.deps.flights, usage=ctx.usage) return result.output @master_agent.output_validator async def validate_output( ctx: RunContext[Deps], output: FlightDetails | NoFlightFound ) -> FlightDetails | NoFlightFound: """ 校验主控智能体输出 """ if isinstance(output, NoFlightFound): return output errors = "" if output.date != ctx.deps.date: errors += f"航班日期应为 {ctx.deps.date}, 不是 {output.date}\n" if output.origin != ctx.deps.origin: errors += f"航班出发机场应为 {ctx.deps.origin}, 不是 {output.origin}\n" if output.destination != ctx.deps.destination: errors += ( f"航班到达机场应为 {ctx.deps.destination}, 不是 {output.destination}\n" ) if errors: raise ModelRetry(errors) return output class SeatPreference(BaseModel): row: int = Field(ge=1, le=30) seat: Literal["A", "B", "C", "D", "E", "F"] class Failed(BaseModel): """Unable to extract a seat selection.""" # 选座智能体 seat_selection_agent = Agent[object, SeatPreference | Failed]( model=DEEPSEEK_V4_FLASH_MODEL, output_type=SeatPreference | Failed, system_prompt=( "Extract the user's seat preference. " "Seats A and F are window seats. " "Row 1 is the front row and has extra leg room. " "Rows 14, and 20 also have extra leg room. " ), ) flights = """ 1. Flight SFO-AK123 - Price: $350 - Origin: San Francisco International Airport (SFO) - Destination: Ted Stevens Anchorage International Airport (ANC) - Date: January 10, 2025 2. Flight SFO-AK456 - Price: $370 - Origin: San Francisco International Airport (SFO) - Destination: Fairbanks International Airport (FAI) - Date: January 10, 2025 3. Flight SFO-AK789 - Price: $400 - Origin: San Francisco International Airport (SFO) - Destination: Juneau International Airport (JNU) - Date: January 20, 2025 4. Flight NYC-LA101 - Price: $250 - Origin: San Francisco International Airport (SFO) - Destination: Ted Stevens Anchorage International Airport (ANC) - Date: January 10, 2025 5. Flight CHI-MIA202 - Price: $200 - Origin: Chicago O'Hare International Airport (ORD) - Destination: Miami International Airport (MIA) - Date: January 12, 2025 6. Flight BOS-SEA303 - Price: $120 - Origin: Boston Logan International Airport (BOS) - Destination: Ted Stevens Anchorage International Airport (ANC) - Date: January 12, 2025 7. Flight DFW-DEN404 - Price: $150 - Origin: Dallas/Fort Worth International Airport (DFW) - Destination: Denver International Airport (DEN) - Date: January 10, 2025 8. Flight ATL-HOU505 - Price: $180 - Origin: Hartsfield-Jackson Atlanta International Airport (ATL) - Destination: George Bush Intercontinental Airport (IAH) - Date: January 10, 2025 """ def init_task() -> Task: return Task( type=TaskType.BOOK_FLIGHT, node=TaskNode.EXECUTION, deps=Deps( flights=flights, date=datetime.date(2025, 1, 10), origin="SFO", destination="ANC", ), usage_limits=usage_limits_to_dict(UsageLimits(request_limit=5)), ) async def run_stream_events( task: Task, user_prompt: str | None = None, message_history: list[ModelMessage] | None = None, ) -> AsyncGenerator: result = None while True: if task.node == TaskNode.EXECUTION: async with master_agent.run_stream_events( user_prompt=user_prompt, deps=task.deps, message_history=message_history, usage=usage_to_object(task.usage), usage_limits=usage_limits_to_object(task.usage_limits), ) as events: async for event in events: if not isinstance(event, AgentRunResultEvent): yield event else: result = event.result # 更新任务使用量 task.usage = usage_to_dict(result.usage) if isinstance(result.output, FlightDetails): content = "\n---\n已查询到航班，请回复 buy 购票 / search 重新查询\n" # 更新任务节点为待用户输入 task.node = TaskNode.PENDING_INPUT else: content = "\n---\n未找到符合条件的航班，流程结束\n" # 更新任务节点为结束 task.node = TaskNode.FINISH yield TaskResultEvent( task=task, content=content, ) yield event return if task.node == TaskNode.PENDING_INPUT: if user_prompt == "buy": yield TaskResultEvent( task=task, content="请和我说下你的座位偏好吧：\nA、F 座位是靠窗位；1 排、14 排、20 排腿部空间更大、更舒展，你更想要靠窗座位，宽敞大空间座位", ) return elif user_prompt == "search": # 更新任务节点为执行 task.node = TaskNode.EXECUTION else: async with seat_selection_agent.run_stream_events( user_prompt=user_prompt, message_history=message_history, usage=usage_to_object(task.usage), usage_limits=usage_limits_to_object(task.usage_limits), ) as events: async for event in events: if not isinstance(event, AgentRunResultEvent): yield event else: result = event.result # 更新任务使用量 task.usage = usage_to_dict(result.usage) # 更新任务节点为结束 task.node = TaskNode.FINISH yield TaskResultEvent( task=task, content="已为您预定好座位，流程结束", ) yield event return