diff --git a/litellm/llms/bedrock/vector_stores/transformation.py b/litellm/llms/bedrock/vector_stores/transformation.py index 8f1efafb0bb..e571c55c2c9 100644 --- a/litellm/llms/bedrock/vector_stores/transformation.py +++ b/litellm/llms/bedrock/vector_stores/transformation.py @@ -10,6 +10,8 @@ import httpx from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_logger +from litellm.exceptions import BadRequestError +from litellm.litellm_core_utils.core_helpers import as_str_mapping, drops_unsupported_params from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM @@ -47,6 +49,11 @@ _BEDROCK_IMAGE_FORMATS: Final[Mapping[VectorStoreSearchImageMediaType, BedrockKB _BASE64_DATA_URI: Final = re.compile(r"data:(?P[^;,]*)[^,]*;base64,(?P.*)", re.DOTALL) _SOURCE_FILE_MIME_TYPE_KEY: Final = "x-amz-bedrock-kb-source-file-mime-type" _UNKNOWN_MEDIA_TYPE: Final = "application/octet-stream" +_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): @@ -82,7 +89,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 @@ -90,7 +97,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) @@ -275,14 +282,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")) + caller_params: Final = as_str_mapping(litellm_params) or {} + 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: dict[str, Any] = {} @@ -398,10 +405,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 diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 6164bafdf96..97c7a3cab57 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -80,7 +80,6 @@ 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: @@ -1844,20 +1843,12 @@ def _get_vector_store_request_for_spend_logs_payload( ) response_data = vector_store_search_response.get("data", []) or [] for response_item in response_data: - response_item["content"] = [ - _redacted_vector_store_content(content_item) for content_item in response_item.get("content", []) or [] - ] + for content_item in response_item.get("content", []) or []: + if "text" in content_item: + content_item["text"] = REDACTED_BY_LITELM_STRING return vector_store_request_metadata -def _redacted_vector_store_content(content: VectorStoreResultContent) -> VectorStoreResultContent: - without_text: Final[VectorStoreResultContent] = ( - {**content, "text": REDACTED_BY_LITELM_STRING} if "text" in content else content - ) - source: Final = content.get("source") - return without_text if source is None else {**without_text, "source": {**source, "data": REDACTED_BY_LITELM_STRING}} - - def _get_response_for_spend_logs_payload( payload: StandardLoggingPayload | None, kwargs: dict | None = None, diff --git a/litellm/rag/rag_query.py b/litellm/rag/rag_query.py index 830d042b83a..9f608dbc99a 100644 --- a/litellm/rag/rag_query.py +++ b/litellm/rag/rag_query.py @@ -1,4 +1,4 @@ -from collections.abc import Iterator, Sequence +from collections.abc import Sequence from typing import Any, Final from typing_extensions import NotRequired, ReadOnly, TypedDict @@ -42,13 +42,6 @@ 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" @@ -133,17 +126,24 @@ 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", [])} - return [text for _, text in _text_documents(search_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 @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", [])} - document_results: Final = tuple(result for result, _ in _text_documents(search_data)) + original_results: Final = search_data["results"] 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(original_results): + top_chunks.append(original_results[index]) + return top_chunks diff --git a/litellm/responses/file_search/emulated_handler.py b/litellm/responses/file_search/emulated_handler.py index dd0f9ab29cd..e2bb2ae40ea 100644 --- a/litellm/responses/file_search/emulated_handler.py +++ b/litellm/responses/file_search/emulated_handler.py @@ -166,7 +166,7 @@ async def _run_vector_searches( ) results_data = response.get("data") if isinstance(response, dict) else getattr(response, "data", None) if results_data: - all_results.extend(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: diff --git a/litellm/vector_stores/main.py b/litellm/vector_stores/main.py index 96c739c9b44..b4e1a0117f9 100644 --- a/litellm/vector_stores/main.py +++ b/litellm/vector_stores/main.py @@ -293,7 +293,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, @@ -303,6 +302,8 @@ async def asearch( # LiteLLM specific params, custom_llm_provider: str | None = None, router: "Router | None" = None, + *, + query_image: VectorStoreSearchQueryImage | None = None, **kwargs, ) -> VectorStoreSearchResponse: """ @@ -368,7 +369,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, @@ -378,6 +378,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]: """ diff --git a/litellm/vector_stores/search_query.py b/litellm/vector_stores/search_query.py index c9b47cdd4fc..8574ce1a4bc 100644 --- a/litellm/vector_stores/search_query.py +++ b/litellm/vector_stores/search_query.py @@ -49,10 +49,17 @@ SearchQuery: TypeAlias = TextSearchQuery[_Query] | ImageSearchQuery SearchQueryError: TypeAlias = QueryAndImageBothSent | NoQuerySent | InvalidQueryImage | QueryImageUnsupported -def _has_text(query: str | Sequence[str]) -> bool: +def _as_object_list(value: object) -> Sequence[object] | None: + return value if isinstance(value, list) else None + + +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: diff --git a/litellm/vector_stores/vector_store_registry.py b/litellm/vector_stores/vector_store_registry.py index b93d9874dbc..2365b920814 100644 --- a/litellm/vector_stores/vector_store_registry.py +++ b/litellm/vector_stores/vector_store_registry.py @@ -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 as_str_mapping, 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, @@ -137,19 +136,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 = as_str_mapping(tool) or {} + 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]: """ diff --git a/tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py b/tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py index 25935050ac4..93b7dad9b17 100644 --- a/tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py +++ b/tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py @@ -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"), [ diff --git a/tests/unit/llms/test_file_search_responses.py b/tests/unit/llms/test_file_search_responses.py index ae9d5083c85..9d2a5dbaf59 100644 --- a/tests/unit/llms/test_file_search_responses.py +++ b/tests/unit/llms/test_file_search_responses.py @@ -1046,53 +1046,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"] diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index 7d4b0f37ec7..7291f49af99 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -863,22 +863,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"}} - - stored: Final = _get_vector_store_request_for_spend_logs_payload( - [{"vector_store_search_response": {"data": [{"content": [image_hit]}]}}] - ) - - assert stored is not None - assert stored[0]["vector_store_search_response"]["data"][0]["content"] == [ - {"type": "image", "source": {"type": "base64", "media_type": "image/jpeg", "data": REDACTED_BY_LITELM_STRING}} - ] - - @patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") def test_get_vector_store_request_for_spend_logs_payload_null_input(mock_should_store): # When input is None diff --git a/tests/unit/rag/test_main.py b/tests/unit/rag/test_main.py index 75f0e5b47d8..87599d32c22 100644 --- a/tests/unit/rag/test_main.py +++ b/tests/unit/rag/test_main.py @@ -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(): """ diff --git a/tests/unit/vector_stores/test_main.py b/tests/unit/vector_stores/test_main.py index 473da3f2550..21a4a141bc0 100644 --- a/tests/unit/vector_stores/test_main.py +++ b/tests/unit/vector_stores/test_main.py @@ -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", diff --git a/tests/unit/vector_stores/test_vector_store_registry.py b/tests/unit/vector_stores/test_vector_store_registry.py index 83b59d77aaf..0fed3c0cda2 100644 --- a/tests/unit/vector_stores/test_vector_store_registry.py +++ b/tests/unit/vector_stores/test_vector_store_registry.py @@ -377,6 +377,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"]], diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTester.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTester.test.tsx index cce4d5c8edc..4203a3dbf70 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTester.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTester.test.tsx @@ -93,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(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTester.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTester.tsx index cdd2db7ea05..a9c4e0d1061 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTester.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTester.tsx @@ -8,27 +8,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[]; @@ -190,7 +174,7 @@ export const VectorStoreTester: React.FC = ({ vectorStor {/* Show preview of content when collapsed */} {!isExpanded && result.content && result.content[0] && ( - - {contentPreview(result.content[0])}... + - {result.content[0].text.substring(0, 100)}... )} @@ -210,15 +194,7 @@ export const VectorStoreTester: React.FC = ({ vectorStor Content ({content.type})
- {isImageContent(content) ? ( - {result.filename - ) : ( - content.text - )} + {content.text}
))}