diff --git a/litellm/caching/valkey_semantic_cache.py b/litellm/caching/valkey_semantic_cache.py index 645bd0892f4..02bb77c46b2 100644 --- a/litellm/caching/valkey_semantic_cache.py +++ b/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 diff --git a/litellm/llms/valkey/common_utils.py b/litellm/llms/valkey/common_utils.py index 9691450f3e0..703509904be 100644 --- a/litellm/llms/valkey/common_utils.py +++ b/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 diff --git a/litellm/llms/valkey/vector_stores/transformation.py b/litellm/llms/valkey/vector_stores/transformation.py index cd95d9c95b7..c40133d0764 100644 --- a/litellm/llms/valkey/vector_stores/transformation.py +++ b/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, diff --git a/tests/unit/caching/test_valkey_semantic_cache.py b/tests/unit/caching/test_valkey_semantic_cache.py index 749658784ac..935fe72f3f6 100644 --- a/tests/unit/caching/test_valkey_semantic_cache.py +++ b/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)" diff --git a/tests/unit/llms/valkey/vector_stores/test_valkey_transformation.py b/tests/unit/llms/valkey/vector_stores/test_valkey_transformation.py index 64e15008536..0f723f365cd 100644 --- a/tests/unit/llms/valkey/vector_stores/test_valkey_transformation.py +++ b/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)"