Merge 37e3115850 into 6085538b33

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,

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 ##

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