Is CompiledStateGraph Thread safe and How to Use MemorySaver (checkpointer)

[Mohamed-AtiaB]

I am trying to publish my graph using FastAPI, I have 2 questions:

1. is the CompiledStateGraph Thread safe? graph instance from the below code. Can I use it as a singleton, or I have to create a new instance per request (like in the code below)

app = FastAPI()
memory = MemorySaver()

sessions = {
    
}

@app.post("/test")
def test(request: RequestModel):
    graphConfig = {
        "configurable": {"thread_id": request.session_id, "checkpoint_ns": "" },
        "recursion_limit": 150,
        }

    graph = graph_builder.compile(checkpointer=memory)
    
    checkpoint = sessions.get(request.session_id)
    state = None
    
    if not checkpoint:
        state = State()
        state["messages"] = []
        state["user_data"] = request.user_data
        state["messages"].append(HumanMessage(request.message))
    else:
        state = checkpoint.values
        state['messages'].append(HumanMessage(request.message))
    
    result = graph.invoke(state, graphConfig)

    state = graph.get_state(graphConfig)
    sessions[request.session_id] = state


    return {"message": result.get("messages")[-1].content}

2. I could not find/understand how to use MemorySaver across different requests.

   • on receiving the first user message in the thread/session, store/insert the snapshot/checkpoint/state in the MemorySaver
   • on receiving subsequent user messages on the same thread, load state, append new message, and invoke graph

   The previous code sinppet used a simple way, simply maintain the state dictionary (i can later store it in redis manually), but how do I use default implementation.
   What is the difference between snapshot/checkpoint/state?
   Note: i am using MemorySaver for simplicity, but optimally i want to use Redis as storage.

@app.post("/test2")
def test2(request: RequestModel):
    graphConfig = {
        "configurable": {"thread_id": request.session_id, "checkpoint_ns": "" },
        "recursion_limit": 150,
        }

    result = None
    graph = None
    checkpoint:Checkpoint = memory.get(graphConfig)
    
    if checkpoint:
        graph = graph_builder.compile(checkpointer=memory)
        checkpoint["channel_values"]["messages"].append(HumanMessage(request.message))
        result = graph.invoke(checkpoint, graphConfig)
    else:
        graph = graph_builder.compile(checkpointer=memory)
        state = State()
        state["messages"] = []
        state["node_history"] = []
        state["user_data"] = request.user_data
        state["messages"].append(HumanMessage(request.message))
        result = graph.invoke(state, graphConfig)
    

    
    checkpoint:Checkpoint = memory.get(graphConfig)
    graphConfig["configurable"]["checkpoint_id"] = checkpoint["id"]
    metadata = CheckpointMetadata(source="input")
    memory.put(graphConfig, checkpoint, metadata, None)

    return {"message": result.get("messages")[-1].content}

[vbarda]

@Mohamed-AtiaB

1. yes, compiled state graph is thread safe, and you should be able to just define a single instance of the compiled graph
2. the way you use any checkpointer across different requests is by passing a config with thread_id field, like you're already doing in graphConfig.

importantly, you have an anti-pattern in your code of getting the latest state from sessions/checkpointer, modifying it and passing it to graph.invoke. you should not do that -- as long as you're using a checkpointer, graph will automatically have access to the state as of the last execution (for a given thread ID). instead, you should just simply invoke it with new inputs like graph.invoke(inputs, config)

i would recommend checking out these sections of documentation https://langchain-ai.github.io/langgraph/concepts/low_level/#threads and https://langchain-ai.github.io/langgraph/how-tos/persistence/

[nimishdm-airindia]

I'm using FastAPI with AsyncMongoDBSaver. Is it okay to keep the checkpointer as singleton as well? What's the recommended approach?

DB_URI = "URI"
graph = None
checkpointer_cm = None
checkpointer = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    global graph, checkpointer_cm, checkpointer
    checkpointer_cm = AsyncMongoDBSaver.from_conn_string(DB_URI)
    checkpointer = await checkpointer_cm.__aenter__()

    builder = StateGraph(State)
    graph = builder.compile(checkpointer=checkpointer)
    yield
    if checkpointer_cm:
        await checkpointer_cm.__aexit__(None, None, None)

app = FastAPI(lifespan=lifespan)

async def invoke_graph(input_obj, config) -> AsyncGenerator[bytes, None]:
    async for event in graph.astream(
        input_obj,
        config=config,
        stream_mode="values",
    ):
        message = event["messages"][-1]
        message.pretty_print()
        if isinstance(message, AIMessage):
            yield f"data: {message.content}\n\n"

@app.post("/messages")
async def process_message(message: ChatMessage):
    thread_config = {"configurable": {"thread_id": message.conversation_id}}
    snapshot = await graph.aget_state(thread_config)
    input_obj = {"messages": [{"role": "user", "content": message.content}]}

    # Handle interrupt
    if len(snapshot.next) > 0:
        input_obj = Command(resume=message.content)

    return StreamingResponse(
        invoke_graph(input_obj, thread_config),
        status_code=200,
        media_type="text/event-stream",
    )

[Mohamed-AtiaB]

Thanks vbarda, You are right. I ended up doing it as below:

@app.post("/test2")
def test2(request: RequestModel):
    graphConfig = {
        "configurable": {"thread_id": request.session_id, "checkpoint_ns": "" },
        "recursion_limit": 150,
        }

    result = None
    graph = None
    checkpoint:Checkpoint = memory.get(graphConfig)
    
    if checkpoint := checkpointer.get(config):# check if existing session or new
        result = graph.invoke({"messages":[(request.role, request.message)]}, config)
    else:
        state = State()
        state["messages"] = []
        state["node_history"] = []
        state["user_data"] = request.user_data
        state["messages"].append(HumanMessage(request.message))
        result = graph.invoke(state, graphConfig)
    
    return {"message": result.get("messages")[-1].content}