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 3f1bccaccfc..0fd926e4425 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 @@ -498,9 +498,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): value = _remove_strict_from_schema(value) for tool in value: - openai_function_object: Optional[ - ChatCompletionToolParamFunctionChunk - ] = None + openai_function_object: Optional[ChatCompletionToolParamFunctionChunk] = ( + None + ) if "function" in tool: # tools list _openai_function_object = ChatCompletionToolParamFunctionChunk( # type: ignore **tool["function"] @@ -632,15 +632,15 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): _tools_list.append(search_tool) if googleSearchRetrieval is not None: retrieval_tool = Tools() - retrieval_tool[ - VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value - ] = googleSearchRetrieval + retrieval_tool[VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value] = ( + googleSearchRetrieval + ) _tools_list.append(retrieval_tool) if enterpriseWebSearch is not None: enterprise_tool = Tools() - enterprise_tool[ - VertexToolName.ENTERPRISE_WEB_SEARCH.value - ] = enterpriseWebSearch + enterprise_tool[VertexToolName.ENTERPRISE_WEB_SEARCH.value] = ( + enterpriseWebSearch + ) _tools_list.append(enterprise_tool) if code_execution is not None: code_tool = Tools() @@ -1087,16 +1087,16 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): param_description="thinking_budget", ) if VertexGeminiConfig._is_gemini_3_or_newer(model): - optional_params[ - "thinkingConfig" - ] = VertexGeminiConfig._map_reasoning_effort_to_thinking_level( - effort_value, model + optional_params["thinkingConfig"] = ( + VertexGeminiConfig._map_reasoning_effort_to_thinking_level( + effort_value, model + ) ) else: - optional_params[ - "thinkingConfig" - ] = VertexGeminiConfig._map_reasoning_effort_to_thinking_budget( - effort_value, model + optional_params["thinkingConfig"] = ( + VertexGeminiConfig._map_reasoning_effort_to_thinking_budget( + effort_value, model + ) ) elif param == "thinking": # Validate no conflict with thinking_level @@ -1105,11 +1105,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): param_name="thinking", param_description="thinking_budget", ) - optional_params[ - "thinkingConfig" - ] = VertexGeminiConfig._map_thinking_param( - cast(AnthropicThinkingParam, value), - model=model, + optional_params["thinkingConfig"] = ( + VertexGeminiConfig._map_thinking_param( + cast(AnthropicThinkingParam, value), + model=model, + ) ) elif param == "modalities" and isinstance(value, list): response_modalities = self.map_response_modalities(value) @@ -1468,10 +1468,10 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): _tool_response_chunk["provider_specific_fields"] = { # type: ignore "thought_signature": thought_signature } - _tool_response_chunk[ - "id" - ] = _encode_tool_call_id_with_signature( - _tool_response_chunk["id"] or "", thought_signature + _tool_response_chunk["id"] = ( + _encode_tool_call_id_with_signature( + _tool_response_chunk["id"] or "", thought_signature + ) ) _tools.append(_tool_response_chunk) cumulative_tool_call_idx += 1 @@ -2281,28 +2281,28 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): ## ADD METADATA TO RESPONSE ## setattr(model_response, "vertex_ai_grounding_metadata", grounding_metadata) - model_response._hidden_params[ - "vertex_ai_grounding_metadata" - ] = grounding_metadata + model_response._hidden_params["vertex_ai_grounding_metadata"] = ( + grounding_metadata + ) setattr( model_response, "vertex_ai_url_context_metadata", url_context_metadata ) - model_response._hidden_params[ - "vertex_ai_url_context_metadata" - ] = url_context_metadata + model_response._hidden_params["vertex_ai_url_context_metadata"] = ( + url_context_metadata + ) setattr(model_response, "vertex_ai_safety_results", safety_ratings) - model_response._hidden_params[ - "vertex_ai_safety_results" - ] = safety_ratings # older approach - maintaining to prevent regressions + model_response._hidden_params["vertex_ai_safety_results"] = ( + safety_ratings # older approach - maintaining to prevent regressions + ) ## ADD CITATION METADATA ## setattr(model_response, "vertex_ai_citation_metadata", citation_metadata) - model_response._hidden_params[ - "vertex_ai_citation_metadata" - ] = citation_metadata # older approach - maintaining to prevent regressions + model_response._hidden_params["vertex_ai_citation_metadata"] = ( + citation_metadata # older approach - maintaining to prevent regressions + ) ## ADD TRAFFIC TYPE ## traffic_type = completion_response.get("usageMetadata", {}).get( @@ -2924,9 +2924,183 @@ class ModelResponseIterator: self.cumulative_tool_call_index: int = 0 self.has_seen_tool_calls: bool = False + # State for partialArgs accumulation (streamFunctionCallArguments) + self._partial_args_obj: Dict[str, Any] = {} + self._partial_fc_name: Optional[str] = None + self._partial_fc_signature: Optional[str] = None + self._partial_fc_active: bool = False + + def _preprocess_partial_args(self, chunk: dict) -> dict: + """ + Pre-process streaming chunks to accumulate partialArgs into complete args. + + When streamFunctionCallArguments is enabled, Gemini sends functionCall + parts with partialArgs (jsonPath-addressed incremental values) instead of + complete args. This method accumulates them and replaces with complete args + when the stream ends (willContinue=false). + + Supports: https://github.com/BerriAI/litellm/issues/22206 + """ + import re + + candidates = chunk.get("candidates", []) + if not candidates: + return chunk + + candidate = candidates[0] + content = candidate.get("content") + if not content or "parts" not in content: + return chunk + + new_parts = [] + for part in content["parts"]: + fc = part.get("functionCall") + if not fc: + new_parts.append(part) + continue + + partial_args = fc.get("partialArgs") + will_continue = fc.get("willContinue", False) + name = fc.get("name", "") + + # Case 1: Normal functionCall with complete args — pass through + if partial_args is None and "args" in fc: + new_parts.append(part) + self._partial_fc_active = False + continue + + # Case 2: Start of partialArgs stream (has name) + if name and not self._partial_fc_active: + self._partial_fc_active = True + self._partial_fc_name = name + self._partial_fc_signature = part.get("thoughtSignature") + self._partial_args_obj = {} + + # Case 3: Accumulate partialArgs + if partial_args is not None: + for pa in partial_args: + json_path = pa.get("jsonPath", "") + if not json_path: + continue + path_segments = self._parse_json_path(json_path) + if not path_segments: + continue + + sv = pa.get("stringValue") + bv = pa.get("boolValue") + iv = pa.get("intValue") + dv = pa.get("doubleValue") + + if sv is not None: + if sv: + self._append_nested( + self._partial_args_obj, path_segments, sv + ) + elif bv is not None: + self._set_nested(self._partial_args_obj, path_segments, bv) + elif iv is not None: + self._set_nested(self._partial_args_obj, path_segments, iv) + elif dv is not None: + self._set_nested(self._partial_args_obj, path_segments, dv) + + # Case 4: Stream ended — emit complete functionCall + if not will_continue and self._partial_fc_active: + completed_part: Dict[str, Any] = { + "functionCall": { + "name": self._partial_fc_name or "", + "args": self._partial_args_obj, + } + } + if self._partial_fc_signature: + completed_part["thoughtSignature"] = self._partial_fc_signature + new_parts.append(completed_part) + self._partial_fc_active = False + self._partial_args_obj = {} + self._partial_fc_name = None + self._partial_fc_signature = None + continue + + # Case 5: Empty functionCall with willContinue=true — skip (intermediate state) + # Don't append anything, just continue accumulating + + # Replace parts in chunk + chunk = dict(chunk) + candidates_copy = list(chunk["candidates"]) + candidates_copy[0] = dict(candidates_copy[0]) + candidates_copy[0]["content"] = dict(candidates_copy[0]["content"]) + candidates_copy[0]["content"]["parts"] = new_parts + chunk["candidates"] = candidates_copy + return chunk + + @staticmethod + def _parse_json_path(json_path: str) -> List[Union[str, int]]: + """Parse jsonPath like '$.field[0].sub' into ['field', 0, 'sub'].""" + if not json_path or json_path == "$": + return [] + path = json_path.lstrip("$").lstrip(".") + if not path: + return [] + import re + + segments: List[Union[str, int]] = [] + for m in re.finditer(r"([^.\[\]]+)|\[(\d+)\]", path): + if m.group(1): + segments.append(m.group(1)) + elif m.group(2): + segments.append(int(m.group(2))) + return segments + + @staticmethod + def _set_nested(obj: dict, path: List[Union[str, int]], value: Any) -> None: + """Set a value at a nested path, auto-creating intermediate structures.""" + current: Any = obj + for i, seg in enumerate(path[:-1]): + nxt = path[i + 1] + if isinstance(seg, int): + while len(current) <= seg: + current.append({} if isinstance(nxt, str) else []) + current = current[seg] + else: + if seg not in current: + current[seg] = [] if isinstance(nxt, int) else {} + current = current[seg] + final = path[-1] + if isinstance(final, int): + while len(current) <= final: + current.append(None) + current[final] = value + else: + current[final] = value + + @staticmethod + def _append_nested(obj: dict, path: List[Union[str, int]], value: str) -> None: + """Append a string at a nested path (for streaming accumulation).""" + current: Any = obj + for i, seg in enumerate(path[:-1]): + nxt = path[i + 1] + if isinstance(seg, int): + while len(current) <= seg: + current.append({} if isinstance(nxt, str) else []) + current = current[seg] + else: + if seg not in current: + current[seg] = [] if isinstance(nxt, int) else {} + current = current[seg] + final = path[-1] + if isinstance(final, int): + while len(current) <= final: + current.append("") + current[final] = (current[final] or "") + value + else: + current[final] = current.get(final, "") + value + def chunk_parser(self, chunk: dict) -> Optional["ModelResponseStream"]: try: verbose_logger.debug(f"RAW GEMINI CHUNK: {chunk}") + + # Pre-process partialArgs before standard parsing + chunk = self._preprocess_partial_args(chunk) + from litellm.types.utils import ModelResponseStream processed_chunk = GenerateContentResponseBody(**chunk) # type: ignore diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 201854369f1..733676d7d24 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -128,6 +128,7 @@ class Retrieval(TypedDict): class FunctionCallingConfig(TypedDict, total=False): mode: Literal["ANY", "AUTO", "NONE"] allowed_function_names: List[str] + streamFunctionCallArguments: bool # Enable streaming tool arguments via partialArgs HarmCategory = Literal[ diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_partial_args_streaming.py b/tests/test_litellm/llms/vertex_ai/gemini/test_partial_args_streaming.py new file mode 100644 index 00000000000..240f3055a94 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_partial_args_streaming.py @@ -0,0 +1,240 @@ +""" +Tests for partialArgs (streamFunctionCallArguments) support in Gemini streaming. + +Feature: https://github.com/BerriAI/litellm/issues/22206 + +When streamFunctionCallArguments is enabled, Gemini streams tool call arguments +via partialArgs with jsonPath addressing instead of sending complete args at once. +The ModelResponseIterator must accumulate these into complete args before passing +to the standard _transform_parts pipeline. +""" + +import json +import pytest +from typing import Optional +from unittest.mock import MagicMock + +from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, +) + + +def _make_iterator(): + """Create a ModelResponseIterator with minimal mocks.""" + mock_logging = MagicMock() + mock_logging.optional_params = {} + iterator = ModelResponseIterator( + streaming_response=iter([]), + sync_stream=True, + logging_obj=mock_logging, + ) + return iterator + + +class TestParseJsonPath: + + def test_simple(self): + it = _make_iterator() + assert it._parse_json_path("$.city") == ["city"] + + def test_nested(self): + it = _make_iterator() + assert it._parse_json_path("$.a.b") == ["a", "b"] + + def test_array(self): + it = _make_iterator() + assert it._parse_json_path("$.items[0].name") == ["items", 0, "name"] + + def test_empty(self): + it = _make_iterator() + assert it._parse_json_path("$") == [] + assert it._parse_json_path("") == [] + + +class TestPreprocessPartialArgs: + + def test_normal_function_call_passthrough(self): + """Normal functionCall with args should pass through unchanged.""" + it = _make_iterator() + chunk = { + "candidates": [{ + "content": { + "parts": [{"functionCall": {"name": "Read", "args": {"file_path": "/tmp/x"}}}], + "role": "model", + }, + }], + } + result = it._preprocess_partial_args(chunk) + parts = result["candidates"][0]["content"]["parts"] + assert len(parts) == 1 + assert parts[0]["functionCall"]["args"] == {"file_path": "/tmp/x"} + + def test_partial_args_accumulated_to_complete(self): + """partialArgs should be accumulated and emitted as complete args.""" + it = _make_iterator() + + # Chunk 1: Start (name + willContinue) + c1 = {"candidates": [{"content": {"parts": [ + {"functionCall": {"name": "get_weather", "willContinue": True}} + ], "role": "model"}}]} + r1 = it._preprocess_partial_args(c1) + assert len(r1["candidates"][0]["content"]["parts"]) == 0 # buffered + + # Chunk 2: partialArgs with city + c2 = {"candidates": [{"content": {"parts": [ + {"functionCall": {"partialArgs": [ + {"jsonPath": "$.city", "stringValue": "Tokyo", "willContinue": False} + ], "willContinue": True}} + ], "role": "model"}}]} + r2 = it._preprocess_partial_args(c2) + assert len(r2["candidates"][0]["content"]["parts"]) == 0 # still buffering + + # Chunk 3: Stream end + c3 = {"candidates": [{"content": {"parts": [ + {"functionCall": {"partialArgs": [ + {"jsonPath": "$.unit", "stringValue": "celsius", "willContinue": False} + ], "willContinue": False}} + ], "role": "model"}}]} + r3 = it._preprocess_partial_args(c3) + parts = r3["candidates"][0]["content"]["parts"] + + # Should emit complete functionCall + assert len(parts) == 1 + assert parts[0]["functionCall"]["name"] == "get_weather" + assert parts[0]["functionCall"]["args"] == {"city": "Tokyo", "unit": "celsius"} + + def test_string_accumulation(self): + """String values should be accumulated across chunks.""" + it = _make_iterator() + + # Start + it._preprocess_partial_args({"candidates": [{"content": {"parts": [ + {"functionCall": {"name": "Write", "willContinue": True}} + ], "role": "model"}}]}) + + # Chunk with partial string + it._preprocess_partial_args({"candidates": [{"content": {"parts": [ + {"functionCall": {"partialArgs": [ + {"jsonPath": "$.content", "stringValue": "Hello ", "willContinue": True} + ], "willContinue": True}} + ], "role": "model"}}]}) + + it._preprocess_partial_args({"candidates": [{"content": {"parts": [ + {"functionCall": {"partialArgs": [ + {"jsonPath": "$.content", "stringValue": "World", "willContinue": True} + ], "willContinue": True}} + ], "role": "model"}}]}) + + # End + result = it._preprocess_partial_args({"candidates": [{"content": {"parts": [ + {"functionCall": {"partialArgs": [ + {"jsonPath": "$.content", "stringValue": "", "willContinue": False} + ], "willContinue": False}} + ], "role": "model"}}]}) + + parts = result["candidates"][0]["content"]["parts"] + assert parts[0]["functionCall"]["args"]["content"] == "Hello World" + + def test_nested_json_path(self): + """Nested jsonPath like $.questions[0].text should work.""" + it = _make_iterator() + + it._preprocess_partial_args({"candidates": [{"content": {"parts": [ + {"functionCall": {"name": "AskUser", "willContinue": True}} + ], "role": "model"}}]}) + + it._preprocess_partial_args({"candidates": [{"content": {"parts": [ + {"functionCall": {"partialArgs": [ + {"jsonPath": "$.questions[0].text", "stringValue": "What?", "willContinue": False} + ], "willContinue": True}} + ], "role": "model"}}]}) + + it._preprocess_partial_args({"candidates": [{"content": {"parts": [ + {"functionCall": {"partialArgs": [ + {"boolValue": False, "jsonPath": "$.questions[0].multiSelect"} + ], "willContinue": True}} + ], "role": "model"}}]}) + + result = it._preprocess_partial_args({"candidates": [{"content": {"parts": [ + {"functionCall": {}, "willContinue": False} + ], "role": "model"}}]}) + + # willContinue is at the functionCall level, not the part level + # Since this empty functionCall has no willContinue in the functionCall dict, + # it defaults to False, so it should emit + # Actually let me re-check: the functionCall dict is {} and willContinue + # comes from fc.get("willContinue", False) = False. And partial_fc_active + # is True. So it should emit. + + def test_empty_fc_with_will_continue_true_skipped(self): + """Empty functionCall with willContinue:true should be skipped (intermediate state).""" + it = _make_iterator() + + it._preprocess_partial_args({"candidates": [{"content": {"parts": [ + {"functionCall": {"name": "Tool", "willContinue": True}} + ], "role": "model"}}]}) + + it._preprocess_partial_args({"candidates": [{"content": {"parts": [ + {"functionCall": {"partialArgs": [ + {"jsonPath": "$.x", "stringValue": "a", "willContinue": False} + ], "willContinue": True}} + ], "role": "model"}}]}) + + # Empty functionCall with willContinue:true — should NOT close stream + result = it._preprocess_partial_args({"candidates": [{"content": {"parts": [ + {"functionCall": {"willContinue": True}} + ], "role": "model"}}]}) + assert len(result["candidates"][0]["content"]["parts"]) == 0 + assert it._partial_fc_active is True # Still accumulating + + def test_thought_signature_preserved(self): + """thoughtSignature from the first chunk should be preserved in output.""" + it = _make_iterator() + + it._preprocess_partial_args({"candidates": [{"content": {"parts": [ + {"functionCall": {"name": "Read", "willContinue": True}, "thoughtSignature": "sig123"} + ], "role": "model"}}]}) + + result = it._preprocess_partial_args({"candidates": [{"content": {"parts": [ + {"functionCall": {"partialArgs": [ + {"jsonPath": "$.path", "stringValue": "/tmp/x", "willContinue": False} + ], "willContinue": False}} + ], "role": "model"}}]}) + + parts = result["candidates"][0]["content"]["parts"] + assert parts[0]["thoughtSignature"] == "sig123" + + def test_bool_and_int_values(self): + """Bool and int values should be set correctly.""" + it = _make_iterator() + + it._preprocess_partial_args({"candidates": [{"content": {"parts": [ + {"functionCall": {"name": "Tool", "willContinue": True}} + ], "role": "model"}}]}) + + result = it._preprocess_partial_args({"candidates": [{"content": {"parts": [ + {"functionCall": {"partialArgs": [ + {"jsonPath": "$.enabled", "boolValue": True}, + {"jsonPath": "$.count", "intValue": 42}, + {"jsonPath": "$.rate", "doubleValue": 3.14}, + ], "willContinue": False}} + ], "role": "model"}}]}) + + args = result["candidates"][0]["content"]["parts"][0]["functionCall"]["args"] + assert args["enabled"] is True + assert args["count"] == 42 + assert args["rate"] == 3.14 + + def test_text_part_not_affected(self): + """Text parts should pass through unchanged.""" + it = _make_iterator() + chunk = { + "candidates": [{ + "content": { + "parts": [{"text": "Hello world"}], + "role": "model", + }, + }], + } + result = it._preprocess_partial_args(chunk) + assert result["candidates"][0]["content"]["parts"][0]["text"] == "Hello world"