liubiren
/
Python
1
1
Fork 0
Code Issues Pull Requests Packages Projects Releases Activity

1

This commit is contained in:
liubiren 2026-08-31 20:28:18 +08:00
parent 16c4072c62
commit 38b5b8289a
3 changed files with 156 additions and 141 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
agent/application/states/conversation.py
View File

@ -30,6 +30,7 @@ from application.states.models import (
Work,
WorkType,
TaskStatus,
AgentRunWarmUpEvent,
ConversationHistoryItem,
usage_validate_python,
MessageHistoryItem,
@ -308,8 +309,8 @@ class ConversationState(rx.State):
# 对话
if not conversation.work:
from application.workshop.talk import (
run_stream_events,
)
run_stream_events,
)
else:
# 匹配工作类型
match conversation.work.type:
@ -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):
# 更新使用量

12
agent/application/states/models.py
View File

@ -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):
"""
工作类型枚举

263
agent/application/workshop/book_flight.py
View File

@ -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():