257 lines
8.2 KiB
Python
257 lines
8.2 KiB
Python
# -*- 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
|