Skip to content

Commit f98149b

Browse files
tylerhutchersontylerhutcherson
committed
refactor to use redisvl 0.4.1+
1 parent fed21c6 commit f98149b

File tree

20 files changed

+409
-209
lines changed

20 files changed

+409
-209
lines changed

‎.github/workflows/test.yml‎

Lines changed: 21 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -7,29 +7,26 @@ on:
77
branches:
88
- main
99

10+
workflow_dispatch:
11+
1012
env:
1113
POETRY_VERSION: "1.8.3"
1214

1315
jobs:
1416
test:
15-
name: Python ${{ matrix.python-version }} - ${{ matrix.connection }} [redis-stack ${{matrix.redis-stack-version}}]
17+
name: Python ${{ matrix.python-version }} - [redis ${{ matrix.redis-version }}]
1618
runs-on: ubuntu-latest
1719

1820
strategy:
1921
fail-fast: false
2022
matrix:
21-
python-version: [3.9, '3.10', 3.11, 3.12]
22-
connection: ['hiredis', 'plain']
23-
redis-stack-version: ['6.2.6-v9', 'latest', 'edge']
24-
25-
services:
26-
redis:
27-
image: redis/redis-stack-server:${{matrix.redis-stack-version}}
28-
ports:
29-
- 6379:6379
23+
python-version: [3.9, 3.11, 3.13]
24+
redis-version: ['6.2.6-v9', 'latest', '8.0-M03']
3025

3126
steps:
32-
- uses: actions/checkout@v2
27+
- name: Check out repository
28+
uses: actions/checkout@v3
29+
3330
- name: Set up Python ${{ matrix.python-version }}
3431
uses: actions/setup-python@v4
3532
with:
@@ -45,15 +42,22 @@ jobs:
4542
run: |
4643
poetry install --all-extras
4744
48-
- name: Install hiredis if needed
49-
if: matrix.connection == 'hiredis'
45+
- name: Set Redis image name
5046
run: |
51-
poetry add hiredis
52-
53-
- name: Set Redis version
47+
if [[ "${{ matrix.redis-version }}" == "8.0-M03" ]]; then
48+
echo "REDIS_IMAGE=redis:${{ matrix.redis-version }}" >> $GITHUB_ENV
49+
else
50+
echo "REDIS_IMAGE=redis/redis-stack-server:${{ matrix.redis-version }}" >> $GITHUB_ENV
51+
fi
52+
53+
- name: Run API tests
54+
if: matrix.redis-version == 'latest'
55+
env:
56+
OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
5457
run: |
55-
echo "REDIS_VERSION=${{ matrix.redis-stack-version }}" >> $GITHUB_ENV
58+
make test-all
5659
5760
- name: Run tests
61+
if: matrix.redis-version != 'latest'
5862
run: |
5963
make test

‎Makefile‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
.PHONY: install format lint test clean redis-start redis-stop check-types check
1+
.PHONY: install format lint test test-all clean redis-start redis-stop check-types check
22

33
install:
44
poetry install --all-extras
@@ -21,6 +21,9 @@ lint: format check-types
2121
test:
2222
poetry run test-verbose
2323

24+
test-all:
25+
poetry run test-verbose --run-api-tests
26+
2427
check: lint test
2528

2629
clean:

‎langgraph/checkpoint/redis/__init__.py‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -51,20 +51,20 @@ def configure_client(
5151
) -> None:
5252
"""Configure the Redis client."""
5353
self._owns_its_client = redis_client is None
54-
5554
self._redis = redis_client or RedisConnectionFactory.get_redis_connection(
5655
redis_url, **connection_args
5756
)
5857

