diff --git a/litellm/caching/valkey_semantic_cache.py b/litellm/caching/valkey_semantic_cache.py index aa10d91fc66..cf6fab780fc 100644 --- a/litellm/caching/valkey_semantic_cache.py +++ b/litellm/caching/valkey_semantic_cache.py @@ -17,7 +17,6 @@ RedisSemanticCache since those are backend agnostic. import asyncio import hashlib import os -import struct from dataclasses import dataclass from typing import Any, Final @@ -29,6 +28,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 .redis_semantic_cache import RedisSemanticCache @@ -92,19 +92,17 @@ class ValkeySemanticCache(RedisSemanticCache): @staticmethod def _build_valkey_url(host: str | None, port: str | None, password: str | None, ssl: bool = False) -> str: - host = host or os.environ.get("VALKEY_HOST") or os.environ.get("REDIS_HOST") - port = port or os.environ.get("VALKEY_PORT") or os.environ.get("REDIS_PORT") - password = password or os.environ.get("VALKEY_PASSWORD") or os.environ.get("REDIS_PASSWORD") + resolved_host: Final = host or os.environ.get("VALKEY_HOST") or os.environ.get("REDIS_HOST") + resolved_port: Final = port or os.environ.get("VALKEY_PORT") or os.environ.get("REDIS_PORT") + resolved_password: Final = password or os.environ.get("VALKEY_PASSWORD") or os.environ.get("REDIS_PASSWORD") - if not host or not port: + if not resolved_host or not resolved_port: raise ValueError( "Missing required Valkey configuration. Provide host and port " "(or VALKEY_HOST/VALKEY_PORT), or pass redis_url." ) - credentials: Final = f":{password}@" if password else "" - scheme: Final = "rediss" if ssl else "redis" - return f"{scheme}://{credentials}{host}:{port}" + return build_valkey_url(host=resolved_host, port=resolved_port, password=resolved_password, ssl=ssl) @classmethod def _scope_tag(cls, key: str) -> str: @@ -116,7 +114,7 @@ class ValkeySemanticCache(RedisSemanticCache): @staticmethod def _embedding_to_bytes(embedding: list[float]) -> bytes: - return struct.pack(f"<{len(embedding)}f", *embedding) + return pack_vector(embedding) def _index_schema(self, dim: int) -> tuple[TagField, VectorField]: return ( diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index 8083d2485ba..723d1375e1e 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -1,5 +1,6 @@ from abc import abstractmethod -from typing import TYPE_CHECKING, Any +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Any, NoReturn import httpx @@ -154,3 +155,73 @@ class BaseVectorStoreConfig: response: VectorStoreSearchResponse, ) -> tuple[float, float]: return 0.0, 0.0 + + +class BaseDirectVectorStoreConfig(BaseVectorStoreConfig): + """ + Base config for vector store providers whose datastore has no HTTP API + (e.g. Valkey over RESP). Instead of transforming to an httpx request, the + config executes the search itself via (a)execute_search_vector_store_request. + """ + + @abstractmethod + def execute_search_vector_store_request( + self, + vector_store_id: str, + query: str | Sequence[str], + vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + litellm_logging_obj: LiteLLMLoggingObj, + litellm_params: Mapping[str, object], + ) -> VectorStoreSearchResponse: + pass + + @abstractmethod + async def aexecute_search_vector_store_request( + self, + vector_store_id: str, + query: str | Sequence[str], + vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + litellm_logging_obj: LiteLLMLoggingObj, + litellm_params: Mapping[str, object], + ) -> VectorStoreSearchResponse: + pass + + def transform_search_vector_store_request( + self, + vector_store_id: str, + query: str | Sequence[str], + vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + api_base: str, + litellm_logging_obj: LiteLLMLoggingObj, + litellm_params: Mapping[str, object], + extra_body: Mapping[str, object] | None = None, + ) -> NoReturn: + raise NotImplementedError("Direct vector store providers execute the search themselves; no HTTP request shape") + + def transform_search_vector_store_response( + self, response: httpx.Response, litellm_logging_obj: LiteLLMLoggingObj + ) -> NoReturn: + raise NotImplementedError("Direct vector store providers execute the search themselves; no HTTP response shape") + + def transform_create_vector_store_request( + self, + vector_store_create_optional_params: VectorStoreCreateOptionalRequestParams, + api_base: str, + ) -> NoReturn: + raise NotImplementedError + + def transform_create_vector_store_response(self, response: httpx.Response) -> NoReturn: + raise NotImplementedError + + def get_complete_url( + self, + api_base: str | None, + litellm_params: Mapping[str, object], + ) -> str: + return api_base or "" + + def get_auth_credentials(self, litellm_params: Mapping[str, object]) -> BaseVectorStoreAuthCredentials: + return BaseVectorStoreAuthCredentials() + + def get_vector_store_endpoints_by_type(self) -> VectorStoreIndexEndpoints: + return VectorStoreIndexEndpoints(read=[], write=[]) # mutable-ok: the TypedDict declares list fields diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 721b9545ac1..6e5825b5ca7 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -55,7 +55,10 @@ from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfi from litellm.llms.base_llm.search.transformation import BaseSearchConfig, SearchResponse from litellm.llms.base_llm.skills.transformation import BaseSkillsAPIConfig from litellm.llms.base_llm.text_to_speech.transformation import BaseTextToSpeechConfig -from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig +from litellm.llms.base_llm.vector_store.transformation import ( + BaseDirectVectorStoreConfig, + BaseVectorStoreConfig, +) from litellm.llms.base_llm.vector_store_files.transformation import ( BaseVectorStoreFilesConfig, ) @@ -9411,6 +9414,23 @@ class BaseLLMHTTPHandler: client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, ) -> VectorStoreSearchResponse: + if isinstance(vector_store_provider_config, BaseDirectVectorStoreConfig): + logging_obj.pre_call( + input="", + api_key="", + additional_args={ # mutable-ok: pre_call's additional_args contract is a dict + "query": query, + "vector_store_id": vector_store_id, + }, + ) + return await vector_store_provider_config.aexecute_search_vector_store_request( + vector_store_id=vector_store_id, + query=query, + vector_store_search_optional_params=vector_store_search_optional_params, + litellm_logging_obj=logging_obj, + litellm_params=dict(litellm_params), # mutable-ok: snapshot GenericLiteLLMParams into the Mapping shape + ) + if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( llm_provider=litellm.LlmProviders(custom_llm_provider), @@ -9524,6 +9544,23 @@ class BaseLLMHTTPHandler: client=client, ) + if isinstance(vector_store_provider_config, BaseDirectVectorStoreConfig): + logging_obj.pre_call( + input="", + api_key="", + additional_args={ # mutable-ok: pre_call's additional_args contract is a dict + "query": query, + "vector_store_id": vector_store_id, + }, + ) + return vector_store_provider_config.execute_search_vector_store_request( + vector_store_id=vector_store_id, + query=query, + vector_store_search_optional_params=vector_store_search_optional_params, + litellm_logging_obj=logging_obj, + litellm_params=dict(litellm_params), # mutable-ok: snapshot GenericLiteLLMParams into the Mapping shape + ) + if client is None or not isinstance(client, HTTPHandler): sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: diff --git a/litellm/llms/valkey/__init__.py b/litellm/llms/valkey/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/valkey/common_utils.py b/litellm/llms/valkey/common_utils.py new file mode 100644 index 00000000000..9691450f3e0 --- /dev/null +++ b/litellm/llms/valkey/common_utils.py @@ -0,0 +1,18 @@ +"""Shared helpers for Valkey integrations (semantic cache, vector stores).""" + +import struct +from collections.abc import Sequence +from typing import Final +from urllib.parse import quote + + +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 + host, so an env-sourced password would be sent to a caller-chosen server.""" + credentials: Final = f":{quote(password, safe='')}@" if password else "" + scheme: Final = "rediss" if ssl else "redis" + return f"{scheme}://{credentials}{host}:{port}" + + +def pack_vector(embedding: Sequence[float]) -> bytes: + return struct.pack(f"<{len(embedding)}f", *embedding) diff --git a/litellm/llms/valkey/vector_stores/__init__.py b/litellm/llms/valkey/vector_stores/__init__.py new file mode 100644 index 00000000000..c826607a800 --- /dev/null +++ b/litellm/llms/valkey/vector_stores/__init__.py @@ -0,0 +1,3 @@ +from litellm.llms.valkey.vector_stores.transformation import ValkeyVectorStoreConfig + +__all__ = ("ValkeyVectorStoreConfig",) diff --git a/litellm/llms/valkey/vector_stores/transformation.py b/litellm/llms/valkey/vector_stores/transformation.py new file mode 100644 index 00000000000..67ce02d4c69 --- /dev/null +++ b/litellm/llms/valkey/vector_stores/transformation.py @@ -0,0 +1,274 @@ +""" +Valkey vector store provider. + +Valkey's vector search (the valkey-search module) speaks RESP only, no HTTP +API, so this config extends BaseDirectVectorStoreConfig and executes the +FT.SEARCH KNN query itself via redis-py instead of shaping an httpx request. +Documents are HASHes indexed by an FT index named after the vector_store_id. +""" + +from collections.abc import Awaitable, Callable, Mapping, Sequence +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, NoReturn + +import httpx +from pydantic import BaseModel, ConfigDict + +import litellm +from litellm.llms.base_llm.vector_store.transformation import BaseDirectVectorStoreConfig +from litellm.llms.valkey.common_utils import build_valkey_url, pack_vector +from litellm.types.utils import EmbeddingResponse +from litellm.types.vector_stores import ( + VectorStoreCreateOptionalRequestParams, + VectorStoreResultContent, + VectorStoreSearchOptionalRequestParams, + VectorStoreSearchResponse, + VectorStoreSearchResult, +) + +if TYPE_CHECKING: + from redis import Redis + from redis.asyncio import Redis as AsyncRedis + from redis.commands.search.document import Document + from redis.commands.search.query import Query + from redis.commands.search.result import Result + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + +DEFAULT_VALKEY_PORT: Final = 6379 +DEFAULT_MAX_NUM_RESULTS: Final = 10 +MIN_MAX_NUM_RESULTS: Final = 1 +MAX_MAX_NUM_RESULTS: Final = 50 +DEFAULT_EMBEDDING_FIELD_NAME: Final = "embedding" +DEFAULT_TEXT_FIELD_NAME: Final = "text" +DISTANCE_FIELD_NAME: Final = "vector_distance" + +_EMPTY_EMBEDDING_CONFIG: Final = MappingProxyType({}) +_REDIS_INSTALL_HINT: Final = ( + "The Valkey vector store requires the 'redis' package. Run 'pip install redis' to install it." +) +_SEARCH_ONLY_MESSAGE: Final = "Valkey vector store is search-only; create indexes with FT.CREATE directly" + + +def _import_sync_redis() -> "type[Redis]": + try: + from redis import Redis as SyncRedisClient + except ImportError as e: + raise ValueError(_REDIS_INSTALL_HINT) from e + return SyncRedisClient + + +def _import_async_redis() -> "type[AsyncRedis]": + try: + from redis.asyncio import Redis as AsyncRedisClient + except ImportError as e: + raise ValueError(_REDIS_INSTALL_HINT) from e + return AsyncRedisClient + + +def _import_query() -> "type[Query]": + try: + from redis.commands.search.query import Query as RedisQuery + except ImportError as e: + raise ValueError(_REDIS_INSTALL_HINT) from e + return RedisQuery + + +class _ValkeySearchParams(BaseModel): + """Typed view over the vector store's litellm_params; unrelated keys are ignored.""" + + model_config = ConfigDict(frozen=True, extra="ignore") + + litellm_embedding_model: str | None = None + litellm_embedding_config: Mapping[str, object] | None = None + valkey_host: str | None = None + valkey_port: int | None = None + valkey_password: str | None = None + valkey_ssl: bool | None = None + valkey_text_field: str | None = None + valkey_embedding_field: str | None = None + + @property + def text_field(self) -> str: + return self.valkey_text_field or DEFAULT_TEXT_FIELD_NAME + + @property + def embedding_field(self) -> str: + return self.valkey_embedding_field or DEFAULT_EMBEDDING_FIELD_NAME + + def require_embedding_model(self) -> str: + if not self.litellm_embedding_model: + raise ValueError( + "litellm_embedding_model is required in litellm_params for the Valkey vector store. " + "Example: litellm_params['litellm_embedding_model'] = 'openai/text-embedding-3-small'" + ) + return self.litellm_embedding_model + + def connection_url(self) -> str: + if not self.valkey_host: + raise ValueError( + "valkey_host is required in litellm_params for the Valkey vector store. " + "Set it on the vector store's litellm_params, e.g. valkey_host: my-valkey.example.com" + ) + return build_valkey_url( + host=self.valkey_host, + port=str(self.valkey_port or DEFAULT_VALKEY_PORT), + password=self.valkey_password, + ssl=bool(self.valkey_ssl), + ) + + +class ValkeyVectorStoreConfig(BaseDirectVectorStoreConfig): + def __init__( + self, + sync_client: "Redis | None" = None, + async_client: "AsyncRedis | None" = None, + embedding_fn: Callable[..., EmbeddingResponse] | None = None, + aembedding_fn: Callable[..., Awaitable[EmbeddingResponse]] | None = None, + ) -> None: + super().__init__() + self.sync_client = sync_client + self.async_client = async_client + self.embedding_fn = embedding_fn if embedding_fn is not None else litellm.embedding + self.aembedding_fn = aembedding_fn if aembedding_fn is not None else litellm.aembedding + + @staticmethod + def _query_text(query: str | Sequence[str]) -> str: + if isinstance(query, str): + return query + if not query: + raise ValueError("query must not be empty") + return query[0] + + @staticmethod + def _knn_limit(vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams) -> int: + requested: Final = vector_store_search_optional_params.get("max_num_results") + if requested is None: + return DEFAULT_MAX_NUM_RESULTS + if not MIN_MAX_NUM_RESULTS <= requested <= MAX_MAX_NUM_RESULTS: + raise ValueError( + f"max_num_results must be between {MIN_MAX_NUM_RESULTS} and {MAX_MAX_NUM_RESULTS}, got {requested}" + ) + return requested + + @classmethod + def _knn_query( + cls, + vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + embedding_field: str, + text_field: str, + ) -> "Query": + if vector_store_search_optional_params.get("filters") is not None: + raise ValueError("Valkey vector store does not support the filters parameter yet") + k: Final = cls._knn_limit(vector_store_search_optional_params) + query_cls: Final = _import_query() + knn_expr: Final = f"*=>[KNN {k} @{embedding_field} $vec AS {DISTANCE_FIELD_NAME}]" + # valkey-search rejects SORTBY on the KNN distance alias, so results are + # re-ordered client-side in _to_response instead. + return query_cls(knn_expr).return_fields(text_field, DISTANCE_FIELD_NAME).paging(0, k).dialect(2) + + @staticmethod + def _to_result(doc: "Document", text_field: str) -> VectorStoreSearchResult: + content: Final = [ # mutable-ok: VectorStoreSearchResult declares a list of content parts + VectorStoreResultContent(text=str(getattr(doc, text_field, "")), type="text") + ] + return VectorStoreSearchResult( + score=1.0 - float(getattr(doc, DISTANCE_FIELD_NAME)), + content=content, + file_id=getattr(doc, "id", None), + filename=getattr(doc, "id", None), + ) + + @classmethod + def _to_response(cls, search_result: "Result", query_text: str, text_field: str) -> VectorStoreSearchResponse: + docs: Final = getattr(search_result, "docs", None) or () + data: Final = sorted( + (cls._to_result(doc, text_field) for doc in docs), + key=lambda result: result.get("score") or 0.0, + reverse=True, + ) + return VectorStoreSearchResponse( + object="vector_store.search_results.page", + search_query=query_text, + data=data, + ) + + def execute_search_vector_store_request( + self, + vector_store_id: str, + query: str | Sequence[str], + vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + litellm_logging_obj: "LiteLLMLoggingObj", + litellm_params: Mapping[str, object], + ) -> VectorStoreSearchResponse: + params: Final = _ValkeySearchParams.model_validate(litellm_params) + query_text: Final = self._query_text(query) + knn: Final = self._knn_query( + vector_store_search_optional_params, + embedding_field=params.embedding_field, + text_field=params.text_field, + ) + embedding_response: Final = self.embedding_fn( + model=params.require_embedding_model(), + input=[query_text], # mutable-ok: litellm.embedding's input contract is a list + **(params.litellm_embedding_config or _EMPTY_EMBEDDING_CONFIG), + ) + vec_params: Final = {"vec": pack_vector(embedding_response.data[0]["embedding"])} # mutable-ok: redis-py API + + if self.sync_client is not None: + raw: Final = self.sync_client.ft(vector_store_id).search(knn, query_params=vec_params) + return self._to_response(raw, query_text, params.text_field) + + client: Final = _import_sync_redis().from_url(params.connection_url()) + 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) + finally: + client.close() + + async def aexecute_search_vector_store_request( + self, + vector_store_id: str, + query: str | Sequence[str], + vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + litellm_logging_obj: "LiteLLMLoggingObj", + litellm_params: Mapping[str, object], + ) -> VectorStoreSearchResponse: + params: Final = _ValkeySearchParams.model_validate(litellm_params) + query_text: Final = self._query_text(query) + knn: Final = self._knn_query( + vector_store_search_optional_params, + embedding_field=params.embedding_field, + text_field=params.text_field, + ) + embedding_response: Final = await self.aembedding_fn( + model=params.require_embedding_model(), + input=[query_text], # mutable-ok: litellm.embedding's input contract is a list + **(params.litellm_embedding_config or _EMPTY_EMBEDDING_CONFIG), + ) + vec_params: Final = {"vec": pack_vector(embedding_response.data[0]["embedding"])} # mutable-ok: redis-py API + + if self.async_client is not None: + raw: Final = await self.async_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 + knn, query_params=vec_params + ) + return self._to_response(raw, query_text, params.text_field) + + client: Final = _import_async_redis().from_url(params.connection_url()) + 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 + knn, query_params=vec_params + ) + return self._to_response(raw_result, query_text, params.text_field) + finally: + await client.aclose() + + def transform_create_vector_store_request( + self, + vector_store_create_optional_params: VectorStoreCreateOptionalRequestParams, + api_base: str, + ) -> NoReturn: + raise NotImplementedError(_SEARCH_ONLY_MESSAGE) + + def transform_create_vector_store_response(self, response: httpx.Response) -> NoReturn: + raise NotImplementedError(_SEARCH_ONLY_MESSAGE) diff --git a/litellm/proxy/_experimental/out/assets/logos/valkey.svg b/litellm/proxy/_experimental/out/assets/logos/valkey.svg new file mode 100644 index 00000000000..0e97e680df4 --- /dev/null +++ b/litellm/proxy/_experimental/out/assets/logos/valkey.svg @@ -0,0 +1,6 @@ + + + + + + diff --git a/litellm/types/router.py b/litellm/types/router.py index 217364c48b7..11a64fb9e85 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -349,6 +349,12 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): milvus_text_field: str | None = None milvus_db_name: str | None = None milvus_partition_names: list[str] | None = None + valkey_host: str | None = None + valkey_port: int | None = None + valkey_password: str | None = None + valkey_ssl: bool | None = None + valkey_text_field: str | None = None + valkey_embedding_field: str | None = None @model_validator(mode="before") @classmethod diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 220826ccbca..c9c4a213a47 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3680,6 +3680,7 @@ class LlmProviders(str, Enum): NSCALE = "nscale" PG_VECTOR = "pg_vector" S3_VECTORS = "s3_vectors" + VALKEY = "valkey" HELICONE = "helicone" HYPERBOLIC = "hyperbolic" RECRAFT = "recraft" diff --git a/litellm/utils.py b/litellm/utils.py index f011c1ff62f..836a3f7b3f2 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8734,6 +8734,12 @@ class ProviderConfigManager: ) return S3VectorsVectorStoreConfig() + elif litellm.LlmProviders.VALKEY == provider: + from litellm.llms.valkey.vector_stores.transformation import ( + ValkeyVectorStoreConfig, + ) + + return ValkeyVectorStoreConfig() return None @staticmethod diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 0712e8e383d..ec0b1c27344 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -2809,6 +2809,13 @@ "vector_stores_search": true } }, + "valkey": { + "display_name": "Valkey (`valkey`)", + "url": "https://docs.litellm.ai/docs/providers/valkey_vector_stores", + "endpoints": { + "vector_stores_search": true + } + }, "helicone": { "display_name": "Helicone (`helicone`)", "url": "https://docs.litellm.ai/docs/providers/helicone", diff --git a/tests/documentation_tests/test_readme_providers.py b/tests/documentation_tests/test_readme_providers.py index f9de25bc85b..d3b4e22180b 100644 --- a/tests/documentation_tests/test_readme_providers.py +++ b/tests/documentation_tests/test_readme_providers.py @@ -16,6 +16,7 @@ EXCLUDED_PROVIDERS = { "langfuse", # observability, not LLM provider "humanloop", # observability, not LLM provider "pg_vector", # database, not LLM provider + "valkey", # database, not LLM provider "dotprompt", # prompt management, not provider "vertex_ai_beta", # beta variant, not needed in main table } diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index fddd8d09dfc..b5bd5ed3356 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -2071,3 +2071,86 @@ async def test_anthropic_invalid_thinking_signature_retry_resigns_bedrock_reques retry_authorization = posts[1]["headers"]["Authorization"] assert retry_authorization.startswith("AWS4-HMAC-SHA256") assert retry_authorization != first_attempt_headers["Authorization"] + + +def _make_stub_direct_vector_store_config(response): + from litellm.llms.base_llm.vector_store.transformation import ( + BaseDirectVectorStoreConfig, + ) + + class StubDirectVectorStoreConfig(BaseDirectVectorStoreConfig): + def __init__(self): + super().__init__() + self.sync_calls = [] + self.async_calls = [] + + def execute_search_vector_store_request(self, **kwargs): + self.sync_calls.append(kwargs) + return response + + async def aexecute_search_vector_store_request(self, **kwargs): + self.async_calls.append(kwargs) + return response + + return StubDirectVectorStoreConfig() + + +def test_vector_store_search_handler_direct_config_sync_skips_http(): + handler = BaseLLMHTTPHandler() + stub_response = {"object": "vector_store.search_results.page", "search_query": "q", "data": []} + config = _make_stub_direct_vector_store_config(stub_response) + logging_obj = Mock() + + with patch("litellm.llms.custom_httpx.llm_http_handler._get_httpx_client") as mock_get_client: + result = handler.vector_store_search_handler( + vector_store_id="vs_direct", + query="q", + vector_store_search_optional_params={"max_num_results": 4}, + vector_store_provider_config=config, + custom_llm_provider="valkey", + litellm_params=GenericLiteLLMParams(valkey_host="localhost"), + logging_obj=logging_obj, + _is_async=False, + ) + + assert result is stub_response + mock_get_client.assert_not_called() + assert len(config.sync_calls) == 1 + call = config.sync_calls[0] + assert call["vector_store_id"] == "vs_direct" + assert call["query"] == "q" + assert call["vector_store_search_optional_params"] == {"max_num_results": 4} + assert isinstance(call["litellm_params"], dict) + assert call["litellm_params"]["valkey_host"] == "localhost" + pre_call_args = logging_obj.pre_call.call_args.kwargs["additional_args"] + assert pre_call_args["query"] == "q" + assert pre_call_args["vector_store_id"] == "vs_direct" + + +@pytest.mark.asyncio +async def test_vector_store_search_handler_direct_config_async_skips_http(): + handler = BaseLLMHTTPHandler() + stub_response = {"object": "vector_store.search_results.page", "search_query": "q", "data": []} + config = _make_stub_direct_vector_store_config(stub_response) + logging_obj = Mock() + + with patch("litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client") as mock_get_client: + result = await handler.vector_store_search_handler( + vector_store_id="vs_direct", + query=["q1", "q2"], + vector_store_search_optional_params={}, + vector_store_provider_config=config, + custom_llm_provider="valkey", + litellm_params=GenericLiteLLMParams(valkey_host="localhost"), + logging_obj=logging_obj, + _is_async=True, + ) + + assert result is stub_response + mock_get_client.assert_not_called() + assert len(config.async_calls) == 1 + assert config.async_calls[0]["query"] == ["q1", "q2"] + assert config.async_calls[0]["litellm_params"]["valkey_host"] == "localhost" + pre_call_args = logging_obj.pre_call.call_args.kwargs["additional_args"] + assert pre_call_args["query"] == ["q1", "q2"] + assert pre_call_args["vector_store_id"] == "vs_direct" diff --git a/tests/test_litellm/llms/valkey/vector_stores/test_valkey_transformation.py b/tests/test_litellm/llms/valkey/vector_stores/test_valkey_transformation.py new file mode 100644 index 00000000000..188a1982b8b --- /dev/null +++ b/tests/test_litellm/llms/valkey/vector_stores/test_valkey_transformation.py @@ -0,0 +1,360 @@ +import struct +import sys +from types import SimpleNamespace +from typing import Final +from unittest.mock import MagicMock, patch +from urllib.parse import unquote, urlsplit + +import pytest + +from litellm.llms.valkey.vector_stores.transformation import ( + ValkeyVectorStoreConfig, + _ValkeySearchParams, +) +from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager + + +class FakeSearchIndex: + def __init__(self, result): + self.result = result + self.searched_query = None + self.searched_query_params = None + + def search(self, query, query_params=None): + self.searched_query = query + self.searched_query_params = query_params + return self.result + + +class FakeRedis: + def __init__(self, result=None): + self.index = FakeSearchIndex(result if result is not None else SimpleNamespace(docs=[])) + self.ft_index_name = None + + def ft(self, index_name): + self.ft_index_name = index_name + return self.index + + +class FakeAsyncSearchIndex(FakeSearchIndex): + async def search(self, query, query_params=None): + self.searched_query = query + self.searched_query_params = query_params + return self.result + + +class FakeAsyncRedis(FakeRedis): + def __init__(self, result=None): + super().__init__(result) + self.index = FakeAsyncSearchIndex(self.index.result) + + +class FakeEmbeddingFn: + def __init__(self, embedding): + self.embedding = embedding + self.captured_kwargs = None + + def __call__(self, **kwargs): + self.captured_kwargs = kwargs + return SimpleNamespace(data=[{"embedding": self.embedding}]) + + +class FakeAsyncEmbeddingFn(FakeEmbeddingFn): + async def __call__(self, **kwargs): + self.captured_kwargs = kwargs + return SimpleNamespace(data=[{"embedding": self.embedding}]) + + +def _doc(doc_id, distance, **fields): + return SimpleNamespace(id=doc_id, vector_distance=str(distance), **fields) + + +def _search(config, client=None, query="what is litellm", optional_params=None, litellm_params=None): + return config.execute_search_vector_store_request( + vector_store_id="my_index", + query=query, + vector_store_search_optional_params=optional_params or {}, + litellm_logging_obj=MagicMock(), + litellm_params={"litellm_embedding_model": "openai/text-embedding-3-small", **(litellm_params or {})}, + ) + + +def test_sync_search_builds_knn_query_with_packed_vector(): + embedding_fn = FakeEmbeddingFn([0.1, 0.2, 0.3]) + client = FakeRedis() + config = ValkeyVectorStoreConfig(sync_client=client, embedding_fn=embedding_fn) + + _search(config, optional_params={"max_num_results": 5}) + + assert client.ft_index_name == "my_index" + assert client.index.searched_query.query_string() == "*=>[KNN 5 @embedding $vec AS vector_distance]" + args = client.index.searched_query.get_args() + assert args[args.index("DIALECT") + 1] == 2 + assert args[args.index("LIMIT") : args.index("LIMIT") + 3] == ["LIMIT", 0, 5] + return_args = args[args.index("RETURN") : args.index("RETURN") + 4] + assert return_args == ["RETURN", 2, "text", "vector_distance"] + assert client.index.searched_query_params == {"vec": struct.pack("<3f", 0.1, 0.2, 0.3)} + + +def test_sync_search_defaults_to_10_results(): + client = FakeRedis() + config = ValkeyVectorStoreConfig(sync_client=client, embedding_fn=FakeEmbeddingFn([1.0])) + + _search(config) + + assert client.index.searched_query.query_string() == "*=>[KNN 10 @embedding $vec AS vector_distance]" + + +def test_sync_search_honors_custom_field_names(): + client = FakeRedis(result=SimpleNamespace(docs=[_doc("doc:1", 0.5, chunk="custom text")])) + config = ValkeyVectorStoreConfig(sync_client=client, embedding_fn=FakeEmbeddingFn([1.0])) + + response = _search( + config, + litellm_params={"valkey_embedding_field": "emb", "valkey_text_field": "chunk"}, + ) + + assert client.index.searched_query.query_string() == "*=>[KNN 10 @emb $vec AS vector_distance]" + assert "chunk" in client.index.searched_query.get_args() + assert response["data"][0]["content"][0]["text"] == "custom text" + + +def test_sync_search_maps_response_with_inverted_score_sorted_best_first(): + client = FakeRedis( + result=SimpleNamespace(docs=[_doc("doc:2", 0.75, text="bye"), _doc("doc:1", 0.25, text="hello world")]) + ) + config = ValkeyVectorStoreConfig(sync_client=client, embedding_fn=FakeEmbeddingFn([1.0])) + + response = _search(config) + + assert response["object"] == "vector_store.search_results.page" + assert response["search_query"] == "what is litellm" + assert response["data"][0]["score"] == pytest.approx(0.75) + assert response["data"][0]["content"] == [{"text": "hello world", "type": "text"}] + assert response["data"][0]["file_id"] == "doc:1" + assert response["data"][0]["filename"] == "doc:1" + assert response["data"][1]["score"] == pytest.approx(0.25) + assert response["data"][1]["file_id"] == "doc:2" + + +def test_sync_search_list_query_takes_first_element(): + embedding_fn = FakeEmbeddingFn([1.0]) + config = ValkeyVectorStoreConfig(sync_client=FakeRedis(), embedding_fn=embedding_fn) + + response = _search(config, query=["first query", "second query"]) + + assert embedding_fn.captured_kwargs["input"] == ["first query"] + assert response["search_query"] == "first query" + + +def test_sync_search_expands_embedding_config_into_kwargs(): + embedding_fn = FakeEmbeddingFn([1.0]) + config = ValkeyVectorStoreConfig(sync_client=FakeRedis(), embedding_fn=embedding_fn) + + _search( + config, + litellm_params={"litellm_embedding_config": {"api_key": "sk-test", "api_base": "https://embed.example.com"}}, + ) + + assert embedding_fn.captured_kwargs == { + "model": "openai/text-embedding-3-small", + "input": ["what is litellm"], + "api_key": "sk-test", + "api_base": "https://embed.example.com", + } + + +def test_sync_search_requires_embedding_model(): + config = ValkeyVectorStoreConfig(sync_client=FakeRedis(), embedding_fn=FakeEmbeddingFn([1.0])) + + with pytest.raises(ValueError, match="litellm_embedding_model is required"): + config.execute_search_vector_store_request( + vector_store_id="my_index", + query="q", + vector_store_search_optional_params={}, + litellm_logging_obj=MagicMock(), + litellm_params={}, + ) + + +def test_sync_search_requires_valkey_host_without_injected_client(monkeypatch): + monkeypatch.delenv("VALKEY_HOST", raising=False) + monkeypatch.delenv("REDIS_HOST", raising=False) + config = ValkeyVectorStoreConfig(embedding_fn=FakeEmbeddingFn([1.0])) + + with pytest.raises(ValueError, match="valkey_host is required"): + _search(config) + + +_VALKEY_ENV_VARS: Final = ( + "VALKEY_HOST", + "VALKEY_PORT", + "VALKEY_PASSWORD", + "REDIS_HOST", + "REDIS_PORT", + "REDIS_PASSWORD", +) + + +def test_connection_url_building(monkeypatch): + for var in _VALKEY_ENV_VARS: + monkeypatch.delenv(var, raising=False) + + full: Final = _ValkeySearchParams.model_validate( + {"valkey_host": "h", "valkey_port": 6380, "valkey_password": "p", "valkey_ssl": True} + ) + assert full.connection_url() == "rediss://:p@h:6380" + minimal: Final = _ValkeySearchParams.model_validate({"valkey_host": "h", "valkey_password": ""}) + assert minimal.connection_url() == "redis://h:6379" + + +def test_connection_url_never_borrows_gateway_credentials_from_the_environment(monkeypatch): + monkeypatch.setenv("VALKEY_HOST", "gateway-valkey.internal") + monkeypatch.setenv("VALKEY_PORT", "6380") + monkeypatch.setenv("VALKEY_PASSWORD", "gateway-secret") + monkeypatch.setenv("REDIS_HOST", "gateway-redis.internal") + monkeypatch.setenv("REDIS_PORT", "6381") + monkeypatch.setenv("REDIS_PASSWORD", "gateway-redis-secret") + + caller_controlled: Final = _ValkeySearchParams.model_validate({"valkey_host": "attacker.example.com"}) + + assert caller_controlled.connection_url() == "redis://attacker.example.com:6379" + + +def test_connection_url_percent_encodes_the_password(): + password: Final = "p@ss/w#rd%1:x" + params: Final = _ValkeySearchParams.model_validate({"valkey_host": "h", "valkey_password": password}) + + parsed: Final = urlsplit(params.connection_url()) + + assert parsed.hostname == "h" + assert parsed.port == 6379 + assert unquote(parsed.password or "") == password + + +def test_connection_url_accepts_string_booleans_from_the_ui_select(): + params: Final = _ValkeySearchParams.model_validate( + {"valkey_host": "h", "valkey_port": "6380", "valkey_ssl": "true"} + ) + + assert params.connection_url() == "rediss://h:6380" + assert _ValkeySearchParams.model_validate({"valkey_host": "h", "valkey_ssl": "false"}).connection_url() == ( + "redis://h:6379" + ) + + +def test_search_rejects_filters(): + embedding_fn = FakeEmbeddingFn([1.0]) + config = ValkeyVectorStoreConfig(sync_client=FakeRedis(), embedding_fn=embedding_fn) + + with pytest.raises(ValueError, match="does not support the filters parameter"): + _search(config, optional_params={"filters": {"category": "docs"}}) + + assert embedding_fn.captured_kwargs is None + + +@pytest.mark.asyncio +async def test_async_search_rejects_filters(): + aembedding_fn = FakeAsyncEmbeddingFn([1.0]) + config = ValkeyVectorStoreConfig(async_client=FakeAsyncRedis(), aembedding_fn=aembedding_fn) + + with pytest.raises(ValueError, match="does not support the filters parameter"): + await config.aexecute_search_vector_store_request( + vector_store_id="my_index", + query="q", + vector_store_search_optional_params={"filters": {"category": "docs"}}, + litellm_logging_obj=MagicMock(), + litellm_params={"litellm_embedding_model": "openai/text-embedding-3-small"}, + ) + + assert aembedding_fn.captured_kwargs is None + + +def test_search_rejects_empty_query(): + config = ValkeyVectorStoreConfig(sync_client=FakeRedis(), embedding_fn=FakeEmbeddingFn([1.0])) + + with pytest.raises(ValueError, match="query must not be empty"): + _search(config, query=[]) + + +@pytest.mark.parametrize("max_num_results", [0, -1, 51]) +def test_search_rejects_out_of_range_max_num_results(max_num_results): + embedding_fn = FakeEmbeddingFn([1.0]) + config = ValkeyVectorStoreConfig(sync_client=FakeRedis(), embedding_fn=embedding_fn) + + with pytest.raises(ValueError, match="max_num_results must be between 1 and 50"): + _search(config, optional_params={"max_num_results": max_num_results}) + + assert embedding_fn.captured_kwargs is None + + +def test_search_allows_max_num_results_at_the_upper_bound(): + client = FakeRedis() + config = ValkeyVectorStoreConfig(sync_client=client, embedding_fn=FakeEmbeddingFn([1.0])) + + _search(config, optional_params={"max_num_results": 50}) + + assert client.index.searched_query.query_string() == "*=>[KNN 50 @embedding $vec AS vector_distance]" + + +def test_search_treats_an_explicit_null_max_num_results_as_the_default(): + client = FakeRedis() + config = ValkeyVectorStoreConfig(sync_client=client, embedding_fn=FakeEmbeddingFn([1.0])) + + _search(config, optional_params={"max_num_results": None}) + + assert client.index.searched_query.query_string() == "*=>[KNN 10 @embedding $vec AS vector_distance]" + + +def test_missing_redis_dependency_raises_actionable_error(): + config = ValkeyVectorStoreConfig(sync_client=FakeRedis(), embedding_fn=FakeEmbeddingFn([1.0])) + blocked = {name: None for name in list(sys.modules) if name == "redis" or name.startswith("redis.")} + + with patch.dict(sys.modules, blocked): + with pytest.raises(ValueError, match="pip install redis"): + _search(config) + + +@pytest.mark.asyncio +async def test_async_search_builds_knn_query_and_maps_response(): + aembedding_fn = FakeAsyncEmbeddingFn([0.5, 0.5]) + client = FakeAsyncRedis(result=SimpleNamespace(docs=[_doc("doc:9", 0.1, text="async hit")])) + config = ValkeyVectorStoreConfig(async_client=client, aembedding_fn=aembedding_fn) + + response = await config.aexecute_search_vector_store_request( + vector_store_id="my_index", + query=["async query", "ignored"], + vector_store_search_optional_params={"max_num_results": 3}, + litellm_logging_obj=MagicMock(), + litellm_params={ + "litellm_embedding_model": "openai/text-embedding-3-small", + "litellm_embedding_config": {"api_key": "sk-async"}, + }, + ) + + assert client.ft_index_name == "my_index" + assert client.index.searched_query.query_string() == "*=>[KNN 3 @embedding $vec AS vector_distance]" + assert client.index.searched_query_params == {"vec": struct.pack("<2f", 0.5, 0.5)} + assert aembedding_fn.captured_kwargs == { + "model": "openai/text-embedding-3-small", + "input": ["async query"], + "api_key": "sk-async", + } + assert response["search_query"] == "async query" + assert response["data"][0]["score"] == pytest.approx(0.9) + assert response["data"][0]["content"] == [{"text": "async hit", "type": "text"}] + assert response["data"][0]["file_id"] == "doc:9" + + +def test_create_vector_store_is_not_supported(): + config = ValkeyVectorStoreConfig() + + with pytest.raises(NotImplementedError, match="search-only"): + config.transform_create_vector_store_request(vector_store_create_optional_params={}, api_base="") + + +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) diff --git a/ui/litellm-dashboard/public/assets/logos/valkey.svg b/ui/litellm-dashboard/public/assets/logos/valkey.svg new file mode 100644 index 00000000000..0e97e680df4 --- /dev/null +++ b/ui/litellm-dashboard/public/assets/logos/valkey.svg @@ -0,0 +1,6 @@ + + + + + + diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/CreateVectorStore.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/CreateVectorStore.test.tsx index 6c7a647e2d3..377bd2c4b56 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/CreateVectorStore.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/CreateVectorStore.test.tsx @@ -23,18 +23,21 @@ vi.mock("@/components/vector_store_providers", () => ({ OPENAI: "OpenAI", AZURE_OPENAI: "Azure OpenAI", S3Vectors: "AWS S3 Vectors", + Valkey: "Valkey", }, vectorStoreProviderMap: { BEDROCK: "bedrock", OPENAI: "openai", AZURE_OPENAI: "azure_openai", S3Vectors: "s3_vectors", + Valkey: "valkey", }, vectorStoreProviderLogoMap: { "Amazon Bedrock": "https://example.com/bedrock.png", OpenAI: "https://example.com/openai.png", "Azure OpenAI": "https://example.com/azure.png", "AWS S3 Vectors": "https://example.com/aws.png", + Valkey: "https://example.com/valkey.svg", }, getProviderSpecificFields: vi.fn((provider: string) => { if (provider === "s3_vectors") { @@ -206,6 +209,21 @@ describe("CreateVectorStore", () => { }); }); + it("should exclude valkey from the provider dropdown since it has no RAG ingestion", async () => { + render(); + + const providerSelect = screen.getByRole("combobox"); + + await act(async () => { + fireEvent.mouseDown(providerSelect); + }); + + await waitFor(() => { + expect(screen.getByText("AWS S3 Vectors")).toBeInTheDocument(); + }); + expect(screen.queryByText("Valkey")).not.toBeInTheDocument(); + }); + it("should display S3 Vectors provider-specific fields when selected", async () => { render(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/CreateVectorStore.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/CreateVectorStore.tsx index 79ba9b77c14..4eaa6779bfd 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/CreateVectorStore.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/CreateVectorStore.tsx @@ -20,6 +20,8 @@ import S3VectorsConfig from "./S3VectorsConfig"; const { Dragger } = Upload; +const RAG_INGEST_UNSUPPORTED_PROVIDERS = new Set(["valkey"]); + interface CreateVectorStoreProps { accessToken: string | null; onSuccess?: (vectorStoreId: string) => void; @@ -290,20 +292,24 @@ const CreateVectorStore: React.FC = ({ accessToken, onSu size="large" style={{ width: "100%" }} > - {Object.entries(VectorStoreProviders).map(([providerEnum, providerDisplayName]) => { - return ( - -
- - {providerDisplayName} -
-
- ); - })} + {Object.entries(VectorStoreProviders) + .filter( + ([providerEnum]) => !RAG_INGEST_UNSUPPORTED_PROVIDERS.has(vectorStoreProviderMap[providerEnum]), + ) + .map(([providerEnum, providerDisplayName]) => { + return ( + +
+ + {providerDisplayName} +
+
+ ); + })} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreForm.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreForm.test.tsx index 3eb013d47ad..c9adefbc7b4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreForm.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreForm.test.tsx @@ -3,7 +3,7 @@ import { describe, expect, it, vi } from "vitest"; import { CredentialItem } from "@/components/networking"; import { Providers, providerLogoMap } from "@/components/provider_info_helpers"; import { VectorStoreProviders } from "@/components/vector_store_providers"; -import VectorStoreForm from "./VectorStoreForm"; +import VectorStoreForm, { buildVectorStoreLitellmParams } from "./VectorStoreForm"; vi.mock("@/components/networking"); @@ -32,3 +32,56 @@ describe("VectorStoreForm", () => { expect(logo.getAttribute("src")).toBe(providerLogoMap[Providers.Bedrock]); }); }); + +describe("buildVectorStoreLitellmParams", () => { + it("renames embedding_model to litellm_embedding_model for valkey", () => { + const valkeyFormValues = { + valkey_host: "my-valkey.example.com", + valkey_port: "6379", + valkey_password: "secret", + valkey_ssl: "true", + embedding_model: "text-embedding-3-small", + valkey_text_field: "text", + valkey_embedding_field: "vector", + }; + + const params = buildVectorStoreLitellmParams("valkey", valkeyFormValues); + + const expectedParams = { + valkey_host: "my-valkey.example.com", + valkey_port: "6379", + valkey_password: "secret", + valkey_ssl: "true", + litellm_embedding_model: "text-embedding-3-small", + valkey_text_field: "text", + valkey_embedding_field: "vector", + }; + expect(params).toEqual(expectedParams); + expect(params).not.toHaveProperty("embedding_model"); + }); + + it("renames embedding_model to litellm_embedding_model for milvus", () => { + const params = buildVectorStoreLitellmParams("milvus", { + api_key: "user:pass", + api_base: "https://my-milvus-endpoint.com/", + embedding_model: "text-embedding-3-small", + }); + + expect(params).toEqual({ + api_key: "user:pass", + api_base: "https://my-milvus-endpoint.com/", + litellm_embedding_model: "text-embedding-3-small", + }); + }); + + it("keeps embedding_model as-is for providers outside the rename set", () => { + const params = buildVectorStoreLitellmParams("s3_vectors", { + vector_bucket_name: "my-vector-bucket", + aws_region_name: "us-west-2", + embedding_model: "text-embedding-3-small", + }); + + expect(params.embedding_model).toBe("text-embedding-3-small"); + expect(params).not.toHaveProperty("litellm_embedding_model"); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreForm.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreForm.tsx index 67ba469f68d..4753e26a05f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreForm.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreForm.tsx @@ -14,6 +14,21 @@ import { Logo } from "@/components/molecules/logo/Logo"; import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models"; import NotificationsManager from "@/components/molecules/notifications_manager"; +const EMBEDDING_MODEL_RENAME_PROVIDERS = new Set(["milvus", "valkey"]); + +export const buildVectorStoreLitellmParams = ( + provider: string, + formValues: Record, +): Record => + Object.fromEntries( + getProviderSpecificFields(provider).map((field) => [ + EMBEDDING_MODEL_RENAME_PROVIDERS.has(provider) && field.name === "embedding_model" + ? "litellm_embedding_model" + : field.name, + formValues[field.name], + ]), + ); + interface VectorStoreFormProps { isVisible: boolean; onCancel: () => void; @@ -74,22 +89,7 @@ const VectorStoreForm: React.FC = ({ litellm_credential_name: formValues.litellm_credential_name, }; - // pass all provider fields as litellm params dict - const providerFields = getProviderSpecificFields(formValues.custom_llm_provider); - const litellmParams = providerFields.reduce( - (acc, field) => { - // Special handling for Milvus: rename embedding_model to litellm_embedding_model - if (formValues.custom_llm_provider === "milvus" && field.name === "embedding_model") { - acc["litellm_embedding_model"] = formValues[field.name]; - } else { - acc[field.name] = formValues[field.name]; - } - return acc; - }, - {} as Record, - ); - - payload["litellm_params"] = litellmParams; + payload["litellm_params"] = buildVectorStoreLitellmParams(formValues.custom_llm_provider, formValues); await vectorStoreCreateCall(accessToken, payload); NotificationsManager.success("Vector store created successfully"); @@ -169,6 +169,26 @@ const VectorStoreForm: React.FC = ({ /> )} + {/* Valkey Setup Instructions */} + {selectedProvider === "valkey" && ( + +

+ Your Valkey instance must have the valkey-search module loaded, with an FT index already created over + HASH documents that includes a vector field whose name matches the Vector Field Name setting below + (defaults to "embedding"). LiteLLM embeds your query with the embedding model below and runs + a KNN search against that index. +

+ + } + type="info" + showIcon + style={{ marginBottom: "16px" }} + /> + )} + {/* Vertex RAG Engine Setup Instructions */} {selectedProvider === "vertex_rag_engine" && ( = ({ ? vertexEngineId ? "Any identifier you'll use to reference this in LiteLLM" : 'my-datastore_1234567890 (data store ID from Vertex AI / "Agent Search" console)' - : "Enter vector store ID from your provider" + : selectedProvider === "valkey" + ? "my-search-index (FT index name in Valkey)" + : "Enter vector store ID from your provider" } /> @@ -326,6 +348,7 @@ const VectorStoreForm: React.FC = ({ } name={field.name} + initialValue={field.initialValue} rules={ field.required ? [{ required: true, message: `Please input the ${field.label.toLowerCase()}` }] : [] } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTable.test.tsx index 7446d0efe3c..7aae01254da 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTable.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTable.test.tsx @@ -6,8 +6,8 @@ import { VectorStore } from "@/components/vector_store_management/types"; import VectorStoreTable from "./VectorStoreTable"; -vi.mock("@/components/provider_info_helpers", () => ({ - getProviderLogoAndName: (provider: string) => { +vi.mock("@/components/vector_store_providers", () => ({ + getVectorStoreProviderLogoAndName: (provider: string) => { const providerMap: Record = { openai: { displayName: "OpenAI", logo: "/openai-logo.png" }, azure: { displayName: "Azure", logo: "/azure-logo.png" }, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTableColumns.tsx index cf162578177..a4c18e48e1a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTableColumns.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTableColumns.tsx @@ -5,7 +5,7 @@ import { Copy, MoreHorizontal, Pencil, Trash2 } from "lucide-react"; import { DataTableSortHeader } from "@/components/shared/DataTable"; import { CellTooltip, DateCell, IdentityCell } from "@/components/shared/table_cells"; -import { getProviderLogoAndName } from "@/components/provider_info_helpers"; +import { getVectorStoreProviderLogoAndName } from "@/components/vector_store_providers"; import { buttonVariants } from "@/components/ui/button"; import { DropdownMenu, @@ -19,7 +19,7 @@ import { cn } from "@/lib/cva.config"; import { copyToClipboard } from "@/utils/dataUtils"; function VectorStoreProviderCell({ provider }: { provider: string }) { - const { displayName, logo } = getProviderLogoAndName(provider); + const { displayName, logo } = getVectorStoreProviderLogoAndName(provider); return (
{logo ? ( diff --git a/ui/litellm-dashboard/src/components/vector_store_providers.test.tsx b/ui/litellm-dashboard/src/components/vector_store_providers.test.tsx index 73faad09d44..782314cc4c9 100644 --- a/ui/litellm-dashboard/src/components/vector_store_providers.test.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_providers.test.tsx @@ -1,6 +1,12 @@ import { describe, expect, it } from "vitest"; import { Providers, providerLogoMap } from "@/components/provider_info_helpers"; -import { getVectorStoreProviderLogoAndName, VectorStoreProviders } from "./vector_store_providers"; +import { + getProviderSpecificFields, + getVectorStoreProviderLogoAndName, + VectorStoreProviders, + vectorStoreProviderLogoMap, + vectorStoreProviderMap, +} from "./vector_store_providers"; describe("getVectorStoreProviderLogoAndName", () => { it("resolves vector-store-only slugs to their own logo and display name", () => { @@ -16,6 +22,38 @@ describe("getVectorStoreProviderLogoAndName", () => { logo: expect.stringContaining("s3_vector"), displayName: VectorStoreProviders.S3Vectors, }); + expect(getVectorStoreProviderLogoAndName("valkey")).toEqual({ + logo: expect.stringContaining("valkey"), + displayName: VectorStoreProviders.Valkey, + }); + }); + + it("registers valkey in the provider, logo, and field maps", () => { + expect(vectorStoreProviderMap.Valkey).toBe("valkey"); + expect(vectorStoreProviderLogoMap[VectorStoreProviders.Valkey]).toContain("valkey"); + expect(getProviderSpecificFields("valkey").map((field) => field.name)).toEqual([ + "valkey_host", + "valkey_port", + "valkey_password", + "valkey_ssl", + "embedding_model", + "valkey_text_field", + "valkey_embedding_field", + ]); + }); + + it("offers valkey_ssl as a false/true select defaulting to false", () => { + const sslField = getProviderSpecificFields("valkey").find((field) => field.name === "valkey_ssl"); + + expect(sslField).toMatchObject({ + type: "select", + required: false, + initialValue: "false", + }); + expect(sslField?.options).toEqual([ + { value: "false", label: "false" }, + { value: "true", label: "true" }, + ]); }); it("resolves shared slugs to the same bundled logo as the provider map", () => { diff --git a/ui/litellm-dashboard/src/components/vector_store_providers.tsx b/ui/litellm-dashboard/src/components/vector_store_providers.tsx index d8ee2eb4720..ed91988f397 100644 --- a/ui/litellm-dashboard/src/components/vector_store_providers.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_providers.tsx @@ -2,6 +2,7 @@ import { getProviderLogoAndName, Providers, providerLogoMap } from "@/components import milvusLogo from "../../public/assets/logos/milvus.svg"; import postgresqlLogo from "../../public/assets/logos/postgresql.svg"; import s3VectorLogo from "../../public/assets/logos/s3_vector.png"; +import valkeyLogo from "../../public/assets/logos/valkey.svg"; export enum VectorStoreProviders { Bedrock = "Amazon Bedrock", @@ -12,6 +13,7 @@ export enum VectorStoreProviders { OpenAI = "OpenAI", Azure = "Azure OpenAI", Milvus = "Milvus", + Valkey = "Valkey", } export const vectorStoreProviderMap: Record = { @@ -23,6 +25,7 @@ export const vectorStoreProviderMap: Record = { Azure: "azure", Milvus: "milvus", S3Vectors: "s3_vectors", + Valkey: "valkey", }; export const vectorStoreProviderLogoMap: Record = { @@ -34,6 +37,7 @@ export const vectorStoreProviderLogoMap: Record = { [VectorStoreProviders.Azure]: providerLogoMap[Providers.Azure] ?? "", [VectorStoreProviders.Milvus]: milvusLogo.src, [VectorStoreProviders.S3Vectors]: s3VectorLogo.src, + [VectorStoreProviders.Valkey]: valkeyLogo.src, }; // Define field types for provider-specific configurations @@ -165,6 +169,71 @@ export const vectorStoreProviderFields: Record type: "select", }, ], + valkey: [ + { + name: "valkey_host", + label: "Valkey Host", + tooltip: "Hostname of your Valkey server with the valkey-search module loaded (no scheme or port)", + placeholder: "my-valkey.example.com", + required: true, + type: "text", + }, + { + name: "valkey_port", + label: "Valkey Port", + tooltip: "Port of your Valkey server", + placeholder: "6379", + required: false, + type: "text", + initialValue: "6379", + }, + { + name: "valkey_password", + label: "Valkey Password", + tooltip: "Password for your Valkey server (leave blank if authentication is disabled)", + required: false, + type: "password", + }, + { + name: "valkey_ssl", + label: "Use TLS", + tooltip: + "Connect over TLS with rediss:// instead of redis://. Required for AWS ElastiCache clusters that have in-transit encryption enabled.", + required: false, + type: "select", + options: [ + { value: "false", label: "false" }, + { value: "true", label: "true" }, + ], + initialValue: "false", + }, + { + name: "embedding_model", + label: "Embedding Model", + tooltip: "Embedding model used to embed queries before searching the Valkey index", + placeholder: "text-embedding-3-small", + required: true, + type: "text", + }, + { + name: "valkey_text_field", + label: "Text Field", + tooltip: "Name of the hash field that stores the document text", + placeholder: "text", + required: false, + type: "text", + initialValue: "text", + }, + { + name: "valkey_embedding_field", + label: "Vector Field Name", + tooltip: "Name of the HASH field in your FT index that stores the document vector", + placeholder: "embedding", + required: false, + type: "text", + initialValue: "embedding", + }, + ], s3_vectors: [ { name: "vector_bucket_name",