diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index 6885183aeba..4924c5a5acc 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -226,6 +226,9 @@ class BaseVectorStoreConfig: ``get_complete_url`` adds a query string, like Azure, override this.""" return f"{api_base}/{encoded_vector_store_id}" + def get_vector_store_list_url(self, api_base: str) -> str: + return api_base + def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: from ..chat.transformation import BaseLLMException diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index afa12b58b29..00833390542 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -11055,7 +11055,7 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - url: Final = api_base + url: Final = vector_store_provider_config.get_vector_store_list_url(api_base) params: Final[dict[str, object]] = {} if after is not None: @@ -11134,7 +11134,7 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - url: Final = api_base + url: Final = vector_store_provider_config.get_vector_store_list_url(api_base) params: Final[dict[str, object]] = {} if after is not None: diff --git a/litellm/llms/opensearch/vector_stores/transformation.py b/litellm/llms/opensearch/vector_stores/transformation.py index fbef67df368..2d862bc8ea6 100644 --- a/litellm/llms/opensearch/vector_stores/transformation.py +++ b/litellm/llms/opensearch/vector_stores/transformation.py @@ -11,6 +11,7 @@ from pydantic import ConfigDict, Field, TypeAdapter, ValidationError from litellm._logging import verbose_logger from litellm.exceptions import AuthenticationError, BadRequestError +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.litellm_core_utils.optional_imports import ensure_optional_import from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.vector_store.transformation import ( @@ -51,7 +52,7 @@ _QUERY_VECTOR: Final = TypeAdapter(tuple[float, ...]) _ATTRIBUTE_VALUE_TYPES: Final = (str, int, float, bool) _SEARCH_ONLY_MESSAGE: Final = ( "The OpenSearch vector store is search-only. Create the k-NN index in OpenSearch, " - "then register it here with the index name as vector_store_id." + "then register it here with its name as opensearch_index_name or vector_store_id." ) _PER_ACCOUNT_MESSAGE: Final = ( "api_base is an OpenSearch Serverless per-account endpoint, which needs a collection header the vector store " @@ -92,6 +93,7 @@ class _OpenSearchParams(LiteLLMBaseModel): litellm_embedding_model: str | None = None opensearch_embedding_field: str | None = None + opensearch_index_name: str | None = None opensearch_text_field: str | None = None @property @@ -102,6 +104,9 @@ class _OpenSearchParams(LiteLLMBaseModel): def text_field(self) -> str: return self.opensearch_text_field or DEFAULT_TEXT_FIELD + def index_name(self, vector_store_id: str) -> str: + return self.opensearch_index_name or vector_store_id + def require_embedding_model(self) -> str: if not self.litellm_embedding_model: raise _bad_request( @@ -221,8 +226,10 @@ class OpenSearchVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, SignsRequ return {**headers, "Content-Type": "application/json"} def _signing_region(self, params: Mapping[str, object], api_base: str) -> str: - configured_region: Final = _SigningRegion.model_validate(params).aws_region_name - return self._aws.get_aws_region_name_for_non_llm_api_calls(configured_region or aws_signing_target(api_base)[1]) + host_region: Final = aws_signing_target(api_base)[1] + return self._aws.get_aws_region_name_for_non_llm_api_calls( + host_region or _SigningRegion.model_validate(params).aws_region_name + ) def _signing_credentials(self, params: Mapping[str, object], region: str) -> "Credentials": if params.get("api_key"): @@ -260,6 +267,9 @@ class OpenSearchVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, SignsRequ def get_vector_store_url(self, api_base: str, encoded_vector_store_id: str) -> NoReturn: raise _bad_request(_SEARCH_ONLY_MESSAGE) + def get_vector_store_list_url(self, api_base: str) -> NoReturn: + raise _bad_request(_SEARCH_ONLY_MESSAGE) + @staticmethod def _params( litellm_params: Mapping[str, object], @@ -280,16 +290,17 @@ class OpenSearchVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, SignsRequ options: Final = _SearchOptions.model_validate(optional_params) except ValidationError as error: raise _bad_request(f"Invalid vector store search parameters: {error}") from None - if options.unsupported: + if options.unsupported and not drops_unsupported_params(litellm_params.get("drop_params")): raise _bad_request( f"The OpenSearch vector store does not support {', '.join(options.unsupported)} yet. " - "Remove it from the request." + "Remove it from the request, or set drop_params to drop it." ) try: return _OpenSearchParams.model_validate(litellm_params), options.k except ValidationError: raise _bad_request( - "opensearch_embedding_field, opensearch_text_field, and litellm_embedding_model must be strings." + "opensearch_embedding_field, opensearch_index_name, opensearch_text_field, and litellm_embedding_model " + "must be strings." ) from None @staticmethod @@ -302,7 +313,7 @@ class OpenSearchVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, SignsRequ ) -> tuple[str, dict[str, object]]: # mutable-ok: the provider contract returns a writable JSON request body vector: Final = _QUERY_VECTOR.validate_python(embedding.data[0]["embedding"]) # pyright: ignore[reportUnknownMemberType] # EmbeddingResponse.data is an untyped list, the vector is validated here return ( - f"{api_base}/{quote(vector_store_id, safe='')}/_search", + f"{api_base}/{quote(params.index_name(vector_store_id), safe='')}/_search", { "size": k, "query": {"knn": {params.embedding_field: {"vector": vector, "k": k}}}, diff --git a/tests/unit/llms/opensearch/vector_stores/test_opensearch_transformation.py b/tests/unit/llms/opensearch/vector_stores/test_opensearch_transformation.py index 7f33b234d1a..07836d57af0 100644 --- a/tests/unit/llms/opensearch/vector_stores/test_opensearch_transformation.py +++ b/tests/unit/llms/opensearch/vector_stores/test_opensearch_transformation.py @@ -156,7 +156,7 @@ async def _search( "api_base,region_params,service,region", [ (SERVERLESS, {}, "aoss", "us-east-1"), - (SERVERLESS, {"aws_region_name": "us-west-2"}, "aoss", "us-west-2"), + (SERVERLESS, {"aws_region_name": "us-west-2"}, "aoss", "us-east-1"), ("https://search-docs-xyz.eu-west-1.es.amazonaws.com/", {}, "es", "eu-west-1"), ("https://vpc-docs-xyz.cn-north-1.es.amazonaws.com.cn", {}, "es", "cn-north-1"), ("https://abc123.us-gov-west-1.aoss-fips.amazonaws.com", {}, "aoss", "us-gov-west-1"), @@ -247,6 +247,7 @@ async def test_opensearch_error_keeps_its_status(asynchronous: bool, status: int ({"litellm_embedding_model": None}, "litellm_embedding_model"), ({"opensearch_text_feild": "passage"}, "opensearch_text_feild"), ({"opensearch_text_field": 7}, "opensearch_text_field"), + ({"opensearch_index_name": ["docs"]}, "opensearch_index_name"), ({"extra_body": {"min_score": 0.5}}, "extra_body"), ({"api_base": None}, "api_base"), ({"api_base": "abc123.us-east-1.aoss.amazonaws.com"}, "api_base"), @@ -352,6 +353,46 @@ def test_failing_credential_source_is_a_401_that_keeps_its_output_out_of_the_rep assert CREDENTIAL_SOURCE_OUTPUT not in str(error.value) +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.asyncio +async def test_opensearch_index_name_picks_the_index_so_the_store_id_can_differ(asynchronous: bool) -> None: + respond: Final = MagicMock(return_value=httpx.Response(200, json=_opensearch_hits("text", "embedding"))) + + response: Final = await _search( + asynchronous, + respond, + RecordingEmbeddingExecutor(), + {"vector_store_id": "east-docs", "opensearch_index_name": "docs/v2"}, + ) + + assert response["data"] == EXPECTED_RESULTS + assert str(_sent(respond)[0].url) == f"{SERVERLESS}/docs%2Fv2/_search" + + +@pytest.mark.parametrize( + "unsupported", + [ + {"filters": {"type": "eq", "key": "category", "value": "hr"}}, + {"ranking_options": {"score_threshold": 0.5}}, + {"rewrite_query": True}, + ], +) +@pytest.mark.asyncio +async def test_unsupported_search_params_are_dropped_under_drop_params(unsupported: Mapping[str, object]) -> None: + respond: Final = MagicMock(return_value=httpx.Response(200, json=_opensearch_hits("text", "embedding"))) + + response: Final = await _search( + True, respond, RecordingEmbeddingExecutor(), {**unsupported, "drop_params": True, "max_num_results": 3} + ) + + assert response["data"] == EXPECTED_RESULTS + assert json.loads(_sent(respond)[0].content) == { + "size": 3, + "query": {"knn": {"embedding": {"vector": QUERY_VECTOR, "k": 3}}}, + "_source": {"excludes": ["embedding"]}, + } + + @pytest.mark.asyncio async def test_no_op_values_of_unsupported_params_are_accepted() -> None: def respond(request: httpx.Request) -> httpx.Response: @@ -394,6 +435,28 @@ async def test_store_operations_besides_search_send_nothing_to_the_index(asynchr respond.assert_not_called() +async def _list_stores(asynchronous: bool, respond: Callable[[httpx.Request], httpx.Response]) -> object: + request_kwargs: Final = {**BASE_PARAMS, "custom_llm_provider": "opensearch"} + if asynchronous: + return await vector_stores.alist( + **request_kwargs, client=AsyncHTTPHandler(transport=httpx.MockTransport(respond)) + ) + with httpx.Client(transport=httpx.MockTransport(respond)) as transport: + return vector_stores.list(**request_kwargs, client=HTTPHandler(client=transport)) + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.asyncio +async def test_listing_stores_sends_nothing_to_the_endpoint(asynchronous: bool) -> None: + respond: Final = MagicMock(return_value=httpx.Response(200, json={"cluster_name": "docs", "tagline": "OpenSearch"})) + + with pytest.raises(litellm.BadRequestError, match="search-only") as error: + await _list_stores(asynchronous, respond) + + assert error.value.status_code == 400 + respond.assert_not_called() + + @pytest.mark.parametrize( "response,status", [ 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 42664f54c97..f571ec6b9a0 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 @@ -154,6 +154,7 @@ describe("VectorStoreForm submit payload", () => { litellm_embedding_model: "text-embedding-3-small", opensearch_embedding_field: "embedding", opensearch_text_field: "text", + opensearch_index_name: undefined, }; await vi.waitFor(() => expect(mockCreate).toHaveBeenCalledTimes(1)); @@ -162,6 +163,26 @@ describe("VectorStoreForm submit payload", () => { expect(createdPayload().litellm_params).toStrictEqual(expectedLitellmParams); }); + it("registers an opensearch index under its own name when the vector store id differs", async () => { + const user = setupUser(); + renderForm(); + + await chooseProvider(user, "Amazon OpenSearch"); + fireEvent.change(screen.getByPlaceholderText("my-knn-index (OpenSearch index name)"), { + target: { value: "east-docs" }, + }); + fireEvent.change(screen.getByPlaceholderText("https://abc123.us-east-1.aoss.amazonaws.com"), { + target: { value: "https://yzatq55tpsuvlwozcal4.us-east-1.aoss.amazonaws.com" }, + }); + fireEvent.change(screen.getByPlaceholderText("my-knn-index"), { target: { value: "docs" } }); + await chooseFromSelect(user, 1, "text-embedding-3-small"); + await submit(user); + + await vi.waitFor(() => expect(mockCreate).toHaveBeenCalledTimes(1)); + expect(createdPayload().vector_store_id).toBe("east-docs"); + expect(createdPayload().litellm_params).toMatchObject({ opensearch_index_name: "docs" }); + }); + it("blocks an opensearch submit that has no endpoint or embedding model", async () => { const user = setupUser(); renderForm(); 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 6658971b968..bf5b925eb16 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 @@ -85,6 +85,7 @@ const PROVIDER_FIELD_NAMES = [ "valkey_embedding_field", "opensearch_embedding_field", "opensearch_text_field", + "opensearch_index_name", ] as const; type ProviderFieldName = (typeof PROVIDER_FIELD_NAMES)[number]; @@ -123,6 +124,7 @@ const vectorStoreShape = { valkey_embedding_field: optionalText, opensearch_embedding_field: optionalText, opensearch_text_field: optionalText, + opensearch_index_name: optionalText, }; const vectorStoreSchema = z.object(vectorStoreShape).superRefine((values, ctx) => { @@ -396,7 +398,7 @@ const VectorStoreForm: React.FC = ({
  1. The index has a knn_vector field holding each document's embedding. Enter the index name as - the Vector Store ID + the Vector Store ID, or under Index Name when another endpoint already uses that ID
  2. You know which embedding model created those embeddings. That model must be added to this proxy diff --git a/ui/litellm-dashboard/src/components/vector_store_providers.test.tsx b/ui/litellm-dashboard/src/components/vector_store_providers.test.tsx index 151f3549ad8..b4d890475c7 100644 --- a/ui/litellm-dashboard/src/components/vector_store_providers.test.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_providers.test.tsx @@ -82,6 +82,7 @@ describe("getVectorStoreProviderLogoAndName", () => { "embedding_model", "opensearch_embedding_field", "opensearch_text_field", + "opensearch_index_name", ]); expect(byName("api_base")).toMatchObject({ required: true }); expect(byName("aws_region_name")).toMatchObject({ required: false }); @@ -89,6 +90,8 @@ describe("getVectorStoreProviderLogoAndName", () => { expect(byName("embedding_model")).not.toHaveProperty("options"); expect(byName("opensearch_embedding_field")).toMatchObject({ required: false, initialValue: "embedding" }); expect(byName("opensearch_text_field")).toMatchObject({ required: false, initialValue: "text" }); + expect(byName("opensearch_index_name")).toMatchObject({ required: false }); + expect(byName("opensearch_index_name")).not.toHaveProperty("initialValue"); }); it("registers valkey in the provider, logo, and field maps", () => { diff --git a/ui/litellm-dashboard/src/components/vector_store_providers.tsx b/ui/litellm-dashboard/src/components/vector_store_providers.tsx index ec247b243b5..e953d6e3db6 100644 --- a/ui/litellm-dashboard/src/components/vector_store_providers.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_providers.tsx @@ -365,6 +365,15 @@ export const vectorStoreProviderFields: Record type: "text", initialValue: "text", }, + { + name: "opensearch_index_name", + label: "Index Name (optional)", + tooltip: + "The OpenSearch index to search. Leave blank to use the Vector Store ID, and set it to register same-named indexes from two endpoints under different IDs", + placeholder: "my-knn-index", + required: false, + type: "text", + }, ], s3_vectors: [ {