mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
fix(vertex_ai): return chunk content, extractive text, and structData from search_api vector store hits (#43100)
* fix(vertex_ai): return chunk content, extractive text, and structData from search_api vector store hits * fix(vertex_ai): report a chunk hit's relevanceScore as the search result score * test(vertex_ai): type the search response helper and parametrized case --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
57bead9842
commit
1a4a9c5ab3
2 changed files with 324 additions and 90 deletions
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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"] == []
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue