mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-05 07:37:50 +00:00
fix(gemini): keep the grounding citation for a span that starts at byte 0
Google leaves startIndex out of a grounding segment when its value is 0, so the support for an answer's first sentence arrives with only endIndex and text. The annotation builder required startIndex and skipped that support, so grounded answers lost their first url_citation on unary and streamed responses alike. A missing startIndex now reads as 0, which is how long_decode already treats it
This commit is contained in:
parent
aa33b052d2
commit
a9de9fdf40
2 changed files with 66 additions and 1 deletions
|
|
@ -2736,7 +2736,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
# Process each grounding support to create annotations
|
||||
for support in grounding_supports:
|
||||
segment = support.get("segment", {})
|
||||
start_index = segment.get("startIndex")
|
||||
start_index = segment.get("startIndex", 0)
|
||||
end_index = segment.get("endIndex")
|
||||
|
||||
# Get the chunk indices for this support
|
||||
|
|
|
|||
|
|
@ -2417,6 +2417,71 @@ def test_gemini_with_grounding(respx_mock: respx.MockRouter) -> None:
|
|||
assert complete_response.usage.prompt_tokens_details.web_search_requests == 1
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("_gemini_httpx_transport")
|
||||
@pytest.mark.respx(assert_all_called=True)
|
||||
@pytest.mark.parametrize("stream", [False, True], ids=["unary", "stream"])
|
||||
def test_gemini_grounding_support_without_start_index_is_cited_from_zero(
|
||||
respx_mock: respx.MockRouter, stream: bool
|
||||
) -> None:
|
||||
body: Final = {
|
||||
"candidates": [
|
||||
{
|
||||
"content": {"role": "model", "parts": [{"text": "Spain won Euro 2024. The final was in Berlin."}]},
|
||||
"finishReason": "STOP",
|
||||
"index": 0,
|
||||
"groundingMetadata": {
|
||||
"groundingChunks": [
|
||||
{"web": {"uri": "https://example.com/uefa", "title": "uefa.com"}},
|
||||
{"web": {"uri": "https://example.com/news", "title": "news.com"}},
|
||||
],
|
||||
"groundingSupports": [
|
||||
{"segment": {"endIndex": 20, "text": "Spain won Euro 2024."}, "groundingChunkIndices": [0]},
|
||||
{
|
||||
"segment": {"startIndex": 21, "endIndex": 45, "text": "The final was in Berlin."},
|
||||
"groundingChunkIndices": [1],
|
||||
},
|
||||
],
|
||||
},
|
||||
}
|
||||
],
|
||||
"usageMetadata": {"promptTokenCount": 6, "candidatesTokenCount": 12, "totalTokenCount": 18},
|
||||
}
|
||||
method: Final = "streamGenerateContent?alt=sse" if stream else "generateContent"
|
||||
respx_mock.post(f"{GEMINI_MODELS_URL}/gemini-2.5-flash:{method}").mock(
|
||||
return_value=_gemini_sse(body) if stream else httpx.Response(200, json=body)
|
||||
)
|
||||
|
||||
response: Final = litellm.completion(
|
||||
model="gemini/gemini-2.5-flash",
|
||||
messages=[{"role": "user", "content": "Who won Euro 2024?"}],
|
||||
web_search_options={},
|
||||
stream=stream,
|
||||
api_key="test-api-key",
|
||||
)
|
||||
message: Final = (litellm.stream_chunk_builder(list(response)) if stream else response).choices[0].message
|
||||
|
||||
assert message.annotations == [
|
||||
{
|
||||
"type": "url_citation",
|
||||
"url_citation": {
|
||||
"start_index": 0,
|
||||
"end_index": 20,
|
||||
"url": "https://example.com/uefa",
|
||||
"title": "uefa.com",
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "url_citation",
|
||||
"url_citation": {
|
||||
"start_index": 21,
|
||||
"end_index": 45,
|
||||
"url": "https://example.com/news",
|
||||
"title": "news.com",
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("status_code", "expected_exception"),
|
||||
[
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue