From 22252073e7f0b0c18236b10ac4cc5d6060d04992 Mon Sep 17 00:00:00 2001 From: zoveki <349028329+zoveki@users.noreply.github.com> Date: Wed, 11 Nov 2026 16:19:34 -0800 Subject: [PATCH] feat(bedrock): search knowledge bases by image and return image hits (#51063) * feat(bedrock): search knowledge bases by image and return image hits vector_stores.search and asearch take a query_image (base64 jpeg, png, gif or webp), and query is optional when it is sent. Bedrock sends it as an IMAGE retrievalQuery, and a caller's own retrievalQuery in extra_body still wins Bedrock IMAGE hits now come back as image content items instead of being dropped. The media type comes from the source file mime type metadata, then the data URI label, then application/octet-stream A request with both query and query_image, with neither, or with a malformed image gets a 400. So does query_image on a provider without image search, unless drop_params is on. The proxy accepts query_image in place of query on /v1/vector_stores/{id}/search Spend logs redact image data when prompts are not stored. Emulated file_search skips image-only hits, RAG rerank maps its indices to the text documents it ranked, and the dashboard tester renders image hits * fix(vector-stores): keep image hits to image searches and reject a caller retrievalQuery Text searches on a multimodal Bedrock knowledge base return the same text hits as before, so chat, RAG and logging payloads stay unchanged. Image hits come back only when the caller sent query_image A caller retrievalQuery now answers 400 pointing at query_image, or is dropped under drop_params, instead of replacing the built query unchecked A file_search chat tool that carries query_image or rewrite_query searches again instead of failing, and a non-string query next to an image answers 400 instead of 500. query_image is keyword-only on search and asearch * fix(vector-stores): redact query images and keep store search params out of image checks * fix(vector-stores): redact images nested in extra_body and keep a store's extra_body retrievalQuery out of searches * fix(vector_stores): keep a caller's image query through store and router params A deployment's retrievalQuery no longer reaches a routed search, a caller's query fields survive a store's own extra_body, and an extra_body image stops before the provider. file_search forwards rewrite_query, both retrieval query spellings are redacted, and the cost error stops echoing the response when message logging is off * fix(cost): leave image search responses out of the cost error text * fix(vector_stores): accept line-wrapped base64 in query_image data Bedrock takes inlineContent wrapped with \n or \r\n line breaks, as plain base64 writes it, so the check ignores only those before decoding and the data goes to Bedrock as sent * fix(vector_stores): drop rewrite_query false from the chat file_search tool False asks for the default search every store already runs, and MongoDB rejects the field outright, so forwarding it made a MongoDB chat lose the store's context for a no-op flag. True is still forwarded. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/cost_calculator.py | 21 +- .../vector_store_pre_call_hook.py | 3 +- litellm/litellm_core_utils/redact_messages.py | 37 +++- .../bedrock/vector_stores/transformation.py | 116 ++++++++-- litellm/proxy/route_llm_request.py | 5 +- .../proxy/vector_store_endpoints/endpoints.py | 14 +- litellm/router_rust/invoke/endpoints.py | 13 +- .../integrations/rag/bedrock_knowledgebase.py | 12 +- litellm/types/vector_stores.py | 26 ++- litellm/vector_stores/main.py | 35 ++- litellm/vector_stores/search_query.py | 205 ++++++++++++++++++ .../vector_stores/vector_store_registry.py | 27 +-- .../test_vector_store_pre_call_hook.py | 34 ++- .../test_redact_messages.py | 63 ++++++ ...est_bedrock_vector_store_transformation.py | 156 +++++++++++++ .../test_mongodb_transformation.py | 37 ++++ tests/unit/proxy/test_route_llm_request.py | 4 + .../test_vector_store_endpoints.py | 80 +++++++ tests/unit/test_cost_calculator.py | 36 +++ ...test_endpoints_characterization_generic.py | 23 ++ tests/unit/vector_stores/test_main.py | 141 ++++++++++++ .../test_vector_store_registry.py | 23 ++ 22 files changed, 1036 insertions(+), 75 deletions(-) create mode 100644 litellm/vector_stores/search_query.py diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 233ff48e8a3..a37160b4528 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -64,6 +64,7 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import ( select_cost_metric_for_model, ) from litellm.litellm_core_utils.llm_response_utils.response_metadata import response_ms_since_logging_start +from litellm.litellm_core_utils.redact_messages import should_redact_message_logging from litellm.llms.anthropic.cost_calculation import ( cost_per_token as anthropic_cost_per_token, ) @@ -1665,6 +1666,14 @@ def _response_time_ms_for_cost( return litellm_logging_obj.get_response_ms() +def _model_is_none_error(completion_response_text: str, litellm_logging_obj: LitellmLoggingObject | None) -> ValueError: + message: Final = "Model is None and does not exist in passed completion_response." + model_call_details: Final = {} if litellm_logging_obj is None else litellm_logging_obj.model_call_details + if should_redact_message_logging(model_call_details) or model_call_details.get("query_image") is not None: + return ValueError(message) + return ValueError(f"{message} Passed completion_response={completion_response_text}, model=None") + + def _breakdown_total(litellm_logging_obj: LitellmLoggingObject | None) -> float | None: breakdown: Final = None if litellm_logging_obj is None else litellm_logging_obj.cost_breakdown total: Final[object] = None if breakdown is None else breakdown.get("total_cost") @@ -2013,9 +2022,7 @@ def completion_cost( else: if model is None: - raise ValueError( - f"Model is None and does not exist in passed completion_response. Passed completion_response={completion_response}, model={model}" - ) + raise _model_is_none_error(f"{completion_response}", litellm_logging_obj) if len(messages) > 0: prompt_tokens = token_counter(model=model, messages=messages) elif len(prompt) > 0: @@ -2029,9 +2036,7 @@ def completion_cost( return A2ACostCalculator.calculate_a2a_cost(litellm_logging_obj=litellm_logging_obj) if model is None: - raise ValueError( - f"Model is None and does not exist in passed completion_response. Passed completion_response={completion_response}, model={model}" - ) + raise _model_is_none_error(f"{completion_response}", litellm_logging_obj) if custom_llm_provider is None: try: model, custom_llm_provider, _, _ = litellm.get_llm_provider( @@ -2276,9 +2281,7 @@ def completion_cost( return get_replicate_completion_pricing(completion_response, response_time_ms) if model is None: - raise ValueError( - f"Model is None and does not exist in passed completion_response. Passed completion_response={completion_response}, model={model}" - ) + raise _model_is_none_error(f"{completion_response}", litellm_logging_obj) if custom_llm_provider is not None and custom_llm_provider == "vertex_ai": # Calculate the prompt characters + response characters diff --git a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py index f40a166e67d..80341cb282d 100644 --- a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py +++ b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py @@ -35,6 +35,7 @@ from litellm.types.vector_stores import ( VectorStoreSearchResponse, VectorStoreSearchResult, ) +from litellm.vector_stores.search_query import store_search_params if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -327,7 +328,7 @@ class VectorStorePreCallHook(CustomLogger): "query": query, "custom_llm_provider": custom_llm_provider, "metadata": request_metadata, - **litellm_params_for_vector_store, + **store_search_params(litellm_params_for_vector_store), }, ) except Exception as search_error: diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index 3760d7f3059..530339f1db2 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -34,6 +34,7 @@ from litellm.llms.vertex_ai.common_utils import ( from litellm.secret_managers.main import str_to_bool from litellm.types.router import BaselineRouteStamp from litellm.types.utils import StandardCallbackDynamicParams +from litellm.types.vector_stores import RETRIEVAL_QUERY_PARAMS if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import ( @@ -86,8 +87,8 @@ def _redacted_entries_for_opted_out_logger(model_call_details: Mapping[str, obje "messages": [{"role": "user", "content": REDACTED_BY_LITELLM}], "prompt": "", "input": "", - **({"instructions": REDACTED_BY_LITELLM} if model_call_details.get("instructions") is not None else {}), - **({"optional_params": _with_redacted_instructions(optional_params)} if optional_params is not None else {}), + **_request_field_redactions(model_call_details), + **({"optional_params": _with_redacted_request_fields(optional_params)} if optional_params is not None else {}), **({"litellm_params": redacted_litellm_params(litellm_params)} if litellm_params is not None else {}), **( {"standard_logging_object": _redact_standard_logging_object(standard_logging_object)} @@ -464,8 +465,24 @@ def _redact_responses_api_output_dict(output_items: Sequence[object], redacted_s return [_redacted_responses_api_output_item(item, redacted_str) for item in output_items] -def _with_redacted_instructions(params: Mapping[str, object]) -> Mapping[str, object]: - return params if params.get("instructions") is None else {**params, "instructions": REDACTED_BY_LITELLM} +_REDACTED_REQUEST_FIELDS: Final = ("instructions", "query_image", *RETRIEVAL_QUERY_PARAMS) +_NESTED_REQUEST_FIELDS: Final = ("extra_body", "additional_args", "complete_input_dict") + + +def _request_field_redactions(params: Mapping[str, object]) -> Mapping[str, object]: + return { + **{field: REDACTED_BY_LITELLM for field in _REDACTED_REQUEST_FIELDS if params.get(field) is not None}, + **{ + field: _with_redacted_request_fields(nested) + for field in _NESTED_REQUEST_FIELDS + if (nested := _as_mapping(params.get(field))) is not None + }, + } + + +def _with_redacted_request_fields(params: Mapping[str, object]) -> Mapping[str, object]: + redactions: Final = _request_field_redactions(params) + return {**params, **redactions} if redactions else params def _redacted_responses_api_response( @@ -478,7 +495,7 @@ def _redacted_responses_api_response( if isinstance(output, list) else output, } - return _with_redacted_instructions(redacted) if redact_instructions else redacted + return _with_redacted_request_fields(redacted) if redact_instructions else redacted def redacted_standard_logging_payload(payload: Mapping[str, object]) -> Mapping[str, object]: @@ -502,7 +519,7 @@ def _redact_standard_logging_object( standard_logging_object["messages"] = [{"role": "user", "content": redacted_str}] model_parameters: Final = _as_mapping(standard_logging_object.get("model_parameters")) if model_parameters is not None and redact_instructions: - standard_logging_object["model_parameters"] = _with_redacted_instructions(model_parameters) + standard_logging_object["model_parameters"] = _with_redacted_request_fields(model_parameters) response: Final = standard_logging_object.get("response") if response is not None: @@ -599,7 +616,7 @@ def redacted_litellm_params(params: Mapping[str, object], *, redact_request_body """ request: Final = _as_mapping(params.get("proxy_server_request")) return { - **(_with_redacted_instructions(params) if redact_request_body else params), + **(_with_redacted_request_fields(params) if redact_request_body else params), **{ key: _redacted_baseline_metadata(value) for key, value in params.items() @@ -613,7 +630,7 @@ def redacted_litellm_params(params: Mapping[str, object], *, redact_request_body } -def perform_redaction(model_call_details: dict, result, redact_streaming_responses: bool = True): +def perform_redaction(model_call_details: dict[str, object], result, redact_streaming_responses: bool = True): """ Performs the actual redaction on the logging object and result. @@ -638,8 +655,8 @@ def perform_redaction(model_call_details: dict, result, redact_streaming_respons if redact_streaming_responses: model_call_details.update( { - **({"instructions": REDACTED_BY_LITELLM} if model_call_details.get("instructions") is not None else {}), - **({"optional_params": _with_redacted_instructions(optional_params)} if optional_params else {}), + **_request_field_redactions(model_call_details), + **({"optional_params": _with_redacted_request_fields(optional_params)} if optional_params else {}), } ) model_call_details.pop(SERVED_OUTPUT_TEXTS_KEY, None) diff --git a/litellm/llms/bedrock/vector_stores/transformation.py b/litellm/llms/bedrock/vector_stores/transformation.py index afcbe61cd75..b4cac70fee2 100644 --- a/litellm/llms/bedrock/vector_stores/transformation.py +++ b/litellm/llms/bedrock/vector_stores/transformation.py @@ -1,20 +1,25 @@ +import re from collections.abc import Mapping from copy import deepcopy +from itertools import product +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol, cast, runtime_checkable from urllib.parse import urlparse import httpx from pydantic import TypeAdapter -from typing_extensions import TypedDict +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_logger -from litellm.litellm_core_utils.core_helpers import decoded_json_object, keyword_arguments +from litellm.exceptions import BadRequestError +from litellm.litellm_core_utils.core_helpers import decoded_json_object, drops_unsupported_params, keyword_arguments from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.bedrock.common_utils import BedrockError from litellm.types.integrations.rag.bedrock_knowledgebase import ( BedrockKBContent, + BedrockKBImageFormat, BedrockKBResponse, BedrockKBRetrievalConfiguration, BedrockKBRetrievalQuery, @@ -22,11 +27,14 @@ from litellm.types.integrations.rag.bedrock_knowledgebase import ( ) from litellm.types.router import GenericLiteLLMParams from litellm.types.vector_stores import ( + RETRIEVAL_QUERY_PARAMS, VECTOR_STORE_OPENAI_PARAMS, BaseVectorStoreAuthCredentials, VectorStoreIndexEndpoints, VectorStoreResultContent, + VectorStoreSearchImageMediaType, VectorStoreSearchOptionalRequestParams, + VectorStoreSearchQueryImage, VectorStoreSearchResponse, VectorStoreSearchResult, ) @@ -70,6 +78,64 @@ class _KBResponseFields(TypedDict, total=False, extra_items=object): pass +_BEDROCK_IMAGE_FORMATS: Final[Mapping[VectorStoreSearchImageMediaType, BedrockKBImageFormat]] = MappingProxyType( + {"image/jpeg": "jpeg", "image/png": "png", "image/gif": "gif", "image/webp": "webp"} +) +_BASE64_DATA_URI: Final = re.compile(r"data:(?P[^;,]*)[^,]*;base64,(?P.*)", re.DOTALL) +_SOURCE_FILE_MIME_TYPE_KEY: Final = "x-amz-bedrock-kb-source-file-mime-type" +_UNKNOWN_MEDIA_TYPE: Final = "application/octet-stream" +_RETRIEVAL_QUERY_UNSUPPORTED: Final = ( + "bedrock vector stores do not take retrievalQuery. Send query, or query_image to search by image. To drop it, " + "set `litellm.drop_params=True` or for proxy:\n\n`litellm_settings:\n drop_params: true`\n" +) + + +class _QueryImageParams(TypedDict, total=False): + query_image: ReadOnly[VectorStoreSearchQueryImage | None] + + +def _retrieval_query(query: str, params: _QueryImageParams) -> BedrockKBRetrievalQuery: + query_image: Final = params.get("query_image") + if query_image is None: + return BedrockKBRetrievalQuery(text=query) + return BedrockKBRetrievalQuery( + type="IMAGE", + image={"format": _BEDROCK_IMAGE_FORMATS[query_image["media_type"]], "inlineContent": query_image["data"]}, + ) + + +def _image_content(byte_content: str, metadata: Mapping[str, object]) -> VectorStoreResultContent: + data_uri: Final = _BASE64_DATA_URI.fullmatch(byte_content) + declared_media_type: Final = data_uri.group("media_type") if data_uri is not None else "" + source_file_media_type: Final = metadata.get(_SOURCE_FILE_MIME_TYPE_KEY) + return VectorStoreResultContent( + type="image", + source={ + "type": "base64", + "media_type": ( + source_file_media_type + if isinstance(source_file_media_type, str) and source_file_media_type.startswith("image/") + else declared_media_type or _UNKNOWN_MEDIA_TYPE + ), + "data": data_uri.group("data") if data_uri is not None else byte_content, + }, + ) + + +def _result_content( + content: BedrockKBContent | None, metadata: Mapping[str, object], include_images: bool +) -> VectorStoreResultContent | None: + if content is None: + return None + text: Final = content.get("text") + if text is not None: + return VectorStoreResultContent(text=text, type="text") + byte_content: Final = content.get("byteContent") + if not include_images or content.get("type") != "IMAGE" or not byte_content: + return None + return _image_content(byte_content, metadata) + + class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): """Vector store configuration for AWS Bedrock Knowledge Bases.""" @@ -95,7 +161,7 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): } def get_supported_openai_params(self, model: str) -> list[VECTOR_STORE_OPENAI_PARAMS]: - return ["filters", "max_num_results", "ranking_options"] + return ["filters", "max_num_results", "ranking_options", "query_image"] def _map_operator_to_aws(self, operator: str) -> str: """ @@ -250,8 +316,14 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): encoded_vector_store_id: Final = encode_url_path_segment(vector_store_id, field_name="vector_store_id") url: Final = join_api_base(api_base, f"{encoded_vector_store_id}/retrieve") + caller_params: Final = decoded_json_object(litellm_params) + caller_sources: Final = self._caller_sources(extra_body=extra_body, litellm_params=caller_params) + if self._caller_value(caller_sources, RETRIEVAL_QUERY_PARAMS) is not None and not drops_unsupported_params( + caller_params.get("drop_params") + ): + raise BadRequestError(message=_RETRIEVAL_QUERY_UNSUPPORTED, model="", llm_provider="bedrock") request_body: Final[dict[str, object]] = { - "retrievalQuery": BedrockKBRetrievalQuery(text=query), + "retrievalQuery": _retrieval_query(query, vector_store_search_optional_params), } retrieval_config: Final[object] = ( @@ -283,7 +355,7 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): ] = filters if retrieval_config: request_body["retrievalConfiguration"] = cast(BedrockKBRetrievalConfiguration, retrieval_config) - user_context: Final = self._user_context(extra_body=extra_body, litellm_params=litellm_params) + user_context: Final = self._user_context(caller_sources) if user_context is not None: request_body["userContext"] = user_context @@ -291,19 +363,18 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): return url, request_body @staticmethod - def _user_context( + def _caller_sources( extra_body: Mapping[str, object] | None, litellm_params: Mapping[str, object] - ) -> BedrockKBUserContext | None: - sources: Final = tuple(source for source in (extra_body, litellm_params) if isinstance(source, Mapping)) - found: Final = next( - ( - source[key] - for source in sources - for key in ("userContext", "user_context") - if source.get(key) is not None - ), - None, - ) + ) -> tuple[Mapping[str, object], ...]: + return tuple(source for source in (extra_body, litellm_params) if isinstance(source, Mapping)) + + @staticmethod + def _caller_value(sources: tuple[Mapping[str, object], ...], keys: tuple[str, ...]) -> object: + return next((source[key] for source, key in product(sources, keys) if source.get(key) is not None), None) + + @staticmethod + def _user_context(sources: tuple[Mapping[str, object], ...]) -> BedrockKBUserContext | None: + found: Final = BedrockVectorStoreConfig._caller_value(sources, ("userContext", "user_context")) return None if found is None else cast(BedrockKBUserContext, found) def sign_request( @@ -375,15 +446,14 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): **keyword_arguments(_DECODED_JSON.validate_python(response.json()), BedrockKBResponse) ) ) + include_images: Final = litellm_logging_obj.model_call_details.get("query_image") is not None results: Final[list[VectorStoreSearchResult]] = [] for item in response_data.get("retrievalResults", []) or []: - content: BedrockKBContent | None = item.get("content") - text = content.get("text") if content else None - if text is None: + metadata = item.get("metadata", {}) or {} + result_content = _result_content(item.get("content"), metadata, include_images) + if result_content is None: continue - # Extract metadata and use helper functions - metadata = item.get("metadata", {}) or {} file_id = self._get_file_id_from_metadata(metadata) filename = self._get_filename_from_metadata(metadata) attributes = self._get_attributes_from_metadata(metadata) @@ -391,7 +461,7 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): results.append( VectorStoreSearchResult( score=item.get("score"), - content=[VectorStoreResultContent(text=text, type="text")], + content=[result_content], file_id=file_id, filename=filename, attributes=attributes, diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index b1ebd042f3d..5810a2f71ab 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -280,7 +280,10 @@ REQUIRED_ONE_OF_BODY_PARAMS_BY_ROUTE: Final[Mapping[str, tuple[str, str]]] = Map ) BODY_PARAM_ALTERNATIVES_BY_ROUTE: Final[Mapping[str, Mapping[str, tuple[str, ...]]]] = MappingProxyType( - {"acreate_interaction": MappingProxyType({"input": ("continuation_token",)})} + { + "acreate_interaction": MappingProxyType({"input": ("continuation_token",)}), + "avector_store_search": MappingProxyType({"query": ("query_image",)}), + } ) JSON_OBJECT_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object]) diff --git a/litellm/proxy/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py index a6f9ba2ef14..0065f8e4002 100644 --- a/litellm/proxy/vector_store_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_endpoints/endpoints.py @@ -36,6 +36,7 @@ from litellm.proxy.vector_store_endpoints.utils import ( from litellm.repositories.table_repositories import ManagedVectorStoreIndexRepository from litellm.types.utils import LlmProviders from litellm.types.vector_stores import IndexCreateRequest, IndexListResponse +from litellm.vector_stores.search_query import store_search_params, with_caller_query_params from litellm.vector_stores.utils import VectorStoreWriteBody, callers_metadata_snapshot from litellm.vector_stores.vector_store_registry import VectorStoreIndexRegistry @@ -113,8 +114,8 @@ def build_request_data_from_managed_vector_store( if key in vector_store } ) - litellm_params: Final = vector_store.get("litellm_params") or MappingProxyType({}) - return MappingProxyType({**top_level, **litellm_params}) + litellm_params: Final[Mapping[str, object]] = vector_store.get("litellm_params") or MappingProxyType({}) + return MappingProxyType({**top_level, **store_search_params(litellm_params)}) async def _update_request_data_with_litellm_managed_vector_store_registry( @@ -187,8 +188,13 @@ async def vector_store_search( caller_body: Final = JSON_OBJECT_ADAPTER.validate_python(data) # Check for legacy vector store registry (non-managed vector stores) - data = await _update_request_data_with_litellm_managed_vector_store_registry( - data=data, vector_store_id=vector_store_id, user_api_key_dict=user_api_key_dict + data = dict( # rebind-ok: pre-existing rebinding of the request body + with_caller_query_params( + await _update_request_data_with_litellm_managed_vector_store_registry( + data=data, vector_store_id=vector_store_id, user_api_key_dict=user_api_key_dict + ), + caller_body.get("extra_body"), + ) ) reject_caller_opensearch_index_name( caller_body=caller_body, diff --git a/litellm/router_rust/invoke/endpoints.py b/litellm/router_rust/invoke/endpoints.py index fd3d055f707..c170aaab6e6 100644 --- a/litellm/router_rust/invoke/endpoints.py +++ b/litellm/router_rust/invoke/endpoints.py @@ -53,6 +53,7 @@ from litellm.router_strategy.complexity_router.context_compaction import ( surface_for_call, ) from litellm.router_utils.endpoint_eligibility import endpoint_for_call, endpoint_selection +from litellm.vector_stores.search_query import store_search_params if TYPE_CHECKING: from litellm.caching.dual_cache import DualCache @@ -854,6 +855,16 @@ async def compacted_attempt( return await compact_to_fit(ctx, deployment, call.kwargs, _compaction_surface(core, plan)) +_VECTOR_STORE_SEARCH_OPS: Final = frozenset({"vector_stores.main.asearch", "vector_stores.main.search"}) + + +def _deployment_params( + plan: Plan, deployment: Mapping[str, object] +) -> dict[str, object]: # mutable-ok: the attempt's own copy of the deployment params + params: Final = as_dict(deployment.get("litellm_params")) + return as_dict(store_search_params(params)) if plan.host_op in _VECTOR_STORE_SEARCH_OPS else params + + def attempt_call( core: EndpointCore, plan: Plan, @@ -869,7 +880,7 @@ def attempt_call( shaping: Final = core.shaping kwargs: Final = without(request, *ROUTING_KEYS) named: Final = {name: kwargs.pop(name) for name in plan.named if name in kwargs} - data: Final = as_dict(deployment.get("litellm_params")) + data: Final = _deployment_params(plan, deployment) model_name: Final = str(data.get("model")) cache: Final = core.routing_cache caching: Final = shaping.cache_responses diff --git a/litellm/types/integrations/rag/bedrock_knowledgebase.py b/litellm/types/integrations/rag/bedrock_knowledgebase.py index 7288eca03b4..f472cdc078f 100644 --- a/litellm/types/integrations/rag/bedrock_knowledgebase.py +++ b/litellm/types/integrations/rag/bedrock_knowledgebase.py @@ -1,4 +1,4 @@ -from typing import Literal +from typing import Literal, TypeAlias from typing_extensions import ReadOnly, TypedDict @@ -114,10 +114,20 @@ class BedrockKBRetrievalConfiguration(TypedDict, total=False): vectorSearchConfiguration: BedrockKBVectorSearchConfiguration | None +BedrockKBImageFormat: TypeAlias = Literal["png", "jpeg", "gif", "webp"] + + +class BedrockKBInputImage(TypedDict): + format: ReadOnly[BedrockKBImageFormat] + inlineContent: ReadOnly[str] + + class BedrockKBRetrievalQuery(TypedDict, total=False): """Query structure for retrieval.""" text: str | None + type: ReadOnly[Literal["TEXT", "IMAGE"]] + image: ReadOnly[BedrockKBInputImage] class BedrockKBGuardrailConfiguration(TypedDict, total=False): diff --git a/litellm/types/vector_stores.py b/litellm/types/vector_stores.py index 2259339e479..b154d161e13 100644 --- a/litellm/types/vector_stores.py +++ b/litellm/types/vector_stores.py @@ -2,12 +2,14 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass from datetime import datetime from enum import Enum -from typing import Any, Literal +from typing import Any, Final, Literal, TypeAlias from typing_extensions import ReadOnly, TypedDict from litellm.types.llms.base import LiteLLMBaseModel +RETRIEVAL_QUERY_PARAMS: Final = ("retrievalQuery", "retrieval_query") + class SupportedVectorStoreIntegrations(str, Enum): """Supported vector store integrations.""" @@ -74,11 +76,29 @@ class VectorStoreInfoRequest(LiteLLMBaseModel): vector_store_id: str +VectorStoreSearchImageMediaType: TypeAlias = Literal["image/jpeg", "image/png", "image/gif", "image/webp"] + + +class VectorStoreSearchQueryImage(TypedDict): + """Base64 image to search a multimodal vector store with, in place of a text query""" + + type: ReadOnly[Literal["base64"]] + media_type: ReadOnly[VectorStoreSearchImageMediaType] + data: ReadOnly[str] + + +class VectorStoreResultImageSource(TypedDict): + type: ReadOnly[Literal["base64"]] + media_type: ReadOnly[str] + data: ReadOnly[str] + + class VectorStoreResultContent(TypedDict, total=False): """Content of a vector store result""" text: str | None type: str | None + source: ReadOnly[VectorStoreResultImageSource] class VectorStoreSearchResult(TypedDict, total=False): @@ -117,6 +137,7 @@ class VectorStoreSearchOptionalRequestParams(TypedDict, total=False): max_num_results: int | None ranking_options: dict | None rewrite_query: bool | None + query_image: ReadOnly[VectorStoreSearchQueryImage | None] class VectorStoreSearchRequest(VectorStoreSearchOptionalRequestParams, total=False): @@ -320,6 +341,7 @@ VECTOR_STORE_OPENAI_PARAMS = Literal[ "max_num_results", "ranking_options", "rewrite_query", + "query_image", ] @@ -330,6 +352,7 @@ class VectorStoreToolParams: filters: dict | None = None max_num_results: int | None = None ranking_options: dict | None = None + rewrite_query: bool | None = None def to_dict(self) -> dict: """Convert to dict, excluding None values""" @@ -339,6 +362,7 @@ class VectorStoreToolParams: "filters": self.filters, "max_num_results": self.max_num_results, "ranking_options": self.ranking_options, + "rewrite_query": self.rewrite_query or None, }.items() if v is not None } diff --git a/litellm/vector_stores/main.py b/litellm/vector_stores/main.py index 9f53accfa0b..88bc477249b 100644 --- a/litellm/vector_stores/main.py +++ b/litellm/vector_stores/main.py @@ -13,6 +13,7 @@ import httpx import litellm from litellm.constants import request_timeout +from litellm.litellm_core_utils.core_helpers import as_str_mapping, drops_unsupported_params from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.vector_store.transformation import ( @@ -28,10 +29,12 @@ from litellm.types.vector_stores import ( VectorStoreFileCounts, VectorStoreResultContent, VectorStoreSearchOptionalRequestParams, + VectorStoreSearchQueryImage, VectorStoreSearchResponse, VectorStoreSearchResult, ) from litellm.utils import ProviderConfigManager, client, load_credentials_from_list +from litellm.vector_stores.search_query import searched_query, sent_query_image from litellm.vector_stores.utils import VectorStoreRequestUtils if TYPE_CHECKING: @@ -289,7 +292,7 @@ def create( @client async def asearch( vector_store_id: str, - query: str | builtins.list[str], + query: str | builtins.list[str] | None = None, filters: dict | None = None, max_num_results: int | None = None, ranking_options: dict | None = None, @@ -303,6 +306,8 @@ async def asearch( # LiteLLM specific params, custom_llm_provider: str | None = None, router: "Router | None" = None, + *, + query_image: VectorStoreSearchQueryImage | None = None, **kwargs, ) -> VectorStoreSearchResponse: """ @@ -329,6 +334,7 @@ async def asearch( max_num_results=max_num_results, ranking_options=ranking_options, rewrite_query=rewrite_query, + query_image=query_image, extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, @@ -362,7 +368,7 @@ async def asearch( @client def search( vector_store_id: str, - query: str | builtins.list[str], + query: str | builtins.list[str] | None = None, filters: dict | None = None, max_num_results: int | None = None, ranking_options: dict | None = None, @@ -376,6 +382,8 @@ def search( # LiteLLM specific params, custom_llm_provider: str | None = None, router: "Router | None" = None, + *, + query_image: VectorStoreSearchQueryImage | None = None, **kwargs, ) -> VectorStoreSearchResponse | Coroutine[object, object, VectorStoreSearchResponse]: """ @@ -383,11 +391,12 @@ def search( Args: vector_store_id: The ID of the vector store to search. - query: A query string or array for the search. + query: A query string or array for the search. Optional when query_image is sent. filters: Optional filter to apply based on file attributes. max_num_results: Maximum number of results to return (1-50, default 10). ranking_options: Optional ranking options for search. rewrite_query: Whether to rewrite the natural language query for vector search. + query_image: A base64 image to search with in place of query, on providers that support it. Returns: VectorStoreSearchResponse containing the search results. @@ -444,21 +453,31 @@ def search( local_vars.update(kwargs) + search_query: Final = searched_query( + query=query, + query_image=sent_query_image(query_image, extra_body), + custom_llm_provider=custom_llm_provider, + supports_query_image="query_image" in vector_store_provider_config.get_supported_openai_params(model=""), + drop_params=drops_unsupported_params(local_vars.get("drop_params")), + ) + # Get VectorStoreSearchOptionalRequestParams with only valid parameters vector_store_search_optional_params: Final[VectorStoreSearchOptionalRequestParams] = ( VectorStoreRequestUtils.get_requested_vector_store_search_optional_param( - local_vars, + {**local_vars, "query_image": search_query.image}, vector_store_provider_config=vector_store_provider_config, ) ) + sent_extra_body: Final = as_str_mapping(extra_body) + # Pre Call logging litellm_logging_obj.update_from_kwargs( kwargs=kwargs, model=api_type, optional_params={ "vector_store_id": vector_store_id, - "query": query, + "query": search_query.query, **vector_store_search_optional_params, }, litellm_params={ @@ -471,7 +490,7 @@ def search( response: Final = base_llm_http_handler.vector_store_search_handler( vector_store_id=vector_store_id, - query=query, + query=search_query.query, vector_store_search_optional_params=vector_store_search_optional_params, vector_store_provider_config=vector_store_provider_config, custom_llm_provider=custom_llm_provider, @@ -479,7 +498,9 @@ def search( logging_obj=litellm_logging_obj, embedding_executor=embedding_executor, extra_headers=extra_headers, - extra_body=extra_body, + extra_body=extra_body + if sent_extra_body is None + else dict(_items_without_key(sent_extra_body, "query_image")), timeout=timeout or request_timeout, _is_async=_is_async, client=kwargs.get("client"), diff --git a/litellm/vector_stores/search_query.py b/litellm/vector_stores/search_query.py new file mode 100644 index 00000000000..b4ec0b6707b --- /dev/null +++ b/litellm/vector_stores/search_query.py @@ -0,0 +1,205 @@ +import base64 +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final, Generic, Literal, TypeAlias, TypeVar + +from pydantic import TypeAdapter, ValidationError +from pydantic_core import ErrorDetails +from typing_extensions import assert_never + +import litellm +from litellm.litellm_core_utils.core_helpers import as_object_list, as_str_mapping +from litellm.types.vector_stores import RETRIEVAL_QUERY_PARAMS, VectorStoreSearchQueryImage + +_QUERY_IMAGE: Final = TypeAdapter(VectorStoreSearchQueryImage) +_Query: Final = TypeVar("_Query", bound=str | Sequence[str]) +_CALLER_QUERY_PARAMS: Final = ("query_image", *RETRIEVAL_QUERY_PARAMS) +_BASE64_LINE_BREAKS: Final = str.maketrans("", "", "\r\n") + + +@dataclass(frozen=True, slots=True) +class TextSearchQuery(Generic[_Query]): + query: _Query + image: None = None + + +@dataclass(frozen=True, slots=True) +class ImageSearchQuery: + image: VectorStoreSearchQueryImage + query: Literal[""] = "" + + +@dataclass(frozen=True, slots=True) +class QueryAndImageBothSent: + pass + + +@dataclass(frozen=True, slots=True) +class NoQuerySent: + pass + + +@dataclass(frozen=True, slots=True) +class InvalidQueryImage: + problems: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class QueryImageUnsupported: + custom_llm_provider: str + + +@dataclass(frozen=True, slots=True) +class QueryImageDroppedWithoutQuery: + custom_llm_provider: str + + +SearchQuery: TypeAlias = TextSearchQuery[_Query] | ImageSearchQuery +SearchQueryError: TypeAlias = ( + QueryAndImageBothSent | NoQuerySent | InvalidQueryImage | QueryImageUnsupported | QueryImageDroppedWithoutQuery +) + + +def _has_text(query: object) -> bool: + if isinstance(query, str): + return bool(query.strip()) + items: Final = as_object_list(query) + if items is None: + return True + return any(not isinstance(item, str) or item.strip() for item in items) + + +def _problem(detail: ErrorDetails) -> str: + location: Final = ".".join(("query_image", *(str(part) for part in detail["loc"]))) + return f"{location}: {detail['msg']}" + + +def _is_base64(data: str) -> bool: + try: + base64.b64decode(data.translate(_BASE64_LINE_BREAKS), validate=True) + except ValueError: + return False + return True + + +def _validated_query_image(query_image: object) -> ImageSearchQuery | InvalidQueryImage: + try: + image: Final = _QUERY_IMAGE.validate_python(query_image) + except ValidationError as error: + return InvalidQueryImage( + problems=tuple(_problem(detail) for detail in error.errors(include_url=False, include_input=False)) + ) + if not _is_base64(image["data"]): + return InvalidQueryImage(problems=("query_image.data: not valid base64",)) + return ImageSearchQuery(image=image) + + +def sent_query_image(query_image: object, extra_body: object) -> object: + if query_image is not None: + return query_image + body: Final = as_str_mapping(extra_body) + return None if body is None else body.get("query_image") + + +def _without_retrieval_query(params: Mapping[str, object]) -> Mapping[str, object]: + return {key: value for key, value in params.items() if key not in RETRIEVAL_QUERY_PARAMS} + + +def store_search_params(litellm_params: Mapping[str, object]) -> Mapping[str, object]: + extra_body: Final = as_str_mapping(litellm_params.get("extra_body")) + return MappingProxyType( + { + **_without_retrieval_query(litellm_params), + **({} if extra_body is None else {"extra_body": _without_retrieval_query(extra_body)}), + } + ) + + +def with_caller_query_params(request: Mapping[str, object], caller_extra_body: object) -> Mapping[str, object]: + caller: Final = as_str_mapping(caller_extra_body) + store: Final = as_str_mapping(request.get("extra_body")) + if caller is None or store is None: + return request + query_params: Final = {key: caller[key] for key in _CALLER_QUERY_PARAMS if key in caller} + return {**request, "extra_body": {**store, **query_params}} + + +def resolve_search_query( + query: _Query | None, + query_image: object, + custom_llm_provider: str, + supports_query_image: bool, + drop_params: bool, +) -> SearchQuery[_Query] | SearchQueryError: + if query_image is None: + return NoQuerySent() if query is None else TextSearchQuery(query=query) + text_query: Final = query if query is not None and _has_text(query) else None + if not supports_query_image and not drop_params: + return QueryImageUnsupported(custom_llm_provider=custom_llm_provider) + if not supports_query_image: + return ( + QueryImageDroppedWithoutQuery(custom_llm_provider=custom_llm_provider) + if text_query is None + else TextSearchQuery(query=text_query) + ) + if text_query is not None: + return QueryAndImageBothSent() + return _validated_query_image(query_image) + + +def _error_message(error: SearchQueryError) -> str: + match error: + case QueryAndImageBothSent(): + return "Send either query or query_image, not both. Leave query empty when you send query_image." + case NoQuerySent(): + return "query or query_image is required." + case InvalidQueryImage(problems=problems): + return ( + f"Invalid query_image: {'; '.join(problems)}. Expected " + '{"type": "base64", "media_type": "image/jpeg" | "image/png" | "image/gif" | "image/webp", ' + '"data": ""}.' + ) + case QueryImageUnsupported(custom_llm_provider=provider): + return ( + f"{provider} vector stores do not support query_image. To drop it, set `litellm.drop_params=True` " + "or for proxy:\n\n`litellm_settings:\n drop_params: true`\n" + ) + case QueryImageDroppedWithoutQuery(custom_llm_provider=provider): + return ( + f"{provider} vector stores do not support query_image and drop_params dropped it, " + "which leaves no query to search with. Send query." + ) + case _: + assert_never(error) + + +def searched_query( + query: _Query | None, + query_image: object, + custom_llm_provider: str, + supports_query_image: bool, + drop_params: bool, +) -> SearchQuery[_Query]: + resolved: Final = resolve_search_query( + query=query, + query_image=query_image, + custom_llm_provider=custom_llm_provider, + supports_query_image=supports_query_image, + drop_params=drop_params, + ) + match resolved: + case TextSearchQuery() | ImageSearchQuery(): + return resolved + case ( + QueryAndImageBothSent() + | NoQuerySent() + | InvalidQueryImage() + | QueryImageUnsupported() + | QueryImageDroppedWithoutQuery() + ): + raise litellm.BadRequestError( + message=_error_message(resolved), model=None, llm_provider=custom_llm_provider + ) + case _: + assert_never(resolved) diff --git a/litellm/vector_stores/vector_store_registry.py b/litellm/vector_stores/vector_store_registry.py index 2adaf17ac44..31fc111edab 100644 --- a/litellm/vector_stores/vector_store_registry.py +++ b/litellm/vector_stores/vector_store_registry.py @@ -1,6 +1,7 @@ # litellm/proxy/vector_stores/vector_store_registry.py import json from collections.abc import Iterable, Mapping, Sequence +from dataclasses import fields, replace from datetime import datetime, timezone from itertools import chain from typing import ( @@ -8,20 +9,18 @@ from typing import ( Any, # noqa: TID251 # untyped non_default_params dict is the only source of the unknown key type Final, cast, # noqa: TID251 # untyped non_default_params dict is the only source of the unknown key type - get_args, ) from pydantic import ConfigDict, TypeAdapter, ValidationError from typing_extensions import TypedDict from litellm._logging import verbose_logger -from litellm.litellm_core_utils.core_helpers import remove_items_at_indices +from litellm.litellm_core_utils.core_helpers import decoded_json_object, remove_items_at_indices from litellm.repositories.table_repositories import ( ManagedVectorStoreIndexRepository, ManagedVectorStoresRepository, ) from litellm.types.vector_stores import ( - VECTOR_STORE_OPENAI_PARAMS, LiteLLM_ManagedVectorStore, LiteLLM_ManagedVectorStoreIndex, LiteLLM_ManagedVectorStoreListResponse, @@ -150,19 +149,15 @@ class VectorStoreRegistry: self.vector_store_ids_to_vector_store_map: dict[str, LiteLLM_ManagedVectorStore] = {} def _extract_tool_params(self, tool: dict) -> VectorStoreToolParams: - """ - Extract supported parameters from a tool definition. - - Dynamically extracts all parameters defined in VECTOR_STORE_OPENAI_PARAMS. - """ - # Get the list of supported param names from the Literal type - declared_params: Final[tuple[object, ...]] = get_args(VECTOR_STORE_OPENAI_PARAMS) - supported_params: Final = tuple(param for param in declared_params if isinstance(param, str)) - - # Extract only the params that exist in the tool - kwargs: Final = {param: tool.get(param) for param in supported_params if param in tool} - - return VectorStoreToolParams(**kwargs) + tool_fields: Final = decoded_json_object(tool) + return replace( + VectorStoreToolParams(), + **{ + field.name: tool_fields[field.name] + for field in fields(VectorStoreToolParams) + if field.name in tool_fields + }, + ) def get_vector_store_ids_to_run(self, non_default_params: dict, tools: object = None) -> list[str]: """ diff --git a/tests/unit/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py b/tests/unit/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py index 726ce2f7f1d..48d726aa4df 100644 --- a/tests/unit/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py +++ b/tests/unit/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py @@ -2,7 +2,7 @@ import logging from collections.abc import Callable, Iterable, Iterator, Mapping from dataclasses import dataclass, field from types import MappingProxyType -from typing import Literal, Protocol +from typing import Final, Literal, Protocol import pytest from fastapi import HTTPException @@ -157,6 +157,38 @@ async def _run_hook( ) +@pytest.mark.asyncio +@pytest.mark.parametrize("retrieval_query_param", ["retrievalQuery", "retrieval_query"]) +async def test_a_store_configured_with_a_raw_retrieval_query_still_runs_its_chat_search( + monkeypatch: pytest.MonkeyPatch, retrieval_query_param: str +) -> None: + monkeypatch.setattr( + litellm, + "vector_store_registry", + VectorStoreRegistry( + vector_stores=[ + LiteLLM_ManagedVectorStore( + vector_store_id="vs-raw", + custom_llm_provider="bedrock", + litellm_params={"aws_region_name": "us-west-2", retrieval_query_param: {"text": "from the store"}}, + ) + ], + ), + ) + router: Final = RecordingRouter() + + await _run_hook( + VectorStorePreCallHook(proxy_runtime=FakeProxyRuntime(router=router)), + ["vs-raw"], + FakeLoggingObj({}), + ) + + assert len(router.calls) == 1 + assert router.calls[0]["query"] == "what is litellm?" + assert router.calls[0]["aws_region_name"] == "us-west-2" + assert retrieval_query_param not in router.calls[0] + + @pytest.mark.asyncio async def test_hook_searches_through_the_injected_router_with_the_request_metadata( registry_with: RegisterStores, diff --git a/tests/unit/litellm_core_utils/test_redact_messages.py b/tests/unit/litellm_core_utils/test_redact_messages.py index 8e4bed02d7e..888eb9c9507 100644 --- a/tests/unit/litellm_core_utils/test_redact_messages.py +++ b/tests/unit/litellm_core_utils/test_redact_messages.py @@ -2011,6 +2011,69 @@ def test_request_body_redaction_replaces_instructions_only_when_sent() -> None: assert "instructions" not in chat_body +def _vector_store_search_details(query_image: Mapping[str, str]) -> dict[str, object]: + return { + "query_image": query_image, + "optional_params": {"query_image": query_image}, + "litellm_params": {"proxy_server_request": {"body": {"query_image": query_image}}}, + "standard_logging_object": {"model_parameters": {"query_image": query_image}, "messages": [], "response": {}}, + } + + +def test_redaction_removes_a_vector_store_query_image_from_every_logged_copy() -> None: + details: Final = _vector_store_search_details({"type": "base64", "media_type": "image/jpeg", "data": SECRET}) + + perform_redaction(details, result=None) + + assert details["query_image"] == REDACTED_BY_LITELLM + assert SECRET not in json.dumps(details, default=str) + + +def test_redaction_removes_a_query_image_sent_inside_extra_body() -> None: + image: Final = {"type": "base64", "media_type": "image/jpeg", "data": SECRET} + details: Final[dict[str, object]] = { + "optional_params": {"extra_body": {"query_image": image}}, + "litellm_params": {"proxy_server_request": {"body": {"extra_body": {"query_image": image}}}}, + } + + perform_redaction(details, result=None) + + assert SECRET not in json.dumps(details, default=str) + + +@pytest.mark.parametrize("retrieval_query_param", ["retrievalQuery", "retrieval_query"]) +def test_redaction_removes_a_raw_retrieval_query_sent_inside_extra_body(retrieval_query_param: str) -> None: + retrieval_query: Final = {"type": "IMAGE", "image": {"format": "jpeg", "inlineContent": SECRET}} + details: Final[dict[str, object]] = { + "litellm_params": {"proxy_server_request": {"body": {"extra_body": {retrieval_query_param: retrieval_query}}}}, + } + + perform_redaction(details, result=None) + + assert SECRET not in json.dumps(details, default=str) + + +def test_redaction_removes_the_image_from_the_provider_request_copy() -> None: + retrieval_query: Final = {"type": "IMAGE", "image": {"format": "jpeg", "inlineContent": SECRET}} + details: Final[dict[str, object]] = { + "additional_args": {"complete_input_dict": {"retrievalQuery": retrieval_query, "knowledgeBaseId": "kb"}}, + } + + perform_redaction(details, result=None) + + assert SECRET not in json.dumps(details, default=str) + assert details["additional_args"]["complete_input_dict"]["knowledgeBaseId"] == "kb" + + +def test_one_callback_opting_out_redacts_a_query_image_only_for_itself() -> None: + details: Final = _vector_store_search_details({"type": "base64", "media_type": "image/jpeg", "data": SECRET}) + + redacted: Final = redacted_model_call_details_for_custom_logger(details, TestCustomLogger(message_logging=False)) + + assert SECRET not in json.dumps(redacted, default=str) + assert SECRET in json.dumps(details, default=str) + + @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.asyncio async def test_redaction_responses_api_stream(): diff --git a/tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py b/tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py index e69fb34038f..e20dfaffcc6 100644 --- a/tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py +++ b/tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py @@ -6,6 +6,7 @@ import httpx import pytest import respx +import litellm from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.vector_stores.transformation import BedrockVectorStoreConfig @@ -378,6 +379,161 @@ async def test_bedrock_search_with_credentials_managed_registry( assert search_response["data"][0]["content"][0]["text"] == "Mock search result" +def _image_search_body(media_type: str, extra_body: dict[str, object] | None) -> dict[str, object]: + config: Final = BedrockVectorStoreConfig() + mock_log: Final = MagicMock() + mock_log.model_call_details = {} + _, body = config.transform_search_vector_store_request( + vector_store_id="kb123", + query="", + vector_store_search_optional_params={ + "max_num_results": 3, + "query_image": {"type": "base64", "media_type": media_type, "data": "aW1hZ2UtYnl0ZXM="}, + }, + api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases", + litellm_logging_obj=mock_log, + litellm_params={}, + extra_body=extra_body, + ) + return body + + +@pytest.mark.parametrize( + ("media_type", "bedrock_format"), + [("image/jpeg", "jpeg"), ("image/png", "png"), ("image/gif", "gif"), ("image/webp", "webp")], +) +def test_transform_search_request_sends_query_image_as_an_image_retrieval_query( + media_type: str, bedrock_format: str +) -> None: + body: Final = _image_search_body(media_type, extra_body=None) + + assert body["retrievalQuery"] == { + "type": "IMAGE", + "image": {"format": bedrock_format, "inlineContent": "aW1hZ2UtYnl0ZXM="}, + } + assert body["retrievalConfiguration"] == {"vectorSearchConfiguration": {"numberOfResults": 3}} + + +CALLER_RETRIEVAL_QUERY: Final = {"type": "IMAGE", "image": {"format": "png", "inlineContent": "Y2FsbGVy"}} + + +@pytest.mark.parametrize("key", ["retrievalQuery", "retrieval_query"]) +@pytest.mark.parametrize("source", ["extra_body", "litellm_params"]) +def test_transform_search_request_rejects_a_callers_retrieval_query(key: str, source: str) -> None: + caller_fields: Final = {key: CALLER_RETRIEVAL_QUERY} + + with pytest.raises(litellm.BadRequestError) as exc_info: + _search_body( + extra_body=caller_fields if source == "extra_body" else None, + litellm_params=caller_fields if source == "litellm_params" else {}, + ) + + assert exc_info.value.status_code == 400 + assert "query_image" in str(exc_info.value) + + +def test_transform_search_request_drops_a_callers_retrieval_query_when_drop_params_is_on() -> None: + body: Final = _search_body( + extra_body={"retrievalQuery": CALLER_RETRIEVAL_QUERY}, litellm_params={"drop_params": True} + ) + + assert body["retrievalQuery"] == {"text": "hello"} + + +QUERY_IMAGE: Final = {"type": "base64", "media_type": "image/jpeg", "data": "/9j/4AAQSkZJRg=="} +IMAGE_HIT_METADATA: Final = { + "x-amz-bedrock-kb-source-uri": "s3://kb-images/shapes/red-circle.jpg", + "x-amz-bedrock-kb-source-file-mime-type": "image/jpeg", + "x-amz-bedrock-kb-source-file-modality": "IMAGE", +} +TEXT_HIT_METADATA: Final = {"x-amz-bedrock-kb-source-uri": "s3://kb-docs/shapes.txt"} +MIXED_RETRIEVAL_RESULTS: Final = [ + { + "content": {"type": "IMAGE", "byteContent": "data:image/png;base64,/9j/4AAQSkZJRg=="}, + "metadata": IMAGE_HIT_METADATA, + "location": {"type": "S3", "s3Location": {"uri": "s3://kb-images/shapes/red-circle.jpg"}}, + "score": 0.998, + }, + {"content": {"type": "TEXT", "text": "a red circle"}, "metadata": TEXT_HIT_METADATA, "score": 0.5}, + {"content": {"type": "ROW", "row": []}, "metadata": {}, "score": 0.4}, + {"content": {"type": "AUDIO", "byteContent": "data:audio/mpeg;base64,SUQz"}, "metadata": {}, "score": 0.3}, +] +TEXT_HIT: Final = { + "score": 0.5, + "content": [{"text": "a red circle", "type": "text"}], + "file_id": "s3://kb-docs/shapes.txt", + "filename": "shapes.txt", + "attributes": TEXT_HIT_METADATA, +} + + +def _search_response( + retrieval_results: list[dict[str, object]], query: str = "", query_image: dict[str, str] | None = QUERY_IMAGE +) -> dict[str, object]: + config: Final = BedrockVectorStoreConfig() + mock_log: Final = MagicMock() + mock_log.model_call_details = {"query": query, "query_image": query_image} + return config.transform_search_vector_store_response( + response=httpx.Response(200, json={"retrievalResults": retrieval_results}), + litellm_logging_obj=mock_log, + ) + + +def test_transform_search_response_to_an_image_search_returns_image_hits_next_to_text_hits() -> None: + response: Final = _search_response(MIXED_RETRIEVAL_RESULTS) + + assert response["search_query"] == "" + assert response["data"] == [ + { + "score": 0.998, + "content": [ + { + "type": "image", + "source": {"type": "base64", "media_type": "image/jpeg", "data": "/9j/4AAQSkZJRg=="}, + } + ], + "file_id": "s3://kb-images/shapes/red-circle.jpg", + "filename": "red-circle.jpg", + "attributes": IMAGE_HIT_METADATA, + }, + TEXT_HIT, + ] + + +def test_transform_search_response_to_a_text_search_returns_only_text_hits() -> None: + response: Final = _search_response(MIXED_RETRIEVAL_RESULTS, query="red circle", query_image=None) + + assert response["search_query"] == "red circle" + assert response["data"] == [TEXT_HIT] + + +@pytest.mark.parametrize( + ("byte_content", "metadata", "media_type", "data"), + [ + ("data:image/webp;base64,UklGRg==", {}, "image/webp", "UklGRg=="), + ("data:image/png;base64,iVBORw0K", {"x-amz-bedrock-kb-source-file-mime-type": ""}, "image/png", "iVBORw0K"), + ( + "data:image/png;base64,iVBORw0K", + {"x-amz-bedrock-kb-source-file-mime-type": "application/pdf"}, + "image/png", + "iVBORw0K", + ), + ("/9j/4AAQ", {"x-amz-bedrock-kb-source-file-mime-type": "image/jpeg"}, "image/jpeg", "/9j/4AAQ"), + ("/9j/4AAQ", {}, "application/octet-stream", "/9j/4AAQ"), + ], +) +def test_transform_search_response_image_hit_media_type_falls_back_to_the_data_uri( + byte_content: str, metadata: dict[str, object], media_type: str, data: str +) -> None: + response: Final = _search_response( + [{"content": {"type": "IMAGE", "byteContent": byte_content}, "metadata": metadata, "score": 0.9}] + ) + + assert response["data"][0]["content"] == [ + {"type": "image", "source": {"type": "base64", "media_type": media_type, "data": data}} + ] + + class _SearchLogging: def __init__(self, query: str) -> None: self.model_call_details: Final = {"query": query} diff --git a/tests/unit/llms/mongodb/vector_stores/test_mongodb_transformation.py b/tests/unit/llms/mongodb/vector_stores/test_mongodb_transformation.py index 9de473fb1f0..d7d9b21ed20 100644 --- a/tests/unit/llms/mongodb/vector_stores/test_mongodb_transformation.py +++ b/tests/unit/llms/mongodb/vector_stores/test_mongodb_transformation.py @@ -220,3 +220,40 @@ async def test_public_sdk_preserves_http_errors_response_and_timeout( else: assert await search() == RESULT executor.call.assert_called_once_with("embedding-alias", "travel policy", {}) + + +QUERY_IMAGE: Final = {"type": "base64", "media_type": "image/jpeg", "data": "/9j/4AAQSkZJRg=="} + + +def _search_with_extra_body(extra_body: dict[str, object], **search_kwargs: object) -> VectorStoreSearchResponse: + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json=RESULT) + + with httpx.Client(transport=httpx.MockTransport(respond)) as transport: + return litellm.vector_stores.search( + **BASE_PARAMS, + vector_store_id="policy_index", + query="travel policy", + custom_llm_provider="mongodb", + extra_body=extra_body, + _direct_vector_store_embedding_executor=RecordingEmbeddingExecutor(), + client=HTTPHandler(client=transport), + **search_kwargs, + ) + + +@pytest.mark.parametrize("drop_params_source", ["request", "global"]) +def test_search_with_drop_params_drops_an_extra_body_query_image( + drop_params_source: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "drop_params", drop_params_source == "global") + drop_kwargs: Final = {"drop_params": True} if drop_params_source == "request" else {} + + assert _search_with_extra_body({"query_image": QUERY_IMAGE}, **drop_kwargs) == RESULT + + +def test_search_with_drop_params_still_rejects_other_extra_body_overrides(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "drop_params", True) + + with pytest.raises(litellm.BadRequestError, match="does not support extra_body overrides"): + _search_with_extra_body({"query_image": QUERY_IMAGE, "other": "value"}) diff --git a/tests/unit/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index 75a17d49d3a..fa53f422f36 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -1899,6 +1899,10 @@ def test_raise_if_required_body_param_missing_names_each_missing_param(route_typ ("aspeech", {"model": "gpt-4o-mini-tts", "input": "hi", "voice": "alloy"}), ("amoderation", {"model": "omni-moderation-latest", "input": ""}), ("asearch", {"model": "perplexity-search", "query": "litellm"}), + ( + "avector_store_search", + {"vector_store_id": "kb1", "query_image": {"type": "base64", "media_type": "image/png", "data": "iVBO"}}, + ), ( "acreate_batch", {"input_file_id": "file-abc", "endpoint": "/v1/chat/completions", "completion_window": "24h"}, diff --git a/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py b/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py index fa4a1656bc3..2d0f3679144 100644 --- a/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py +++ b/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py @@ -1,5 +1,6 @@ import json from datetime import datetime, timezone +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -15,8 +16,10 @@ from litellm.llms.base_llm.vector_store.transformation import ( RouterVectorStoreEmbeddingExecutor, ) from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth +from litellm.vector_stores.search_query import with_caller_query_params from litellm.proxy.vector_store_endpoints.endpoints import ( _update_request_data_with_litellm_managed_vector_store_registry, + build_request_data_from_managed_vector_store, index_create, index_list, ) @@ -3103,6 +3106,83 @@ def test_vector_store_search_rejects_caller_opensearch_index_name_only_on_opense assert fake_search.await_count == expected_searches +def test_vector_store_search_accepts_a_query_image_without_a_query(): + from fastapi.testclient import TestClient + + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.proxy_server import app + from litellm.types.vector_stores import VectorStoreSearchResponse + + query_image: Final = {"type": "base64", "media_type": "image/png", "data": "iVBORw0KGgo="} + store: Final = { + "vector_store_id": "kb", + "custom_llm_provider": "bedrock", + "litellm_params": {"vector_store_id": "kb", "custom_llm_provider": "bedrock"}, + } + registry: Final = MagicMock() + registry.get_litellm_managed_vector_store_from_registry.return_value = store + fake_search: Final = AsyncMock( + return_value=VectorStoreSearchResponse(object="vector_store.search_results.page", search_query="", data=[]) + ) + mock_auth: Final = UserAPIKeyAuth(user_id="test_internal_user", user_role=LitellmUserRoles.INTERNAL_USER.value) + original_overrides: Final = app.dependency_overrides.copy() + app.dependency_overrides[user_api_key_auth] = lambda: mock_auth + try: + with ( + patch("litellm.vector_stores.main.asearch", new=fake_search), # test-quality-ok: the search boundary under test + patch.object(litellm, "vector_store_registry", registry), # test-quality-ok: seeds the store under test + patch( # test-quality-ok: store access is not under test, so the request reaches the search boundary + "litellm.proxy.vector_store_endpoints.utils.can_user_access_vector_store", + new=AsyncMock(return_value=True), + ), + ): + router: Final = litellm.Router(model_list=[]) + with patch("litellm.proxy.proxy_server.llm_router", router): # test-quality-ok: router built over the fake search + response: Final = TestClient(app).post("/v1/vector_stores/kb/search", json={"query_image": query_image}) + finally: + app.dependency_overrides = original_overrides + + assert response.status_code == 200, response.json() + assert fake_search.await_args.kwargs["query_image"] == query_image + + +@pytest.mark.parametrize("retrieval_query_param", ["retrievalQuery", "retrieval_query"]) +def test_a_stored_raw_retrieval_query_never_reaches_the_search_request(retrieval_query_param: str): + request_data: Final = build_request_data_from_managed_vector_store( + { + "vector_store_id": "kb", + "custom_llm_provider": "bedrock", + "litellm_params": {"aws_region_name": "us-west-2", retrieval_query_param: {"text": "from the store"}}, + } + ) + + assert request_data == {"custom_llm_provider": "bedrock", "aws_region_name": "us-west-2"} + + +@pytest.mark.parametrize("retrieval_query_param", ["retrievalQuery", "retrieval_query"]) +def test_a_raw_retrieval_query_in_a_stores_extra_body_never_reaches_the_search_request(retrieval_query_param: str): + request_data: Final = build_request_data_from_managed_vector_store( + { + "vector_store_id": "kb", + "custom_llm_provider": "bedrock", + "litellm_params": {"extra_body": {retrieval_query_param: {"text": "from the store"}, "kept": "value"}}, + } + ) + + assert request_data == {"custom_llm_provider": "bedrock", "extra_body": {"kept": "value"}} + + +@pytest.mark.parametrize("caller_param", ["query_image", "retrievalQuery", "retrieval_query"]) +def test_a_callers_query_survives_a_stores_own_extra_body(caller_param: str): + caller_value: Final = {"type": "base64", "media_type": "image/png", "data": "iVBORw0KGgo="} + request_data: Final = with_caller_query_params( + {"custom_llm_provider": "bedrock", "extra_body": {"retrievalConfiguration": "from the store"}}, + {caller_param: caller_value, "retrievalConfiguration": "from the caller", "other": "from the caller"}, + ) + + assert request_data["extra_body"] == {"retrievalConfiguration": "from the store", caller_param: caller_value} + + class TestConfigOwnedVectorStores: """Stores declared under ``vector_store_registry`` in config.yaml are owned by the config file""" diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index ad1e15330d2..71efa1352e4 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -11900,3 +11900,39 @@ def test_per_call_custom_prices_keep_the_regional_and_speed_multipliers( ) assert response._hidden_params["response_cost"] == pytest.approx((1000 * 2e-06 + 1000 * 4e-06) * multiplier) + + +@pytest.mark.parametrize( + ("message_logging_off", "query_image", "echoes_response"), + [ + (True, None, False), + (False, None, True), + (False, {"type": "base64", "media_type": "image/jpeg", "data": "cXVlcnk="}, False), + ], +) +def test_cost_failure_for_an_unpriced_response_echoes_only_a_text_search_response_with_message_logging_on( + message_logging_off: bool, + query_image: Mapping[str, str] | None, + echoes_response: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "turn_off_message_logging", message_logging_off) + image_data: Final = "c2VjcmV0LWltYWdlLWJ5dGVz" + logging_obj: Final = Logging( + model="", + messages=[], + stream=False, + call_type="avector_store_search", + start_time=datetime.datetime.now(), + litellm_call_id="unpriced-search", + function_id="f", + ) + logging_obj.model_call_details["query_image"] = query_image + + with pytest.raises(ValueError, match="Model is None") as exc_info: + completion_cost( + completion_response={"data": [{"content": [{"source": {"data": image_data}}]}]}, + litellm_logging_obj=logging_obj, + ) + + assert (image_data in str(exc_info.value)) is echoes_response diff --git a/tests/unit/test_router/core/test_endpoints_characterization_generic.py b/tests/unit/test_router/core/test_endpoints_characterization_generic.py index f7669bff80c..9a78f5d4fb1 100644 --- a/tests/unit/test_router/core/test_endpoints_characterization_generic.py +++ b/tests/unit/test_router/core/test_endpoints_characterization_generic.py @@ -431,6 +431,29 @@ async def test_a_plain_response_id_without_a_model_fails_with_type_error(make_co assert host.calls == [] +@pytest.mark.parametrize("retrieval_query_param", ["retrievalQuery", "retrieval_query"]) +@pytest.mark.parametrize("endpoint", ["avector_store_search", "vector_store_search"]) +async def test_a_deployments_raw_retrieval_query_never_reaches_a_routed_vector_store_search( + make_core: MakeCore, host: EndpointHost, endpoint: str, retrieval_query_param: str +): + from_deployment: Final = {"text": "from the deployment"} + core: Final = make_core( + model_list=[ + deployment( + "g", + "d1", + **{retrieval_query_param: from_deployment, "extra_body": {retrieval_query_param: from_deployment, "kept": 1}}, + ) + ] + ) + + await _call(core, endpoint, {"model": "g", "vector_store_id": "v", "query": "q"}) + + sent: Final = host.calls[0].kwargs + assert retrieval_query_param not in sent + assert sent["extra_body"] == {"kept": 1} + + # --- container ids ------------------------------------------------------------------------ ENCODED_CONTAINER_ID: Final = ResponsesAPIRequestUtils.build_container_id("azure", "d2", "cntr_upstream") diff --git a/tests/unit/vector_stores/test_main.py b/tests/unit/vector_stores/test_main.py index 21b1eee3a3e..b89806e6377 100644 --- a/tests/unit/vector_stores/test_main.py +++ b/tests/unit/vector_stores/test_main.py @@ -120,6 +120,147 @@ def test_search_forwards_top_level_user_context_to_bedrock_retrieve(): assert posted["retrievalQuery"] == {"text": "q"} +QUERY_IMAGE: Final = {"type": "base64", "media_type": "image/jpeg", "data": "/9j/4AAQSkZJRg=="} + + +def _search_with_recorded_post(**search_kwargs: object) -> MagicMock: + client: Final = MagicMock(spec=HTTPHandler) + client.post.return_value = MagicMock( + status_code=200, + json=MagicMock( + return_value=( + {"object": "vector_store.search_results.page", "search_query": "", "data": []} + if search_kwargs.get("custom_llm_provider") == "openai" + else {"retrievalResults": []} + ) + ), + ) + credentials: Final = ( + {"api_key": "sk-test"} + if search_kwargs.get("custom_llm_provider") == "openai" + else {"aws_region_name": "us-west-2", "aws_access_key_id": "test-key-id", "aws_secret_access_key": "secret"} + ) + search(vector_store_id="kb123", client=client, litellm_logging_obj=MagicMock(), **credentials, **search_kwargs) + return client + + +@pytest.mark.parametrize("query", [None, "", " ", ["", " "]]) +def test_search_by_query_image_sends_a_bedrock_image_retrieval_query(query: str | list[str] | None): + query_kwargs: Final = {} if query is None else {"query": query} + + client: Final = _search_with_recorded_post(custom_llm_provider="bedrock", query_image=QUERY_IMAGE, **query_kwargs) + + posted: Final = json.loads(client.post.call_args.kwargs["data"]) + assert posted["retrievalQuery"] == { + "type": "IMAGE", + "image": {"format": "jpeg", "inlineContent": "/9j/4AAQSkZJRg=="}, + } + + +@pytest.mark.parametrize( + ("search_kwargs", "message"), + [ + ( + {"custom_llm_provider": "bedrock", "query": "red circle", "query_image": QUERY_IMAGE}, + "Send either query or query_image, not both.", + ), + ({"custom_llm_provider": "bedrock"}, "query or query_image is required."), + ( + {"custom_llm_provider": "bedrock", "query": 5, "query_image": QUERY_IMAGE}, + "Send either query or query_image, not both.", + ), + ( + {"custom_llm_provider": "bedrock", "query": [1, 2], "query_image": QUERY_IMAGE}, + "Send either query or query_image, not both.", + ), + ( + {"custom_llm_provider": "openai", "query_image": QUERY_IMAGE}, + "openai vector stores do not support query_image. To drop it", + ), + ( + {"custom_llm_provider": "openai", "query": "red circle", "query_image": QUERY_IMAGE}, + "openai vector stores do not support query_image. To drop it", + ), + ( + {"custom_llm_provider": "openai", "query": "", "query_image": QUERY_IMAGE, "drop_params": True}, + "drop_params dropped it", + ), + ], +) +def test_search_rejects_an_unusable_query_with_a_bad_request(search_kwargs: dict[str, object], message: str): + with pytest.raises(litellm.BadRequestError) as exc_info: + _search_with_recorded_post(**search_kwargs) + + assert exc_info.value.status_code == 400 + assert message in str(exc_info.value) + + +def test_search_rejects_an_invalid_query_image_without_echoing_its_data(): + image_data: Final = "c2VjcmV0LWltYWdlLWJ5dGVz" + + with pytest.raises(litellm.BadRequestError) as exc_info: + _search_with_recorded_post( + custom_llm_provider="bedrock", + query_image={"type": "base64", "media_type": "image/tiff", "data": image_data}, + ) + + assert exc_info.value.status_code == 400 + assert "query_image.media_type" in str(exc_info.value) + assert image_data not in str(exc_info.value) + + +@pytest.mark.parametrize("data", ["not base64!", "/9j/4AAQ SkZJRg==", "/9j/4AAQ\tSkZJRg=="]) +def test_search_rejects_query_image_data_that_is_not_base64(data: str): + with pytest.raises(litellm.BadRequestError) as exc_info: + _search_with_recorded_post( + custom_llm_provider="bedrock", + query_image={"type": "base64", "media_type": "image/jpeg", "data": data}, + ) + + assert exc_info.value.status_code == 400 + assert "query_image.data: not valid base64" in str(exc_info.value) + assert data not in str(exc_info.value) + + +@pytest.mark.parametrize("line_break", ["\n", "\r\n"]) +def test_search_sends_line_wrapped_query_image_data_to_bedrock_as_sent(line_break: str): + data: Final = line_break.join(("/9j/4AAQ", "SkZJRg==")) + + client: Final = _search_with_recorded_post( + custom_llm_provider="bedrock", + query_image={"type": "base64", "media_type": "image/jpeg", "data": data}, + ) + + posted: Final = json.loads(client.post.call_args.kwargs["data"]) + assert posted["retrievalQuery"]["image"]["inlineContent"] == data + + +def test_search_by_a_query_image_sent_in_extra_body_sends_a_bedrock_image_retrieval_query(): + client: Final = _search_with_recorded_post(custom_llm_provider="bedrock", extra_body={"query_image": QUERY_IMAGE}) + + posted: Final = json.loads(client.post.call_args.kwargs["data"]) + assert posted["retrievalQuery"] == { + "type": "IMAGE", + "image": {"format": "jpeg", "inlineContent": "/9j/4AAQSkZJRg=="}, + } + + +@pytest.mark.parametrize("drop_params_source", ["request", "global"]) +def test_search_drops_query_image_for_a_provider_without_image_search_when_drop_params_is_on( + drop_params_source: str, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "drop_params", drop_params_source == "global") + drop_kwargs: Final = {"drop_params": True} if drop_params_source == "request" else {} + + client: Final = _search_with_recorded_post( + custom_llm_provider="openai", query="red circle", query_image=QUERY_IMAGE, **drop_kwargs + ) + + posted: Final = json.loads(client.post.call_args.kwargs["data"]) + assert posted["query"] == "red circle" + assert "query_image" not in posted + + PROXY_STAMPED_METADATA: Final = { "team": "llmproxy", "user_api_key": "hashed-key", diff --git a/tests/unit/vector_stores/test_vector_store_registry.py b/tests/unit/vector_stores/test_vector_store_registry.py index a6e28f258cb..588dbf35678 100644 --- a/tests/unit/vector_stores/test_vector_store_registry.py +++ b/tests/unit/vector_stores/test_vector_store_registry.py @@ -497,6 +497,29 @@ def test_vector_store_ids_to_run_reads_every_store_id_a_malformed_vector_store_i assert VectorStoreRegistry().get_vector_store_ids_to_run(non_default_params={}, tools=tools) == ["vs_requested"] +@pytest.mark.parametrize( + ("tool_field", "tool_value", "searched_params"), + [ + ("query_image", {"type": "base64", "media_type": "image/jpeg", "data": "/9j/"}, {"max_num_results": 3}), + ("rewrite_query", True, {"max_num_results": 3, "rewrite_query": True}), + ("rewrite_query", False, {"max_num_results": 3}), + ], +) +def test_file_search_tool_passes_its_store_only_the_fields_a_store_search_takes( + tool_field: str, tool_value: object, searched_params: dict[str, object] +) -> None: + now: Final = datetime.now(timezone.utc) + registry: Final = VectorStoreRegistry( + [LiteLLM_ManagedVectorStore(vector_store_id="vs_kb", created_at=now, updated_at=now)] + ) + tools: Final = [{"type": "file_search", "vector_store_ids": ["vs_kb"], "max_num_results": 3, tool_field: tool_value}] + + to_run: Final = registry.pop_vector_stores_to_run(non_default_params={}, tools=tools) + + assert [store["vector_store_id"] for store in to_run] == ["vs_kb"] + assert to_run[0]["litellm_params"] == searched_params + + @pytest.mark.parametrize( "vector_store_ids", ["vs_allowed", ["vs_allowed", 12345], [12345], ["vs_allowed"]],