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 + )}
))}