使用时间旅行功能

%%capture --no-stderr
%pip install --quiet -U langgraph langchain_anthropic

import getpass
import os


def _set_env(var: str):
    if not os.environ.get(var):
        os.environ[var] = getpass.getpass(f"{var}: ")


_set_env("ANTHROPIC_API_KEY")

import uuid

from typing_extensions import TypedDict, NotRequired
from langgraph.graph import StateGraph, START, END
from langchain.chat_models import init_chat_model
from langgraph.checkpoint.memory import InMemorySaver


class State(TypedDict):
    topic: NotRequired[str]
    joke: NotRequired[str]


llm = init_chat_model(
    "anthropic:claude-3-7-sonnet-latest",
    temperature=0,
)


def generate_topic(state: State):
    """调用 LLM 生成笑话主题"""
    msg = llm.invoke("给我一个有趣的笑话主题")
    return {"topic": msg.content}


def write_joke(state: State):
    """调用 LLM 根据主题编写笑话"""
    msg = llm.invoke(f"写一个关于 {state['topic']} 的短笑话")
    return {"joke": msg.content}


# 构建工作流
workflow = StateGraph(State)

# 添加节点
workflow.add_node("generate_topic", generate_topic)
workflow.add_node("write_joke", write_joke)

# 添加边以连接节点
workflow.add_edge(START, "generate_topic")
workflow.add_edge("generate_topic", "write_joke")
workflow.add_edge("write_joke", END)

# 编译
checkpointer = InMemorySaver()
graph = workflow.compile(checkpointer=checkpointer)
graph

config = {
    "configurable": {
        "thread_id": uuid.uuid4(),
    }
}
state = graph.invoke({}, config)

print(state["topic"])
print()
print(state["joke"])

“洗衣机里袜子的秘密生活”怎么样？你知道的，探讨一下这个神秘现象：为什么袜子成对进入洗衣机，却总是单只出来？它们去哪儿了？是不是在其他地方开始了新生活？是否存在我们不知道的“袜子天堂”？这个日常谜团蕴含着丰富的喜剧潜力，能引起所有人的共鸣！

# 洗衣机里袜子的秘密生活

我终于发现了所有失踪袜子的去向。原来它们根本没有失踪——只是和洗衣店别人的袜子私奔了，一起开始新生活。

我的蓝色菱形花纹袜现在正和一只红色波点袜在百慕大度假，还在“Sockstagram”上晒照片，并寄给我一些绒毛作为赡养费。

# 状态按逆时间顺序返回。
states = list(graph.get_state_history(config))

for state in states:
    print(state.next)
    print(state.config["configurable"]["checkpoint_id"])
    print()

()
1f02ac4a-ec9f-6524-8002-8f7b0bbeed0e

('write_joke',)
1f02ac4a-ce2a-6494-8001-cb2e2d651227

('generate_topic',)
1f02ac4a-a4e0-630d-8000-b73c254ba748

('__start__',)
1f02ac4a-a4dd-665e-bfff-e6c8c44315d9

# 这是倒数第二个状态（状态按时间顺序列出）
selected_state = states[1]
print(selected_state.next)
print(selected_state.values)

('write_joke',)
{'topic': '“洗衣机里袜子的秘密生活”怎么样？你知道的，探讨一下这个神秘现象：为什么袜子成对进入洗衣机，却总是单只出来？它们去哪儿了？是不是在其他地方开始了新生活？是否存在我们不知道的“袜子天堂”？这个日常谜团蕴含着丰富的喜剧潜力，能引起所有人的共鸣！'}

new_config = graph.update_state(selected_state.config, values={"topic": "chickens"})
print(new_config)

{'configurable': {'thread_id': 'c62e2e03-c27b-4cb6-8cea-ea9bfedae006', 'checkpoint_ns': '', 'checkpoint_id': '1f02ac4a-ecee-600b-8002-a1d21df32e4c'}}

graph.invoke(None, new_config)

{'topic': 'chickens',
 'joke': '为什么鸡要加入乐队？\n\n因为它有很棒的鼓槌！'}