mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
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
This commit is contained in:
mateo-berri
parent
0b070d25d3
commit
55587c7012
18 changed files with 756 additions and 56 deletions
|
|
@ -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<media_type>[^;,]*)[^,]*;base64,(?P<data>.*)", 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,
|
||||
|
|
|
|||
|
|
@ -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])
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
134
litellm/vector_stores/search_query.py
Normal file
134
litellm/vector_stores/search_query.py
Normal file
|
|
@ -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": "<base64>"}.'
|
||||
)
|
||||
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)
|
||||
|
|
@ -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}}
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"},
|
||||
|
|
|
|||
|
|
@ -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"""
|
||||
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
|
|
|
|||
|
|
@ -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<VectorStoreTesterProps> = ({ vectorStor
|
|||
{/* Show preview of content when collapsed */}
|
||||
{!isExpanded && result.content && result.content[0] && (
|
||||
<span className="ml-2 max-w-md truncate text-xs text-muted-foreground">
|
||||
- {result.content[0].text.substring(0, 100)}...
|
||||
- {contentPreview(result.content[0])}...
|
||||
</span>
|
||||
)}
|
||||
</div>
|
||||
|
|
@ -198,7 +214,15 @@ export const VectorStoreTester: React.FC<VectorStoreTesterProps> = ({ vectorStor
|
|||
{t("vectorStores.tester.content", { type: content.type })}
|
||||
</div>
|
||||
<div className="max-h-40 overflow-y-auto rounded-sm border bg-muted/50 p-3 text-sm text-foreground">
|
||||
{content.text}
|
||||
{isImageContent(content) ? (
|
||||
<img
|
||||
src={`data:${content.source.media_type};base64,${content.source.data}`}
|
||||
alt={result.filename || content.source.media_type}
|
||||
className="max-h-36 rounded-sm"
|
||||
/>
|
||||
) : (
|
||||
content.text
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
))}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue