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: [ {