phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 19:24:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat(vertex_httpx.py): support logging citation metadata

Browse source 
Closes https://github.com/BerriAI/litellm/issues/3230
This commit is contained in:
Krrish Dholakia 2024-07-26 20:54:59 -07:00
parent fe7f78fbf6
commit a264d1ca8c
1 changed files with 9 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

9
litellm/llms/vertex_httpx.py
View file

@ -689,6 +689,7 @@ class VertexLLM(BaseLLM):
## CHECK IF GROUNDING METADATA IN REQUEST ## CHECK IF GROUNDING METADATA IN REQUEST
grounding_metadata: List[dict] = [] grounding_metadata: List[dict] = []
safety_ratings: List = [] safety_ratings: List = []
citation_metadata: List = []
## GET TEXT ## ## GET TEXT ##
chat_completion_message = {"role": "assistant"} chat_completion_message = {"role": "assistant"}
content_str = "" content_str = ""
@ -702,6 +703,9 @@ class VertexLLM(BaseLLM):
if "safetyRatings" in candidate: if "safetyRatings" in candidate:
safety_ratings.append(candidate["safetyRatings"]) safety_ratings.append(candidate["safetyRatings"])
if "citationMetadata" in candidate:
citation_metadata.append(candidate["citationMetadata"])
if "text" in candidate["content"]["parts"][0]: if "text" in candidate["content"]["parts"][0]:
content_str = candidate["content"]["parts"][0]["text"] content_str = candidate["content"]["parts"][0]["text"]
@ -756,6 +760,11 @@ class VertexLLM(BaseLLM):
## ADD SAFETY RATINGS ## ## ADD SAFETY RATINGS ##
model_response._hidden_params["vertex_ai_safety_results"] = 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: except Exception as e:
raise VertexAIError( raise VertexAIError(
message="Received={}, Error converting to valid response block={}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues".format( message="Received={}, Error converting to valid response block={}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues".format(