mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
fix(vector-stores): keep image hits to image searches and reject a caller retrievalQuery
Text searches on a multimodal Bedrock knowledge base return the same text hits as before, so chat, RAG and logging payloads stay unchanged. Image hits come back only when the caller sent query_image A caller retrievalQuery now answers 400 pointing at query_image, or is dropped under drop_params, instead of replacing the built query unchecked A file_search chat tool that carries query_image or rewrite_query searches again instead of failing, and a non-string query next to an image answers 400 instead of 500. query_image is keyword-only on search and asearch
This commit is contained in:
mateo-berri
parent
dbfba27ee6
commit
65a7206199
15 changed files with 151 additions and 283 deletions
|
|
@ -10,7 +10,8 @@ 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
|
||||
from litellm.exceptions import BadRequestError
|
||||
from litellm.litellm_core_utils.core_helpers import decoded_json_object, drops_unsupported_params
|
||||
from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base
|
||||
from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
|
|
@ -74,6 +75,11 @@ _BEDROCK_IMAGE_FORMATS: Final[Mapping[VectorStoreSearchImageMediaType, BedrockKB
|
|||
_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"
|
||||
_RETRIEVAL_QUERY_KEYS: Final = ("retrievalQuery", "retrieval_query")
|
||||
_RETRIEVAL_QUERY_UNSUPPORTED: Final = (
|
||||
"bedrock vector stores do not take retrievalQuery. Send query, or query_image to search by image. To drop it, "
|
||||
"set `litellm.drop_params=True` or for proxy:\n\n`litellm_settings:\n drop_params: true`\n"
|
||||
)
|
||||
|
||||
|
||||
class _QueryImageParams(TypedDict, total=False):
|
||||
|
|
@ -109,7 +115,7 @@ def _image_content(byte_content: str, metadata: Mapping[str, object]) -> VectorS
|
|||
|
||||
|
||||
def _result_content(
|
||||
content: BedrockKBContent | None, metadata: Mapping[str, object]
|
||||
content: BedrockKBContent | None, metadata: Mapping[str, object], include_images: bool
|
||||
) -> VectorStoreResultContent | None:
|
||||
if content is None:
|
||||
return None
|
||||
|
|
@ -117,7 +123,7 @@ def _result_content(
|
|||
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:
|
||||
if not include_images or content.get("type") != "IMAGE" or not byte_content:
|
||||
return None
|
||||
return _image_content(byte_content, metadata)
|
||||
|
||||
|
|
@ -302,14 +308,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"))
|
||||
caller_params: Final = decoded_json_object(litellm_params)
|
||||
caller_sources: Final = self._caller_sources(extra_body=extra_body, litellm_params=caller_params)
|
||||
if self._caller_value(caller_sources, _RETRIEVAL_QUERY_KEYS) is not None and not drops_unsupported_params(
|
||||
caller_params.get("drop_params")
|
||||
):
|
||||
raise BadRequestError(message=_RETRIEVAL_QUERY_UNSUPPORTED, model="", llm_provider="bedrock")
|
||||
request_body: Final[dict[str, object]] = {
|
||||
"retrievalQuery": (
|
||||
_retrieval_query(query, vector_store_search_optional_params)
|
||||
if caller_retrieval_query is None
|
||||
else caller_retrieval_query
|
||||
),
|
||||
"retrievalQuery": _retrieval_query(query, vector_store_search_optional_params),
|
||||
}
|
||||
|
||||
retrieval_config: Final[object] = (
|
||||
|
|
@ -428,10 +434,11 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
|
|||
) -> VectorStoreSearchResponse:
|
||||
try:
|
||||
response_data: Final = BedrockKBResponse(**response.json())
|
||||
include_images: Final = litellm_logging_obj.model_call_details.get("query_image") is not None
|
||||
results: Final[list[VectorStoreSearchResult]] = []
|
||||
for item in response_data.get("retrievalResults", []) or []:
|
||||
metadata = item.get("metadata", {}) or {}
|
||||
result_content = _result_content(item.get("content"), metadata)
|
||||
result_content = _result_content(item.get("content"), metadata, include_images)
|
||||
if result_content is None:
|
||||
continue
|
||||
|
||||
|
|
|
|||
|
|
@ -2126,11 +2126,7 @@ def _get_vector_store_request_for_spend_logs_payload(
|
|||
|
||||
|
||||
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}}
|
||||
return {**content, "text": REDACTED_BY_LITELM_STRING} if "text" in content else content
|
||||
|
||||
|
||||
def _redacted_vector_store_result(result: VectorStoreSearchResult) -> VectorStoreSearchResult:
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Any, Final, TypeVar
|
||||
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
|
@ -51,12 +51,10 @@ def _iterable_items(items: Sequence[_ItemT] | None) -> Sequence[_ItemT]:
|
|||
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
|
||||
def _sized_items(items: Sequence[_ItemT] | None) -> Sequence[_ItemT]:
|
||||
if items is None:
|
||||
raise TypeError("object of type 'NoneType' has no len()")
|
||||
return items
|
||||
|
||||
|
||||
class RAGQuery:
|
||||
|
|
@ -144,17 +142,24 @@ class RAGQuery:
|
|||
search_response: VectorStoreSearchResponse,
|
||||
) -> list[str | dict[str, object]]:
|
||||
"""Extract text documents from vector store search response."""
|
||||
return [text for _, text in _text_documents(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
|
||||
|
||||
@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."""
|
||||
document_results: Final = tuple(result for result, _ in _text_documents(search_response))
|
||||
top_chunks: Final[list[_StoreResultView]] = []
|
||||
original_results: Final = search_response.get("data", [])
|
||||
reranked: Final[_RerankResultsView] = {"results": rerank_response.get("results", [])}
|
||||
return [
|
||||
document_results[index]
|
||||
for result in reranked["results"]
|
||||
if (index := result.get("index")) is not None and index < len(document_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
|
||||
|
|
|
|||
|
|
@ -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(result for result in results_data if not _is_image_only_hit(result))
|
||||
all_results.extend(results_data)
|
||||
except Exception as exc:
|
||||
verbose_logger.warning(
|
||||
"file_search emulated: search failed for query='%s', vector_store_id='%s': %s",
|
||||
|
|
@ -190,24 +190,13 @@ 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."""
|
||||
text_chunks: Final = [
|
||||
c.get("text", "") if isinstance(c, dict) else getattr(c, "text", "") for c in _content_items(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]
|
||||
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:
|
||||
|
|
|
|||
|
|
@ -297,7 +297,6 @@ async def asearch(
|
|||
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,
|
||||
|
|
@ -307,6 +306,8 @@ async def asearch(
|
|||
# LiteLLM specific params,
|
||||
custom_llm_provider: str | None = None,
|
||||
router: "Router | None" = None,
|
||||
*,
|
||||
query_image: VectorStoreSearchQueryImage | None = None,
|
||||
**kwargs,
|
||||
) -> VectorStoreSearchResponse:
|
||||
"""
|
||||
|
|
@ -372,7 +373,6 @@ def search(
|
|||
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,
|
||||
|
|
@ -382,6 +382,8 @@ def search(
|
|||
# LiteLLM specific params,
|
||||
custom_llm_provider: str | None = None,
|
||||
router: "Router | None" = None,
|
||||
*,
|
||||
query_image: VectorStoreSearchQueryImage | None = None,
|
||||
**kwargs,
|
||||
) -> VectorStoreSearchResponse | Coroutine[object, object, VectorStoreSearchResponse]:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ from pydantic_core import ErrorDetails
|
|||
from typing_extensions import assert_never
|
||||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.core_helpers import as_object_list
|
||||
from litellm.types.vector_stores import VectorStoreSearchQueryImage
|
||||
|
||||
_QUERY_IMAGE: Final = TypeAdapter(VectorStoreSearchQueryImage)
|
||||
|
|
@ -49,10 +50,13 @@ SearchQuery: TypeAlias = TextSearchQuery[_Query] | ImageSearchQuery
|
|||
SearchQueryError: TypeAlias = QueryAndImageBothSent | NoQuerySent | InvalidQueryImage | QueryImageUnsupported
|
||||
|
||||
|
||||
def _has_text(query: str | Sequence[str]) -> bool:
|
||||
def _has_text(query: object) -> bool:
|
||||
if isinstance(query, str):
|
||||
return bool(query.strip())
|
||||
return any(item.strip() for item in query)
|
||||
items: Final = as_object_list(query)
|
||||
if items is None:
|
||||
return True
|
||||
return any(not isinstance(item, str) or item.strip() for item in items)
|
||||
|
||||
|
||||
def _problem(detail: ErrorDetails) -> str:
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# litellm/proxy/vector_stores/vector_store_registry.py
|
||||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import fields, replace
|
||||
from datetime import datetime, timezone
|
||||
from itertools import chain
|
||||
from typing import (
|
||||
|
|
@ -8,19 +9,17 @@ from typing import (
|
|||
Any, # noqa: TID251 # untyped non_default_params dict is the only source of the unknown key type
|
||||
Final,
|
||||
cast, # noqa: TID251 # untyped non_default_params dict is the only source of the unknown key type
|
||||
get_args,
|
||||
)
|
||||
|
||||
from pydantic import ConfigDict, TypeAdapter, ValidationError
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.core_helpers import remove_items_at_indices
|
||||
from litellm.litellm_core_utils.core_helpers import decoded_json_object, remove_items_at_indices
|
||||
from litellm.repositories.table_repositories import (
|
||||
ManagedVectorStoreIndexRepository,
|
||||
ManagedVectorStoresRepository,
|
||||
)
|
||||
from litellm.types.vector_stores import (
|
||||
VECTOR_STORE_OPENAI_PARAMS,
|
||||
LiteLLM_ManagedVectorStore,
|
||||
LiteLLM_ManagedVectorStoreIndex,
|
||||
LiteLLM_ManagedVectorStoreListResponse,
|
||||
|
|
@ -141,19 +140,15 @@ class VectorStoreRegistry:
|
|||
self.vector_store_ids_to_vector_store_map: dict[str, LiteLLM_ManagedVectorStore] = {}
|
||||
|
||||
def _extract_tool_params(self, tool: dict) -> VectorStoreToolParams:
|
||||
"""
|
||||
Extract supported parameters from a tool definition.
|
||||
|
||||
Dynamically extracts all parameters defined in VECTOR_STORE_OPENAI_PARAMS.
|
||||
"""
|
||||
# Get the list of supported param names from the Literal type
|
||||
declared_params: Final[tuple[object, ...]] = get_args(VECTOR_STORE_OPENAI_PARAMS)
|
||||
supported_params: Final = tuple(param for param in declared_params if isinstance(param, str))
|
||||
|
||||
# Extract only the params that exist in the tool
|
||||
kwargs: Final = {param: tool.get(param) for param in supported_params if param in tool}
|
||||
|
||||
return VectorStoreToolParams(**kwargs)
|
||||
tool_fields: Final = decoded_json_object(tool)
|
||||
return replace(
|
||||
VectorStoreToolParams(),
|
||||
**{
|
||||
field.name: tool_fields[field.name]
|
||||
for field in fields(VectorStoreToolParams)
|
||||
if field.name in tool_fields
|
||||
},
|
||||
)
|
||||
|
||||
def get_vector_store_ids_to_run(self, non_default_params: dict, tools: object = None) -> list[str]:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ import httpx
|
|||
import pytest
|
||||
import respx
|
||||
|
||||
import litellm
|
||||
from litellm.llms.bedrock.vector_stores.transformation import BedrockVectorStoreConfig
|
||||
|
||||
|
||||
|
|
@ -412,61 +413,73 @@ def test_transform_search_request_sends_query_image_as_an_image_retrieval_query(
|
|||
assert body["retrievalConfiguration"] == {"vectorSearchConfiguration": {"numberOfResults": 3}}
|
||||
|
||||
|
||||
CALLER_RETRIEVAL_QUERY: Final = {"type": "IMAGE", "image": {"format": "png", "inlineContent": "Y2FsbGVy"}}
|
||||
|
||||
|
||||
@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}
|
||||
def test_transform_search_request_rejects_a_callers_retrieval_query(key: str, source: str) -> None:
|
||||
caller_fields: Final = {key: CALLER_RETRIEVAL_QUERY}
|
||||
|
||||
with pytest.raises(litellm.BadRequestError) as exc_info:
|
||||
_search_body(
|
||||
extra_body=caller_fields if source == "extra_body" else None,
|
||||
litellm_params=caller_fields if source == "litellm_params" else {},
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
assert "query_image" in str(exc_info.value)
|
||||
|
||||
|
||||
def test_transform_search_request_drops_a_callers_retrieval_query_when_drop_params_is_on() -> None:
|
||||
body: Final = _search_body(
|
||||
extra_body=caller_fields if source == "extra_body" else None,
|
||||
litellm_params=caller_fields if source == "litellm_params" else {},
|
||||
extra_body={"retrievalQuery": CALLER_RETRIEVAL_QUERY}, litellm_params={"drop_params": True}
|
||||
)
|
||||
|
||||
assert body["retrievalQuery"] == caller_query
|
||||
assert body["retrievalQuery"] == {"text": "hello"}
|
||||
|
||||
|
||||
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
|
||||
QUERY_IMAGE: Final = {"type": "base64", "media_type": "image/jpeg", "data": "/9j/4AAQSkZJRg=="}
|
||||
IMAGE_HIT_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_HIT_METADATA: Final = {"x-amz-bedrock-kb-source-uri": "s3://kb-docs/shapes.txt"}
|
||||
MIXED_RETRIEVAL_RESULTS: Final = [
|
||||
{
|
||||
"content": {"type": "IMAGE", "byteContent": "data:image/png;base64,/9j/4AAQSkZJRg=="},
|
||||
"metadata": IMAGE_HIT_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_HIT_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},
|
||||
]
|
||||
TEXT_HIT: Final = {
|
||||
"score": 0.5,
|
||||
"content": [{"text": "a red circle", "type": "text"}],
|
||||
"file_id": "s3://kb-docs/shapes.txt",
|
||||
"filename": "shapes.txt",
|
||||
"attributes": TEXT_HIT_METADATA,
|
||||
}
|
||||
|
||||
|
||||
def _search_response(retrieval_results: list[dict[str, object]], query: str = "") -> dict[str, object]:
|
||||
def _search_response(
|
||||
retrieval_results: list[dict[str, object]], query: str = "", query_image: dict[str, str] | None = QUERY_IMAGE
|
||||
) -> dict[str, object]:
|
||||
config: Final = BedrockVectorStoreConfig()
|
||||
mock_log: Final = MagicMock()
|
||||
mock_log.model_call_details = {"query": query}
|
||||
mock_log.model_call_details = {"query": query, "query_image": query_image}
|
||||
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},
|
||||
]
|
||||
)
|
||||
def test_transform_search_response_to_an_image_search_returns_image_hits_next_to_text_hits() -> None:
|
||||
response: Final = _search_response(MIXED_RETRIEVAL_RESULTS)
|
||||
|
||||
assert response["search_query"] == ""
|
||||
assert response["data"] == [
|
||||
|
|
@ -480,18 +493,19 @@ def test_transform_search_response_returns_image_hits_next_to_text_hits() -> Non
|
|||
],
|
||||
"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,
|
||||
"attributes": IMAGE_HIT_METADATA,
|
||||
},
|
||||
TEXT_HIT,
|
||||
]
|
||||
|
||||
|
||||
def test_transform_search_response_to_a_text_search_returns_only_text_hits() -> None:
|
||||
response: Final = _search_response(MIXED_RETRIEVAL_RESULTS, query="red circle", query_image=None)
|
||||
|
||||
assert response["search_query"] == "red circle"
|
||||
assert response["data"] == [TEXT_HIT]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("byte_content", "metadata", "media_type", "data"),
|
||||
[
|
||||
|
|
|
|||
|
|
@ -1047,53 +1047,3 @@ 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,22 +1302,6 @@ 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
|
||||
|
|
|
|||
|
|
@ -441,60 +441,6 @@ 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():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -165,6 +165,14 @@ def test_search_by_query_image_sends_a_bedrock_image_retrieval_query(query: str
|
|||
"Send either query or query_image, not both.",
|
||||
),
|
||||
({"custom_llm_provider": "bedrock"}, "query or query_image is required."),
|
||||
(
|
||||
{"custom_llm_provider": "bedrock", "query": 5, "query_image": QUERY_IMAGE},
|
||||
"Send either query or query_image, not both.",
|
||||
),
|
||||
(
|
||||
{"custom_llm_provider": "bedrock", "query": [1, 2], "query_image": QUERY_IMAGE},
|
||||
"Send either query or query_image, not both.",
|
||||
),
|
||||
(
|
||||
{"custom_llm_provider": "openai", "query_image": QUERY_IMAGE},
|
||||
"openai vector stores do not support query_image",
|
||||
|
|
|
|||
|
|
@ -497,6 +497,25 @@ def test_vector_store_ids_to_run_reads_every_store_id_a_malformed_vector_store_i
|
|||
assert VectorStoreRegistry().get_vector_store_ids_to_run(non_default_params={}, tools=tools) == ["vs_requested"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("tool_field", "tool_value"),
|
||||
[("query_image", {"type": "base64", "media_type": "image/jpeg", "data": "/9j/"}), ("rewrite_query", True)],
|
||||
)
|
||||
def test_file_search_tool_with_a_field_the_store_search_does_not_take_still_runs_its_store(
|
||||
tool_field: str, tool_value: object
|
||||
) -> None:
|
||||
now: Final = datetime.now(timezone.utc)
|
||||
registry: Final = VectorStoreRegistry(
|
||||
[LiteLLM_ManagedVectorStore(vector_store_id="vs_kb", created_at=now, updated_at=now)]
|
||||
)
|
||||
tools: Final = [{"type": "file_search", "vector_store_ids": ["vs_kb"], "max_num_results": 3, tool_field: tool_value}]
|
||||
|
||||
to_run: Final = registry.pop_vector_stores_to_run(non_default_params={}, tools=tools)
|
||||
|
||||
assert [store["vector_store_id"] for store in to_run] == ["vs_kb"]
|
||||
assert to_run[0]["litellm_params"] == {"max_num_results": 3}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"vector_store_ids",
|
||||
["vs_allowed", ["vs_allowed", 12345], [12345], ["vs_allowed"]],
|
||||
|
|
|
|||
|
|
@ -93,33 +93,6 @@ 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,27 +9,11 @@ import { Separator } from "@/components/ui/separator";
|
|||
import { Textarea } from "@/components/ui/textarea";
|
||||
import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner";
|
||||
|
||||
interface VectorStoreTextContent {
|
||||
interface VectorStoreContent {
|
||||
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[];
|
||||
|
|
@ -194,7 +178,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">
|
||||
- {contentPreview(result.content[0])}...
|
||||
- {result.content[0].text.substring(0, 100)}...
|
||||
</span>
|
||||
)}
|
||||
</div>
|
||||
|
|
@ -214,15 +198,7 @@ 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">
|
||||
{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
|
||||
)}
|
||||
{content.text}
|
||||
</div>
|
||||
</div>
|
||||
))}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue