From 55587c70122fc0fa4369cc23886e753a91b47ee0 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 --- .../bedrock/vector_stores/transformation.py | 106 +++++++++++--- litellm/proxy/route_llm_request.py | 5 +- .../spend_tracking/spend_tracking_utils.py | 6 +- litellm/rag/rag_query.py | 33 ++--- .../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 | 93 ++++++++++++ .../_components/VectorStoreTester.test.tsx | 27 ++++ .../_components/VectorStoreTester.tsx | 30 +++- 18 files changed, 756 insertions(+), 56 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 b2b3a9b0faf..ebe6bf8001f 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, Protocol, cast, runtime_checkable from urllib.parse import urlparse import httpx +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_logger from litellm.litellm_core_utils.core_helpers import decoded_json_object @@ -13,6 +17,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, @@ -24,7 +29,9 @@ from litellm.types.vector_stores import ( BaseVectorStoreAuthCredentials, VectorStoreIndexEndpoints, VectorStoreResultContent, + VectorStoreSearchImageMediaType, VectorStoreSearchOptionalRequestParams, + VectorStoreSearchQueryImage, VectorStoreSearchResponse, VectorStoreSearchResult, ) @@ -61,6 +68,60 @@ def _vector_search_configuration(retrieval_config: object) -> Mapping[str, objec return decoded_json_object(decoded_json_object(retrieval_config).get("vectorSearchConfiguration", {})) +_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.""" @@ -86,7 +147,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: """ @@ -241,8 +302,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_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: Final[object] = ( @@ -274,7 +341,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 @@ -282,19 +349,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( @@ -364,13 +430,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) @@ -378,7 +442,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/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 15b2d18eeae..9761e4ea7e9 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -2126,7 +2126,11 @@ def _get_vector_store_request_for_spend_logs_payload( def _redacted_vector_store_content(content: VectorStoreResultContent) -> VectorStoreResultContent: - return {**content, "text": REDACTED_BY_LITELM_STRING} if "text" in content else content + 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 _redacted_vector_store_result(result: VectorStoreSearchResult) -> VectorStoreSearchResult: diff --git a/litellm/rag/rag_query.py b/litellm/rag/rag_query.py index 66686beb982..fa8af883688 100644 --- a/litellm/rag/rag_query.py +++ b/litellm/rag/rag_query.py @@ -1,4 +1,4 @@ -from collections.abc import Mapping, Sequence +from collections.abc import Iterator, Mapping, Sequence from typing import Any, Final, TypeVar from typing_extensions import NotRequired, ReadOnly, TypedDict @@ -51,10 +51,12 @@ def _iterable_items(items: Sequence[_ItemT] | None) -> Sequence[_ItemT]: return items -def _sized_items(items: Sequence[_ItemT] | None) -> Sequence[_ItemT]: - if items is None: - raise TypeError("object of type 'NoneType' has no len()") - return items +def _text_documents(search_response: VectorStoreSearchResponse) -> Iterator[tuple[_StoreResultView, str]]: + results: Final[Sequence[_StoreResultView]] = _iterable_items(search_response.get("data", [])) + for result in results: + for content in _iterable_items(result.get("content", [])): + if content.get("type") == "text" and (content_text := content.get("text")): + yield result, content_text class RAGQuery: @@ -142,24 +144,17 @@ class RAGQuery: search_response: VectorStoreSearchResponse, ) -> list[str | dict[str, object]]: """Extract text documents from vector store search response.""" - documents: Final[list[str | dict[str, object]]] = [] - results: Final[Sequence[_StoreResultView]] = _iterable_items(search_response.get("data", [])) - for result in results: - for content in _iterable_items(result.get("content", [])): - if content.get("type") == "text" and (content_text := content.get("text")): - documents.append(content_text) - return documents + return [text for _, text in _text_documents(search_response)] @staticmethod def get_top_chunks_from_rerank( search_response: VectorStoreSearchResponse, rerank_response: Any ) -> list[_StoreResultView]: """Get the original search results corresponding to the top reranked results.""" - top_chunks: Final[list[_StoreResultView]] = [] - original_results: Final = search_response.get("data", []) + document_results: Final = tuple(result for result, _ in _text_documents(search_response)) 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(sized_results := _sized_items(original_results)): - top_chunks.append(sized_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 f29f6c83ef0..b123d263986 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 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..2533f8452e7 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 9f53accfa0b..454a73e63c2 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, load_credentials_from_list +from litellm.vector_stores.search_query import searched_query from litellm.vector_stores.utils import VectorStoreRequestUtils if TYPE_CHECKING: @@ -289,11 +292,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, @@ -329,6 +333,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,11 +367,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, @@ -383,11 +389,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,10 +451,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, ) ) @@ -458,7 +473,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={ @@ -471,7 +486,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 424aaad1fd1..72f84d18749 100644 --- a/tests/unit/llms/test_file_search_responses.py +++ b/tests/unit/llms/test_file_search_responses.py @@ -1047,3 +1047,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 af8bc254e2c..b2f6f015ce2 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -1302,6 +1302,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"}} + shared: Final = [{"vector_store_search_response": {"data": [{"content": [image_hit]}]}}] + + stored: Final = _get_vector_store_request_for_spend_logs_payload(shared) + + 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}} + ] + assert shared[0]["vector_store_search_response"]["data"][0]["content"][0]["source"]["data"] == "/9j/4AAQ" + + @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 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..732f97e4ab8 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 @@ -3103,6 +3103,46 @@ 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 = {"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 21b1eee3a3e..5c1faab32db 100644 --- a/tests/unit/vector_stores/test_main.py +++ b/tests/unit/vector_stores/test_main.py @@ -120,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 fd77834cf30..ffb91a31814 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 @@ -9,11 +9,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[]; @@ -178,7 +194,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])}... )} @@ -198,7 +214,15 @@ export const VectorStoreTester: React.FC = ({ vectorStor {t("vectorStores.tester.content", { type: content.type })}
- {content.text} + {isImageContent(content) ? ( + {result.filename + ) : ( + content.text + )}
))}