Merge 1feb1aa37d into a392a39627

litellm/caching/valkey_semantic_cache.py

@@ -30,7 +30,7 @@ from redis.commands.search.query import Query
 
 from litellm._logging import print_verbose
 from litellm._uuid import uuid
-from litellm.llms.valkey.common_utils import build_valkey_url, pack_vector
+from litellm.llms.valkey.common_utils import VALKEY_LIB_NAME, build_valkey_url, pack_vector
 
 from ._embedding_router import resolve_embedding_timeout
 from .redis_semantic_cache import RedisSemanticCache
@@ -96,8 +96,14 @@ class ValkeySemanticCache(RedisSemanticCache):
             self.async_client = async_client
         else:
             resolved_url: Final = redis_url or self._build_valkey_url(host, port, password, ssl)
-            self.sync_client = sync_client if sync_client is not None else Redis.from_url(resolved_url)
-            self.async_client = async_client if async_client is not None else AsyncRedis.from_url(resolved_url)
+            self.sync_client = (
+                sync_client if sync_client is not None else Redis.from_url(resolved_url, lib_name=VALKEY_LIB_NAME)  # pyright: ignore[reportCallIssue]  # types-redis 4.6 stubs shadow redis 5.3.1 and omit from_url(lib_name=...); redis-py accepts it at runtime
+            )
+            self.async_client = (
+                async_client
+                if async_client is not None
+                else AsyncRedis.from_url(resolved_url, lib_name=VALKEY_LIB_NAME)  # pyright: ignore[reportCallIssue]  # types-redis 4.6 stubs shadow redis 5.3.1 and omit from_url(lib_name=...); redis-py accepts it at runtime
+            )
 
         print_verbose(f"Valkey semantic-cache initializing index - {self.index_name}")
 
@@ -172,26 +178,26 @@ class ValkeySemanticCache(RedisSemanticCache):
         if self._index_dim == dim:
             return
         try:
