diff --git a/langgraph/store/redis/__init__.py b/langgraph/store/redis/__init__.py index a9172e0..2db0557 100644 --- a/langgraph/store/redis/__init__.py +++ b/langgraph/store/redis/__init__.py @@ -27,7 +27,6 @@ from redisvl.index import SearchIndex from redisvl.query import FilterQuery, VectorQuery from redisvl.redis.connection import RedisConnectionFactory -from redisvl.utils.token_escaper import TokenEscaper from ulid import ULID from langgraph.checkpoint.redis.util import unwrap_connection_errors @@ -43,11 +42,11 @@ _namespace_to_text, _row_to_item, _row_to_search_item, + _tag_escaper, ) from .token_unescaper import TokenUnescaper -_token_escaper = TokenEscaper() _token_unescaper = TokenUnescaper() logger = logging.getLogger(__name__) @@ -324,7 +323,7 @@ def _batch_put_ops( # First delete any existing documents that are being updated/deleted for _, op in put_ops: namespace = _namespace_to_text(op.namespace) - query = f"@prefix:{namespace} @key:{{{_token_escaper.escape(op.key)}}}" + query = f"@prefix:{namespace} @key:{{{_tag_escaper.escape(op.key)}}}" results = self.store_index.search(query) if self.cluster_mode: diff --git a/langgraph/store/redis/aio.py b/langgraph/store/redis/aio.py index 75adcdc..b72c1fc 100644 --- a/langgraph/store/redis/aio.py +++ b/langgraph/store/redis/aio.py @@ -26,7 +26,6 @@ from redisvl.index import AsyncSearchIndex from redisvl.query import FilterQuery, VectorQuery from redisvl.redis.connection import RedisConnectionFactory -from redisvl.utils.token_escaper import TokenEscaper from ulid import ULID from langgraph.checkpoint.redis.util import unwrap_connection_errors @@ -42,12 +41,12 @@ _namespace_to_text, _row_to_item, _row_to_search_item, + _tag_escaper, logger, ) from .token_unescaper import TokenUnescaper -_token_escaper = TokenEscaper() _token_unescaper = TokenUnescaper() from redis.asyncio.cluster import RedisCluster as AsyncRedisCluster @@ -499,7 +498,7 @@ async def _aprepare_batch_PUT_queries( # Delete matching documents for op in deletes: prefix = _namespace_to_text(op.namespace) - query = f"(@prefix:{prefix} @key:{{{_token_escaper.escape(op.key)}}})" + query = f"(@prefix:{prefix} @key:{{{_tag_escaper.escape(op.key)}}})" results = await self.store_index.search(query) for doc in results.docs: await self._redis.delete(doc.id) @@ -560,7 +559,7 @@ async def _batch_put_ops( # First delete any existing documents that are being updated/deleted for _, op in put_ops: namespace = _namespace_to_text(op.namespace) - query = f"@prefix:{namespace} @key:{{{_token_escaper.escape(op.key)}}}" + query = f"@prefix:{namespace} @key:{{{_tag_escaper.escape(op.key)}}}" results = await self.store_index.search(query) if self.cluster_mode: diff --git a/langgraph/store/redis/base.py b/langgraph/store/redis/base.py index 76c4f83..3d8be4e 100644 --- a/langgraph/store/redis/base.py +++ b/langgraph/store/redis/base.py @@ -5,6 +5,7 @@ import copy import json import logging +import re import threading from collections import defaultdict from datetime import datetime, timedelta, timezone @@ -39,7 +40,6 @@ from redis.asyncio import Redis as AsyncRedis from redis.exceptions import ResponseError from redisvl.index import SearchIndex -from redisvl.query.filter import Tag, Text from redisvl.utils.token_escaper import TokenEscaper from langgraph.checkpoint.redis.jsonplus_redis import JsonPlusRedisSerializer @@ -48,6 +48,11 @@ from .types import IndexType, RedisClientType _token_escaper = TokenEscaper() +# For tag values. The default set leaves "|", which separates tag alternatives, +# and whitespace other than a space, which ends the value. +_tag_escaper = TokenEscaper( + re.compile(TokenEscaper.DEFAULT_ESCAPED_CHARS[:-1] + r"|\n\r\t\v\f]") +) _token_unescaper = TokenUnescaper() logger = logging.getLogger(__name__) @@ -494,12 +499,12 @@ def _get_batch_GET_ops_queries( results: list[tuple[str, Sequence, tuple[str, ...], list]] = [] for namespace, items in namespace_groups.items(): _, keys = zip(*items) - # Use Tag helper to properly escape all special characters - prefix_filter = Text("prefix") == _namespace_to_text(namespace) - filter_str = f"({prefix_filter} " + # redisvl's Text filter turns a double quote into a space, so it + # can't match the escaped namespace that put() stored. + filter_str = f'(@prefix:("{_namespace_to_text(namespace)}") ' if keys: - key_filter = Tag("key") == list(keys) - filter_str += f"{key_filter})" + key_filter = "|".join(_tag_escaper.escape(key) for key in keys) + filter_str += f"@key:{{{key_filter}}})" else: filter_str += ")" results.append((filter_str, [], namespace, items)) @@ -532,7 +537,7 @@ def _prepare_batch_PUT_queries( # Delete matching documents for op in deletes: prefix = _namespace_to_text(op.namespace) - query = f"(@prefix:{prefix} @key:{{{_token_escaper.escape(op.key)}}})" + query = f"(@prefix:{prefix} @key:{{{_tag_escaper.escape(op.key)}}})" results = self.store_index.search(query) for doc in results.docs: self._redis.delete(doc.id) diff --git a/tests/test_issue_237_store_get_escaping.py b/tests/test_issue_237_store_get_escaping.py new file mode 100644 index 0000000..f0031a9 --- /dev/null +++ b/tests/test_issue_237_store_get_escaping.py @@ -0,0 +1,79 @@ +"""Test for issue #237 - get() must find every item that search() finds. + +``get()`` built its query with redisvl's ``Text`` and ``Tag`` filters. ``Text`` +turned a double quote in a namespace label into a space, and ``Tag`` left a +newline in a key unescaped, so ``get()`` returned ``None`` for items that +``put()`` had stored and ``search()`` could find. + +``delete()`` and the update path of ``put()`` had the same problem with keys, +and with ``|`` too: ``delete()`` left the item in place, and ``put()`` kept the +old document next to the new one. +""" + +import pytest +from langgraph.store.base import GetOp + +from langgraph.store.redis import AsyncRedisStore, RedisStore + +ITEMS = [ + (("issue237", 'q"t'), "k"), + (("issue237", "newline"), "k\n1"), + (("issue237", "carriage-return"), "k\r1"), + (("issue237", "tab"), "k\t1"), + (("issue237", "vertical-tab"), "k\x0b1"), + (("issue237", "form-feed"), "k\x0c1"), + (("issue237", "pipe"), "k|1"), +] + + +@pytest.mark.parametrize("namespace,key", ITEMS) +def test_sync_store_get_update_delete(redis_url, namespace, key) -> None: + with RedisStore.from_conn_string(redis_url) as store: + store.setup() + store.put(namespace, key, {"v": 1}) + assert store.get(namespace, key).value == {"v": 1} + + store.put(namespace, key, {"v": 2}) + assert store.get(namespace, key).value == {"v": 2} + assert [item.value for item in store.search(namespace)] == [{"v": 2}] + + store.delete(namespace, key) + assert store.get(namespace, key) is None + assert store.search(namespace) == [] + + +@pytest.mark.parametrize("namespace,key", ITEMS) +async def test_async_store_get_update_delete(redis_url, namespace, key) -> None: + async with AsyncRedisStore.from_conn_string(redis_url) as store: + await store.setup() + await store.aput(namespace, key, {"v": 1}) + assert (await store.aget(namespace, key)).value == {"v": 1} + + await store.aput(namespace, key, {"v": 2}) + assert (await store.aget(namespace, key)).value == {"v": 2} + assert [item.value for item in await store.asearch(namespace)] == [{"v": 2}] + + await store.adelete(namespace, key) + assert await store.aget(namespace, key) is None + assert await store.asearch(namespace) == [] + + +def test_batched_get_with_several_keys(redis_url) -> None: + namespace = ("issue237", "batch") + keys = ["plain", "new\nline", "pi|pe"] + with RedisStore.from_conn_string(redis_url) as store: + store.setup() + for key in keys: + store.put(namespace, key, {"key": key}) + + items = store.batch([GetOp(namespace, key) for key in keys]) + + assert [item.value["key"] for item in items] == keys + + +def test_get_stays_within_the_exact_namespace(redis_url) -> None: + with RedisStore.from_conn_string(redis_url) as store: + store.setup() + store.put(("issue237", "parent", "child"), "k", {"v": 1}) + + assert store.get(("issue237", "parent"), "k") is None