diff --git a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py index 0bcf16ee06f..aa2cc8575f9 100644 --- a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py +++ b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py @@ -1,4 +1,5 @@ -from collections.abc import Mapping +from collections.abc import Iterable, Mapping, Sequence +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol import httpx @@ -57,23 +58,57 @@ class VertexSearchSnippet(TypedDict, total=False): htmlSnippet: ReadOnly[str] +class VertexSearchExtractiveContent(TypedDict, total=False): + """One ``extractive_answers`` or ``extractive_segments`` entry (opt-in via ``extractiveContentSpec``).""" + + content: ReadOnly[str] + pageNumber: ReadOnly[str] + + class VertexSearchDerivedStructData(TypedDict, total=False): - """The ``derivedStructData`` blob Discovery Engine attaches to each search hit.""" + """The ``derivedStructData`` blob Discovery Engine attaches to each document hit.""" title: ReadOnly[str] link: ReadOnly[str] displayLink: ReadOnly[str] formattedUrl: ReadOnly[str] snippets: ReadOnly[list[VertexSearchSnippet]] + extractive_answers: ReadOnly[list[VertexSearchExtractiveContent]] + extractive_segments: ReadOnly[list[VertexSearchExtractiveContent]] class VertexSearchDocument(TypedDict, total=False): + id: ReadOnly[str] + structData: ReadOnly[Mapping[str, object]] derivedStructData: ReadOnly[VertexSearchDerivedStructData] +class VertexSearchChunkDocumentMetadata(TypedDict, total=False): + uri: ReadOnly[str] + title: ReadOnly[str] + structData: ReadOnly[Mapping[str, object]] + + +class VertexSearchChunkPageSpan(TypedDict, total=False): + pageStart: ReadOnly[int] + pageEnd: ReadOnly[int] + + +class VertexSearchChunk(TypedDict, total=False): + """A hit when ``searchResultMode`` is ``CHUNKS``; such hits carry no ``document`` and no top-level ``id``.""" + + id: ReadOnly[str] + name: ReadOnly[str] + content: ReadOnly[str] + documentMetadata: ReadOnly[VertexSearchChunkDocumentMetadata] + pageSpan: ReadOnly[VertexSearchChunkPageSpan] + relevanceScore: ReadOnly[float] + + class VertexSearchHit(TypedDict, total=False): id: ReadOnly[str] document: ReadOnly[VertexSearchDocument] + chunk: ReadOnly[VertexSearchChunk] class VertexSearchApiResponse(TypedDict, total=False): @@ -98,6 +133,97 @@ def _vertex_search_payload(response: _VertexSearchApiSource) -> VertexSearchApiR return response.json() +_UNKNOWN_DOCUMENT: Final = "Unknown Document" +_EMPTY_DOCUMENT: Final[VertexSearchDocument] = {} +_EMPTY_DERIVED_STRUCT_DATA: Final[VertexSearchDerivedStructData] = {} +_EMPTY_CHUNK_DOCUMENT_METADATA: Final[VertexSearchChunkDocumentMetadata] = {} + + +def _joined_content(entries: Sequence[VertexSearchExtractiveContent]) -> str: + return "\n\n".join(content for entry in entries if (content := entry.get("content"))) + + +def _snippet_text(snippets: Sequence[VertexSearchSnippet]) -> str: + return " ".join(snippet.get("snippet", snippet.get("htmlSnippet", "")) for snippet in snippets) + + +def _document_text(derived: VertexSearchDerivedStructData) -> str: + candidates: Final = ( + _joined_content(derived.get("extractive_segments", ())), + _joined_content(derived.get("extractive_answers", ())), + _snippet_text(derived.get("snippets", ())), + derived.get("title", ""), + ) + return next((text for text in candidates if text), "") + + +def _document_id_from_chunk_name(name: str) -> str: + return name.partition("/documents/")[2].partition("/")[0] + + +def _non_empty_attributes(pairs: Iterable[tuple[str, object]]) -> Mapping[str, object]: + return MappingProxyType({key: value for key, value in pairs if value}) + + +def _chunk_result(chunk: VertexSearchChunk, positional_score: float) -> VectorStoreSearchResult: + metadata: Final = chunk.get("documentMetadata", _EMPTY_CHUNK_DOCUMENT_METADATA) + uri: Final = metadata.get("uri", "") + title: Final = metadata.get("title", "") + document_id: Final = _document_id_from_chunk_name(chunk.get("name", "")) + return VectorStoreSearchResult( + score=chunk.get("relevanceScore", positional_score), + content=[VectorStoreResultContent(text=chunk.get("content", ""), type="text")], + file_id=uri or document_id, + filename=title or _UNKNOWN_DOCUMENT, + attributes={ + "document_id": document_id, + **_non_empty_attributes( + ( + ("chunk_id", chunk.get("id", "")), + ("link", uri), + ("title", title), + ("structData", metadata.get("structData")), + ("pageSpan", chunk.get("pageSpan")), + ) + ), + }, + ) + + +def _document_result(hit: VertexSearchHit, score: float) -> VectorStoreSearchResult: + document: Final = hit.get("document", _EMPTY_DOCUMENT) + derived: Final = document.get("derivedStructData", _EMPTY_DERIVED_STRUCT_DATA) + link: Final = derived.get("link", "") + title: Final = derived.get("title", "") + document_id: Final = hit.get("id", "") + return VectorStoreSearchResult( + score=score, + content=[VectorStoreResultContent(text=_document_text(derived), type="text")], + file_id=link or document_id, + filename=title or _UNKNOWN_DOCUMENT, + attributes={ + "document_id": document_id, + **_non_empty_attributes( + ( + ("link", link), + ("title", title), + ("displayLink", derived.get("displayLink", "")), + ("formattedUrl", derived.get("formattedUrl", "")), + ("structData", document.get("structData")), + ) + ), + }, + ) + + +def _search_result(hit: VertexSearchHit, position: int) -> VectorStoreSearchResult: + score: Final = 1.0 / (position + 1) + chunk: Final = hit.get("chunk") + if chunk is not None: + return _chunk_result(chunk, score) + return _document_result(hit, score) + + class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase): """ Configuration for Vertex AI Search API Vector Store @@ -285,98 +411,19 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase): self, response: httpx.Response, litellm_logging_obj: LiteLLMLoggingObj ) -> VectorStoreSearchResponse: """ - Transform Vertex AI Search API response to standard vector store search response + Transform a Discovery Engine ``:search`` response into the standard vector store search response. - Handles the format from Discovery Engine Search API which returns: - { - "results": [ - { - "id": "...", - "document": { - "derivedStructData": { - "title": "...", - "link": "...", - "snippets": [...] - } - } - } - ] - } + Document hits (``results[].document``) take their text from ``derivedStructData`` in a fixed order: + ``extractive_segments``, then ``extractive_answers``, then ``snippets``, then ``title``; ``structData`` + and the link metadata land in ``attributes``. Chunk hits (``results[].chunk``, returned when the + caller sets ``contentSearchSpec.searchResultMode`` to ``CHUNKS`` via ``extra_body``) take their text + from ``chunk.content`` and their file id and name from ``chunk.documentMetadata``. """ try: response_json: Final = _vertex_search_payload(response) - - # Extract results from Vertex AI Search API response - results: Final = response_json.get("results", []) - - # Transform results to standard format - search_results: Final[list[VectorStoreSearchResult]] = [] - for result in results: - document: VertexSearchDocument = result.get("document", {}) - derived_data: VertexSearchDerivedStructData = document.get("derivedStructData", {}) - - # Extract text content from snippets - snippets = derived_data.get("snippets", []) - text_content = "" - - if snippets: - # Combine all snippets into one text - text_parts = [snippet.get("snippet", snippet.get("htmlSnippet", "")) for snippet in snippets] - text_content = " ".join(text_parts) - - # If no snippets, use title as fallback - if not text_content: - text_content = derived_data.get("title", "") - - content = [ - VectorStoreResultContent( - text=text_content, - type="text", - ) - ] - - # Extract file/document information - document_link = derived_data.get("link", "") - document_title = derived_data.get("title", "") - document_id = result.get("id", "") - - # Use link as file_id if available, otherwise use document ID - file_id = document_link if document_link else document_id - filename = document_title if document_title else "Unknown Document" - - # Build attributes with available metadata - attributes = { - "document_id": document_id, - } - - if document_link: - attributes["link"] = document_link - if document_title: - attributes["title"] = document_title - - # Add display link if available - display_link = derived_data.get("displayLink", "") - if display_link: - attributes["displayLink"] = display_link - - # Add formatted URL if available - formatted_url = derived_data.get("formattedUrl", "") - if formatted_url: - attributes["formattedUrl"] = formatted_url - - # Note: Search API doesn't provide explicit scores in the response - # You can use the position/rank as an implicit score - score = 1.0 / (float(search_results.__len__() + 1)) # Decreasing score based on position - - result_obj = VectorStoreSearchResult( - score=score, - content=content, - file_id=file_id, - filename=filename, - attributes=attributes, - ) - search_results.append(result_obj) - + search_results: Final = [ + _search_result(hit, position) for position, hit in enumerate(response_json.get("results", ())) + ] query_view: Final[_SearchQueryView] = {"query": litellm_logging_obj.model_call_details.get("query", "")} return VectorStoreSearchResponse( object="vector_store.search_results.page", diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_search_vector_store_transformation.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_search_vector_store_transformation.py index 034f85f5a0b..f3a276f7e46 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_search_vector_store_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_search_vector_store_transformation.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from types import SimpleNamespace import pytest @@ -6,6 +7,7 @@ from litellm.exceptions import BadRequestError from litellm.llms.vertex_ai.vector_stores.search_api.transformation import ( VertexSearchAPIVectorStoreConfig, ) +from litellm.types.vector_stores import VectorStoreSearchResponse def test_should_encode_vertex_search_vector_store_id_in_complete_url(): @@ -297,3 +299,188 @@ def test_search_request_logs_effective_query_when_extra_body_overrides_query(): assert body["query"] == "from-extra-body" assert log.model_call_details["query"] == "from-extra-body" + + +_CHUNK_NAME = ( + "projects/p/locations/global/collections/default_collection/dataStores/ds-1/" + "branches/0/documents/policy/chunks/c3" +) + + +def _search_response(payload: Mapping[str, object]) -> VectorStoreSearchResponse: + return VertexSearchAPIVectorStoreConfig().transform_search_vector_store_response( + response=SimpleNamespace(json=lambda: payload, status_code=200, headers={}), + litellm_logging_obj=SimpleNamespace(model_call_details={"query": "hello"}), + ) + + +def test_chunk_hit_uses_chunk_content_and_document_metadata(): + payload = { + "results": [ + { + "chunk": { + "id": "c3", + "name": _CHUNK_NAME, + "content": "Refunds are available within 14 days.", + "documentMetadata": { + "uri": "gs://bucket/policy.pdf", + "title": "Refund policy", + "structData": {"department": "billing"}, + }, + "pageSpan": {"pageStart": 2, "pageEnd": 2}, + "relevanceScore": 0.91, + } + } + ] + } + + result = _search_response(payload)["data"][0] + + assert result["content"] == [ + {"text": "Refunds are available within 14 days.", "type": "text"} + ] + assert result["score"] == 0.91 + assert result["file_id"] == "gs://bucket/policy.pdf" + assert result["filename"] == "Refund policy" + assert result["attributes"] == { + "document_id": "policy", + "chunk_id": "c3", + "link": "gs://bucket/policy.pdf", + "title": "Refund policy", + "structData": {"department": "billing"}, + "pageSpan": {"pageStart": 2, "pageEnd": 2}, + } + + +def test_chunk_hit_without_uri_or_title_falls_back_to_document_id(): + payload = { + "results": [ + { + "chunk": { + "id": "c1", + "name": _CHUNK_NAME.replace("policy/chunks/c3", "handbook/chunks/c1"), + "content": "Guest Services Handbook", + "documentMetadata": {"structData": {"title": "Handbook"}}, + } + } + ] + } + + result = _search_response(payload)["data"][0] + + assert result["content"] == [{"text": "Guest Services Handbook", "type": "text"}] + assert result["score"] == 1.0 + assert result["file_id"] == "handbook" + assert result["filename"] == "Unknown Document" + assert result["attributes"] == { + "document_id": "handbook", + "chunk_id": "c1", + "structData": {"title": "Handbook"}, + } + + +@pytest.mark.parametrize( + ("derived", "expected_text"), + [ + ( + { + "extractive_segments": [{"content": "seg one"}, {"content": "seg two"}], + "extractive_answers": [{"content": "ans"}], + "snippets": [{"snippet": "snip"}], + "title": "policy.pdf", + }, + "seg one\n\nseg two", + ), + ( + { + "extractive_answers": [{"content": "ans one"}, {"content": "ans two"}], + "snippets": [{"snippet": "snip"}], + "title": "policy.pdf", + }, + "ans one\n\nans two", + ), + ( + { + "snippets": [{"snippet": "snip a"}, {"htmlSnippet": "snip b"}], + "title": "policy.pdf", + }, + "snip a snip b", + ), + ( + {"extractive_segments": [{"pageNumber": "1"}, {"content": "seg", "pageNumber": "2"}]}, + "seg", + ), + ({"title": "policy.pdf"}, "policy.pdf"), + ], + ids=["segments", "answers", "snippets", "content_less_segment", "title"], +) +def test_document_hit_text_prefers_extractive_content(derived: Mapping[str, object], expected_text: str) -> None: + payload = {"results": [{"id": "doc-1", "document": {"derivedStructData": derived}}]} + + result = _search_response(payload)["data"][0] + + assert result["content"] == [{"text": expected_text, "type": "text"}] + + +def test_document_hit_surfaces_struct_data_in_attributes(): + payload = { + "results": [ + { + "id": "attr-1", + "document": { + "structData": {"title": "Thunder Loop", "waitMinutes": 45}, + "derivedStructData": {"clearbox_escorer_score": 0.5}, + }, + }, + {"id": "attr-2", "document": {"structData": {}, "derivedStructData": {}}}, + ] + } + + first, second = _search_response(payload)["data"] + + assert first["content"] == [{"text": "", "type": "text"}] + assert first["file_id"] == "attr-1" + assert first["filename"] == "Unknown Document" + assert first["attributes"] == { + "document_id": "attr-1", + "structData": {"title": "Thunder Loop", "waitMinutes": 45}, + } + assert second["attributes"] == {"document_id": "attr-2"} + + +def test_search_response_keeps_link_metadata_and_positional_scores(): + payload = { + "results": [ + { + "id": "doc-1", + "document": { + "derivedStructData": { + "title": "Terms", + "link": "gs://bucket/terms.pdf", + "displayLink": "bucket", + "formattedUrl": "https://bucket/terms.pdf", + "snippets": [{"snippet": "snip"}], + } + }, + }, + {"chunk": {"name": _CHUNK_NAME, "content": "chunk text"}}, + ] + } + + response = _search_response(payload) + + assert response["object"] == "vector_store.search_results.page" + assert response["search_query"] == "hello" + assert [result["score"] for result in response["data"]] == [1.0, 0.5] + assert response["data"][0]["file_id"] == "gs://bucket/terms.pdf" + assert response["data"][0]["attributes"] == { + "document_id": "doc-1", + "link": "gs://bucket/terms.pdf", + "title": "Terms", + "displayLink": "bucket", + "formattedUrl": "https://bucket/terms.pdf", + } + + +def test_search_response_without_results_key_is_empty(): + assert _search_response({"totalSize": 0})["data"] == []