211 lines
6.8 KiB
Python
211 lines
6.8 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
生成产品需求文档智能体
|
||
"""
|
||
from dataclasses import dataclass
|
||
import datetime
|
||
from typing import AsyncGenerator, Optional, Literal, List
|
||
|
||
from pydantic import BaseModel, Field, Optional
|
||
from pydantic_ai import (
|
||
Agent,
|
||
ModelMessage,
|
||
ModelRetry,
|
||
RunContext,
|
||
RunUsage,
|
||
UsageLimits,
|
||
)
|
||
from pydantic_ai.usage import RunUsage
|
||
from pydantic_ai.run import AgentRunResultEvent
|
||
from models import DEEPSEEK_V4_FLASH_MODEL
|
||
|
||
|
||
class Deps(BaseModel):
|
||
web_page_text: str
|
||
req_origin: str
|
||
req_destination: str
|
||
req_date: datetime.date
|
||
|
||
|
||
class FlightDetails(BaseModel):
|
||
"""
|
||
Details of the most suitable flight
|
||
"""
|
||
|
||
flight_number: str
|
||
price: int
|
||
origin: str = Field(description="Three-letter airport code")
|
||
destination: str = Field(description="Three-letter airport code")
|
||
date: datetime.date
|
||
|
||
|
||
class NoFlightFound(BaseModel):
|
||
"""
|
||
When no valid flight is found
|
||
"""
|
||
|
||
|
||
# This agent is responsible for controlling the flow of the conversation.
|
||
search_agent = Agent[Deps, FlightDetails | NoFlightFound](
|
||
model=DEEPSEEK_V4_FLASH_MODEL,
|
||
deps_type=Deps,
|
||
output_type=FlightDetails | NoFlightFound,
|
||
retries=3,
|
||
system_prompt=(
|
||
"Your job is to find the cheapest flight for the user on the given date. "
|
||
),
|
||
)
|
||
|
||
# This agent is responsible for extracting flight details from web page text.
|
||
extraction_agent = Agent(
|
||
model=DEEPSEEK_V4_FLASH_MODEL,
|
||
output_type=List[FlightDetails],
|
||
system_prompt="Extract all the flight details from the given text.",
|
||
)
|
||
|
||
|
||
@search_agent.tool
|
||
async def extract_flights(ctx: RunContext[Deps]) -> List[FlightDetails]:
|
||
"""Get details of all flights."""
|
||
# we pass the usage to the search agent so requests within this agent are counted
|
||
result = await extraction_agent.run(ctx.deps.web_page_text, usage=ctx.usage)
|
||
return result.output
|
||
|
||
|
||
@search_agent.output_validator
|
||
async def validate_output(
|
||
ctx: RunContext[Deps], output: FlightDetails | NoFlightFound
|
||
) -> FlightDetails | NoFlightFound:
|
||
"""Procedural validation that the flight meets the constraints."""
|
||
if isinstance(output, NoFlightFound):
|
||
return output
|
||
errors: list[str] = []
|
||
if output.origin != ctx.deps.req_origin:
|
||
errors.append(
|
||
f"Flight should have origin {ctx.deps.req_origin}, not {output.origin}"
|
||
)
|
||
if output.destination != ctx.deps.req_destination:
|
||
errors.append(
|
||
f"Flight should have destination {ctx.deps.req_destination}, not {output.destination}"
|
||
)
|
||
if output.date != ctx.deps.req_date:
|
||
errors.append(f"Flight should be on {ctx.deps.req_date}, not {output.date}")
|
||
if errors:
|
||
raise ModelRetry("\n".join(errors))
|
||
else:
|
||
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."""
|
||
|
||
|
||
# This agent is responsible for extracting the user's seat selection
|
||
seat_preference_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. "
|
||
),
|
||
)
|
||
# in reality this would be downloaded from a booking site,
|
||
# potentially using another agent to navigate the site
|
||
flights_web_page = """
|
||
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
|
||
"""
|
||
# restrict how many requests this app can make to the LLM
|
||
usage_limits = UsageLimits(request_limit=5)
|
||
|
||
|
||
async def flight_booking(
|
||
deps: Deps, user_prompt: Optional[str] = None
|
||
) -> AsyncGenerator:
|
||
if deps.stage == FlowStage.EXEC:
|
||
prompt = f"Find me a flight from {deps.req_origin} to {deps.req_destination} on {deps.req_date}"
|
||
run_result = None
|
||
|
||
# 1. 模型流式查询航班
|
||
async with search_agent.run_stream_events(
|
||
user_prompt=prompt,
|
||
deps=deps,
|
||
message_history=message_history,
|
||
usage_limits=usage_limits,
|
||
) as events:
|
||
async for event in events:
|
||
if isinstance(event, AgentRunResultEvent):
|
||
run_result = event
|
||
yield event
|
||
|
||
# 2. 保存本轮模型对话到数据库（核心：上下文持久化，防止断裂）
|
||
if run_result is not None:
|
||
# 写入对话历史，下一轮get_message_history可以读到航班内容
|
||
await db_state.create_result_record(
|
||
conversation_id=state.conversation_id,
|
||
dialog_id=state.dialog_id,
|
||
new_messages=run_result.new_messages(),
|
||
)
|
||
# 用量回填
|
||
state.usage = run_result.usage
|
||
|
||
# 3. 判断业务结果分支
|
||
if isinstance(run_result.output, NoFlightFound):
|
||
# 无航班场景
|
||
state.stage = FlowStage.FINISH
|
||
yield AgentStreamEvent.text_event("未找到符合条件的航班，预订流程结束")
|
||
return
|
||
else:
|
||
# ✅ 查询到航班，推送业务选择提示（前端展示按钮/文字提示）
|
||
tip_text = "已查询到航班，请回复 buy 购票 / search 重新查询"
|
||
yield AgentStreamEvent.text_event(tip_text)
|
||
|
||
# 阶段挂起，等待用户输入buy/search，不直接结束流程
|
||
state.stage = FlowStage.WAIT_USER_INPUT
|
||
return
|