[Python] How to rewind time for human to edit message?

There is a bug which I understand but cannot solve. When I enter e to rewind, sometimes it doesn't work properly. This is because each rewind starts things at the chatbot node, creating 2 chatbot nodes after the call, screwing things up. This could be solved by adding a rewind counter to state and adding this counter to the rewind function to retrieve the correct message, but I wonder if there is a more elegant solution.

The more important issue is that old states are still present in the graph. This makes it much more difficult to create any logic to skip through these 'old states'. I need a way to clear the state history for the old states that are redundant now that I rewound time.

main.py:

from langchain_core.messages import HumanMessage, SystemMessage, trim_messages
from langgraph.types import Command
from build_graph import graph

user_id = "1"

global_config = {"configurable": {"thread_id": "1"}}

def create_prompt(user_type):
return f"You are an agent called Ethanbot, Ethan's web portfolio manager. You are speaking to {user_type}"

def rewind(num_rewind:int, config, user_input):
global graph
num_encountered = 0
for state in graph.get_state_history(config):
if "chatbot" in state.next:
num_encountered+=1
if num_rewind == num_encountered:
config = state.config
last_message = state.values["messages"][-1]
print("last message: ", last_message)
new_message = HumanMessage(
content=user_input,
id=last_message.id
)
graph.update_state(config, {"messages": [new_message]})
break
# return config

def stream_graph_updates(user_input: str, user_type: str, user_id: str, num_rewind: int, config: dict):
state = {
"messages": [SystemMessage(content=create_prompt(user_type)), {"role": "user", "content": user_input}],
"user_type": user_type,
"user_id": user_id,
}
if num_rewind != 0:
rewind(int(num_rewind), config, user_input)
for event in graph.stream(state, config):
if "__interrupt__" in event:

action = input("continue or no: ")
for resume_event in graph.stream(Command(resume={"action": action}), config):
# print("RESUME_EVENT: ", resume_event)
try:
is_chatbot = resume_event.get("chatbot", False)
except:
is_chatbot = False
if is_chatbot:
msg = resume_event["chatbot"]["messages"][-1].content
if msg:
print("ASSISTANT:", msg, "\n")

else:
for value in event.values():
# print("Value: ", value)
msg = value["messages"][-1].content
if msg:
print("ASSISTANT:", msg, "\n")

while True:
num_rewind = 0

config = global_config

for state in graph.get_state_history(config):
print("Num Messages: ", len(state.values["messages"]), "Next: ", state.next)
print("-" * 80)

user_input = input("User ('q' to exit, 'e' to edit previous messages): ")

if user_input.lower() in ["quit", "q"]:
break

elif user_input.lower() in ["edit", "e"]:
num_rewind = input("Number of messages to rewind (num pls): ")
user_input = input("User: ")

user_type = input("User Type: ").lower()
if user_type == "":
user_type = "viewer"

stream_graph_updates(user_input, user_type, user_id, num_rewind, config)


build_graph.py:

from langchain_openai import ChatOpenAI
from langchain_core.messages import ToolMessage
from typing import Annotated

from typing_extensions import TypedDict, Literal

from langgraph.checkpoint.memory import MemorySaver

from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langgraph.types import Command, interrupt
from dotenv import load_dotenv


import os

# Load environment variables
load_dotenv()

# Constants and configuration
NVIDIA_API_KEY = os.getenv("NVIDIA_API_KEY")
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

llm = ChatOpenAI(
model="gpt-4o-mini",
temperature=0.8,
max_tokens=5000,
timeout=20,
max_retries=2,
api_key=OPENAI_API_KEY,
)


from langchain_core.tools import tool, InjectedToolArg
from typing import Annotated

@tool
def ban_user(user_id: Annotated[str, InjectedToolArg]) -> str:
"""Ban the user from the website for 24 hours"""
return f"User {user_id} has been banned"

tools = [ban_user]

memory = MemorySaver()

llm_with_tools = llm.bind_tools(tools)

class State(TypedDict):
"""Add attributes that are mutable via nodes, for example if the user type can change from guest to user with the help of
a node in the graph, we should add user type as an attribute"""

# Messages have the type "list". The `add_messages` function
# in the annotation defines how this state key should be updated
# (in this case, it appends messages to the list, rather than overwriting them)
messages: Annotated[list, add_messages]
user_type: str
user_id: str


graph_builder = StateGraph(State)

from copy import deepcopy

def chatbot(state: State):
message = llm_with_tools.invoke(state["messages"])
assert len(message.tool_calls) <= 1
tool_calls = []
for tool_call in message.tool_calls:
if tool_call["name"] == "ban_user":
tool_call_copy = deepcopy(tool_call)
tool_call_copy["args"]["user_id"] = state["user_id"]
tool_calls.append(tool_call_copy)
else:
tool_calls.append(tool_call)

message.tool_calls = tool_calls

return {"messages": [message], "user_type": state["user_type"]}

def tool_node(state):
new_messages = []
tools = {"ban_user": ban_user}
tool_calls = state["messages"][-1].tool_calls
for tool_call in tool_calls:
tool = tools[tool_call["name"]]
result = tool.invoke(tool_call["args"])
print("result: ", result)
new_messages.append(
{
"role": "tool",
"name": tool_call["name"],
"content": result,
"tool_call_id": tool_call["id"],
}
)
return {"messages": new_messages}

def route_after_llm(state) -> Literal[END, "human_review_node"]:
if len(state["messages"][-1].tool_calls) == 0:
return END
else:
return "human_review_node"

def human_review_node(state) -> Command[Literal["chatbot", "tools"]]:
last_message = state["messages"][-1]
tool_call = last_message.tool_calls[-1]

# this is the value we'll be providing via Command(resume=<human_review>)
human_review = interrupt(
{
"question": "Is this correct?",
# Surface tool calls for review
"tool_call": tool_call,
}
)

review_action = human_review["action"]
review_data = human_review.get("data")

# if approved, call the tool
if review_action == "continue":
return Command(goto="tools")
else:
tool_message = ToolMessage(
role="tool",
name=tool_call["name"],
tool_call_id=tool_call["id"],
content="Tool execution skipped.",
)
return Command(goto="chatbot", update={"messages": [tool_message]})

graph_builder.add_node("chatbot", chatbot)
graph_builder.add_node("tools", tool_node)
graph_builder.add_node(human_review_node)
graph_builder.add_conditional_edges(
"chatbot",
route_after_llm,
)
graph_builder.add_edge("tools", "chatbot")
graph_builder.add_edge(START, "chatbot")

graph = graph_builder.compile(checkpointer=memory)

try:
with open("graph_output.png", "wb") as f:
f.write(graph.get_graph().draw_mermaid_png())
except Exception:
# This requires some extra dependencies and is optional
pass

Continue reading...