diff --git a/litellm/llms/opensearch/vector_stores/transformation.py b/litellm/llms/opensearch/vector_stores/transformation.py index 904e77a730b..809adfc43f5 100644 --- a/litellm/llms/opensearch/vector_stores/transformation.py +++ b/litellm/llms/opensearch/vector_stores/transformation.py @@ -45,11 +45,16 @@ _AWS_HOSTS: Final = ( re.compile(r"\.(?P[a-z0-9-]+)\.(?Paoss|aoss-fips|es)\.amazonaws\.com(?:\.cn)?\Z"), re.compile(r"\.(?Paoss|aos)\.(?P[a-z0-9-]+)\.on\.aws\Z"), ) +_ACCOUNT_ID_LABEL: Final = re.compile(r"\A\d{12}\.") _QUERY_VECTOR: Final = TypeAdapter(tuple[float, ...]) _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." ) +_PER_ACCOUNT_MESSAGE: Final = ( + "api_base is an OpenSearch Serverless per-account endpoint, which needs a collection header the vector store " + "does not send. Use the collection endpoint instead, for example https://..aoss.amazonaws.com." +) _SIGV4_ONLY_MESSAGE: Final = ( "The OpenSearch vector store signs every request with AWS SigV4 and does not support api_key. " "Remove api_key and give the vector store AWS credentials instead." @@ -218,6 +223,9 @@ class OpenSearchVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, SignsRequ def _signing_credentials(self, params: Mapping[str, object], region: str) -> "Credentials": if params.get("api_key"): raise _bad_request(_SIGV4_ONLY_MESSAGE) + auth: Final = AwsAuthParams.model_validate(params) + if bool(auth.aws_access_key_id) != bool(auth.aws_secret_access_key): + raise _unauthenticated("aws_access_key_id and aws_secret_access_key must be set together") ensure_optional_import("botocore") from botocore.exceptions import ( CredentialRetrievalError, @@ -227,7 +235,7 @@ class OpenSearchVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, SignsRequ ) try: - return _usable_credentials(self._aws.resolve_credentials(AwsAuthParams.model_validate(params), region)) + return _usable_credentials(self._aws.resolve_credentials(auth, region)) except (CredentialRetrievalError, NoCredentialsError, PartialCredentialsError, ProfileNotFound) as error: raise _unauthenticated(str(error)) from None @@ -238,8 +246,13 @@ class OpenSearchVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, SignsRequ "api_base must be the OpenSearch endpoint URL, for example " "https://..aoss.amazonaws.com or https://." ) + if _ACCOUNT_ID_LABEL.match(parsed.hostname) and aws_signing_target(api_base or "")[0] == "aoss": + raise _bad_request(_PER_ACCOUNT_MESSAGE) return (api_base or "").rstrip("/") + def get_vector_store_url(self, api_base: str, encoded_vector_store_id: str) -> NoReturn: + raise _bad_request(_SEARCH_ONLY_MESSAGE) + @staticmethod def _params( litellm_params: Mapping[str, object], 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 4ffffba862d..11da283566b 100644 --- a/tests/unit/llms/opensearch/vector_stores/test_opensearch_transformation.py +++ b/tests/unit/llms/opensearch/vector_stores/test_opensearch_transformation.py @@ -21,6 +21,7 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.opensearch.vector_stores.transformation import OpenSearchVectorStoreConfig from litellm.types.utils import EmbeddingResponse from litellm.types.vector_stores import VectorStoreSearchResponse +from litellm.vector_stores import main as vector_stores ACCESS_KEY: Final = "AKIDOPENSEARCHTEST" SECRET_KEY: Final = "opensearch-test-secret" @@ -238,6 +239,7 @@ async def test_opensearch_error_keeps_its_status(asynchronous: bool, status: int ({"extra_body": {"min_score": 0.5}}, "extra_body"), ({"api_base": None}, "api_base"), ({"api_base": "abc123.us-east-1.aoss.amazonaws.com"}, "api_base"), + ({"api_base": "https://123456789012.aoss.us-east-1.on.aws"}, "per-account endpoint"), ({"api_key": "opensearch-basic-auth"}, "api_key"), ], ) @@ -279,6 +281,16 @@ _AMBIENT_AWS_ENV: Final = ( pytest.param({"aws_session_name": "opensearch-no-credentials"}, {}, id="none-anywhere"), pytest.param({"aws_profile_name": "opensearch-missing-profile"}, {}, id="missing-profile"), pytest.param({}, {"AWS_ACCESS_KEY_ID": "AKIDPARTIAL"}, id="key-without-secret"), + pytest.param( + {"aws_access_key_id": "AKIDSTOREONLY"}, + {"AWS_ACCESS_KEY_ID": "AKIDAMBIENT", "AWS_SECRET_ACCESS_KEY": "ambient-secret"}, + id="store-key-without-secret-beside-ambient-credentials", + ), + pytest.param( + {"aws_secret_access_key": "store-secret-only"}, + {"AWS_ACCESS_KEY_ID": "AKIDAMBIENT", "AWS_SECRET_ACCESS_KEY": "ambient-secret"}, + id="store-secret-without-key-beside-ambient-credentials", + ), ], ) @pytest.mark.asyncio @@ -331,6 +343,31 @@ def test_create_is_rejected_because_the_store_is_search_only() -> None: litellm.vector_stores.create(custom_llm_provider="opensearch") +async def _store_operation( + asynchronous: bool, operation: str, respond: Callable[[httpx.Request], httpx.Response] +) -> object: + request_kwargs: Final = {**BASE_PARAMS, "vector_store_id": "litellm-qa-docs", "custom_llm_provider": "opensearch"} + if asynchronous: + return await getattr(vector_stores, f"a{operation}")( + **request_kwargs, client=AsyncHTTPHandler(transport=httpx.MockTransport(respond)) + ) + with httpx.Client(transport=httpx.MockTransport(respond)) as transport: + return getattr(vector_stores, operation)(**request_kwargs, client=HTTPHandler(client=transport)) + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("operation", ["retrieve", "update", "delete"]) +@pytest.mark.asyncio +async def test_store_operations_besides_search_send_nothing_to_the_index(asynchronous: bool, operation: str) -> None: + respond: Final = MagicMock(return_value=httpx.Response(200, json={"acknowledged": True})) + + with pytest.raises(litellm.BadRequestError, match="search-only") as error: + await _store_operation(asynchronous, operation, respond) + + assert error.value.status_code == 400 + respond.assert_not_called() + + @pytest.mark.parametrize( "response,status", [