Skip to content

Commit 934d269

Browse files
Add client info reporting (#155)
We need to be able to track the usage of clients for measuring adoption and impact. The `CLIENT SETINFO` command can be used for this purpose. The required string format is documented [here](https://redis.io/docs/latest/commands/client-setinfo/#:~:text=The%20CLIENT%20SETINFO%20command%20assigns,CLIENT%20LIST%20and%20CLIENT%20INFO%20). ### Variants - Standalone RedisVL: ``` CLIENT SETINFO LIB-NAME redis-py(redisvl_v0.2.0) CLIENT SETINFO LIB-VER 5.0.4 ``` - Abstraction layers (like LlamaIndex or LangChain): ``` CLIENT SETINFO LIB-NAME redis-py(redisvl_v0.2.0;llama-index-vector-stores-redis_v0.1.0) CLIENT SETINFO LIB-VER 5.0.4 ``` ### Constraints - RedisVL uses both Async connection and standard connection instances from Redis Py - So the technique to run this command needs to properly handle this... - Other wrappers around RedisVL like LangChain and LlamaIndex will need to pass through their lib_name - These libraries generally support the notion of providing your own client instance OR providing the connection string and performing the connection on your behalf -- which also adds some difficulty. ### Planned Route In order to build the proper name string, all clients that wrap redis-py will need to use the following format: ``` {package-name}_v{version} ``` RedisVL will implement the default which is `redisvl_v0.x.x`, but outer wrappers of RedisVL can implement their own by using a lib_name `kwarg` on the index class like ``` SearchIndex(lib_name="langchain_v0.2.1") ```
1 parent c4e3017 commit 934d269

File tree

5 files changed

+173
-65
lines changed

5 files changed

+173
-65
lines changed

.gitignore

+1
Original file line numberDiff line numberDiff line change
@@ -9,5 +9,6 @@ scratch
99
wiki_schema.yaml
1010
docs/_build/
1111
.venv
12+
.env
1213
coverage.xml
1314
dist/

conftest.py

-6
Original file line numberDiff line numberDiff line change
@@ -6,12 +6,6 @@
66
from testcontainers.compose import DockerCompose
77

88

9-
# @pytest.fixture(scope="session")
10-
# def event_loop():
11-
# loop = asyncio.get_event_loop_policy().new_event_loop()
12-
# yield loop
13-
# loop.close()
14-
159

1610
@pytest.fixture(scope="session", autouse=True)
1711
def redis_container():

redisvl/index/index.py

+14-20
Original file line numberDiff line numberDiff line change
@@ -84,25 +84,12 @@ def _process(doc: "Document") -> Dict[str, Any]:
8484
return [_process(doc) for doc in results.docs]
8585

8686

87-
def check_modules_present():
87+
def setup_redis():
8888
def decorator(func):
8989
@wraps(func)
9090
def wrapper(self, *args, **kwargs):
9191
result = func(self, *args, **kwargs)
92-
RedisConnectionFactory.validate_redis_modules(self._redis_client)
93-
return result
94-
95-
return wrapper
96-
97-
return decorator
98-
99-
100-
def check_async_modules_present():
101-
def decorator(func):
102-
@wraps(func)
103-
def wrapper(self, *args, **kwargs):
104-
result = func(self, *args, **kwargs)
105-
RedisConnectionFactory.validate_async_redis_modules(self._redis_client)
92+
RedisConnectionFactory.validate_redis(self._redis_client, self._lib_name)
10693
return result
10794

10895
return wrapper
@@ -175,6 +162,9 @@ def __init__(
175162

176163
self.schema = schema
177164

165+
# set custom lib name
166+
self._lib_name: Optional[str] = kwargs.pop("lib_name", None)
167+
178168
# set up redis connection
179169
self._redis_client: Optional[Union[redis.Redis, aredis.Redis]] = None
180170
if redis_client is not None:
@@ -350,11 +340,13 @@ def connect(self, redis_url: Optional[str] = None, **kwargs):
350340
index.connect(redis_url="redis://localhost:6379")
351341
352342
"""
353-
client = RedisConnectionFactory.connect(redis_url, use_async=False, **kwargs)
343+
client = RedisConnectionFactory.connect(
344+
redis_url=redis_url, use_async=False, **kwargs
345+
)
354346
return self.set_client(client)
355347

356-
@check_modules_present()
357-
def set_client(self, client: redis.Redis):
348+
@setup_redis()
349+
def set_client(self, client: redis.Redis, **kwargs):
358350
"""Manually set the Redis client to use with the search index.
359351
360352
This method configures the search index to use a specific Redis or
@@ -729,10 +721,12 @@ def connect(self, redis_url: Optional[str] = None, **kwargs):
729721
index.connect(redis_url="redis://localhost:6379")
730722
731723
"""
732-
client = RedisConnectionFactory.connect(redis_url, use_async=True, **kwargs)
724+
client = RedisConnectionFactory.connect(
725+
redis_url=redis_url, use_async=True, **kwargs
726+
)
733727
return self.set_client(client)
734728

735-
@check_async_modules_present()
729+
@setup_redis()
736730
def set_client(self, client: aredis.Redis):
737731
"""Manually set the Redis client to use with the search index.
738732

redisvl/redis/connection.py

+110-31
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
1+
import asyncio
12
import os
2-
from typing import Any, Dict, List, Optional, Type
3+
from typing import Any, Dict, List, Optional, Type, Union
34

45
from redis import Redis
56
from redis.asyncio import Redis as AsyncRedis
@@ -10,9 +11,11 @@
1011
ConnectionPool,
1112
SSLConnection,
1213
)
14+
from redis.exceptions import ResponseError
1315

1416
from redisvl.redis.constants import REDIS_REQUIRED_MODULES
1517
from redisvl.redis.utils import convert_bytes
18+
from redisvl.version import __version__
1619

1720

1821
def get_address_from_env() -> str:
@@ -26,6 +29,20 @@ def get_address_from_env() -> str:
2629
return os.environ["REDIS_URL"]
2730

2831

32+
def make_lib_name(*args) -> str:
33+
"""Build the lib name to be reported through the Redis client setinfo
34+
command.
35+
36+
Returns:
37+
str: Redis client library name
38+
"""
39+
custom_libs = f"redisvl_v{__version__}"
40+
for arg in args:
41+
if arg:
42+
custom_libs += f";{arg}"
43+
return f"redis-py({custom_libs})"
44+
45+
2946
class RedisConnectionFactory:
3047
"""Builds connections to a Redis database, supporting both synchronous and
3148
asynchronous clients.
@@ -108,54 +125,116 @@ def get_async_redis_connection(url: Optional[str] = None, **kwargs) -> AsyncRedi
108125
return AsyncRedis.from_url(get_address_from_env(), **kwargs)
109126

110127
@staticmethod
111-
def validate_redis_modules(
112-
client: Redis, redis_required_modules: Optional[List[Dict[str, Any]]] = None
128+
def validate_redis(
129+
client: Union[Redis, AsyncRedis],
130+
lib_name: Optional[str] = None,
131+
redis_required_modules: Optional[List[Dict[str, Any]]] = None,
113132
) -> None:
114-
"""Validates if the required Redis modules are installed.
133+
"""Validates the Redis connection.
115134
116135
Args:
117-
client (Redis): Synchronous Redis client.
136+
client (Redis or AsyncRedis): Redis client.
137+
lib_name (str): Library name to set on the Redis client.
138+
redis_required_modules (List[Dict[str, Any]]): List of required modules and their versions.
118139
119140
Raises:
120141
ValueError: If required Redis modules are not installed.
121142
"""
122-
RedisConnectionFactory._validate_redis_modules(
123-
convert_bytes(client.module_list()), redis_required_modules
124-
)
143+
if isinstance(client, AsyncRedis):
144+
RedisConnectionFactory._run_async(
145+
RedisConnectionFactory._validate_async_redis,
146+
client,
147+
lib_name,
148+
redis_required_modules,
149+
)
150+
else:
151+
RedisConnectionFactory._validate_sync_redis(
152+
client, lib_name, redis_required_modules
153+
)
154+
155+
@staticmethod
156+
def _validate_sync_redis(
157+
client: Redis,
158+
lib_name: Optional[str],
159+
redis_required_modules: Optional[List[Dict[str, Any]]],
160+
) -> None:
161+
"""Validates the sync client."""
162+
# Set client library name
163+
_lib_name = make_lib_name(lib_name)
164+
try:
165+
client.client_setinfo("LIB-NAME", _lib_name) # type: ignore
166+
except ResponseError:
167+
# Fall back to a simple log echo
168+
client.echo(_lib_name)
169+
170+
# Get list of modules
171+
modules_list = convert_bytes(client.module_list())
172+
173+
# Validate available modules
174+
RedisConnectionFactory._validate_modules(modules_list, redis_required_modules)
125175

126176
@staticmethod
127-
def validate_async_redis_modules(
177+
async def _validate_async_redis(
128178
client: AsyncRedis,
129-
redis_required_modules: Optional[List[Dict[str, Any]]] = None,
179+
lib_name: Optional[str],
180+
redis_required_modules: Optional[List[Dict[str, Any]]],
130181
) -> None:
182+
"""Validates the async client."""
183+
# Set client library name
184+
_lib_name = make_lib_name(lib_name)
185+
try:
186+
await client.client_setinfo("LIB-NAME", _lib_name) # type: ignore
187+
except ResponseError:
188+
# Fall back to a simple log echo
189+
await client.echo(_lib_name)
190+
191+
# Get list of modules
192+
modules_list = convert_bytes(await client.module_list())
193+
194+
# Validate available modules
195+
RedisConnectionFactory._validate_modules(modules_list, redis_required_modules)
196+
197+
@staticmethod
198+
def _run_async(coro, *args, **kwargs):
131199
"""
132-
Validates if the required Redis modules are installed.
200+
Runs an asynchronous function in the appropriate event loop context.
201+
202+
This method checks if there is an existing event loop running. If there is,
203+
it schedules the coroutine to be run within the current loop using `asyncio.ensure_future`.
204+
If no event loop is running, it creates a new event loop, runs the coroutine,
205+
and then closes the loop to avoid resource leaks.
133206
134207
Args:
135-
client (AsyncRedis): Asynchronous Redis client.
208+
coro (coroutine): The coroutine function to be run.
209+
*args: Positional arguments to pass to the coroutine function.
210+
**kwargs: Keyword arguments to pass to the coroutine function.
136211
137-
Raises:
138-
ValueError: If required Redis modules are not installed.
212+
Returns:
213+
The result of the coroutine if a new event loop is created,
214+
otherwise a task object representing the coroutine execution.
139215
"""
140-
# pick the right connection class
141-
connection_class: Type[AbstractConnection] = (
142-
SSLConnection
143-
if client.connection_pool.connection_class == ASSLConnection
144-
else Connection
145-
)
146-
# set up a temp sync client
147-
temp_client = Redis(
148-
connection_pool=ConnectionPool(
149-
connection_class=connection_class,
150-
**client.connection_pool.connection_kwargs,
151-
)
152-
)
153-
RedisConnectionFactory.validate_redis_modules(
154-
temp_client, redis_required_modules
155-
)
216+
try:
217+
# Try to get the current running event loop
218+
loop = asyncio.get_running_loop()
219+
except RuntimeError: # No running event loop
220+
loop = None
221+
222+
if loop and loop.is_running():
223+
# If an event loop is running, schedule the coroutine to run in the existing loop
224+
return asyncio.ensure_future(coro(*args, **kwargs))
225+
else:
226+
# No event loop is running, create a new event loop
227+
loop = asyncio.new_event_loop()
228+
asyncio.set_event_loop(loop)
229+
try:
230+
# Run the coroutine in the new event loop and wait for it to complete
231+
return loop.run_until_complete(coro(*args, **kwargs))
232+
finally:
233+
# Close the event loop to release resources
234+
loop.close()
156235

157236
@staticmethod
158-
def _validate_redis_modules(
237+
def _validate_modules(
159238
installed_modules, redis_required_modules: Optional[List[Dict[str, Any]]] = None
160239
) -> None:
161240
"""

tests/integration/test_connection.py

+48-8
Original file line numberDiff line numberDiff line change
@@ -6,13 +6,45 @@
66
from redis.exceptions import ConnectionError
77

88
from redisvl.redis.connection import RedisConnectionFactory, get_address_from_env
9+
from redisvl.version import __version__
10+
11+
EXPECTED_LIB_NAME = f"redis-py(redisvl_v{__version__})"
12+
13+
14+
def compare_versions(version1, version2):
15+
"""
16+
Compare two Redis version strings numerically.
17+
18+
Parameters:
19+
version1 (str): The first version string (e.g., "7.2.4").
20+
version2 (str): The second version string (e.g., "6.2.1").
21+
22+
Returns:
23+
int: -1 if version1 < version2, 0 if version1 == version2, 1 if version1 > version2.
24+
"""
25+
v1_parts = list(map(int, version1.split(".")))
26+
v2_parts = list(map(int, version2.split(".")))
27+
28+
for v1, v2 in zip(v1_parts, v2_parts):
29+
if v1 < v2:
30+
return False
31+
elif v1 > v2:
32+
return True
33+
34+
# If the versions are equal so far, compare the lengths of the version parts
35+
if len(v1_parts) < len(v2_parts):
36+
return False
37+
elif len(v1_parts) > len(v2_parts):
38+
return True
39+
40+
return True
941

1042

1143
def test_get_address_from_env(redis_url):
1244
assert get_address_from_env() == redis_url
1345

1446

15-
def test_sync_redis_connection(redis_url):
47+
def test_sync_redis_connect(redis_url):
1648
client = RedisConnectionFactory.connect(redis_url)
1749
assert client is not None
1850
assert isinstance(client, Redis)
@@ -21,7 +53,7 @@ def test_sync_redis_connection(redis_url):
2153

2254

2355
@pytest.mark.asyncio
24-
async def test_async_redis_connection(redis_url):
56+
async def test_async_redis_connect(redis_url):
2557
client = RedisConnectionFactory.connect(redis_url, use_async=True)
2658
assert client is not None
2759
assert isinstance(client, AsyncRedis)
@@ -49,11 +81,19 @@ def test_unknown_redis():
4981
bad_client.ping()
5082

5183

52-
def test_required_modules(client):
53-
RedisConnectionFactory.validate_redis_modules(client)
84+
def test_validate_redis(client):
85+
redis_version = client.info()["redis_version"]
86+
if not compare_versions(redis_version, "7.2.0"):
87+
pytest.skip("Not using a late enough version of Redis")
88+
RedisConnectionFactory.validate_redis(client)
89+
lib_name = client.client_info()
90+
assert lib_name["lib-name"] == EXPECTED_LIB_NAME
5491

5592

56-
@pytest.mark.asyncio
57-
async def test_async_required_modules(async_client):
58-
client = await async_client
59-
RedisConnectionFactory.validate_async_redis_modules(client)
93+
def test_validate_redis_custom_lib_name(client):
94+
redis_version = client.info()["redis_version"]
95+
if not compare_versions(redis_version, "7.2.0"):
96+
pytest.skip("Not using a late enough version of Redis")
97+
RedisConnectionFactory.validate_redis(client, "langchain_v0.1.0")
98+
lib_name = client.client_info()
99+
assert lib_name["lib-name"] == f"redis-py(redisvl_v{__version__};langchain_v0.1.0)"

0 commit comments

Comments
 (0)