Langgraph笔记1 实现一个简单demo

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77

import os

from langchain_tavily import TavilySearch
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import add_messages, StateGraph
from langgraph.prebuilt import ToolNode, tools_condition
from typing_extensions import TypedDict
from typing import Annotated
from langchain_deepseek import ChatDeepSeek

# API 环境变量
os.environ["LANGCHAIN_API_KEY"] = 
os.environ["DEEPSEEK_API_KEY"] = 
os.environ["TAVILY_API_KEY"] = 

# 初始化大模型
llm = ChatDeepSeek(
    model="deepseek-chat",
    temperature=0.1,
    streaming=True,
)

# Langgraph状态
class State(TypedDict):
    messages: Annotated[list, add_messages]

# 构建graph
graph_builder = StateGraph(State)

# 内存保存器
checkpointer = MemorySaver()

# tools 初始化Tavily搜索工具
tool = TavilySearch(max_results=2)
tools = [tool]

# bind tools to llm
llm_with_tools = llm.bind_tools(tools)
def chatbot(state: State):
    return {"messages":[llm_with_tools.invoke(state["messages"])]}

graph_builder.add_node("chatbot", chatbot)

# add tool node
tool_node = ToolNode(tools=[tool])
graph_builder.add_node("tools", tool_node)

graph_builder.add_conditional_edges(
    "chatbot",
    tools_condition,
)
graph_builder.add_edge("tools", "chatbot")
graph_builder.set_entry_point("chatbot")
graph = graph_builder.compile(checkpointer=checkpointer)

print(graph.get_graph().draw_mermaid())
def stream_graph_updates(user_input: str, thread_id: str):
    config = {"configurable": {"thread_id": thread_id, "recursion_limit": 2}}
    initial_state = {"messages":[{"role":"user", "content": user_input}]}
    for event in graph.stream(initial_state, config = config):
        for value in event.values():
            print("Assistant:", value["messages"][-1].content)

# 测试id
thread_id = "user_123"

# 执行逻辑
while True:
    try:
        user_input = input("User: ")
        if user_input.lower() in ["quit", "exit", "q"]:
            print("Goodbye!")
            break
        stream_graph_updates(user_input, thread_id)
    except KeyboardInterrupt:
        print("\nGoodbye!")
        break