diff --git a/agent/application/states/conversation.py b/agent/application/states/conversation.py index 60cb262..912c5ff 100644 --- a/agent/application/states/conversation.py +++ b/agent/application/states/conversation.py @@ -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 diff --git a/agent/application/states/models.py b/agent/application/states/models.py index 315d4f5..9c57645 100644 --- a/agent/application/states/models.py +++ b/agent/application/states/models.py @@ -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="任务结果内容") diff --git a/agent/application/tasks/__init__.py b/agent/application/tasks/__init__.py index 09f73f8..bd9108f 100644 --- a/agent/application/tasks/__init__.py +++ b/agent/application/tasks/__init__.py @@ -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 diff --git a/agent/application/tasks/book_flight.py b/agent/application/tasks/book_flight.py index 512efe7..e891865 100644 --- a/agent/application/tasks/book_flight.py +++ b/agent/application/tasks/book_flight.py @@ -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 diff --git a/agent/database.db b/agent/database.db index b609ef3..e377289 100644 Binary files a/agent/database.db and b/agent/database.db differ