-            self.sync_client.ft(self.index_name).create_index(
+            self.sync_client.ft(self.index_name).create_index(  # pyright: ignore[reportUnknownMemberType]  # the from_url(lib_name=...) overload mismatch leaves sync_client's type unknown; .ft/.create_index inherit that
                 self._index_schema(dim), definition=self._index_definition()
             )
         except Exception as exc:
             if not self._is_index_exists_error(exc):
                 raise
-            self._assert_dim_matches(self.sync_client.ft(self.index_name).info(), dim)
+            self._assert_dim_matches(self.sync_client.ft(self.index_name).info(), dim)  # pyright: ignore[reportUnknownMemberType]  # the from_url(lib_name=...) overload mismatch leaves sync_client's type unknown; .ft/.info inherit that
         self._index_dim = dim
 
     async def _ensure_index_async(self, dim: int) -> None:
         if self._index_dim == dim:
             return
         try:
-            await self.async_client.ft(self.index_name).create_index(
+            await self.async_client.ft(self.index_name).create_index(  # pyright: ignore[reportUnknownMemberType]  # the from_url(lib_name=...) overload mismatch leaves async_client's type unknown; .ft/.create_index inherit that
                 self._index_schema(dim), definition=self._index_definition()
             )
         except Exception as exc:
             if not self._is_index_exists_error(exc):
                 raise
-            info: Final = await self.async_client.ft(self.index_name).info()
+            info: Final = await self.async_client.ft(self.index_name).info()  # pyright: ignore[reportUnknownMemberType]  # the from_url(lib_name=...) overload mismatch leaves async_client's type unknown; .ft/.info inherit that
             self._assert_dim_matches(info, dim)
         self._index_dim = dim
 
@@ -218,7 +224,7 @@ class ValkeySemanticCache(RedisSemanticCache):
 
     async def _async_search(self, key: str, embedding: Sequence[float]) -> object:
         """Run the KNN query on the async client, stopping the untyped search surface here."""
-        return await self.async_client.ft(self.index_name).search(
+        return await self.async_client.ft(self.index_name).search(  # pyright: ignore[reportUnknownMemberType]  # the from_url(lib_name=...) overload mismatch leaves async_client's type unknown; .ft/.search inherit that
             self._knn_query(key),
             query_params={"vec": self._embedding_to_bytes(embedding)},  # pyright: ignore[reportArgumentType]  # redis stubs omit bytes; KNN vectors are raw bytes at runtime
         )
@@ -270,10 +276,10 @@ class ValkeySemanticCache(RedisSemanticCache):
             self._ensure_index_sync(len(embedding))
 
             doc_key: Final = self._doc_key(key)
-            self.sync_client.hset(doc_key, mapping=self._doc_mapping(key, prompt, str(value), embedding))
+            self.sync_client.hset(doc_key, mapping=self._doc_mapping(key, prompt, str(value), embedding))  # pyright: ignore[reportUnknownMemberType]  # the from_url(lib_name=...) overload mismatch leaves sync_client's type unknown; .hset inherits that
             ttl: Final = self._get_ttl(**kwargs)
             if ttl is not None:
-                self.sync_client.expire(doc_key, ttl)
+                self.sync_client.expire(doc_key, ttl)  # pyright: ignore[reportUnknownMemberType]  # the from_url(lib_name=...) overload mismatch leaves sync_client's type unknown; .expire inherits that
         except Exception as e:
             print_verbose(f"Error in Valkey semantic-cache set_cache: {e}")
 
@@ -288,7 +294,7 @@ class ValkeySemanticCache(RedisSemanticCache):
             embedding: Final = self._get_embedding(prompt)
             self._ensure_index_sync(len(embedding))
 
-            search_result: Final = self.sync_client.ft(self.index_name).search(
+            search_result: Final = self.sync_client.ft(self.index_name).search(  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]  # the from_url(lib_name=...) overload mismatch leaves sync_client's type unknown; .ft/.search and their result inherit that
                 self._knn_query(key),
                 query_params={"vec": self._embedding_to_bytes(embedding)},  # pyright: ignore[reportArgumentType]  # redis stubs omit bytes; KNN vectors are raw bytes at runtime
             )
@@ -309,10 +315,10 @@ class ValkeySemanticCache(RedisSemanticCache):
             await self._ensure_index_async(len(embedding))
 
             doc_key: Final = self._doc_key(key)
-            await self.async_client.hset(doc_key, mapping=self._doc_mapping(key, prompt, str(value), embedding))
+            await self.async_client.hset(doc_key, mapping=self._doc_mapping(key, prompt, str(value), embedding))  # pyright: ignore[reportUnknownMemberType]  # the from_url(lib_name=...) overload mismatch leaves async_client's type unknown; .hset inherits that
             ttl: Final = self._get_ttl(**kwargs)
             if ttl is not None:
-                await self.async_client.expire(doc_key, ttl)
+                await self.async_client.expire(doc_key, ttl)  # pyright: ignore[reportUnknownMemberType]  # the from_url(lib_name=...) overload mismatch leaves async_client's type unknown; .expire inherits that
         except Exception as e:
             print_verbose(f"Error in async Valkey semantic-cache set_cache: {e}")
 
@@ -340,4 +346,4 @@ class ValkeySemanticCache(RedisSemanticCache):
             print_verbose(f"Error in Valkey semantic-cache async_set_cache_pipeline: {e}")
 
     async def _index_info(self) -> Mapping[str, object]:
-        return await self.async_client.ft(self.index_name).info()
+        return await self.async_client.ft(self.index_name).info()  # pyright: ignore[reportUnknownMemberType]  # the from_url(lib_name=...) overload mismatch leaves async_client's type unknown; .ft/.info inherit that

litellm/llms/valkey/common_utils.py

@@ -5,6 +5,11 @@ from collections.abc import Sequence
 from typing import Final
 from urllib.parse import quote
 
+# Sent via CLIENT SETINFO LIB-NAME so server-side client listings can attribute
+# connections to LiteLLM. redis-py's lib_name replaces its default "redis-py",
+# so the base name is kept and the tag appended in the conventional base(tag) form.
+VALKEY_LIB_NAME: Final = "redis-py(litellm)"
+
 
 def build_valkey_url(host: str, port: str, password: str | None = None, ssl: bool = False) -> str:
     """Deliberately reads no environment: callers of the vector store control the

litellm/llms/valkey/vector_stores/transformation.py

@@ -19,7 +19,7 @@ from litellm.llms.base_llm.vector_store.transformation import (
     BaseDirectVectorStoreConfig,
     VectorStoreEmbeddingExecutor,
 )
-from litellm.llms.valkey.common_utils import build_valkey_url, pack_vector
+from litellm.llms.valkey.common_utils import VALKEY_LIB_NAME, build_valkey_url, pack_vector
 from litellm.types.llms.base import LiteLLMBaseModel
 from litellm.types.utils import EmbeddingResponse
 from litellm.types.vector_stores import (
@@ -245,16 +245,17 @@ class ValkeyVectorStoreConfig(BaseDirectVectorStoreConfig):
             return self._to_response(raw, query_text, params.text_field)
 
         connect_timeout, op_timeout = self._socket_timeouts(timeout)
-        client: Final = _import_sync_redis().from_url(
+        client: Final = _import_sync_redis().from_url(  # pyright: ignore[reportCallIssue, reportUnknownVariableType]  # types-redis 4.6 stubs shadow redis 5.3.1 and omit from_url(lib_name=...); redis-py accepts it at runtime, but the overload mismatch leaves the return type unknown
             params.connection_url(),
             socket_connect_timeout=connect_timeout,
             socket_timeout=op_timeout,
+            lib_name=VALKEY_LIB_NAME,
         )
         try:
-            raw_result: Final = client.ft(vector_store_id).search(knn, query_params=vec_params)
-            return self._to_response(raw_result, query_text, params.text_field)
+            raw_result: Final = client.ft(vector_store_id).search(knn, query_params=vec_params)  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]  # client's type is unknown (see from_url above), so .ft/.search and their result inherit that
+            return self._to_response(raw_result, query_text, params.text_field)  # pyright: ignore[reportUnknownArgumentType]  # raw_result flows from the from_url client whose return type the types-redis 4.6 stubs leave unknown; the shape is validated inside _to_response
         finally:
-            client.close()
+            client.close()  # pyright: ignore[reportUnknownMemberType]  # client's type is unknown (see from_url above)
 
     async def aexecute_search_vector_store_request(
         self,
@@ -295,18 +296,19 @@ class ValkeyVectorStoreConfig(BaseDirectVectorStoreConfig):
             return self._to_response(raw, query_text, params.text_field)
 
         connect_timeout, op_timeout = self._socket_timeouts(timeout)
-        client: Final = _import_async_redis().from_url(
+        client: Final = _import_async_redis().from_url(  # pyright: ignore[reportCallIssue, reportUnknownVariableType]  # types-redis 4.6 stubs shadow redis 5.3.1 and omit from_url(lib_name=...); redis-py accepts it at runtime, but the overload mismatch leaves the return type unknown
             params.connection_url(),
             socket_connect_timeout=connect_timeout,
             socket_timeout=op_timeout,
+            lib_name=VALKEY_LIB_NAME,
         )
         try:
-            raw_result: Final = await client.ft(vector_store_id).search(  # pyright: ignore[reportGeneralTypeIssues]  # types-redis 4.6 stubs shadow redis 5.3.1 and type the async client's ft() as the sync Search, so search() returns a non-awaitable Result; it is a coroutine at runtime
+            raw_result: Final = await client.ft(vector_store_id).search(  # pyright: ignore[reportGeneralTypeIssues, reportUnknownMemberType, reportUnknownVariableType]  # types-redis 4.6 stubs shadow redis 5.3.1 and type the async client's ft() as the sync Search, so search() returns a non-awaitable Result; it is a coroutine at runtime. client's type is unknown (see from_url above), so .ft/.search and their result inherit that
                 knn, query_params=vec_params
             )
-            return self._to_response(raw_result, query_text, params.text_field)
+            return self._to_response(raw_result, query_text, params.text_field)  # pyright: ignore[reportUnknownArgumentType]  # raw_result flows from the from_url client whose return type the types-redis 4.6 stubs leave unknown; the shape is validated inside _to_response
         finally:
-            await client.aclose()
+            await client.aclose()  # pyright: ignore[reportUnknownMemberType]  # client's type is unknown (see from_url above)
 
     def transform_create_vector_store_request(
         self,

tests/unit/caching/test_valkey_semantic_cache.py

@@ -534,3 +534,12 @@ def test_importing_caching_does_not_require_redis():
     )
     assert result.returncode == 0, result.stderr
     assert "ok" in result.stdout
+
+
+def test_self_built_clients_identify_as_litellm_via_lib_name():
+    # from_url builds a pool without connecting, so the kwargs redis-py will send
+    # on CLIENT SETINFO LIB-NAME can be read straight off the real clients.
+    cache = ValkeySemanticCache(similarity_threshold=0.8, redis_url="redis://valkey-host:6379")
+
+    assert cache.sync_client.connection_pool.connection_kwargs["lib_name"] == "redis-py(litellm)"
+    assert cache.async_client.connection_pool.connection_kwargs["lib_name"] == "redis-py(litellm)"

tests/unit/llms/valkey/vector_stores/test_valkey_transformation.py

@@ -398,3 +398,52 @@ def test_provider_config_manager_returns_valkey_config():
     config = ProviderConfigManager.get_provider_vector_stores_config(provider=LlmProviders.VALKEY, api_type=None)
 
     assert isinstance(config, ValkeyVectorStoreConfig)
+
+
+class _RecordingRedisFactory:
+    """Stands in for the redis-py client class so the search path's from_url kwargs can be asserted."""
+
+    def __init__(self, client):
+        self.client = client
+        self.from_url_kwargs = None
+
+    def from_url(self, url, **kwargs):
+        self.from_url_kwargs = kwargs
+        return self.client
+
+
+class _ClosableFakeRedis(FakeRedis):
+    def close(self):
+        pass
+
+
+class _ClosableFakeAsyncRedis(FakeAsyncRedis):
+    async def aclose(self):
+        pass
+
+
+def test_sync_search_identifies_connection_as_litellm_via_lib_name(monkeypatch):
+    factory = _RecordingRedisFactory(_ClosableFakeRedis())
+    monkeypatch.setattr("litellm.llms.valkey.vector_stores.transformation._import_sync_redis", lambda: factory)
+    config = ValkeyVectorStoreConfig(embedding_fn=FakeEmbeddingFn([1.0]))
+
+    _search(config, litellm_params={"valkey_host": "valkey-host"})
+
+    assert factory.from_url_kwargs["lib_name"] == "redis-py(litellm)"
+
+
+@pytest.mark.asyncio
+async def test_async_search_identifies_connection_as_litellm_via_lib_name(monkeypatch):
+    factory = _RecordingRedisFactory(_ClosableFakeAsyncRedis())
+    monkeypatch.setattr("litellm.llms.valkey.vector_stores.transformation._import_async_redis", lambda: factory)
+    config = ValkeyVectorStoreConfig(aembedding_fn=FakeAsyncEmbeddingFn([1.0]))
+
+    await config.aexecute_search_vector_store_request(
+        vector_store_id="my_index",
+        query="q",
+        vector_store_search_optional_params={},
+        litellm_logging_obj=MagicMock(),
+        litellm_params={"litellm_embedding_model": "openai/text-embedding-3-small", "valkey_host": "valkey-host"},
+    )
+
+    assert factory.from_url_kwargs["lib_name"] == "redis-py(litellm)"