From 2930e510ff51959a6264805958c94bf5757f1f3b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 10 Nov 2026 10:32:24 -0800 Subject: [PATCH] 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 (cherry picked from commit 55587c70122fc0fa4369cc23886e753a91b47ee0) --- .../bedrock/vector_stores/transformation.py | 106 +++++++++++--- litellm/proxy/route_llm_request.py | 9 +- .../spend_tracking/spend_tracking_utils.py | 15 +- litellm/rag/rag_query.py | 30 ++-- .../responses/file_search/emulated_handler.py | 17 ++- .../integrations/rag/bedrock_knowledgebase.py | 12 +- litellm/types/vector_stores.py | 22 ++- litellm/vector_stores/main.py | 27 +++- litellm/vector_stores/search_query.py | 134 +++++++++++++++++ ...est_bedrock_vector_store_transformation.py | 136 ++++++++++++++++++ tests/unit/llms/test_file_search_responses.py | 50 +++++++ .../test_spend_tracking_utils.py | 16 +++ tests/unit/proxy/test_route_llm_request.py | 4 + .../test_vector_store_endpoints.py | 40 ++++++ tests/unit/rag/test_main.py | 54 +++++++ tests/unit/vector_stores/test_main.py | 94 ++++++++++++ .../_components/VectorStoreTester.test.tsx | 27 ++++ .../_components/VectorStoreTester.tsx | 30 +++- 18 files changed, 769 insertions(+), 54 deletions(-) create mode 100644 litellm/vector_stores/search_query.py diff --git a/litellm/llms/bedrock/vector_stores/transformation.py b/litellm/llms/bedrock/vector_stores/transformation.py index ba8ce7e5625..8f1efafb0bb 100644 --- a/litellm/llms/bedrock/vector_stores/transformation.py +++ b/litellm/llms/bedrock/vector_stores/transformation.py @@ -1,9 +1,13 @@ +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, cast from urllib.parse import urlparse import httpx +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_logger from litellm.litellm_core_utils.url_utils import encode_url_path_segment @@ -12,6 +16,7 @@ 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, @@ -23,7 +28,9 @@ from litellm.types.vector_stores import ( BaseVectorStoreAuthCredentials, VectorStoreIndexEndpoints, VectorStoreResultContent, + VectorStoreSearchImageMediaType, VectorStoreSearchOptionalRequestParams, + VectorStoreSearchQueryImage, VectorStoreSearchResponse, VectorStoreSearchResult, ) @@ -34,6 +41,60 @@ else: LiteLLMLoggingObj = Any +_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" + + +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 + 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] +) -> 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 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.""" @@ -59,7 +120,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: """ @@ -214,8 +275,14 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): encoded_vector_store_id: Final = encode_url_path_segment(vector_store_id, field_name="vector_store_id") url: Final = f"{api_base}/{encoded_vector_store_id}/retrieve" + caller_sources: Final = self._caller_sources(extra_body=extra_body, litellm_params=litellm_params) + caller_retrieval_query: Final = self._caller_value(caller_sources, ("retrievalQuery", "retrieval_query")) request_body: Final[dict[str, object]] = { - "retrievalQuery": BedrockKBRetrievalQuery(text=query), + "retrievalQuery": ( + _retrieval_query(query, vector_store_search_optional_params) + if caller_retrieval_query is None + else caller_retrieval_query + ), } retrieval_config: dict[str, Any] = {} @@ -244,7 +311,7 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): retrieval_config.setdefault("vectorSearchConfiguration", {})["filter"] = 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 @@ -252,19 +319,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( @@ -334,13 +400,11 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): response_data: Final = BedrockKBResponse(**response.json()) 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) + 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) @@ -348,7 +412,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 6f93bcd82dc..936c4604ba1 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -203,6 +203,10 @@ REQUIRED_ONE_OF_BODY_PARAMS_BY_ROUTE: Final[Mapping[str, tuple[str, str]]] = Map {"acreate_interaction": ("model", "agent")} ) +BODY_PARAM_ALTERNATIVES_BY_ROUTE: Final[Mapping[str, Mapping[str, tuple[str, ...]]]] = MappingProxyType( + {"avector_store_search": MappingProxyType({"query": ("query_image",)})} +) + JSON_OBJECT_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object]) @@ -240,8 +244,11 @@ def _find_missing_required_body_param( ) if missing_merge_base_param is not None: return MissingBodyParam(name=missing_merge_base_param, model_deployments_loaded=True) + alternatives: Final = BODY_PARAM_ALTERNATIVES_BY_ROUTE.get(route_type, {}) missing_present_params: Final = tuple( - param for param in REQUIRED_PRESENT_BODY_PARAMS_BY_ROUTE.get(route_type, ()) if param not in data + param + for param in REQUIRED_PRESENT_BODY_PARAMS_BY_ROUTE.get(route_type, ()) + if param not in data and not any(alternative in data for alternative in alternatives.get(param, ())) ) if not missing_present_params: return None diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 97c7a3cab57..6164bafdf96 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -80,6 +80,7 @@ from litellm.types.utils import ( StandardLoggingVectorStoreRequest, VectorStoreSearchResponse, ) +from litellm.types.vector_stores import VectorStoreResultContent from litellm.utils import get_end_user_id_for_cost_tracking if TYPE_CHECKING: @@ -1843,12 +1844,20 @@ def _get_vector_store_request_for_spend_logs_payload( ) response_data = vector_store_search_response.get("data", []) or [] for response_item in response_data: - for content_item in response_item.get("content", []) or []: - if "text" in content_item: - content_item["text"] = REDACTED_BY_LITELM_STRING + response_item["content"] = [ + _redacted_vector_store_content(content_item) for content_item in response_item.get("content", []) or [] + ] return vector_store_request_metadata +def _redacted_vector_store_content(content: VectorStoreResultContent) -> VectorStoreResultContent: + without_text: Final[VectorStoreResultContent] = ( + {**content, "text": REDACTED_BY_LITELM_STRING} if "text" in content else content + ) + source: Final = content.get("source") + return without_text if source is None else {**without_text, "source": {**source, "data": REDACTED_BY_LITELM_STRING}} + + def _get_response_for_spend_logs_payload( payload: StandardLoggingPayload | None, kwargs: dict | None = None, diff --git a/litellm/rag/rag_query.py b/litellm/rag/rag_query.py index 9f608dbc99a..830d042b83a 100644 --- a/litellm/rag/rag_query.py +++ b/litellm/rag/rag_query.py @@ -1,4 +1,4 @@ -from collections.abc import Sequence +from collections.abc import Iterator, Sequence from typing import Any, Final from typing_extensions import NotRequired, ReadOnly, TypedDict @@ -42,6 +42,13 @@ class _MessageView(TypedDict): message: ReadOnly[object] +def _text_documents(search_data: _SearchDataView) -> Iterator[tuple[_SearchResultView, str]]: + for result in search_data["results"]: + for content in result.get("content", []): + if content.get("type") == "text" and (content_text := content.get("text")): + yield result, content_text + + class RAGQuery: CONTENT_PREFIX_STRING = "Context:\n\n" @@ -126,24 +133,17 @@ class RAGQuery: search_response: Any, ) -> list[str | dict[str, object]]: """Extract text documents from vector store search response.""" - documents: Final[list[str | dict[str, object]]] = [] search_data: Final[_SearchDataView] = {"results": search_response.get("data", [])} - for result in search_data["results"]: - content_list = result.get("content", []) - for content in content_list: - if content.get("type") == "text" and content.get("text"): - documents.append(content["text"]) - return documents + return [text for _, text in _text_documents(search_data)] @staticmethod def get_top_chunks_from_rerank(search_response: Any, rerank_response: Any) -> list[_SearchResultView]: """Get the original search results corresponding to the top reranked results.""" - top_chunks: Final[list[_SearchResultView]] = [] search_data: Final[_SearchDataView] = {"results": search_response.get("data", [])} - original_results: Final = search_data["results"] + document_results: Final = tuple(result for result, _ in _text_documents(search_data)) reranked: Final[_RerankResultsView] = {"results": rerank_response.get("results", [])} - for result in reranked["results"]: - index = result.get("index") - if index is not None and index < len(original_results): - top_chunks.append(original_results[index]) - return top_chunks + return [ + document_results[index] + for result in reranked["results"] + if (index := result.get("index")) is not None and index < len(document_results) + ] diff --git a/litellm/responses/file_search/emulated_handler.py b/litellm/responses/file_search/emulated_handler.py index e2bb2ae40ea..dd0f9ab29cd 100644 --- a/litellm/responses/file_search/emulated_handler.py +++ b/litellm/responses/file_search/emulated_handler.py @@ -166,7 +166,7 @@ async def _run_vector_searches( ) results_data = response.get("data") if isinstance(response, dict) else getattr(response, "data", None) if results_data: - all_results.extend(results_data) + all_results.extend(result for result in results_data if not _is_image_only_hit(result)) except Exception as exc: verbose_logger.warning( "file_search emulated: search failed for query='%s', vector_store_id='%s': %s", @@ -190,13 +190,24 @@ def _get_field(result: object, key: str, default: object = None) -> object: return getattr(result, key, default) +def _content_items(result: object) -> Iterable[object]: + return cast(Iterable[object], _get_field(result, "content") or []) # cast-ok: iterated as today + + def _joined_content_text(result: object) -> str: """Concatenate the text of every content chunk on a search result.""" - content_items: Final = cast(Iterable[object], _get_field(result, "content") or []) # cast-ok: iterated as today - text_chunks: Final = [c.get("text", "") if isinstance(c, dict) else getattr(c, "text", "") for c in content_items] + text_chunks: Final = [ + c.get("text", "") if isinstance(c, dict) else getattr(c, "text", "") for c in _content_items(result) + ] return " ".join(t for t in text_chunks if t) +def _is_image_only_hit(result: object) -> bool: + return not _joined_content_text(result) and any( + _get_field(item, "type") == "image" for item in _content_items(result) + ) + + def _format_search_results_as_tool_output( results: list[VectorStoreSearchResult], ) -> str: diff --git a/litellm/types/integrations/rag/bedrock_knowledgebase.py b/litellm/types/integrations/rag/bedrock_knowledgebase.py index 7156d8101e1..7600eea3c70 100644 --- a/litellm/types/integrations/rag/bedrock_knowledgebase.py +++ b/litellm/types/integrations/rag/bedrock_knowledgebase.py @@ -1,4 +1,4 @@ -from typing import Any, Literal +from typing import Any, 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 5e7d517ece2..9a051b28a2c 100644 --- a/litellm/types/vector_stores.py +++ b/litellm/types/vector_stores.py @@ -2,7 +2,7 @@ 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, Literal, TypeAlias from typing_extensions import ReadOnly, TypedDict @@ -74,11 +74,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 +135,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 +339,7 @@ VECTOR_STORE_OPENAI_PARAMS = Literal[ "max_num_results", "ranking_options", "rewrite_query", + "query_image", ] diff --git a/litellm/vector_stores/main.py b/litellm/vector_stores/main.py index 4d945310293..96c739c9b44 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 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 +from litellm.vector_stores.search_query import searched_query from litellm.vector_stores.utils import VectorStoreRequestUtils if TYPE_CHECKING: @@ -285,11 +288,12 @@ 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, rewrite_query: bool | None = None, + query_image: VectorStoreSearchQueryImage | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. extra_headers: dict[str, object] | None = None, @@ -325,6 +329,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, @@ -358,11 +363,12 @@ 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, rewrite_query: bool | None = None, + query_image: VectorStoreSearchQueryImage | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. extra_headers: dict[str, object] | None = None, @@ -379,11 +385,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. @@ -439,10 +446,18 @@ def search( local_vars.update(kwargs) + search_query: Final = searched_query( + query=query, + query_image=query_image, + 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, ) ) @@ -453,7 +468,7 @@ def search( model=api_type, optional_params={ "vector_store_id": vector_store_id, - "query": query, + "query": search_query.query, **vector_store_search_optional_params, }, litellm_params={ @@ -466,7 +481,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, diff --git a/litellm/vector_stores/search_query.py b/litellm/vector_stores/search_query.py new file mode 100644 index 00000000000..c9b47cdd4fc --- /dev/null +++ b/litellm/vector_stores/search_query.py @@ -0,0 +1,134 @@ +from collections.abc import Sequence +from dataclasses import dataclass +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.types.vector_stores import VectorStoreSearchQueryImage + +_QUERY_IMAGE: Final = TypeAdapter(VectorStoreSearchQueryImage) +_Query: Final = TypeVar("_Query", bound=str | Sequence[str]) + + +@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 + + +SearchQuery: TypeAlias = TextSearchQuery[_Query] | ImageSearchQuery +SearchQueryError: TypeAlias = QueryAndImageBothSent | NoQuerySent | InvalidQueryImage | QueryImageUnsupported + + +def _has_text(query: str | Sequence[str]) -> bool: + if isinstance(query, str): + return bool(query.strip()) + return any(item.strip() for item in query) + + +def _problem(detail: ErrorDetails) -> str: + location: Final = ".".join(("query_image", *(str(part) for part in detail["loc"]))) + return f"{location}: {detail['msg']}" + + +def _validated_query_image(query_image: object) -> ImageSearchQuery | InvalidQueryImage: + try: + return ImageSearchQuery(image=_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)) + ) + + +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 NoQuerySent() 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 _: + 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(): + raise litellm.BadRequestError( + message=_error_message(resolved), model=None, llm_provider=custom_llm_provider + ) + case _: + assert_never(resolved) 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 87c57023c17..25935050ac4 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 @@ -375,3 +375,139 @@ async def test_bedrock_search_with_credentials_managed_registry( assert search_response["data"][0]["score"] == 0.9 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}} + + +@pytest.mark.parametrize("key", ["retrievalQuery", "retrieval_query"]) +@pytest.mark.parametrize("source", ["extra_body", "litellm_params"]) +def test_transform_search_request_sends_the_callers_retrieval_query_in_place_of_the_built_one( + key: str, source: str +) -> None: + caller_query: Final = {"type": "IMAGE", "image": {"format": "png", "inlineContent": "Y2FsbGVy"}} + caller_fields: Final = {key: caller_query} + + body: Final = _search_body( + extra_body=caller_fields if source == "extra_body" else None, + litellm_params=caller_fields if source == "litellm_params" else {}, + ) + + assert body["retrievalQuery"] == caller_query + + +def test_transform_search_request_prefers_the_callers_retrieval_query_over_query_image() -> None: + caller_query: Final = {"text": "from extra_body"} + + body: Final = _image_search_body("image/jpeg", extra_body={"retrievalQuery": caller_query}) + + assert body["retrievalQuery"] == caller_query + + +def _search_response(retrieval_results: list[dict[str, object]], query: str = "") -> dict[str, object]: + config: Final = BedrockVectorStoreConfig() + mock_log: Final = MagicMock() + mock_log.model_call_details = {"query": query} + 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_returns_image_hits_next_to_text_hits() -> None: + image_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_metadata: Final = {"x-amz-bedrock-kb-source-uri": "s3://kb-docs/shapes.txt"} + + response: Final = _search_response( + [ + { + "content": {"type": "IMAGE", "byteContent": "data:image/png;base64,/9j/4AAQSkZJRg=="}, + "metadata": image_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_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}, + ] + ) + + 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_metadata, + }, + { + "score": 0.5, + "content": [{"text": "a red circle", "type": "text"}], + "file_id": "s3://kb-docs/shapes.txt", + "filename": "shapes.txt", + "attributes": text_metadata, + }, + ] + + +@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"), + ("/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}} + ] diff --git a/tests/unit/llms/test_file_search_responses.py b/tests/unit/llms/test_file_search_responses.py index 9d2a5dbaf59..ae9d5083c85 100644 --- a/tests/unit/llms/test_file_search_responses.py +++ b/tests/unit/llms/test_file_search_responses.py @@ -1046,3 +1046,53 @@ class TestEmulatedFileSearchHandler: assert set(searched_ids) == {"vs_two"}, ( "A request-listed id the model picks should narrow the search to that store only" ) + + @pytest.mark.asyncio + async def test_H18_image_only_hits_are_left_out_of_the_answer_and_citations(self): + from litellm.responses.file_search.emulated_handler import ( + aresponses_with_emulated_file_search, + ) + + first_resp = self._make_mock_responses_api_response(include_function_call=True) + final_resp = self._make_mock_responses_api_response(text="A red circle.") + image_hit = { + "file_id": "file-image", + "filename": "circle.jpg", + "score": 0.99, + "content": [{"type": "image", "source": {"type": "base64", "media_type": "image/jpeg", "data": "/9j/"}}], + } + text_hit = { + "file_id": "file-text", + "filename": "shapes.txt", + "score": 0.8, + "content": [{"type": "text", "text": "The logo is a red circle."}], + } + mock_call_aresponses = AsyncMock(side_effect=[first_resp, final_resp]) + + with ( + patch.object( + import_module("litellm.responses.file_search.emulated_handler"), + "_call_aresponses", + new=mock_call_aresponses, + ), + patch( # test-quality-ok: the vector store returns the hits under test + "litellm.vector_stores.main.asearch", + new=AsyncMock(return_value={"data": [image_hit, text_hit]}), + ), + ): + result = await aresponses_with_emulated_file_search( + input="What is the logo?", + model="anthropic/claude-3-5-sonnet", + tools=[{"type": "file_search", "vector_store_ids": ["vs_001"]}], + include=["file_search_call.results"], + ) + + def _get(item, key): + return item[key] if isinstance(item, dict) else getattr(item, key, None) + + follow_up_input = str(mock_call_aresponses.call_args_list[1].kwargs["input"]) + assert "file-image" not in follow_up_input + assert "file-text" in follow_up_input + assert [_get(hit, "file_id") for hit in _get(result.output[0], "search_results")] == ["file-text"] + annotations = _get(_get(result.output[1], "content")[0], "annotations") + assert [_get(annotation, "file_id") for annotation in annotations] == ["file-text"] diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index 7291f49af99..7d4b0f37ec7 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -863,6 +863,22 @@ def test_get_vector_store_request_for_spend_logs_payload_store_prompts_false( assert result[0]["vector_store_search_response"]["data"][0]["content"][0]["type"] == "text" +def test_get_vector_store_request_for_spend_logs_payload_redacts_image_hit_data(monkeypatch: pytest.MonkeyPatch): + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "general_settings", {"store_prompts_in_spend_logs": False}) + image_hit: Final = {"type": "image", "source": {"type": "base64", "media_type": "image/jpeg", "data": "/9j/4AAQ"}} + + stored: Final = _get_vector_store_request_for_spend_logs_payload( + [{"vector_store_search_response": {"data": [{"content": [image_hit]}]}}] + ) + + assert stored is not None + assert stored[0]["vector_store_search_response"]["data"][0]["content"] == [ + {"type": "image", "source": {"type": "base64", "media_type": "image/jpeg", "data": REDACTED_BY_LITELM_STRING}} + ] + + @patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") def test_get_vector_store_request_for_spend_logs_payload_null_input(mock_should_store): # When input is None diff --git a/tests/unit/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index 67f29220cb6..c6388de67d6 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -1416,6 +1416,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 48642a9b50d..2b8f303a43b 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 @@ -3066,6 +3066,46 @@ def test_vector_store_search_rejects_caller_embedding_selection_params(blocked_k assert blocked_key in str(response.json()) +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 = {"type": "base64", "media_type": "image/png", "data": "iVBORw0KGgo="} + store = { + "vector_store_id": "kb", + "custom_llm_provider": "bedrock", + "litellm_params": {"vector_store_id": "kb", "custom_llm_provider": "bedrock"}, + } + registry = MagicMock() + registry.get_litellm_managed_vector_store_from_registry.return_value = store + fake_search = AsyncMock( + return_value=VectorStoreSearchResponse(object="vector_store.search_results.page", search_query="", data=[]) + ) + mock_auth = UserAPIKeyAuth(user_id="test_internal_user", user_role=LitellmUserRoles.INTERNAL_USER.value) + original_overrides = 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 = litellm.Router(model_list=[]) + with patch("litellm.proxy.proxy_server.llm_router", router): # test-quality-ok: router built over the fake search + response = TestClient(app).post("/v1/vector_stores/kb/search", json={"query_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 + + class TestConfigOwnedVectorStores: """Stores declared under ``vector_store_registry`` in config.yaml are owned by the config file""" diff --git a/tests/unit/rag/test_main.py b/tests/unit/rag/test_main.py index 87599d32c22..75f0e5b47d8 100644 --- a/tests/unit/rag/test_main.py +++ b/tests/unit/rag/test_main.py @@ -441,6 +441,60 @@ async def test_aquery_forwards_filters_to_vector_store_search( assert request_body["max_num_results"] == 50 +@pytest.mark.asyncio +async def test_aquery_rerank_picks_the_text_hit_the_reranker_ranked_when_image_hits_come_first(monkeypatch): + from litellm.types.rerank import RerankResponse + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + image_hit: Final = { + "file_id": "file-image", + "filename": "circle.jpg", + "score": 0.99, + "content": [{"type": "image", "source": {"type": "base64", "media_type": "image/jpeg", "data": "/9j/"}}], + } + text_hit: Final = { + "file_id": "file-text", + "filename": "shapes.txt", + "score": 0.8, + "content": [{"type": "text", "text": "The logo is a red circle."}], + } + reranked_documents: Final[list[object]] = [] + + async def fake_arerank(**kwargs: object) -> RerankResponse: + reranked_documents.append(kwargs["documents"]) + return RerankResponse(id="rr_1", results=[{"index": 0, "relevance_score": 0.9}], meta={}) + + with respx.mock(assert_all_called=True) as respx_mock, patch("litellm.arerank", side_effect=fake_arerank): + respx_mock.post("https://example.com/v1/vector_stores/vs_test_123/search").mock( + return_value=httpx.Response( + 200, + json={"object": "vector_store.search_results.page", "search_query": "q", "data": [image_hit, text_hit]}, + ) + ) + completion_route: Final = respx_mock.post("https://example.com/v1/chat/completions").mock( + return_value=httpx.Response( + 200, + content=( + '{"id":"chatcmpl-test","object":"chat.completion","created":1,"model":"gpt-4o-mini",' + '"choices":[{"index":0,"message":{"role":"assistant","content":"answer"},"finish_reason":"stop"}],' + '"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}' + ), + ) + ) + await litellm.aquery( + model="openai/gpt-4o-mini", + messages=[{"role": "user", "content": "What is the logo?"}], + retrieval_config={"vector_store_id": "vs_test_123", "custom_llm_provider": "openai"}, + rerank={"enabled": True, "model": "cohere/rerank-english-v3.0", "top_n": 1}, + api_key="sk-test", + api_base="https://example.com/v1", + ) + completion_body: Final = completion_route.calls.last.request.content.decode() + + assert reranked_documents == [["The logo is a red circle."]] + assert "The logo is a red circle." in completion_body + + @pytest.mark.asyncio async def test_aquery_forwards_provider_retrieval_config_and_router_to_search(): """ diff --git a/tests/unit/vector_stores/test_main.py b/tests/unit/vector_stores/test_main.py index 25a378270b1..473da3f2550 100644 --- a/tests/unit/vector_stores/test_main.py +++ b/tests/unit/vector_stores/test_main.py @@ -15,6 +15,7 @@ from unittest.mock import MagicMock, patch import httpx import pytest +import litellm import litellm.vector_stores.main as vector_stores_main from litellm.llms.base_llm.vector_store.transformation import ( RouterVectorStoreEmbeddingExecutor, @@ -119,6 +120,99 @@ 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": "openai", "query_image": QUERY_IMAGE}, + "openai vector stores do not support query_image", + ), + ( + {"custom_llm_provider": "openai", "query": "", "query_image": QUERY_IMAGE, "drop_params": True}, + "query or query_image is required.", + ), + ], +) +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("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/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTester.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTester.test.tsx index 4203a3dbf70..cce4d5c8edc 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTester.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTester.test.tsx @@ -93,6 +93,33 @@ describe("VectorStoreTester", () => { expect(screen.getByText("Filename:").parentElement).toHaveTextContent("notes.txt"); }); + it("shows an image hit as an image when the result is expanded", async () => { + const user = userEvent.setup(); + mockSearch.mockResolvedValue({ + object: "vector_store.search_results.page", + search_query: "", + data: [ + { + score: 0.88, + content: [{ type: "image", source: { type: "base64", media_type: "image/jpeg", data: "/9j/4AAQ" } }], + file_id: "file-2", + filename: "red-circle.jpg", + attributes: {}, + }, + ], + }); + renderTester(); + + await user.type(queryInput(), "red circle"); + await user.click(searchButton()); + await user.click(await screen.findByText("Result 1")); + + expect(screen.getByRole("img", { name: "red-circle.jpg" })).toHaveAttribute( + "src", + "data:image/jpeg;base64,/9j/4AAQ", + ); + }); + it("warns instead of searching when the query is only whitespace", async () => { const user = userEvent.setup(); renderTester(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTester.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTester.tsx index a9c4e0d1061..cdd2db7ea05 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTester.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTester.tsx @@ -8,11 +8,27 @@ import { Separator } from "@/components/ui/separator"; import { Textarea } from "@/components/ui/textarea"; import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; -interface VectorStoreContent { +interface VectorStoreTextContent { text: string; type: string; } +interface VectorStoreImageContent { + type: "image"; + source: { + type: "base64"; + media_type: string; + data: string; + }; +} + +type VectorStoreContent = VectorStoreTextContent | VectorStoreImageContent; + +const isImageContent = (content: VectorStoreContent): content is VectorStoreImageContent => "source" in content; + +const contentPreview = (content: VectorStoreContent): string => + isImageContent(content) ? content.source.media_type : content.text.substring(0, 100); + interface VectorStoreResult { score: number; content: VectorStoreContent[]; @@ -174,7 +190,7 @@ export const VectorStoreTester: React.FC = ({ vectorStor {/* Show preview of content when collapsed */} {!isExpanded && result.content && result.content[0] && ( - - {result.content[0].text.substring(0, 100)}... + - {contentPreview(result.content[0])}... )} @@ -194,7 +210,15 @@ export const VectorStoreTester: React.FC = ({ vectorStor Content ({content.type})
- {content.text} + {isImageContent(content) ? ( + {result.filename + ) : ( + content.text + )}
))}