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

1

This commit is contained in:
liubiren 2026-08-10 14:36:19 +08:00
parent 72dae054f1
commit 3dde59d8be
6 changed files with 264 additions and 237 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

30
agent/application/pages/index.py
View File

@ -58,7 +58,7 @@ def nav_button(
gap="4px", gap="4px",
cursor="pointer", cursor="pointer",
# 点击事件：将指定导航按钮设置为激活的导航按钮 # 点击事件：将指定导航按钮设置为激活的导航按钮
on_click=AuthState.switch_activated_nav_button(nav_button), on_click=AuthState.set_activated_nav_button(nav_button),
) )
@ -234,8 +234,8 @@ def login_popup() -> rx.Component:
value=AuthState.email, value=AuthState.email,
# 值变化事件：设置邮箱 # 值变化事件：设置邮箱
on_change=AuthState.set_email, on_change=AuthState.set_email,
# 获取焦点事件：重置登录错误信息 # 获取焦点事件：重置错误信息
on_focus=AuthState.reset_login_error_message, on_focus=AuthState.reset_error,
width="100%", width="100%",
height="45px", height="45px",
padding="4px", padding="4px",
@ -248,12 +248,12 @@ def login_popup() -> rx.Component:
rx.input( rx.input(
placeholder="请输入验证码", placeholder="请输入验证码",
# 值绑定验证码 # 值绑定验证码
value=AuthState.captcha, value=AuthState.captcha_code,
max_length=6, max_length=6,
# 值变化事件：设置验证码 # 值变化事件：设置验证码
on_change=AuthState.set_captcha, on_change=AuthState.set_captcha_code,
# 获取焦点事件：重置登录错误信息 # 获取焦点事件：重置错误信息
on_focus=AuthState.reset_login_error_message, on_focus=AuthState.reset_error,
flex=1, flex=1,
height="100%", height="100%",
padding="4px", padding="4px",
@ -264,25 +264,25 @@ def login_popup() -> rx.Component:
rx.button( rx.button(
# 若再次发送验证码倒计时大于0则显示倒计时，否则显示获取验证码 # 若再次发送验证码倒计时大于0则显示倒计时，否则显示获取验证码
rx.cond( rx.cond(
AuthState.resend_captcha_countdown > 0, AuthState.resend_captcha_code_countdown > 0,
f"{AuthState.resend_captcha_countdown} s", f"{AuthState.resend_captcha_code_countdown} s",
"获取验证码", "获取验证码",
), ),
# 点击事件：发送验证码 # 点击事件：发送验证码
on_click=AuthState.send_captcha, on_click=AuthState.send_captcha_code,
# 是否禁用绑定验证码发送禁用状态 # 是否禁用绑定验证码发送禁用状态
disabled=AuthState.is_captcha_sending_disabled, disabled=AuthState.is_captcha_code_sending_button_disabled,
margin="0 8px", margin="0 8px",
background_color="var(--prismui-background-color-1)", background_color="var(--prismui-background-color-1)",
line_height="16px", line_height="16px",
# 若验证码发送禁用状态为是则颜色设置为 color-4，否则设置为 color-5（主题色） # 若验证码发送禁用状态为是则颜色设置为 color-4，否则设置为 color-5（主题色）
color=rx.cond( color=rx.cond(
AuthState.is_captcha_sending_disabled, AuthState.is_captcha_code_sending_button_disabled,
"var(--prismui-color-4)", "var(--prismui-color-4)",
"var(--prismui-color-5)", "var(--prismui-color-5)",
), ),
cursor=rx.cond( cursor=rx.cond(
AuthState.is_captcha_sending_disabled, AuthState.is_captcha_code_sending_button_disabled,
"not-allowed", "not-allowed",
"pointer", "pointer",
), ),
@ -294,7 +294,7 @@ def login_popup() -> rx.Component:
border_radius="var(--prismui-border-radius-4)", border_radius="var(--prismui-border-radius-4)",
), ),
rx.text( rx.text(
AuthState.login_error_message, AuthState.error,
width="100%", width="100%",
min_height="20px", min_height="20px",
text_align="left", text_align="left",
@ -304,7 +304,7 @@ def login_popup() -> rx.Component:
), ),
rx.box( rx.box(
rx.cond( rx.cond(
AuthState.is_captcha_sent, AuthState.is_captcha_code_sent,
rx.hstack( rx.hstack(
rx.icon( rx.icon(
"circle-check", "circle-check",

150
agent/application/states/auth.py
View File

@ -17,7 +17,7 @@ from application.states.database import DatabaseState
# 验证码HTML模板 # 验证码HTML模板
CAPTCHA_HTML_TEMPLATE = """ CAPTCHA_HTML_TEMPLATE = """
<p>验证码：<strong>{captcha}</strong></p> <p>验证码：<strong>{captcha_code}</strong></p>
""" """
@ -32,34 +32,43 @@ class AuthState(rx.State):
# 邮箱 # 邮箱
email: str = "" email: str = ""
# 验证码 # 验证码
captcha: str = "" captcha_code: str = ""
# 验证码已发送
is_captcha_code_sent: bool = False
# 再次发送验证码倒计时
resend_captcha_code_countdown: int = 0
# 协议同意状态 # 协议同意状态
is_policies_agreed: bool = True is_policies_agreed: bool = True
# 错误信息
# 登录错误信息 error: str = ""
login_error_message: str = ""
# 再次发送验证码倒计时
resend_captcha_countdown: int = 0
# 验证码已发送状态
is_captcha_sent: bool = False
# 登录中状态 # 登录中状态
is_logging_in: bool = False is_logging_in: bool = False
# 设置悬停卡片打开状态
is_settings_hover_card_open: bool = False
# 激活的导航按钮 # 激活的导航按钮
activated_nav_button: str = "conversation" activated_nav_button: str = "conversation"
# 设置悬停卡片打开状态
is_settings_hover_card_open: bool = False
# 数据库状态 # 数据库状态
_db_state: Optional[DatabaseState] = None _db_state: Optional[DatabaseState] = None
@rx.event async def get_db_state(self) -> DatabaseState:
def reset_login_error_message(self) -> None:
""" """
重置登录错误信息 获取当前数据库状态
:return: 当前数据库状态
"""
# 若数据库状态为空则获取数据库状态，否则直接返回当前数据库状态
if not self._db_state:
self._db_state = await self.get_state(DatabaseState)
return self._db_state
@rx.event
def reset_error(self) -> None:
"""
重置错误信息
:return: None :return: None
""" """
self.login_error_message = "" self.error = ""
@rx.event @rx.event
def set_email(self, email: str) -> None: def set_email(self, email: str) -> None:
@ -73,11 +82,11 @@ class AuthState(rx.State):
def validate_email(self) -> bool: def validate_email(self) -> bool:
""" """
校验邮箱 校验邮箱
:return: 校验是否通过，True 表示通过，False 表示不通过 :return: 校验是否通过
""" """
# 邮箱非空校验 # 邮箱非空校验
if not self.email: if not self.email:
self.login_error_message = "请输入邮箱" self.error = "请输入邮箱"
return False return False
# 邮箱格式校验 # 邮箱格式校验
@ -92,51 +101,50 @@ class AuthState(rx.State):
r"^[a-zA-Z0-9]+([_\-+.][a-zA-Z0-9_\-+]+)*@[a-zA-Z0-9]+(\-[a-zA-Z0-9\-]+)*(\.[a-zA-Z0-9]+(\-[a-zA-Z0-9\-]+)*)*\.[a-zA-Z0-9]{2,}$", r"^[a-zA-Z0-9]+([_\-+.][a-zA-Z0-9_\-+]+)*@[a-zA-Z0-9]+(\-[a-zA-Z0-9\-]+)*(\.[a-zA-Z0-9]+(\-[a-zA-Z0-9\-]+)*)*\.[a-zA-Z0-9]{2,}$",
self.email, self.email,
): ):
self.login_error_message = "邮箱格式不正确，请重新输入" self.error = "邮箱格式不正确，请重新输入"
return False return False
self.login_error_message = "" self.error = ""
return True return True
@rx.var @rx.var
def is_captcha_sending_disabled(self) -> bool: def is_captcha_code_sending_button_disabled(self) -> bool:
""" """
验证码发送禁用状态 验证码发送按钮不可点击状态
邮箱为空或再次发送验证码倒计时大于0时禁用 邮箱为空或再次发送验证码倒计时大于0时禁用
:return: 禁用状态，True 表示禁用，False 表示未禁用 :return: bool
""" """
return not self.email or self.resend_captcha_countdown > 0 return not self.email or self.resend_captcha_code_countdown > 0
@rx.event(background=True) @rx.event(background=True)
async def countdown(self) -> None: async def countdown_background(self) -> None:
""" """
倒计时（后台任务） 倒计时（后台任务）
:return: None :return: None
""" """
async with self: async with self:
self.resend_captcha_countdown = 59 self.resend_captcha_code_countdown = 59
while self.resend_captcha_countdown > 0: while self.resend_captcha_code_countdown > 0:
await sleep(1) await sleep(1)
async with self: async with self:
self.resend_captcha_countdown -= 1 self.resend_captcha_code_countdown -= 1
@rx.event(background=True) @rx.event(background=True)
async def send_captcha_background(self) -> None: async def send_captcha_code_background(self) -> None:
""" """
发送验证码（后台任务） 发送验证码（后台任务）
:return: None :return: None
""" """
if not self._db_state: db_state = await self.get_db_state()
return
async with self: async with self:
# 创建验证码记录并获取验证码 # 创建验证记录
captcha = await self._db_state.create_captcha_record(email=self.email) captcha_code = await db_state.create_captcha_record(email=self.email)
# 构建验证码邮件 # 构建验证码邮件
message = MIMEText( message = MIMEText(
CAPTCHA_HTML_TEMPLATE.format(captcha=captcha), CAPTCHA_HTML_TEMPLATE.format(captcha_code=captcha_code),
"html", "html",
"utf-8", "utf-8",
) )
@ -158,7 +166,7 @@ class AuthState(rx.State):
pass pass
@rx.event @rx.event
async def send_captcha(self) -> AsyncGenerator[EventCallback, None]: async def send_captcha_code(self) -> AsyncGenerator[EventCallback, None]:
""" """
发送验证码 发送验证码
:return: AsyncGenerator[EventCallback, None] :return: AsyncGenerator[EventCallback, None]
@ -167,27 +175,24 @@ class AuthState(rx.State):
if not self.validate_email(): if not self.validate_email():
return return
if self.resend_captcha_countdown > 0: if self.resend_captcha_code_countdown > 0:
return return
self.is_captcha_sent = True # 标记验证码已发送
self.is_captcha_code_sent = True
# 获取数据库状态 # 同时触发发送验证码和倒计时后台事件
if not self._db_state: yield type(self).send_captcha_code_background()
self._db_state = await self.get_state(DatabaseState) yield type(self).countdown_background()
# 将倒计时和发送验证码事件添加至后台任务队列
yield type(self).countdown()
yield type(self).send_captcha_background()
@rx.event @rx.event
def set_captcha(self, captcha: str) -> None: def set_captcha_code(self, captcha_code: str) -> None:
""" """
设置验证码 设置验证码
:param captcha: 验证码 :param captcha_code: 验证码
:return: None :return: None
""" """
self.captcha = captcha.strip() self.captcha_code = captcha_code.strip()
@rx.event @rx.event
def toggle_policies_agreed(self) -> None: def toggle_policies_agreed(self) -> None:
@ -204,7 +209,7 @@ class AuthState(rx.State):
""" """
# 恢复会话状态 # 恢复会话状态
conversation_state = await self.get_state(ConversationState) conversation_state = await self.get_state(ConversationState)
await conversation_state.resume(user_id=self.user_id) await conversation_state.resume_conversation_state(user_id=self.user_id)
@rx.event @rx.event
async def login(self) -> None: async def login(self) -> None:
@ -217,77 +222,76 @@ class AuthState(rx.State):
return return
# 校验验证码非空 # 校验验证码非空
if not self.captcha: if not self.captcha_code:
self.login_error_message = "请输入验证码" self.error = "请输入验证码"
return return
# 校验验证码格式 # 校验验证码格式
if not re.fullmatch(r"^\d{6}$", self.captcha): if not re.fullmatch(r"^\d{6}$", self.captcha_code):
self.login_error_message = "验证码格式不正确，请重新输入" self.error = "验证码格式不正确，请重新输入"
return return
# 校验协议同意状态 # 校验协议同意状态
if not self.is_policies_agreed: if not self.is_policies_agreed:
self.login_error_message = "请先阅读并同意协议和政策" self.error = "请先阅读并同意协议和政策"
return return
self.is_logging_in = True self.is_logging_in = True
# 获取数据库状态 # 获取数据库状态
if not self._db_state: db_state = await self.get_db_state()
return
# 核验验证码 # 核验验证码
if not await self._db_state.verify_captcha( if not await db_state.verify_captcha_code(
email=self.email, captcha=self.captcha email=self.email, captcha_code=self.captcha_code
): ):
self.is_logging_in = False self.is_logging_in = False
self.login_error_message = "验证码错误" self.error = "验证码错误"
return return
# 创建用户记录 # 创建用户记录
user_id = await self._db_state.create_user_record(email=self.email) user_id = await db_state.create_user_record(email=self.email)
# 恢复当前用户会话状态 # 恢复当前用户会话状态
conversation_state = await self.get_state(ConversationState) conversation_state = await self.get_state(ConversationState)
await conversation_state.resume(user_id=user_id) await conversation_state.resume_conversation_state(user_id=user_id)
self.user_id = user_id self.user_id = user_id
self.email = "" self.email = ""
self.captcha = "" self.captcha = ""
self.is_captcha_code_sent = False
self.resend_captcha_code_countdown = 0
self.is_policies_agreed = True self.is_policies_agreed = True
self.login_error_message = "" self.error = ""
self.resend_captcha_countdown = 0
self.is_captcha_sent = False
self.is_logging_in = False self.is_logging_in = False
@rx.event @rx.event
def set_settings_hover_card_open(self, is_settings_hover_card_open: bool) -> None: def set_settings_hover_card_open(self, is_settings_hover_card_open: bool) -> None:
""" """
设置设置悬停卡片打开状态 设置设置悬停卡片打开
:param is_settings_hover_card_open: 设置悬停卡片打开状态，True 表示打开，False 表示关闭 :param is_settings_hover_card_open: 设置悬停卡片打开状态，True 表示打开，False 表示关闭
:return: None :return: None
""" """
self.is_settings_hover_card_open = is_settings_hover_card_open self.is_settings_hover_card_open = is_settings_hover_card_open
@rx.event
def set_activated_nav_button(self, nav_button: str):
"""
设置激活导航按钮
"""
self.activated_nav_button = nav_button
@rx.event @rx.event
async def logout(self) -> None: async def logout(self) -> None:
""" """
退出登录 退出登录
:return: None :return: None
""" """
self.is_settings_hover_card_open = False self.is_settings_hover_card_open = False # 设置设置悬停卡片关闭
conversation_state = await self.get_state(ConversationState) conversation_state = await self.get_state(ConversationState)
conversation_state.user_id = "" conversation_state.user_id = ""
conversation_state.conversations = {} conversation_state.conversations = {}
conversation_state.conversation_id = "" conversation_state.conversation_id = ""
conversation_state.is_conversation_history_shown = False conversation_state.is_conversation_history_shown = False
self.user_id = "" self.user_id = ""
@rx.event
def switch_activated_nav_button(self, button: str):
"""
将指定导航按钮设置为激活的导航按钮
"""
self.activated_nav_button = button

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

@ -21,19 +21,22 @@ from pydantic_ai.messages import (
) )
from pydantic_ai.run import AgentRunResultEvent from pydantic_ai.run import AgentRunResultEvent
import reflex as rx import reflex as rx
from pydantic_ai._uuid import uuid7
from application.states.database import DatabaseState from application.states.database import DatabaseState
from application.states.models import ( from application.states.models import (
Conversation, Conversation,
Message,
MessageType,
Work,
ConversationHistoryItem, ConversationHistoryItem,
usage_to_dict, usage_to_dict,
) )
from application.tasks import run_stream_events from application.tasks import run_stream_events
def format_conversation_history_item_created_at(created_at: datetime) -> str: def format_conversation_created_at(created_at: datetime) -> str:
""" """
格式化会话历史项创建日期时间 格式化会话创建日期时间
:param created_at: 创建日期时间 :param created_at: 创建日期时间
:return: 格式化后的日期时间字符串 :return: 格式化后的日期时间字符串
""" """
@ -55,8 +58,9 @@ class ConversationState(rx.State):
# 当前用户唯一标识 # 当前用户唯一标识
user_id: str = "" user_id: str = ""
# 当前用户的会话字典（私有变量） # 当前用户的会话字典（私有变量）
# 私有变量：reflex 约定以 _ 开头的变量为私有变量，后端不序列化，前端不可使用
_conversations: dict[str, Conversation] = {} # 键为会话唯一标识，值为会话实例 # 当前用户的会话字典
conversations: dict[str, Conversation] = {} # 键为会话唯一标识，值为会话实例
# 激活会话唯一标识 # 激活会话唯一标识
actived_conversation_id: str = "" actived_conversation_id: str = ""
@ -65,10 +69,11 @@ class ConversationState(rx.State):
# 会话历史项显示悬停气泡，True 表示显示，False 表示隐藏 # 会话历史项显示悬停气泡，True 表示显示，False 表示隐藏
is_conversation_history_item_popover_shown: bool = False is_conversation_history_item_popover_shown: bool = False
# 用户提示词 # 当前工作实例
user_prompt: str = "" work: Work | None = None
# 当前数据库状态（私有变量） # 当前数据库状态（私有变量）
# 私有变量：reflex 约定以 _ 开头的变量为私有变量，后端不序列化，前端不可使用
_db_state: DatabaseState | None = None _db_state: DatabaseState | None = None
async def get_db_state(self) -> DatabaseState: async def get_db_state(self) -> DatabaseState:
@ -94,14 +99,14 @@ class ConversationState(rx.State):
# 获取当前数据库状态 # 获取当前数据库状态
db_state = await self.get_db_state() db_state = await self.get_db_state()
# 获取当前用户的会话字典 # 获取当前用户的会话字典
self._conversations = await db_state.get_conversations(self.user_id) self.conversations = await db_state.get_conversations(self.user_id)
# 若当前用户的会话字典为空则先创建会话记录再添加会话实例 # 若当前用户的会话字典为空则先创建会话记录再添加会话实例
if not self._conversations: if not self.conversations:
self._conversations.update( self.conversations.update(
await db_state.create_conversation_record(self.user_id) await db_state.create_conversation_record(self.user_id)
) )
# 将最后一个会话的唯一标识作为激活会话唯一标识 # 将最后一个会话的唯一标识作为激活会话唯一标识
self.actived_conversation_id = next(reversed(self._conversations.keys())) self.actived_conversation_id = next(reversed(self.conversations.keys()))
@rx.event @rx.event
def toggle_conversation_history_shown(self) -> None: def toggle_conversation_history_shown(self) -> None:
@ -120,11 +125,9 @@ class ConversationState(rx.State):
return [ return [
ConversationHistoryItem( ConversationHistoryItem(
description=conversation.description, description=conversation.description,
created_at=format_conversation_history_item_created_at( created_at=format_conversation_created_at(conversation.created_at),
conversation.created_at
),
) )
for conversation in reversed(self._conversations.values()) for conversation in reversed(self.conversations.values())
] ]
@rx.event @rx.event
@ -145,17 +148,17 @@ class ConversationState(rx.State):
# 获取当前数据库状态 # 获取当前数据库状态
db_state = await self.get_db_state() db_state = await self.get_db_state()
await db_state.delete_conversation_record(conversation_id) await db_state.delete_conversation_record(conversation_id)
del self._conversations[conversation_id] del self.conversations[conversation_id]
# 删除后，若当前用户的会话字典为空则创建会话 # 删除后，若当前用户的会话字典为空则创建会话
if not self._conversations: if not self.conversations:
self._conversations.update( self.conversations.update(
await db_state.create_conversation_record(self.user_id) await db_state.create_conversation_record(self.user_id)
) )
# 删除后，若激活会话唯一标识不存在则将最后一个会话的唯一标识作为激活会话唯一标识 # 删除后，若激活会话唯一标识不存在则将最后一个会话的唯一标识作为激活会话唯一标识
if self.actived_conversation_id not in self._conversations: if self.actived_conversation_id not in self.conversations:
self.actived_conversation_id = next(reversed(self._conversations.keys())) self.actived_conversation_id = next(reversed(self.conversations.keys()))
@rx.event @rx.event
def set_actived_conversation(self, conversation_id: str) -> None: def set_actived_conversation(self, conversation_id: str) -> None:
@ -175,11 +178,11 @@ class ConversationState(rx.State):
# 获取当前数据库状态 # 获取当前数据库状态
db_state = await self.get_db_state() db_state = await self.get_db_state()
# 创建会话记录再添加会话实例 # 创建会话记录再添加会话实例
self._conversations.update( self.conversations.update(
await db_state.create_conversation_record(self.user_id) await db_state.create_conversation_record(self.user_id)
) )
# 将最后一个会话的唯一标识作为激活会话唯一标识 # 将最后一个会话的唯一标识作为激活会话唯一标识
self.actived_conversation_id = next(reversed(self._conversations.keys())) self.actived_conversation_id = next(reversed(self.conversations.keys()))
@rx.event @rx.event
def set_user_prompt(self, user_prompt: str) -> None: def set_user_prompt(self, user_prompt: str) -> None:
@ -188,81 +191,75 @@ class ConversationState(rx.State):
:param user_prompt: 用户提示词 :param user_prompt: 用户提示词
:return: None :return: None
""" """
self.user_prompt = user_prompt.strip() # 当前会话
conversation = self.conversations.get(self.actived_conversation_id)
if not conversation:
return
conversation.user_prompt = user_prompt.strip()
@rx.var @rx.var
def is_user_prompt_send_button_disabled(self) -> bool: def is_user_prompt_send_button_disabled(self) -> bool:
""" """
用户提示词发送按钮不可点击 用户提示词发送按钮不可点击状态
:return: bool :return: bool
""" """
# 当前会话 # 当前会话
conversation = self._conversations.get(self.actived_conversation_id) conversation = self.conversations.get(self.actived_conversation_id)
if not conversation: if not conversation:
return True return True
return conversation.is_running or not self.user_prompt return conversation.is_running or not conversation.user_prompt
@rx.var @rx.var
def running_status(self) -> bool: def is_running(self) -> bool:
""" """
获取当前会话的运行状态 获取当前会话的运行状态
:return: 当前会话的运行状态（True 表示正在运行，False 表示运行完成） :return: 当前会话的运行状态（True 表示正在运行，False 表示运行完成）
""" """
# 当前会话 # 当前会话
conversation = self.conversations.get(self.conversation_id) conversation = self.conversations.get(self.actived_conversation_id)
if not conversation: if not conversation:
return False return False
return conversation.is_running return conversation.is_running
@rx.var
def dialogs(self) -> dict[str, Dialog]:
"""
当前会话的对话字典
:return: 当前会话的对话字典
"""
# 当前会话
conversation = self.conversations.get(self.conversation_id)
if not conversation:
return {}
return conversation.dialogs
@rx.event @rx.event
async def handle_user_prompt(self) -> AsyncGenerator[None]: async def run(self) -> AsyncGenerator[None]:
""" """
处理用户提示词 运行
:return: AsyncGenerator[None] :return: AsyncGenerator[None]
""" """
if not self.user_prompt: # 当前会话
conversation = self.conversations.get(self.actived_conversation_id)
if not conversation:
return
# 若用户提示词为空则直接返回
if not conversation.user_prompt:
return return
# 当前会话
conversation = self.conversations[self.conversation_id]
# 将正在运行设置为是 # 将正在运行设置为是
conversation.is_running = True conversation.is_running = True
# 当前对话 # 构建用户提示词消息实例
dialog = Dialog(user_prompt=self.user_prompt) message = Message(
type=MessageType.USER_PROMPT, content=conversation.user_prompt
)
# 清空用户提示词 # 清空用户提示词
self.user_prompt = "" conversation.user_prompt = ""
# 添加对话实例 # 添加至消息字典
conversation.dialogs.update({dialog.id: dialog}) user_prompt = conversation.messages.setdefault(message.id, message)
# 强制更新会话并推送前端 yield # 通知前端更新渲染
self.conversations[self.conversation_id] = conversation
yield
# 获取数据库状态 # 获取数据库状态
db_state = await self.get_db_state() 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] = {} index_map_to_message_id: dict[int, str] = {}
# 获取运行流式输出事件 # 获取运行流式输出事件
async for event in run_stream_events( async for event in run_stream_events(
task=conversation.task, work=self.work,
user_prompt=dialog.user_prompt, user_prompt=user_prompt,
message_history=message_history, message_history=await db_state.get_message_history(
conversation_id=self.actived_conversation_id
),
): ):
match event: match event:
# ========== 开始事件 ========== # ========== 开始事件 ==========
@ -274,12 +271,14 @@ class ConversationState(rx.State):
match part: match part:
# 思考分片开始事件 # 思考分片开始事件
case ThinkingPart(content=content): case ThinkingPart(content=content):
# 若上一分片种类为空则将正在思考设置为正在思考 message = Message(type=MessageType.THINKING)
# 将消息实例唯一标识与片段索引映射
message_id_map_to_index[index] = message.id
# 若上一分片种类为空则将正在思考设置为是
if not previous_part_kind: if not previous_part_kind:
dialog.is_thinking = True # 构建消息实例
dialog.thoughts[index] = Thought( message.is_thinking = True
type="thinking", content=content message.content = content
)
# 工具检索分片开始事件 # 工具检索分片开始事件
case ToolSearchCallPart(tool_call_id=tool_call_id): case ToolSearchCallPart(tool_call_id=tool_call_id):

147
agent/application/states/database.py
View File

@ -27,13 +27,13 @@ from application.states.models import (
) )
class VerificationCodeRecord(SQLModel, table=True, table_name="verification_code"): class CaptchaRecord(SQLModel, table=True, table_name="captcha_code"):
""" """
验证码记录 验证码记录
""" """
email: str = Field(..., primary_key=True, description="邮箱") email: str = Field(..., primary_key=True, description="邮箱")
verification_code: str = Field( code: str = Field(
default_factory=lambda: "".join(choices("0123456789", k=6)), default_factory=lambda: "".join(choices("0123456789", k=6)),
primary_key=True, primary_key=True,
description="验证码", description="验证码",
@ -109,55 +109,69 @@ class MessageRecord(SQLModel, table=True, table_name="message"):
content: str = Field(default="", description="消息内容") content: str = Field(default="", description="消息内容")
class RunRecord(SQLModel, table=True, table_name="run_result"):
"""
运行记录
"""
id: str = Field(
default_factory=lambda: str(uuid7()),
primary_key=True,
description="运行唯一标识",
)
conversation_id: str = Field(..., index=True, description="会话唯一标识")
new_messages: str = Field(
default_factory=list, sa_type=JSON, description="新增消息"
)
class DatabaseState(rx.State): class DatabaseState(rx.State):
""" """
数据库状态 数据库状态
""" """
async def create_verification_code_record(self, email: str) -> str: async def create_captcha_record(self, email: str) -> str:
""" """
创建验证码记录 创建验证记录
:param email: 邮箱 :param email: 邮箱
:return: 验证码 :return: 验证码
""" """
async with rx.asession() as session: async with rx.asession() as session:
# 先将该邮箱有效、未核验的验证码记录设置为无效 # 先将该邮箱有效、未核验的验证码记录设置为无效
await session.exec( await session.exec(
update(VerificationCodeRecord) update(CaptchaRecord)
.where( .where(
VerificationCodeRecord.email == email, # type: ignore CaptchaRecord.email == email, # type: ignore
VerificationCodeRecord.is_valid == True, # type: ignore CaptchaRecord.is_valid == True, # type: ignore
VerificationCodeRecord.is_verified == False, # type: ignore CaptchaRecord.is_verified == False, # type: ignore
) )
.values(is_valid=False) .values(is_valid=False)
) )
await session.flush() await session.flush()
# 创建验证码记录 # 创建验证记录
record = VerificationCodeRecord( record = CaptchaRecord(
email=email, email=email,
) )
session.add(record) session.add(record)
await session.commit() await session.commit()
await session.refresh(record) await session.refresh(record)
return record.verification_code return record.code
async def verify_verification_code( async def verify_captcha_code(self, email: str, captcha_code: str) -> bool:
self, email: str, verification_code: str
) -> bool:
""" """
核验验证码 核验验证码
:param email: 邮箱 :param email: 邮箱
:param verification_code: 验证码 :param code: 验证码
:return: 是否核验成功，True 表示核验成功，False 表示核验失败（根据邮箱和验证码未查询到有效且未核验的记录，或失效） :return: 核验是否成功，True 表示核验成功，False 表示核验失败
""" """
async with rx.asession() as session: async with rx.asession() as session:
result = await session.exec( result = await session.exec(
select(VerificationCodeRecord).where( select(CaptchaRecord).where(
VerificationCodeRecord.email == email, CaptchaRecord.email == email,
VerificationCodeRecord.verification_code == verification_code, CaptchaRecord.code == captcha_code,
VerificationCodeRecord.is_valid == True, CaptchaRecord.is_valid == True,
VerificationCodeRecord.is_verified == False, CaptchaRecord.is_verified == False,
VerificationCodeRecord.expired_at > datetime.now(), CaptchaRecord.expired_at > datetime.now(),
) )
) )
record = result.first() record = result.first()
@ -214,7 +228,8 @@ class DatabaseState(rx.State):
Conversation( Conversation(
id=conversation_record.id, id=conversation_record.id,
description=conversation_record.description, description=conversation_record.description,
usage=usage_validate_python(conversation_record.usage), usage=conversation_record.usage,
messages={},
created_at=conversation_record.created_at, created_at=conversation_record.created_at,
), ),
) )
@ -246,6 +261,8 @@ class DatabaseState(rx.State):
record.id: Conversation( record.id: Conversation(
id=record.id, id=record.id,
description=record.description, description=record.description,
usage=record.usage,
messages={}, # 新会话默认消息字典为空
created_at=record.created_at, created_at=record.created_at,
) )
} }
@ -257,30 +274,30 @@ class DatabaseState(rx.State):
:return: None :return: None
""" """
async with rx.asession() as session: async with rx.asession() as session:
record = await session.get(ConversationTable, conversation_id) record = await session.get(ConversationRecord, conversation_id)
if not record: if not record:
return return
record.is_deleted = True record.is_deleted = True
await session.commit() await session.commit()
async def create_dialog_record( async def save_new_messages_record1(
self, self,
conversation_id: str, conversation_id: str,
id: str, run_id: str,
user_prompt: str, user_prompt: str,
thoughts: dict[int, Any], thoughts: dict[int, Any],
result_output: str, result_output: str,
usage: dict[str, Any], usage: dict[str, Any],
) -> dict[str, Dialog]: ) -> dict[str, Message]:
""" """
创建对话记录 保存新增消息
:param conversation_id: 会话唯一标识 :param conversation_id: 会话唯一标识
:param user_prompt: 用户提示词 :param user_prompt: 用户提示词
:param result_output: 结果输出 :param result_output: 结果输出
:return: 创建对话记录的唯一标识 :return: 保存消息记录的唯一标识
""" """
async with rx.asession() as session: async with rx.asession() as session:
record = DialogRecord( record = MessageRecord(
id=id, id=id,
conversation_id=conversation_id, conversation_id=conversation_id,
user_prompt=user_prompt, user_prompt=user_prompt,
@ -301,48 +318,48 @@ class DatabaseState(rx.State):
) )
} }
async def create_result_record(
self,
conversation_id: str,
dialog_id: str,
new_messages: list[ModelMessage],
) -> None:
"""
创建结果记录
:param conversation_id: 会话唯一标识
:param dialog_id: 对话唯一标识
:param new_messages: 新增消息
:return: None
"""
async with rx.asession() as session:
session.add(
ResultRecord(
conversation_id=conversation_id,
dialog_id=dialog_id,
new_messages=ModelMessagesTypeAdapter.dump_json(
new_messages
).decode(
"utf-8"
), # 序列化为 JSON 字符串
)
)
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: 会话唯一标识 :param conversation_id: 会话唯一标识
:return: 消息历史 :return: 消息历史
""" """
records: list[ModelMessage] = [] message_history: list[ModelMessage] = []
async with rx.asession() as session: async with rx.asession() as session:
result = await session.exec( result = await session.exec(
select(ResultRecord) select(RunRecord)
.where(ResultRecord.conversation_id == conversation_id) .where(RunRecord.conversation_id == conversation_id)
.order_by(desc(ResultRecord.dialog_id)) .order_by(RunRecord.id)
) )
for record in result.all(): for run in result.all():
records.extend( message_history.extend(
ModelMessagesTypeAdapter.validate_json(record.new_messages) ModelMessagesTypeAdapter.validate_python(
run.new_messages
) # 将 run.new_messages 由 Python 类型反序列化为 List[ModelMessage]
) )
return records return message_history
async def save_new_messages_record(
self,
id: str,
conversation_id: str,
new_messages: list[ModelMessage],
) -> None:
"""
保存新增消息
:param id: 运行唯一标识
:param conversation_id: 会话唯一标识
:param new_messages: 新增消息
:return: None
"""
async with rx.asession() as session:
session.add(
RunRecord(
id=id,
conversation_id=conversation_id,
new_messages=ModelMessagesTypeAdapter.dump_python(
new_messages
), # 将 messages 由 List[ModelMessage] 序列化为 Python 类型
)
)
await session.commit()

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

@ -31,8 +31,11 @@ class Message(BaseModel):
type: MessageType = Field(..., description="消息类型") type: MessageType = Field(..., description="消息类型")
title: str = Field(default="", description="消息标题") title: str = Field(default="", description="消息标题")
content: str = Field(default="", description="消息内容") content: str = Field(default="", description="消息内容")
is_thinking: bool = Field(
default=False, description="正在思考，True 表示正在思考，False 表示思考完成"
)
is_expanded: bool = Field( is_expanded: bool = Field(
default=True, description="展开，True 表示展开，False 表示折叠" default=False, description="展开组件，True 表示展开，False 表示折叠"
) )
@ -60,12 +63,11 @@ class Run(BaseModel):
) )
class TaskType(StrEnum): class WorkType(StrEnum):
""" """
任务类型枚举 工作类型枚举
""" """
CHAT = "聊天"
BOOK_FLIGHT = "预定航班" BOOK_FLIGHT = "预定航班"
@ -85,15 +87,15 @@ class Deps(BaseModel):
... ...
class Task(BaseModel): class Work(BaseModel):
""" """
任务类 工作类
""" """
type: TaskType = Field(default=TaskType.CHAT, description="任务类型") type: WorkType = Field(..., description="工作类型")
status: TaskStatus = Field(default=TaskStatus.NONE, description="任务状态") status: TaskStatus = Field(default=TaskStatus.NONE, description="工作状态")
deps: Deps | None = Field(default=None, description="任务依赖项") deps: Deps | None = Field(default=None, description="工作依赖项")
usage: RunUsage = Field(default=RunUsage(), description="任务使用量") usage: RunUsage = Field(default=RunUsage(), description="工作使用量")
usage_limits: UsageLimits | None = Field( usage_limits: UsageLimits | None = Field(
default=None, default=None,
description="任务使用量限制", description="任务使用量限制",
@ -105,12 +107,17 @@ class Conversation(BaseModel):
会话类 会话类
""" """
id: str = Field(default_factory=lambda: str(uuid7()), description="会话唯一标识") id: str = Field(..., description="会话唯一标识")
description: str = Field(default="新会话", description="会话描述") description: str = Field(..., description="会话描述")
usage: RunUsage = Field(default=RunUsage(), description="会话使用量") user_prompt: str = Field(default="", description="用户提示词")
messages: dict[str, Message] = Field(default_factory=dict, description="运行字典") usage: dict[str, Any] = Field(..., description="会话使用量")
messages: dict[str, Message] = Field(..., description="消息字典")
created_at: datetime = Field(..., description="会话创建日期时间") created_at: datetime = Field(..., description="会话创建日期时间")
is_running: bool = Field(
default=False, description="正在运行，True 表示正在运行， False 表示运行结束"
)
class ConversationHistoryItem(BaseModel): class ConversationHistoryItem(BaseModel):
""" """

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

@ -8,7 +8,7 @@ from pydantic_ai import Agent, ModelMessage
from pydantic_ai.messages import AgentStreamEvent from pydantic_ai.messages import AgentStreamEvent
from pydantic_ai.run import AgentRunResultEvent from pydantic_ai.run import AgentRunResultEvent
from application.states.models import TaskNodeResultEvent, TaskType, Task from application.states.models import TaskNodeResultEvent, TaskType, Task, Message
from application.tasks.models import DEEPSEEK_V4_FLASH_MODEL from application.tasks.models import DEEPSEEK_V4_FLASH_MODEL
@ -30,8 +30,8 @@ instruction = """
async def run_stream_events( async def run_stream_events(
task: Task | None, work: Work | None,
user_prompt: str, user_prompt: Message,
message_history: List[ModelMessage], message_history: List[ModelMessage],
) -> AsyncGenerator[AgentStreamEvent | AgentRunResultEvent | TaskNodeResultEvent, None]: ) -> AsyncGenerator[AgentStreamEvent | AgentRunResultEvent | TaskNodeResultEvent, None]:
""" """