Skip to content

Commit 6523f76

Browse files
feat: Support ElastiCache Valkey LangGraph Checkpointer (#697)
PR adds support for using ElastiCache [Valkey](https://valkey.io/) datastores as LangGraph checkpointer to store short-term (thread-scoped) memory.
1 parent 9cb856b commit 6523f76

File tree

19 files changed

+6016
-36
lines changed

19 files changed

+6016
-36
lines changed

README.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ The following packages are hosted in this repository:
2020

2121
### LangGraph
2222

23-
- **Checkpointers**: Provides a custom checkpointing solution for LangGraph agents using either the [AgentCore Memory Service](https://docs.aws.amazon.com/bedrock-agentcore/latest/devguide/memory.html) or the [AWS Bedrock Session Management Service](https://docs.aws.amazon.com/bedrock/latest/userguide/sessions.html).
23+
- **Checkpointers**: Provides a custom checkpointing solution for LangGraph agents using the [AgentCore Memory Service](https://docs.aws.amazon.com/bedrock-agentcore/latest/devguide/memory.html), the [AWS Bedrock Session Management Service](https://docs.aws.amazon.com/bedrock/latest/userguide/sessions.html), or the [ElastiCache Valkey Service](https://aws.amazon.com/elasticache/).
2424
- **Memory Stores** - Provides a memory store solution for saving, processing, and retrieving intelligent long term memories using the [AgentCore Memory Service](https://docs.aws.amazon.com/bedrock-agentcore/latest/devguide/memory.html).
2525

2626
...and more to come. This repository will continue to expand and offer additional components for various AWS services as development progresses.

libs/langgraph-checkpoint-aws/README.md

Lines changed: 204 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,33 +1,43 @@
11
# LangGraph Checkpoint AWS
2-
3-
A custom LangChain checkpointer implementation that uses Bedrock AgentCore Memory to enable stateful and resumable LangGraph agents through efficient state persistence and retrieval.
2+
A custom AWS-based persistence solution for LangGraph agents that provides multiple storage backends including Bedrock AgentCore Memory and high-performance Valkey (Redis-compatible) storage.
43

54
## Overview
5+
This package provides multiple persistence solutions for LangGraph agents:
66

7-
This package provides a custom checkpointing solution for LangGraph agents using AWS Bedrock AgentCore Memory Service. It enables:
8-
7+
### AWS Bedrock AgentCore Memory Service
98
1. Stateful conversations and interactions
109
2. Resumable agent sessions
1110
3. Efficient state persistence and retrieval
1211
4. Seamless integration with AWS Bedrock
1312

13+
### Valkey Storage Solutions
14+
1. **Checkpoint storage** with Valkey (Redis-compatible)
15+
1416
## Installation
1517

1618
You can install the package using pip:
1719

1820
```bash
21+
# Base package (includes Bedrock AgentCore Memory components)
1922
pip install langgraph-checkpoint-aws
20-
```
2123

22-
## Requirements
24+
# With Valkey support
25+
pip install 'langgraph-checkpoint-aws[valkey]'
2326

24-
```text
25-
Python >=3.9
26-
langgraph >=0.2.55
27-
boto3 >=1.39.7
2827
```
2928

30-
## Usage - Checkpointer
29+
## Components
30+
31+
This package provides three main components:
32+
33+
1. **AgentCoreMemorySaver** - AWS Bedrock-based checkpoint storage
34+
2. **ValkeySaver** - Valkey checkpoint storage
35+
3. **AgentCoreMemoryStore** - AWS Bedrock-based document store
36+
37+
38+
## Usage
39+
40+
### 1. Bedrock Session Management
3141

3242
```python
3343
# Import LangGraph and LangChain components
@@ -70,7 +80,7 @@ response = graph.invoke(
7080
)
7181
```
7282

73-
## Usage - Memory Store
83+
### 2. Bedrock Memory Store
7484

7585
```python
7686
# Import LangGraph and LangChain components
@@ -138,6 +148,92 @@ response = graph.invoke(
138148
)
139149
```
140150

151+
### 3. Valkey Checkpoint Storage
152+
153+
```python
154+
from langgraph.graph import StateGraph
155+
from langgraph_checkpoint_aws import ValkeySaver
156+
157+
# Using connection string
158+
with ValkeySaver.from_conn_string(
159+
"valkey://localhost:6379",
160+
ttl_seconds=3600, # 1 hour TTL
161+
pool_size=10
162+
) as checkpointer:
163+
# Create your graph
164+
builder = StateGraph(int)
165+
builder.add_node("add_one", lambda x: x + 1)
166+
builder.set_entry_point("add_one")
167+
builder.set_finish_point("add_one")
168+
169+
graph = builder.compile(checkpointer=checkpointer)
170+
config = {"configurable": {"thread_id": "session-1"}}
171+
result = graph.invoke(1, config)
172+
```
173+
174+
## Async Usage
175+
176+
All components support async operations:
177+
178+
```python
179+
from langgraph_checkpoint_aws.async_saver import AsyncBedrockSessionSaver
180+
from langgraph_checkpoint_aws.checkpoint.valkey import AsyncValkeySaver
181+
182+
# Async Bedrock usage
183+
session_saver = AsyncBedrockSessionSaver(region_name="us-west-2")
184+
session_id = (await session_saver.session_client.create_session()).session_id
185+
186+
# Async Valkey usage
187+
async with AsyncValkeySaver.from_conn_string("valkey://localhost:6379") as checkpointer:
188+
graph = builder.compile(checkpointer=checkpointer)
189+
result = await graph.ainvoke(1, {"configurable": {"thread_id": "session-1"}})
190+
```
191+
192+
## Configuration Options
193+
194+
### Bedrock Session Saver
195+
196+
`BedrockSessionSaver` and `AsyncBedrockSessionSaver` accept the following parameters:
197+
198+
```python
199+
def __init__(
200+
client: Optional[Any] = None,
201+
session: Optional[boto3.Session] = None,
202+
region_name: Optional[str] = None,
203+
credentials_profile_name: Optional[str] = None,
204+
aws_access_key_id: Optional[SecretStr] = None,
205+
aws_secret_access_key: Optional[SecretStr] = None,
206+
aws_session_token: Optional[SecretStr] = None,
207+
endpoint_url: Optional[str] = None,
208+
config: Optional[Config] = None,
209+
)
210+
```
211+
212+
### Valkey Components
213+
214+
Valkey components support these common configuration options:
215+
216+
#### Connection Options
217+
- **Connection String**: `valkey://localhost:6379` or `valkeys://localhost:6380` (SSL). Refer [connection examples](https://valkey-py.readthedocs.io/en/latest/examples/connection_examples.html).
218+
- **Connection Pool**: Reusable connection pools for better performance
219+
- **Pool Size**: Maximum number of connections (default: 10)
220+
- **SSL Support**: Secure connections with certificate validation
221+
222+
#### Performance Options
223+
- **TTL (Time-to-Live)**: Automatic expiration of stored data
224+
- **Batch Operations**: Efficient bulk operations for better throughput
225+
- **Async Support**: Non-blocking operations for high concurrency
226+
227+
#### ValkeySaver Options
228+
```python
229+
valkey_client = Valkey.from_url("valkey://localhost:6379")
230+
ValkeySaver(
231+
client: valkey_client,
232+
ttl: float | None = None, # TTL in seconds
233+
serde: SerializerProtocol | None = None # Custom serialization
234+
)
235+
```
236+
141237
## Development
142238

143239
Setting Up Development Environment
@@ -187,7 +283,9 @@ make spell_check # Check spelling
187283
make clean # Remove all generated files
188284
```
189285

190-
## AWS Configuration
286+
## Infrastructure Setup
287+
288+
### AWS Configuration (for Bedrock components)
191289

192290
Ensure you have AWS credentials configured using one of these methods:
193291

@@ -196,14 +294,14 @@ Ensure you have AWS credentials configured using one of these methods:
196294
3. IAM roles
197295
4. Direct credential injection via constructor parameters
198296

199-
## Required AWS permissions
297+
Required AWS permissions for Bedrock Session Management:
200298

201299
```json
202300
{
203301
"Version": "2012-10-17",
204302
"Statement": [
205303
{
206-
"Sid": "Statement1",
304+
"Sid": "BedrockSessionManagement",
207305
"Effect": "Allow",
208306
"Action": [
209307
"bedrock-agentcore:CreateEvent",
@@ -325,11 +423,10 @@ def __init__(
325423
"bedrock:GetInvocationStep",
326424
"bedrock:ListInvocationSteps"
327425
],
328-
"Resource": [
329-
"*"
330-
]
426+
"Resource": ["*"]
331427
},
332428
{
429+
"Sid": "KMSAccess",
333430
"Effect": "Allow",
334431
"Action": [
335432
"kms:Decrypt",
@@ -338,20 +435,68 @@ def __init__(
338435
"kms:DescribeKey"
339436
],
340437
"Resource": "arn:aws:kms:{region}:{account}:key/{kms-key-id}"
341-
},
342-
{
343-
"Effect": "Allow",
344-
"Action": [
345-
"bedrock:TagResource",
346-
"bedrock:UntagResource",
347-
"bedrock:ListTagsForResource"
348-
],
349-
"Resource": "arn:aws:bedrock:{region}:{account}:session/*"
350438
}
351439
]
352440
}
353441
```
354442

443+
### Valkey Setup
444+
445+
#### Using AWS ElastiCache for Valkey (Recommended)
446+
```python
447+
# Connect to AWS ElastiCache from host running inside VPC with access to cache
448+
from langgraph_checkpoint_aws.checkpoint.valkey import ValkeySaver
449+
450+
with ValkeySaver.from_conn_string(
451+
"valkeys://your-elasticache-cluster.amazonaws.com:6379",
452+
pool_size=20
453+
) as checkpointer:
454+
pass
455+
```
456+
If you want to connect to cache from a host outside of VPC, use ElastiCache console to setup a jump host so you could create SSH tunnel to access cache locally.
457+
458+
#### Using Docker
459+
```bash
460+
# Start Valkey with required modules
461+
docker run --name valkey-bundle -p 6379:6379 -d valkey/valkey-bundle:latest
462+
463+
# Or with custom configuration
464+
docker run --name valkey-custom \
465+
-p 6379:6379 \
466+
-v $(pwd)/valkey.conf:/etc/valkey/valkey.conf \
467+
-d valkey/valkey-bundle:latest
468+
```
469+
470+
## Performance and Best Practices
471+
472+
### Valkey Performance Optimization
473+
474+
#### Connection Pooling
475+
```python
476+
# Use connection pools for better performance
477+
from valkey.connection import ConnectionPool
478+
479+
pool = ConnectionPool.from_url(
480+
"valkey://localhost:6379",
481+
max_connections=20,
482+
retry_on_timeout=True
483+
)
484+
485+
with ValkeySaver.from_pool(pool) as checkpointer:
486+
# Reuse connections across operations
487+
pass
488+
```
489+
490+
#### TTL Strategy
491+
```python
492+
# Configure appropriate TTL values
493+
with ValkeySaver.from_conn_string(
494+
"valkey://localhost:6379",
495+
ttl_seconds=3600 # 1 hour for active sessions
496+
) as checkpointer:
497+
pass
498+
```
499+
355500
## Security Considerations
356501

357502
* Never commit AWS credentials
@@ -361,6 +506,37 @@ def __init__(
361506
* Use IAM roles and temporary credentials when possible
362507
* Implement proper access controls for session management
363508

509+
### Valkey Security
510+
* Use SSL/TLS for production deployments (`valkeys://` protocol), refer [SSL connection examples](https://valkey-py.readthedocs.io/en/latest/examples/ssl_connection_examples.html#Connect-to-a-Valkey-instance-via-SSL,-and-validate-OCSP-stapled-certificates)
511+
* Configure authentication with strong passwords
512+
* Implement network security (VPC, security groups)
513+
* Regular security updates and monitoring
514+
* Use AWS ElastiCache for managed Valkey with encryption
515+
516+
```python
517+
# Secure connection example
518+
import os
519+
import valkey
520+
521+
pki_dir = os.path.join("..", "..", "dockers", "stunnel", "keys")
522+
523+
valkey_client = valkey.Valkey(
524+
host="localhost",
525+
port=6666,
526+
ssl=True,
527+
ssl_certfile=os.path.join(pki_dir, "client-cert.pem"),
528+
ssl_keyfile=os.path.join(pki_dir, "client-key.pem"),
529+
ssl_cert_reqs="required",
530+
ssl_ca_certs=os.path.join(pki_dir, "ca-cert.pem"),
531+
)
532+
533+
checkpointer = ValkeySaver(valkey_client)
534+
```
535+
536+
## Examples and Samples
537+
538+
Comprehensive examples are available in the `samples/memory/` directory:
539+
364540
## Contributing
365541

366542
* Fork the repository
@@ -377,5 +553,5 @@ This project is licensed under the MIT License - see the [LICENSE](LICENSE) file
377553
## Acknowledgments
378554

379555
* LangChain team for the base LangGraph framework
380-
381556
* AWS Bedrock team for the session management service
557+
* Valkey team for the Redis-compatible storage

libs/langgraph-checkpoint-aws/langgraph_checkpoint_aws/__init__.py

Lines changed: 29 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,19 +3,47 @@
33
Bedrock Session Management Service.
44
"""
55

6+
from importlib.metadata import version
7+
68
from langgraph_checkpoint_aws.agentcore.saver import (
79
AgentCoreMemorySaver,
810
)
911
from langgraph_checkpoint_aws.agentcore.store import (
1012
AgentCoreMemoryStore,
1113
)
1214

13-
__version__ = "1.0.0"
15+
# Conditional imports for checkpoint functionality
16+
try:
17+
from langgraph_checkpoint_aws.checkpoint import AsyncValkeySaver, ValkeySaver
18+
19+
valkey_available = True
20+
except ImportError:
21+
# If checkpoint dependencies are not available, create placeholder classes
22+
from typing import Any
23+
24+
def _missing_checkpoint_dependencies_error(*args: Any, **kwargs: Any) -> Any:
25+
raise ImportError(
26+
"Valkey checkpoint functionality requires optional dependencies. "
27+
"Install them with: pip install 'langgraph-checkpoint-aws[valkey]'"
28+
)
29+
30+
# Create placeholder classes that raise helpful errors
31+
AsyncValkeySaver: type[Any] = _missing_checkpoint_dependencies_error # type: ignore[assignment,no-redef]
32+
ValkeySaver: type[Any] = _missing_checkpoint_dependencies_error # type: ignore[assignment,no-redef]
33+
valkey_available = False
34+
35+
try:
36+
__version__ = version("langgraph-checkpoint-aws")
37+
except Exception:
38+
# Fallback version if package is not installed
39+
__version__ = "1.0.0a1"
1440
SDK_USER_AGENT = f"LangGraphCheckpointAWS#{__version__}"
1541

1642
# Expose the saver class at the package level
1743
__all__ = [
1844
"AgentCoreMemorySaver",
1945
"AgentCoreMemoryStore",
46+
"AsyncValkeySaver",
47+
"ValkeySaver",
2048
"SDK_USER_AGENT",
2149
]

0 commit comments

Comments
 (0)