1

agent/alembic/versions/0d3f238840fe_.py

@@ -1,8 +1,8 @@
 """empty message
 
-Revision ID: 4269426de639
+Revision ID: 0d3f238840fe
 Revises: 
-Create Date: 2026-07-27 10:10:56.902466
+Create Date: 2026-08-06 13:35:59.116161
 
 """
 from typing import Sequence, Union
@@ -12,7 +12,7 @@ import sqlalchemy as sa
 import sqlmodel
 
 # revision identifiers, used by Alembic.
-revision: str = '4269426de639'
+revision: str = '0d3f238840fe'
 down_revision: Union[str, Sequence[str], None] = None
 branch_labels: Union[str, Sequence[str], None] = None
 depends_on: Union[str, Sequence[str], None] = None
@@ -51,6 +51,7 @@ def upgrade() -> None:
     sa.Column('user_prompt', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
     sa.Column('thoughts', sa.JSON(), nullable=False),
     sa.Column('result_output', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
+    sa.Column('usage', sa.JSON(), nullable=False),
     sa.PrimaryKeyConstraint('id')
     )
     with op.batch_alter_table('dialogrecord', schema=None) as batch_op:

agent/application/domain_models.py

@@ -1,62 +0,0 @@
-# -*- coding: utf-8 -*-
-"""
-领域模型
-"""
-from datetime import datetime
-from typing import Dict, List, Optional, Callable, Any
-
-from pydantic import BaseModel, Field
-from pydantic_ai import Agent, UsageLimits
-from pydantic_ai._uuid import uuid7
-from enum import StrEnum
-
-
-class Thought(BaseModel):
-    """
-    思考节点领域模型
-    """
-
-    type: str = Field(..., description="思考类型")
-    content: str = Field(..., description="思考内容")
-
-
-class Dialog(BaseModel):
-    """
-    对话领域模型
-    """
-
-    id: str = Field(default_factory=lambda: str(uuid7()), description="对话唯一标识")
-    user_prompt: str = Field(default="", description="用户提示词")
-    thoughts: Dict[int, Thought] = Field(default_factory=dict, description="思考列表")
-    result_output: str = Field(default="", description="结果输出")
-    is_thinking: bool = Field(
-        default=False, description="正在思考，True 表示正在思考，False 表示未正在思考"
-    )
-    is_expanded: bool = Field(
-        default=False, description="思考折叠面板展开状态，True 表示展开，False 表示折叠"
-    )
-
-
-class Conversation(BaseModel):
-    """
-    会话领域模型
-    """
-
-    id: str = Field(default_factory=lambda: str(uuid7()), description="会话唯一标识")
-    description: str = Field(default="新会话", description="会话描述")
-    is_running: bool = Field(
-        default=False, description="正在运行，True 表示正在运行，False 表示未正在运行"
-    )
-    dialogs: Dict[str, Dialog] = Field(default_factory=dict, description="对话列表")
-    created_at: str = Field(
-        ...,
-        description="创建时间",
-    )
-
-
-class TaskType(StrEnum):
-    """
-    任务类型枚举
-    """
-
-    CHAT = "chat"

agent/application/pages/conversation.py

@@ -5,7 +5,7 @@
 import reflex as rx
 from typing import Tuple
 
-from application.domain_models import (
+from application.states.models import (
     Conversation,
     Dialog,
     Thought,

agent/application/states/conversation.py

@@ -2,7 +2,8 @@
 """
 会话状态
 """
-from typing import AsyncGenerator, Dict, List, Optional, Literal
+from typing import AsyncGenerator
+
 from pydantic_ai.messages import (
     FunctionToolCallEvent,
     FunctionToolResultEvent,
@@ -17,17 +18,19 @@ from pydantic_ai.messages import (
     ToolCallPart,
     ToolSearchCallPart,
 )
-
 from pydantic_ai.run import AgentRunResultEvent
 import reflex as rx
 
-
 from application.states.database import DatabaseState
-from application.domain_models import (
+from application.states.models import (
     Conversation,
     Dialog,
-    Thought,
+    Task,
+    TaskResultEvent,
     TaskType,
+    Thought,
+    thoughts_to_dict,
+    usage_to_dict,
 )
 from application.tasks import run_stream_events
 
@@ -40,7 +43,7 @@ class ConversationState(rx.State):
     # 当前用户唯一标识
     user_id: str = ""
     # 键为会话唯一标识，值为会话实例的会话字典
-    conversations: Dict[str, Conversation] = {}  # 按照会话唯一标识顺序排序
+    conversations: dict[str, Conversation] = {}  # 按照会话唯一标识顺序排序
     # 当前会话唯一标识
     conversation_id: str = ""
 
@@ -50,12 +53,12 @@ class ConversationState(rx.State):
     shown_more_conversation_id: str = ""
 
     # 任务类型
-    task_type: TaskType = TaskType.CHAT
+    task_type: TaskType = TaskType.NONE
     # 用户提示词
     user_prompt: str = ""
 
     # 数据库状态（私有变量，不予序列化）
-    _db_state: Optional[DatabaseState] = None
+    _db_state: DatabaseState | None = None
 
     async def get_db_state(self) -> DatabaseState:
         """
@@ -116,7 +119,7 @@ class ConversationState(rx.State):
         # 获取数据库状态
         db_state = await self.get_db_state()
         # 先设置会话记录为已删除再在会话字典中删除会话实例
-        await db_state.set_conversations_record_deleted(conversation_id=conversation_id)
+        await db_state.delete_conversations_record(conversation_id=conversation_id)
         del self.conversations[conversation_id]
 
         # 删除后，若会话字典为空则先创建会话记录再添加会话实例
@@ -189,7 +192,7 @@ class ConversationState(rx.State):
         return conversation.is_running
 
     @rx.var
-    def dialogs(self) -> Dict[str, Dialog]:
+    def dialogs(self) -> dict[str, Dialog]:
         """
         当前会话的对话字典
         :return: 当前会话的对话字典
@@ -209,33 +212,29 @@ class ConversationState(rx.State):
         if not self.user_prompt:
             return
 
-        # 获取数据库状态
-        db_state = await self.get_db_state()
         # 当前会话
         conversation = self.conversations[self.conversation_id]
         # 将正在运行设置为是
         conversation.is_running = True
-        # 创建对话记录再添加对话实例
+        # 当前对话
-        conversation.dialogs.update(
+        dialog = Dialog(user_prompt=self.user_prompt)
-            await db_state.create_dialog_record(
+        # 清空用户提示词
-                conversation_id=self.conversation_id, user_prompt=self.user_prompt
-            )
-        )
-        # 将最后一个对话作为当前对话
-        dialog = next(reversed(conversation.dialogs.values()))
-        # 清空前端用户提示词
         self.user_prompt = ""
+        # 添加对话实例
+        conversation.dialogs.update({dialog.id: dialog})
         # 强制更新会话并推送前端
         self.conversations[self.conversation_id] = conversation
         yield
 
+        # 获取数据库状态
+        db_state = await self.get_db_state()
         # 获取消息历史列表
         message_history = await db_state.get_message_history(
             conversation_id=self.conversation_id
         )
 
         # 初始化工具调用唯一标识和片段索引映射字典
-        tool_call_ids: Dict[str, int] = {}
+        tool_call_ids: dict[str, int] = {}
         # 获取运行流式输出事件
         async for event in run_stream_events(
             task_type=self.task_type,
@@ -363,14 +362,21 @@ class ConversationState(rx.State):
                         case "tool_call":
                             dialog.thoughts[index].content = f"已调用 {content}"
 
+                case TaskResultEvent(task=task, content=content):
+                    # 更新任务
+                    conversation.task = task
+                    dialog.result_output += content
+
                 # ========== 智能体运行结果事件 ==========
                 case AgentRunResultEvent(result=result):
-                    # 获取数据库状态
+                    # 创建对话记录
-                    # 补全对话记录
+                    await db_state.create_dialog_record(
-                    await db_state.complete_dialog_record(
+                        conversation_id=self.conversation_id,
                         id=dialog.id,
-                        thoughts=dialog.thoughts,
+                        user_prompt=self.user_prompt,
+                        thoughts=thoughts_to_dict(dialog.thoughts),
                         result_output=dialog.result_output,
+                        usage=usage_to_dict(result.usage),
                     )
                     # 创建结果记录
                     await db_state.create_result_record(
@@ -378,30 +384,29 @@ class ConversationState(rx.State):
                         dialog_id=dialog.id,
                         new_messages=result.new_messages(),
                     )
-                    # 将正在运行设置为否
-                    conversation.is_running = False
 
             # 强制更新会话并推送前端
             self.conversations[self.conversation_id] = conversation
             yield
 
+        # 将正在运行设置为否
+        conversation.is_running = False
+        self.conversations[self.conversation_id] = conversation
+        yield
+
     @rx.event
-    async def generate_prd(self) -> AsyncGenerator[None]:
+    async def generate_prd(self) -> None:
         """
-        生成产品需求文档
+        预订航班
         """
+        from application.tasks.book_flight import init_task
+
         # 当前会话
         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} 的航班"
-        db_state = await self.get_db_state()
-        # 创建对话记录再添加对话实例
-        conversation.dialogs.update(
-            await db_state.create_dialog_record(
-                conversation_id=self.conversation_id, result_output="请输入产品需求"
-            )
-        )
-        yield
 
     @rx.event
     def toggle_collapse(self, dialog_id: str) -> None:

agent/application/states/database.py

@@ -4,16 +4,15 @@
 """
 from datetime import datetime, timedelta
 from random import choices
-from typing import Dict, List
+from typing import Any
 
-from pydantic import TypeAdapter
 from pydantic_ai import ModelMessage, ModelMessagesTypeAdapter
 from pydantic_ai._uuid import uuid7
 import reflex as rx
 from sqlalchemy import desc
 from sqlmodel import Field, JSON, SQLModel, select, update
 
-from application.domain_models import Conversation, Dialog, Thought
+from application.states.models import Conversation, Dialog
 
 
 class CaptchaRecord(SQLModel, table=True):
@@ -80,36 +79,25 @@ class ConversationRecord(SQLModel, table=True):
     )
 
 
-class ThoughtRecord(SQLModel):
-    """
-    思考记录（不创建数据表）
-    """
-
-    id: str = Field(
-        default_factory=lambda: str(uuid7()),
-        primary_key=True,
-        description="思考唯一标识",
-    )
-    type: str = Field(description="思考类型")
-    content: str = Field(default="", description="思考内容")
-
-
 class DialogRecord(SQLModel, table=True):
     """
     对话记录
     """
 
     id: str = Field(
-        default_factory=lambda: str(uuid7()),
+        ...,
         primary_key=True,
         description="对话唯一标识",
     )
     conversation_id: str = Field(..., index=True, description="会话唯一标识")
     user_prompt: str = Field(..., description="用户提示词")
-    thoughts: Dict[int, ThoughtRecord] = Field(
+    thoughts: dict[int, Any] = Field(
         default_factory=dict, sa_type=JSON, description="思考列表"
     )
     result_output: str = Field(default="", description="结果输出")
+    usage: dict[str, Any] = Field(
+        default_factory=dict, sa_type=JSON, description="使用量"
+    )
 
 
 class ResultRecord(SQLModel, table=True):
@@ -138,10 +126,6 @@ def format_at(at: datetime) -> str:
     return formatted_at
 
 
-# 思考领域模型类型适配器
-ThoughtTypeAdapter = TypeAdapter(Dict[int, Thought])
-
-
 class DatabaseState(rx.State):
     """
     数据库状态
@@ -219,13 +203,13 @@ class DatabaseState(rx.State):
             await session.refresh(record)
             return record.id
 
-    async def get_conversations(self, user_id: str) -> Dict[str, Conversation]:
+    async def get_conversations(self, user_id: str) -> dict[str, Conversation]:
         """
         获取会话字典
         :param user_id: 用户唯一标识
         :return: 会话字典
         """
-        records: Dict[str, Conversation] = {}
+        records: dict[str, Conversation] = {}
         async with rx.asession() as session:
             result = await session.exec(
                 select(ConversationRecord, DialogRecord)
@@ -237,7 +221,6 @@ class DatabaseState(rx.State):
                 .order_by(ConversationRecord.id, DialogRecord.id)
             )
             for conversation_record, dialog_record in result.all():
-
                 record = records.setdefault(
                     conversation_record.id,
                     Conversation(
@@ -253,10 +236,9 @@ class DatabaseState(rx.State):
                         dialog_record.id: Dialog(
                             id=dialog_record.id,
                             user_prompt=dialog_record.user_prompt,
-                            thoughts=ThoughtTypeAdapter.validate_python(
+                            thoughts=dialog_record.thoughts,
-                                dialog_record.thoughts
-                            ),
                             result_output=dialog_record.result_output,
+                            usage=dialog_record.usage,
                         )
                     }
                 )
@@ -264,7 +246,7 @@ class DatabaseState(rx.State):
 
     async def create_conversations_record(
         self, user_id: str, description: str = "新会话"
-    ) -> Dict[str, Conversation]:
+    ) -> dict[str, Conversation]:
         """
         创建会话记录
         :param user_id: 用户唯一标识
@@ -284,9 +266,9 @@ class DatabaseState(rx.State):
             )
         }
 
-    async def set_conversations_record_deleted(self, conversation_id: str) -> None:
+    async def delete_conversations_record(self, conversation_id: str) -> None:
         """
-        设置会话记录为已删除
+        删除会话记录（逻辑删除）
         :param conversation_id: 指定会话唯一标识
         :return: None
         """
@@ -298,8 +280,14 @@ class DatabaseState(rx.State):
             await session.commit()
 
     async def create_dialog_record(
-        self, conversation_id: str, user_prompt: str = "", result_output: str = ""
+        self,
-    ) -> Dict[str, Dialog]:
+        conversation_id: str,
+        id: str,
+        user_prompt: str,
+        thoughts: dict[int, Any],
+        result_output: str,
+        usage: dict[str, Any],
+    ) -> dict[str, Dialog]:
         """
         创建对话记录
         :param conversation_id: 会话唯一标识
@@ -309,45 +297,31 @@ class DatabaseState(rx.State):
         """
         async with rx.asession() as session:
             record = DialogRecord(
+                id=id,
                 conversation_id=conversation_id,
                 user_prompt=user_prompt,
+                thoughts=thoughts,
                 result_output=result_output,
+                usage=usage,
             )
             session.add(record)
             await session.commit()
             await session.refresh(record)
         return {
             record.id: Dialog(
-                id=record.id, user_prompt=user_prompt, result_output=result_output
+                id=record.id,
+                user_prompt=record.user_prompt,
+                thoughts=record.thoughts,
+                result_output=record.result_output,
+                usage=record.usage,
             )
         }
 
-    async def complete_dialog_record(
-        self,
-        id: str,
-        thoughts: Dict[int, Thought],
-        result_output: str,
-    ) -> None:
-        """
-        补全对话记录
-        :param id: 对话唯一标识
-        :param thoughts: 思考列表
-        :param result_output: 结果输出
-        :return: None
-        """
-        async with rx.asession() as session:
-            record = await session.get(DialogRecord, id)  # 通过主键查询记录
-            if not record:
-                return
-            record.thoughts = ThoughtTypeAdapter.dump_python(thoughts)
-            record.result_output = result_output
-            await session.commit()
-
     async def create_result_record(
         self,
         conversation_id: str,
         dialog_id: str,
-        new_messages: List[ModelMessage],
+        new_messages: list[ModelMessage],
     ) -> None:
         """
         创建结果记录
@@ -370,13 +344,13 @@ class DatabaseState(rx.State):
             )
             await session.commit()
 
-    async def get_message_history(self, conversation_id: str) -> List[ModelMessage]:
+    async def get_message_history(self, conversation_id: str) -> list[ModelMessage]:
         """
         获取消息历史列表
         :param conversation_id: 会话唯一标识
         :return: 消息历史
         """
-        records: List[ModelMessage] = []
+        records: list[ModelMessage] = []
         async with rx.asession() as session:
             result = await session.exec(
                 select(ResultRecord)

agent/application/states/models.py

@@ -0,0 +1,151 @@
+# -*- coding: utf-8 -*-
+"""
+面向 reflex.state 的类
+"""
+from datetime import datetime
+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
+
+
+class Thought(BaseModel):
+    """
+    思考类
+    """
+
+    type: str = Field(..., description="思考类型")
+    content: str = Field(..., description="思考内容")
+
+
+# 思考领域模型类型适配器
+ThoughtsAdapter = TypeAdapter(dict[int, Thought])
+
+
+def thoughts_to_dict(thoughts: dict[int, Thought]) -> dict:
+    """
+    Thought 转为字典
+    """
+    return ThoughtsAdapter.dump_python(thoughts)
+
+
+class Dialog(BaseModel):
+    """
+    对话类
+    """
+
+    id: str = Field(default_factory=lambda: str(uuid7()), description="对话唯一标识")
+    user_prompt: str = Field(default="", description="用户提示词")
+    thoughts: dict[int, Thought] = Field(default_factory=dict, description="思考列表")
+    result_output: str = Field(default="", description="结果输出")
+    usage: dict[str, Any] = Field(default_factory=dict, description="对话使用量")
+    is_thinking: bool = Field(
+        default=False, description="正在思考，True 表示正在思考，False 表示未正在思考"
+    )
+    is_expanded: bool = Field(
+        default=False, description="思考折叠面板展开状态，True 表示展开，False 表示折叠"
+    )
+
+
+# Usage 适配器
+UsageAdapter = TypeAdapter(RunUsage)
+
+
+def usage_to_object(usage: dict) -> RunUsage:
+    """
+    Usage 转为对象
+    """
+    return UsageAdapter.validate_python(usage) if usage else RunUsage()
+
+
+def usage_to_dict(usage: RunUsage) -> dict:
+    """
+    Usage 转为字典
+    """
+    return UsageAdapter.dump_python(usage)
+
+
+# UsageLimits  适配器
+UsageLimitsAdapter = TypeAdapter(UsageLimits)
+
+
+def usage_limits_to_object(usage_limits: dict) -> UsageLimits:
+    """
+    UsageLimits 转为对象
+    """
+    return (
+        UsageLimitsAdapter.validate_python(usage_limits)
+        if usage_limits
+        else UsageLimits()
+    )
+
+
+def usage_limits_to_dict(usage_limits: UsageLimits) -> dict:
+    """
+    UsageLimits 转为字典
+    """
+    return UsageLimitsAdapter.dump_python(usage_limits)
+
+
+class TaskType(StrEnum):
+    """
+    任务类型枚举
+    """
+
+    NONE = "none"
+    BOOK_FLIGHT = "book_flight"
+
+
+class TaskNode(StrEnum):
+    """
+    任务节点枚举
+    """
+
+    EXECUTION = "execution"
+    PENDING_INPUT = "pending_input"
+    FINISH = "finish"
+
+
+class Task(BaseModel):
+    """
+    任务类
+    """
+
+    id: str = Field(default_factory=lambda: str(uuid7()), description="任务唯一标识")
+    type: TaskType = Field(..., description="任务类型")
+    node: TaskNode = Field(default=TaskNode.EXECUTION, description="任务节点")
+    deps: Any = Field(default=None, description="任务依赖项")
+    usage: dict[str, Any] = Field(default_factory=dict, description="任务使用量")
+    usage_limits: dict[str, Any] = Field(
+        default_factory=dict, description="任务使用量限制"
+    )
+
+
+class TaskResultEvent(BaseModel):
+    """
+    任务结果事件类
+    """
+
+    task: Task = Field(..., description="任务实例")
+    content: str = Field(default="", description="任务结果内容")
+
+
+class Conversation(BaseModel):
+    """
+    会话类
+    """
+
+    id: str = Field(default_factory=lambda: str(uuid7()), description="会话唯一标识")
+    description: str = Field(default="新会话", description="会话描述")
+    is_running: bool = Field(
+        default=False,
+        description="会话正在运行，True 表示正在运行，False 表示未正在运行",
+    )
+    dialogs: dict[str, Dialog] = Field(default_factory=dict, description="对话字典")
+    task: Task | None = Field(default=None, description="会话任务")
+    created_at: str = Field(
+        ...,
+        description="创建时间",
+    )

agent/application/tasks/__init__.py

@@ -10,8 +10,9 @@ from pydantic_ai import Agent, ModelMessage
 from pydantic_ai.models.openai import OpenAIChatModel
 from pydantic_ai.providers.openai import OpenAIProvider
 
-from application.domain_models import Dialog, TaskType
+from application.states.models import Dialog, TaskType
-from models import DEEPSEEK_V4_FLASH_MODEL
+from application.tasks.models import DEEPSEEK_V4_FLASH_MODEL
+from application.states.models import TaskNode, Task
 
 instruction = """
 # 角色
@@ -39,7 +40,7 @@ async def run_stream_events(
     以流式事件模式运行
     """
     match task_type:
-        case TaskType.CHAT:
+        case TaskType.NONE:
             agent = Agent(
                 model=DEEPSEEK_V4_FLASH_MODEL,
                 instructions=instruction,
@@ -51,5 +52,12 @@ async def run_stream_events(
                 async for event in events:
                     yield event
 
-        case "flight":
+        case TaskType.BOOK_FLIGHT:
-            yield "未知任务类型"
+            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,
+            ):
+                yield event

agent/application/tasks/book_flight.py

@@ -0,0 +1,256 @@
+# -*- 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

agent/application/tasks/generating_prd_agent.py

@@ -1,210 +0,0 @@
-# -*- 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