diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index 8efeca71aa9..8dc3fa9df94 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -1,7 +1,7 @@ from __future__ import annotations from abc import abstractmethod -from collections.abc import Mapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, runtime_checkable @@ -62,6 +62,32 @@ class LiteLLMVectorStoreEmbeddingExecutor: ) +class EmbedOnceExecutor: + def __init__(self, executor: VectorStoreEmbeddingExecutor) -> None: + self._executor: Final = executor + self._embedded: tuple[str, str, EmbeddingResponse] | None = None + + def _reused(self, model: str, query: str) -> EmbeddingResponse | None: + embedded: Final = self._embedded + return embedded[2] if embedded is not None and embedded[:2] == (model, query) else None + + def embed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse: + reused: Final = self._reused(model, query) + if reused is not None: + return reused + response: Final = self._executor.embed(model, query, configuration) + self._embedded = (model, query, response) + return response + + async def aembed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse: + reused: Final = self._reused(model, query) + if reused is not None: + return reused + response: Final = await self._executor.aembed(model, query, configuration) + self._embedded = (model, query, response) + return response + + _REQUEST_METADATA: Final = TypeAdapter(dict[str, object]) @@ -260,6 +286,36 @@ class BaseVectorStoreConfig: """ return headers, None + def resolve_search_params( + self, + vector_store_id: str, + vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + api_base: str, + litellm_params: Mapping[str, object], + headers: Mapping[str, object], + get: Callable[[str, Mapping[str, object]], httpx.Response], + ) -> VectorStoreSearchOptionalRequestParams: + return vector_store_search_optional_params + + async def aresolve_search_params( + self, + vector_store_id: str, + vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + api_base: str, + litellm_params: Mapping[str, object], + headers: Mapping[str, object], + get: Callable[[str, Mapping[str, object]], Awaitable[httpx.Response]], + ) -> VectorStoreSearchOptionalRequestParams: + return vector_store_search_optional_params + + def recover_rejected_search( + self, + rejection: httpx.Response, + vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + litellm_params: Mapping[str, object], + ) -> VectorStoreSearchOptionalRequestParams | BaseLLMException | None: + return None + def calculate_vector_store_cost( self, response: VectorStoreSearchResponse, diff --git a/litellm/llms/bedrock/web_identity_session_policy.py b/litellm/llms/bedrock/web_identity_session_policy.py index 31187953dbc..8f31cb23171 100644 --- a/litellm/llms/bedrock/web_identity_session_policy.py +++ b/litellm/llms/bedrock/web_identity_session_policy.py @@ -264,7 +264,9 @@ SAGEMAKER_SESSION_POLICY: Final = _policy( POLLY_SESSION_POLICY: Final = _policy(_allow("PollyLiteLLM", ("polly:SynthesizeSpeech",))) -OPENSEARCH_SESSION_POLICY: Final = _policy(_allow("OpenSearchLiteLLM", ("es:ESHttpPost", "aoss:APIAccessAll"))) +OPENSEARCH_SESSION_POLICY: Final = _policy( + _allow("OpenSearchLiteLLM", ("es:ESHttpGet", "es:ESHttpPost", "aoss:APIAccessAll")) +) OPENSEARCH_INDEX_SESSION_POLICY: Final = _policy( _allow( diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 9c3d428d2f6..4c4a9c3fec2 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -124,6 +124,7 @@ from litellm.llms.base_llm.vector_store.transformation import ( BaseDirectVectorStoreConfig, BaseQueryEmbeddingVectorStoreConfig, BaseVectorStoreConfig, + EmbedOnceExecutor, VectorStoreEmbeddingExecutor, ) from litellm.llms.base_llm.vector_store_files.transformation import ( @@ -437,6 +438,12 @@ def _send_with_goaway_retry(send: Callable[[], httpx.Response]) -> httpx.Respons return send() +def _vector_store_search_timeout(provider_config: BaseVectorStoreConfig) -> BaseLLMException: + return provider_config.get_error_class( + error_message="Vector store search exceeded the caller timeout.", status_code=408, headers=httpx.Headers() + ) + + def _is_sized_json_body(response: httpx.Response) -> bool: headers: Final = httpx.Headers(response.headers) media_types: Final = (value.split(";")[0].strip().lower() for value in headers.get_list("content-type")) @@ -10675,7 +10682,7 @@ class BaseLLMHTTPHandler: else: async_httpx_client = client - headers = await asyncio.to_thread( + headers: Final[dict[str, object]] = await asyncio.to_thread( vector_store_provider_config.validate_environment, headers=extra_headers or {}, litellm_params=litellm_params, @@ -10691,74 +10698,131 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - if isinstance(vector_store_provider_config, BaseQueryEmbeddingVectorStoreConfig): - ( - url, - request_body, - ) = await vector_store_provider_config.atransform_search_vector_store_request( - vector_store_id=vector_store_id, - query=query, - vector_store_search_optional_params=vector_store_search_optional_params, - api_base=api_base, - litellm_logging_obj=logging_obj, - litellm_params=MappingProxyType(dict(litellm_params, timeout=timeout)), - extra_body=extra_body, - embedding_executor=embedding_executor, - ) - else: - ( - url, - request_body, - ) = await vector_store_provider_config.atransform_search_vector_store_request( - vector_store_id=vector_store_id, - query=query, - vector_store_search_optional_params=vector_store_search_optional_params, - api_base=api_base, - litellm_logging_obj=logging_obj, - litellm_params=dict(litellm_params), - extra_body=extra_body, - ) - all_optional_params: Final[dict[str, object]] = dict(litellm_params) - all_optional_params.update(vector_store_search_optional_params or {}) - headers, signed_json_body = await run_off_loop_if_aws( - vector_store_provider_config, - vector_store_provider_config.sign_request, + async def lookup(url: str, lookup_headers: Mapping[str, object]) -> httpx.Response: + try: + return await _asend_with_goaway_retry( + partial( + async_httpx_client.get, + url, + headers=dict(lookup_headers), + timeout=timeout, + ) + ) + except httpx.TimeoutException: + raise _vector_store_search_timeout(vector_store_provider_config) from None + except (httpx.TransportError, TypeError, UnicodeEncodeError) as e: + raise self._handle_error(e=e, provider_config=vector_store_provider_config) + + search_params: Final = await vector_store_provider_config.aresolve_search_params( + vector_store_id=vector_store_id, + vector_store_search_optional_params=vector_store_search_optional_params, + api_base=api_base, + litellm_params=MappingProxyType(dict(litellm_params, timeout=timeout)), headers=headers, - optional_params=all_optional_params, - request_data=request_body, - api_base=url, + get=lookup, + ) + search_embedding_executor: Final = EmbedOnceExecutor( + BaseQueryEmbeddingVectorStoreConfig.query_embedding_executor(embedding_executor, None) ) - logging_obj.pre_call( - input="", - api_key="", - additional_args={ - "complete_input_dict": request_body, - "api_base": api_base, - "headers": headers, - }, - ) + async def send(params: VectorStoreSearchOptionalRequestParams) -> httpx.Response: + if isinstance(vector_store_provider_config, BaseQueryEmbeddingVectorStoreConfig): + ( + url, + request_body, + ) = await vector_store_provider_config.atransform_search_vector_store_request( + vector_store_id=vector_store_id, + query=query, + vector_store_search_optional_params=params, + api_base=api_base, + litellm_logging_obj=logging_obj, + litellm_params=MappingProxyType(dict(litellm_params, timeout=timeout)), + extra_body=extra_body, + embedding_executor=search_embedding_executor, + ) + else: + ( + url, + request_body, + ) = await vector_store_provider_config.atransform_search_vector_store_request( + vector_store_id=vector_store_id, + query=query, + vector_store_search_optional_params=params, + api_base=api_base, + litellm_logging_obj=logging_obj, + litellm_params=dict(litellm_params), + extra_body=extra_body, + ) + signed_headers, signed_json_body = await run_off_loop_if_aws( + vector_store_provider_config, + vector_store_provider_config.sign_request, + headers=headers, + optional_params={**dict(litellm_params), **(params or {})}, + request_data=request_body, + api_base=url, + ) - request_data: Final = json.dumps(request_body) if signed_json_body is None else signed_json_body + logging_obj.pre_call( + input="", + api_key="", + additional_args={ + "complete_input_dict": request_body, + "api_base": api_base, + "headers": signed_headers, + }, + ) + + request_data: Final = json.dumps(request_body) if signed_json_body is None else signed_json_body + + try: + return await _asend_with_goaway_retry( + partial( + async_httpx_client.post, url=url, headers=signed_headers, data=request_data, timeout=timeout + ) + ) + except httpx.TimeoutException: + raise _vector_store_search_timeout(vector_store_provider_config) from None + except httpx.HTTPStatusError: + raise + except Exception as e: + raise self._handle_error(e=e, provider_config=vector_store_provider_config) try: - response: Final = await _asend_with_goaway_retry( - partial(async_httpx_client.post, url=url, headers=headers, data=request_data, timeout=timeout) + response: Final = await send(search_params) + except httpx.HTTPStatusError as rejection: + retry_params: Final = self._vector_store_search_retry_params( + rejection, vector_store_provider_config, search_params, litellm_params + ) + try: + retry_response: Final = await send(retry_params) + except httpx.HTTPStatusError as retry_rejection: + raise self._handle_error(e=retry_rejection, provider_config=vector_store_provider_config) from None + return vector_store_provider_config.transform_search_vector_store_response( + response=retry_response, + litellm_logging_obj=logging_obj, ) - except httpx.TimeoutException: - raise vector_store_provider_config.get_error_class( - error_message="Vector store search exceeded the caller timeout.", - status_code=408, - headers=httpx.Headers(), - ) from None - except Exception as e: - raise self._handle_error(e=e, provider_config=vector_store_provider_config) return vector_store_provider_config.transform_search_vector_store_response( response=response, litellm_logging_obj=logging_obj, ) + def _vector_store_search_retry_params( + self, + rejection: httpx.HTTPStatusError, + vector_store_provider_config: BaseVectorStoreConfig, + vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + litellm_params: GenericLiteLLMParams, + ) -> VectorStoreSearchOptionalRequestParams: + recovery: Final = vector_store_provider_config.recover_rejected_search( + rejection.response, vector_store_search_optional_params, MappingProxyType(dict(litellm_params)) + ) + if isinstance(recovery, BaseLLMException): + raise recovery from None + if recovery is None: + raise self._handle_error(e=rejection, provider_config=vector_store_provider_config) + return recovery + def vector_store_search_handler( self, vector_store_id: str, @@ -10813,7 +10877,7 @@ class BaseLLMHTTPHandler: else: sync_httpx_client = client - headers = vector_store_provider_config.validate_environment( + headers: Final[dict[str, object]] = vector_store_provider_config.validate_environment( headers=extra_headers or {}, litellm_params=litellm_params ) @@ -10825,68 +10889,106 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - if isinstance(vector_store_provider_config, BaseQueryEmbeddingVectorStoreConfig): - ( - url, - request_body, - ) = vector_store_provider_config.transform_search_vector_store_request( - vector_store_id=vector_store_id, - query=query, - vector_store_search_optional_params=vector_store_search_optional_params, - api_base=api_base, - litellm_logging_obj=logging_obj, - litellm_params=MappingProxyType(dict(litellm_params, timeout=timeout)), - extra_body=extra_body, - embedding_executor=embedding_executor, - ) - else: - ( - url, - request_body, - ) = vector_store_provider_config.transform_search_vector_store_request( - vector_store_id=vector_store_id, - query=query, - vector_store_search_optional_params=vector_store_search_optional_params, - api_base=api_base, - litellm_logging_obj=logging_obj, - litellm_params=dict(litellm_params), - extra_body=extra_body, - ) + def lookup(url: str, lookup_headers: Mapping[str, object]) -> httpx.Response: + try: + return _send_with_goaway_retry( + partial( + sync_httpx_client.get, + url, + headers=dict(lookup_headers), + timeout=timeout, + ) + ) + except httpx.TimeoutException: + raise _vector_store_search_timeout(vector_store_provider_config) from None + except (httpx.TransportError, TypeError, UnicodeEncodeError) as e: + raise self._handle_error(e=e, provider_config=vector_store_provider_config) - all_optional_params: Final[dict[str, object]] = dict(litellm_params) - all_optional_params.update(vector_store_search_optional_params or {}) - - headers, signed_json_body = vector_store_provider_config.sign_request( + search_params: Final = vector_store_provider_config.resolve_search_params( + vector_store_id=vector_store_id, + vector_store_search_optional_params=vector_store_search_optional_params, + api_base=api_base, + litellm_params=MappingProxyType(dict(litellm_params, timeout=timeout)), headers=headers, - optional_params=all_optional_params, - request_data=request_body, - api_base=url, + get=lookup, + ) + search_embedding_executor: Final = EmbedOnceExecutor( + BaseQueryEmbeddingVectorStoreConfig.query_embedding_executor(embedding_executor, None) ) - logging_obj.pre_call( - input="", - api_key="", - additional_args={ - "complete_input_dict": request_body, - "api_base": api_base, - "headers": headers, - }, - ) + def send(params: VectorStoreSearchOptionalRequestParams) -> httpx.Response: + if isinstance(vector_store_provider_config, BaseQueryEmbeddingVectorStoreConfig): + ( + url, + request_body, + ) = vector_store_provider_config.transform_search_vector_store_request( + vector_store_id=vector_store_id, + query=query, + vector_store_search_optional_params=params, + api_base=api_base, + litellm_logging_obj=logging_obj, + litellm_params=MappingProxyType(dict(litellm_params, timeout=timeout)), + extra_body=extra_body, + embedding_executor=search_embedding_executor, + ) + else: + ( + url, + request_body, + ) = vector_store_provider_config.transform_search_vector_store_request( + vector_store_id=vector_store_id, + query=query, + vector_store_search_optional_params=params, + api_base=api_base, + litellm_logging_obj=logging_obj, + litellm_params=dict(litellm_params), + extra_body=extra_body, + ) - request_data: Final = json.dumps(request_body) if signed_json_body is None else signed_json_body + signed_headers, signed_json_body = vector_store_provider_config.sign_request( + headers=headers, + optional_params={**dict(litellm_params), **(params or {})}, + request_data=request_body, + api_base=url, + ) + + logging_obj.pre_call( + input="", + api_key="", + additional_args={ + "complete_input_dict": request_body, + "api_base": api_base, + "headers": signed_headers, + }, + ) + + request_data: Final = json.dumps(request_body) if signed_json_body is None else signed_json_body + + try: + return _send_with_goaway_retry( + partial(sync_httpx_client.post, url=url, headers=signed_headers, data=request_data, timeout=timeout) + ) + except httpx.TimeoutException: + raise _vector_store_search_timeout(vector_store_provider_config) from None + except httpx.HTTPStatusError: + raise + except Exception as e: + raise self._handle_error(e=e, provider_config=vector_store_provider_config) try: - response: Final = _send_with_goaway_retry( - partial(sync_httpx_client.post, url=url, headers=headers, data=request_data, timeout=timeout) + response: Final = send(search_params) + except httpx.HTTPStatusError as rejection: + retry_params: Final = self._vector_store_search_retry_params( + rejection, vector_store_provider_config, search_params, litellm_params + ) + try: + retry_response: Final = send(retry_params) + except httpx.HTTPStatusError as retry_rejection: + raise self._handle_error(e=retry_rejection, provider_config=vector_store_provider_config) from None + return vector_store_provider_config.transform_search_vector_store_response( + response=retry_response, + litellm_logging_obj=logging_obj, ) - except httpx.TimeoutException: - raise vector_store_provider_config.get_error_class( - error_message="Vector store search exceeded the caller timeout.", - status_code=408, - headers=httpx.Headers(), - ) from None - except Exception as e: - raise self._handle_error(e=e, provider_config=vector_store_provider_config) return vector_store_provider_config.transform_search_vector_store_response( response=response, diff --git a/litellm/llms/opensearch/vector_stores/transformation.py b/litellm/llms/opensearch/vector_stores/transformation.py index f7ba28e7f3f..30543a261c6 100644 --- a/litellm/llms/opensearch/vector_stores/transformation.py +++ b/litellm/llms/opensearch/vector_stores/transformation.py @@ -1,13 +1,15 @@ import hashlib import json import re -from collections.abc import Mapping, Sequence -from functools import partial -from typing import TYPE_CHECKING, ClassVar, Final, Literal, NoReturn -from urllib.parse import quote, urlsplit +from collections.abc import Awaitable, Callable, Mapping, Sequence +from dataclasses import dataclass +from functools import partial, reduce +from itertools import accumulate +from typing import TYPE_CHECKING, Annotated, ClassVar, Final, Literal, NoReturn, TypeAlias +from urllib.parse import quote, urlencode, urlsplit import httpx -from pydantic import ConfigDict, Field, TypeAdapter, ValidationError +from pydantic import ConfigDict, Field, StrictBool, StrictFloat, StrictInt, StrictStr, TypeAdapter, ValidationError from litellm._logging import verbose_logger from litellm.exceptions import AuthenticationError, BadRequestError, InternalServerError @@ -18,7 +20,7 @@ from litellm.llms.base_llm.vector_store.transformation import ( BaseQueryEmbeddingVectorStoreConfig, VectorStoreEmbeddingExecutor, ) -from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, SignsRequestsWithAWS, sign_aws_request +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, SignsRequestsWithAWS, run_off_loop_if_aws, sign_aws_request from litellm.llms.bedrock.web_identity_session_policy import ( OPENSEARCH_INDEX_SESSION_POLICY, OPENSEARCH_SESSION_POLICY, @@ -54,6 +56,11 @@ _AWS_HOSTS: Final = ( re.compile(r"\.(?Paoss|aos)\.(?P[a-z0-9-]+)\.on\.aws\Z"), ) _ACCOUNT_ID_LABEL: Final = re.compile(r"\A\d{12}\.") +_UNFILTERABLE_KNN: Final = ("does not support filters", "unknown token [START_OBJECT] after [filter]") +_ANALYZED_STRING_TYPES: Final = frozenset({"text", "match_only_text", "search_as_you_type"}) +_STRING_TYPES: Final = _ANALYZED_STRING_TYPES | {"keyword", "wildcard", "constant_keyword"} +_CONTAINER_TYPES: Final = frozenset({"object", "nested"}) +_FORBIDDEN: Final = 403 _QUERY_VECTOR: Final = TypeAdapter(tuple[float, ...]) _ATTRIBUTE_VALUE_TYPES: Final = (str, int, float, bool) _SEARCH_ONLY_MESSAGE: Final = ( @@ -78,6 +85,11 @@ def _bad_request(message: str) -> BadRequestError: return BadRequestError(message=message, model=None, llm_provider="opensearch") +def _with_signed_headers(headers: Mapping[str, object], signed: Mapping[str, str]) -> Mapping[str, object]: + signed_names: Final = frozenset(name.lower() for name in signed) + return {**{name: value for name, value in headers.items() if name.lower() not in signed_names}, **signed} + + def _unauthenticated(reason: str) -> AuthenticationError: return AuthenticationError( message=f"No usable AWS credentials to sign the OpenSearch request ({reason}). {_CREDENTIALS_HINT}", @@ -109,7 +121,9 @@ class _OpenSearchParams(LiteLLMBaseModel): litellm_embedding_model: str | None = None opensearch_embedding_field: str | None = None + opensearch_hybrid_search: bool | None = None opensearch_index_name: str | None = None + opensearch_search_pipeline: str | None = None opensearch_text_field: str | None = None @property @@ -138,6 +152,258 @@ class _SigningRegion(LiteLLMBaseModel): aws_region_name: str | None = None +_FiniteFloat: TypeAlias = Annotated[StrictFloat, Field(allow_inf_nan=False)] + + +class _Equality(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + type: Literal["eq", "ne"] + key: str + value: StrictStr | StrictInt | _FiniteFloat | StrictBool + + +class _Range(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + type: Literal["gt", "gte", "lt", "lte"] + key: str + value: StrictStr | StrictInt | _FiniteFloat + + +class _Membership(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + type: Literal["in", "nin"] + key: str + value: tuple[StrictStr | StrictInt | _FiniteFloat, ...] + + +class _Compound(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + type: Literal["and", "or"] + filters: "tuple[_Filter, ...]" + + +_Filter: TypeAlias = Annotated[_Equality | _Range | _Membership | _Compound, Field(discriminator="type")] +_Compound.model_rebuild() +_FILTER: Final = TypeAdapter[_Equality | _Range | _Membership | _Compound](_Filter) + + +class _FieldMapping(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + type: str = "object" + ignore_above: int | None = None + normalizer: str | None = None + path: str | None = None + include_in_parent: bool = False + include_in_root: bool = False + properties: "Mapping[str, _FieldMapping]" = Field(default_factory=dict) + fields: "Mapping[str, _FieldMapping]" = Field(default_factory=dict) + + +class _IndexMapping(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + mappings: _FieldMapping + + +_INDEX_MAPPINGS: Final = TypeAdapter(Mapping[str, _IndexMapping]) + + +@dataclass(frozen=True, slots=True) +class _ExactField: + name: str + ignore_above: int | None + + +@dataclass(frozen=True, slots=True) +class _InexactKey: + key: str + reason: str + remedy: str | None = None + + +def _child_field(field: _FieldMapping | None, name: str) -> _FieldMapping | None: + if field is None: + return None + return field.properties.get(name) or field.fields.get(name) + + +def _mapped_field(mapping: _FieldMapping, key: str) -> _FieldMapping | None: + root: Final[_FieldMapping | None] = mapping + return reduce(_child_field, key.split("."), root) + + +def _resolved_field(mapping: _FieldMapping, key: str) -> tuple[str, _FieldMapping] | None: + field: Final = _mapped_field(mapping, key) + if field is None: + return None + if field.type != "alias" or field.path is None: + return key, field + target: Final = _mapped_field(mapping, field.path) + return None if target is None else (field.path, target) + + +def _nested_segments(mapping: _FieldMapping, path: str) -> tuple[tuple[str, _FieldMapping], ...]: + parts: Final = path.split(".") + prefixes: Final = (".".join(parts[:end]) for end in range(1, len(parts) + 1)) + segments: Final = ((prefix, _mapped_field(mapping, prefix)) for prefix in prefixes) + return tuple((prefix, field) for prefix, field in segments if field is not None and field.type == "nested") + + +def _copied_to_root(parent_copied_to_root: bool, nested: _FieldMapping) -> bool: + return nested.include_in_root or (nested.include_in_parent and parent_copied_to_root) + + +def _nested_path(mapping: _FieldMapping, path: str) -> str | None: + segments: Final = _nested_segments(mapping, path) + copied: Final = tuple(accumulate((field for _, field in segments), _copied_to_root, initial=True))[1:] + if not segments or copied[-1]: + return None + return next(prefix for (prefix, _), reaches in zip(segments, copied, strict=True) if not reaches) + + +def _exact_field( + key: str, field: _FieldMapping, comparison: _Equality | _Range | _Membership +) -> _ExactField | _InexactKey: + if isinstance(comparison, _Range) and not isinstance(comparison.value, str) and field.type in _STRING_TYPES: + return _InexactKey( + key=key, + reason=f"the index maps {key} as {field.type}, which compares numbers as text", + remedy=f"Map {key} as a numeric field", + ) + if field.type not in _ANALYZED_STRING_TYPES and field.normalizer is None: + return _ExactField(name=key, ignore_above=field.ignore_above) + keywords: Final = sorted( + (name != "keyword", name, subfield.ignore_above) + for name, subfield in field.fields.items() + if subfield.type == "keyword" and subfield.normalizer is None + ) + if keywords: + _, name, ignore_above = keywords[0] + return _ExactField(name=f"{key}.{name}", ignore_above=ignore_above) + if field.normalizer is not None or any(subfield.type == "keyword" for subfield in field.fields.values()): + return _InexactKey(key=key, reason=f"the index normalizes every keyword mapping of {key}") + return _InexactKey(key=key, reason=f"the index maps {key} as {field.type} with no keyword subfield") + + +def _string_values(comparison: _Equality | _Range | _Membership) -> tuple[str, ...]: + values: Final = comparison.value if isinstance(comparison, _Membership) else (comparison.value,) + return tuple(value for value in values if isinstance(value, str)) + + +def _queried_field( + key: str, mapping: _FieldMapping, comparison: _Equality | _Range | _Membership +) -> _ExactField | _InexactKey | None: + resolved: Final = _resolved_field(mapping, key) + if resolved is None: + return None + path, field = resolved + if field.type in _CONTAINER_TYPES: + return _InexactKey( + key=key, + reason=f"the index maps {path} as {field.type}, which holds fields rather than values", + remedy=f"Filter on a field inside {path}", + ) + nested: Final = _nested_path(mapping, path) + if nested is not None: + return _InexactKey( + key=key, + reason=f"the index maps {nested} as nested, which needs a nested query", + remedy=f"Map {nested} as object", + ) + exact: Final = _exact_field(path, field, comparison) + if isinstance(exact, _ExactField) and exact.name == path: + return _ExactField(name=key, ignore_above=exact.ignore_above) + return exact + + +def _exact_comparison( + comparison: _Equality | _Range | _Membership, mappings: Mapping[str, _IndexMapping] +) -> _Equality | _Range | _Membership | _InexactKey: + key: Final = comparison.key + queried: Final = tuple(_queried_field(key, index.mappings, comparison) for index in mappings.values()) + exact: Final = tuple(field for field in queried if field is not None) + differently: Final = _InexactKey( + key=key, + reason=f"the indexes behind this store map {key} differently", + remedy=f"Map {key} to the same keyword field in every index behind the store", + ) + if 0 < len(exact) < len(queried): + return differently + inexact: Final = next((field for field in exact if isinstance(field, _InexactKey)), None) + if inexact is not None: + return inexact + fields: Final = tuple(field for field in exact if isinstance(field, _ExactField)) + names: Final = frozenset(field.name for field in fields) + name: Final = next(iter(names), key) + if len(names) > 1: + return differently + ignore_above: Final = min((field.ignore_above for field in fields if field.ignore_above), default=None) + if ignore_above is not None and any(len(value) > ignore_above for value in _string_values(comparison)): + return _InexactKey(key=key, reason=f"the index ignores {key} values longer than {ignore_above} characters") + return comparison.model_copy(update={"key": name}) + + +def _read_index_mappings(mapping: httpx.Response) -> Mapping[str, _IndexMapping] | BaseLLMException: + if mapping.status_code == _FORBIDDEN: + return BaseLLMException( + status_code=mapping.status_code, + message=( + "OpenSearch refused the index mapping read that filters need. That read needs aoss:DescribeIndex on " + "Serverless or es:ESHttpGet on a domain. Set drop_params to drop the filters instead." + ), + headers=mapping.headers, + ) + if not mapping.is_success: + return BaseLLMException( + status_code=mapping.status_code, + message=f"OpenSearch answered {mapping.status_code} to the index mapping read that filters need.", + headers=mapping.headers, + ) + try: + return _INDEX_MAPPINGS.validate_json(mapping.content) + except ValidationError: + return BaseLLMException( + status_code=502, + message="OpenSearch answered the index mapping read without index mappings.", + headers=mapping.headers, + ) + + +def _exact_filter(search_filter: _Filter, mappings: Mapping[str, _IndexMapping]) -> _Filter | _InexactKey: + match search_filter: + case _Compound(filters=filters): + children: Final = tuple(_exact_filter(child, mappings) for child in filters) + inexact: Final = next((child for child in children if isinstance(child, _InexactKey)), None) + return inexact or search_filter.model_copy(update={"filters": children}) + case _: + return _exact_comparison(search_filter, mappings) + + +def _opensearch_filter(search_filter: _Filter) -> Mapping[str, object]: + match search_filter: + case _Compound(type="and", filters=filters): + return {"bool": {"filter": tuple(_opensearch_filter(child) for child in filters)}} + case _Compound(filters=filters): + return { + "bool": {"should": tuple(_opensearch_filter(child) for child in filters), "minimum_should_match": 1} + } + case _Equality(type="eq", key=key, value=value): + return {"term": {key: value}} + case _Equality(key=key, value=value): + return {"bool": {"must_not": ({"term": {key: value}},)}} + case _Membership(type="in", key=key, value=values): + return {"terms": {key: values}} + case _Membership(key=key, value=values): + return {"bool": {"must_not": ({"terms": {key: values}},)}} + case _Range(type=operator, key=key, value=value): + return {"range": {key: {operator: value}}} + + class _SearchOptions(LiteLLMBaseModel): model_config = ConfigDict(frozen=True, extra="ignore") @@ -149,12 +415,25 @@ class _SearchOptions(LiteLLMBaseModel): @property def unsupported(self) -> tuple[str, ...]: requested: Final = ( - ("filters", self.filters), ("ranking_options", self.ranking_options), ("rewrite_query", self.rewrite_query), ) return tuple(name for name, value in requested if value) + def search_filter(self, drop_unsupported: bool) -> _Filter | None: + if self.filters is None or self.filters == {}: + return None + try: + return _FILTER.validate_python(self.filters) + except ValidationError: + if drop_unsupported: + return None + raise _bad_request( + "filters must be an OpenAI vector store filter, either a comparison with a key, a value, and type eq, " + 'ne, gt, gte, lt, lte, in, or nin, or a compound with a filters list and type "and" or "or". Fix it, ' + "or set drop_params to drop it." + ) from None + @property def k(self) -> int: if self.max_num_results is None: @@ -169,6 +448,48 @@ class _SearchOptions(LiteLLMBaseModel): _KNOWN_PARAMS: Final = frozenset(name for name in _OpenSearchParams.model_fields if name.startswith(_PARAM_PREFIX)) +@dataclass(frozen=True, slots=True) +class _Search: + params: _OpenSearchParams + k: int + search_filter: _Filter | None + + @property + def hybrid(self) -> bool: + return bool(self.params.opensearch_hybrid_search and self.params.opensearch_search_pipeline) + + @property + def opensearch_filter(self) -> Mapping[str, object] | None: + return None if self.search_filter is None else _opensearch_filter(self.search_filter) + + @property + def path(self) -> str: + pipeline: Final = self.params.opensearch_search_pipeline + return f"/_search?{urlencode({'search_pipeline': pipeline}, quote_via=quote)}" if pipeline else "/_search" + + def query(self, query_text: str, vector: tuple[float, ...]) -> Mapping[str, object]: + opensearch_filter: Final = self.opensearch_filter + knn: Final[Mapping[str, object]] = {"vector": vector, "k": self.k} + knn_query: Final[Mapping[str, object]] = { + "knn": { + self.params.embedding_field: knn if opensearch_filter is None else {**knn, "filter": opensearch_filter} + } + } + if not self.hybrid: + return knn_query + keyword_query: Final[Mapping[str, object]] = {"match": {self.params.text_field: {"query": query_text}}} + return { + "hybrid": { + "queries": ( + keyword_query + if opensearch_filter is None + else {"bool": {"must": (keyword_query,), "filter": (opensearch_filter,)}}, + knn_query, + ) + } + } + + class _LoggedSearch(LiteLLMBaseModel): model_config = ConfigDict(frozen=True, extra="ignore") @@ -309,7 +630,7 @@ class OpenSearchVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, SignsRequ litellm_params: Mapping[str, object], optional_params: VectorStoreSearchOptionalRequestParams, extra_body: Mapping[str, object] | None, - ) -> tuple[_OpenSearchParams, int]: + ) -> _Search: unknown: Final = sorted( key for key in litellm_params if key.startswith(_PARAM_PREFIX) and key not in _KNOWN_PARAMS ) @@ -324,34 +645,42 @@ 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 and not drops_unsupported_params(litellm_params.get("drop_params")): + drop_unsupported: Final = drops_unsupported_params(litellm_params.get("drop_params")) + if options.unsupported and not drop_unsupported: raise _bad_request( f"The OpenSearch vector store does not support {', '.join(options.unsupported)} yet. " "Remove it from the request, or set drop_params to drop it." ) try: - return _OpenSearchParams.model_validate(litellm_params), options.k + params: Final = _OpenSearchParams.model_validate(litellm_params) except ValidationError: raise _bad_request( - "opensearch_embedding_field, opensearch_index_name, opensearch_text_field, and litellm_embedding_model " - "must be strings." + "opensearch_embedding_field, opensearch_index_name, opensearch_search_pipeline, opensearch_text_field, " + "and litellm_embedding_model must be strings, and opensearch_hybrid_search must be a boolean." ) from None + if params.opensearch_hybrid_search and not params.opensearch_search_pipeline and not drop_unsupported: + raise _bad_request( + "opensearch_hybrid_search needs opensearch_search_pipeline, the name of a search pipeline with a " + "normalization-processor, since OpenSearch returns unnormalized hybrid scores without one. Set it, or " + "set drop_params to run a k-NN search instead." + ) + return _Search(params=params, k=options.k, search_filter=options.search_filter(drop_unsupported)) @staticmethod def _search_request( vector_store_id: str, + query_text: str, embedding: EmbeddingResponse, - k: int, - params: _OpenSearchParams, + search: _Search, api_base: str, ) -> 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(params.index_name(vector_store_id), safe='')}/_search", + f"{api_base}/{quote(search.params.index_name(vector_store_id), safe='')}{search.path}", { - "size": k, - "query": {"knn": {params.embedding_field: {"vector": vector, "k": k}}}, - "_source": {"excludes": [params.embedding_field]}, + "size": search.k, + "query": search.query(query_text, vector), + "_source": {"excludes": [search.params.embedding_field]}, }, ) @@ -366,11 +695,12 @@ class OpenSearchVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, SignsRequ extra_body: Mapping[str, object] | None = None, embedding_executor: VectorStoreEmbeddingExecutor | None = None, ) -> tuple[str, dict[str, object]]: # mutable-ok: the provider contract returns a writable JSON request body - params, k = self._params(litellm_params, vector_store_search_optional_params, extra_body) + search: Final = self._params(litellm_params, vector_store_search_optional_params, extra_body) + query_text: Final = self.query_text(query) embedding: Final = self.query_embedding_executor(embedding_executor, None).embed( - params.require_embedding_model(), self.query_text(query), self.query_embedding_configuration(litellm_params) + search.params.require_embedding_model(), query_text, self.query_embedding_configuration(litellm_params) ) - return self._search_request(vector_store_id, embedding, k, params, api_base) + return self._search_request(vector_store_id, query_text, embedding, search, api_base) async def atransform_search_vector_store_request( self, @@ -383,28 +713,121 @@ class OpenSearchVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, SignsRequ extra_body: Mapping[str, object] | None = None, embedding_executor: VectorStoreEmbeddingExecutor | None = None, ) -> tuple[str, dict[str, object]]: # mutable-ok: the provider contract returns a writable JSON request body - params, k = self._params(litellm_params, vector_store_search_optional_params, extra_body) + search: Final = self._params(litellm_params, vector_store_search_optional_params, extra_body) + query_text: Final = self.query_text(query) embedding: Final = await self.query_embedding_executor(embedding_executor, None).aembed( - params.require_embedding_model(), self.query_text(query), self.query_embedding_configuration(litellm_params) + search.params.require_embedding_model(), query_text, self.query_embedding_configuration(litellm_params) + ) + return self._search_request(vector_store_id, query_text, embedding, search, api_base) + + def _filter_mapping_url( + self, + vector_store_id: str, + vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + api_base: str, + litellm_params: Mapping[str, object], + ) -> tuple[_Filter, str] | None: + search: Final = self._params(litellm_params, vector_store_search_optional_params, None) + if search.search_filter is None: + return None + return search.search_filter, f"{api_base}/{quote(search.params.index_name(vector_store_id), safe='')}/_mapping" + + def _signed_mapping_read_headers( + self, url: str, headers: Mapping[str, object], litellm_params: Mapping[str, object] + ) -> Mapping[str, object]: + signed: Final = self.sign_http_request("GET", url, b"", "application/json", litellm_params) + return _with_signed_headers(headers, dict(signed.headers.items())) + + def _exact_search_params( + self, + vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + search_filter: _Filter, + mapping: httpx.Response, + litellm_params: Mapping[str, object], + ) -> VectorStoreSearchOptionalRequestParams: + mappings: Final = _read_index_mappings(mapping) + exact: Final = mappings if isinstance(mappings, BaseLLMException) else _exact_filter(search_filter, mappings) + if not isinstance(exact, BaseLLMException | _InexactKey): + return {**vector_store_search_optional_params, "filters": exact.model_dump()} + if isinstance(exact, BaseLLMException) and exact.status_code != _FORBIDDEN: + raise exact + if drops_unsupported_params(litellm_params.get("drop_params")): + return {**vector_store_search_optional_params, "filters": None} + if isinstance(exact, BaseLLMException): + raise exact + remedy: Final = ( + exact.remedy or f"Map {exact.key} as keyword without a normalizer or give it such a keyword subfield" + ) + raise _bad_request( + f"filters cannot match {exact.key} exactly because {exact.reason}. {remedy}, or set drop_params to drop " + "the filters." + ) + + def resolve_search_params( + self, + vector_store_id: str, + vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + api_base: str, + litellm_params: Mapping[str, object], + headers: Mapping[str, object], + get: Callable[[str, Mapping[str, object]], httpx.Response], + ) -> VectorStoreSearchOptionalRequestParams: + lookup: Final = self._filter_mapping_url( + vector_store_id, vector_store_search_optional_params, api_base, litellm_params + ) + if lookup is None: + return vector_store_search_optional_params + search_filter, url = lookup + mapping: Final = get(url, self._signed_mapping_read_headers(url, headers, litellm_params)) + return self._exact_search_params(vector_store_search_optional_params, search_filter, mapping, litellm_params) + + async def aresolve_search_params( + self, + vector_store_id: str, + vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + api_base: str, + litellm_params: Mapping[str, object], + headers: Mapping[str, object], + get: Callable[[str, Mapping[str, object]], Awaitable[httpx.Response]], + ) -> VectorStoreSearchOptionalRequestParams: + lookup: Final = self._filter_mapping_url( + vector_store_id, vector_store_search_optional_params, api_base, litellm_params + ) + if lookup is None: + return vector_store_search_optional_params + search_filter, url = lookup + signed: Final = await run_off_loop_if_aws(self, self._signed_mapping_read_headers, url, headers, litellm_params) + mapping: Final = await get(url, signed) + return self._exact_search_params(vector_store_search_optional_params, search_filter, mapping, litellm_params) + + def recover_rejected_search( + self, + rejection: httpx.Response, + vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + litellm_params: Mapping[str, object], + ) -> VectorStoreSearchOptionalRequestParams | BaseLLMException | None: + if rejection.status_code != 400 or not any(marker in rejection.text for marker in _UNFILTERABLE_KNN): + return None + if drops_unsupported_params(litellm_params.get("drop_params")): + return {**vector_store_search_optional_params, "filters": None} + return BaseLLMException( + status_code=400, + message="This OpenSearch index cannot apply filters inside a k-NN query, which needs the Lucene engine " + "on OpenSearch 2.4+ or Faiss on 2.9+. Remove the filters, or set drop_params to drop them.", + headers=httpx.Headers(), ) - return self._search_request(vector_store_id, embedding, k, params, api_base) def sign_request( self, - headers: Mapping[str, str], + headers: Mapping[str, object], optional_params: Mapping[str, object], request_data: Mapping[str, object], api_base: str, api_key: str | None = None, - ) -> tuple[dict[str, str], bytes]: # mutable-ok: the shared HTTP handler requires writable headers + ) -> tuple[dict[str, object], bytes]: # mutable-ok: the shared HTTP handler requires writable headers body: Final = json.dumps(request_data).encode() signed: Final = self.sign_http_request("POST", api_base, body, "application/json", optional_params) - signed_headers: Final = dict(signed.headers.items()) - signed_names: Final = frozenset(name.lower() for name in signed_headers) - return { - **{name: value for name, value in headers.items() if name.lower() not in signed_names}, - **signed_headers, - }, body + return dict(_with_signed_headers(headers, dict(signed.headers.items()))), body def sign_http_request( self, method: str, url: str, body: bytes, content_type: str, litellm_params: Mapping[str, object] diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index 38041ed1862..c96bc14712f 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -67,7 +67,7 @@ from litellm.proxy.rag_endpoints.upload_security import ( from litellm.proxy.vector_store_endpoints.endpoints import ( build_request_data_from_managed_vector_store, reject_caller_embedding_selection_params, - reject_caller_opensearch_index_name, + reject_caller_opensearch_store_params, ) from litellm.proxy.vector_store_endpoints.utils import ( assert_user_can_access_vector_store_id, @@ -1153,7 +1153,7 @@ async def rag_query( **retrieval_config, **store_data, } - reject_caller_opensearch_index_name( + reject_caller_opensearch_store_params( caller_body=retrieval_config, request_fields=merged_retrieval_config, source="retrieval_config", diff --git a/litellm/proxy/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py index 0065f8e4002..a8836b85a91 100644 --- a/litellm/proxy/vector_store_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_endpoints/endpoints.py @@ -81,20 +81,24 @@ def reject_caller_embedding_selection_params(payload: Mapping[str, object], sour ) -def reject_caller_opensearch_index_name( +_REGISTRATION_ONLY_OPENSEARCH_PARAMS: Final = frozenset( + ("opensearch_index_name", "opensearch_search_pipeline", "opensearch_hybrid_search") +) + + +def reject_caller_opensearch_store_params( caller_body: Mapping[str, object], request_fields: Mapping[str, object], source: str ) -> None: provider: Final = request_fields.get("custom_llm_provider") - if ( - isinstance(provider, str) - and provider.split("/", 1)[0] == LlmProviders.OPENSEARCH.value - and "opensearch_index_name" in caller_body - ): + if not isinstance(provider, str) or provider.split("/", 1)[0] != LlmProviders.OPENSEARCH.value: + return + blocked: Final = next((key for key in caller_body if key in _REGISTRATION_ONLY_OPENSEARCH_PARAMS), None) + if blocked is not None: raise HTTPException( status_code=400, detail={ - "error": f"'opensearch_index_name' cannot be set in {source}. " - "The index comes from the vector store's server-side registration." + "error": f"'{blocked}' cannot be set in {source}. " + "It comes from the vector store's server-side registration." }, ) @@ -196,7 +200,7 @@ async def vector_store_search( caller_body.get("extra_body"), ) ) - reject_caller_opensearch_index_name( + reject_caller_opensearch_store_params( caller_body=caller_body, request_fields=JSON_OBJECT_ADAPTER.validate_python(data), source="the search request body", diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreForm.integration.test.tsx b/litellm/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreForm.integration.test.tsx index f571ec6b9a0..00e5b337daf 100644 --- a/litellm/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreForm.integration.test.tsx +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreForm.integration.test.tsx @@ -155,6 +155,8 @@ describe("VectorStoreForm submit payload", () => { opensearch_embedding_field: "embedding", opensearch_text_field: "text", opensearch_index_name: undefined, + opensearch_search_pipeline: undefined, + opensearch_hybrid_search: undefined, }; await vi.waitFor(() => expect(mockCreate).toHaveBeenCalledTimes(1)); @@ -183,6 +185,60 @@ describe("VectorStoreForm submit payload", () => { expect(createdPayload().litellm_params).toMatchObject({ opensearch_index_name: "docs" }); }); + it("registers an opensearch store for hybrid search with its search pipeline", async () => { + const user = setupUser(); + renderForm(); + + await chooseProvider(user, "Amazon OpenSearch"); + fireEvent.change(screen.getByPlaceholderText("my-knn-index (OpenSearch index name)"), { + target: { value: "cav-lucene" }, + }); + fireEvent.change(screen.getByPlaceholderText("https://abc123.us-east-1.aoss.amazonaws.com"), { + target: { value: "http://127.0.0.1:27069" }, + }); + await chooseFromSelect(user, 1, "text-embedding-3-small"); + fireEvent.change(screen.getByPlaceholderText("my-hybrid-search-pipeline"), { + target: { value: "litellm-qa-norm" }, + }); + await user.click(screen.getByRole("switch", { name: "Hybrid Search" })); + await submit(user); + + const expectedLitellmParams = { + api_base: "http://127.0.0.1:27069", + aws_region_name: undefined, + litellm_embedding_model: "text-embedding-3-small", + opensearch_embedding_field: "embedding", + opensearch_text_field: "text", + opensearch_index_name: undefined, + opensearch_search_pipeline: "litellm-qa-norm", + opensearch_hybrid_search: true, + }; + + await vi.waitFor(() => expect(mockCreate).toHaveBeenCalledTimes(1)); + expect(createdPayload().litellm_params).toStrictEqual(expectedLitellmParams); + }); + + it("blocks a hybrid opensearch submit that has no search pipeline", async () => { + const user = setupUser(); + renderForm(); + + await chooseProvider(user, "Amazon OpenSearch"); + fireEvent.change(screen.getByPlaceholderText("my-knn-index (OpenSearch index name)"), { + target: { value: "cav-lucene" }, + }); + fireEvent.change(screen.getByPlaceholderText("https://abc123.us-east-1.aoss.amazonaws.com"), { + target: { value: "http://127.0.0.1:27069" }, + }); + await chooseFromSelect(user, 1, "text-embedding-3-small"); + await user.click(screen.getByRole("switch", { name: "Hybrid Search" })); + await submit(user); + + expect( + await screen.findByText("Hybrid search needs a search pipeline with a normalization-processor"), + ).toBeInTheDocument(); + expect(mockCreate).not.toHaveBeenCalled(); + }); + it("blocks an opensearch submit that has no endpoint or embedding model", async () => { const user = setupUser(); renderForm(); diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreForm.tsx b/litellm/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreForm.tsx index 48a9f0ba37e..1b40949ca8d 100644 --- a/litellm/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreForm.tsx +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreForm.tsx @@ -33,6 +33,7 @@ import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/u import { Input } from "@/components/ui/input"; import { InputGroup, InputGroupAddon, InputGroupButton, InputGroupInput } from "@/components/ui/input-group"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import { Switch } from "@/components/ui/switch"; import { Textarea } from "@/components/ui/textarea"; import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; import { useZodForm } from "@/lib/forms/useZodForm"; @@ -65,7 +66,7 @@ interface VectorStoreFormProps { credentials: CredentialItem[]; } -const PROVIDER_FIELD_NAMES = [ +const TEXT_PROVIDER_FIELD_NAMES = [ "api_base", "api_key", "vertex_project", @@ -90,12 +91,25 @@ const PROVIDER_FIELD_NAMES = [ "opensearch_embedding_field", "opensearch_text_field", "opensearch_index_name", + "opensearch_search_pipeline", ] as const; -type ProviderFieldName = (typeof PROVIDER_FIELD_NAMES)[number]; +const BOOLEAN_PROVIDER_FIELD_NAMES = ["opensearch_hybrid_search"] as const; + +type TextProviderFieldName = (typeof TEXT_PROVIDER_FIELD_NAMES)[number]; + +type BooleanProviderFieldName = (typeof BOOLEAN_PROVIDER_FIELD_NAMES)[number]; + +type ProviderFieldName = TextProviderFieldName | BooleanProviderFieldName; + +const isTextProviderFieldName = (name: string): name is TextProviderFieldName => + (TEXT_PROVIDER_FIELD_NAMES as readonly string[]).includes(name); + +const isBooleanProviderFieldName = (name: string): name is BooleanProviderFieldName => + (BOOLEAN_PROVIDER_FIELD_NAMES as readonly string[]).includes(name); const isProviderFieldName = (name: string): name is ProviderFieldName => - (PROVIDER_FIELD_NAMES as readonly string[]).includes(name); + isTextProviderFieldName(name) || isBooleanProviderFieldName(name); const optionalText = z.string().optional(); @@ -129,6 +143,8 @@ const vectorStoreShape = { opensearch_embedding_field: optionalText, opensearch_text_field: optionalText, opensearch_index_name: optionalText, + opensearch_search_pipeline: optionalText, + opensearch_hybrid_search: z.boolean().optional(), }; const vectorStoreSchema = z.object(vectorStoreShape).superRefine((values, ctx) => { @@ -144,6 +160,17 @@ const vectorStoreSchema = z.object(vectorStoreShape).superRefine((values, ctx) = : i18n.t("vectorStores.form.errors.inputField", { field: i18n.t(field.label).toLowerCase() }), }), ); + if ( + values.custom_llm_provider === "opensearch" && + values.opensearch_hybrid_search === true && + !values.opensearch_search_pipeline + ) { + ctx.addIssue({ + code: "custom", + path: ["opensearch_search_pipeline"], + message: i18n.t("vectorStores.form.errors.hybridSearchNeedsPipeline"), + }); + } }); type VectorStoreFormValues = z.output; @@ -540,10 +567,12 @@ const VectorStoreForm: React.FC = ({ ); }; -type SupportedProviderField = VectorStoreFieldConfig & { name: ProviderFieldName }; +type SupportedProviderField = + | (VectorStoreFieldConfig & { type: "boolean"; name: BooleanProviderFieldName }) + | (VectorStoreFieldConfig & { type?: "text" | "password" | "select"; name: TextProviderFieldName }); const isSupportedProviderField = (field: VectorStoreFieldConfig): field is SupportedProviderField => - isProviderFieldName(field.name); + field.type === "boolean" ? isBooleanProviderFieldName(field.name) : isTextProviderFieldName(field.name); interface ProviderFieldProps { field: SupportedProviderField; @@ -556,6 +585,14 @@ const ProviderField: React.FC = ({ field, control, modelInfo const label = labelWithHint(t(field.label), t(field.tooltip)); const placeholder = field.placeholderKey ? t(field.placeholderKey) : field.placeholder; + if (field.type === "boolean") { + return ( + + {({ id, value, onChange }) => } + + ); + } + if (field.type === "select") { const selectOptions = field.options ?? diff --git a/litellm/ui/litellm-dashboard/src/components/vector_store_providers.test.tsx b/litellm/ui/litellm-dashboard/src/components/vector_store_providers.test.tsx index b4d890475c7..ce9a24327ae 100644 --- a/litellm/ui/litellm-dashboard/src/components/vector_store_providers.test.tsx +++ b/litellm/ui/litellm-dashboard/src/components/vector_store_providers.test.tsx @@ -68,7 +68,7 @@ describe("getVectorStoreProviderLogoAndName", () => { expect(byName("mongodb_num_candidates")).toMatchObject({ required: false }); }); - it("registers opensearch with its endpoint, optional region, proxy embedding model, and default field names", () => { + it("registers opensearch with its endpoint, optional region, proxy embedding model, default field names, and hybrid search", () => { const fields = getProviderSpecificFields("opensearch"); const byName = (name: string) => fields.find((field) => field.name === name); @@ -83,6 +83,8 @@ describe("getVectorStoreProviderLogoAndName", () => { "opensearch_embedding_field", "opensearch_text_field", "opensearch_index_name", + "opensearch_search_pipeline", + "opensearch_hybrid_search", ]); expect(byName("api_base")).toMatchObject({ required: true }); expect(byName("aws_region_name")).toMatchObject({ required: false }); @@ -92,6 +94,8 @@ describe("getVectorStoreProviderLogoAndName", () => { 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"); + expect(byName("opensearch_search_pipeline")).toMatchObject({ type: "text", required: false }); + expect(byName("opensearch_hybrid_search")).toMatchObject({ type: "boolean", required: false }); }); it("registers valkey in the provider, logo, and field maps", () => { diff --git a/litellm/ui/litellm-dashboard/src/components/vector_store_providers.tsx b/litellm/ui/litellm-dashboard/src/components/vector_store_providers.tsx index b43357e076b..8311e32998c 100644 --- a/litellm/ui/litellm-dashboard/src/components/vector_store_providers.tsx +++ b/litellm/ui/litellm-dashboard/src/components/vector_store_providers.tsx @@ -57,7 +57,7 @@ export interface VectorStoreFieldConfig { placeholder?: string; placeholderKey?: ParseKeys; required: boolean; - type?: "text" | "password" | "select"; + type?: "text" | "password" | "select" | "boolean"; options?: { value: string; label: string }[]; initialValue?: string; } @@ -360,6 +360,21 @@ export const vectorStoreProviderFields: Record required: false, type: "text", }, + { + name: "opensearch_search_pipeline", + label: "vectorStores.providerFields.labels.searchPipelineOptional", + tooltip: "vectorStores.providerFields.tooltips.opensearch.searchPipeline", + placeholder: "my-hybrid-search-pipeline", + required: false, + type: "text", + }, + { + name: "opensearch_hybrid_search", + label: "vectorStores.providerFields.labels.hybridSearch", + tooltip: "vectorStores.providerFields.tooltips.opensearch.hybridSearch", + required: false, + type: "boolean", + }, ], s3_vectors: [ { diff --git a/litellm/ui/litellm-dashboard/src/i18n/locales/en.json b/litellm/ui/litellm-dashboard/src/i18n/locales/en.json index 28cb5d97c1c..f4e30c873ee 100644 --- a/litellm/ui/litellm-dashboard/src/i18n/locales/en.json +++ b/litellm/ui/litellm-dashboard/src/i18n/locales/en.json @@ -976,7 +976,8 @@ "providerRequired": "Please select a provider", "vectorStoreIdRequired": "Please input the vector store ID from your api provider", "selectField": "Please select the {{field}}", - "inputField": "Please input the {{field}}" + "inputField": "Please input the {{field}}", + "hybridSearchNeedsPipeline": "Hybrid search needs a search pipeline with a normalization-processor" }, "toasts": { "invalidMetadata": "Invalid JSON in metadata field", @@ -1062,6 +1063,8 @@ "opensearchEndpoint": "OpenSearch Endpoint", "awsRegionOptional": "AWS Region (optional)", "indexNameOptional": "Index Name (optional)", + "searchPipelineOptional": "Search Pipeline (optional)", + "hybridSearch": "Hybrid Search", "vectorBucketName": "Vector Bucket Name", "indexName": "Index Name", "awsRegion": "AWS Region" @@ -1114,7 +1117,9 @@ "embeddingModel": "The embedding model on this proxy that created the vectors already stored in your index. LiteLLM embeds every search query with it, so it must be the same model. Add it under Models first if it is not listed", "embeddingField": "The top-level knn_vector field in your index mapping that holds each document's embedding (default: embedding)", "textField": "The top-level field in each document that holds its readable text. LiteLLM returns this text in search results (default: text)", - "indexName": "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" + "indexName": "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", + "searchPipeline": "The name of a search pipeline in your OpenSearch cluster. LiteLLM runs every search through it. Hybrid search needs one with a normalization-processor, since OpenSearch returns unnormalized hybrid scores without one", + "hybridSearch": "Turn on to combine a BM25 keyword match on the Text Field with the k-NN vector search. Needs a Search Pipeline with a normalization-processor" }, "s3Vectors": { "bucketName": "S3 bucket name for vector storage (will be auto-created if it doesn't exist)", diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index 9ca064e308a..ab39c2ad774 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -64,6 +64,8 @@ IGNORE_FUNCTIONS = [ "_resolve", # OCI: $ref resolver bounded by `resolving_stack` cycle guard. "resolve_oci_schema_anyof", # OCI: bounded by JSON-schema tree depth (no cycles possible in well-formed input). "sanitize_oci_schema", # OCI: bounded by JSON-schema tree depth. + "_opensearch_filter", # OpenSearch: runs only on a filter pydantic validated, whose recursion guard stops at depth 255. + "_exact_filter", # OpenSearch: runs only on a filter pydantic validated, whose recursion guard stops at depth 255. "_freeze_for_dedupe", # OTEL: max depth set (default 16, _FREEZE_MAX_DEPTH); fails closed by returning repr(value) at the cap. "apply_json_merge_patch", # max depth set (_MAX_MERGE_DEPTH=64); fails closed by raising ValueError at the cap. "_inlined", # max depth set (_MAX_INLINE_DEPTH=32); passes the schema node through untouched at the cap. Walks a Pydantic JSON schema once at import time. diff --git a/tests/unit/llms/bedrock/test_web_identity_session_policy.py b/tests/unit/llms/bedrock/test_web_identity_session_policy.py index a181087cba9..9d304b45e09 100644 --- a/tests/unit/llms/bedrock/test_web_identity_session_policy.py +++ b/tests/unit/llms/bedrock/test_web_identity_session_policy.py @@ -840,8 +840,8 @@ class TestEachCallerGetsItsOwnCeiling: _OPENSEARCH_ENDPOINT_ACTIONS: Final = MappingProxyType( { - "https://search-docs-abc123.us-east-1.es.amazonaws.com": "es:ESHttpPost", - "https://abc123.us-east-1.aoss.amazonaws.com": "aoss:APIAccessAll", + "https://search-docs-abc123.us-east-1.es.amazonaws.com": frozenset({"es:ESHttpGet", "es:ESHttpPost"}), + "https://abc123.us-east-1.aoss.amazonaws.com": frozenset({"aoss:APIAccessAll"}), } ) @@ -868,13 +868,15 @@ def _opensearch_policy(api_base: str) -> WebIdentitySessionPolicy: class TestOpenSearchSearchGetsItsOwnCeiling: - @pytest.mark.parametrize(("api_base", "action"), sorted(_OPENSEARCH_ENDPOINT_ACTIONS.items())) - def test_search_action_is_granted_by_the_opensearch_ceiling(self, api_base: str, action: str): - assert action in _granted_actions(_opensearch_policy(api_base)) + @pytest.mark.parametrize(("api_base", "actions"), sorted(_OPENSEARCH_ENDPOINT_ACTIONS.items())) + def test_search_and_mapping_read_actions_are_granted_by_the_opensearch_ceiling( + self, api_base: str, actions: frozenset[str] + ): + assert actions <= _granted_actions(_opensearch_policy(api_base)) @pytest.mark.parametrize("api_base", sorted(_OPENSEARCH_ENDPOINT_ACTIONS)) def test_opensearch_ceiling_grants_only_opensearch_services(self, api_base: str): - allowed: Final = frozenset(_service(action) for action in _OPENSEARCH_ENDPOINT_ACTIONS.values()) + allowed: Final = frozenset(map(_service, frozenset[str]().union(*_OPENSEARCH_ENDPOINT_ACTIONS.values()))) extra: Final = sorted(a for a in _granted_actions(_opensearch_policy(api_base)) if _service(a) not in allowed) assert not extra, f"the OpenSearch session policy also grants {extra}" diff --git a/tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py b/tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py index e20dfaffcc6..2309149b00c 100644 --- a/tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py +++ b/tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py @@ -1,14 +1,18 @@ import json +from collections.abc import Mapping from typing import Final from unittest.mock import MagicMock import httpx +import openai import pytest import respx import litellm from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.vector_stores.transformation import BedrockVectorStoreConfig +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.types.vector_stores import VectorStoreSearchResponse def test_transform_search_request(): @@ -379,6 +383,55 @@ async def test_bedrock_search_with_credentials_managed_registry( assert search_response["data"][0]["content"][0]["text"] == "Mock search result" +async def _search_knowledge_base( + asynchronous: bool, respond: MagicMock, extra_headers: Mapping[str, object] +) -> VectorStoreSearchResponse: + request_kwargs: Final = { + "vector_store_id": "T37J8R4WTM", + "query": "what happens after we add a model", + "custom_llm_provider": "bedrock", + "aws_access_key_id": "test_access_key", + "aws_secret_access_key": "test_secret_key", + "aws_region_name": "us-east-1", + "extra_headers": extra_headers, + } + if asynchronous: + return await litellm.vector_stores.asearch( + **request_kwargs, client=AsyncHTTPHandler(transport=httpx.MockTransport(respond)) + ) + with httpx.Client(transport=httpx.MockTransport(respond)) as transport: + return litellm.vector_stores.search(**request_kwargs, client=HTTPHandler(client=transport)) + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.asyncio +async def test_a_null_extra_header_still_searches_the_knowledge_base(asynchronous: bool) -> None: + respond: Final = MagicMock(return_value=httpx.Response(200, json=BEDROCK_RETRIEVE_RESPONSE)) + + response: Final = await _search_knowledge_base( + asynchronous, respond, {"x-trace-id": None, "x-gateway-key": "gateway-secret"} + ) + + assert response["data"][0]["content"][0]["text"] == "Mock search result" + (call,) = respond.call_args_list + request: Final = call.args[0] + assert str(request.url) == BEDROCK_RETRIEVE_URL + assert "AWS4-HMAC-SHA256 Credential=test_access_key" in request.headers["authorization"] + assert request.headers["x-gateway-key"] == "gateway-secret" + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.asyncio +async def test_a_non_string_extra_header_error_does_not_echo_its_value(asynchronous: bool) -> None: + respond: Final = MagicMock(return_value=httpx.Response(200, json=BEDROCK_RETRIEVE_RESPONSE)) + + with pytest.raises(openai.APIError) as error: + await _search_knowledge_base(asynchronous, respond, {"x-tenant-pin": 918273645}) + + assert "918273645" not in str(error.value) + respond.assert_not_called() + + def _image_search_body(media_type: str, extra_body: dict[str, object] | None) -> dict[str, object]: config: Final = BedrockVectorStoreConfig() mock_log: Final = MagicMock() 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 c55393f2e06..6afa447c294 100644 --- a/tests/unit/llms/opensearch/vector_stores/test_opensearch_transformation.py +++ b/tests/unit/llms/opensearch/vector_stores/test_opensearch_transformation.py @@ -3,11 +3,13 @@ import json import re from collections.abc import Callable, Mapping from datetime import datetime, timezone +from functools import partial from pathlib import Path from types import MappingProxyType from typing import Final from unittest.mock import MagicMock +import httpcore import httpx import openai import pytest @@ -15,6 +17,7 @@ from botocore.auth import SigV4Auth from botocore.awsrequest import AWSRequest from botocore.credentials import Credentials, RefreshableCredentials from botocore.exceptions import CredentialRetrievalError +from h2.events import ConnectionTerminated import litellm from litellm.litellm_core_utils.litellm_logging import Logging @@ -242,10 +245,31 @@ async def test_opensearch_error_keeps_its_status(asynchronous: bool, status: int assert error.value.status_code == status +MALFORMED_FILTERS: Final = ( + {"filters": {"type": "between", "key": "revision", "value": [6, 8]}}, + {"filters": {"type": "gte", "key": "published", "value": True}}, + {"filters": {"type": "in", "key": "category", "value": "hr"}}, + {"filters": {"type": "eq", "key": "category"}}, + {"filters": {"type": "eq", "key": "category", "value": "hr", "case_sensitive": False}}, + {"filters": {"type": "and", "filters": [{"type": "eq", "key": "category", "value": "hr"}, {"type": "near"}]}}, + {"filters": {"type": "gt", "key": "score", "value": float("nan")}}, + {"filters": {"type": "in", "key": "score", "value": [0.5, float("inf")]}}, +) +NON_OBJECT_FILTERS: Final = ({"filters": "category = hr"}, {"filters": []}, {"filters": 5}) + + @pytest.mark.parametrize( "overrides,message", [ - ({"filters": {"type": "eq", "key": "category", "value": "hr"}}, "filters"), + *((malformed, "filters must be an OpenAI vector store filter") for malformed in MALFORMED_FILTERS), + *((non_object, "Invalid vector store search parameters") for non_object in NON_OBJECT_FILTERS), + *( + ({**non_object, "drop_params": True}, "Invalid vector store search parameters") + for non_object in NON_OBJECT_FILTERS + ), + ({"opensearch_hybrid_search": True}, "opensearch_hybrid_search needs opensearch_search_pipeline"), + ({"opensearch_hybrid_search": "sometimes"}, "opensearch_hybrid_search must be a boolean"), + ({"opensearch_search_pipeline": ["hybrid-norm"]}, "opensearch_search_pipeline"), ({"ranking_options": {"score_threshold": 0.5}}, "ranking_options"), ({"rewrite_query": True}, "rewrite_query"), ({"max_num_results": 0}, "max_num_results"), @@ -293,7 +317,6 @@ _AMBIENT_AWS_ENV: Final = ( ) - def _isolate_from_ambient_aws(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: for name in _AMBIENT_AWS_ENV: monkeypatch.delenv(name, raising=False) @@ -417,9 +440,7 @@ class _CredentialSourceThatCannotRefresh(BaseAWSLLM): ) -@pytest.mark.parametrize( - "refresh", [_refresh_returning_expired_credentials, _refresh_after_the_token_file_is_gone] -) +@pytest.mark.parametrize("refresh", [_refresh_returning_expired_credentials, _refresh_after_the_token_file_is_gone]) @pytest.mark.parametrize( "sign", [ @@ -450,7 +471,9 @@ def test_credentials_that_fail_to_refresh_at_signing_are_a_500( @pytest.mark.asyncio async def test_the_index_pass_through_answers_500_when_credentials_fail_to_refresh() -> None: refused: Final = await prepare_index_request( - config=OpenSearchVectorStoreConfig(aws=_CredentialSourceThatCannotRefresh(_refresh_after_the_token_file_is_gone)), + config=OpenSearchVectorStoreConfig( + aws=_CredentialSourceThatCannotRefresh(_refresh_after_the_token_file_is_gone) + ), method="GET", path=IndexPath(index_name="docs", api=("_mapping",)), operation=index_operation("GET", ("_mapping",)), @@ -485,7 +508,8 @@ async def test_opensearch_index_name_picks_the_index_so_the_store_id_can_differ( @pytest.mark.parametrize( "unsupported", [ - {"filters": {"type": "eq", "key": "category", "value": "hr"}}, + *MALFORMED_FILTERS, + {"opensearch_hybrid_search": True}, {"ranking_options": {"score_threshold": 0.5}}, {"rewrite_query": True}, ], @@ -499,6 +523,7 @@ async def test_unsupported_search_params_are_dropped_under_drop_params(unsupport ) assert response["data"] == EXPECTED_RESULTS + assert str(_sent(respond)[0].url) == f"{SERVERLESS}/litellm-qa-docs/_search" assert json.loads(_sent(respond)[0].content) == { "size": 3, "query": {"knn": {"embedding": {"vector": QUERY_VECTOR, "k": 3}}}, @@ -518,6 +543,859 @@ async def test_no_op_values_of_unsupported_params_are_accepted() -> None: assert response["data"] == EXPECTED_RESULTS +DYNAMIC_MAPPING: Final = { + "properties": { + "category": {"type": "text", "fields": {"keyword": {"type": "keyword", "ignore_above": 256}}}, + "published": {"type": "boolean"}, + "published_at": {"type": "date"}, + "revision": {"type": "long"}, + "score": {"type": "float"}, + } +} +HR_EXACT: Final = {"term": {"category.keyword": "hr"}} + + +def _index_mappings(mappings: Mapping[str, Mapping[str, object]]) -> httpx.Response: + return httpx.Response(200, json={index: {"mappings": mapping} for index, mapping in mappings.items()}) + + +def _index(mapping: Mapping[str, object], *searches: httpx.Response) -> MagicMock: + answers: Final = iter(searches) + return MagicMock( + side_effect=lambda request: ( + _index_mappings({"litellm-qa-docs": mapping}) if request.method == "GET" else next(answers) + ) + ) + + +def _searches(respond: MagicMock) -> tuple[httpx.Request, ...]: + return tuple(request for request in _sent(respond) if request.method == "POST") + + +def _mapping_reads(respond: MagicMock) -> tuple[httpx.Request, ...]: + return tuple(request for request in _sent(respond) if request.method == "GET") + + +@pytest.mark.parametrize( + "search_filter,opensearch_filter", + [ + pytest.param({"type": "eq", "key": "category", "value": "hr"}, HR_EXACT, id="eq"), + pytest.param({"type": "eq", "key": "published", "value": True}, {"term": {"published": True}}, id="eq-bool"), + pytest.param({"type": "ne", "key": "category", "value": "hr"}, {"bool": {"must_not": [HR_EXACT]}}, id="ne"), + pytest.param( + {"type": "ne", "key": "revision", "value": 7}, + {"bool": {"must_not": [{"term": {"revision": 7}}]}}, + id="ne-int", + ), + *( + pytest.param( + {"type": operator, "key": "revision", "value": 7}, {"range": {"revision": {operator: 7}}}, id=operator + ) + for operator in ("gt", "gte", "lt", "lte") + ), + pytest.param({"type": "gte", "key": "score", "value": 0.5}, {"range": {"score": {"gte": 0.5}}}, id="gte-float"), + pytest.param( + {"type": "lt", "key": "published_at", "value": "2026-01-01"}, + {"range": {"published_at": {"lt": "2026-01-01"}}}, + id="lt-date", + ), + pytest.param( + {"type": "in", "key": "category", "value": ["hr", "legal"]}, + {"terms": {"category.keyword": ["hr", "legal"]}}, + id="in", + ), + pytest.param({"type": "in", "key": "revision", "value": [6, 7]}, {"terms": {"revision": [6, 7]}}, id="in-int"), + pytest.param( + {"type": "nin", "key": "revision", "value": [6, 7]}, + {"bool": {"must_not": [{"terms": {"revision": [6, 7]}}]}}, + id="nin", + ), + pytest.param( + { + "type": "and", + "filters": [ + {"type": "eq", "key": "category", "value": "hr"}, + { + "type": "or", + "filters": [ + {"type": "gte", "key": "revision", "value": 7}, + {"type": "eq", "key": "published", "value": False}, + ], + }, + ], + }, + { + "bool": { + "filter": [ + HR_EXACT, + { + "bool": { + "should": [{"range": {"revision": {"gte": 7}}}, {"term": {"published": False}}], + "minimum_should_match": 1, + } + }, + ] + } + }, + id="and-of-or", + ), + ], +) +@pytest.mark.asyncio +async def test_filters_narrow_the_knn_query_with_the_matching_opensearch_filter( + search_filter: Mapping[str, object], opensearch_filter: Mapping[str, object] +) -> None: + respond: Final = _index(DYNAMIC_MAPPING, httpx.Response(200, json=_opensearch_hits("text", "embedding"))) + + response: Final = await _search( + True, respond, RecordingEmbeddingExecutor(), {"filters": search_filter, "max_num_results": 3} + ) + + assert response["data"] == EXPECTED_RESULTS + (search,) = _searches(respond) + assert json.loads(search.content) == { + "size": 3, + "query": {"knn": {"embedding": {"vector": QUERY_VECTOR, "k": 3, "filter": opensearch_filter}}}, + "_source": {"excludes": ["embedding"]}, + } + + +NESTED_TEAM: Final = {"team": {"type": "keyword"}} + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize( + "mapping,search_filter,opensearch_filter", + [ + pytest.param( + {"properties": {"tenant": {"type": "keyword"}}}, + {"type": "eq", "key": "tenant", "value": "acme"}, + {"term": {"tenant": "acme"}}, + id="keyword-field", + ), + pytest.param( + {"properties": {"tenant": {"type": "text", "fields": {"raw": {"type": "keyword"}}}}}, + {"type": "ne", "key": "tenant", "value": "acme"}, + {"bool": {"must_not": [{"term": {"tenant.raw": "acme"}}]}}, + id="text-with-a-keyword-subfield", + ), + pytest.param( + {"properties": {"tenant": {"type": "text", "fields": {"keyword": {"type": "keyword"}}}}}, + {"type": "gte", "key": "tenant", "value": "b"}, + {"range": {"tenant.keyword": {"gte": "b"}}}, + id="string-range", + ), + pytest.param( + {"properties": {"owner": {"properties": {"team": {"type": "keyword"}}}}}, + {"type": "in", "key": "owner.team", "value": ["people-ops"]}, + {"terms": {"owner.team": ["people-ops"]}}, + id="object-subfield", + ), + pytest.param( + { + "properties": { + "tenant": { + "type": "text", + "fields": { + "keyword": {"type": "keyword", "normalizer": "lowercase"}, + "raw": {"type": "keyword"}, + }, + } + } + }, + {"type": "eq", "key": "tenant", "value": "Acme"}, + {"term": {"tenant.raw": "Acme"}}, + id="skips-a-normalized-keyword-subfield", + ), + pytest.param( + { + "properties": { + "tenant": {"type": "keyword", "normalizer": "lowercase", "fields": {"raw": {"type": "keyword"}}} + } + }, + {"type": "eq", "key": "tenant", "value": "Acme"}, + {"term": {"tenant.raw": "Acme"}}, + id="normalized-keyword-with-a-plain-subfield", + ), + pytest.param( + {"properties": {"category": {"type": "keyword"}}}, + {"type": "eq", "key": "tenant", "value": "acme"}, + {"term": {"tenant": "acme"}}, + id="unmapped", + ), + pytest.param( + {"properties": {"ref": {"type": "text", "fields": {"keyword": {"type": "keyword"}}}}}, + {"type": "eq", "key": "ref", "value": 7}, + {"term": {"ref.keyword": 7}}, + id="number-on-text-with-a-keyword-subfield", + ), + pytest.param( + {"properties": {"ref_code": {"type": "keyword"}, "ref": {"type": "alias", "path": "ref_code"}}}, + {"type": "eq", "key": "ref", "value": "FAC-017"}, + {"term": {"ref": "FAC-017"}}, + id="alias-to-a-keyword-field", + ), + pytest.param( + { + "properties": { + "owner": {"properties": {"team": {"type": "text", "fields": {"raw": {"type": "keyword"}}}}}, + "team": {"type": "alias", "path": "owner.team"}, + } + }, + {"type": "in", "key": "team", "value": ["People Ops"]}, + {"terms": {"owner.team.raw": ["People Ops"]}}, + id="alias-to-text-with-a-keyword-subfield", + ), + pytest.param( + {"properties": {"owner": {"type": "nested", "include_in_root": True, "properties": NESTED_TEAM}}}, + {"type": "eq", "key": "owner.team", "value": "people-ops"}, + {"term": {"owner.team": "people-ops"}}, + id="nested-object-copied-to-the-root", + ), + pytest.param( + {"properties": {"owner": {"type": "nested", "include_in_parent": True, "properties": NESTED_TEAM}}}, + {"type": "in", "key": "owner.team", "value": ["people-ops"]}, + {"terms": {"owner.team": ["people-ops"]}}, + id="top-level-nested-object-copied-to-its-parent", + ), + pytest.param( + { + "properties": { + "org": { + "type": "nested", + "properties": {"owner": {"type": "nested", "include_in_root": True, "properties": NESTED_TEAM}}, + } + } + }, + {"type": "eq", "key": "org.owner.team", "value": "people-ops"}, + {"term": {"org.owner.team": "people-ops"}}, + id="nested-in-nested-copied-to-the-root", + ), + pytest.param( + { + "properties": { + "org": { + "type": "nested", + "include_in_parent": True, + "properties": { + "owner": {"type": "nested", "include_in_parent": True, "properties": NESTED_TEAM} + }, + } + } + }, + {"type": "eq", "key": "org.owner.team", "value": "people-ops"}, + {"term": {"org.owner.team": "people-ops"}}, + id="nested-in-nested-copied-up-through-each-parent", + ), + ], +) +@pytest.mark.asyncio +async def test_filters_compare_the_exact_field_the_index_maps( + asynchronous: bool, + mapping: Mapping[str, object], + search_filter: Mapping[str, object], + opensearch_filter: Mapping[str, object], +) -> None: + respond: Final = _index(mapping, httpx.Response(200, json=_opensearch_hits("text", "embedding"))) + + await _search(asynchronous, respond, RecordingEmbeddingExecutor(), {"filters": search_filter, "max_num_results": 3}) + + (mapping_read,) = _mapping_reads(respond) + assert str(mapping_read.url) == f"{SERVERLESS}/litellm-qa-docs/_mapping" + assert _signature_scope(mapping_read) == ("aoss", "us-east-1") + (search,) = _searches(respond) + assert json.loads(search.content)["query"]["knn"]["embedding"]["filter"] == opensearch_filter + + +TEXT_ONLY: Final = {"litellm-qa-docs": {"properties": {"tenant": {"type": "text"}}}} + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize( + "mappings,search_filter,reason", + [ + pytest.param( + TEXT_ONLY, + {"type": "eq", "key": "tenant", "value": "team"}, + "the index maps tenant as text with no keyword subfield", + id="text-only", + ), + pytest.param( + TEXT_ONLY, + { + "type": "or", + "filters": [ + {"type": "eq", "key": "revision", "value": 7}, + {"type": "nin", "key": "tenant", "value": ["team other"]}, + ], + }, + "the index maps tenant as text with no keyword subfield", + id="text-only-in-a-compound", + ), + pytest.param( + { + "litellm-qa-docs": { + "properties": { + "tenant": {"type": "text", "fields": {"keyword": {"type": "keyword", "ignore_above": 8}}} + } + } + }, + {"type": "eq", "key": "tenant", "value": "acme quarterly"}, + "the index ignores tenant values longer than 8 characters", + id="longer-than-ignore-above", + ), + pytest.param( + { + "docs-2025": {"properties": {"tenant": {"type": "keyword"}}}, + "docs-2026": {"properties": {"tenant": {"type": "text", "fields": {"keyword": {"type": "keyword"}}}}}, + }, + {"type": "eq", "key": "tenant", "value": "acme"}, + "the indexes behind this store map tenant differently", + id="alias-indexes-disagree", + ), + pytest.param( + { + "docs-2025": { + "properties": { + "owner": {"properties": {"team": {"type": "text", "fields": {"raw": {"type": "keyword"}}}}}, + "team": {"type": "alias", "path": "owner.team"}, + } + }, + "docs-2026": { + "properties": { + "owner": {"properties": {"team": {"type": "text", "fields": {"raw": {"type": "keyword"}}}}} + } + }, + }, + {"type": "eq", "key": "team", "value": "People Ops"}, + "the indexes behind this store map team differently", + id="alias-rewrite-reaches-an-index-without-the-alias", + ), + pytest.param( + { + "docs-2025": {"properties": {"team": {"type": "keyword"}}}, + "docs-2026": {"properties": {"category": {"type": "keyword"}}}, + }, + {"type": "eq", "key": "team", "value": "HR"}, + "the indexes behind this store map team differently. Map team to the same keyword field in every " + "index behind the store", + id="only-some-indexes-map-the-key", + ), + pytest.param( + { + "docs-2025": { + "properties": {"category": {"type": "keyword"}, "team": {"type": "alias", "path": "category"}} + }, + "docs-2026": {"properties": {"category": {"type": "keyword"}}}, + }, + {"type": "eq", "key": "team", "value": "HR"}, + "the indexes behind this store map team differently. Map team to the same keyword field in every " + "index behind the store", + id="only-some-indexes-alias-the-key", + ), + pytest.param( + {"litellm-qa-docs": {"properties": {"tenant": {"type": "keyword", "normalizer": "lowercase"}}}}, + {"type": "eq", "key": "tenant", "value": "Acme"}, + "the index normalizes every keyword mapping of tenant", + id="normalized-keyword", + ), + pytest.param( + { + "litellm-qa-docs": { + "properties": { + "tenant": { + "type": "text", + "fields": {"keyword": {"type": "keyword", "normalizer": "lowercase"}}, + } + } + } + }, + {"type": "eq", "key": "tenant", "value": "Acme"}, + "the index normalizes every keyword mapping of tenant", + id="text-with-only-a-normalized-keyword-subfield", + ), + pytest.param( + {"litellm-qa-docs": {"properties": {"year": {"type": "keyword"}}}}, + {"type": "gte", "key": "year", "value": 2026}, + "the index maps year as keyword, which compares numbers as text. Map year as a numeric field", + id="number-range-on-a-keyword-field", + ), + pytest.param( + {"litellm-qa-docs": {"properties": {"year": {"type": "text", "fields": {"keyword": {"type": "keyword"}}}}}}, + {"type": "lt", "key": "year", "value": 2026.5}, + "the index maps year as text, which compares numbers as text", + id="number-range-on-text-with-a-keyword-subfield", + ), + pytest.param( + TEXT_ONLY, + {"type": "eq", "key": "tenant", "value": 7}, + "the index maps tenant as text with no keyword subfield", + id="number-on-text-only", + ), + pytest.param( + {"litellm-qa-docs": {"properties": {"title": {"type": "search_as_you_type"}}}}, + {"type": "eq", "key": "title", "value": "HR"}, + "the index maps title as search_as_you_type with no keyword subfield", + id="search-as-you-type", + ), + pytest.param( + { + "litellm-qa-docs": { + "properties": {"tenant": {"type": "text"}, "org": {"type": "alias", "path": "tenant"}} + } + }, + {"type": "ne", "key": "org", "value": "acme"}, + "the index maps tenant as text with no keyword subfield", + id="alias-to-text-only", + ), + pytest.param( + {"litellm-qa-docs": {"properties": {"owner": {"type": "nested", "properties": NESTED_TEAM}}}}, + {"type": "eq", "key": "owner.team", "value": "people-ops"}, + "the index maps owner as nested, which needs a nested query. Map owner as object", + id="field-inside-a-nested-object", + ), + pytest.param( + { + "litellm-qa-docs": { + "properties": { + "owner": {"type": "nested", "properties": NESTED_TEAM}, + "team": {"type": "alias", "path": "owner.team"}, + } + } + }, + {"type": "in", "key": "team", "value": ["people-ops"]}, + "the index maps owner as nested, which needs a nested query", + id="alias-into-a-nested-object", + ), + pytest.param( + {"litellm-qa-docs": {"properties": {"owner": {"type": "nested", "properties": NESTED_TEAM}}}}, + {"type": "ne", "key": "owner", "value": "people-ops"}, + "the index maps owner as nested, which holds fields rather than values. Filter on a field inside owner", + id="the-nested-object-itself", + ), + pytest.param( + { + "litellm-qa-docs": { + "properties": {"org": {"properties": {"owner": {"type": "nested", "properties": NESTED_TEAM}}}} + } + }, + {"type": "eq", "key": "org.owner.team", "value": "people-ops"}, + "the index maps org.owner as nested, which needs a nested query. Map org.owner as object", + id="nested-object-inside-a-plain-object", + ), + pytest.param( + { + "litellm-qa-docs": { + "properties": { + "org": { + "type": "nested", + "properties": { + "owner": {"type": "nested", "include_in_parent": True, "properties": NESTED_TEAM} + }, + } + } + } + }, + {"type": "eq", "key": "org.owner.team", "value": "people-ops"}, + "the index maps org as nested, which needs a nested query. Map org as object", + id="nested-object-copied-only-into-a-nested-parent", + ), + pytest.param( + { + "litellm-qa-docs": { + "properties": { + "org": { + "type": "nested", + "include_in_root": True, + "properties": {"owner": {"type": "nested", "properties": NESTED_TEAM}}, + } + } + } + }, + {"type": "eq", "key": "org.owner.team", "value": "people-ops"}, + "the index maps org.owner as nested, which needs a nested query. Map org.owner as object", + id="nested-object-inside-one-copied-to-the-root", + ), + pytest.param( + {"litellm-qa-docs": {"properties": {"owner": {"properties": NESTED_TEAM}}}}, + {"type": "eq", "key": "owner", "value": "people-ops"}, + "the index maps owner as object, which holds fields rather than values. Filter on a field inside owner", + id="an-object", + ), + ], +) +@pytest.mark.asyncio +async def test_a_filter_the_index_cannot_match_exactly_is_a_400_unless_drop_params_drops_the_filters( + asynchronous: bool, mappings: Mapping[str, Mapping[str, object]], search_filter: Mapping[str, object], reason: str +) -> None: + rejected: Final = MagicMock(return_value=_index_mappings(mappings)) + executor: Final = RecordingEmbeddingExecutor() + + with pytest.raises(openai.BadRequestError) as error: + await _search(asynchronous, rejected, executor, {"filters": search_filter}) + + assert reason in str(error.value) + assert not _searches(rejected) + executor.call.assert_not_called() + + dropped: Final = MagicMock( + side_effect=[_index_mappings(mappings), httpx.Response(200, json=_opensearch_hits("text", "embedding"))] + ) + response: Final = await _search( + asynchronous, dropped, RecordingEmbeddingExecutor(), {"filters": search_filter, "drop_params": True} + ) + + assert response["data"] == EXPECTED_RESULTS + (search,) = _searches(dropped) + assert json.loads(search.content)["query"] == {"knn": {"embedding": {"vector": QUERY_VECTOR, "k": 10}}} + + +def _mapping_read_error(status: int) -> httpx.Response: + return httpx.Response(status, json={"status": status, "error": {"reason": "no mapping for you", "type": "x"}}) + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.asyncio +async def test_a_refused_index_mapping_read_names_the_permission_it_needs_and_sends_no_search( + asynchronous: bool, +) -> None: + respond: Final = MagicMock(return_value=_mapping_read_error(403)) + executor: Final = RecordingEmbeddingExecutor() + + with pytest.raises(openai.APIStatusError) as error: + await _search(asynchronous, respond, executor, {"filters": {"type": "eq", "key": "tenant", "value": "acme"}}) + + assert error.value.status_code == 403 + assert "es:ESHttpGet" in str(error.value) + assert "no mapping for you" not in str(error.value) + assert not _searches(respond) + executor.call.assert_not_called() + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("drop_params", [False, True]) +@pytest.mark.parametrize( + "mapping_read,status", + [ + *(pytest.param(partial(_mapping_read_error, code), code, id=str(code)) for code in (404, 429, 500, 503)), + pytest.param(partial(httpx.Response, 200, json={"acknowledged": True}), 502, id="unreadable"), + ], +) +@pytest.mark.asyncio +async def test_a_failed_index_mapping_read_keeps_the_filters_and_its_status_even_under_drop_params( + asynchronous: bool, drop_params: bool, mapping_read: Callable[[], httpx.Response], status: int +) -> None: + respond: Final = MagicMock(side_effect=lambda request: mapping_read()) + executor: Final = RecordingEmbeddingExecutor() + + with pytest.raises(openai.APIStatusError) as error: + await _search( + asynchronous, + respond, + executor, + {"filters": {"type": "eq", "key": "tenant", "value": "acme"}, "drop_params": drop_params}, + ) + + assert error.value.status_code == status + assert "es:ESHttpGet" not in str(error.value) + assert "no mapping for you" not in str(error.value) + assert not _searches(respond) + executor.call.assert_not_called() + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.asyncio +async def test_a_refused_index_mapping_read_under_drop_params_drops_the_filters_and_searches_once( + asynchronous: bool, +) -> None: + respond: Final = MagicMock( + side_effect=[_mapping_read_error(403), httpx.Response(200, json=_opensearch_hits("text", "embedding"))] + ) + + response: Final = await _search( + asynchronous, + respond, + RecordingEmbeddingExecutor(), + {"filters": {"type": "eq", "key": "tenant", "value": "acme"}, "drop_params": True}, + ) + + assert response["data"] == EXPECTED_RESULTS + (search,) = _searches(respond) + assert json.loads(search.content)["query"] == {"knn": {"embedding": {"vector": QUERY_VECTOR, "k": 10}}} + + +def _rejection(reason: str) -> httpx.Response: + return httpx.Response( + 400, + json={ + "error": {"root_cause": [{"type": "query_shard_exception", "reason": reason}], "type": "x"}, + "status": 400, + }, + ) + + +ENGINE_REFUSES_FILTERS: Final = "failed to create query: Engine [NMSLIB] does not support filters" +KNN_PARSER_REFUSES_FILTERS: Final = "[knn] unknown token [START_OBJECT] after [filter]" +FILTERED_KNN: Final = {"knn": {"embedding": {"vector": QUERY_VECTOR, "k": 3, "filter": HR_EXACT}}} +UNFILTERED_KNN: Final = {"knn": {"embedding": {"vector": QUERY_VECTOR, "k": 3}}} + + +def _close_with_goaway(request: httpx.Request) -> httpx.Response: + raise httpx.RemoteProtocolError("GOAWAY", request=request) from httpcore.RemoteProtocolError(ConnectionTerminated()) + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.asyncio +async def test_the_index_mapping_read_sends_the_caller_headers_and_retries_a_goaway(asynchronous: bool) -> None: + answers: Final = iter( + ( + _close_with_goaway, + lambda request: _index_mappings({"litellm-qa-docs": DYNAMIC_MAPPING}), + lambda request: httpx.Response(200, json=_opensearch_hits("text", "embedding")), + ) + ) + respond: Final = MagicMock(side_effect=lambda request: next(answers)(request)) + + response: Final = await _search( + asynchronous, + respond, + RecordingEmbeddingExecutor(), + { + "filters": {"type": "eq", "key": "category", "value": "hr"}, + "max_num_results": 3, + "extra_headers": {"x-gateway-key": "gateway-secret", "x-gateway-token": b"gateway-token"}, + }, + ) + + assert response["data"] == EXPECTED_RESULTS + goaway, mapping_read = _mapping_reads(respond) + assert mapping_read.headers["x-gateway-key"] == "gateway-secret" + assert mapping_read.headers["x-gateway-token"] == "gateway-token" + assert _signature_scope(mapping_read) == ("aoss", "us-east-1") + (search,) = _searches(respond) + assert json.loads(search.content)["query"] == FILTERED_KNN + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("value", [None, 7, "café"]) +@pytest.mark.asyncio +async def test_a_header_httpx_cannot_send_fails_a_filtered_search_like_an_unfiltered_one_before_embedding( + asynchronous: bool, value: object +) -> None: + respond: Final = _index(DYNAMIC_MAPPING, httpx.Response(200, json=_opensearch_hits("text", "embedding"))) + executor: Final = RecordingEmbeddingExecutor() + headers: Final = {"extra_headers": {"x-trace-id": value}} + + with pytest.raises(openai.APIError) as unfiltered: + await _search(asynchronous, respond, RecordingEmbeddingExecutor(), headers) + with pytest.raises(openai.APIError) as filtered: + await _search( + asynchronous, respond, executor, {**headers, "filters": {"type": "eq", "key": "category", "value": "hr"}} + ) + + assert type(filtered.value) is type(unfiltered.value) + assert str(filtered.value) == str(unfiltered.value) + assert not _sent(respond) + executor.call.assert_not_called() + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.asyncio +async def test_a_connection_failure_on_the_index_mapping_read_raises_what_one_on_the_search_raises( + asynchronous: bool, +) -> None: + def refuse(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("connection refused", request=request) + + with pytest.raises(openai.APIError) as on_search: + await _search(asynchronous, MagicMock(side_effect=refuse), RecordingEmbeddingExecutor(), {}) + with pytest.raises(openai.APIError) as on_mapping_read: + await _search( + asynchronous, + MagicMock(side_effect=refuse), + RecordingEmbeddingExecutor(), + {"filters": {"type": "eq", "key": "category", "value": "hr"}}, + ) + + assert type(on_mapping_read.value) is type(on_search.value) + assert getattr(on_mapping_read.value, "status_code", None) == getattr(on_search.value, "status_code", None) + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.asyncio +async def test_a_timeout_on_the_index_mapping_read_answers_408_like_one_on_the_search_without_embedding( + asynchronous: bool, +) -> None: + def time_out(request: httpx.Request) -> httpx.Response: + raise httpx.ReadTimeout("timed out", request=request) + + executor: Final = RecordingEmbeddingExecutor() + with pytest.raises(openai.APIError) as on_search: + await _search(asynchronous, MagicMock(side_effect=time_out), RecordingEmbeddingExecutor(), {}) + with pytest.raises(openai.APIError) as on_mapping_read: + await _search( + asynchronous, + MagicMock(side_effect=time_out), + executor, + {"filters": {"type": "eq", "key": "category", "value": "hr"}}, + ) + + assert getattr(on_search.value, "status_code", None) == 408 + assert type(on_mapping_read.value) is type(on_search.value) + assert getattr(on_mapping_read.value, "status_code", None) == 408 + executor.call.assert_not_called() + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("refusal", [ENGINE_REFUSES_FILTERS, KNN_PARSER_REFUSES_FILTERS]) +@pytest.mark.asyncio +async def test_a_filter_the_index_refuses_is_dropped_under_drop_params_without_embedding_again( + asynchronous: bool, refusal: str +) -> None: + respond: Final = _index( + DYNAMIC_MAPPING, _rejection(refusal), httpx.Response(200, json=_opensearch_hits("text", "embedding")) + ) + executor: Final = RecordingEmbeddingExecutor() + + response: Final = await _search( + asynchronous, + respond, + executor, + {"filters": {"type": "eq", "key": "category", "value": "hr"}, "drop_params": True, "max_num_results": 3}, + ) + + assert response["data"] == EXPECTED_RESULTS + sent: Final = _searches(respond) + assert [json.loads(request.content)["query"] for request in sent] == [FILTERED_KNN, UNFILTERED_KNN] + assert _signature_scope(sent[1]) == ("aoss", "us-east-1") + executor.call.assert_called_once_with("embedding-alias", QUERY, {}) + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("refusal", [ENGINE_REFUSES_FILTERS, KNN_PARSER_REFUSES_FILTERS]) +@pytest.mark.asyncio +async def test_a_filter_the_index_refuses_without_drop_params_is_a_400_that_hides_the_opensearch_error( + asynchronous: bool, refusal: str +) -> None: + respond: Final = _index(DYNAMIC_MAPPING, _rejection(refusal)) + + with pytest.raises(openai.BadRequestError) as error: + await _search( + asynchronous, + respond, + RecordingEmbeddingExecutor(), + {"filters": {"type": "eq", "key": "category", "value": "hr"}, "max_num_results": 3}, + ) + + assert "cannot apply filters inside a k-NN query" in str(error.value) + assert "set drop_params to drop them" in str(error.value) + assert refusal not in str(error.value) + assert "query_shard_exception" not in str(error.value) + assert len(_searches(respond)) == 1 + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize( + "drop_params,rejections,sent_queries,reason", + [ + pytest.param( + {"drop_params": True}, + ["failed to create query: unknown field [category]"], + [FILTERED_KNN], + "unknown field", + id="other-400", + ), + pytest.param( + {"drop_params": True}, + [ENGINE_REFUSES_FILTERS, "index is read-only"], + [FILTERED_KNN, UNFILTERED_KNN], + "index is read-only", + id="retry-rejected", + ), + ], +) +@pytest.mark.asyncio +async def test_a_rejected_filtered_search_is_a_400_unless_drop_params_can_drop_the_filter( + asynchronous: bool, + drop_params: Mapping[str, object], + rejections: list[str], + sent_queries: list[Mapping[str, object]], + reason: str, +) -> None: + respond: Final = _index(DYNAMIC_MAPPING, *(_rejection(reason) for reason in rejections)) + + with pytest.raises(openai.BadRequestError) as error: + await _search( + asynchronous, + respond, + RecordingEmbeddingExecutor(), + {"filters": {"type": "eq", "key": "category", "value": "hr"}, "max_num_results": 3, **drop_params}, + ) + + assert reason in str(error.value) + assert [json.loads(request.content)["query"] for request in _searches(respond)] == sent_queries + + +HYBRID: Final = { + "opensearch_hybrid_search": True, + "opensearch_search_pipeline": "hybrid norm", + "opensearch_text_field": "passage", +} +KEYWORD_QUERY: Final = {"match": {"passage": {"query": QUERY}}} +HR_ONLY: Final = {"type": "eq", "key": "category", "value": "hr"} + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize( + "filters,queries", + [ + pytest.param({}, [KEYWORD_QUERY, {"knn": {"embedding": {"vector": QUERY_VECTOR, "k": 3}}}], id="unfiltered"), + pytest.param( + {"filters": HR_ONLY}, + [ + {"bool": {"must": [KEYWORD_QUERY], "filter": [HR_EXACT]}}, + {"knn": {"embedding": {"vector": QUERY_VECTOR, "k": 3, "filter": HR_EXACT}}}, + ], + id="filtered", + ), + ], +) +@pytest.mark.asyncio +async def test_hybrid_search_sends_a_keyword_and_a_knn_query_through_the_signed_search_pipeline( + asynchronous: bool, filters: Mapping[str, object], queries: list[Mapping[str, object]] +) -> None: + executor: Final = RecordingEmbeddingExecutor() + respond: Final = _index(DYNAMIC_MAPPING, httpx.Response(200, json=_opensearch_hits("passage", "embedding"))) + + response: Final = await _search(asynchronous, respond, executor, {**HYBRID, **filters, "max_num_results": 3}) + + assert response["data"] == EXPECTED_RESULTS + sent: Final = _searches(respond) + assert len(sent) == 1 + assert str(sent[0].url) == f"{SERVERLESS}/litellm-qa-docs/_search?search_pipeline=hybrid%20norm" + assert json.loads(sent[0].content) == { + "size": 3, + "query": {"hybrid": {"queries": queries}}, + "_source": {"excludes": ["embedding"]}, + } + assert _signature_scope(sent[0]) == ("aoss", "us-east-1") + executor.call.assert_called_once_with("embedding-alias", QUERY, {}) + + +@pytest.mark.parametrize("hybrid", [{}, {"opensearch_hybrid_search": False}]) +@pytest.mark.asyncio +async def test_a_search_pipeline_without_hybrid_search_runs_the_knn_query_through_it( + hybrid: Mapping[str, object], +) -> None: + respond: Final = MagicMock(return_value=httpx.Response(200, json={"hits": {"hits": []}})) + + await _search(True, respond, RecordingEmbeddingExecutor(), {**hybrid, "opensearch_search_pipeline": "hybrid-norm"}) + + sent: Final = _sent(respond) + assert str(sent[0].url) == f"{SERVERLESS}/litellm-qa-docs/_search?search_pipeline=hybrid-norm" + assert json.loads(sent[0].content)["query"] == {"knn": {"embedding": {"vector": QUERY_VECTOR, "k": 10}}} + + def test_create_is_rejected_because_the_store_is_search_only() -> None: with pytest.raises(litellm.BadRequestError, match="search-only"): litellm.vector_stores.create(custom_llm_provider="opensearch") diff --git a/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py b/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py index 1d5a1f8fa30..bb24cdd3c49 100644 --- a/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py @@ -2471,15 +2471,24 @@ def test_rag_query_rejects_caller_embedding_selection_params(client_internal_use @pytest.mark.parametrize( - ("provider", "expected_status", "expected_searches"), - [("opensearch", 400, 0), ("opensearch/x", 400, 0), ("openai", 200, 1)], + ("store_key", "provider", "expected_status", "expected_searches"), + [ + ("opensearch_index_name", "opensearch", 400, 0), + ("opensearch_index_name", "opensearch/x", 400, 0), + ("opensearch_index_name", "openai", 200, 1), + ("opensearch_search_pipeline", "opensearch", 400, 0), + ("opensearch_hybrid_search", "opensearch/x", 400, 0), + ("opensearch_search_pipeline", "openai", 200, 1), + ("opensearch_text_field", "opensearch", 200, 1), + ("opensearch_embedding_field", "opensearch/x", 200, 1), + ], ) -def test_rag_query_rejects_caller_opensearch_index_name_only_on_opensearch_stores( - client_internal_user, provider, expected_status, expected_searches +def test_rag_query_rejects_caller_opensearch_store_params_only_on_opensearch_stores( + client_internal_user, store_key, provider, expected_status, expected_searches ): """ - Regression: retrieval_config must not pick the index an OpenSearch store searches, - while other providers keep ignoring the key as they did before OpenSearch support. + Regression: retrieval_config must not change which index, pipeline, or hybrid mode an OpenSearch + store searches with, while the field keys and other providers' requests are answered as before. """ import litellm from litellm import Router @@ -2518,12 +2527,13 @@ def test_rag_query_rejects_caller_opensearch_index_name_only_on_opensearch_store json={ "model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hello"}], - "retrieval_config": {"vector_store_id": "kb", "opensearch_index_name": "another-teams-index"}, + "retrieval_config": {"vector_store_id": "kb", store_key: "_none"}, }, ) assert response.status_code == expected_status, response.json() assert fake_search.await_count == expected_searches + assert expected_status == 200 or store_key in str(response.json()) EICAR = r"X5O!P%@AP[4\PZX54(P^)7CC)7}$EICAR-STANDARD-ANTIVIRUS-TEST-FILE!$H+H*" diff --git a/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py b/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py index 2d0f3679144..77994479b30 100644 --- a/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py +++ b/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py @@ -3055,15 +3055,24 @@ def test_vector_store_search_rejects_caller_embedding_selection_params(blocked_k @pytest.mark.parametrize( - ("provider", "expected_status", "expected_searches"), - [("opensearch", 400, 0), ("opensearch/x", 400, 0), ("openai", 200, 1)], + ("store_key", "provider", "expected_status", "expected_searches"), + [ + ("opensearch_index_name", "opensearch", 400, 0), + ("opensearch_index_name", "opensearch/x", 400, 0), + ("opensearch_index_name", "openai", 200, 1), + ("opensearch_search_pipeline", "opensearch", 400, 0), + ("opensearch_hybrid_search", "opensearch/x", 400, 0), + ("opensearch_search_pipeline", "openai", 200, 1), + ("opensearch_text_field", "opensearch", 200, 1), + ("opensearch_embedding_field", "opensearch/x", 200, 1), + ], ) -def test_vector_store_search_rejects_caller_opensearch_index_name_only_on_opensearch_stores( - provider, expected_status, expected_searches +def test_vector_store_search_rejects_caller_opensearch_store_params_only_on_opensearch_stores( + store_key, provider, expected_status, expected_searches ): """ - Regression: a caller must not pick the index an OpenSearch store searches, while - other providers keep ignoring the key as they did before OpenSearch support. + Regression: a caller must not change which index, pipeline, or hybrid mode an OpenSearch store + searches with, while the field keys and other providers' requests pass through as before. """ from fastapi.testclient import TestClient @@ -3097,13 +3106,14 @@ def test_vector_store_search_rejects_caller_opensearch_index_name_only_on_opense with patch("litellm.proxy.proxy_server.llm_router", router): # test-quality-ok: router built over the fake search response = TestClient(app).post( "/v1/vector_stores/kb/search", - json={"query": "q", "opensearch_index_name": "another-teams-index"}, + json={"query": "q", store_key: "_none"}, ) finally: app.dependency_overrides = original_overrides assert response.status_code == expected_status, response.json() - assert fake_search.await_count == expected_searches + assert [call.kwargs.get(store_key) for call in fake_search.await_args_list] == ["_none"] * expected_searches + assert expected_status == 200 or store_key in str(response.json()) def test_vector_store_search_accepts_a_query_image_without_a_query():