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
(cherry picked from commit 55587c7012)
This commit is contained in:
mateo-berri
parent
2a61f2114d
commit
2930e510ff
18 changed files with 769 additions and 54 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, 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<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."""
|
||||
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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 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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
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,
|
||||
|
|
|
|||
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}}
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"},
|
||||
|
|
|
|||
|
|
@ -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"""
|
||||
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
|
|
|
|||
|
|
@ -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<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>
|
||||
|
|
@ -194,7 +210,15 @@ export const VectorStoreTester: React.FC<VectorStoreTesterProps> = ({ vectorStor
|
|||
Content ({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