5958
def create_indexes(self) -> None:
60-
self.checkpoints_index = SearchIndex.from_dict(self.SCHEMAS[0])
61-
self.checkpoint_blobs_index = SearchIndex.from_dict(self.SCHEMAS[1])
62-
self.checkpoint_writes_index = SearchIndex.from_dict(self.SCHEMAS[2])
63-
64-
# Connect Redis client to indices
65-
self.checkpoints_index.set_client(self._redis)
66-
self.checkpoint_blobs_index.set_client(self._redis)
67-
self.checkpoint_writes_index.set_client(self._redis)
59+
self.checkpoints_index = SearchIndex.from_dict(
60+
self.SCHEMAS[0], redis_client=self._redis
61+
)
62+
self.checkpoint_blobs_index = SearchIndex.from_dict(
63+
self.SCHEMAS[1], redis_client=self._redis
64+
)
65+
self.checkpoint_writes_index = SearchIndex.from_dict(
66+
self.SCHEMAS[2], redis_client=self._redis
67+
)
6868

6969
def list(
7070
self,

‎langgraph/checkpoint/redis/aio.py‎

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -88,9 +88,15 @@ def configure_client(
8888

8989
def create_indexes(self) -> None:
9090
"""Create indexes without connecting to Redis."""
91-
self.checkpoints_index = AsyncSearchIndex.from_dict(self.SCHEMAS[0])
92-
self.checkpoint_blobs_index = AsyncSearchIndex.from_dict(self.SCHEMAS[1])
93-
self.checkpoint_writes_index = AsyncSearchIndex.from_dict(self.SCHEMAS[2])
91+
self.checkpoints_index = AsyncSearchIndex.from_dict(
92+
self.SCHEMAS[0], redis_client=self._redis
93+
)
94+
self.checkpoint_blobs_index = AsyncSearchIndex.from_dict(
95+
self.SCHEMAS[1], redis_client=self._redis
96+
)
97+
self.checkpoint_writes_index = AsyncSearchIndex.from_dict(
98+
self.SCHEMAS[2], redis_client=self._redis
99+
)
94100

95101
async def __aenter__(self) -> AsyncRedisSaver:
96102
"""Async context manager enter."""
@@ -116,11 +122,6 @@ async def __aexit__(
116122

117123
async def asetup(self) -> None:
118124
"""Initialize Redis indexes asynchronously."""
119-
# Connect Redis client to indices asynchronously
120-
await self.checkpoints_index.set_client(self._redis)
121-
await self.checkpoint_blobs_index.set_client(self._redis)
122-
await self.checkpoint_writes_index.set_client(self._redis)
123-
124125
# Create indexes in Redis asynchronously
125126
await self.checkpoints_index.create(overwrite=False)
126127
await self.checkpoint_blobs_index.create(overwrite=False)

‎langgraph/checkpoint/redis/ashallow.py‎

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -153,11 +153,6 @@ async def from_conn_string(
153153

154154
async def asetup(self) -> None:
155155
"""Initialize Redis indexes asynchronously."""
156-
# Connect Redis client to indices asynchronously
157-
await self.checkpoints_index.set_client(self._redis)
158-
await self.checkpoint_blobs_index.set_client(self._redis)
159-
await self.checkpoint_writes_index.set_client(self._redis)
160-
161156
# Create indexes in Redis asynchronously
162157
await self.checkpoints_index.create(overwrite=False)
163158
await self.checkpoint_blobs_index.create(overwrite=False)
@@ -557,9 +552,15 @@ def configure_client(
557552

558553
def create_indexes(self) -> None:
559554
"""Create indexes without connecting to Redis."""
560-
self.checkpoints_index = AsyncSearchIndex.from_dict(self.SCHEMAS[0])
561-
self.checkpoint_blobs_index = AsyncSearchIndex.from_dict(self.SCHEMAS[1])
562-
self.checkpoint_writes_index = AsyncSearchIndex.from_dict(self.SCHEMAS[2])
555+
self.checkpoints_index = AsyncSearchIndex.from_dict(
556+
self.SCHEMAS[0], redis_client=self._redis
557+
)
558+
self.checkpoint_blobs_index = AsyncSearchIndex.from_dict(
559+
self.SCHEMAS[1], redis_client=self._redis
560+
)
561+
self.checkpoint_writes_index = AsyncSearchIndex.from_dict(
562+
self.SCHEMAS[2], redis_client=self._redis
563+
)
563564

564565
def setup(self) -> None:
565566
"""Initialize the checkpoint_index in Redis."""

‎langgraph/checkpoint/redis/shallow.py‎

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -388,14 +388,15 @@ def configure_client(
388388
)
389389

390390
def create_indexes(self) -> None:
391-
self.checkpoints_index = SearchIndex.from_dict(self.SCHEMAS[0])
392-
self.checkpoint_blobs_index = SearchIndex.from_dict(self.SCHEMAS[1])
393-
self.checkpoint_writes_index = SearchIndex.from_dict(self.SCHEMAS[2])
394-
395-
# Connect Redis client to indices
396-
self.checkpoints_index.set_client(self._redis)
397-
self.checkpoint_blobs_index.set_client(self._redis)
398-
self.checkpoint_writes_index.set_client(self._redis)
391+
self.checkpoints_index = SearchIndex.from_dict(
392+
self.SCHEMAS[0], redis_client=self._redis
393+
)
394+
self.checkpoint_blobs_index = SearchIndex.from_dict(
395+
self.SCHEMAS[1], redis_client=self._redis
396+
)
397+
self.checkpoint_writes_index = SearchIndex.from_dict(
398+
self.SCHEMAS[2], redis_client=self._redis
399+
)
399400

400401
def put_writes(
401402
self,

‎langgraph/store/redis/__init__.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@
88
from contextlib import contextmanager
99
from datetime import datetime, timezone
1010
from typing import Any, Iterable, Iterator, Optional, Sequence, cast
11-
from ulid import ULID
1211

1312
from langgraph.store.base import (
1413
BaseStore,
@@ -26,6 +25,7 @@
2625
from redisvl.query import FilterQuery, VectorQuery
2726
from redisvl.redis.connection import RedisConnectionFactory
2827
from redisvl.utils.token_escaper import TokenEscaper
28+
from ulid import ULID
2929

3030
from langgraph.store.redis.aio import AsyncRedisStore
3131
from langgraph.store.redis.base import (

‎langgraph/store/redis/aio.py‎

Lines changed: 24 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@
77
from datetime import datetime, timezone
88
from types import TracebackType
99
from typing import Any, AsyncIterator, Iterable, Optional, Sequence, cast
10-
from ulid import ULID
1110

1211
from langgraph.store.base import (
1312
BaseStore,
@@ -29,6 +28,7 @@
2928
from redisvl.query import FilterQuery, VectorQuery
3029
from redisvl.redis.connection import RedisConnectionFactory
3130
from redisvl.utils.token_escaper import TokenEscaper
31+
from ulid import ULID
3232

3333
from langgraph.store.redis.base import (
3434
REDIS_KEY_SEPARATOR,
@@ -57,14 +57,15 @@ class AsyncRedisStore(
5757

5858
store_index: AsyncSearchIndex
5959
vector_index: AsyncSearchIndex
60-
_owns_client: bool
60+
_owns_its_client: bool
6161

6262
def __init__(
6363
self,
6464
redis_url: Optional[str] = None,
6565
*,
6666
redis_client: Optional[AsyncRedis] = None,
6767
index: Optional[IndexConfig] = None,
68+
connection_args: Optional[dict[str, Any]] = None,
6869
) -> None:
6970
"""Initialize store with Redis connection and optional index config."""
7071
if redis_url is None and redis_client is None:
@@ -94,10 +95,16 @@ def __init__(
9495
]
9596

9697
# Configure client
97-
self.configure_client(redis_url=redis_url, redis_client=redis_client)
98+
self.configure_client(
99+
redis_url=redis_url,
100+
redis_client=redis_client,
101+
connection_args=connection_args or {},
102+
)
98103

99104
# Create store index
100-
self.store_index = AsyncSearchIndex.from_dict(self.SCHEMAS[0])
105+
self.store_index = AsyncSearchIndex.from_dict(
106+
self.SCHEMAS[0], redis_client=self._redis
107+
)
101108

102109
# Configure vector index if needed
103110
if self.index_config:
@@ -131,7 +138,9 @@ def __init__(
131138
vector_field["attrs"].update(self.index_config["ann_index_config"])
132139

133140
try:
134-
self.vector_index = AsyncSearchIndex.from_dict(vector_schema)
141+
self.vector_index = AsyncSearchIndex.from_dict(
142+
vector_schema, redis_client=self._redis
143+
)
135144
except Exception as e:
136145
raise ValueError(
137146
f"Failed to create vector index with schema: {vector_schema}. Error: {str(e)}"
@@ -145,11 +154,12 @@ def configure_client(
145154
self,
146155
redis_url: Optional[str] = None,
147156
redis_client: Optional[AsyncRedis] = None,
157+
connection_args: Optional[dict[str, Any]] = None,
148158
) -> None:
149159
"""Configure the Redis client."""
150-
self._owns_client = redis_client is None
160+
self._owns_its_client = redis_client is None
151161
self._redis = redis_client or RedisConnectionFactory.get_async_redis_connection(
152-
redis_url
162+
redis_url, **connection_args
153163
)
154164

155165
async def setup(self) -> None:
@@ -160,11 +170,6 @@ async def setup(self) -> None:
160170
self.index_config.get("embed"),
161171
)
162172

163-
# Now connect Redis client to indices
164-
await self.store_index.set_client(self._redis)
165-
if self.index_config:
166-
await self.vector_index.set_client(self._redis)
167-
168173
# Create indices in Redis
169174
await self.store_index.create(overwrite=False)
170175
if self.index_config:
@@ -188,9 +193,13 @@ async def from_conn_string(
188193

189194
def create_indexes(self) -> None:
190195
"""Create async indices."""
191-
self.store_index = AsyncSearchIndex.from_dict(self.SCHEMAS[0])
196+
self.store_index = AsyncSearchIndex.from_dict(
197+
self.SCHEMAS[0], redis_client=self._redis
198+
)
192199
if self.index_config:
193-
self.vector_index = AsyncSearchIndex.from_dict(self.SCHEMAS[1])
200+
self.vector_index = AsyncSearchIndex.from_dict(
201+
self.SCHEMAS[1], redis_client=self._redis
202+
)
194203

195204
async def __aenter__(self) -> AsyncRedisStore:
196205
"""Async context manager enter."""
@@ -210,7 +219,7 @@ async def __aexit__(
210219
except asyncio.CancelledError:
211220
pass
212221

213-
if self._owns_client:
222+
if self._owns_its_client:
214223
await self._redis.aclose() # type: ignore[attr-defined]
215224
await self._redis.connection_pool.disconnect()
216225

‎langgraph/store/redis/base.py‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -121,8 +121,9 @@ def __init__(
121121
]
122122

123123
# Initialize search indices
124-
self.store_index = SearchIndex.from_dict(self.SCHEMAS[0])
125-
self.store_index.set_client(self._redis)
124+
self.store_index = SearchIndex.from_dict(
125+
self.SCHEMAS[0], redis_client=self._redis
126+
)
126127

127128
# Configure vector index if needed
128129
if self.index_config:
@@ -156,8 +157,9 @@ def __init__(
156157
if "ann_index_config" in self.index_config:
157158
vector_field["attrs"].update(self.index_config["ann_index_config"])
158159

159-
self.vector_index = SearchIndex.from_dict(vector_schema)
160-
self.vector_index.set_client(self._redis)
160+
self.vector_index = SearchIndex.from_dict(
161+
vector_schema, redis_client=self._redis
162+
)
161163

162164
def _get_batch_GET_ops_queries(
163165
self,

0 commit comments

Comments
 (0)