feat(opensearch): map vector store search filters to k-NN filters and add hybrid search (#51131)

* feat(opensearch): map vector store search filters to k-NN filters and add hybrid search

* test(code-quality): allow the bounded OpenSearch filter recursion

* fix(opensearch): retry an unfilterable engine without filters under drop_params, match strings on keyword subfields, keep per-request field keys

Under drop_params, a 400 'does not support filters' from an nmslib or older
engine re-runs the search once without filters, as filtered searches behaved
before. String eq, ne, in and nin match the key's .keyword subfield, falling
back to the key where no subfield is mapped. Only the index name, search
pipeline and hybrid flag are registration-only, so opensearch_text_field and
opensearch_embedding_field pass through per request again

* fix(opensearch): match string filters on the exact keyword field and embed once on the filter retry

String filters now read the index mapping and compare the exact keyword field, refusing a key the index cannot match exactly. The drop_params retry reuses the first embedding and also covers OpenSearch before 2.4. An index that refuses filters without drop_params answers litellm's own 400, and a filters value that is not an object stays a 400 under drop_params. The Admin UI store form gains the search pipeline and hybrid search fields

* test(ui): expect the OpenSearch search pipeline and hybrid search fields in the provider field list

* fix(opensearch): drop string filters under drop_params when the mapping read fails, and skip normalized keyword fields

* fix(opensearch): name the permissions a refused index mapping read needs instead of returning OpenSearch's body

* fix(opensearch): check the index mapping for every filter, resolve aliases, and keep failed mapping reads from dropping filters

* fix(opensearch): keep an alias filter's own name and refuse rewrites that reach an index without the alias

* fix(vector_stores): sign the caller's headers as given and reject filters on keys only some OpenSearch indexes map

A null or non-str extra_headers value no longer fails Bedrock Knowledge Base, S3 Vectors, or other provider searches before signing. The mapping read gets only the string-valued headers.

A filter on a key that some of a store's indexes map and others do not answers 400 (drops under drop_params), since OpenSearch returns its k-NN hits on some calls and none on others.

* fix(opensearch): send the caller's headers on the index mapping read as given and name the fix for keys only some indexes map

A header httpx cannot send now fails the filtered search on the mapping read, the same way it fails the search, before the query is embedded. Bytes headers now reach the mapping read too.

* test(opensearch): pin that a timeout on the index mapping read answers 408 without embedding the query

* fix(opensearch): answer 400 for a filter key inside a nested field or naming an object

A root-level term query on a field under a nested mapping, or on an object field, matches nothing, so the search answered 200 with no hits. These keys now answer 400 naming the mapping, or drop the filters under drop_params

* fix(opensearch): filter on a nested field whose values OpenSearch copies to the root

A nested field mapped with include_in_root, or include_in_parent on it and
every nested field above it, has its values in the root document, so a root
term query matches it. Those filters now search instead of answering 400.
A key that names an object or nested field answers that it holds fields
rather than values before the nested check, so its remedy no longer leads
into a second 400

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
zoveki 2026-11-13 19:30:55 -08:00 • committed by GitHub
parent 752eb7c5b7
commit 01de1a686d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
17 changed files with 1852 additions and 193 deletions

View file

@ -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,

View file

@ -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(

View file

@ -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,

View file

@ -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"\.(?P<service>aoss|aos)\.(?P<region>[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]

View file

@ -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",

View file

@ -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",

View file

@ -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();

View file

@ -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<typeof vectorStoreSchema>;
@ -540,10 +567,12 @@ const VectorStoreForm: React.FC<VectorStoreFormProps> = ({
);
};
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<ProviderFieldProps> = ({ 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 (
<FormField control={control} name={field.name} label={label}>
{({ id, value, onChange }) => <Switch id={id} checked={value === true} onCheckedChange={onChange} />}
</FormField>
);
}
if (field.type === "select") {
const selectOptions =
field.options ??

View file

@ -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", () => {

View file

@ -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<string, VectorStoreFieldConfig[]>
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: [
{

View file

@ -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)",

View file

@ -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.

View file

@ -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}"

View file

@ -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()

View file

@ -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")

View file

@ -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*"

View file

@ -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():