From a264d1ca8cedd15c1494cb047b3376b0d68cf1db Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 26 Jul 2024 20:54:59 -0700 Subject: [PATCH] feat(vertex_httpx.py): support logging citation metadata Closes https://github.com/BerriAI/litellm/issues/3230 --- litellm/llms/vertex_httpx.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/litellm/llms/vertex_httpx.py b/litellm/llms/vertex_httpx.py index bd3367830a..142cabbe96 100644 --- a/litellm/llms/vertex_httpx.py +++ b/litellm/llms/vertex_httpx.py @@ -689,6 +689,7 @@ class VertexLLM(BaseLLM): ## CHECK IF GROUNDING METADATA IN REQUEST grounding_metadata: List[dict] = [] safety_ratings: List = [] + citation_metadata: List = [] ## GET TEXT ## chat_completion_message = {"role": "assistant"} content_str = "" @@ -702,6 +703,9 @@ class VertexLLM(BaseLLM): if "safetyRatings" in candidate: safety_ratings.append(candidate["safetyRatings"]) + + if "citationMetadata" in candidate: + citation_metadata.append(candidate["citationMetadata"]) if "text" in candidate["content"]["parts"][0]: content_str = candidate["content"]["parts"][0]["text"] @@ -756,6 +760,11 @@ class VertexLLM(BaseLLM): ## ADD SAFETY RATINGS ## model_response._hidden_params["vertex_ai_safety_results"] = safety_ratings + ## ADD CITATION METADATA ## + model_response._hidden_params["vertex_ai_citation_metadata"] = ( + citation_metadata + ) + 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(