This commit is contained in:
liubiren 2026-07-13 09:04:51 +08:00
parent 78175a1ec4
commit 5a68247656
7 changed files with 345 additions and 221 deletions

View File

@ -1,50 +1,50 @@
# -*- coding: utf-8 -*- # -*- coding: utf-8 -*-
""" """
数据模型 数据模型
""" """
from datetime import datetime from datetime import datetime
from typing import Any
from pydantic_ai._uuid import uuid7 from pydantic_ai._uuid import uuid7
from pydantic_ai.messages import ModelMessage, ModelMessagesTypeAdapter from sqlmodel import Field, JSON, SQLModel
from sqlmodel import Field, SQLModel, JSON
# 会话数据表模型 # 会话数据表模型
class Conversations(SQLModel, table=True): class Conversation(SQLModel, table=True):
id: str = Field( id: str = Field(
default_factory=lambda: str(uuid7()), default_factory=lambda: str(uuid7()),
primary_key=True, primary_key=True,
description="会话唯一标识", description="会话唯一标识",
) )
user_id: str = Field(index=True, description="用户唯一标识") user_id: str = Field(..., index=True, description="用户唯一标识")
description: str = Field(description="会话描述") description: str = Field(..., description="会话描述")
is_deleted: bool = Field(default=False, description="是否删除") is_deleted: bool = Field(default=False, description="是否删除")
dialogs: dict[str, dict[str, Any]] = Field(default_factory=dict, sa_type=JSON, description="对话字典")
created_at: datetime = Field(default_factory=datetime.now, description="创建时间") created_at: datetime = Field(default_factory=datetime.now, description="创建时间")
# 对话数据表模型
class Dialog(SQLModel, table=True):
id: str = Field(
default_factory=lambda: str(uuid7()),
primary_key=True,
description="对话唯一标识",
)
conversation_id: str = Field(..., index=True, description="会话唯一标识")
question: str = Field(..., description="问题")
thought_nodes: list[ThoughtNode] = Field(
default_factory=list, sa_type=JSON, description="思考节点列表"
)
answer: str = Field(default="", description="回答")
# 思考节点数据表模型
class ThoughtNode(SQLModel):
kind: str = Field(description="思考节点类型")
content: str = Field(default="", description="思考节点内容")
# 运行结果数据表模型 # 运行结果数据表模型
class RunResults(SQLModel, table=True): class RunResult(SQLModel, table=True):
conversation_id: str = Field(primary_key=True, description="会话唯一标识") conversation_id: str = Field(primary_key=True, description="会话唯一标识")
dialog_id: str = Field(primary_key=True, description="对话唯一标识") dialog_id: str = Field(primary_key=True, description="对话唯一标识")
new_messages: str = Field(description="新增消息") new_messages: str = Field(description="新增消息")
@staticmethod
def adapt(
conversation_id: str, dialog_id: str, new_messages: list[ModelMessage]
) -> "RunResults":
"""
适配为运行结果记录
:param conversation_id: 会话唯一标识
:param dialog_id: 对话唯一标识
:param new_messages: 新增消息
:return: 运行结果记录
"""
return RunResults(
conversation_id=conversation_id,
dialog_id=dialog_id,
new_messages=ModelMessagesTypeAdapter.dump_json(new_messages).decode(
"utf-8"
), # 序列化为 JSON 字符串
)

View File

