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] 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 a2130c50706..763596336a0 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -17,8 +17,11 @@ from litellm.integrations.custom_logger import CustomLogger 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 e6e39651109..85c23d8603c 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -12,7 +12,11 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import unpack_defs 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 @@ class VertexAIError(BaseLLMException): 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]]: