From d791df59b67ea487e4c7200dd3a627eb48f25dd8 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 8 Jun 2026 09:49:31 +0530 Subject: [PATCH 01/10] fix(vertex): propagate Vertex AI metadata in streaming success callbacks Streaming calls assembled via stream_chunk_builder were missing vertex_ai_grounding_metadata and vertex_ai_url_context_metadata in standard_logging_object.response. Merge metadata from chunks into the assembled response and mirror non-streaming hidden_params on Gemini chunks. Co-authored-by: Cursor --- .../streaming_chunk_builder_utils.py | 44 ++++++++++ .../vertex_and_google_ai_studio_gemini.py | 8 ++ litellm/main.py | 2 + .../test_litellm_logging.py | 35 ++++++++ .../test_streaming_chunk_builder_utils.py | 82 +++++++++++++++++++ 5 files changed, 171 insertions(+) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index fe7c62c38425..99e610d69ad0 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -32,6 +32,14 @@ ) +VERTEX_AI_STREAM_METADATA_FIELDS = ( + "vertex_ai_grounding_metadata", + "vertex_ai_url_context_metadata", + "vertex_ai_safety_ratings", + "vertex_ai_citation_metadata", +) + + class ChunkProcessor: def __init__(self, chunks: List, messages: Optional[list] = None): self.chunks = self._sort_chunks(chunks) @@ -79,6 +87,42 @@ def update_model_response_with_hidden_params( model_response._hidden_params = chunk.get("_hidden_params", {}) return model_response + @staticmethod + def _get_chunk_attr(chunk: Any, field_name: str) -> Any: + if isinstance(chunk, dict): + value = chunk.get(field_name) + if value is not None: + return value + model_extra = chunk.get("model_extra") + if isinstance(model_extra, dict): + return model_extra.get(field_name) + return None + return getattr(chunk, field_name, None) + + @staticmethod + def propagate_vertex_ai_metadata_from_chunks( + response: ModelResponse, chunks: List[Any] + ) -> None: + """ + Merge Vertex AI metadata from streaming chunks into the assembled response. + + Gemini/Vertex streaming sets these fields on individual chunks but + stream_chunk_builder must propagate them for logging callbacks. + """ + for field_name in VERTEX_AI_STREAM_METADATA_FIELDS: + merged: List[Any] = [] + for chunk in chunks: + value = ChunkProcessor._get_chunk_attr(chunk, field_name) + if not value: + continue + if isinstance(value, list): + merged.extend(value) + else: + merged.append(value) + if merged: + setattr(response, field_name, merged) + response._hidden_params[field_name] = merged + @staticmethod def _get_chunk_id(chunks: List[Dict[str, Any]]) -> str: """ 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 5cd02293f142..d6436f744c47 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 @@ -3386,9 +3386,17 @@ def _apply_stream_candidates( choice.finish_reason = "tool_calls" setattr(model_response, "vertex_ai_grounding_metadata", grounding_metadata) # type: ignore + model_response._hidden_params["vertex_ai_grounding_metadata"] = ( + grounding_metadata + ) setattr(model_response, "vertex_ai_url_context_metadata", url_context_metadata) # type: ignore + model_response._hidden_params["vertex_ai_url_context_metadata"] = ( + url_context_metadata + ) setattr(model_response, "vertex_ai_safety_ratings", safety_ratings) # type: ignore + model_response._hidden_params["vertex_ai_safety_ratings"] = safety_ratings setattr(model_response, "vertex_ai_citation_metadata", citation_metadata) # type: ignore + model_response._hidden_params["vertex_ai_citation_metadata"] = citation_metadata return ( grounding_metadata, diff --git a/litellm/main.py b/litellm/main.py index 64891e2def95..c7f0b4a2091c 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -7761,6 +7761,7 @@ def stream_chunk_builder( # noqa: PLR0915 "cost", logging_obj._response_cost_calculator(result=response), ) + processor.propagate_vertex_ai_metadata_from_chunks(response, chunks) return response tool_call_chunks = [ @@ -7940,6 +7941,7 @@ def stream_chunk_builder( # noqa: PLR0915 usage, "cost", logging_obj._response_cost_calculator(result=response) ) + processor.propagate_vertex_ai_metadata_from_chunks(response, chunks) return response except Exception as e: verbose_logger.exception( diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index d57d8dafdbdc..34edd6eccf30 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -2165,6 +2165,41 @@ def test_get_assembled_streaming_response_returns_result_for_streaming(): assert assembled is result +def test_streaming_success_handler_includes_vertex_ai_metadata_in_standard_logging(): + """Assembled streaming responses should include Vertex AI metadata in logging payload.""" + import datetime + + from litellm.types.utils import Choices, Message + + logging_obj = _make_logging_obj(stream=True) + grounding_metadata = [{"webSearchQueries": ["weather in SF"]}] + url_context_metadata = [{"urlMetadata": [{"retrievedUrl": "https://example.com"}]}] + result = ModelResponse( + id="resp-1", + choices=[ + Choices( + index=0, + message=Message(role="assistant", content="hello"), + finish_reason="stop", + ) + ], + model="gemini-2.5-flash", + ) + setattr(result, "vertex_ai_grounding_metadata", grounding_metadata) + setattr(result, "vertex_ai_url_context_metadata", url_context_metadata) + result._hidden_params["vertex_ai_grounding_metadata"] = grounding_metadata + result._hidden_params["vertex_ai_url_context_metadata"] = url_context_metadata + + start = datetime.datetime.now() + end = datetime.datetime.now() + logging_obj.success_handler(result=result, start_time=start, end_time=end) + + payload = logging_obj.model_call_details.get("standard_logging_object") + assert payload is not None + assert payload["response"]["vertex_ai_grounding_metadata"] == grounding_metadata + assert payload["response"]["vertex_ai_url_context_metadata"] == url_context_metadata + + def test_get_assembled_streaming_response_returns_none_for_non_streaming_text_completion(): """Non-streaming TextCompletionResponse should also return None.""" import datetime diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py index e40a0817fd94..22383deb5697 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -613,3 +613,85 @@ def test_stream_chunk_builder_dict_snapshot_preserves_hidden_provider_fields(): assert ( response._hidden_params["provider_specific_fields"]["traffic_type"] == "default" ) + + +def test_stream_chunk_builder_propagates_vertex_ai_metadata_from_chunks(): + """Vertex AI metadata on streaming chunks must appear on assembled response.""" + grounding_metadata = [{"webSearchQueries": ["weather in SF"]}] + url_context_metadata = [{"urlMetadata": [{"retrievedUrl": "https://example.com"}]}] + + chunk1 = ModelResponseStream( + id="chatcmpl-vertex-1", + created=1, + model="gemini-2.5-flash", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(content="The weather", role="assistant"), + ) + ], + ) + setattr(chunk1, "vertex_ai_grounding_metadata", grounding_metadata) + chunk1._hidden_params["vertex_ai_grounding_metadata"] = grounding_metadata + + chunk2 = ModelResponseStream( + id="chatcmpl-vertex-1", + created=1, + model="gemini-2.5-flash", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(content=" is sunny.", role="assistant"), + ) + ], + ) + setattr(chunk2, "vertex_ai_url_context_metadata", url_context_metadata) + chunk2._hidden_params["vertex_ai_url_context_metadata"] = url_context_metadata + + response = stream_chunk_builder(chunks=[chunk1, chunk2]) + assert response is not None + assert getattr(response, "vertex_ai_grounding_metadata") == grounding_metadata + assert getattr(response, "vertex_ai_url_context_metadata") == url_context_metadata + assert response._hidden_params["vertex_ai_grounding_metadata"] == grounding_metadata + assert ( + response._hidden_params["vertex_ai_url_context_metadata"] + == url_context_metadata + ) + + dumped = response.model_dump() + assert dumped["vertex_ai_grounding_metadata"] == grounding_metadata + assert dumped["vertex_ai_url_context_metadata"] == url_context_metadata + + +def test_stream_chunk_builder_propagates_vertex_ai_metadata_from_dict_chunks(): + """Dict snapshot chunks (model_dump) should also propagate Vertex AI metadata.""" + chunk_dict = ModelResponseStream( + id="chatcmpl-vertex-2", + created=1, + model="gemini-2.5-flash", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(content="hello", role="assistant"), + ) + ], + ).model_dump() + chunk_dict["vertex_ai_grounding_metadata"] = [{"webSearchQueries": ["test query"]}] + chunk_dict["_hidden_params"] = { + "vertex_ai_grounding_metadata": [{"webSearchQueries": ["test query"]}] + } + + response = stream_chunk_builder(chunks=[chunk_dict]) + assert response is not None + assert getattr(response, "vertex_ai_grounding_metadata") == [ + {"webSearchQueries": ["test query"]} + ] + assert response.model_dump()["vertex_ai_grounding_metadata"] == [ + {"webSearchQueries": ["test query"]} + ] From 1add616c3ec48974173a9fe4006e4fb1334a49c5 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 8 Jun 2026 10:06:59 +0530 Subject: [PATCH 02/10] refactor(vertex): move streaming metadata merge into provider config hook Address review feedback by delegating assembled-stream metadata propagation to VertexGeminiConfig via BaseConfig.apply_assembled_streaming_response_metadata, and only write chunk hidden_params when metadata is non-empty. Co-authored-by: Cursor --- .../streaming_chunk_builder_utils.py | 85 ++++++++++--------- litellm/llms/base_llm/chat/transformation.py | 8 ++ .../vertex_and_google_ai_studio_gemini.py | 82 +++++++++++++++--- litellm/main.py | 8 +- 4 files changed, 130 insertions(+), 53 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 99e610d69ad0..f1de30fa72d6 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -32,14 +32,6 @@ ) -VERTEX_AI_STREAM_METADATA_FIELDS = ( - "vertex_ai_grounding_metadata", - "vertex_ai_url_context_metadata", - "vertex_ai_safety_ratings", - "vertex_ai_citation_metadata", -) - - class ChunkProcessor: def __init__(self, chunks: List, messages: Optional[list] = None): self.chunks = self._sort_chunks(chunks) @@ -88,40 +80,53 @@ def update_model_response_with_hidden_params( return model_response @staticmethod - def _get_chunk_attr(chunk: Any, field_name: str) -> Any: - if isinstance(chunk, dict): - value = chunk.get(field_name) - if value is not None: - return value - model_extra = chunk.get("model_extra") - if isinstance(model_extra, dict): - return model_extra.get(field_name) - return None - return getattr(chunk, field_name, None) - - @staticmethod - def propagate_vertex_ai_metadata_from_chunks( - response: ModelResponse, chunks: List[Any] + def apply_provider_assembled_streaming_metadata( + response: ModelResponse, + chunks: List[Any], + logging_obj: Optional[Any] = None, ) -> None: - """ - Merge Vertex AI metadata from streaming chunks into the assembled response. + if not chunks: + return - Gemini/Vertex streaming sets these fields on individual chunks but - stream_chunk_builder must propagate them for logging callbacks. - """ - for field_name in VERTEX_AI_STREAM_METADATA_FIELDS: - merged: List[Any] = [] - for chunk in chunks: - value = ChunkProcessor._get_chunk_attr(chunk, field_name) - if not value: - continue - if isinstance(value, list): - merged.extend(value) - else: - merged.append(value) - if merged: - setattr(response, field_name, merged) - response._hidden_params[field_name] = merged + first_chunk = chunks[0] + model = ( + first_chunk.get("model") + if isinstance(first_chunk, dict) + else getattr(first_chunk, "model", None) + ) + if not model: + return + + custom_llm_provider = None + if logging_obj is not None: + custom_llm_provider = logging_obj.model_call_details.get( + "custom_llm_provider" + ) + + try: + from litellm.litellm_core_utils.get_llm_provider_logic import ( + get_llm_provider, + ) + from litellm.types.utils import LlmProviders + from litellm.utils import ProviderConfigManager + + if custom_llm_provider: + provider = LlmProviders(custom_llm_provider) + else: + _, provider_str, _, _ = get_llm_provider(model) + provider = LlmProviders(provider_str) + + provider_config = ProviderConfigManager.get_provider_chat_config( + model=model, + provider=provider, + ) + if provider_config is not None: + provider_config.apply_assembled_streaming_response_metadata( + response=response, + chunks=chunks, + ) + except Exception: + return @staticmethod def _get_chunk_id(chunks: List[Dict[str, Any]]) -> str: diff --git a/litellm/llms/base_llm/chat/transformation.py b/litellm/llms/base_llm/chat/transformation.py index 5f35a58ce1f1..8f9d5cad7c41 100644 --- a/litellm/llms/base_llm/chat/transformation.py +++ b/litellm/llms/base_llm/chat/transformation.py @@ -442,6 +442,14 @@ def post_stream_processing(self, stream: Any) -> Any: """Hook for providers to post-process streaming responses. Default: pass-through.""" return stream + def apply_assembled_streaming_response_metadata( + self, + response: "ModelResponse", + chunks: List[Any], + ) -> None: + """Hook for providers to merge chunk metadata into assembled streaming responses.""" + return None + def calculate_additional_costs( self, model: str, prompt_tokens: int, completion_tokens: int ) -> Optional[dict]: 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 d6436f744c47..0a7ba70dacc1 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 @@ -2253,6 +2253,71 @@ def _extract_candidate_metadata( citation_metadata, ) + _STREAM_METADATA_FIELDS = ( + "vertex_ai_grounding_metadata", + "vertex_ai_url_context_metadata", + "vertex_ai_safety_ratings", + "vertex_ai_citation_metadata", + ) + + @staticmethod + def _get_stream_chunk_attr(chunk: Any, field_name: str) -> Any: + if isinstance(chunk, dict): + value = chunk.get(field_name) + if value is not None: + return value + model_extra = chunk.get("model_extra") + if isinstance(model_extra, dict): + return model_extra.get(field_name) + return None + return getattr(chunk, field_name, None) + + @staticmethod + def _set_stream_metadata_on_response( + model_response: Any, + grounding_metadata: List[dict], + url_context_metadata: List[dict], + safety_ratings: List[dict], + citation_metadata: List[dict], + ) -> None: + setattr(model_response, "vertex_ai_grounding_metadata", grounding_metadata) # type: ignore + if grounding_metadata: + model_response._hidden_params["vertex_ai_grounding_metadata"] = ( + grounding_metadata + ) + setattr(model_response, "vertex_ai_url_context_metadata", url_context_metadata) # type: ignore + if url_context_metadata: + model_response._hidden_params["vertex_ai_url_context_metadata"] = ( + url_context_metadata + ) + setattr(model_response, "vertex_ai_safety_ratings", safety_ratings) # type: ignore + if safety_ratings: + model_response._hidden_params["vertex_ai_safety_ratings"] = safety_ratings + setattr(model_response, "vertex_ai_citation_metadata", citation_metadata) # type: ignore + if citation_metadata: + model_response._hidden_params["vertex_ai_citation_metadata"] = ( + citation_metadata + ) + + def apply_assembled_streaming_response_metadata( + self, + response: ModelResponse, + chunks: List[Any], + ) -> None: + for field_name in VertexGeminiConfig._STREAM_METADATA_FIELDS: + merged: List[Any] = [] + for chunk in chunks: + value = VertexGeminiConfig._get_stream_chunk_attr(chunk, field_name) + if not value: + continue + if isinstance(value, list): + merged.extend(value) + else: + merged.append(value) + if merged: + setattr(response, field_name, merged) + response._hidden_params[field_name] = merged + @staticmethod def _convert_grounding_metadata_to_annotations( grounding_metadata: List[dict], @@ -3385,18 +3450,13 @@ def _apply_stream_candidates( if choice.finish_reason == "stop": choice.finish_reason = "tool_calls" - setattr(model_response, "vertex_ai_grounding_metadata", grounding_metadata) # type: ignore - model_response._hidden_params["vertex_ai_grounding_metadata"] = ( - grounding_metadata - ) - setattr(model_response, "vertex_ai_url_context_metadata", url_context_metadata) # type: ignore - model_response._hidden_params["vertex_ai_url_context_metadata"] = ( - url_context_metadata + VertexGeminiConfig._set_stream_metadata_on_response( + model_response, + grounding_metadata, + url_context_metadata, + safety_ratings, + citation_metadata, ) - setattr(model_response, "vertex_ai_safety_ratings", safety_ratings) # type: ignore - model_response._hidden_params["vertex_ai_safety_ratings"] = safety_ratings - setattr(model_response, "vertex_ai_citation_metadata", citation_metadata) # type: ignore - model_response._hidden_params["vertex_ai_citation_metadata"] = citation_metadata return ( grounding_metadata, diff --git a/litellm/main.py b/litellm/main.py index c7f0b4a2091c..1a0d0312d73a 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -7761,7 +7761,9 @@ def stream_chunk_builder( # noqa: PLR0915 "cost", logging_obj._response_cost_calculator(result=response), ) - processor.propagate_vertex_ai_metadata_from_chunks(response, chunks) + processor.apply_provider_assembled_streaming_metadata( + response, chunks, logging_obj + ) return response tool_call_chunks = [ @@ -7941,7 +7943,9 @@ def stream_chunk_builder( # noqa: PLR0915 usage, "cost", logging_obj._response_cost_calculator(result=response) ) - processor.propagate_vertex_ai_metadata_from_chunks(response, chunks) + processor.apply_provider_assembled_streaming_metadata( + response, chunks, logging_obj + ) return response except Exception as e: verbose_logger.exception( From 9960d1c6b9eb7cea0779b6ea88344956a0bfe448 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 8 Jun 2026 12:29:20 +0530 Subject: [PATCH 03/10] fix(redaction): scrub Vertex provider metadata when message logging is off Clear vertex_ai_grounding_metadata and related fields from standard logging responses and assembled streaming ModelResponse objects so turn_off_message_logging cannot leak prompt-derived web search queries. Co-authored-by: Cursor --- litellm/litellm_core_utils/redact_messages.py | 33 ++++++++++ .../test_redact_messages.py | 63 +++++++++++++++++++ 2 files changed, 96 insertions(+) diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index dbc9cabdc7a6..f552f89f1591 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -29,6 +29,14 @@ else: LiteLLMLoggingObject = Any +VERTEX_PROVIDER_METADATA_FIELDS = ( + "vertex_ai_grounding_metadata", + "vertex_ai_url_context_metadata", + "vertex_ai_safety_ratings", + "vertex_ai_safety_results", + "vertex_ai_citation_metadata", +) + def redact_message_input_output_from_custom_logger( litellm_logging_obj: LiteLLMLoggingObject, result, custom_logger: CustomLogger @@ -100,6 +108,26 @@ def _redact_responses_api_output_dict(output_items, redacted_str: str): summary_item["text"] = redacted_str +def _redact_vertex_provider_metadata(obj: Any) -> None: + if isinstance(obj, dict): + for field in VERTEX_PROVIDER_METADATA_FIELDS: + if field in obj: + obj[field] = [] + hidden_params = obj.get("_hidden_params") + if isinstance(hidden_params, dict): + for field in VERTEX_PROVIDER_METADATA_FIELDS: + hidden_params.pop(field, None) + return + + for field in VERTEX_PROVIDER_METADATA_FIELDS: + if hasattr(obj, field): + setattr(obj, field, []) + hidden_params = getattr(obj, "_hidden_params", None) + if isinstance(hidden_params, dict): + for field in VERTEX_PROVIDER_METADATA_FIELDS: + hidden_params.pop(field, None) + + def _redact_standard_logging_object(model_call_details: dict): """Redact messages and response inside standard_logging_object if present.""" standard_logging_object = model_call_details.get("standard_logging_object") @@ -119,10 +147,12 @@ def _redact_standard_logging_object(model_call_details: dict): # ResponsesAPIResponse format - redact content in output items if isinstance(response.get("output"), list): _redact_responses_api_output_dict(response["output"], redacted_str) + _redact_vertex_provider_metadata(response) elif isinstance(response, dict) and "choices" in response: # ModelResponse dict format - redact content in choices if isinstance(response.get("choices"), list): _redact_model_response_dict_choices(response["choices"], redacted_str) + _redact_vertex_provider_metadata(response) elif isinstance(response, str): standard_logging_object["response"] = redacted_str else: @@ -174,6 +204,7 @@ def perform_redaction(model_call_details: dict, result): if hasattr(_streaming_response, "choices"): for choice in _streaming_response.choices: _redact_choice_content(choice) + _redact_vertex_provider_metadata(_streaming_response) elif hasattr(_streaming_response, "output"): _redact_responses_api_output(_streaming_response.output) # Redact reasoning field in ResponsesAPIResponse @@ -200,12 +231,14 @@ def perform_redaction(model_call_details: dict, result): if hasattr(_result, "choices") and _result.choices is not None: for choice in _result.choices: _redact_choice_content(choice) + _redact_vertex_provider_metadata(_result) elif isinstance(_result, dict) and "choices" in _result: # Handle dict representation of ModelResponse (e.g., from model_dump()) if _result.get("choices") is not None: _redact_model_response_dict_choices( _result["choices"], "redacted-by-litellm" ) + _redact_vertex_provider_metadata(_result) elif isinstance(_result, dict) and "output" in _result: if isinstance(_result.get("output"), list): _redact_responses_api_output_dict( diff --git a/tests/test_litellm/litellm_core_utils/test_redact_messages.py b/tests/test_litellm/litellm_core_utils/test_redact_messages.py index 60cfff6e4a05..c2ceb570688a 100644 --- a/tests/test_litellm/litellm_core_utils/test_redact_messages.py +++ b/tests/test_litellm/litellm_core_utils/test_redact_messages.py @@ -349,3 +349,66 @@ def test_redacts_responses_api_response_object(self): assert redacted.output[0].content[0].text == "redacted-by-litellm" assert response.output[0].content[0].text == "sensitive output" + + def test_redacts_vertex_provider_metadata_in_standard_logging_response(self): + details = { + "standard_logging_object": { + "messages": [{"role": "user", "content": "sensitive prompt"}], + "response": { + "choices": [ + { + "message": { + "content": "sensitive answer", + "role": "assistant", + } + } + ], + "vertex_ai_grounding_metadata": [ + {"webSearchQueries": ["sensitive search term"]} + ], + "vertex_ai_url_context_metadata": [ + {"urlMetadata": [{"retrievedUrl": "https://example.com"}]} + ], + }, + } + } + + perform_redaction(details, None) + + response = details["standard_logging_object"]["response"] + assert response["choices"][0]["message"]["content"] == "redacted-by-litellm" + assert response["vertex_ai_grounding_metadata"] == [] + assert response["vertex_ai_url_context_metadata"] == [] + + def test_redacts_vertex_provider_metadata_on_streaming_model_response(self): + response = litellm.ModelResponse( + id="resp-1", + choices=[ + litellm.Choices( + message=litellm.Message( + content="sensitive answer", + role="assistant", + ) + ) + ], + model="gemini-2.5-flash", + ) + setattr( + response, + "vertex_ai_grounding_metadata", + [{"webSearchQueries": ["sensitive search term"]}], + ) + response._hidden_params["vertex_ai_grounding_metadata"] = [ + {"webSearchQueries": ["sensitive search term"]} + ] + + details = { + "stream": True, + "complete_streaming_response": response, + } + + perform_redaction(details, response) + + assert response.choices[0].message.content == "redacted-by-litellm" + assert getattr(response, "vertex_ai_grounding_metadata") == [] + assert "vertex_ai_grounding_metadata" not in response._hidden_params From 098cca6d5dc6f7079423613e3dfe66363e65141c Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Mon, 8 Jun 2026 07:40:10 +0000 Subject: [PATCH 04/10] Use assembled model for streaming metadata hook --- .../streaming_chunk_builder_utils.py | 7 +--- .../test_streaming_chunk_builder_utils.py | 38 +++++++++++++++++++ 2 files changed, 39 insertions(+), 6 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index f1de30fa72d6..2cfe7e7af5f7 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -88,12 +88,7 @@ def apply_provider_assembled_streaming_metadata( if not chunks: return - first_chunk = chunks[0] - model = ( - first_chunk.get("model") - if isinstance(first_chunk, dict) - else getattr(first_chunk, "model", None) - ) + model = getattr(response, "model", None) if not model: return diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py index 22383deb5697..85661a9ecb9a 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -667,6 +667,44 @@ def test_stream_chunk_builder_propagates_vertex_ai_metadata_from_chunks(): assert dumped["vertex_ai_url_context_metadata"] == url_context_metadata +def test_stream_chunk_builder_uses_assembled_model_for_provider_metadata(): + grounding_metadata = [{"webSearchQueries": ["weather in SF"]}] + + chunk1 = ModelResponseStream( + id="chatcmpl-vertex-router", + created=1, + model="gpt-4o", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(content="The weather", role="assistant"), + ) + ], + ) + chunk2 = ModelResponseStream( + id="chatcmpl-vertex-router", + created=1, + model="gemini-2.5-flash", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(content=" is sunny.", role=None), + ) + ], + ) + setattr(chunk2, "vertex_ai_grounding_metadata", grounding_metadata) + chunk2._hidden_params["vertex_ai_grounding_metadata"] = grounding_metadata + + response = stream_chunk_builder(chunks=[chunk1, chunk2]) + assert response is not None + assert response.model == "gemini-2.5-flash" + assert getattr(response, "vertex_ai_grounding_metadata") == grounding_metadata + + def test_stream_chunk_builder_propagates_vertex_ai_metadata_from_dict_chunks(): """Dict snapshot chunks (model_dump) should also propagate Vertex AI metadata.""" chunk_dict = ModelResponseStream( From 5ffea2f1d710f17b0f65666d0aede806a03b1d6a Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 8 Jun 2026 18:24:02 +0530 Subject: [PATCH 05/10] Fix Vertex metadata redaction bypass in logging callbacks. Scrub Vertex provider fields from litellm_params.metadata.hidden_params during perform_redaction so streaming success_handler merges do not leak prompt-derived metadata when message logging is disabled. Co-authored-by: Cursor --- litellm/litellm_core_utils/redact_messages.py | 25 ++++++++++++++++ .../test_redact_messages.py | 30 +++++++++++++++++++ 2 files changed, 55 insertions(+) diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index f552f89f1591..dbaf6b594652 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -128,6 +128,30 @@ def _redact_vertex_provider_metadata(obj: Any) -> None: hidden_params.pop(field, None) +def _redact_vertex_provider_metadata_from_litellm_params( + model_call_details: dict, +) -> None: + """ + Scrub Vertex provider metadata copied into litellm_params metadata. + + success_handler() merges response._hidden_params into + litellm_params.metadata['hidden_params'] before perform_redaction() runs. + """ + litellm_params = model_call_details.get("litellm_params") + if not isinstance(litellm_params, dict): + return + + for metadata_key in ("metadata", "litellm_metadata"): + metadata = litellm_params.get(metadata_key) + if not isinstance(metadata, dict): + continue + hidden_params = metadata.get("hidden_params") + if not isinstance(hidden_params, dict): + continue + for field in VERTEX_PROVIDER_METADATA_FIELDS: + hidden_params.pop(field, None) + + def _redact_standard_logging_object(model_call_details: dict): """Redact messages and response inside standard_logging_object if present.""" standard_logging_object = model_call_details.get("standard_logging_object") @@ -194,6 +218,7 @@ def perform_redaction(model_call_details: dict, result): model_call_details["prompt"] = "" model_call_details["input"] = "" _redact_standard_logging_object(model_call_details) + _redact_vertex_provider_metadata_from_litellm_params(model_call_details) # Redact streaming response if ( diff --git a/tests/test_litellm/litellm_core_utils/test_redact_messages.py b/tests/test_litellm/litellm_core_utils/test_redact_messages.py index c2ceb570688a..36f220f9a2cc 100644 --- a/tests/test_litellm/litellm_core_utils/test_redact_messages.py +++ b/tests/test_litellm/litellm_core_utils/test_redact_messages.py @@ -412,3 +412,33 @@ def test_redacts_vertex_provider_metadata_on_streaming_model_response(self): assert response.choices[0].message.content == "redacted-by-litellm" assert getattr(response, "vertex_ai_grounding_metadata") == [] assert "vertex_ai_grounding_metadata" not in response._hidden_params + + def test_redacts_vertex_provider_metadata_from_metadata_hidden_params(self): + """Streaming success_handler copies _hidden_params into metadata before redaction.""" + details = { + "stream": True, + "litellm_params": { + "metadata": { + "hidden_params": { + "response_cost": 0.01, + "vertex_ai_grounding_metadata": [ + {"webSearchQueries": ["sensitive search term"]} + ], + "vertex_ai_url_context_metadata": [ + {"urlMetadata": [{"retrievedUrl": "https://example.com"}]} + ], + "vertex_ai_safety_ratings": [{"category": "HARM"}], + "vertex_ai_citation_metadata": [{"citations": ["source"]}], + } + } + }, + } + + perform_redaction(details, None) + + hidden_params = details["litellm_params"]["metadata"]["hidden_params"] + assert hidden_params["response_cost"] == 0.01 + assert "vertex_ai_grounding_metadata" not in hidden_params + assert "vertex_ai_url_context_metadata" not in hidden_params + assert "vertex_ai_safety_ratings" not in hidden_params + assert "vertex_ai_citation_metadata" not in hidden_params From 15507be90a5a7cbe82ed662760f13ae99e5346bb Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Mon, 8 Jun 2026 13:16:49 +0000 Subject: [PATCH 06/10] Fix Vertex streaming metadata from hidden params --- .../vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py | 7 ++++++- .../test_streaming_chunk_builder_utils.py | 1 - 2 files changed, 6 insertions(+), 2 deletions(-) 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 0a7ba70dacc1..372b79638359 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 @@ -2268,7 +2268,12 @@ def _get_stream_chunk_attr(chunk: Any, field_name: str) -> Any: return value model_extra = chunk.get("model_extra") if isinstance(model_extra, dict): - return model_extra.get(field_name) + value = model_extra.get(field_name) + if value is not None: + return value + hidden_params = chunk.get("_hidden_params") + if isinstance(hidden_params, dict): + return hidden_params.get(field_name) return None return getattr(chunk, field_name, None) diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py index 85661a9ecb9a..867f4cc7a0fa 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -720,7 +720,6 @@ def test_stream_chunk_builder_propagates_vertex_ai_metadata_from_dict_chunks(): ) ], ).model_dump() - chunk_dict["vertex_ai_grounding_metadata"] = [{"webSearchQueries": ["test query"]}] chunk_dict["_hidden_params"] = { "vertex_ai_grounding_metadata": [{"webSearchQueries": ["test query"]}] } From 8702ccfad24306bd6a821b5cf859732c5df49fbc Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 8 Jun 2026 20:45:59 +0000 Subject: [PATCH 07/10] fix(vertex): mirror vertex_ai_safety_results on assembled streaming responses The non-streaming transform_response stores safety data under vertex_ai_safety_results, but the streaming path only wrote vertex_ai_safety_ratings. Assembled streaming responses therefore never carried vertex_ai_safety_results, so any consumer reading that field saw a silent difference between streaming and non-streaming calls. Set vertex_ai_safety_results alongside vertex_ai_safety_ratings in the shared stream metadata setter and add it to the assembled metadata field list so it propagates through stream_chunk_builder. --- .../vertex_and_google_ai_studio_gemini.py | 3 ++ .../test_streaming_chunk_builder_utils.py | 31 +++++++++++++++++++ ...test_vertex_and_google_ai_studio_gemini.py | 20 ++++++++++++ 3 files changed, 54 insertions(+) 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 372b79638359..2d9965670582 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 @@ -2257,6 +2257,7 @@ def _extract_candidate_metadata( "vertex_ai_grounding_metadata", "vertex_ai_url_context_metadata", "vertex_ai_safety_ratings", + "vertex_ai_safety_results", "vertex_ai_citation_metadata", ) @@ -2296,8 +2297,10 @@ def _set_stream_metadata_on_response( url_context_metadata ) setattr(model_response, "vertex_ai_safety_ratings", safety_ratings) # type: ignore + setattr(model_response, "vertex_ai_safety_results", safety_ratings) # type: ignore if safety_ratings: model_response._hidden_params["vertex_ai_safety_ratings"] = safety_ratings + model_response._hidden_params["vertex_ai_safety_results"] = safety_ratings setattr(model_response, "vertex_ai_citation_metadata", citation_metadata) # type: ignore if citation_metadata: model_response._hidden_params["vertex_ai_citation_metadata"] = ( diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py index 867f4cc7a0fa..77765340c614 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -705,6 +705,37 @@ def test_stream_chunk_builder_uses_assembled_model_for_provider_metadata(): assert getattr(response, "vertex_ai_grounding_metadata") == grounding_metadata +def test_stream_chunk_builder_propagates_vertex_ai_safety_results(): + """Assembled response must expose safety data under the non-streaming field name.""" + safety_ratings = [ + [{"category": "HARM_CATEGORY_HATE_SPEECH", "probability": "NEGLIGIBLE"}] + ] + + chunk = ModelResponseStream( + id="chatcmpl-vertex-safety", + created=1, + model="gemini-2.5-flash", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(content="hello", role="assistant"), + ) + ], + ) + setattr(chunk, "vertex_ai_safety_ratings", safety_ratings) + setattr(chunk, "vertex_ai_safety_results", safety_ratings) + chunk._hidden_params["vertex_ai_safety_ratings"] = safety_ratings + chunk._hidden_params["vertex_ai_safety_results"] = safety_ratings + + response = stream_chunk_builder(chunks=[chunk]) + assert response is not None + assert getattr(response, "vertex_ai_safety_results") == safety_ratings + assert response._hidden_params["vertex_ai_safety_results"] == safety_ratings + assert response.model_dump()["vertex_ai_safety_results"] == safety_ratings + + def test_stream_chunk_builder_propagates_vertex_ai_metadata_from_dict_chunks(): """Dict snapshot chunks (model_dump) should also propagate Vertex AI metadata.""" chunk_dict = ModelResponseStream( diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 0d02521433a8..671d7355e8fb 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -1459,6 +1459,26 @@ def test_vertex_ai_process_candidates_with_grounding_metadata(): assert len(result[0]) == 1 +def test_set_stream_metadata_mirrors_non_streaming_safety_field_names(): + safety_ratings = [ + [{"category": "HARM_CATEGORY_HATE_SPEECH", "probability": "NEGLIGIBLE"}] + ] + + model_response = ModelResponse() + VertexGeminiConfig._set_stream_metadata_on_response( + model_response=model_response, + grounding_metadata=[], + url_context_metadata=[], + safety_ratings=safety_ratings, + citation_metadata=[], + ) + + assert getattr(model_response, "vertex_ai_safety_ratings") == safety_ratings + assert getattr(model_response, "vertex_ai_safety_results") == safety_ratings + assert model_response._hidden_params["vertex_ai_safety_ratings"] == safety_ratings + assert model_response._hidden_params["vertex_ai_safety_results"] == safety_ratings + + def test_vertex_ai_tool_call_id_format(): """ Test that tool call IDs have the correct format and length. From c3bcd7630c43a1b28c9f16fa489c900ef222f7f0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 8 Jun 2026 20:57:47 +0000 Subject: [PATCH 08/10] fix(streaming): log provider streaming metadata hook failures instead of swallowing them --- .../litellm_core_utils/streaming_chunk_builder_utils.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 2cfe7e7af5f7..6257cce9aecd 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -20,6 +20,7 @@ ServerToolUse, Usage, ) +from litellm._logging import verbose_logger from litellm.utils import print_verbose, token_counter if TYPE_CHECKING: @@ -120,8 +121,12 @@ def apply_provider_assembled_streaming_metadata( response=response, chunks=chunks, ) - except Exception: - return + except Exception as e: + verbose_logger.debug( + "apply_provider_assembled_streaming_metadata failed for model=%s: %s", + model, + e, + ) @staticmethod def _get_chunk_id(chunks: List[Dict[str, Any]]) -> str: From 514c0577722e92c65ac8c0da37b01adbafc3b053 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 8 Jun 2026 21:13:14 +0000 Subject: [PATCH 09/10] refactor(vertex): share single Vertex metadata field tuple across redaction and streaming --- litellm/litellm_core_utils/redact_messages.py | 19 ++++++------------- .../vertex_and_google_ai_studio_gemini.py | 11 ++--------- litellm/types/llms/vertex_ai.py | 9 +++++++++ 3 files changed, 17 insertions(+), 22 deletions(-) diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index dbaf6b594652..a2130c507064 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -18,6 +18,7 @@ get_metadata_variable_name_from_kwargs, ) from litellm.secret_managers.main import str_to_bool +from litellm.types.llms.vertex_ai import VERTEX_AI_PROVIDER_METADATA_FIELDS from litellm.types.utils import StandardCallbackDynamicParams if TYPE_CHECKING: @@ -29,14 +30,6 @@ else: LiteLLMLoggingObject = Any -VERTEX_PROVIDER_METADATA_FIELDS = ( - "vertex_ai_grounding_metadata", - "vertex_ai_url_context_metadata", - "vertex_ai_safety_ratings", - "vertex_ai_safety_results", - "vertex_ai_citation_metadata", -) - def redact_message_input_output_from_custom_logger( litellm_logging_obj: LiteLLMLoggingObject, result, custom_logger: CustomLogger @@ -110,21 +103,21 @@ def _redact_responses_api_output_dict(output_items, redacted_str: str): def _redact_vertex_provider_metadata(obj: Any) -> None: if isinstance(obj, dict): - for field in VERTEX_PROVIDER_METADATA_FIELDS: + for field in VERTEX_AI_PROVIDER_METADATA_FIELDS: if field in obj: obj[field] = [] hidden_params = obj.get("_hidden_params") if isinstance(hidden_params, dict): - for field in VERTEX_PROVIDER_METADATA_FIELDS: + for field in VERTEX_AI_PROVIDER_METADATA_FIELDS: hidden_params.pop(field, None) return - for field in VERTEX_PROVIDER_METADATA_FIELDS: + for field in VERTEX_AI_PROVIDER_METADATA_FIELDS: if hasattr(obj, field): setattr(obj, field, []) hidden_params = getattr(obj, "_hidden_params", None) if isinstance(hidden_params, dict): - for field in VERTEX_PROVIDER_METADATA_FIELDS: + for field in VERTEX_AI_PROVIDER_METADATA_FIELDS: hidden_params.pop(field, None) @@ -148,7 +141,7 @@ def _redact_vertex_provider_metadata_from_litellm_params( hidden_params = metadata.get("hidden_params") if not isinstance(hidden_params, dict): continue - for field in VERTEX_PROVIDER_METADATA_FIELDS: + for field in VERTEX_AI_PROVIDER_METADATA_FIELDS: hidden_params.pop(field, None) 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 2d9965670582..c919a367fbf6 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 @@ -63,6 +63,7 @@ OpenAIChatCompletionFinishReason, ) from litellm.types.llms.vertex_ai import ( + VERTEX_AI_PROVIDER_METADATA_FIELDS, VERTEX_CREDENTIALS_TYPES, Candidates, ContentType, @@ -2253,14 +2254,6 @@ def _extract_candidate_metadata( citation_metadata, ) - _STREAM_METADATA_FIELDS = ( - "vertex_ai_grounding_metadata", - "vertex_ai_url_context_metadata", - "vertex_ai_safety_ratings", - "vertex_ai_safety_results", - "vertex_ai_citation_metadata", - ) - @staticmethod def _get_stream_chunk_attr(chunk: Any, field_name: str) -> Any: if isinstance(chunk, dict): @@ -2312,7 +2305,7 @@ def apply_assembled_streaming_response_metadata( response: ModelResponse, chunks: List[Any], ) -> None: - for field_name in VertexGeminiConfig._STREAM_METADATA_FIELDS: + for field_name in VERTEX_AI_PROVIDER_METADATA_FIELDS: merged: List[Any] = [] for chunk in chunks: value = VertexGeminiConfig._get_stream_chunk_attr(chunk, field_name) diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 51429d0769e7..ea5d4471e82f 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -757,3 +757,12 @@ class VertexPartnerProvider(str, Enum): llama = "llama" ai21 = "ai21" claude = "claude" + + +VERTEX_AI_PROVIDER_METADATA_FIELDS = ( + "vertex_ai_grounding_metadata", + "vertex_ai_url_context_metadata", + "vertex_ai_safety_ratings", + "vertex_ai_safety_results", + "vertex_ai_citation_metadata", +) From e411597e386d8f99203211a91e388efa6f2c693f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 8 Jun 2026 21:31:26 +0000 Subject: [PATCH 10/10] refactor(vertex): move Vertex metadata redaction helpers into llms/vertex_ai --- litellm/litellm_core_utils/redact_messages.py | 61 +++---------------- litellm/llms/vertex_ai/common_utils.py | 47 +++++++++++++- 2 files changed, 56 insertions(+), 52 deletions(-) diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index a2130c507064..763596336a08 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -17,8 +17,11 @@ from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, ) +from litellm.llms.vertex_ai.common_utils import ( + redact_vertex_ai_metadata_from_litellm_params, + redact_vertex_ai_metadata_from_logged_object, +) from litellm.secret_managers.main import str_to_bool -from litellm.types.llms.vertex_ai import VERTEX_AI_PROVIDER_METADATA_FIELDS from litellm.types.utils import StandardCallbackDynamicParams if TYPE_CHECKING: @@ -101,50 +104,6 @@ def _redact_responses_api_output_dict(output_items, redacted_str: str): summary_item["text"] = redacted_str -def _redact_vertex_provider_metadata(obj: Any) -> None: - if isinstance(obj, dict): - for field in VERTEX_AI_PROVIDER_METADATA_FIELDS: - if field in obj: - obj[field] = [] - hidden_params = obj.get("_hidden_params") - if isinstance(hidden_params, dict): - for field in VERTEX_AI_PROVIDER_METADATA_FIELDS: - hidden_params.pop(field, None) - return - - for field in VERTEX_AI_PROVIDER_METADATA_FIELDS: - if hasattr(obj, field): - setattr(obj, field, []) - hidden_params = getattr(obj, "_hidden_params", None) - if isinstance(hidden_params, dict): - for field in VERTEX_AI_PROVIDER_METADATA_FIELDS: - hidden_params.pop(field, None) - - -def _redact_vertex_provider_metadata_from_litellm_params( - model_call_details: dict, -) -> None: - """ - Scrub Vertex provider metadata copied into litellm_params metadata. - - success_handler() merges response._hidden_params into - litellm_params.metadata['hidden_params'] before perform_redaction() runs. - """ - litellm_params = model_call_details.get("litellm_params") - if not isinstance(litellm_params, dict): - return - - for metadata_key in ("metadata", "litellm_metadata"): - metadata = litellm_params.get(metadata_key) - if not isinstance(metadata, dict): - continue - hidden_params = metadata.get("hidden_params") - if not isinstance(hidden_params, dict): - continue - for field in VERTEX_AI_PROVIDER_METADATA_FIELDS: - hidden_params.pop(field, None) - - def _redact_standard_logging_object(model_call_details: dict): """Redact messages and response inside standard_logging_object if present.""" standard_logging_object = model_call_details.get("standard_logging_object") @@ -164,12 +123,12 @@ def _redact_standard_logging_object(model_call_details: dict): # ResponsesAPIResponse format - redact content in output items if isinstance(response.get("output"), list): _redact_responses_api_output_dict(response["output"], redacted_str) - _redact_vertex_provider_metadata(response) + redact_vertex_ai_metadata_from_logged_object(response) elif isinstance(response, dict) and "choices" in response: # ModelResponse dict format - redact content in choices if isinstance(response.get("choices"), list): _redact_model_response_dict_choices(response["choices"], redacted_str) - _redact_vertex_provider_metadata(response) + redact_vertex_ai_metadata_from_logged_object(response) elif isinstance(response, str): standard_logging_object["response"] = redacted_str else: @@ -211,7 +170,7 @@ def perform_redaction(model_call_details: dict, result): model_call_details["prompt"] = "" model_call_details["input"] = "" _redact_standard_logging_object(model_call_details) - _redact_vertex_provider_metadata_from_litellm_params(model_call_details) + redact_vertex_ai_metadata_from_litellm_params(model_call_details) # Redact streaming response if ( @@ -222,7 +181,7 @@ def perform_redaction(model_call_details: dict, result): if hasattr(_streaming_response, "choices"): for choice in _streaming_response.choices: _redact_choice_content(choice) - _redact_vertex_provider_metadata(_streaming_response) + redact_vertex_ai_metadata_from_logged_object(_streaming_response) elif hasattr(_streaming_response, "output"): _redact_responses_api_output(_streaming_response.output) # Redact reasoning field in ResponsesAPIResponse @@ -249,14 +208,14 @@ def perform_redaction(model_call_details: dict, result): if hasattr(_result, "choices") and _result.choices is not None: for choice in _result.choices: _redact_choice_content(choice) - _redact_vertex_provider_metadata(_result) + redact_vertex_ai_metadata_from_logged_object(_result) elif isinstance(_result, dict) and "choices" in _result: # Handle dict representation of ModelResponse (e.g., from model_dump()) if _result.get("choices") is not None: _redact_model_response_dict_choices( _result["choices"], "redacted-by-litellm" ) - _redact_vertex_provider_metadata(_result) + redact_vertex_ai_metadata_from_logged_object(_result) elif isinstance(_result, dict) and "output" in _result: if isinstance(_result.get("output"), list): _redact_responses_api_output_dict( diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index e6e396511096..85c23d8603c1 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -12,7 +12,11 @@ from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.llms.openai import AllMessageValues -from litellm.types.llms.vertex_ai import PartType, Schema +from litellm.types.llms.vertex_ai import ( + VERTEX_AI_PROVIDER_METADATA_FIELDS, + PartType, + Schema, +) from litellm.types.utils import TokenCountResponse from litellm.utils import supports_response_schema, supports_system_messages @@ -27,6 +31,47 @@ def __init__( super().__init__(message=message, status_code=status_code, headers=headers) +def redact_vertex_ai_metadata_from_logged_object(obj: Any) -> None: + if isinstance(obj, dict): + for field in VERTEX_AI_PROVIDER_METADATA_FIELDS: + if field in obj: + obj[field] = [] + hidden_params = obj.get("_hidden_params") + if isinstance(hidden_params, dict): + for field in VERTEX_AI_PROVIDER_METADATA_FIELDS: + hidden_params.pop(field, None) + return + + for field in VERTEX_AI_PROVIDER_METADATA_FIELDS: + if hasattr(obj, field): + setattr(obj, field, []) + hidden_params = getattr(obj, "_hidden_params", None) + if isinstance(hidden_params, dict): + for field in VERTEX_AI_PROVIDER_METADATA_FIELDS: + hidden_params.pop(field, None) + + +def redact_vertex_ai_metadata_from_litellm_params(model_call_details: dict) -> None: + """ + success_handler() merges response._hidden_params into + litellm_params.metadata['hidden_params'] before redaction runs, so the Vertex + metadata must be scrubbed from that copy too. + """ + litellm_params = model_call_details.get("litellm_params") + if not isinstance(litellm_params, dict): + return + + for metadata_key in ("metadata", "litellm_metadata"): + metadata = litellm_params.get(metadata_key) + if not isinstance(metadata, dict): + continue + hidden_params = metadata.get("hidden_params") + if not isinstance(hidden_params, dict): + continue + for field in VERTEX_AI_PROVIDER_METADATA_FIELDS: + hidden_params.pop(field, None) + + def vertex_request_labels_from_litellm_params( litellm_params: Optional[dict], ) -> Optional[Dict[str, str]]: