diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index b5f32d57061..7ee32260fe0 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -2217,6 +2217,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): continue image_response: list[ImageURLListItem] | None = None + content: str | None = None chat_completion_message: ChatCompletionResponseMessage = {"role": "assistant"} chat_completion_logprobs: ChoiceLogprobs | None = None tools: list[ChatCompletionToolCallChunk] | None = None @@ -2278,13 +2279,6 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): if reasoning_content is not None: chat_completion_message["reasoning_content"] = reasoning_content - if candidate_grounding_metadata: - annotations = VertexGeminiConfig._convert_grounding_metadata_to_annotations( - grounding_metadata=candidate_grounding_metadata, - content_text=content, - ) - if annotations: - chat_completion_message["annotations"] = annotations ( functions, tools, @@ -2295,6 +2289,14 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): is_function_call=is_function_call(standard_optional_params), ) + if candidate_grounding_metadata: + annotations = VertexGeminiConfig._convert_grounding_metadata_to_annotations( + grounding_metadata=candidate_grounding_metadata, + content_text=content, + ) + if annotations: + chat_completion_message["annotations"] = annotations + if "logprobsResult" in candidate: chat_completion_logprobs = VertexGeminiConfig._transform_logprobs( logprobs_result=candidate["logprobsResult"] diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index fd735afb16e..c01353a3568 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -16,6 +16,7 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.gemini.chat.transformation import GoogleAIStudioGeminiConfig from litellm.llms.vertex_ai.common_utils import VertexAIError from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, VertexGeminiConfig, ) from litellm.types.llms.vertex_ai import GeminiFinishReason, UsageMetadata @@ -6356,3 +6357,141 @@ def test_gemini_multi_candidate_messages_do_not_share_state(): assert resp.choices[1].message.tool_calls is None assert getattr(resp.choices[1].message, "reasoning_content", None) is None assert resp.choices[1].provider_specific_fields["native_finish_reason"] == "STOP" + + +_GROUNDING_METADATA: Final = { + "webSearchQueries": ["current price of gold"], + "groundingChunks": [ + {"web": {"uri": "https://vertexaisearch.cloud.google.com/grounding-api-redirect/AbC123", "title": "kitco"}}, + {"web": {"uri": "https://vertexaisearch.cloud.google.com/grounding-api-redirect/DeF456", "title": "reuters"}}, + ], + "groundingSupports": [ + {"segment": {"startIndex": 0, "endIndex": 42}, "groundingChunkIndices": [0]}, + {"segment": {"startIndex": 43, "endIndex": 80}, "groundingChunkIndices": [1]}, + ], +} +_GROUNDED_TEXT: Final = "As of today spot gold trades near $4,270" +_GROUNDING_URIS: Final = ( + "https://vertexaisearch.cloud.google.com/grounding-api-redirect/AbC123", + "https://vertexaisearch.cloud.google.com/grounding-api-redirect/DeF456", +) + + +def _gemini_stream_iterator() -> ModelResponseIterator: + logging_obj = MagicMock() + logging_obj.optional_params = {} + return ModelResponseIterator(streaming_response=iter([]), sync_stream=True, logging_obj=logging_obj) + + +def _delta_annotations(model_response: object) -> tuple[dict, ...]: + return tuple( + annotation + for choice in getattr(model_response, "choices", ()) + for annotation in (getattr(getattr(choice, "delta", None), "annotations", None) or ()) + ) + + +def _citation_urls(annotations: tuple[dict, ...]) -> tuple[str, ...]: + return tuple(annotation["url_citation"]["url"] for annotation in annotations) + + +def test_streaming_grounding_on_the_final_chunk_produces_annotations() -> None: + iterator = _gemini_stream_iterator() + iterator.chunk_parser({"candidates": [{"index": 0, "content": {"role": "model", "parts": [{"text": _GROUNDED_TEXT}]}}]}) + + final = iterator.chunk_parser( + {"candidates": [{"index": 0, "finishReason": "STOP", "groundingMetadata": _GROUNDING_METADATA}]} + ) + + assert _citation_urls(_delta_annotations(final)) == _GROUNDING_URIS + + +def test_streaming_grounding_alongside_text_produces_annotations() -> None: + iterator = _gemini_stream_iterator() + + chunk = iterator.chunk_parser( + { + "candidates": [ + { + "index": 0, + "content": {"role": "model", "parts": [{"text": _GROUNDED_TEXT}]}, + "groundingMetadata": _GROUNDING_METADATA, + "finishReason": "STOP", + } + ] + } + ) + + annotations = _delta_annotations(chunk) + assert _citation_urls(annotations) == _GROUNDING_URIS + assert annotations[0]["type"] == "url_citation" + assert annotations[0]["url_citation"]["start_index"] == 0 + + +def test_streaming_without_grounding_carries_no_annotations() -> None: + iterator = _gemini_stream_iterator() + + chunk = iterator.chunk_parser( + {"candidates": [{"index": 0, "content": {"role": "model", "parts": [{"text": "Paris."}]}, "finishReason": "STOP"}]} + ) + + assert _delta_annotations(chunk) == () + + +def test_streaming_grounding_without_web_chunks_carries_no_annotations() -> None: + iterator = _gemini_stream_iterator() + + chunk = iterator.chunk_parser( + { + "candidates": [ + { + "index": 0, + "finishReason": "STOP", + "groundingMetadata": { + "groundingChunks": [{"maps": {"placeId": "abc"}}], + "groundingSupports": [ + {"segment": {"startIndex": 0, "endIndex": 5}, "groundingChunkIndices": [0]} + ], + }, + } + ] + } + ) + + assert _delta_annotations(chunk) == () + + +def test_non_streaming_grounding_annotations_are_unchanged() -> None: + model_response = ModelResponse() + VertexGeminiConfig._process_candidates( + [ + { + "index": 0, + "content": {"role": "model", "parts": [{"text": _GROUNDED_TEXT}]}, + "groundingMetadata": _GROUNDING_METADATA, + "finishReason": "STOP", + } + ], + model_response, + {}, + ) + + annotations = getattr(model_response.choices[-1].message, "annotations", None) + assert _citation_urls(tuple(annotations or ())) == _GROUNDING_URIS + + +def test_streamed_grounding_survives_reassembly() -> None: + iterator = _gemini_stream_iterator() + chunks = [ + iterator.chunk_parser( + {"candidates": [{"index": 0, "content": {"role": "model", "parts": [{"text": _GROUNDED_TEXT}]}}]} + ), + iterator.chunk_parser( + {"candidates": [{"index": 0, "finishReason": "STOP", "groundingMetadata": _GROUNDING_METADATA}]} + ), + ] + + rebuilt = litellm.stream_chunk_builder(chunks=[chunk for chunk in chunks if chunk is not None]) + + assert rebuilt is not None + assert _citation_urls(tuple(getattr(rebuilt.choices[0].message, "annotations", None) or ())) == _GROUNDING_URIS