diff --git a/litellm/llms/s3_vectors/vector_stores/transformation.py b/litellm/llms/s3_vectors/vector_stores/transformation.py index 2b0de3dc9ea..2e872ee22e8 100644 --- a/litellm/llms/s3_vectors/vector_stores/transformation.py +++ b/litellm/llms/s3_vectors/vector_stores/transformation.py @@ -3,6 +3,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.vector_store.transformation import ( BaseQueryEmbeddingVectorStoreConfig, @@ -28,6 +29,8 @@ else: _DEFAULT_QUERY_EMBEDDING_MODEL: Final = "text-embedding-3-small" _DEFAULT_TOP_K: Final = 5 +# The metadata key `/rag/ingest` writes each chunk's text to. +_DEFAULT_TEXT_FIELD: Final = "source_text" S3_VECTORS_STORE_ID_ERROR: Final = ( "vector_store_id must be in format 'bucket_name:index_name' for S3 Vectors, " "or vector_bucket_name must be provided in litellm_params" @@ -59,6 +62,13 @@ def s3_vectors_ingest_target(vector_store_config: Mapping[str, object]) -> tuple return explicit_bucket_name or derived_bucket_name, explicit_index_name or derived_index_name +def s3_vectors_text_field(litellm_params: object) -> str: + """The metadata key a search hit's text is read from: `s3_vectors_text_field`, else the key `/rag/ingest` writes.""" + params: Final = as_str_mapping(litellm_params) + configured: Final = _non_empty_str(params.get("s3_vectors_text_field")) if params is not None else None + return configured or _DEFAULT_TEXT_FIELD + + def s3_vectors_configured_embedding_model(litellm_params: Mapping[str, object]) -> str | None: return _non_empty_str(litellm_params.get("litellm_embedding_model") or litellm_params.get("embedding_model")) @@ -212,13 +222,13 @@ class S3VectorsVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, BaseAWSLLM try: response_data: Final = response.json() results: Final[list[VectorStoreSearchResult]] = [] + text_field: Final = s3_vectors_text_field(litellm_logging_obj.model_call_details.get("litellm_params")) for item in response_data.get("vectors", []) or []: metadata = item.get("metadata", {}) or {} - source_text = metadata.get("source_text", "") - - if not source_text: - continue + # A hit without the text field is still a hit: it is returned with empty text and its metadata, + # as the Milvus and Azure AI Search stores do, rather than dropped from the results. + text = metadata.get(text_field) chunk_index = metadata.get("chunk_index", "0") file_id = f"s3-vectors-chunk-{chunk_index}" @@ -232,7 +242,7 @@ class S3VectorsVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, BaseAWSLLM results.append( VectorStoreSearchResult( score=score, - content=[VectorStoreResultContent(text=source_text, type="text")], + content=[VectorStoreResultContent(text=text if isinstance(text, str) else "", type="text")], file_id=file_id, filename=filename, attributes=metadata, diff --git a/litellm/types/router.py b/litellm/types/router.py index a9f28fbc673..cc8edbacb7c 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -855,6 +855,7 @@ class LiteLLMParamsTypedDict(TypedDict, total=False): vector_bucket_name: str | None index_name: str | None embedding_model: str | None + s3_vectors_text_field: ReadOnly[str | None] ## IBM WATSONX ## watsonx_region_name: str | None ## CUSTOM PRICING ## diff --git a/tests/unit/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py b/tests/unit/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py index 3ef09c5e514..724de9d9244 100644 --- a/tests/unit/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py +++ b/tests/unit/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py @@ -3,6 +3,7 @@ from unittest.mock import AsyncMock, MagicMock, Mock, patch import httpx import pytest +import respx from litellm.llms.base_llm.vector_store.transformation import ( RouterVectorStoreEmbeddingExecutor, @@ -12,6 +13,7 @@ from litellm.llms.s3_vectors.vector_stores.transformation import ( ) from litellm.types.utils import EmbeddingResponse from litellm.types.vector_stores import VectorStoreSearchResponse +from litellm.vector_stores import search as vector_store_search from tests.unit.proxy.router_entry_point_stubs import route_entry_points_to_stubs QUERY_VECTOR = [0.1, 0.2, 0.3] @@ -310,3 +312,43 @@ class TestS3VectorsVectorStoreConfig: result = config.map_openai_params(non_default_params, optional_params, False) assert result["maxResults"] == 5 + + +_QUERY_VECTORS_URL = "https://s3vectors.us-west-2.api.aws/QueryVectors" + + +@pytest.mark.parametrize( + ("store_params", "metadata", "expected_text"), + [ + pytest.param({}, {"source_text": "ingested chunk"}, "ingested chunk", id="default-field"), + pytest.param({"s3_vectors_text_field": "title"}, {"title": "Bus timings"}, "Bus timings", id="configured-field"), + pytest.param({}, {"title": "Bus timings"}, "", id="hit-without-text-is-kept"), + ], +) +def test_search_reads_text_from_the_configured_metadata_field( + respx_mock: respx.MockRouter, store_params: dict[str, str], metadata: dict[str, str], expected_text: str +): + """A hit's text comes from `s3_vectors_text_field`, default `source_text`, and a hit without it isn't dropped. + + Indexes built outside `/rag/ingest` keep their text under their own keys; before, every one of their hits was + dropped and a search came back empty. + """ + respx_mock.post(_QUERY_VECTORS_URL).mock( + return_value=httpx.Response(200, json={"vectors": [{"key": "k1", "distance": 0.25, "metadata": metadata}]}) + ) + + response = vector_store_search( + vector_store_id="test-bucket:test-index", + query="when does the bus leave?", + custom_llm_provider="s3_vectors", + aws_region_name="us-west-2", + aws_access_key_id="test-key", + aws_secret_access_key="test-secret", + _direct_vector_store_embedding_executor=_RecordingExecutor(), + **store_params, + ) + + [hit] = response["data"] + assert hit["content"] == [{"text": expected_text, "type": "text"}] + assert hit["attributes"] == metadata + assert hit["score"] == 0.75