LangGraph异步化sqlite checkpoint

作者：CSDN博客
安装

1. pip install langgraph-checkpoint-sqlite

复制代码

异步checkpiont初始化：

1. from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
2. conn = aiosqlite.connect(":memory:", check_same_thread=False)
3. memory = AsyncSqliteSaver(conn)

复制代码

如果使用异步流式应对，需要确保llm节点或者相关节点也转成异步化操作

1. asyncdefllm(self, state: AgentState):
2.         llm_msgs = state['messages']if self.systemMessage:
3.             llm_msgs = self.systemMessage + state['messages']print(f'ask llm to handle request msg, msg: {llm_msgs}')try:# 关键修复：await 异步调用并直接获取结果
4.             msg =await self.model.ainvoke(llm_msgs)print(f'msg={msg}')return{'messages':[msg]}# 确保返回的是消息对象而非协程except Exception as e:print(f"Model invocation error: {e}")# 返回错误提示消息（需符合Message类型）from langchain_core.messages import AIMessage
5.             return{'messages':[AIMessage(content=f"Error: {str(e)}")]}asyncdeftake_action_tool(self, state: AgentState):
6.         current_tools: List[ToolCall]= state['messages'][-1].tool_calls
7.         results =[]for t in current_tools:
8.             tool_result =await self.tools[t['name']].ainvoke(t['args'])
9.             results.append(ToolMessage(
10.                 tool_call_id=t['id'],
11.                 content=str(tool_result),
12.                 name=t['name'],))print(f'Back to model')return{'messages': results}

复制代码

最后的完整代码如下：

1. import asyncio
2. from typing import Annotated, List, TypedDict
3. import os
5. import aiosqlite
6. from langchain_community.chat_models import ChatTongyi
7. from langchain_core.language_models import BaseChatModel
8. from langchain_core.messages import AnyMessage, HumanMessage, SystemMessage, ToolMessage, ToolCall
9. from dotenv import load_dotenv
10. from langchain_community.tools.tavily_search import TavilySearchResults
11. from langchain_core.tools import BaseTool
12. from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
13. from langgraph.constants import END, START
14. from langgraph.graph import add_messages, StateGraph
16. conn = aiosqlite.connect(":memory:", check_same_thread=False)
17. load_dotenv(dotenv_path='../keys.env')
18. ts_tool = TavilySearchResults(max_results=2)classAgentState(TypedDict):
19.     messages: Annotated[List[AnyMessage], add_messages]classAgent:def__init__(
20.             self,
21.             model: BaseChatModel,
22.             systemMessage: List[SystemMessage],
23.             tools: List[BaseTool],
24.             memory,):assertall(isinstance(t, BaseTool)for t in tools),'tools must implement BASEcALL'
25.         graph = StateGraph(AgentState)
27.         graph.add_node('llm', self.llm)
28.         graph.add_node('take_action_tool', self.take_action_tool)
29.         graph.add_conditional_edges('llm',
30.             self.exist_action,{True:'take_action_tool',False: END
31.             })
32.         graph.set_entry_point('llm')
33.         graph.add_edge('take_action_tool','llm')
35.         self.app = graph.compile(checkpointer=memory)
37.         self.tools ={t.name: t for t in tools}
38.         self.model = model.bind_tools(tools)
39.         self.systemMessage = systemMessage
41.     defexist_action(self, state: AgentState):
42.         tool_calls = state['messages'][-1].tool_calls
43.         print(f'tool_calls size {len(tool_calls)}')returnlen(tool_calls)>0asyncdefllm(self, state: AgentState):
44.         llm_msgs = state['messages']if self.systemMessage:
45.             llm_msgs = self.systemMessage + state['messages']print(f'ask llm to handle request msg, msg: {llm_msgs}')try:# 关键修复：await 异步调用并直接获取结果
46.             msg =await self.model.ainvoke(llm_msgs)print(f'msg={msg}')return{'messages':[msg]}# 确保返回的是消息对象而非协程except Exception as e:print(f"Model invocation error: {e}")# 返回错误提示消息（需符合Message类型）from langchain_core.messages import AIMessage
47.             return{'messages':[AIMessage(content=f"Error: {str(e)}")]}asyncdeftake_action_tool(self, state: AgentState):
48.         current_tools: List[ToolCall]= state['messages'][-1].tool_calls
49.         results =[]for t in current_tools:
50.             tool_result =await self.tools[t['name']].ainvoke(t['args'])
51.             results.append(ToolMessage(
52.                 tool_call_id=t['id'],
53.                 content=str(tool_result),
54.                 name=t['name'],))print(f'Back to model')return{'messages': results}asyncdefwork():
55.     prompt ="""You are a smart research assistant. Use the search engine to look up information. \
56.     You are allowed to make multiple calls (either together or in sequence). \
57.     Only look up information when you are sure of what you want. \
58.     If you need to look up some information before asking a follow up question, you are allowed to do that!
59.     """
61.     qwen_model = ChatTongyi(
62.         model=os.getenv('model'),
63.         api_key=os.getenv('api_key'),
64.         base_url=os.getenv('base_url'),)# reduce inference cost
65.     memory = AsyncSqliteSaver(conn)
66.     agent = Agent(model=qwen_model, tools=[ts_tool], systemMessage=[SystemMessage(content=prompt)], memory=memory)
68.     messages =[HumanMessage("who is the popular football star in the world?")]
69.     configurable ={"configurable":{"thread_id":"1"}}asyncfor event in agent.app.astream_events({"messages": messages}, configurable, version="v1"):
70.         kind = event["event"]# print(f"kind = {kind}")if kind =="on_chat_model_stream":
71.             content = event["data"]["chunk"].content
72.             if content:# Empty content in the context of OpenAI means# that the model is asking for a tool to be invoked.# So we only print non-empty contentprint(content, end="|")if __name__ =='__main__':
73.     asyncio.run(work())

复制代码

原文地址：https://blog.csdn.net/zhangkai1992/article/details/147075196