This commit is contained in:
parent
107c7ed03f
commit
ae05009acd
|
|
@ -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:
|
||||
|
|
@ -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"
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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="创建时间",
|
||||
)
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
Binary file not shown.
Loading…
Reference in New Issue