diff --git a/agent/alembic/versions/4269426de639_.py b/agent/alembic/versions/0d3f238840fe_.py similarity index 96% rename from agent/alembic/versions/4269426de639_.py rename to agent/alembic/versions/0d3f238840fe_.py index e383932..56c1673 100644 --- a/agent/alembic/versions/4269426de639_.py +++ b/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: diff --git a/agent/application/domain_models.py b/agent/application/domain_models.py deleted file mode 100644 index ba95753..0000000 --- a/agent/application/domain_models.py +++ /dev/null @@ -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" diff --git a/agent/application/pages/conversation.py b/agent/application/pages/conversation.py index ddef791..ef7ccb8 100644 --- a/agent/application/pages/conversation.py +++ b/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, diff --git a/agent/application/states/conversation.py b/agent/application/states/conversation.py index 6c6b1db..60cb262 100644 --- a/agent/application/states/conversation.py +++ b/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( - await db_state.create_dialog_record( - conversation_id=self.conversation_id, user_prompt=self.user_prompt - ) - ) - # 将最后一个对话作为当前对话 - dialog = next(reversed(conversation.dialogs.values())) - # 清空前端用户提示词 + # 当前对话 + dialog = Dialog(user_prompt=self.user_prompt) + # 清空用户提示词 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.complete_dialog_record( + # 创建对话记录 + await db_state.create_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() - # 获取数据库状态 - db_state = await self.get_db_state() - # 创建对话记录再添加对话实例 - conversation.dialogs.update( - await db_state.create_dialog_record( - conversation_id=self.conversation_id, result_output="请输入产品需求" - ) - ) - yield + self.user_prompt = f"帮我找一班从 {conversation.task.deps.origin} 到 {conversation.task.deps.destination} 在 {conversation.task.deps.date} 的航班" @rx.event def toggle_collapse(self, dialog_id: str) -> None: diff --git a/agent/application/states/database.py b/agent/application/states/database.py index 2758e3e..85ba999 100644 --- a/agent/application/states/database.py +++ b/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( - dialog_record.thoughts - ), + 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 = "" - ) -> Dict[str, Dialog]: + self, + 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) diff --git a/agent/application/states/models.py b/agent/application/states/models.py new file mode 100644 index 0000000..315d4f5 --- /dev/null +++ b/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="创建时间", + ) diff --git a/agent/application/tasks/__init__.py b/agent/application/tasks/__init__.py index ae53288..09f73f8 100644 --- a/agent/application/tasks/__init__.py +++ b/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 models import DEEPSEEK_V4_FLASH_MODEL +from application.states.models import Dialog, TaskType +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": - yield "未知任务类型" + 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, + ): + yield event diff --git a/agent/application/tasks/book_flight.py b/agent/application/tasks/book_flight.py new file mode 100644 index 0000000..512efe7 --- /dev/null +++ b/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 diff --git a/agent/application/tasks/generating_prd_agent.py b/agent/application/tasks/generating_prd_agent.py deleted file mode 100644 index a08b29a..0000000 --- a/agent/application/tasks/generating_prd_agent.py +++ /dev/null @@ -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 diff --git a/agent/database.db b/agent/database.db index cf1a978..b609ef3 100644 Binary files a/agent/database.db and b/agent/database.db differ