langgraph-checkpoint-aws

import os import redis import boto3 from langgraph.graph import StateGraph, START from langgraph.checkpoint.base import Checkpoint from langgraph_checkpoint_aws import AWSSessionManagerValkeySaver # 1. Setup AWS client (ensure AWS credentials are configured, e.g., via ~/.aws/credentials or env vars) # The Checkpointer primarily uses DynamoDB via Bedrock Session Management Service. aws_region = os.environ.get("AWS_REGION", "us-east-1") boto3_session = boto3.Session(region_name=aws_region) # 2. Setup Valkey/Redis client # Ensure a Redis/Valkey instance is running and accessible (e.g., local or ElastiCache) redis_url = os.environ.get("REDIS_URL", "redis://localhost:6379/0") try: redis_client = redis.from_url(redis_url) redis_client.ping() # Test connection except redis.exceptions.ConnectionError: print(f"WARNING: Could not connect to Redis at {redis_url}. Please ensure a Redis/Valkey instance is running.") redis_client = None # Mark as failed # 3. Define a simple LangGraph state class AgentState(dict): """A dictionary-like state for the graph.""" pass # 4. Instantiate the AWSSessionManagerValkeySaver # 'application_id' is used to partition state within Bedrock SMS/DynamoDB. application_id = os.environ.get("AWS_APPLICATION_ID", "my_langgraph_app_dev") if redis_client: checkpointer = AWSSessionManagerValkeySaver( boto3_session=boto3_session, redis_client=redis_client, application_id=application_id # Optional: You can specify a custom DynamoDB table name if not using the default Bedrock SMS table. # table_name="MyCustomLangGraphCheckpoints" ) # 5. Build a simple LangGraph with the checkpointer graph_builder = StateGraph(AgentState) graph_builder.add_node("start_node", lambda state: {"message": "Hello from LangGraph!"}) graph_builder.set_entry_point("start_node") graph_builder.set_finish_point("start_node") # Compile the graph with the checkpointer app = graph_builder.compile(checkpointer=checkpointer) # 6. Run the graph with a configurable thread_id for checkpointing thread_id = "unique_session_123" config = {"configurable": {"thread_id": thread_id}} print(f"Running LangGraph with checkpointer for thread_id: {thread_id}") # The first run will create a checkpoint result = app.invoke({}, config=config) print(f"First run output: {result}") # Subsequent runs with the same thread_id will load the checkpoint # (in a more complex graph, state would be loaded and updated) result_again = app.invoke({}, config=config) print(f"Second run output (should be same for this simple graph): {result_again}") print("\nSuccessfully configured AWSSessionManagerValkeySaver and ran a simple graph.") print("Check your AWS Bedrock Session Manager (DynamoDB) and Valkey/Redis instance for stored state.") else: print("Checkpointer instantiation skipped due to Redis connection failure.") print("Please ensure REDIS_URL and AWS credentials are correctly configured.")