From fe7f78fbf6ca18963578ff07ac8858d2de71eae1 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Fri, 26 Jul 2024 20:50:18 -0700
Subject: [PATCH] feat(vertex_httpx.py): support logging vertex ai safety
 results to langfuse

Closes https://github.com/BerriAI/litellm/issues/3230
---
 litellm/litellm_core_utils/litellm_logging.py | 20 +++++++++++++++++++
 litellm/llms/vertex_httpx.py                  |  7 +++++++
 2 files changed, 27 insertions(+)

diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py
index 85b6adc1e..852f1a2d9 100644
--- a/litellm/litellm_core_utils/litellm_logging.py
+++ b/litellm/litellm_core_utils/litellm_logging.py
@@ -529,6 +529,7 @@ class Logging:
                     or isinstance(result, TextCompletionResponse)
                     or isinstance(result, HttpxBinaryResponseContent)  # tts
                 ):
+                    ## RESPONSE COST ##
                     custom_pricing = use_custom_pricing_for_model(
                         litellm_params=self.litellm_params
                     )
@@ -548,6 +549,25 @@ class Logging:
                             custom_pricing=custom_pricing,
                         )
                     )
+
+                    ## HIDDEN PARAMS ##
+                    if hasattr(result, "_hidden_params"):
+                        # add to metadata for logging
+                        if self.model_call_details.get("litellm_params") is not None:
+                            self.model_call_details["litellm_params"].setdefault(
+                                "metadata", {}
+                            )
+                            if (
+                                self.model_call_details["litellm_params"]["metadata"]
+                                is None
+                            ):
+                                self.model_call_details["litellm_params"][
+                                    "metadata"
+                                ] = {}
+
+                            self.model_call_details["litellm_params"]["metadata"][
+                                "hidden_params"
+                            ] = result._hidden_params
             else:  # streaming chunks + image gen.
                 self.model_call_details["response_cost"] = None
 
diff --git a/litellm/llms/vertex_httpx.py b/litellm/llms/vertex_httpx.py
index 93d8f4282..bd3367830 100644
--- a/litellm/llms/vertex_httpx.py
+++ b/litellm/llms/vertex_httpx.py
@@ -688,6 +688,7 @@ class VertexLLM(BaseLLM):
         try:
             ## CHECK IF GROUNDING METADATA IN REQUEST
             grounding_metadata: List[dict] = []
+            safety_ratings: List = []
             ## GET TEXT ##
             chat_completion_message = {"role": "assistant"}
             content_str = ""
@@ -699,6 +700,8 @@ class VertexLLM(BaseLLM):
                 if "groundingMetadata" in candidate:
                     grounding_metadata.append(candidate["groundingMetadata"])
 
+                if "safetyRatings" in candidate:
+                    safety_ratings.append(candidate["safetyRatings"])
                 if "text" in candidate["content"]["parts"][0]:
                     content_str = candidate["content"]["parts"][0]["text"]
 
@@ -749,6 +752,10 @@ class VertexLLM(BaseLLM):
             model_response._hidden_params["vertex_ai_grounding_metadata"] = (
                 grounding_metadata
             )
+
+            ## ADD SAFETY RATINGS ##
+            model_response._hidden_params["vertex_ai_safety_results"] = safety_ratings
+
         except Exception as e:
             raise VertexAIError(
                 message="Received={}, Error converting to valid response block={}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues".format(