1

agent/application/workshop/book_flight.py

@@ -7,7 +7,7 @@ import datetime
 from typing import AsyncGenerator
 
 from logfire_api.variables import ValueDoesNotEqual
-from pydantic import BaseModel, Field, field_validator
+from pydantic import BaseModel, Field, field_validator, TypeAdapter
 from pydantic_ai import (
     Agent,
     ApprovalRequired,
@@ -23,6 +23,7 @@ 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
 
 
 DEEPSEEK_V4_FLASH_MODEL = OpenAIChatModel(
@@ -38,17 +39,6 @@ MODEL_SETTINGS = ModelSettings(
 )  # 禁用思考模式
 
 
-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="来源资料")
-
-
 class Flight(BaseModel):
     """
     航班类
@@ -79,18 +69,35 @@ class FlightNoFound(BaseModel):
     """
 
 
-system_prompt = """你的任务是根据用户提供的日期、出发机场和到达机场选择航班。
+class Deps(BaseModel):
-1. 使用 extract_flights 从来源资料提取所有航班信息；
+    """
-2. 使用 select_flight 从所有航班信息中选择匹配航班；
+    依赖类
-"""
+    """
 
-# 主智能体
+    date: datetime.date = Field(..., description="日期")
-agent = Agent[Deps, Flight | FlightNoFound](
+    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="提取到的所有航班信息"
+    )
+    selected_flight: Flight | FlightNoFound | None = Field(
+        default=None, description="查询到的航班"
+    )
+
+
+# 航班查询与预定智能体
+agent = Agent[Deps, Flight | FlightNoFound | DeferredToolRequests](
     model=DEEPSEEK_V4_FLASH_MODEL,
     model_settings=MODEL_SETTINGS,
     deps_type=Deps,
-    output_type=Flight | FlightNoFound,
+    output_type=Flight | FlightNoFound | DeferredToolRequests,
-    system_prompt=system_prompt,
+    system_prompt=(
+        "你的任务是帮助用户查询并预定航班，**必须按照下述步骤执行**：",
+        "1. 使用 extract_flights 提取所有航班信息；",
+        "2. 使用 select_flight 查询航班；",
+        "3. 若查询到航班则使用 book_flight 预定航班；若未查询到航班则返回 FlightNoFound。",
+    ),
 )
 
 # 航班信息提取智能体
@@ -99,64 +106,78 @@ extraction_agent = Agent(
     model_settings=MODEL_SETTINGS,
     output_type=list[Flight],
     system_prompt=(
-        "你的任务是根据给定的来源资料提取出所有航班信息，包括航班号、日期、出发机场代码、到达机场代码和机票价格。",
+        "你的任务是根据提供的来源资料提取出所有航班信息，包括航班号、日期、出发机场代码、到达机场代码和机票价格。",
         "其中，出发机场代码和到达机场代码须为英文、大写、3位字符串，例如 San Francisco International Airport (SFO) 的机场代码为 SFO。",
         "若某航班信息不全则跳过该航班，若没有任何航班信息则返回空列表。",
         "禁止编造。",
     ),
-)  # 通过提示词约束输出
-
-
-# 航班选择智能体
-selection_agent = Agent(
-    model=DEEPSEEK_V4_FLASH_MODEL,
-    model_settings=MODEL_SETTINGS,
-    output_type=Flight | FlightNoFound,
-    system_prompt="你的任务是从所有航班信息中选择匹配用户提供的日期、出发机场和到达机场的航班：若有多个航班则选择最便宜的航班，选择完成需用户确认是否预定；若没有则返回 FlightNoFound。",
 )
 
-"""
-如何提高提取准确率
-1. 系统提示词规则约束
-2. 降低模型温度
-3. 输出类型约定和输出模型校验
-4. 业务规则约束（例如，输出不可为空列表）
-"""
-
 
 @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_limits=UsageLimits(request_limit=3),
         usage=ctx.usage,
+        usage_limits=UsageLimits(request_limit=10),
     )
-    if len(output := result.output) != 8:
+
-        raise ModelRetry("提取到的所有航班信息应为 8 条")
+    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:
     """
-    选择航班
+    查询航班
     """
-    result = await selection_agent.run(ctx.deps.source_material, usage=ctx.usage)
+    extracted_flights = ctx.deps.extracted_flights
-    if isinstance(result.output, Flight):
+    if not extracted_flights:
-        if not ctx.tool_call_approved:
+        raise ModelRetry("须先使用 extract_flights 提取所有航班信息")
-            raise ApprovalRequired(metadata={"reason": "protected"})
+
-    return result.output
+    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
 
 
-@agent.output_validator
+@selection_agent.output_validator
 async def validate_output(
     ctx: RunContext[Deps], output: Flight | FlightNoFound
 ) -> Flight | FlightNoFound:
     """
-    输出校验
+    航班查询智能体输出校验
     """
     # 不校验未查询到航班
     if isinstance(output, FlightNoFound):
@@ -165,17 +186,39 @@ async def validate_output(
     errors = []
     if output.date != ctx.deps.date:
         errors.append(f"航班日期应为 {ctx.deps.date}, 不是 {output.date}")
-    if output.origin != ctx.deps.origin:
+    if output.origin_airport_code != ctx.deps.origin_airport_code:
-        errors.append(f"航班出发机场应为 {ctx.deps.origin}, 不是 {output.origin}")
-    if output.destination != ctx.deps.destination:
         errors.append(
-            f"航班到达机场应为 {ctx.deps.destination}, 不是 {output.destination}"
+            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
@@ -221,8 +264,8 @@ source_material = """
 
 
 async def run_stream_events(
-    user_prompt: str,
     deps: Deps,
+    user_prompt: str,
     message_history: list[ModelMessage],
     usage: RunUsage,
     deferred_tool_results: DeferredToolResults | None = None,
@@ -239,30 +282,42 @@ async def run_stream_events(
 
 
 async def main():
+    # 实例化依赖
     deps = Deps(
         source_material=source_material,
         date=datetime.date(2025, 1, 10),
-        origin="SFO",
+        origin_airport_code="SFO",
-        destination="ANC",
+        destination_airport_code="ANC",
+        selected_flight=None,
     )
-    # message_history 传空列表，不要传 None
+    user_prompt = f"帮我查询并预定在 {deps.date.strftime('%Y-%m-%d')} 从 {deps.origin_airport_code} 到 {deps.destination_airport_code} 的航班"
-    msg_history: list[ModelMessage] = []
+    message_history: list[ModelMessage] = []
-    prompt = f"Find me a flight from {deps.origin} to {deps.destination} on {deps.date}"
+    events = run_stream_events(
-
-    gen = run_stream_events(
-        user_prompt=prompt,
         deps=deps,
-        message_history=msg_history,
+        user_prompt=user_prompt,
+        message_history=message_history,
         usage=RunUsage(),
     )
     # 消费异步生成器
-    async for ev in gen:
+    async for event in events:
-        print(ev)
+        print(event)
+        if isinstance(event, AgentRunResultEvent):
+            # 保存本次运行新增消息
+            message_history.extend(event.result.new_messages())
 
 
 asyncio.run(main())
 
 """
+@selection_agent.tool
+async def select_flight(ctx: RunContext[Deps]) -> Flight | FlightNoFound:
+
+    result = await selection_agent.run(ctx.deps.source_material, usage=ctx.usage)
+    if isinstance(result.output, Flight):
+        if not ctx.tool_call_approved:
+            raise ApprovalRequired(metadata={"reason": "protected"})
+    return result.output
+
 if not isinstance(event, AgentRunResultEvent):
                     yield event
                 else: