diff --git a/src/oss/langchain/short-term-memory.mdx b/src/oss/langchain/short-term-memory.mdx index fb771648f9..9138212245 100644 --- a/src/oss/langchain/short-term-memory.mdx +++ b/src/oss/langchain/short-term-memory.mdx @@ -107,7 +107,40 @@ const checkpointer = PostgresSaver.fromConnString(DB_URI); ``` ::: +You can also use Azure Cosmos DB as a checkpoint saver: +:::python +```shell +pip install langgraph-checkpoint-cosmosdb +``` + +```python +from langgraph.graph import StateGraph, MessagesState, START +from langgraph_checkpoint_cosmosdb import CosmosDBSaver # [!code highlight] +import os + +# Set environment variables for authentication +os.environ["COSMOSDB_ENDPOINT"] = "your_cosmosdb_endpoint" +os.environ["COSMOSDB_KEY"] = "your_cosmosdb_key" + +# Database and Container are created if they don't exist +checkpointer = CosmosDBSaver( # [!code highlight] + database_name="your_database", # [!code highlight] + container_name="your_container" # [!code highlight] +) # [!code highlight] + +# Build your graph with the checkpointer +builder = StateGraph(MessagesState) +builder.add_node("call_model", call_model) +builder.add_edge(START, "call_model") +graph = builder.compile(checkpointer=checkpointer) # [!code highlight] +``` + +::: + + + Azure Cosmos DB checkpointer supports both sync and async operations. If the database and container already exist, you can use default RBAC credentials (e.g., `az login`) instead of setting the endpoint and key. + ## Customizing agent memory By default, agents use @[`AgentState`] to manage short term memory, specifically the conversation history via a `messages` key. @@ -679,7 +712,7 @@ def update_user_info( runtime: ToolRuntime[CustomContext, CustomState], ) -> Command: """Look up and update user info.""" - user_id = runtime.context.user_id + user_id = runtime.context.user_id name = "John Smith" if user_id == "user_123" else "Unknown user" return Command(update={ # [!code highlight] "user_name": name, @@ -704,7 +737,7 @@ agent = create_agent( model="gpt-5-nano", tools=[update_user_info, greet], state_schema=CustomState, # [!code highlight] - context_schema=CustomContext, + context_schema=CustomContext, ) agent.invoke(