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

1

This commit is contained in:
liubiren 2026-08-07 01:16:53 +08:00
parent ae05009acd
commit 0280261508
5 changed files with 166 additions and 177 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -25,9 +25,7 @@ from application.states.database import DatabaseState
from application.states.models import (
Conversation,
Dialog,
Task,
TaskResultEvent,
TaskType,
TaskNodeResultEvent,
Thought,
thoughts_to_dict,
usage_to_dict,
@ -52,8 +50,6 @@ class ConversationState(rx.State):
# 显示更多的会话唯一标识
shown_more_conversation_id: str = ""
# 任务类型
task_type: TaskType = TaskType.NONE
# 用户提示词
user_prompt: str = ""
@ -237,7 +233,7 @@ class ConversationState(rx.State):
tool_call_ids: dict[str, int] = {}
# 获取运行流式输出事件
async for event in run_stream_events(
task_type=self.task_type,
task=conversation.task,
user_prompt=dialog.user_prompt,
message_history=message_history,
):
@ -362,7 +358,7 @@ class ConversationState(rx.State):
case "tool_call":
dialog.thoughts[index].content = f"已调用 {content}"
case TaskResultEvent(task=task, content=content):
case TaskNodeResultEvent(task=task, content=content):
# 更新任务
conversation.task = task
dialog.result_output += content
@ -373,7 +369,7 @@ class ConversationState(rx.State):
await db_state.create_dialog_record(
conversation_id=self.conversation_id,
id=dialog.id,
user_prompt=self.user_prompt,
user_prompt=dialog.user_prompt,
thoughts=thoughts_to_dict(dialog.thoughts),
result_output=dialog.result_output,
usage=usage_to_dict(result.usage),
@ -405,7 +401,6 @@ class ConversationState(rx.State):
conversation = self.conversations[self.conversation_id]
# 初始化预定航班任务
conversation.task = init_task()
self.user_prompt = f"帮我找一班从 {conversation.task.deps.origin} 到 {conversation.task.deps.destination} 在 {conversation.task.deps.date} 的航班"
@rx.event

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

@ -2,12 +2,12 @@
"""
面向 reflex.state 的类
"""
from datetime import datetime
from enum import StrEnum
from typing import Any
from pydantic import BaseModel, Field, TypeAdapter
from pydantic_ai import RunUsage
from pydantic_ai._uuid import uuid7
from enum import StrEnum
from pydantic_ai.usage import UsageLimits
@ -53,14 +53,14 @@ class Dialog(BaseModel):
UsageAdapter = TypeAdapter(RunUsage)
def usage_to_object(usage: dict) -> RunUsage:
def usage_to_object(usage: dict[str, Any]) -> RunUsage:
"""
Usage 转为对象
"""
return UsageAdapter.validate_python(usage) if usage else RunUsage()
def usage_to_dict(usage: RunUsage) -> dict:
def usage_to_dict(usage: RunUsage) -> dict[str, Any]:
"""
Usage 转为字典
"""
@ -71,7 +71,7 @@ def usage_to_dict(usage: RunUsage) -> dict:
UsageLimitsAdapter = TypeAdapter(UsageLimits)
def usage_limits_to_object(usage_limits: dict) -> UsageLimits:
def usage_limits_to_object(usage_limits: dict[str, Any]) -> UsageLimits:
"""
UsageLimits 转为对象
"""
@ -82,7 +82,7 @@ def usage_limits_to_object(usage_limits: dict) -> UsageLimits:
)
def usage_limits_to_dict(usage_limits: UsageLimits) -> dict:
def usage_limits_to_dict(usage_limits: UsageLimits) -> dict[str, Any]:
"""
UsageLimits 转为字典
"""
@ -94,20 +94,9 @@ class TaskType(StrEnum):
任务类型枚举
"""
NONE = "none"
BOOK_FLIGHT = "book_flight"
class TaskNode(StrEnum):
"""
任务节点枚举
"""
EXECUTION = "execution"
PENDING_INPUT = "pending_input"
FINISH = "finish"
class Task(BaseModel):
"""
任务类
@ -115,7 +104,7 @@ class Task(BaseModel):
id: str = Field(default_factory=lambda: str(uuid7()), description="任务唯一标识")
type: TaskType = Field(..., description="任务类型")
node: TaskNode = Field(default=TaskNode.EXECUTION, description="任务节点")
node: str | None = Field(default="", description="任务节点")
deps: Any = Field(default=None, description="任务依赖项")
usage: dict[str, Any] = Field(default_factory=dict, description="任务使用量")
usage_limits: dict[str, Any] = Field(
@ -123,15 +112,6 @@ class Task(BaseModel):
)
class TaskResultEvent(BaseModel):
"""
任务结果事件类
"""
task: Task = Field(..., description="任务实例")
content: str = Field(default="", description="任务结果内容")
class Conversation(BaseModel):
"""
会话类
@ -149,3 +129,12 @@ class Conversation(BaseModel):
...,
description="创建时间",
)
class TaskNodeResultEvent(BaseModel):
"""
任务节点结果事件类
"""
task: Task | None = Field(..., description="任务实例")
content: str = Field(default="", description="任务结果内容")

53
agent/application/tasks/__init__.py
View File

@ -2,17 +2,15 @@
"""
任务模块
"""
from typing import Any, AsyncGenerator, Callable, Dict, Optional
from typing import AsyncGenerator, List
from pydantic import BaseModel
from pydantic_ai import Agent, ModelMessage
from pydantic_ai.models.openai import OpenAIChatModel
from pydantic_ai.providers.openai import OpenAIProvider
from pydantic_ai.messages import AgentStreamEvent
from pydantic_ai.run import AgentRunResultEvent
from application.states.models import Dialog, TaskType
from application.states.models import TaskNodeResultEvent, TaskType, Task
from application.tasks.models import DEEPSEEK_V4_FLASH_MODEL
from application.states.models import TaskNode, Task
instruction = """
# 角色
@ -32,32 +30,33 @@ instruction = """
async def run_stream_events(
task_type: TaskType,
task: Task | None,
user_prompt: str,
message_history: List[ModelMessage],
) -> AsyncGenerator:
) -> AsyncGenerator[AgentStreamEvent | AgentRunResultEvent | TaskNodeResultEvent, None]:
"""
以流式事件模式运行
"""
match task_type:
case TaskType.NONE:
agent = Agent(
model=DEEPSEEK_V4_FLASH_MODEL,
instructions=instruction,
)
async with agent.run_stream_events(
user_prompt=user_prompt,
message_history=message_history,
) as events:
async for event in events:
if task:
match task.type:
case TaskType.BOOK_FLIGHT:
from application.tasks.book_flight import run_stream_events
async for event in run_stream_events(
task=task,
user_prompt=user_prompt,
message_history=message_history,
):
yield event
case TaskType.BOOK_FLIGHT:
from application.tasks.book_flight import run_stream_events, init_task
async for event in run_stream_events(
task=init_task(),
user_prompt=user_prompt,
message_history=message_history,
):
else:
agent = Agent(
model=DEEPSEEK_V4_FLASH_MODEL,
instructions=instruction,
)
async with agent.run_stream_events(
user_prompt=user_prompt,
message_history=message_history,
) as events:
async for event in events:
yield event

234
agent/application/tasks/book_flight.py
View File

@ -1,42 +1,40 @@
# -*- 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 import Agent, ModelMessage, ModelRetry, RunContext, UsageLimits
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,
TaskNodeResultEvent,
TaskType,
usage_limits_to_dict,
usage_limits_to_object,
usage_to_dict,
usage_to_object,
)
from application.tasks.models import DEEPSEEK_V4_FLASH_MODEL
class Deps(BaseModel):
flights: str = Field(..., description="航班信息")
"""
依赖项类
"""
flight_info: str = Field(..., description="所有航班信息")
date: datetime.date = Field(..., description="航班日期")
origin: str = Field(..., description="出发机场")
destination: str = Field(..., description="到达机场")
class FlightDetails(BaseModel):
class FlightDetail(BaseModel):
"""
航班详情
航班详情类
"""
number: str = Field(description="航班号")
@ -48,83 +46,82 @@ class FlightDetails(BaseModel):
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(
flight_search_agent = Agent[Deps, FlightDetail | NoFlightFound](
model=DEEPSEEK_V4_FLASH_MODEL,
output_type=list[FlightDetails],
system_prompt="从给定文本中提取所有航班详细信息",
deps_type=Deps,
output_type=FlightDetail | NoFlightFound,
system_prompt="你的工作是在给定日期、出发机场和到达机场为用户找到最便宜的航班",
)
# 所有航班详情提取智能体
flight_details_extraction_agent = Agent(
model=DEEPSEEK_V4_FLASH_MODEL,
output_type=list[FlightDetail],
system_prompt="从给定文本中提取所有航班详细，包括航班号、航班日期、出发机场、到达机场和机票价格",
)
@master_agent.tool
async def search_flights(ctx: RunContext[Deps]) -> list[FlightDetails]:
@flight_search_agent.tool
async def extract_flight_details(ctx: RunContext[Deps]) -> list[FlightDetail]:
"""
查询并返回航班详情列表
工具：提取所有航班详情
"""
result = await search_agent.run(ctx.deps.flights, usage=ctx.usage)
result = await flight_details_extraction_agent.run(
ctx.deps.flight_info, usage=ctx.usage
)
return result.output
@master_agent.output_validator
@flight_search_agent.output_validator
async def validate_output(
ctx: RunContext[Deps], output: FlightDetails | NoFlightFound
) -> FlightDetails | NoFlightFound:
ctx: RunContext[Deps], output: FlightDetail | NoFlightFound
) -> FlightDetail | NoFlightFound:
"""
校验主控智能体输出
输出校验：航班详情
"""
if isinstance(output, NoFlightFound):
return output
errors = ""
errors = []
if output.date != ctx.deps.date:
errors += f"航班日期应为 {ctx.deps.date}, 不是 {output.date}\n"
errors.append(f"航班日期应为 {ctx.deps.date}, 不是 {output.date}")
if output.origin != ctx.deps.origin:
errors += f"航班出发机场应为 {ctx.deps.origin}, 不是 {output.origin}\n"
errors.append(f"出发机场应为 {ctx.deps.origin}, 不是 {output.origin}")
if output.destination != ctx.deps.destination:
errors += (
f"航班到达机场应为 {ctx.deps.destination}, 不是 {output.destination}\n"
)
errors.append(f"到达机场应为 {ctx.deps.destination}, 不是 {output.destination}")
if errors:
raise ModelRetry(errors)
raise ModelRetry("\n".join(errors))
return output
class SeatPreference(BaseModel):
row: int = Field(ge=1, le=30)
seat: Literal["A", "B", "C", "D", "E", "F"]
"""
座位偏好类
"""
row: int = Field(ge=1, le=30, description="座位行")
column: Literal["A", "B", "C", "D", "E", "F"] = Field(description="座位列")
class Failed(BaseModel):
"""Unable to extract a seat selection."""
class NoSeatExtracted(BaseModel):
"""
未提取到座位偏好类
"""
# 选座智能体
seat_selection_agent = Agent[object, SeatPreference | Failed](
# 提取座位偏好智能体（无依赖项）
seat_preference_extraction_agent = Agent[object, SeatPreference | NoSeatExtracted](
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. "
),
output_type=SeatPreference | NoSeatExtracted,
system_prompt="提取用户的座位偏好。座位规则说明：A 座、F 座为靠窗座位；第 1 排是前排座位，腿部空间更大；14 排、20 排同样拥有加宽腿部空间",
)
flights = """
flight_info = """
1. Flight SFO-AK123
- Price: $350
- Origin: San Francisco International Airport (SFO)
@ -171,9 +168,9 @@ flights = """
def init_task() -> Task:
return Task(
type=TaskType.BOOK_FLIGHT,
node=TaskNode.EXECUTION,
node="flight_search",
deps=Deps(
flights=flights,
flight_info=flight_info,
date=datetime.date(2025, 1, 10),
origin="SFO",
destination="ANC",
@ -183,58 +180,21 @@ def init_task() -> Task:
async def run_stream_events(
task: Task,
user_prompt: str | None = None,
message_history: list[ModelMessage] | None = None,
task: Task | None,
user_prompt: str,
message_history: list[ModelMessage],
) -> 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 not task:
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(
match task.node:
# 航班查询
case "flight_search":
async with flight_search_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),
@ -246,11 +206,57 @@ async def run_stream_events(
result = event.result
# 更新任务使用量
task.usage = usage_to_dict(result.usage)
# 更新任务节点为结束
task.node = TaskNode.FINISH
yield TaskResultEvent(
if isinstance(result.output, FlightDetail):
content = "\n---\n已查询到航班，请回复 buy 购票 或 search 重新查询\n"
# 更新任务节点为提取座位偏好
task.node = "seat_preference_extraction"
else:
content = (
"\n---\n未查询到满足您需求的航班，流程结束！\n"
)
# 更新任务为空
task = None
# 返回任务节点结果事件
yield TaskNodeResultEvent(
task=task,
content="已为您预定好座位，流程结束",
content=content,
)
yield event
return
# 提取座位偏好
case "seat_preference_extraction":
if user_prompt == "buy":
# 返回任务节点结果事件
yield TaskNodeResultEvent(
task=task,
content="请和我说下您的座位偏好：\nA、F 座位是靠窗位；1 排、14 排、20 排腿部空间更大、更舒展，你更想要靠窗座位，宽敞大空间座位",
)
return
elif user_prompt == "search":
# 更新任务节点为航班查询
task.node = "flight_search"
else:
async with seat_preference_extraction_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 = None
# 返回任务节点结果事件
yield TaskNodeResultEvent(
task=task,
content="已帮您定好座位，流程结束！",
)
yield event
return

BIN
agent/database.db
View File

Binary file not shown.