Problem passing private state between nodes with graph.stream()

[p-lin]

I would like to pass information between two nodes in a way so that the information is not visible in the message state. For example, in node 1 I have updated the state with some sensitive information (e.g. hate speech, personal data) and then want to pass that information to node 2 to process with some other tool. I do not want to see the contents of the updated state in my message history.

I followed the tutorial from here: https://langchain-ai.github.io/langgraph/how-tos/pass_private_state/#define-and-use-the-graph.

However, when I stream the output of "graph.stream()", I am still seeing the updated state that I would like to remain private.

For simplicity, I have replicated the tutorial code but added in "DO NOT WANT THIS OUTPUT TO BE SEEN" as the contents of the private state.

Would appreciate some help. Thanks.

from langgraph.graph import StateGraph, START, END
from typing_extensions import TypedDict
import sqlite3
from langgraph.checkpoint.sqlite import SqliteSaver
from langchain_core.messages import AnyMessage, SystemMessage, HumanMessage, AIMessage, ChatMessage

conn = sqlite3.connect(":memory:", check_same_thread=False)  
memory = SqliteSaver(conn)


# The overall state of the graph (this is the public state shared across nodes)
class OverallState(TypedDict):
    a: str


# Output from node_1 contains private data that is not part of the overall state
class Node1Output(TypedDict):
    private_data: str


# The private data is only shared between node_1 and node_2
def node_1(state: OverallState) -> Node1Output:
    output = {"private_data": "DO NOT WANT THIS OUTPUT TO BE SEEN"}
    #print(f"Entered node `node_1`:\n\tInput: {state}.\n\tReturned: {output}")
    return output


# Node 2 input only requests the private data available after node_1
class Node2Input(TypedDict):
    private_data: str


def node_2(state: Node2Input) -> OverallState:
    output = {"a": "set by node_2"}
    #print(f"Entered node `node_2`:\n\tInput: {state}.\n\tReturned: {output}")
    return output


# Node 3 only has access to the overall state (no access to private data from node_1)
def node_3(state: OverallState) -> OverallState:
    output = {"a": "set by node_3"}
    #print(f"Entered node `node_3`:\n\tInput: {state}.\n\tReturned: {output}")
    return output


# Build the state graph
builder = StateGraph(OverallState)
builder.add_node(node_1)  # node_1 is the first node
builder.add_node(
    node_2
)  # node_2 is the second node and accepts private data from node_1
builder.add_node(node_3)  # node_3 is the third node and does not see the private data
builder.add_edge(START, "node_1")  # Start the graph with node_1
builder.add_edge("node_1", "node_2")  # Pass from node_1 to node_2
builder.add_edge(
    "node_2", "node_3"
)  # Pass from node_2 to node_3 (only overall state is shared)
builder.add_edge("node_3", END)  # End the graph after node_3
graph = builder.compile(checkpointer= memory)

# Invoke the graph with the initial state

thread = {"configurable": {"thread_id": "1"}}
for s in graph.stream({
    "a": "set at start",
    #"max_revisions": 1,
    #"revision_number": 1,
}, thread):
    print(s)
print()
print(f"Output of graph invocation: {response}")