mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
feat(vector_stores): add Valkey as a managed vector store provider (#37002)
* feat(vector_stores): add Valkey as a managed vector store provider Adds a valkey provider for managed vector stores, searchable via the valkey-search module over RESP. Introduces BaseDirectVectorStoreConfig for datastores that execute searches directly instead of building an HTTP request, and refactors the valkey semantic cache to share the new connection URL helper. Registered in the provider enum, router params, proxy config registry, Admin UI Add Vector Store modal, and provider endpoint support matrix. * fix(vector_stores): join list queries and bound valkey socket timeouts Review feedback: multi-string queries are now space-joined like every other embedding-based provider instead of dropping all but the first, and the request timeout is threaded through the direct vector store interface into bounded socket_connect_timeout / socket_timeout values on both redis clients so an unreachable Valkey host cannot pin proxy workers until the OS TCP timeout. * chore(ui): regenerate schema.d.ts for valkey vector store fields * docs(ui): make the Valkey vector store setup note and field tooltips explicit * feat(ui): pick the Valkey embedding model from the proxy's models like Milvus * fix(ui): number the setup steps in the vector store provider alerts
This commit is contained in:
ryan-crabbe-berri
GitHub
parent
3fe0201d40
commit
28266d90e7
25 changed files with 1249 additions and 43 deletions
|
|
@ -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 (
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
0
litellm/llms/valkey/__init__.py
Normal file
0
litellm/llms/valkey/__init__.py
Normal file
18
litellm/llms/valkey/common_utils.py
Normal file
18
litellm/llms/valkey/common_utils.py
Normal file
|
|
@ -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)
|
||||
3
litellm/llms/valkey/vector_stores/__init__.py
Normal file
3
litellm/llms/valkey/vector_stores/__init__.py
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
from litellm.llms.valkey.vector_stores.transformation import ValkeyVectorStoreConfig
|
||||
|
||||
__all__ = ("ValkeyVectorStoreConfig",)
|
||||
299
litellm/llms/valkey/vector_stores/transformation.py
Normal file
299
litellm/llms/valkey/vector_stores/transformation.py
Normal file
|
|
@ -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)
|
||||
6
litellm/proxy/_experimental/out/assets/logos/valkey.svg
Normal file
6
litellm/proxy/_experimental/out/assets/logos/valkey.svg
Normal file
|
|
@ -0,0 +1,6 @@
|
|||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<svg width="64" height="73" viewBox="0 0 64 73" xmlns="http://www.w3.org/2000/svg">
|
||||
<g id="Group-copy">
|
||||
<path id="Path" fill="#123678" fill-rule="evenodd" stroke="none" d="M 13.482285 60.694962 L 0.998384 52.884399 L 0.998384 19.502914 L 31.527868 2.001205 L 61.317604 19.532024 L 61.317604 54.64489 L 31.054855 71.68927 L 20.548372 65.115807 L 20.548372 51.041328 L 20.548372 49.119896 L 14.851504 45.555508 L 14.851504 27.453159 L 31.346497 17.99712 L 47.464485 27.482262 L 47.464485 46.451157 L 34.703495 53.638138 L 34.703495 45.998573 C 38.52874 44.52552 41.274452 40.739189 41.274452 36.270489 C 41.274452 30.510658 36.712814 25.88438 31.158138 25.88438 C 25.603172 25.88438 21.041817 30.510658 21.041817 36.270489 C 21.041817 40.739189 23.787249 44.52552 27.612494 45.998573 L 27.612494 60.473576 L 31.261133 62.756348 L 53.635483 50.15464 L 53.635483 23.924595 L 31.477489 10.884869 L 8.680504 23.953705 L 8.680504 48.628967 L 13.482285 51.633297 L 13.482285 60.694962 Z M 31.158138 31.498383 C 33.671822 31.498383 35.660439 33.664162 35.660439 36.270489 C 35.660439 38.876804 33.671822 41.042587 31.158138 41.042587 C 28.644447 41.042587 26.655558 38.876804 26.655558 36.270489 C 26.655558 33.664162 28.644447 31.498383 31.158138 31.498383 Z" />
|
||||
</g>
|
||||
</svg>
|
||||
|
After Width: | Height: | Size: 1.3 KiB |
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
6
ui/litellm-dashboard/public/assets/logos/valkey.svg
Normal file
6
ui/litellm-dashboard/public/assets/logos/valkey.svg
Normal file
|
|
@ -0,0 +1,6 @@
|
|||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<svg width="64" height="73" viewBox="0 0 64 73" xmlns="http://www.w3.org/2000/svg">
|
||||
<g id="Group-copy">
|
||||
<path id="Path" fill="#123678" fill-rule="evenodd" stroke="none" d="M 13.482285 60.694962 L 0.998384 52.884399 L 0.998384 19.502914 L 31.527868 2.001205 L 61.317604 19.532024 L 61.317604 54.64489 L 31.054855 71.68927 L 20.548372 65.115807 L 20.548372 51.041328 L 20.548372 49.119896 L 14.851504 45.555508 L 14.851504 27.453159 L 31.346497 17.99712 L 47.464485 27.482262 L 47.464485 46.451157 L 34.703495 53.638138 L 34.703495 45.998573 C 38.52874 44.52552 41.274452 40.739189 41.274452 36.270489 C 41.274452 30.510658 36.712814 25.88438 31.158138 25.88438 C 25.603172 25.88438 21.041817 30.510658 21.041817 36.270489 C 21.041817 40.739189 23.787249 44.52552 27.612494 45.998573 L 27.612494 60.473576 L 31.261133 62.756348 L 53.635483 50.15464 L 53.635483 23.924595 L 31.477489 10.884869 L 8.680504 23.953705 L 8.680504 48.628967 L 13.482285 51.633297 L 13.482285 60.694962 Z M 31.158138 31.498383 C 33.671822 31.498383 35.660439 33.664162 35.660439 36.270489 C 35.660439 38.876804 33.671822 41.042587 31.158138 41.042587 C 28.644447 41.042587 26.655558 38.876804 26.655558 36.270489 C 26.655558 33.664162 28.644447 31.498383 31.158138 31.498383 Z" />
|
||||
</g>
|
||||
</svg>
|
||||
|
After Width: | Height: | Size: 1.3 KiB |
|
|
@ -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(<CreateVectorStore accessToken="test-token" />);
|
||||
|
||||
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(<CreateVectorStore accessToken="test-token" />);
|
||||
|
||||
|
|
|
|||
|
|
@ -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<CreateVectorStoreProps> = ({ accessToken, onSu
|
|||
<SelectValue placeholder="Select a provider" />
|
||||
</SelectTrigger>
|
||||
<SelectContent alignItemWithTrigger={false}>
|
||||
{Object.entries(VectorStoreProviders).map(([providerEnum, providerDisplayName]) => (
|
||||
<SelectItem key={providerEnum} value={vectorStoreProviderMap[providerEnum]}>
|
||||
<Logo
|
||||
src={vectorStoreProviderLogoMap[providerDisplayName]}
|
||||
label={providerDisplayName}
|
||||
className="w-5 h-5"
|
||||
/>
|
||||
<span>{providerDisplayName}</span>
|
||||
</SelectItem>
|
||||
))}
|
||||
{Object.entries(VectorStoreProviders)
|
||||
.filter(
|
||||
([providerEnum]) => !RAG_INGEST_UNSUPPORTED_PROVIDERS.has(vectorStoreProviderMap[providerEnum]),
|
||||
)
|
||||
.map(([providerEnum, providerDisplayName]) => (
|
||||
<SelectItem key={providerEnum} value={vectorStoreProviderMap[providerEnum]}>
|
||||
<Logo
|
||||
src={vectorStoreProviderLogoMap[providerDisplayName]}
|
||||
label={providerDisplayName}
|
||||
className="w-5 h-5"
|
||||
/>
|
||||
<span>{providerDisplayName}</span>
|
||||
</SelectItem>
|
||||
))}
|
||||
</SelectContent>
|
||||
</Select>
|
||||
</Field>
|
||||
|
|
|
|||
|
|
@ -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");
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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<string, unknown>,
|
||||
): Record<string, unknown> =>
|
||||
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<VectorStoreFormProps> = ({
|
|||
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<VectorStoreFormProps> = ({
|
|||
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<VectorStoreFormProps> = ({
|
|||
? 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 (
|
||||
<Modal title="Add New Vector Store" open={isVisible} width={1000} footer={null} onCancel={handleCancel}>
|
||||
|
|
@ -291,7 +315,7 @@ const VectorStoreForm: React.FC<VectorStoreFormProps> = ({
|
|||
description={
|
||||
<div>
|
||||
<p>LiteLLM provides a server to connect to PG Vector. To use this provider:</p>
|
||||
<ol style={{ marginLeft: "16px", marginTop: "8px" }}>
|
||||
<ol style={{ marginLeft: "16px", marginTop: "8px", listStyleType: "decimal" }}>
|
||||
<li>
|
||||
Deploy the litellm-pgvector server from:{" "}
|
||||
<a href="https://github.com/BerriAI/litellm-pgvector" target="_blank" rel="noopener noreferrer">
|
||||
|
|
@ -309,6 +333,44 @@ const VectorStoreForm: React.FC<VectorStoreFormProps> = ({
|
|||
/>
|
||||
)}
|
||||
|
||||
{selectedProvider === "valkey" && (
|
||||
<Alert
|
||||
message="Valkey Setup Required"
|
||||
description={
|
||||
<div>
|
||||
<p>
|
||||
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:
|
||||
</p>
|
||||
<ol style={{ marginLeft: "16px", marginTop: "8px", listStyleType: "decimal" }}>
|
||||
<li>
|
||||
Your Valkey server has vector search enabled (the valkey-search module, included in the
|
||||
valkey-bundle image and in AWS ElastiCache / MemoryDB for Valkey)
|
||||
</li>
|
||||
<li>
|
||||
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
|
||||
</li>
|
||||
<li>
|
||||
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
|
||||
</li>
|
||||
<li>
|
||||
You know the field names your documents use for their text and their embedding. If they are not
|
||||
"text" and "embedding", set them below
|
||||
</li>
|
||||
</ol>
|
||||
<p style={{ marginTop: "8px" }}>
|
||||
When a query comes in, LiteLLM converts it to an embedding with the model below and returns the
|
||||
closest matching documents from your index.
|
||||
</p>
|
||||
</div>
|
||||
}
|
||||
type="info"
|
||||
showIcon
|
||||
/>
|
||||
)}
|
||||
|
||||
{selectedProvider === "vertex_rag_engine" && (
|
||||
<Alert
|
||||
message="Vertex AI RAG Engine Setup"
|
||||
|
|
@ -319,7 +381,7 @@ const VectorStoreForm: React.FC<VectorStoreFormProps> = ({
|
|||
Note: Google Cloud has renamed this to "RAG Engine" in its console — the steps below
|
||||
still apply.
|
||||
</p>
|
||||
<ol style={{ marginLeft: "16px", marginTop: "8px" }}>
|
||||
<ol style={{ marginLeft: "16px", marginTop: "8px", listStyleType: "decimal" }}>
|
||||
<li>
|
||||
Set up your Vertex AI RAG Engine corpus following the guide:{" "}
|
||||
<a
|
||||
|
|
@ -354,7 +416,7 @@ const VectorStoreForm: React.FC<VectorStoreFormProps> = ({
|
|||
Note: Google Cloud has renamed this to "Agent Search" in its console — the steps below
|
||||
still apply.
|
||||
</p>
|
||||
<ol style={{ marginLeft: "16px", marginTop: "8px" }}>
|
||||
<ol style={{ marginLeft: "16px", marginTop: "8px", listStyleType: "decimal" }}>
|
||||
<li>
|
||||
Enable the Discovery Engine API on your Google Cloud project and create a data store following
|
||||
the guide:{" "}
|
||||
|
|
|
|||
|
|
@ -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<string, { displayName: string; logo: string }> = {
|
||||
openai: { displayName: "OpenAI", logo: "/openai-logo.png" },
|
||||
azure: { displayName: "Azure", logo: "/azure-logo.png" },
|
||||
|
|
|
|||
|
|
@ -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 (
|
||||
<div className="flex items-center gap-2">
|
||||
{logo ? (
|
||||
|
|
|
|||
|
|
@ -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", () => {
|
||||
|
|
|
|||
|
|
@ -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<string, string> = {
|
||||
|
|
@ -23,6 +25,7 @@ export const vectorStoreProviderMap: Record<string, string> = {
|
|||
Azure: "azure",
|
||||
Milvus: "milvus",
|
||||
S3Vectors: "s3_vectors",
|
||||
Valkey: "valkey",
|
||||
};
|
||||
|
||||
export const vectorStoreProviderLogoMap: Record<string, string> = {
|
||||
|
|
@ -34,6 +37,7 @@ export const vectorStoreProviderLogoMap: Record<string, string> = {
|
|||
[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<string, VectorStoreFieldConfig[]>
|
|||
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",
|
||||
|
|
|
|||
24
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
24
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -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 */
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue