From 3ab17c41aff7918f22785ab8d72346abc20d4965 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 21:22:25 -0700 Subject: [PATCH] fix(vector-stores): return clear errors for Milvus gRPC connection failures Wrap every pymilvus call site so a MilvusException with code 2 comes back as a litellm.APIConnectionError naming api_base and api_key instead of the raw "illegal connection params or server unavailable" text. The async client connects lazily, so the mapping has to cover the search call, not just the constructor. Missing pymilvus and missing litellm_params now raise BadRequestError so the proxy answers 400 rather than 500, non-admin connection updates say whether the change or the reused stored credentials needs an admin, and the vector store form gets a Transport select so gRPC is reachable from the Admin UI. --- .../llms/milvus/vector_stores/connection.py | 37 +- .../vector_stores/grpc_transformation.py | 103 ++++- .../management_endpoints.py | 32 +- litellm/proxy/vector_store_endpoints/utils.py | 2 + litellm/types/router.py | 3 +- litellm/types/vector_stores.py | 5 +- litellm/utils.py | 5 +- .../vector_stores/vector_store_registry.py | 7 + .../vector_stores/test_milvus_vector_store.py | 159 ++++++++ .../test_vector_store_endpoints.py | 355 +++++++++++++++++- .../test_vector_store_registry.py | 39 ++ .../VectorStoreForm.integration.test.tsx | 31 +- .../_components/VectorStoreForm.test.tsx | 24 +- .../_components/VectorStoreForm.tsx | 3 + .../src/components/vector_store_providers.tsx | 16 +- 15 files changed, 767 insertions(+), 54 deletions(-) diff --git a/litellm/llms/milvus/vector_stores/connection.py b/litellm/llms/milvus/vector_stores/connection.py index f734f623450..b2a3204de95 100644 --- a/litellm/llms/milvus/vector_stores/connection.py +++ b/litellm/llms/milvus/vector_stores/connection.py @@ -23,7 +23,16 @@ MILVUS_MANAGED_CONFIGURATION_FIELDS: Final = frozenset( class MilvusConnectionRejection(Enum): - ADMIN_REQUIRED = "Only proxy admins can configure vector store connections. Contact your LiteLLM administrator." + ADMIN_REQUIRED = ( + "Only proxy admins can configure Milvus gRPC vector store connections. Contact your LiteLLM administrator." + ) + ADMIN_CHANGE_REQUIRED = ( + "Only proxy admins can change a Milvus gRPC vector store connection. Contact your LiteLLM administrator." + ) + ADMIN_CREDENTIAL_REUSE = ( + "Only proxy admins can change a vector store connection that keeps its stored credentials. " + "Send the credentials with the update or contact your LiteLLM administrator." + ) ADMIN_SAVE_REQUIRED = "This managed Milvus gRPC connection must be re-saved by a proxy admin before it can be used." @@ -90,6 +99,7 @@ def prepare_connection_for_persistence( litellm_credential_name: object | None = None, existing_litellm_credential_name: object | None = None, litellm_credential_name_supplied: bool = False, + reuses_stored_credentials: bool = False, ) -> Mapping[str, object] | MilvusConnectionRejection: existing: Final = _connection_fields(existing_litellm_params) supplied: Final = _connection_fields(litellm_params) @@ -109,13 +119,26 @@ def prepare_connection_for_persistence( litellm_credential_name != existing_litellm_credential_name ) connection_changed: Final = not is_create and (provider_changed or credential_changed or effective != existing) - missing_marker: Final = effective_is_grpc and ( - not isinstance(existing_litellm_params, Mapping) - or existing_litellm_params.get(MILVUS_ADMIN_CONFIGURED_CONNECTION) is not True + inherits_stored_secrets: Final = ( + reuses_stored_credentials + or not isinstance(litellm_params, Mapping) + or (existing_litellm_credential_name is not None and not credential_changed) ) - if (connection_changed or missing_marker) and not is_proxy_admin: - return MilvusConnectionRejection.ADMIN_REQUIRED - return MappingProxyType({**effective, MILVUS_ADMIN_CONFIGURED_CONNECTION: True}) if effective_is_grpc else effective + if is_proxy_admin: + return ( + MappingProxyType({**effective, MILVUS_ADMIN_CONFIGURED_CONNECTION: True}) + if effective_is_grpc + else effective + ) + if is_create: + return MilvusConnectionRejection.ADMIN_REQUIRED if effective_is_grpc else effective + if not connection_changed: + return MappingProxyType({**existing_litellm_params} if isinstance(existing_litellm_params, Mapping) else {}) + if previous_is_grpc or effective_is_grpc: + return MilvusConnectionRejection.ADMIN_CHANGE_REQUIRED + if inherits_stored_secrets: + return MilvusConnectionRejection.ADMIN_CREDENTIAL_REUSE + return effective def managed_connection_fields(custom_llm_provider: object, litellm_params: object) -> frozenset[str]: diff --git a/litellm/llms/milvus/vector_stores/grpc_transformation.py b/litellm/llms/milvus/vector_stores/grpc_transformation.py index a902a013ea7..618785aaed2 100644 --- a/litellm/llms/milvus/vector_stores/grpc_transformation.py +++ b/litellm/llms/milvus/vector_stores/grpc_transformation.py @@ -1,5 +1,6 @@ import typing -from collections.abc import Mapping, Sequence +from collections.abc import Generator, Mapping, Sequence +from contextlib import contextmanager from types import MappingProxyType from typing import TYPE_CHECKING, Final @@ -32,6 +33,11 @@ _EMPTY_EMBEDDING_CONFIG: Final[Mapping[str, object]] = MappingProxyType({}) _PYMILVUS_INSTALL_HINT: Final = ( "Milvus gRPC transport requires the 'pymilvus' package. Install it with 'pip install litellm[milvus]'." ) +_MILVUS_CONNECT_FAILURE_CODE: Final = 2 +_MILVUS_CONNECTION_HINT: Final = ( + "Milvus gRPC connection failed. Check that api_base points at a reachable gRPC endpoint " + "and that api_key holds a valid 'user:password' token." +) _MILVUS_ENTITY_ADAPTER: Final = TypeAdapter(Mapping[str, object]) _STRING_KEYS_ADAPTER: Final = TypeAdapter(tuple[str, ...]) @@ -91,17 +97,61 @@ class _AsyncMilvusClient(Protocol): async def close(self) -> None: ... +class _MilvusErrorLike(Protocol): + @property + def code(self) -> int: ... + + +class _NeverRaised(Exception): ... + + +def _milvus_error_type() -> type[Exception]: + try: + from pymilvus import ( # pyright: ignore[reportMissingTypeStubs] # pymilvus does not publish typing metadata + MilvusException, + ) + except ImportError: + return _NeverRaised + return MilvusException + + +def _is_connect_failure(cause: Exception) -> bool: + error: Final = typing.cast( # noqa: TID251 # cast-ok: pymilvus lacks typing metadata; MilvusException always carries code + _MilvusErrorLike, cause + ) + return error.code == _MILVUS_CONNECT_FAILURE_CODE + + +@contextmanager +def _milvus_connection_errors_mapped() -> Generator[None, None, None]: + try: + yield + except _milvus_error_type() as e: + if not _is_connect_failure(e): + raise + raise litellm.APIConnectionError( + message=f"{_MILVUS_CONNECTION_HINT} {e}", + model="milvus", + llm_provider="milvus", + ) from e + + def _new_sync_client(uri: str, token: str, db_name: str, timeout: float | None) -> _SyncMilvusClient: try: from pymilvus import ( # pyright: ignore[reportMissingTypeStubs] # pymilvus does not publish typing metadata MilvusClient, ) except ImportError as e: - raise ValueError(_PYMILVUS_INSTALL_HINT) from e - return typing.cast( # noqa: TID251 # cast-ok: pymilvus lacks typing metadata; protocol defines the used surface - _SyncMilvusClient, - MilvusClient(uri=uri, token=token, db_name=db_name, timeout=timeout, dedicated=True), - ) + raise litellm.BadRequestError( + message=_PYMILVUS_INSTALL_HINT, + model="milvus", + llm_provider="milvus", + ) from e + with _milvus_connection_errors_mapped(): + return typing.cast( # noqa: TID251 # cast-ok: pymilvus lacks typing metadata; protocol defines the used surface + _SyncMilvusClient, + MilvusClient(uri=uri, token=token, db_name=db_name, timeout=timeout, dedicated=True), + ) def _new_async_client(uri: str, token: str, db_name: str, timeout: float | None) -> _AsyncMilvusClient: @@ -110,11 +160,16 @@ def _new_async_client(uri: str, token: str, db_name: str, timeout: float | None) AsyncMilvusClient, ) except ImportError as e: - raise ValueError(_PYMILVUS_INSTALL_HINT) from e - return typing.cast( # noqa: TID251 # cast-ok: pymilvus lacks typing metadata; protocol defines the used surface - _AsyncMilvusClient, - AsyncMilvusClient(uri=uri, token=token, db_name=db_name, timeout=timeout, dedicated=True), - ) + raise litellm.BadRequestError( + message=_PYMILVUS_INSTALL_HINT, + model="milvus", + llm_provider="milvus", + ) from e + with _milvus_connection_errors_mapped(): + return typing.cast( # noqa: TID251 # cast-ok: pymilvus lacks typing metadata; protocol defines the used surface + _AsyncMilvusClient, + AsyncMilvusClient(uri=uri, token=token, db_name=db_name, timeout=timeout, dedicated=True), + ) class _MilvusSearchParams(BaseModel): @@ -132,7 +187,11 @@ class _MilvusSearchParams(BaseModel): def uri(self) -> str: uri: Final = self.api_base or get_secret_str("MILVUS_API_BASE") if not uri: - raise ValueError("Milvus API base URL is required. Set MILVUS_API_BASE or pass api_base in litellm_params.") + raise litellm.BadRequestError( + message="Milvus API base URL is required. Set MILVUS_API_BASE or pass api_base in litellm_params.", + model="milvus", + llm_provider="milvus", + ) return uri.rstrip("/") @property @@ -149,9 +208,13 @@ class _MilvusSearchParams(BaseModel): def require_embedding_model(self) -> str: if not self.litellm_embedding_model: - raise ValueError( - "litellm_embedding_model is required in litellm_params for Milvus. " - "Example: litellm_params['litellm_embedding_model'] = 'openai/text-embedding-3-small'" + raise litellm.BadRequestError( + message=( + "litellm_embedding_model is required in litellm_params for Milvus. " + "Example: litellm_params['litellm_embedding_model'] = 'openai/text-embedding-3-small'" + ), + model="milvus", + llm_provider="milvus", ) return self.litellm_embedding_model @@ -380,8 +443,9 @@ class MilvusGRPCVectorStoreConfig(BaseDirectVectorStoreConfig): else _new_sync_client(params.uri, params.token, params.db_name, connection_timeout) ) try: - result: Final = client.search(**arguments) - return self._to_response(result, query_text, params.text_field) + with _milvus_connection_errors_mapped(): + result: Final = client.search(**arguments) + return self._to_response(result, query_text, params.text_field) finally: if self.sync_client is None: client.close() @@ -414,8 +478,9 @@ class MilvusGRPCVectorStoreConfig(BaseDirectVectorStoreConfig): else _new_async_client(params.uri, params.token, params.db_name, connection_timeout) ) try: - result: Final = await client.search(**arguments) - return self._to_response(result, query_text, params.text_field) + with _milvus_connection_errors_mapped(): + result: Final = await client.search(**arguments) + return self._to_response(result, query_text, params.text_field) finally: if self.async_client is None: await client.close() diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index 311850db678..d3e34dddf33 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -13,6 +13,7 @@ from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final from fastapi import APIRouter, Depends, HTTPException +from pydantic import ValidationError if TYPE_CHECKING: from prisma.models import LiteLLM_ManagedVectorStoresTable as _VectorStoreRow @@ -125,12 +126,32 @@ def _restore_redacted_litellm_params(supplied: object, existing: object, _depth: } +def _reuses_redacted_secrets(supplied: object, _depth: int = 0) -> bool: + if supplied == REDACTED_BY_LITELM_STRING: + return True + if _depth >= _REDACT_LITELLM_PARAMS_MAX_DEPTH or not isinstance(supplied, dict): + return False + return any(_reuses_redacted_secrets(value, _depth + 1) for value in deserialize_litellm_params(supplied).values()) + + +def _litellm_params_validation_detail(error: ValidationError) -> str: + return "; ".join( + f"{'.'.join(str(location) for location in issue['loc'])}: {issue['msg']}" for issue in error.errors() + ) + + def _validated_litellm_params( litellm_params: Mapping[str, object], ) -> Mapping[str, object]: from litellm.types.router import GenericLiteLLMParams - return GenericLiteLLMParams.model_validate(litellm_params).model_dump(exclude_none=True) + try: + return GenericLiteLLMParams.model_validate(litellm_params).model_dump(exclude_none=True) + except ValidationError as e: + raise HTTPException( + status_code=400, + detail=f"Invalid litellm_params: {_litellm_params_validation_detail(e)}", + ) from e def _reject_config_vector_store_id(vector_store_id: str) -> None: @@ -147,6 +168,7 @@ async def _fetch_and_authorize_vector_store( vector_store_id: str, user_api_key_dict: UserAPIKeyAuth, prisma_client: "PrismaClient", + reject_config_defined_id: bool = False, ) -> "LiteLLM_ManagedVectorStore": """ Look up a vector store by id and confirm the caller can access it. @@ -155,6 +177,8 @@ async def _fetch_and_authorize_vector_store( """ row: Final = await _vector_store_table(prisma_client).find_unique(where={"vector_store_id": vector_store_id}) if row is None: + if reject_config_defined_id: + _reject_config_vector_store_id(vector_store_id) raise HTTPException( status_code=404, detail=f"Vector store with ID {vector_store_id} not found", @@ -502,11 +526,12 @@ async def delete_vector_store( raise HTTPException(status_code=500, detail="Database not connected") try: - _reject_config_vector_store_id(data.vector_store_id) vector_store, database_exists, memory_exists = await _vector_store_delete_target( data.vector_store_id, prisma_client, ) + if not database_exists: + _reject_config_vector_store_id(data.vector_store_id) if not await _check_vector_store_access(vector_store, user_api_key_dict): raise HTTPException( status_code=403, @@ -626,7 +651,6 @@ async def update_vector_store( try: update_data: Final = data.model_dump(exclude_unset=True) vector_store_id: Final[str] = update_data.pop("vector_store_id") - _reject_config_vector_store_id(vector_store_id) # Per-store access control: anyone authenticated who passes the # premium-feature gate could otherwise update *any* vector store — @@ -635,6 +659,7 @@ async def update_vector_store( vector_store_id=vector_store_id, user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, + reject_config_defined_id=True, ) existing_litellm_params: Final = deserialize_litellm_params(existing_vector_store.get("litellm_params")) @@ -653,6 +678,7 @@ async def update_vector_store( litellm_credential_name=update_data.get("litellm_credential_name"), existing_litellm_credential_name=existing_vector_store.get("litellm_credential_name"), litellm_credential_name_supplied="litellm_credential_name" in update_data, + reuses_stored_credentials=_reuses_redacted_secrets(update_data.get("litellm_params")), ) # Handle metadata serialization diff --git a/litellm/proxy/vector_store_endpoints/utils.py b/litellm/proxy/vector_store_endpoints/utils.py index 7e29e6bcb97..d060f9f379d 100644 --- a/litellm/proxy/vector_store_endpoints/utils.py +++ b/litellm/proxy/vector_store_endpoints/utils.py @@ -92,6 +92,7 @@ def prepare_vector_store_connection_for_persistence( litellm_credential_name: object | None = None, existing_litellm_credential_name: object | None = None, litellm_credential_name_supplied: bool = False, + reuses_stored_credentials: bool = False, ) -> Mapping[str, object]: result: Final = prepare_connection_for_persistence( custom_llm_provider=custom_llm_provider, @@ -102,6 +103,7 @@ def prepare_vector_store_connection_for_persistence( litellm_credential_name=litellm_credential_name, existing_litellm_credential_name=existing_litellm_credential_name, litellm_credential_name_supplied=litellm_credential_name_supplied, + reuses_stored_credentials=reuses_stored_credentials, ) if isinstance(result, MilvusConnectionRejection): raise HTTPException(status_code=403, detail=result.value) diff --git a/litellm/types/router.py b/litellm/types/router.py index c4a21b92a45..3b324a0c4fa 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -27,6 +27,7 @@ from .utils import ( ModelResponse, StandardLoggingRoutingDecision, ) +from .vector_stores import MilvusTransport class ConfigurableClientsideParamsCustomAuth(TypedDict): @@ -379,7 +380,7 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): # Vector Store Params vector_store_id: str | None = None - milvus_transport: Literal["rest", "grpc"] | None = None + milvus_transport: MilvusTransport | None = None milvus_text_field: str | None = None milvus_db_name: str | None = None milvus_partition_names: list[str] | None = None diff --git a/litellm/types/vector_stores.py b/litellm/types/vector_stores.py index 3ad09d27580..f2d3f9e3edb 100644 --- a/litellm/types/vector_stores.py +++ b/litellm/types/vector_stores.py @@ -2,11 +2,14 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass from datetime import datetime from enum import Enum -from typing import Any, Literal +from typing import Any, Final, Literal, TypeAlias, get_args from pydantic import BaseModel from typing_extensions import TypedDict +MilvusTransport: TypeAlias = Literal["rest", "grpc"] +MILVUS_TRANSPORTS: Final[tuple[MilvusTransport, ...]] = get_args(MilvusTransport) + class SupportedVectorStoreIntegrations(str, Enum): """Supported vector store integrations.""" diff --git a/litellm/utils.py b/litellm/utils.py index adf5b21995f..e140d7fcfd7 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -249,6 +249,7 @@ from litellm.types.utils import ( Usage, all_litellm_params, ) +from litellm.types.vector_stores import MilvusTransport _CALL_TYPE_ENUM_MAP: Final[dict] = {ct.value: ct for ct in CallTypes} @@ -8910,7 +8911,7 @@ class ProviderConfigManager: @staticmethod def _get_milvus_vector_stores_config( - transport: Literal["rest", "grpc"] | None, + transport: MilvusTransport | None, ) -> BaseVectorStoreConfig: if transport == "grpc": from litellm.llms.milvus.vector_stores.grpc_transformation import ( @@ -8928,7 +8929,7 @@ class ProviderConfigManager: def get_provider_vector_stores_config( provider: LlmProviders, api_type: str | None = None, - transport: Literal["rest", "grpc"] | None = None, + transport: MilvusTransport | None = None, ) -> BaseVectorStoreConfig | None: """ v2 vector store config, use this for new vector store integrations diff --git a/litellm/vector_stores/vector_store_registry.py b/litellm/vector_stores/vector_store_registry.py index 1ea53cd95fe..d985d6b1c0e 100644 --- a/litellm/vector_stores/vector_store_registry.py +++ b/litellm/vector_stores/vector_store_registry.py @@ -21,6 +21,7 @@ from litellm.repositories.table_repositories import ( ManagedVectorStoresRepository, ) from litellm.types.vector_stores import ( + MILVUS_TRANSPORTS, VECTOR_STORE_OPENAI_PARAMS, LiteLLM_ManagedVectorStore, LiteLLM_ManagedVectorStoreIndex, @@ -469,6 +470,12 @@ class VectorStoreRegistry: raise ValueError( f"custom_llm_provider is required for initializing vector store, got custom_llm_provider={custom_llm_provider}" ) + milvus_transport = vector_store_litellm_params.get("milvus_transport") + if milvus_transport is not None and milvus_transport not in MILVUS_TRANSPORTS: + raise ValueError( + f"milvus_transport must be one of {', '.join(MILVUS_TRANSPORTS)} for vector store " + f"{vector_store_id}, got milvus_transport={milvus_transport}" + ) litellm_managed_vector_store = _MANAGED_VECTOR_STORE_ADAPTER.validate_python( { # mutable-ok: Pydantic validates the config mapping into a managed vector store diff --git a/tests/test_litellm/llms/milvus/vector_stores/test_milvus_vector_store.py b/tests/test_litellm/llms/milvus/vector_stores/test_milvus_vector_store.py index c4162683d20..441a0c62869 100644 --- a/tests/test_litellm/llms/milvus/vector_stores/test_milvus_vector_store.py +++ b/tests/test_litellm/llms/milvus/vector_stores/test_milvus_vector_store.py @@ -4,6 +4,7 @@ Tests for Milvus Vector Store import asyncio import json +import sys from typing import Final, cast from unittest.mock import AsyncMock, MagicMock, patch @@ -622,6 +623,7 @@ class TestMilvusVectorStore: async def test_grpc_client_ownership_after_failure( self, injected: bool, async_mode: bool, failure: str ) -> None: + pytest.importorskip("pymilvus") client: Final = MagicMock() executor: Final = MagicMock() executor.embed.return_value = MOCK_EMBEDDING_RESPONSE @@ -815,6 +817,7 @@ class TestMilvusVectorStore: assert isinstance(config, MilvusVectorStoreConfig) def test_public_grpc_search_passes_connection_settings_to_pymilvus(self): + pytest.importorskip("pymilvus") mock_client = MagicMock() mock_client.search.return_value = [ [ @@ -879,6 +882,7 @@ class TestMilvusVectorStore: @pytest.mark.asyncio async def test_async_grpc_uses_distinct_timeouts_and_releases_dedicated_client(self): + pytest.importorskip("pymilvus") mock_client = MagicMock() mock_client.search = AsyncMock(return_value=[[]]) mock_client.close = AsyncMock() @@ -914,6 +918,7 @@ class TestMilvusVectorStore: mock_client.close.assert_awaited_once_with() def test_http_and_https_targets_get_distinct_dedicated_clients(self): + pytest.importorskip("pymilvus") clients = [MagicMock(), MagicMock()] for client in clients: client.search.return_value = [[]] @@ -950,6 +955,160 @@ class TestMilvusVectorStore: with pytest.raises(ValueError, match="milvus_transport"): GenericLiteLLMParams.model_validate({"milvus_transport": "http"}) + def test_grpc_search_rejects_a_missing_embedding_model_with_a_400(self): + with pytest.raises(litellm.BadRequestError, match="litellm_embedding_model is required") as exc_info: + MilvusGRPCVectorStoreConfig().execute_search_vector_store_request( + vector_store_id="documents", + query="cleanup probe", + vector_store_search_optional_params={}, + litellm_logging_obj=MagicMock(), + litellm_params={"api_base": "http://milvus:19530"}, + embedding_executor=MagicMock(), + ) + + assert exc_info.value.status_code == 400 + + def test_grpc_search_rejects_a_missing_api_base_with_a_400(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("MILVUS_API_BASE", raising=False) + executor: Final = MagicMock() + executor.embed.return_value = MOCK_EMBEDDING_RESPONSE + + with pytest.raises(litellm.BadRequestError, match="Milvus API base URL is required") as exc_info: + MilvusGRPCVectorStoreConfig().execute_search_vector_store_request( + vector_store_id="documents", + query="cleanup probe", + vector_store_search_optional_params={}, + litellm_logging_obj=MagicMock(), + litellm_params={"litellm_embedding_model": "embedding-alias"}, + embedding_executor=executor, + ) + + assert exc_info.value.status_code == 400 + + def test_grpc_search_without_pymilvus_installed_is_a_400(self): + executor: Final = MagicMock() + executor.embed.return_value = MOCK_EMBEDDING_RESPONSE + + with ( + patch.dict(sys.modules, {"pymilvus": None}), + pytest.raises(litellm.BadRequestError, match="pip install litellm") as exc_info, + ): + MilvusGRPCVectorStoreConfig().execute_search_vector_store_request( + vector_store_id="documents", + query="cleanup probe", + vector_store_search_optional_params={}, + litellm_logging_obj=MagicMock(), + litellm_params={ + "api_base": "http://milvus:19530", + "litellm_embedding_model": "embedding-alias", + }, + embedding_executor=executor, + ) + + assert exc_info.value.status_code == 400 + + @pytest.mark.asyncio + async def test_async_grpc_search_without_pymilvus_installed_is_a_400(self): + executor: Final = MagicMock() + executor.aembed = AsyncMock(return_value=MOCK_EMBEDDING_RESPONSE) + + with ( + patch.dict(sys.modules, {"pymilvus": None}), + pytest.raises(litellm.BadRequestError, match="pip install litellm") as exc_info, + ): + await MilvusGRPCVectorStoreConfig().aexecute_search_vector_store_request( + vector_store_id="documents", + query="cleanup probe", + vector_store_search_optional_params={}, + litellm_logging_obj=MagicMock(), + litellm_params={ + "api_base": "http://milvus:19530", + "litellm_embedding_model": "embedding-alias", + }, + embedding_executor=executor, + ) + + assert exc_info.value.status_code == 400 + + def test_grpc_search_maps_a_connect_failure_to_an_api_connection_error(self): + milvus: Final = pytest.importorskip("pymilvus") + client: Final = MagicMock() + client.search.side_effect = milvus.MilvusException( + code=2, + message="Fail connecting to server on REDACTED:19530, illegal connection params or server unavailable", + ) + executor: Final = MagicMock() + executor.embed.return_value = MOCK_EMBEDDING_RESPONSE + + with pytest.raises(litellm.APIConnectionError, match="Milvus gRPC connection failed") as exc_info: + MilvusGRPCVectorStoreConfig(sync_client=client).execute_search_vector_store_request( + vector_store_id="documents", + query="cleanup probe", + vector_store_search_optional_params={}, + litellm_logging_obj=MagicMock(), + litellm_params={ + "api_base": "http://milvus:19530", + "litellm_embedding_model": "embedding-alias", + }, + embedding_executor=executor, + ) + + assert "api_key holds a valid 'user:password' token" in str(exc_info.value) + assert "illegal connection params or server unavailable" in str(exc_info.value) + + @pytest.mark.asyncio + async def test_async_grpc_search_maps_a_connect_failure_to_an_api_connection_error(self): + milvus: Final = pytest.importorskip("pymilvus") + client: Final = MagicMock() + client.search = AsyncMock( + side_effect=milvus.MilvusException( + code=2, + message="Fail connecting to server on REDACTED:19530, illegal connection params or server unavailable", + ) + ) + executor: Final = MagicMock() + executor.aembed = AsyncMock(return_value=MOCK_EMBEDDING_RESPONSE) + + with pytest.raises(litellm.APIConnectionError, match="Milvus gRPC connection failed") as exc_info: + await MilvusGRPCVectorStoreConfig(async_client=client).aexecute_search_vector_store_request( + vector_store_id="documents", + query="cleanup probe", + vector_store_search_optional_params={}, + litellm_logging_obj=MagicMock(), + litellm_params={ + "api_base": "http://milvus:19530", + "litellm_embedding_model": "embedding-alias", + }, + embedding_executor=executor, + ) + + assert "api_base points at a reachable gRPC endpoint" in str(exc_info.value) + + def test_grpc_search_leaves_a_non_connect_milvus_error_alone(self): + milvus: Final = pytest.importorskip("pymilvus") + client: Final = MagicMock() + client.search.side_effect = milvus.MilvusException( + code=100, + message="collection not found[database=default][collection=documents]", + ) + executor: Final = MagicMock() + executor.embed.return_value = MOCK_EMBEDDING_RESPONSE + + with pytest.raises(milvus.MilvusException, match="collection not found") as exc_info: + MilvusGRPCVectorStoreConfig(sync_client=client).execute_search_vector_store_request( + vector_store_id="documents", + query="cleanup probe", + vector_store_search_optional_params={}, + litellm_logging_obj=MagicMock(), + litellm_params={ + "api_base": "http://milvus:19530", + "litellm_embedding_model": "embedding-alias", + }, + embedding_executor=executor, + ) + + assert "Milvus gRPC connection failed" not in str(exc_info.value) + # @pytest.mark.parametrize("sync_mode", [True, False]) # @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py index 5143bb4155b..d62059b6e6d 100644 --- a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py +++ b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py @@ -8,7 +8,7 @@ import pytest from fastapi import HTTPException, Request import litellm -from litellm.constants import MILVUS_ADMIN_CONFIGURED_CONNECTION +from litellm.constants import MILVUS_ADMIN_CONFIGURED_CONNECTION, REDACTED_BY_LITELM_STRING from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import ( LiteLLM_ManagedVectorStore, VectorStorePreCallHook, @@ -959,6 +959,7 @@ async def test_config_vector_store_id_cannot_be_updated_in_database(): registry = VectorStoreRegistry() registry.config_vector_store_ids = frozenset(("configured",)) prisma_client = MagicMock() + prisma_client.db.litellm_managedvectorstorestable.find_unique = AsyncMock(return_value=None) with ( patch.object( # test-quality-ok: update checks collisions against the process-wide registry @@ -982,10 +983,96 @@ async def test_config_vector_store_id_cannot_be_updated_in_database(): ) assert exc_info.value.status_code == 400 - prisma_client.db.litellm_managedvectorstorestable.find_unique.assert_not_called() prisma_client.db.litellm_managedvectorstorestable.update.assert_not_called() +@pytest.mark.asyncio +async def test_database_row_sharing_a_config_vector_store_id_can_still_be_updated(): + from litellm.proxy.vector_store_endpoints.management_endpoints import update_vector_store + from litellm.types.vector_stores import VectorStoreUpdateRequest + + registry = VectorStoreRegistry() + registry.config_vector_store_ids = frozenset(("configured",)) + existing_row = MagicMock() + existing_row.model_dump = MagicMock( + return_value={ + "vector_store_id": "configured", + "custom_llm_provider": "milvus", + "litellm_params": {"api_base": "https://stored-milvus:19530"}, + } + ) + updated_row = MagicMock() + updated_row.model_dump = MagicMock(return_value={"vector_store_id": "configured"}) + prisma_client = MagicMock() + prisma_client.db.litellm_managedvectorstorestable.find_unique = AsyncMock(return_value=existing_row) + prisma_client.db.litellm_managedvectorstorestable.update = AsyncMock(return_value=updated_row) + + with ( + patch.object( # test-quality-ok: update checks collisions against the process-wide registry + litellm, "vector_store_registry", registry + ), + patch( # test-quality-ok: the endpoint reads the proxy database singleton directly + "litellm.proxy.proxy_server.prisma_client", prisma_client + ), + patch( # test-quality-ok: feature entitlement is outside the collision behavior + "litellm.proxy.vector_store_endpoints.management_endpoints.check_feature_access_for_user", + new=AsyncMock(), + ), + ): + await update_vector_store( + data=VectorStoreUpdateRequest( + vector_store_id="configured", + vector_store_description="replacement", + ), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + update_call = prisma_client.db.litellm_managedvectorstorestable.update.call_args + assert update_call.kwargs["where"] == {"vector_store_id": "configured"} + assert update_call.kwargs["data"]["vector_store_description"] == "replacement" + + +@pytest.mark.asyncio +async def test_database_row_sharing_a_config_vector_store_id_can_still_be_deleted(): + from litellm.proxy.vector_store_endpoints.management_endpoints import delete_vector_store + from litellm.types.vector_stores import VectorStoreDeleteRequest + + registry = VectorStoreRegistry() + registry.config_vector_store_ids = frozenset(("configured",)) + existing_row = MagicMock() + existing_row.model_dump = MagicMock( + return_value={ + "vector_store_id": "configured", + "custom_llm_provider": "milvus", + } + ) + prisma_client = MagicMock() + prisma_client.db.litellm_managedvectorstorestable.find_unique = AsyncMock(return_value=existing_row) + prisma_client.db.litellm_managedvectorstorestable.delete = AsyncMock(return_value=existing_row) + + with ( + patch.object( # test-quality-ok: the endpoint reads the process-wide registry directly + litellm, "vector_store_registry", registry + ), + patch( # test-quality-ok: the endpoint reads the proxy database singleton directly + "litellm.proxy.proxy_server.prisma_client", prisma_client + ), + patch( # test-quality-ok: feature entitlement is outside config ownership behavior + "litellm.proxy.vector_store_endpoints.management_endpoints.check_feature_access_for_user", + new=AsyncMock(), + ), + ): + response = await delete_vector_store( + data=VectorStoreDeleteRequest(vector_store_id="configured"), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert response["status"] == "success" + assert "configured" in response["message"] + delete_call = prisma_client.db.litellm_managedvectorstorestable.delete.call_args + assert delete_call.kwargs["where"] == {"vector_store_id": "configured"} + + @pytest.mark.asyncio async def test_config_vector_store_id_cannot_be_deleted(): from litellm.proxy.vector_store_endpoints.management_endpoints import delete_vector_store @@ -1103,17 +1190,19 @@ def test_non_grpc_connection_update_drops_forged_admin_marker(): @pytest.mark.parametrize( - ("custom_llm_provider", "litellm_params", "credential"), + ("custom_llm_provider", "litellm_params", "credential", "reuses_stored_credentials"), ( - ("openai", {"api_key": "sk-real", "api_base": "https://attacker.example"}, None), - ("openai", {"api_key": "sk-real"}, None), - ("openai", {"api_key": "sk-real", "api_base": "https://old.example", "organization": "org-other"}, None), - ("bedrock", None, None), - ("openai", None, "someone-elses-credential"), + ("bedrock", None, None, False), + ("openai", None, "someone-elses-credential", False), + ("openai", {"api_base": "https://attacker.example"}, None, True), + ("openai", {"api_key": REDACTED_BY_LITELM_STRING, "api_base": "https://attacker.example"}, None, True), ), ) -def test_non_admin_cannot_change_a_non_grpc_connection_on_update( - custom_llm_provider: str, litellm_params: dict[str, object] | None, credential: str | None +def test_non_admin_cannot_change_a_connection_that_keeps_its_stored_credentials( + custom_llm_provider: str, + litellm_params: dict[str, object] | None, + credential: str | None, + reuses_stored_credentials: bool, ) -> None: with pytest.raises(HTTPException) as exc_info: prepare_vector_store_connection_for_persistence( @@ -1124,10 +1213,48 @@ def test_non_admin_cannot_change_a_non_grpc_connection_on_update( existing_litellm_params={"api_key": "sk-real", "api_base": "https://old.example"}, litellm_credential_name=credential, litellm_credential_name_supplied=credential is not None, + reuses_stored_credentials=reuses_stored_credentials, ) assert exc_info.value.status_code == 403 - assert "Only proxy admins can configure vector store connections" in exc_info.value.detail + assert "keeps its stored credentials" in exc_info.value.detail + + +def test_non_admin_cannot_repoint_a_connection_that_keeps_its_stored_credential_name() -> None: + with pytest.raises(HTTPException) as exc_info: + prepare_vector_store_connection_for_persistence( + custom_llm_provider="openai", + litellm_params={"api_base": "https://attacker.example"}, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER), + existing_custom_llm_provider="openai", + existing_litellm_params={"api_base": "https://old.example"}, + existing_litellm_credential_name="prod-openai", + ) + + assert exc_info.value.status_code == 403 + assert "keeps its stored credentials" in exc_info.value.detail + + +@pytest.mark.parametrize( + "litellm_params", + ( + {"api_key": "sk-rotated", "api_base": "https://old.example"}, + {"api_key": "sk-rotated"}, + {"api_key": "sk-rotated", "api_base": "https://new.example", "organization": "org-other"}, + ), +) +def test_non_admin_can_change_a_non_grpc_connection_it_supplies_in_full( + litellm_params: dict[str, object], +) -> None: + params = prepare_vector_store_connection_for_persistence( + custom_llm_provider="openai", + litellm_params=litellm_params, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER), + existing_custom_llm_provider="openai", + existing_litellm_params={"api_key": "sk-real", "api_base": "https://old.example"}, + ) + + assert params == litellm_params @pytest.mark.parametrize("litellm_params", (None, {"api_key": "sk-real", "api_base": "https://old.example"})) @@ -1145,6 +1272,106 @@ def test_non_admin_update_leaving_the_non_grpc_connection_alone_is_allowed( assert params == {"api_key": "sk-real", "api_base": "https://old.example"} +_UNAPPROVED_GRPC_PARAMS: Final = { + "milvus_transport": "grpc", + "api_base": "http://internal-milvus:19530", + "litellm_embedding_model": "embedding-alias", +} + + +@pytest.mark.parametrize( + "litellm_params", + ( + None, + _UNAPPROVED_GRPC_PARAMS, + {**_UNAPPROVED_GRPC_PARAMS, MILVUS_ADMIN_CONFIGURED_CONNECTION: True}, + ), +) +def test_non_admin_no_change_update_keeps_an_unapproved_grpc_row_unapproved( + litellm_params: dict[str, object] | None, +) -> None: + params = prepare_vector_store_connection_for_persistence( + custom_llm_provider="milvus", + litellm_params=litellm_params, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER), + existing_custom_llm_provider="milvus", + existing_litellm_params=_UNAPPROVED_GRPC_PARAMS, + ) + + assert params == _UNAPPROVED_GRPC_PARAMS + + +@pytest.mark.parametrize("litellm_params", (None, _UNAPPROVED_GRPC_PARAMS)) +def test_non_admin_no_change_update_keeps_an_admin_approved_grpc_row_approved( + litellm_params: dict[str, object] | None, +) -> None: + approved: Final = {**_UNAPPROVED_GRPC_PARAMS, MILVUS_ADMIN_CONFIGURED_CONNECTION: True} + params = prepare_vector_store_connection_for_persistence( + custom_llm_provider="milvus", + litellm_params=litellm_params, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER), + existing_custom_llm_provider="milvus", + existing_litellm_params=approved, + ) + + assert params == approved + + +@pytest.mark.parametrize("litellm_params", (None, _UNAPPROVED_GRPC_PARAMS)) +def test_admin_no_change_update_approves_an_unapproved_grpc_row( + litellm_params: dict[str, object] | None, +) -> None: + params = prepare_vector_store_connection_for_persistence( + custom_llm_provider="milvus", + litellm_params=litellm_params, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + existing_custom_llm_provider="milvus", + existing_litellm_params=_UNAPPROVED_GRPC_PARAMS, + ) + + assert params == {**_UNAPPROVED_GRPC_PARAMS, MILVUS_ADMIN_CONFIGURED_CONNECTION: True} + + +@pytest.mark.parametrize( + "litellm_params", + ( + {**_UNAPPROVED_GRPC_PARAMS, "api_base": "http://attacker-milvus:19530"}, + {**_UNAPPROVED_GRPC_PARAMS, "milvus_transport": "rest"}, + ), +) +def test_non_admin_cannot_change_an_admin_approved_grpc_connection(litellm_params: dict[str, object]) -> None: + with pytest.raises(HTTPException) as exc_info: + prepare_vector_store_connection_for_persistence( + custom_llm_provider="milvus", + litellm_params=litellm_params, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER), + existing_custom_llm_provider="milvus", + existing_litellm_params={**_UNAPPROVED_GRPC_PARAMS, MILVUS_ADMIN_CONFIGURED_CONNECTION: True}, + ) + + assert exc_info.value.status_code == 403 + assert "Only proxy admins can change a Milvus gRPC vector store connection" in exc_info.value.detail + + +def test_non_admin_cannot_create_a_grpc_store_while_rest_creation_stays_open() -> None: + rest_params = prepare_vector_store_connection_for_persistence( + custom_llm_provider="milvus", + litellm_params={"api_base": "http://internal-milvus:19530", "api_key": "root:Milvus"}, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER), + ) + assert rest_params == {"api_base": "http://internal-milvus:19530", "api_key": "root:Milvus"} + + with pytest.raises(HTTPException) as exc_info: + prepare_vector_store_connection_for_persistence( + custom_llm_provider="milvus", + litellm_params=_UNAPPROVED_GRPC_PARAMS, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER), + ) + + assert exc_info.value.status_code == 403 + assert "Only proxy admins can configure Milvus gRPC vector store connections" in exc_info.value.detail + + def test_non_admin_can_still_create_a_non_grpc_store(): params = prepare_vector_store_connection_for_persistence( custom_llm_provider="openai", @@ -2634,6 +2861,40 @@ async def test_new_vector_store_persists_embedding_reference_without_credentials assert "api_key" not in _serialize_litellm_params(response_vs.get("litellm_params")) +@pytest.mark.asyncio +async def test_new_vector_store_rejects_an_unsupported_milvus_transport(): + from litellm.types.vector_stores import LiteLLM_ManagedVectorStore + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_managedvectorstorestable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_managedvectorstorestable.create = AsyncMock() + + vector_store_data: LiteLLM_ManagedVectorStore = { + "vector_store_id": "test-store-bad-transport", + "custom_llm_provider": "milvus", + "litellm_params": {"api_base": "http://milvus:19530", "milvus_transport": "tcp"}, + } + mock_user_api_key = MagicMock(spec=UserAPIKeyAuth) + mock_user_api_key.user_role = LitellmUserRoles.PROXY_ADMIN + mock_user_api_key.team_id = None + mock_user_api_key.user_id = None + + with ( + patch( # test-quality-ok: the endpoint reads the proxy database singleton directly + "litellm.proxy.proxy_server.prisma_client", mock_prisma_client + ), + patch.object( # test-quality-ok: create checks collisions against the process-wide registry + litellm, "vector_store_registry", MagicMock() + ), + pytest.raises(HTTPException) as exc_info, + ): + await new_vector_store(vector_store=vector_store_data, user_api_key_dict=mock_user_api_key) + + assert exc_info.value.status_code == 400 + assert "milvus_transport" in exc_info.value.detail + mock_prisma_client.db.litellm_managedvectorstorestable.create.assert_not_awaited() + + @pytest.mark.asyncio async def test_new_vector_store_auto_resolves_from_router(): """Test that new_vector_store auto-resolves embedding config from router when model is config-defined.""" @@ -3223,6 +3484,68 @@ class TestUpdateVectorStoreAccessControlAndRedaction: assert exc_info.value.status_code == 403 prisma_client.db.litellm_managedvectorstorestable.update.assert_not_called() + @pytest.mark.parametrize( + "update", + ( + {"vector_store_description": "renamed by its owner"}, + { + "vector_store_description": "renamed by its owner", + "litellm_params": { + "milvus_transport": "grpc", + "api_base": "http://trusted-milvus:19530", + "litellm_embedding_model": "embedding-alias", + }, + }, + ), + ) + @pytest.mark.asyncio + async def test_non_admin_can_edit_an_unapproved_grpc_row_without_approving_it(self, update: dict[str, object]): + import json + + from litellm.proxy.vector_store_endpoints.management_endpoints import update_vector_store + from litellm.types.vector_stores import VectorStoreUpdateRequest + + existing_row = MagicMock() + existing_row.model_dump.return_value = { + "vector_store_id": "vs_owned", + "custom_llm_provider": "milvus", + "team_id": "team-A", + "litellm_params": { + "milvus_transport": "grpc", + "api_base": "http://trusted-milvus:19530", + "litellm_embedding_model": "embedding-alias", + }, + } + prisma_client = MagicMock() + prisma_client.db.litellm_managedvectorstorestable.find_unique = AsyncMock(return_value=existing_row) + prisma_client.db.litellm_managedvectorstorestable.update = AsyncMock(return_value=existing_row) + + with ( + patch( # test-quality-ok: feature entitlement is outside connection authorization behavior + "litellm.proxy.vector_store_endpoints.management_endpoints.check_feature_access_for_user", + new=AsyncMock(), + ), + patch( # test-quality-ok: the endpoint reads the proxy database singleton directly + "litellm.proxy.proxy_server.prisma_client", prisma_client + ), + patch.object( # test-quality-ok: registry synchronization is outside persistence behavior + litellm, "vector_store_registry", None + ), + ): + await update_vector_store( + data=VectorStoreUpdateRequest(vector_store_id="vs_owned", **update), + user_api_key_dict=UserAPIKeyAuth( + user_id="owner", + team_id="team-A", + user_role=LitellmUserRoles.INTERNAL_USER, + ), + ) + + update_data = prisma_client.db.litellm_managedvectorstorestable.update.await_args.kwargs["data"] + assert update_data["vector_store_description"] == "renamed by its owner" + persisted_params = json.loads(update_data.get("litellm_params", "{}")) + assert MILVUS_ADMIN_CONFIGURED_CONNECTION not in persisted_params + @pytest.mark.parametrize( "update", ( @@ -3317,7 +3640,7 @@ class TestUpdateVectorStoreAccessControlAndRedaction: ) assert exc_info.value.status_code == 403 - assert "Only proxy admins can configure vector store connections" in exc_info.value.detail + assert "keeps its stored credentials" in exc_info.value.detail mock_prisma_client.db.litellm_managedvectorstorestable.update.assert_not_called() @pytest.mark.asyncio @@ -3470,8 +3793,12 @@ class TestUpdateVectorStoreAccessControlAndRedaction: mock_prisma_client.db.litellm_managedvectorstorestable.update = AsyncMock(return_value=updated_row) with ( - patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), - patch("litellm.vector_store_registry", None), + patch( # test-quality-ok: the endpoint reads the proxy database singleton directly + "litellm.proxy.proxy_server.prisma_client", mock_prisma_client + ), + patch( # test-quality-ok: registry synchronization is outside redaction-restore behavior + "litellm.vector_store_registry", None + ), ): await update_vector_store( data=VectorStoreUpdateRequest( diff --git a/tests/test_litellm/vector_stores/test_vector_store_registry.py b/tests/test_litellm/vector_stores/test_vector_store_registry.py index 1f78a707104..b0b4417752f 100644 --- a/tests/test_litellm/vector_stores/test_vector_store_registry.py +++ b/tests/test_litellm/vector_stores/test_vector_store_registry.py @@ -118,6 +118,45 @@ def test_fresh_registries_do_not_share_config_loaded_stores(): assert VectorStoreRegistry().get_litellm_managed_vector_store_from_registry("configured") is None +def _milvus_config_entry(transport: object, include_transport: bool = True) -> dict[str, object]: + return { + "vector_store_name": "configured-milvus", + "litellm_params": { + "vector_store_id": "configured-milvus", + "custom_llm_provider": "milvus", + "api_base": "http://milvus:19530", + **({"milvus_transport": transport} if include_transport else {}), + }, + } + + +@pytest.mark.parametrize("transport", ["tcp", "REST", "gRPC", "", "http"]) +def test_config_load_rejects_an_unsupported_milvus_transport(transport: str) -> None: + with pytest.raises(ValueError, match="milvus_transport must be one of rest, grpc"): + VectorStoreRegistry().load_vector_stores_from_config([_milvus_config_entry(transport)]) + + +@pytest.mark.parametrize("transport", ["rest", "grpc"]) +def test_config_load_keeps_the_supported_milvus_transports(transport: str) -> None: + registry: Final = VectorStoreRegistry() + + registry.load_vector_stores_from_config([_milvus_config_entry(transport)]) + + stored: Final = registry.get_litellm_managed_vector_store_from_registry("configured-milvus") + assert stored is not None + assert (stored.get("litellm_params") or {}).get("milvus_transport") == transport + + +def test_config_load_allows_a_milvus_store_without_a_transport() -> None: + registry: Final = VectorStoreRegistry() + + registry.load_vector_stores_from_config([_milvus_config_entry(None, include_transport=False)]) + + stored: Final = registry.get_litellm_managed_vector_store_from_registry("configured-milvus") + assert stored is not None + assert "milvus_transport" not in (stored.get("litellm_params") or {}) + + def test_add_vector_store_to_registry(): """Test that add_vector_store_to_registry adds vector store correctly when there are pre-existing stores""" # Create pre-existing vector stores diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreForm.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreForm.integration.test.tsx index 3b9c714d354..511381716c8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreForm.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreForm.integration.test.tsx @@ -119,14 +119,41 @@ describe("VectorStoreForm submit payload", () => { await submit(user); await vi.waitFor(() => expect(mockCreate).toHaveBeenCalledTimes(1)); - expect(createdPayload().litellm_params).toStrictEqual({ + const expectedRestParams = { api_key: "user:pass", api_base: "https://milvus.example.com", litellm_embedding_model: "text-embedding-3-small", - }); + milvus_transport: "rest", + }; + expect(createdPayload().litellm_params).toStrictEqual(expectedRestParams); expect(createdPayload().custom_llm_provider).toBe("milvus"); }); + it("sends the milvus gRPC transport the admin picked", async () => { + const user = setupUser(); + renderForm(); + + await chooseProvider(user, "Milvus"); + await user.type(screen.getByPlaceholderText("Enter vector store ID from your provider"), "vs-milvus-grpc"); + await user.type(screen.getByPlaceholderText("username:password or api key"), "root:Milvus"); + await user.type( + screen.getByPlaceholderText("https://your-milvus-endpoint.com/"), + "http://milvus.example.com:19530", + ); + await chooseFromSelect(user, 1, "text-embedding-3-small"); + await chooseFromSelect(user, 2, "gRPC"); + await submit(user); + + await vi.waitFor(() => expect(mockCreate).toHaveBeenCalledTimes(1)); + const expectedGrpcParams = { + api_key: "root:Milvus", + api_base: "http://milvus.example.com:19530", + litellm_embedding_model: "text-embedding-3-small", + milvus_transport: "grpc", + }; + expect(createdPayload().litellm_params).toStrictEqual(expectedGrpcParams); + }); + it("sends a provider field's seeded default even when the user never touches it", async () => { const user = setupUser(); renderForm(); 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 84a9314ecce..f107f929d0a 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 @@ -106,17 +106,33 @@ describe("buildVectorStoreLitellmParams", () => { }); it("renames embedding_model to litellm_embedding_model for milvus", () => { - const params = buildVectorStoreLitellmParams("milvus", { + const formValues = { api_key: "user:pass", api_base: "https://my-milvus-endpoint.com/", embedding_model: "text-embedding-3-small", - }); + milvus_transport: "rest", + }; + const params = buildVectorStoreLitellmParams("milvus", formValues); - expect(params).toEqual({ + const expectedParams = { api_key: "user:pass", api_base: "https://my-milvus-endpoint.com/", litellm_embedding_model: "text-embedding-3-small", - }); + milvus_transport: "rest", + }; + expect(params).toEqual(expectedParams); + }); + + it("carries the chosen milvus transport into litellm_params", () => { + const formValues = { + api_key: "user:pass", + api_base: "http://my-milvus-endpoint.com:19530", + embedding_model: "text-embedding-3-small", + milvus_transport: "grpc", + }; + const params = buildVectorStoreLitellmParams("milvus", formValues); + + expect(params.milvus_transport).toBe("grpc"); }); it("renames embedding_model to litellm_embedding_model for mongodb", () => { 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 61da25874a5..6ad76680aa4 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 @@ -67,6 +67,7 @@ const PROVIDER_FIELD_NAMES = [ "vertex_collection_id", "vertex_engine_id", "embedding_model", + "milvus_transport", "vector_bucket_name", "index_name", "aws_region_name", @@ -104,6 +105,7 @@ const vectorStoreShape = { vertex_collection_id: optionalText, vertex_engine_id: optionalText, embedding_model: optionalText, + milvus_transport: optionalText, vector_bucket_name: optionalText, index_name: optionalText, aws_region_name: optionalText, @@ -153,6 +155,7 @@ const EMPTY_VALUES: VectorStoreFormValues = { custom_llm_provider: "bedrock", vector_store_id: "", vertex_location: "global", + milvus_transport: "rest", mongodb_embedding_field: "embedding", mongodb_text_field: "text", valkey_port: "6379", diff --git a/ui/litellm-dashboard/src/components/vector_store_providers.tsx b/ui/litellm-dashboard/src/components/vector_store_providers.tsx index a75f10771a8..9405768f8bf 100644 --- a/ui/litellm-dashboard/src/components/vector_store_providers.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_providers.tsx @@ -159,7 +159,8 @@ export const vectorStoreProviderFields: Record { name: "api_base", label: "API Base", - tooltip: "Enter your Milvus endpoint (e.g., https://your-milvus-endpoint.com/)", + tooltip: + "Enter your Milvus endpoint. The REST transport takes an HTTP address (e.g., https://your-milvus-endpoint.com/) and the gRPC transport takes the gRPC address, port 19530 by default", placeholder: "https://your-milvus-endpoint.com/", required: true, type: "text", @@ -172,6 +173,19 @@ export const vectorStoreProviderFields: Record required: true, type: "select", }, + { + name: "milvus_transport", + label: "Transport", + tooltip: + "How LiteLLM talks to Milvus. REST works with any deployment, and gRPC uses the PyMilvus SDK, which needs the litellm[milvus] extra installed", + required: false, + type: "select", + options: [ + { value: "rest", label: "REST (default)" }, + { value: "grpc", label: "gRPC" }, + ], + initialValue: "rest", + }, ], mongodb: [ {