@ -41,7 +41,7 @@ def conversation() -> rx.Component:
# 渲染列表 # 渲染列表
rx.auto_scroll( rx.auto_scroll(
rx.foreach( rx.foreach(
ConversationState.get_conversation_history, ConversationState.conversation_history,
lambda item, _: render_conversation_history_item( lambda item, _: render_conversation_history_item(
item[0], item[1] item[0], item[1]
), ),

View File

@ -1,20 +1,27 @@
# -*- coding: utf-8 -*- # -*- coding: utf-8 -*-
from application.states.auth import AuthState from application.states.auth import AuthState
from application.states.conversation import ConversationState, Conversation, Dialog, ReasoningPart from application.states.conversation import ConversationState
from application.states.create_conversation import CreateConversationState from application.states.create_conversation import CreateConversationState
from application.states.conversation_history import ConversationHistoryState from application.states.conversation_history import ConversationHistoryState
from application.states.database import DatabaseState from application.states.database import (
DatabaseState,
Conversations,
ConversationMemory as Conversation,
DialogMemory as Dialog,
ThoughtNodeMemory as ThoughtNode,
)
from application.states.sidebar import SidebarState from application.states.sidebar import SidebarState
__all__ = [ __all__ = [
"AuthState", "AuthState",
"ConversationState", "ConversationState",
"Conversations",
"ConversationHistoryState", "ConversationHistoryState",
"CreateConversationState", "CreateConversationState",
"DatabaseState", "DatabaseState",
"SidebarState", "SidebarState",
"Conversation", "Conversation",
"Dialog", "Dialog",
"ReasoningPart", "ThoughtNode",
] ]

View File

@ -9,7 +9,8 @@ class AuthState(rx.State):
""" """
认证状态 认证状态
""" """
# 当前登录的用户唯一标识,绑定 LocalStorage
# 当前登录的用户唯一标识(通过本地存储同步)
user_id = rx.LocalStorage("user_id", sync=True) user_id = rx.LocalStorage("user_id", sync=True)
@rx.event @rx.event
@ -19,4 +20,3 @@ class AuthState(rx.State):
""" """
self.user_id = "1" self.user_id = "1"

View File

@ -2,15 +2,10 @@
""" """
智能体状态 智能体状态
""" """
from datetime import datetime
from enum import StrEnum
from typing import Any, AsyncGenerator, Dict from typing import Any, AsyncGenerator, Dict
from typing import Dict, List from typing import Dict, List
from pydantic import BaseModel, Field
from pydantic_ai import Agent, ThinkingPartDelta from pydantic_ai import Agent, ThinkingPartDelta
from pydantic_ai._uuid import uuid7
from pydantic_ai._uuid import uuid7
from pydantic_ai.messages import ( from pydantic_ai.messages import (
FunctionToolCallEvent, FunctionToolCallEvent,
FunctionToolResultEvent, FunctionToolResultEvent,
@ -24,71 +19,20 @@ from pydantic_ai.messages import (
ToolCallPart, ToolCallPart,
ToolSearchCallPart, ToolSearchCallPart,
) )
from pydantic_ai.messages import ModelMessage, ModelMessagesTypeAdapter
from pydantic_ai.models.openai import OpenAIChatModel from pydantic_ai.models.openai import OpenAIChatModel
from pydantic_ai.providers.openai import OpenAIProvider from pydantic_ai.providers.openai import OpenAIProvider
from pydantic_ai.run import AgentRunResultEvent from pydantic_ai.run import AgentRunResultEvent
import reflex as rx import reflex as rx
from sqlmodel import Field as SqlField, SQLModel
from application.states import CreateConversationState, DatabaseState from application.states import (
AuthState,
Conversation,
class ReasoningPartKind(StrEnum): Conversations,
""" CreateConversationState,
推理分片种类 DatabaseState,
""" Dialog,
ThoughtNode,
THINKING = "thinking" )
TOOL_SEARCH = "tool-search"
CAPABILITY_LOAD = "load-capability"
TOOL_CALL = "tool-call"
TEXT = "text"
TOOL_RETURN = "tool-return"
RETRY_PROMPT = "retry-prompt"
RUN_RETURN = "run-return"
class ReasoningPart(BaseModel):
"""
推理分片内存模型
"""
kind: ReasoningPartKind = Field(..., description="推理分片种类")
content: str = Field(default="", description="推理分片内容")
class Dialog(BaseModel):
"""
对话内存模型
"""
user_prompt: str = Field(..., description="用户提示词")
is_reasoning: bool = Field(
default=False,
description="推理状态True 表示正在推理False 表示推理完成",
)
is_collapse_open: bool = Field(
default=False,
description="推理面板展开状态True 表示推理面板展开False 表示推理面板折叠",
)
reasoning_parts: Dict[int, ReasoningPart] = Field(
default_factory=dict, description="推理分片字典"
)
assistant_content: str = Field(default="", description="回复正文")
is_running: bool = Field(
default=False,
description="运行状态True 表示正在运行False 表示运行完成",
)
class Conversation(BaseModel):
"""
会话内存模型
"""
description: str = Field(default="新会话", description="会话描述")
dialogs: Dict[str, Dialog] = Field(default_factory=dict, description="对话字典")
instructions: str = """ instructions: str = """
@ -127,57 +71,98 @@ class ConversationState(rx.State):
会话状态 会话状态
""" """
# 初始化会话字典 # 用户唯一标识
conversations: Dict[str, Conversation] = {str(uuid7()): Conversation()} user_id: str = ""
# 获取当前会话唯一标识 # 会话字典
conversation_id: str = next(reversed(conversations.keys())) conversations: Conversations = {}
# 当前会话唯一标识
conversation_id: str = ""
@rx.event
async def on_mount(self):
"""
页面挂载时初始化会话列表和当前会话唯一标识
"""
# 获取登录用户唯一标识
auth_state = await self.get_state(AuthState)
self.user_id = auth_state.user_id
if not self.user_id:
return
# 获取数据库状态
database_state = await self.get_state(DatabaseState)
# 获取会话字典
self.conversations = await database_state.get_conversations(
user_id=self.user_id
)
# 若会话字典为空则新建会话
if not self.conversations:
# 先在数据库新建会话,再在会话字典新建会话
conversation_id, created_at = await database_state.create_conversation(
user_id=self.user_id
)
self.conversations[conversation_id] = Conversation(created_at=created_at)
# 默认将最后一个会话的唯一标识设置为当前会话唯一标识
self.conversation_id = next(reversed(self.conversations.keys()))
@rx.var @rx.var
def get_chats(self) -> Dict[str, Conversation]: def conversation_history(self) -> Conversations:
""" """
获取会话字典用于前端渲染会话历史 获取会话历史用于前端渲染会话历史
:return: 会话字典 :return: 会话历史
""" """
return dict(reversed(self.conversations.items())) return dict(reversed(self.conversations.items()))
@rx.event @rx.event
async def create_chat(self, form_data: Dict[str, Any]) -> None: async def create_conversation(self, form_data: Dict[str, Any]) -> None:
""" """
新建会话 新建会话
:param form_data: 表单数据 :param form_data: 表单数据
:return: None :return: None
""" """
# 获取会话描述 # 获取会话描述
description = form_data["description"].strip() description = form_data["description"].strip() or "新会话"
if not description:
description = "新会话"
# 新建会话 # 先在数据库新建会话,再在会话字典新建会话
self.conversations[str(uuid7())] = Conversation(description=description) database_state = await self.get_state(DatabaseState)
# 将末位会话唯一标识设置为当前会话唯一标识 conversation_id, created_at = await database_state.create_conversation(
self.conversation_id = next(reversed(self.conversations.keys())) user_id=self.user_id, description=description
)
self.conversations[conversation_id] = Conversation(
description=description, created_at=created_at
)
self.conversation_id = conversation_id
# 获取新建会话状态
create_chat_state = await self.get_state(CreateConversationState)
# 关闭新建会话模态窗 # 关闭新建会话模态窗
create_chat_state.is_open = False create_conversation_state = await self.get_state(CreateConversationState)
create_conversation_state.is_modal_open = False
@rx.event @rx.event
def delete_conversation(self, conversation_id: str) -> None: async def delete_conversation(self, conversation_id: str) -> None:
""" """
删除会话 删除指定会话
:param conversation_id: 会话唯一标识 :param conversation_id: 指定会话唯一标识
:return: None :return: None
""" """
if conversation_id not in self.conversations: if conversation_id not in self.conversations:
return return
# 先在数据库删除会话,再在会话字典删除会话
database_state = await self.get_state(DatabaseState)
await database_state.delete_conversation(conversation_id=conversation_id)
del self.conversations[conversation_id] del self.conversations[conversation_id]
# 删除后,若会话字典为空则创建会话 # 删除后,若会话字典为空则先在数据库新建会话,再在会话字典新建会话
if not self.conversations: if not self.conversations:
self.conversations[str(uuid7())] = Conversation() conversation_id, created_at = await database_state.create_conversation(
user_id=self.user_id
)
self.conversations[conversation_id] = Conversation(created_at=created_at)
# 删除后,若当前会话唯一标识不存在则将末位会话唯一标识设置为当前会话唯一标识 # 删除后,若当前会话唯一标识不存在则将最后一个会话的唯一标识设置为当前会话唯一标识
if self.conversation_id not in self.conversations: if self.conversation_id not in self.conversations:
self.conversation_id = next(reversed(self.conversations.keys())) self.conversation_id = next(reversed(self.conversations.keys()))
@ -193,39 +178,26 @@ class ConversationState(rx.State):
self.conversation_id = conversation_id self.conversation_id = conversation_id
@rx.var @rx.var
def get_conversation_description(self) -> str: def dialog_history(self) -> Dict[str, Dialog]:
""" """
获取当前会话描述用于前端渲染会话标题 获取当前会话的对话历史用于前端渲染对话历史
:return: 当前会话描述 :return: 当前会话的对话历史
"""
# 当前会话
conversation = self.conversations.get(self.conversation_id)
return conversation.description if conversation else "新会话"
@rx.var
def get_dialogs(self) -> Dict[str, Dialog]:
"""
获取对话字典用于前端渲染对话历史
:return: 对话字典
""" """
# 当前会话 # 当前会话
conversation = self.conversations.get(self.conversation_id) conversation = self.conversations.get(self.conversation_id)
return conversation.dialogs if conversation else {} return conversation.dialogs if conversation else {}
@rx.var @rx.var
def get_running_status(self) -> bool: def running_status(self) -> bool:
""" """
获取当前运行状态 获取当前会话的运行状态
:return: 当前运行状态True 表示正在运行False 表示运行完成 :return: 当前会话的运行状态True 表示正在运行False 表示运行完成
""" """
# 当前会话 # 当前会话
conversation = self.conversations.get(self.conversation_id) conversation = self.conversations.get(self.conversation_id)
if not conversation or not conversation.dialogs: if not conversation:
return False return False
return conversation.is_running
# 末位对话
dialog = next(reversed(conversation.dialogs.values()))
return dialog.is_running
@rx.event @rx.event
async def run(self, form_data: dict[str, Any]) -> AsyncGenerator[None]: async def run(self, form_data: dict[str, Any]) -> AsyncGenerator[None]:
@ -234,37 +206,41 @@ class ConversationState(rx.State):
:param form_data: 表单数据 :param form_data: 表单数据
:return: AsyncGenerator :return: AsyncGenerator
""" """
# 获取问题
question = form_data.get("question", "").strip()
if not question:
return
# 当前会话 # 当前会话
conversation = self.conversations.get(self.conversation_id) conversation = self.conversations.get(self.conversation_id)
if not conversation: if not conversation:
return return
# 获取用户提示词 # 将当前会话的运行状态设置为正在运行
user_prompt = form_data.get("user_prompt", "").strip() conversation.is_running = True
if not user_prompt:
return
# 获取数据库状态 # 获取数据库状态
database_state = await self.get_state(DatabaseState) database_state = await self.get_state(DatabaseState)
# 获取消息历史 # 获取消息历史
message_history = await database_state.get_message_history( message_history = await database_state.get_message_history(
conversation_id=self.conversation_id conversation_id=self.conversation_id
) )
# 先在数据库新建对话,再在当前会话的对话字典新增对话
dialog_id = await database_state.create_dialog(
conversation_id=self.conversation_id, question=question
)
dialog = conversation.dialogs.setdefault(dialog_id, Dialog(question=question))
# 初始化工具调用唯一标识和片段索引映射字典 # 初始化工具调用唯一标识和片段索引映射字典
tool_call_ids: Dict[str, int] = {} tool_call_ids: Dict[str, int] = {}
# 创建对话
conversation.dialogs[str(uuid7())] = Dialog(user_prompt=user_prompt)
# 将末位对话唯一标识、对话设置为当前对话唯一标识、对话
dialog_id, dialog = next(reversed(conversation.dialogs.items()))
# 将运行状态设置为正在运行
dialog.is_running = True
yield # 通知前端渲染 yield # 通知前端渲染
async with agent.run_stream_events( async with agent.run_stream_events(
conversation_id=self.conversation_id, conversation_id=self.conversation_id,
user_prompt=user_prompt, user_prompt=question,
message_history=message_history, message_history=message_history,
) as events: ) as events:
async for event in events: async for event in events:
@ -278,13 +254,13 @@ class ConversationState(rx.State):
match part: match part:
# 思考分片开始事件 # 思考分片开始事件
case ThinkingPart(content=content): case ThinkingPart(content=content):
# 若上一分片种类为空则将推理状态设置为正在推理、推理面板展开状态设置为展开 # 若上一分片种类为空则将思考状态设置为正在思考、思考折叠面板展开状态设置为展开
if not previous_part_kind: if not previous_part_kind:
dialog.is_reasoning = True dialog.is_thinking = True
dialog.is_collapse_open = True dialog.is_collapse_expanded = True
dialog.reasoning_parts[index] = ReasoningPart( dialog.thought_nodes[index] = ThoughtNode(
kind=ReasoningPartKind.THINKING, content=content kind="thinking", content=content
) )
yield yield
@ -293,8 +269,8 @@ class ConversationState(rx.State):
# 创建工具调用唯一标识与片段索引映射 # 创建工具调用唯一标识与片段索引映射
tool_call_ids[tool_call_id] = index tool_call_ids[tool_call_id] = index
dialog.reasoning_parts[index] = ReasoningPart( dialog.thought_nodes[index] = ThoughtNode(
kind=ReasoningPartKind.TOOL_SEARCH, kind="tool_search",
content="正在生成检索关键词", content="正在生成检索关键词",
) )
yield yield
@ -303,8 +279,8 @@ class ConversationState(rx.State):
case LoadCapabilityCallPart(tool_call_id=tool_call_id): case LoadCapabilityCallPart(tool_call_id=tool_call_id):
tool_call_ids[tool_call_id] = index tool_call_ids[tool_call_id] = index
dialog.reasoning_parts[index] = ReasoningPart( dialog.thought_nodes[index] = ThoughtNode(
kind=ReasoningPartKind.CAPABILITY_LOAD, kind="capability_load",
content="正在生成加载参数", content="正在生成加载参数",
) )
yield yield
@ -313,15 +289,15 @@ class ConversationState(rx.State):
case ToolCallPart(tool_call_id=tool_call_id): case ToolCallPart(tool_call_id=tool_call_id):
tool_call_ids[tool_call_id] = index tool_call_ids[tool_call_id] = index
dialog.reasoning_parts[index] = ReasoningPart( dialog.thought_nodes[index] = ThoughtNode(
kind=ReasoningPartKind.TOOL_CALL, kind="tool_call",
content="正在生成调用参数", content="正在生成调用参数",
) )
yield yield
# 文本分片开始事件 # 文本分片开始事件
case TextPart(content=content): case TextPart(content=content):
dialog.assistant_content = content dialog.answer = content
yield yield
# ========== 增量事件 ========== # ========== 增量事件 ==========
@ -330,15 +306,17 @@ class ConversationState(rx.State):
# 思考分片增量事件 # 思考分片增量事件
case ThinkingPartDelta( case ThinkingPartDelta(
content_delta=content_delta, content_delta=content_delta,
): ):
dialog.reasoning_parts[index].content += content_delta or "" dialog.thought_nodes[index].content += (
content_delta or ""
)
yield yield
# 文本分片增量事件 # 文本分片增量事件
case TextPartDelta( case TextPartDelta(
content_delta=content_delta, content_delta=content_delta,
): ):
dialog.assistant_content += content_delta dialog.answer += content_delta
yield yield
# ========== 结束事件 ========== # ========== 结束事件 ==========
@ -350,32 +328,32 @@ class ConversationState(rx.State):
match part: match part:
# 思考分片结束事件 # 思考分片结束事件
case ThinkingPart(part_kind=part_kind, content=content): case ThinkingPart(part_kind=part_kind, content=content):
# 若下一分片种类为文本则将推理状态设置为推理完成、推理面板展开状态设置为折叠 # 若下一分片种类为文本则将思考状态设置为思考完成、思考面板展开状态设置为折叠
if next_part_kind == ReasoningPartKind.TEXT: if next_part_kind == "text":
dialog.is_reasoning = False dialog.is_thinking = False
dialog.is_collapse_open = False dialog.is_collapse_expanded = False
yield yield
# ========== 函数工具调用事件 ========== # ========== 函数工具调用事件 ==========
case FunctionToolCallEvent(tool_call_id=tool_call_id, part=part): case FunctionToolCallEvent(tool_call_id=tool_call_id, part=part):
# 获取分片索引 # 获取分片索引
index = tool_call_ids[tool_call_id] index = tool_call_ids[tool_call_id]
match dialog.reasoning_parts[index].kind: match dialog.thought_nodes[index].kind:
# 工具检索 # 工具检索
case ReasoningPartKind.TOOL_SEARCH: case "tool_search":
dialog.reasoning_parts[index].content = "正在检索" dialog.thought_nodes[index].content = "正在检索"
yield yield
# 能力加载 # 能力加载
case ReasoningPartKind.CAPABILITY_LOAD: case "capability_load":
dialog.reasoning_parts[index].content = ( dialog.thought_nodes[index].content = (
f"正在加载能力 {part.tool_name}" f"正在加载能力 {part.tool_name}"
) )
yield yield
# 工具调用 # 工具调用
case ReasoningPartKind.TOOL_CALL: case "tool_call":
dialog.reasoning_parts[index].content = ( dialog.thought_nodes[index].content = (
f"正在调用工具 {part.tool_name}" f"正在调用工具 {part.tool_name}"
) )
yield yield
@ -386,22 +364,22 @@ class ConversationState(rx.State):
content=content, content=content,
): ):
index = tool_call_ids[tool_call_id] index = tool_call_ids[tool_call_id]
match dialog.reasoning_parts[index].kind: match dialog.thought_nodes[index].kind:
# 工具检索 # 工具检索
case ReasoningPartKind.TOOL_SEARCH: case "tool_search":
dialog.reasoning_parts[index].content = ( dialog.thought_nodes[index].content = (
content if isinstance(content, str) else "" content if isinstance(content, str) else ""
) # 暂仅考虑文本内容 ) # 暂仅考虑文本内容
yield yield
# 能力加载 # 能力加载
case ReasoningPartKind.CAPABILITY_LOAD: case "capability_load":
dialog.reasoning_parts[index].content = "已加载" dialog.thought_nodes[index].content = "已加载"
yield yield
# 工具调用 # 工具调用
case ReasoningPartKind.TOOL_CALL: case "tool_call":
dialog.reasoning_parts[index].content = f"已调用" dialog.thought_nodes[index].content = f"已调用"
yield yield
# ========== 智能体运行结果事件 ========== # ========== 智能体运行结果事件 ==========
@ -412,8 +390,8 @@ class ConversationState(rx.State):
dialog_id=dialog_id, dialog_id=dialog_id,
new_messages=result.new_messages(), new_messages=result.new_messages(),
) )
# 将运行状态设置为运行完成 # 将当前会话的运行状态设置为运行完成
dialog.is_running = False conversation.is_running = False
yield yield
@rx.event @rx.event

View File

@ -3,14 +3,65 @@
数据库状态 数据库状态
""" """
from typing import List, Dict from datetime import datetime
from typing import Dict, List, cast, TypeAlias
from pydantic import BaseModel, Field, TypeAdapter
from pydantic_ai import ModelMessage, ModelMessagesTypeAdapter
import reflex as rx import reflex as rx
from sqlalchemy import desc, and_ from sqlalchemy import ColumnElement, desc
from sqlmodel import select from sqlmodel import select
from application.models import Conversations, RunResults from application.models import Conversation, Dialog, RunResult
from pydantic_ai import ModelMessage, ModelMessagesTypeAdapter
class ConversationMemory(BaseModel):
"""
会话内存模型
"""
description: str = Field(default="新会话", description="会话描述")
is_running: bool = Field(
default=False, description="运行状态True 表示运行中False 表示运行完成"
)
dialogs: Dict[str, DialogMemory] = Field(
default_factory=dict, description="对话字典"
)
created_at: datetime = Field(default_factory=datetime.now, description="创建时间")
class DialogMemory(BaseModel):
"""
对话内存模型
"""
question: str = Field(..., description="问题")
is_thinking: bool = Field(
default=False, description="思考状态True 表示思考中False 表示思考完成"
)
is_collapse_expanded: bool = Field(
default=False, description="思考折叠面板展开状态True 表示展开False 表示折叠"
)
thought_nodes: list[ThoughtNodeMemory] = Field(
default_factory=list, description="思考节点列表"
)
answer: str = Field(default="", description="回答")
class ThoughtNodeMemory(BaseModel):
"""
思考节点内存模型
"""
kind: str
content: str
# 会话字典类型别名
Conversations: TypeAlias = Dict[str, ConversationMemory]
# 思考节点内存模型类型适配器
ThoughtNodeMemoryTypeAdapter = TypeAdapter(list[ThoughtNodeMemory])
class DatabaseState(rx.State): class DatabaseState(rx.State):
@ -18,37 +69,119 @@ class DatabaseState(rx.State):
数据库状态 数据库状态
""" """
async def get_conversations(self, user_id: str) -> Dict[str, Conversation]: async def get_conversations(self, user_id: str) -> Conversations:
""" """
获取指定用户的会话字典 获取会话字典
:param user_id: 用户唯一标识 :param user_id: 用户唯一标识
:return: 指定用户的会话字典 :return: 会话字典
""" """
conversations: Dict[str, Conversation] = {} conversations: Conversations = {}
async with rx.asession() as session: async with rx.asession() as session:
result = await session.exec( result = await session.exec(
select(Conversations, Dialogs, Parts) select(Conversation, Dialog)
.outerjoin(Dialogs, Conversations.id == Dialogs.conversation_id) # type: ignore .outerjoin(
.outerjoin(Parts, Dialogs.id == Parts.dialog_id) # type: ignore Dialog,
cast(
ColumnElement[bool],
Conversation.id == Dialog.conversation_id,
),
)
.where( .where(
models.Conversation.user_id == user_id, Conversation.user_id == user_id,
models.Conversation.is_deleted == False, Conversation.is_deleted == False,
)
.order_by(
models.Conversation.id, models.Dialog.id, models.ReasoningPart.id
) )
.order_by(Conversation.id, Dialog.id)
) )
records = result.all() for conversation, dialog in result.all():
for db_conversation, _, _ in records: conversation_memory = conversations.setdefault(
conversations[conversation.id] = models.Conversation( conversation.id,
user_id=user_id, description=conversation.description ConversationMemory(
description=conversation.description,
created_at=conversation.created_at,
),
) )
if not dialog:
continue
conversation_memory.dialogs[dialog.id] = DialogMemory(
question=dialog.question,
thought_nodes=ThoughtNodeMemoryTypeAdapter.validate_python(
dialog.thought_nodes
), # 反序列化为思考节点内存实例列表
answer=dialog.answer,
)
return conversations return conversations
async def get_message_history( async def create_conversation(
self, conversation_id: str self, user_id: str, description: str = "新会话"
) -> List[ModelMessage]: ) -> tuple[str, datetime]:
"""
新建会话
:param user_id: 用户唯一标识
:param description: 会话描述
:return: 新建会话唯一标识创建时间
"""
async with rx.asession() as session:
conversation = Conversation(user_id=user_id, description=description)
session.add(conversation)
await session.commit()
await session.refresh(conversation)
return conversation.id, conversation.created_at
async def delete_conversation(self, conversation_id: str) -> None:
"""
删除指定会话逻辑删除
:param conversation_id: 指定会话唯一标识
:return: None
"""
async with rx.asession() as session:
conversation = await session.get(Conversation, conversation_id)
if not conversation:
return
conversation.is_deleted = True
await session.commit()
async def create_dialog(self, conversation_id: str, question: str) -> str:
"""
新建对话
:param conversation_id: 会话唯一标识
:param question: 提问
:return: 新建对话唯一标识
"""
async with rx.asession() as session:
dialog = Dialog(
conversation_id=conversation_id,
question=question,
)
session.add(dialog)
await session.commit()
await session.refresh(dialog)
return dialog.id
async def complete_dialog(
self,
dialog_id: str,
thought_nodes: list[ThoughtNodeMemory],
answer: str,
) -> None:
"""
完成对话
:param dialog_id: 对话唯一标识
:param thought_nodes: 思考节点列表
:param answer: 回答
:return: None
"""
async with rx.asession() as session:
dialog = await session.get(Dialog, dialog_id)
if not dialog:
return
dialog.thought_nodes = ThoughtNodeMemoryTypeAdapter.dump_python(
thought_nodes
)
dialog.answer = answer
await session.commit()
async def get_message_history(self, conversation_id: str) -> List[ModelMessage]:
""" """
获取消息历史 获取消息历史
:param conversation_id: 会话唯一标识 :param conversation_id: 会话唯一标识
@ -57,9 +190,9 @@ class DatabaseState(rx.State):
message_history: 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(RunResults) select(RunResult)
.where(RunResults.conversation_id == conversation_id) .where(RunResult.conversation_id == conversation_id)
.order_by(desc(RunResults.dialog_id)) .order_by(desc(RunResult.dialog_id))
) )
for record in result.all(): for record in result.all():
message_history.extend( message_history.extend(
@ -81,10 +214,15 @@ class DatabaseState(rx.State):
:return: None :return: None
""" """
async with rx.asession() as session: async with rx.asession() as session:
record = RunResults.adapt( session.add(
conversation_id=conversation_id, RunResult(
dialog_id=dialog_id, conversation_id=conversation_id,
new_messages=new_messages, dialog_id=dialog_id,
new_messages=ModelMessagesTypeAdapter.dump_json(
new_messages
).decode(
"utf-8"
), # 序列化为 JSON 字符串
)
) )
session.add(record)
await session.commit() await session.commit()

View File

@ -17,6 +17,7 @@ class SidebarState(rx.State):
""" """
生成 LocalStorage 生成 LocalStorage
""" """
# 获取认证状态
auth = await self.get_state(AuthState) auth = await self.get_state(AuthState)
# 获取用户唯一标识 # 获取用户唯一标识
user_id = auth.user_id or "default" user_id = auth.user_id or "default"