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