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..02a51a8bace 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,75 @@ 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], + timeout: float | httpx.Timeout | None = None, + ) -> 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], + timeout: float | httpx.Timeout | None = None, + ) -> 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 d67497dd4da..d5b47466477 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -56,7 +56,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, ) @@ -9416,6 +9419,24 @@ 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 + timeout=timeout, + ) + if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( llm_provider=litellm.LlmProviders(custom_llm_provider), @@ -9529,6 +9550,24 @@ 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 + timeout=timeout, + ) + 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..3cbfca0f1a9 --- /dev/null +++ b/litellm/llms/valkey/vector_stores/transformation.py @@ -0,0 +1,299 @@ +""" +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_SOCKET_CONNECT_TIMEOUT_SECONDS: Final = 5.0 +DEFAULT_SOCKET_TIMEOUT_SECONDS: Final = 30.0 +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 " ".join(query) + + @staticmethod + def _socket_timeouts(timeout: float | httpx.Timeout | None) -> tuple[float, float]: + if isinstance(timeout, httpx.Timeout): + return ( + timeout.connect or DEFAULT_SOCKET_CONNECT_TIMEOUT_SECONDS, + timeout.read or DEFAULT_SOCKET_TIMEOUT_SECONDS, + ) + if timeout is not None: + return (min(float(timeout), DEFAULT_SOCKET_CONNECT_TIMEOUT_SECONDS), float(timeout)) + return (DEFAULT_SOCKET_CONNECT_TIMEOUT_SECONDS, DEFAULT_SOCKET_TIMEOUT_SECONDS) + + @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], + timeout: float | httpx.Timeout | None = None, + ) -> 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) + + connect_timeout, op_timeout = self._socket_timeouts(timeout) + client: Final = _import_sync_redis().from_url( + params.connection_url(), + socket_connect_timeout=connect_timeout, + socket_timeout=op_timeout, + ) + 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], + timeout: float | httpx.Timeout | None = None, + ) -> 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) + + connect_timeout, op_timeout = self._socket_timeouts(timeout) + client: Final = _import_async_redis().from_url( + params.connection_url(), + socket_connect_timeout=connect_timeout, + socket_timeout=op_timeout, + ) + 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 6df757ef98c..7d1dd1358d5 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -351,6 +351,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 cd2ef9dde2c..ae9395fc851 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3711,6 +3711,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 fbef20078f3..73099eb47f4 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..2dda8bf722a 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,90 @@ 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, + timeout=12.5, + _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["timeout"] == 12.5 + 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, + timeout=7.0, + _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" + assert config.async_calls[0]["timeout"] == 7.0 + 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..a2ee2c2bdb1 --- /dev/null +++ b/tests/test_litellm/llms/valkey/vector_stores/test_valkey_transformation.py @@ -0,0 +1,376 @@ +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 httpx +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_joins_all_elements(): + 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 second query"] + assert response["search_query"] == "first query second query" + + +def test_socket_timeouts_default_to_bounded_values(): + assert ValkeyVectorStoreConfig._socket_timeouts(None) == (5.0, 30.0) + + +def test_socket_timeouts_derive_from_numeric_request_timeout(): + assert ValkeyVectorStoreConfig._socket_timeouts(2.0) == (2.0, 2.0) + assert ValkeyVectorStoreConfig._socket_timeouts(120.0) == (5.0, 120.0) + + +def test_socket_timeouts_derive_from_httpx_timeout(): + timeout = httpx.Timeout(connect=3.0, read=7.0, write=1.0, pool=1.0) + + assert ValkeyVectorStoreConfig._socket_timeouts(timeout) == (3.0, 7.0) + + +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", "part two"], + 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 part two"], + "api_key": "sk-async", + } + assert response["search_query"] == "async query part two" + 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 e0bd4feef59..90588860371 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 @@ -16,18 +16,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") { @@ -199,6 +202,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 2765dd687f7..73102965ecc 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 @@ -27,6 +27,8 @@ import S3VectorsConfig from "./S3VectorsConfig"; const { Dragger } = Upload; +const RAG_INGEST_UNSUPPORTED_PROVIDERS = new Set(["valkey"]); + const asText = (value: unknown): string => (typeof value === "string" ? value : ""); const labelWithHint = (label: string, hint: string): React.ReactNode => ( @@ -297,16 +299,20 @@ const CreateVectorStore: React.FC = ({ accessToken, onSu - {Object.entries(VectorStoreProviders).map(([providerEnum, providerDisplayName]) => ( - - - {providerDisplayName} - - ))} + {Object.entries(VectorStoreProviders) + .filter( + ([providerEnum]) => !RAG_INGEST_UNSUPPORTED_PROVIDERS.has(vectorStoreProviderMap[providerEnum]), + ) + .map(([providerEnum, providerDisplayName]) => ( + + + {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 9062dccafd0..09a0c7651f1 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 @@ -4,7 +4,7 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import { CredentialItem, vectorStoreCreateCall } 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"); @@ -68,3 +68,56 @@ describe("VectorStoreForm", () => { expect(vectorStoreCreateCall).not.toHaveBeenCalled(); }); }); + +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 97c633f4f13..05c829d79df 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 @@ -33,6 +33,23 @@ import { Textarea } from "@/components/ui/textarea"; import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; import { useZodForm } from "@/lib/forms/useZodForm"; +const EMBEDDING_MODEL_RENAME_PROVIDERS = new Set(["milvus", "valkey"]); + +export const buildVectorStoreLitellmParams = ( + provider: string, + formValues: Record, +): Record => + Object.fromEntries( + getProviderSpecificFields(provider) + .filter(isSupportedProviderField) + .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; @@ -52,6 +69,12 @@ const PROVIDER_FIELD_NAMES = [ "vector_bucket_name", "index_name", "aws_region_name", + "valkey_host", + "valkey_port", + "valkey_password", + "valkey_ssl", + "valkey_text_field", + "valkey_embedding_field", ] as const; type ProviderFieldName = (typeof PROVIDER_FIELD_NAMES)[number]; @@ -77,6 +100,12 @@ const vectorStoreShape = { vector_bucket_name: optionalText, index_name: optionalText, aws_region_name: optionalText, + valkey_host: optionalText, + valkey_port: optionalText, + valkey_password: optionalText, + valkey_ssl: optionalText, + valkey_text_field: optionalText, + valkey_embedding_field: optionalText, }; const vectorStoreSchema = z.object(vectorStoreShape).superRefine((values, ctx) => { @@ -100,6 +129,10 @@ const EMPTY_VALUES: VectorStoreFormValues = { custom_llm_provider: "bedrock", vector_store_id: "", vertex_location: "global", + valkey_port: "6379", + valkey_ssl: "false", + valkey_text_field: "text", + valkey_embedding_field: "embedding", }; interface CredentialOption { @@ -193,17 +226,6 @@ const VectorStoreForm: React.FC = ({ return; } - const providerFields = getProviderSpecificFields(formValues.custom_llm_provider); - const litellmParams = Object.fromEntries( - providerFields.filter(isSupportedProviderField).map((field) => { - const value = formValues[field.name]; - if (formValues.custom_llm_provider === "milvus" && field.name === "embedding_model") { - return ["litellm_embedding_model", value]; - } - return [field.name, value]; - }), - ); - await vectorStoreCreateCall(accessToken, { vector_store_id: formValues.vector_store_id, custom_llm_provider: formValues.custom_llm_provider, @@ -211,7 +233,7 @@ const VectorStoreForm: React.FC = ({ vector_store_description: formValues.vector_store_description, vector_store_metadata: metadata, litellm_credential_name: formValues.litellm_credential_name, - litellm_params: litellmParams, + litellm_params: buildVectorStoreLitellmParams(formValues.custom_llm_provider, formValues), }); toast.success("Vector store created successfully"); form.reset(EMPTY_VALUES); @@ -237,7 +259,9 @@ const VectorStoreForm: React.FC = ({ ? 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"; return ( @@ -291,7 +315,7 @@ const VectorStoreForm: React.FC = ({ description={

LiteLLM provides a server to connect to PG Vector. To use this provider:

-
    +
    1. Deploy the litellm-pgvector server from:{" "} @@ -309,6 +333,44 @@ const VectorStoreForm: React.FC = ({ /> )} + {selectedProvider === "valkey" && ( + +

      + LiteLLM searches documents you have already stored in Valkey. It does not create the index or + upload documents for you. Before creating this vector store, make sure: +

      +
        +
      1. + Your Valkey server has vector search enabled (the valkey-search module, included in the + valkey-bundle image and in AWS ElastiCache / MemoryDB for Valkey) +
        2. +
      2. + You have already created a search index and loaded your documents and their embeddings into it. + Enter that index name as the Vector Store ID +
        3. +
      3. + You know which embedding model created those stored embeddings. That model must be added to this + proxy under Models so you can pick it below. Using a different model returns wrong results +
        4. +
      4. + You know the field names your documents use for their text and their embedding. If they are not + "text" and "embedding", set them below +
        5. +
      +

      + When a query comes in, LiteLLM converts it to an embedding with the model below and returns the + closest matching documents from your index. +

      +
+ } + type="info" + showIcon + /> + )} + {selectedProvider === "vertex_rag_engine" && ( = ({ Note: Google Cloud has renamed this to "RAG Engine" in its console — the steps below still apply.

-
    +
    1. Set up your Vertex AI RAG Engine corpus following the guide:{" "} = ({ Note: Google Cloud has renamed this to "Agent Search" in its console — the steps below still apply.

      -
        +
        1. Enable the Discovery Engine API on your Google Cloud project and create a data store following the guide:{" "} 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..eaf2a52853f 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,45 @@ 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("picks the valkey embedding model from the proxy's models like milvus does", () => { + const embeddingField = getProviderSpecificFields("valkey").find((field) => field.name === "embedding_model"); + + expect(embeddingField).toMatchObject({ type: "select", required: true }); + expect(embeddingField).not.toHaveProperty("options"); + }); + + 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..35cd5c383f7 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,74 @@ export const vectorStoreProviderFields: Record type: "select", }, ], + valkey: [ + { + name: "valkey_host", + label: "Valkey Host", + tooltip: "Hostname or IP of your Valkey server, without redis:// or a port (e.g. my-valkey.example.com)", + placeholder: "my-valkey.example.com", + required: true, + type: "text", + }, + { + name: "valkey_port", + label: "Valkey Port", + tooltip: "Port your Valkey server listens on. Leave as 6379 unless you changed it", + placeholder: "6379", + required: false, + type: "text", + initialValue: "6379", + }, + { + name: "valkey_password", + label: "Valkey Password", + tooltip: "Password used to log in to your Valkey server. Leave blank if it has no password", + required: false, + type: "password", + }, + { + name: "valkey_ssl", + label: "Use TLS", + tooltip: + "Set to true if your Valkey server requires an encrypted (TLS) connection, for example AWS ElastiCache with in-transit encryption turned on", + required: false, + type: "select", + options: [ + { value: "false", label: "false" }, + { value: "true", label: "true" }, + ], + initialValue: "false", + }, + { + name: "embedding_model", + label: "Embedding Model", + tooltip: + "The embedding model on this proxy that was used to create the embeddings already stored in your Valkey index. LiteLLM uses it to embed each search query, so it must be the same model or results will be wrong. Add it under Models first if it is not listed", + placeholder: "text-embedding-3-small", + required: true, + type: "select", + }, + { + name: "valkey_text_field", + label: "Text Field", + tooltip: + "The field in each stored document that holds its readable text. LiteLLM returns this text in search results. Must match how your documents were stored (default: text)", + placeholder: "text", + required: false, + type: "text", + initialValue: "text", + }, + { + name: "valkey_embedding_field", + label: "Vector Field Name", + tooltip: + "The field in each stored document that holds its embedding. LiteLLM searches against this field, so it must match the field your index was created on (default: embedding)", + placeholder: "embedding", + required: false, + type: "text", + initialValue: "embedding", + }, + ], s3_vectors: [ { name: "vector_bucket_name", diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 3a129890267..0fc2afa948b 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -27559,6 +27559,18 @@ export interface components { * @default false */ use_xai_oauth: boolean | null; + /** Valkey Embedding Field */ + valkey_embedding_field?: string | null; + /** Valkey Host */ + valkey_host?: string | null; + /** Valkey Password */ + valkey_password?: string | null; + /** Valkey Port */ + valkey_port?: number | null; + /** Valkey Ssl */ + valkey_ssl?: boolean | null; + /** Valkey Text Field */ + valkey_text_field?: string | null; /** Vector Store Id */ vector_store_id?: string | null; /** Vertex Credentials */ @@ -36629,6 +36641,18 @@ export interface components { * @default false */ use_xai_oauth: boolean | null; + /** Valkey Embedding Field */ + valkey_embedding_field?: string | null; + /** Valkey Host */ + valkey_host?: string | null; + /** Valkey Password */ + valkey_password?: string | null; + /** Valkey Port */ + valkey_port?: number | null; + /** Valkey Ssl */ + valkey_ssl?: boolean | null; + /** Valkey Text Field */ + valkey_text_field?: string | null; /** Vector Store Id */ vector_store_id?: string | null; /** Vertex Credentials */