langgraph_inmemory_server_quickstart_with_sdk_client.py

1import asyncio
2from langgraph.builder import StateGraph, START
3from langgraph_sdk.runtime.inmem import InMemServer
4from typing import TypedDict
5
6# 1. Define a simple graph
7class State(TypedDict):
8    count: int
9
10def increment(state: State):
11    return {"count": state.get("count", 0) + 1}
12
13builder = StateGraph(State)
14builder.add_node("increment", increment)
15builder.add_edge(START, "increment")
16graph = builder.compile()
17
18async def main():
19    # 2. Initialize the In-Memory Server with your graph
20    # The keys in the dictionary are the assistant IDs
21    async with InMemServer({"my_assistant": graph}) as server:
22        # 3. Create a client to interact with the local server
23        client = server.get_client()
24        
25        # 4. Create a thread
26        thread = await client.threads.create()
27        
28        # 5. Run the graph via the SDK client
29        input_data = {"count": 0}
30        async for chunk in client.runs.stream(
31            thread["thread_id"],
32            "my_assistant",
33            input=input_data
34        ):
35            print(chunk)
36
37if __name__ == "__main__":
38    asyncio.run(main())