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 9acbc9c459c..d124c62b94d 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 @@ -7,7 +7,8 @@ import threading import time from collections.abc import Awaitable, Callable, Mapping, Sequence from copy import deepcopy -from functools import lru_cache, partial +from functools import lru_cache, partial, reduce +from itertools import chain from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, Union, cast, get_args @@ -79,6 +80,7 @@ from litellm.types.utils import ( ChatCompletionAudioResponse, ChatCompletionTokenLogprob, ChoiceLogprobs, + Choices, CompletionTokensDetailsWrapper, PromptTokensDetailsWrapper, TopLogprob, @@ -233,6 +235,49 @@ def _served_model_name(model_version: object) -> str | None: return model_version.split("@", 1)[0] +def _chunk_field(value: object, name: str) -> object: + if isinstance(value, Mapping): + return cast(Mapping[str, object], value).get(name) # cast-ok: streamed chunk dicts are untyped JSON + return getattr(value, name, None) + + +def _delta_provider_fields(chunk: object) -> Mapping[str, object] | None: + choices: Final = _chunk_field(chunk, "choices") + if not isinstance(choices, list) or not choices: + return None + first_choice: Final = cast(Sequence[object], choices)[0] # cast-ok: streamed chunk choices are untyped + return as_str_mapping(_chunk_field(_chunk_field(first_choice, "delta"), "provider_specific_fields")) + + +def _list_field(fields: Mapping[str, object], key: str) -> Sequence[object]: + value: Final = fields.get(key) + return cast(Sequence[object], value) if isinstance(value, list) else () # cast-ok: items forwarded verbatim + + +def _paired_invocation(first: Mapping[str, object], second: Mapping[str, object]) -> Mapping[str, object]: + if ("response" in first) == ("response" in second): + return first + call, result = (first, second) if "response" in second else (second, first) + return { + **{key: value for key, value in call.items() if key not in ("response", "response_thought_signature")}, + "response": result["response"], + **( + {"response_thought_signature": result["response_thought_signature"]} + if "response_thought_signature" in result + else {} + ), + } + + +def _paired_invocations(invocations: Sequence[Mapping[str, object]]) -> tuple[Mapping[str, object], ...]: + ids: Final = tuple(dict.fromkeys(invocation.get("id") for invocation in invocations)) + return tuple( + reduce(_paired_invocation, (invocation for invocation in invocations if invocation.get("id") == invocation_id)) + for invocation_id in ids + if invocation_id is not None + ) + tuple(invocation for invocation in invocations if invocation.get("id") is None) + + _GEMINI_SPEECH_CONFIG_KEY_MAP: Final[Mapping[str, str]] = MappingProxyType( { "multi_speaker_voice_config": "multiSpeakerVoiceConfig", @@ -2706,6 +2751,26 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): if merged: setattr(response, field_name, merged) response.hidden_params[field_name] = merged + VertexGeminiConfig._merge_streamed_message_provider_fields(response, chunks) + + @staticmethod + def _merge_streamed_message_provider_fields(response: ModelResponse, chunks: Sequence[object]) -> None: + # a server-side toolCall and its toolResponse stream in different chunks, pair them as the unary path does + per_chunk: Final = tuple(fields for chunk in chunks if (fields := _delta_provider_fields(chunk)) is not None) + invocations: Final = tuple( + mapping + for invocation in chain.from_iterable( + _list_field(fields, "server_side_tool_invocations") for fields in per_chunk + ) + if (mapping := as_str_mapping(invocation)) is not None + ) + choice: Final = response.choices[0] if response.choices else None + if not invocations or not isinstance(choice, Choices): + return + choice.message.provider_specific_fields = { + **(choice.message.provider_specific_fields or {}), + "server_side_tool_invocations": list(_paired_invocations(invocations)), + } @staticmethod def _convert_grounding_metadata_to_annotations( diff --git a/tests/unit/llms/vertex_ai/gemini/test_context_circulation.py b/tests/unit/llms/vertex_ai/gemini/test_context_circulation.py index 39c40eeae67..70f4a0db036 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_context_circulation.py +++ b/tests/unit/llms/vertex_ai/gemini/test_context_circulation.py @@ -9,11 +9,16 @@ parts for server-side tools (e.g. Google Search). These must be: """ import json -from typing import Any, Dict, List +from collections.abc import Mapping +from typing import Any, Dict, Final, List from unittest.mock import MagicMock +import httpx import pytest +import litellm +from litellm import stream_chunk_builder +from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.llms.vertex_ai.gemini.transformation import ( gemini_convert_messages_with_history, ) @@ -21,6 +26,8 @@ from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, ) from litellm.types.llms.vertex_ai import HttpxPartType +from litellm.types.utils import Choices, ModelResponse +from litellm.utils import CustomStreamWrapper # --- Response extraction tests --- @@ -199,6 +206,102 @@ class TestExtractServerSideToolInvocations: assert result[0]["response_thought_signature"] == "sig_resp_1" +# --- Streaming tests --- + +_SEARCH_CALL: Final = { + "toolCall": {"toolType": "GOOGLE_SEARCH_WEB", "id": "srv_1", "args": {"queries": ["weather paris"]}}, + "thoughtSignature": "sig_call", +} +_SEARCH_RESPONSE: Final = { + "toolResponse": {"toolType": "GOOGLE_SEARCH_WEB", "id": "srv_1", "response": {"results": ["18C, sunny"]}}, + "thoughtSignature": "sig_resp", +} +_ANSWER_TEXT: Final = "It is 18C and sunny in Paris." +_QUESTION: Final = {"role": "user", "content": "Weather in Paris?"} + + +def _gemini_body(*parts: Mapping[str, object], finish: str | None = None) -> Mapping[str, object]: + return { + "candidates": [ + { + "content": {"role": "model", "parts": list(parts)}, + "index": 0, + **({"finishReason": finish} if finish is not None else {}), + } + ], + "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 8, "totalTokenCount": 18}, + } + + +def _gemini_client(response: httpx.Response) -> HTTPHandler: + return HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(lambda request: response))) + + +def _provider_fields(response: object) -> Mapping[str, object]: + assert isinstance(response, ModelResponse) + choice = response.choices[0] + assert isinstance(choice, Choices) + assert choice.message.provider_specific_fields is not None + return choice.message.provider_specific_fields + + +def _streamed_provider_fields() -> Mapping[str, object]: + sse = "".join( + f"data: {json.dumps(chunk)}\r\n\r\n" + for chunk in ( + _gemini_body(_SEARCH_CALL), + _gemini_body(_SEARCH_RESPONSE), + _gemini_body({"text": _ANSWER_TEXT}, finish="STOP"), + ) + ) + stream = litellm.completion( + model="gemini/gemini-3.8-flash", + messages=[_QUESTION], + api_key="fake-gemini-key", + stream=True, + client=_gemini_client( + httpx.Response(200, content=sse.encode(), headers={"content-type": "text/event-stream"}) + ), + ) + assert isinstance(stream, CustomStreamWrapper) + return _provider_fields(stream_chunk_builder(list(stream), messages=[_QUESTION])) + + +class TestStreamedServerSideToolInvocations: + """stream_chunk_builder must rebuild invocations split across chunks (#25869).""" + + def test_rebuilt_stream_matches_the_non_streaming_response(self): + unary = litellm.completion( + model="gemini/gemini-3.8-flash", + messages=[_QUESTION], + api_key="fake-gemini-key", + client=_gemini_client( + httpx.Response( + 200, + json=_gemini_body(_SEARCH_CALL, _SEARCH_RESPONSE, {"text": _ANSWER_TEXT}, finish="STOP"), + ) + ), + ) + + assert ( + _streamed_provider_fields()["server_side_tool_invocations"] + == _provider_fields(unary)["server_side_tool_invocations"] + ) + + def test_next_turn_sends_the_call_with_its_own_args_and_signature(self): + contents = gemini_convert_messages_with_history( + [ + _QUESTION, + {"role": "assistant", "content": _ANSWER_TEXT, "provider_specific_fields": _streamed_provider_fields()}, + {"role": "user", "content": "And tomorrow?"}, + ] + ) + + parts = [part for content in contents if content.get("role") == "model" for part in content["parts"]] + assert [part for part in parts if "toolCall" in part] == [_SEARCH_CALL] + assert [part.get("thoughtSignature") for part in parts if "toolResponse" in part] == ["sig_resp"] + + # --- Input re-injection tests ---