feat: Agent 消息保存在SQLite 中

pyproject.toml

@@ -13,6 +13,7 @@ dependencies = [
     "langchain-huggingface>=0.2.0",
     "langchain-openai>=0.3.14",
     "langgraph>=0.3.34",
+    "langgraph-checkpoint-sqlite>=2.0.7",
     "langsmith>=0.3.37",
     "llama-index>=0.12.33",
     "llama-index-readers-file>=0.4.7",

server/routers/chat_router.py

@@ -165,7 +165,7 @@ async def get_agent(current_user: User = Depends(get_required_user)):
     return {"agents": agents}
 
 @chat.post("/agent/{agent_name}")
-def chat_agent(agent_name: str,
+async def chat_agent(agent_name: str,
                query: str = Body(...),
                config: dict = Body({}),
                meta: dict = Body({}),
@@ -189,7 +189,7 @@ def chat_agent(agent_name: str,
             **kwargs
         }, ensure_ascii=False).encode('utf-8') + b"\n"
 
-    def stream_messages():
+    async def stream_messages():
 
         # 代表服务端已经收到了请求
         yield make_chunk(status="init", meta=meta, msg=HumanMessage(content=query).model_dump())
@@ -212,7 +212,7 @@ def chat_agent(agent_name: str,
         runnable_config = {"configurable": {**config}}
 
         try:
-            for msg, metadata in agent.stream_messages(messages, config_schema=runnable_config):
+            async for msg, metadata in agent.stream_messages(messages, config_schema=runnable_config):
                 logger.debug(f"msg: {msg.model_dump()}, metadata: {metadata}")
                 if isinstance(msg, AIMessageChunk):
                     yield make_chunk(content=msg.content,
@@ -289,7 +289,7 @@ async def get_agent_history(
             raise HTTPException(status_code=404, detail=f"智能体 {agent_name} 不存在")
 
         # 获取历史消息
-        history = agent.get_history(user_id=current_user.id, thread_id=thread_id)
+        history = await agent.get_history(user_id=current_user.id, thread_id=thread_id)
         return {"history": history}
 
     except Exception as e:

src/agents/chatbot/graph.py

@@ -1,16 +1,16 @@
 import os
 import uuid
 from typing import Any
+from pathlib import Path
 from datetime import datetime, timezone
 
+import sqlite3
 from langchain_core.runnables import RunnableConfig
 from langgraph.graph import StateGraph, START, END
 from langgraph.prebuilt import ToolNode, tools_condition
-from langgraph.checkpoint.memory import InMemorySaver
+from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver, aiosqlite
-
-# from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver, aiosqlite
-
 
+from src import config as sys_config
 from src.utils import logger
 from src.agents.registry import State, BaseAgent
 from src.agents.utils import load_chat_model, get_cur_time_with_utc
@@ -26,6 +26,8 @@ class ChatbotAgent(BaseAgent):
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
         self.graph = None
+        self.workdir = Path(sys_config.save_dir) / "agents" / self.name
+        self.workdir.mkdir(parents=True, exist_ok=True)
 
     def _get_tools(self, tools: list[str]):
         """根据配置获取工具。
@@ -56,7 +58,7 @@ class ChatbotAgent(BaseAgent):
         )
         return {"messages": [res]}
 
-    def get_graph(self, config_schema: RunnableConfig = None, **kwargs):
+    async def get_graph(self, config_schema: RunnableConfig = None, **kwargs):
         """构建图"""
         if self.graph:
             return self.graph
@@ -73,18 +75,19 @@ class ChatbotAgent(BaseAgent):
         workflow.add_edge("tools", "chatbot")
         workflow.add_edge("chatbot", END)
 
-        mem_checkpointer = InMemorySaver()
+        # 创建数据库连接
-        graph = workflow.compile(checkpointer=mem_checkpointer)
+        sqlite_checkpointer = AsyncSqliteSaver(await self.get_async_conn())
+        graph = workflow.compile(checkpointer=sqlite_checkpointer)
         self.graph = graph
         return graph
 
-    # async def get_async_conn(self) -> aiosqlite.Connection:
+    async def get_async_conn(self) -> aiosqlite.Connection:
-    #     """获取异步数据库连接"""
+        """获取异步数据库连接"""
-    #     return await aiosqlite.connect(os.path.join(self.db_dir, "aio_history.db"))
+        return await aiosqlite.connect(os.path.join(self.workdir, "aio_history.db"))
 
-    # async def get_aio_memory(self) -> AsyncSqliteSaver:
+    async def get_aio_memory(self) -> AsyncSqliteSaver:
-    #     """获取异步存储实例"""
+        """获取异步存储实例"""
-    #     return AsyncSqliteSaver(await self.get_async_conn())
+        return AsyncSqliteSaver(await self.get_async_conn())
 
 def main():
     agent = ChatbotAgent(ChatbotConfiguration())

src/agents/react/graph.py

@@ -30,7 +30,7 @@ def main():
     config = {"configurable": {"thread_id": thread_id}}
 
     from src.agents.utils import agent_cli
-    agent_cli(agent, config)
+    asyncio.run(agent_cli(agent, config))
 
 
 if __name__ == "__main__":

src/agents/registry.py

@@ -171,28 +171,27 @@ class BaseAgent():
             if requirement not in os.environ:
                 raise ValueError(f"没有配置{requirement} 环境变量，请在 src/.env 文件中配置，并重新启动服务")
 
-    def stream_values(self, messages: list[str], config_schema: RunnableConfig = None, **kwargs):
+    async def stream_values(self, messages: list[str], config_schema: RunnableConfig = None, **kwargs):
-        graph = self.get_graph(config_schema=config_schema, **kwargs)
+        graph = await self.get_graph(config_schema=config_schema, **kwargs)
         logger.debug(f"stream_values: {config_schema}")
-        for event in graph.stream({"messages": messages}, stream_mode="values", config=config_schema):
+        for event in graph.astream({"messages": messages}, stream_mode="values", config=config_schema):
             yield event["messages"]
 
-    def stream_messages(self, messages: list[str], config_schema: RunnableConfig = None, **kwargs):
+    async def stream_messages(self, messages: list[str], config_schema: RunnableConfig = None, **kwargs):
-        graph = self.get_graph(config_schema=config_schema, **kwargs)
+        graph = await self.get_graph(config_schema=config_schema, **kwargs)
         logger.debug(f"stream_messages: {config_schema}")
 
-        for msg, metadata in graph.stream({"messages": messages}, stream_mode="messages", config=config_schema):
+        async for msg, metadata in graph.astream({"messages": messages}, stream_mode="messages", config=config_schema):
             yield msg, metadata
 
-    def get_history(self, user_id, thread_id) -> list[dict]:
+    async def get_history(self, user_id, thread_id) -> list[dict]:
         """获取历史消息"""
         # 获取LangGraph应用实例
-        app = self.get_graph()
+        app = await self.get_graph()
         # 构建配置信息
         config = {"configurable": {"thread_id": thread_id, "user_id": user_id}}
         # 获取状态
-        state = app.get_state(config)
+        state = await app.aget_state(config)
-        # logger.debug(f"获取历史消息: {state}")
 
         result = []
         if state:

src/agents/utils.py

@@ -1,4 +1,5 @@
 from datetime import datetime, timezone
+import asyncio
 
 from src.models import select_model
 from src.agents.registry import BaseAgent
@@ -22,7 +23,7 @@ def load_chat_model(fully_specified_name: str, **kwargs) -> BaseChatModel:
     return model_instance.chat_open_ai
 
 
-def agent_cli(agent: BaseAgent, config: RunnableConfig = None):
+async def agent_cli(agent: BaseAgent, config: RunnableConfig = None):
     config = config or {}
     if "configurable" not in config:
         config["configurable"] = {}
@@ -34,7 +35,7 @@ def agent_cli(agent: BaseAgent, config: RunnableConfig = None):
             break
 
         stream_flag = False
-        for msg, metadata in agent.stream_messages([{"role": "user", "content": user_input}], config):
+        async for msg, metadata in agent.stream_messages([{"role": "user", "content": user_input}], config):
             if isinstance(msg, AIMessageChunk):
                 content = msg.content or msg.tool_calls
 

uv.lock

@@ -138,6 +138,18 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/ec/6a/bc7e17a3e87a2985d3e8f4da4cd0f481060eb78fb08596c42be62c90a4d9/aiosignal-1.3.2-py2.py3-none-any.whl", hash = "sha256:45cde58e409a301715980c2b01d0c28bdde3770d8290b5eb2173759d9acb31a5", size = 7597, upload-time = "2024-12-13T17:10:38.469Z" },
 ]
 
+[[package]]
+name = "aiosqlite"
+version = "0.21.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "typing-extensions" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/13/7d/8bca2bf9a247c2c5dfeec1d7a5f40db6518f88d314b8bca9da29670d2671/aiosqlite-0.21.0.tar.gz", hash = "sha256:131bb8056daa3bc875608c631c678cda73922a2d4ba8aec373b19f18c17e7aa3", size = 13454, upload-time = "2025-02-03T07:30:16.235Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/f5/10/6c25ed6de94c49f88a91fa5018cb4c0f3625f31d5be9f771ebe5cc7cd506/aiosqlite-0.21.0-py3-none-any.whl", hash = "sha256:2549cf4057f95f53dcba16f2b64e8e2791d7e1adedb13197dd8ed77bb226d7d0", size = 15792, upload-time = "2025-02-03T07:30:13.6Z" },
+]
+
 [[package]]
 name = "albucore"
 version = "0.0.23"
@@ -1350,6 +1362,19 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/12/52/bceb5b5348c7a60ef0625ab0a0a0a9ff5d78f0e12aed8cc55c49d5e8a8c9/langgraph_checkpoint-2.0.25-py3-none-any.whl", hash = "sha256:23416a0f5bc9dd712ac10918fc13e8c9c4530c419d2985a441df71a38fc81602", size = 42312, upload-time = "2025-04-26T21:00:42.242Z" },
 ]
 
+[[package]]
+name = "langgraph-checkpoint-sqlite"
+version = "2.0.7"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "aiosqlite" },
+    { name = "langgraph-checkpoint" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/9c/89/125b80e41ddeb8476654a8cda4b76ee999e4fb3d5913ff9c2b45fb9bfff7/langgraph_checkpoint_sqlite-2.0.7.tar.gz", hash = "sha256:344f307c0840a1cbd85a18dcd6daac8e989947979c1a43c2bdc6c6f4ed12084a", size = 9584, upload-time = "2025-05-02T05:44:24.853Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/25/72/86354caec3bd546ea596422dcaf8495a052287400a9961c2b45264536d1e/langgraph_checkpoint_sqlite-2.0.7-py3-none-any.whl", hash = "sha256:b04decd8c3f7c2966ca63b4fa11eb789a03b27001e4d855ccd132c50da59812b", size = 12958, upload-time = "2025-05-02T05:44:23.665Z" },
+]
+
 [[package]]
 name = "langgraph-prebuilt"
 version = "0.1.8"
@@ -4498,6 +4523,7 @@ dependencies = [
     { name = "langchain-huggingface" },
     { name = "langchain-openai" },
     { name = "langgraph" },
+    { name = "langgraph-checkpoint-sqlite" },
     { name = "langsmith" },
     { name = "llama-index" },
     { name = "llama-index-readers-file" },
@@ -4531,6 +4557,7 @@ requires-dist = [
     { name = "langchain-huggingface", specifier = ">=0.2.0" },
     { name = "langchain-openai", specifier = ">=0.3.14" },
     { name = "langgraph", specifier = ">=0.3.34" },
+    { name = "langgraph-checkpoint-sqlite", specifier = ">=2.0.7" },
     { name = "langsmith", specifier = ">=0.3.37" },
     { name = "llama-index", specifier = ">=0.12.33" },
     { name = "llama-index-readers-file", specifier = ">=0.4.7" },