scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-04 02:31:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 69

refactor(vertex): move Vertex metadata redaction helpers into llms/vertex_ai

Browse source
This commit is contained in:
mateo-berri 2026-06-08 21:31:26 +00:00
parent 514c057772
commit e411597e38
No known key found for this signature in database
2 changed files with 56 additions and 52 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

61
litellm/litellm_core_utils/redact_messages.py
View file

@ -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(

47
litellm/llms/vertex_ai/common_utils.py
View file

@ -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]]: