mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
fix(vertex_httpx.py): return grounding metadata
This commit is contained in:
Krrish Dholakia
parent
959c627dd3
commit
023f10cf1c
3 changed files with 145 additions and 10 deletions
|
|
@ -686,6 +686,8 @@ class VertexLLM(BaseLLM):
|
|||
model_response.choices = [] # type: ignore
|
||||
|
||||
try:
|
||||
## CHECK IF GROUNDING METADATA IN REQUEST
|
||||
grounding_metadata: List[dict] = []
|
||||
## GET TEXT ##
|
||||
chat_completion_message = {"role": "assistant"}
|
||||
content_str = ""
|
||||
|
|
@ -694,6 +696,9 @@ class VertexLLM(BaseLLM):
|
|||
if "content" not in candidate:
|
||||
continue
|
||||
|
||||
if "groundingMetadata" in candidate:
|
||||
grounding_metadata.append(candidate["groundingMetadata"])
|
||||
|
||||
if "text" in candidate["content"]["parts"][0]:
|
||||
content_str = candidate["content"]["parts"][0]["text"]
|
||||
|
||||
|
|
@ -739,6 +744,11 @@ class VertexLLM(BaseLLM):
|
|||
)
|
||||
|
||||
setattr(model_response, "usage", usage)
|
||||
|
||||
## ADD GROUNDING METADATA ##
|
||||
model_response._hidden_params["vertex_ai_grounding_metadata"] = (
|
||||
grounding_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(
|
||||
|
|
|
|||
|
|
@ -677,6 +677,128 @@ def test_gemini_pro_vision_base64():
|
|||
pytest.fail(f"An exception occurred - {str(e)}")
|
||||
|
||||
|
||||
def vertex_httpx_grounding_post(*args, **kwargs):
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.headers = {"Content-Type": "application/json"}
|
||||
mock_response.json.return_value = {
|
||||
"candidates": [
|
||||
{
|
||||
"content": {
|
||||
"role": "model",
|
||||
"parts": [
|
||||
{
|
||||
"text": "Argentina won the FIFA World Cup 2022. Argentina defeated France 4-2 on penalties in the FIFA World Cup 2022 final tournament for the first time after 36 years and the third time overall."
|
||||
}
|
||||
],
|
||||
},
|
||||
"finishReason": "STOP",
|
||||
"safetyRatings": [
|
||||
{
|
||||
"category": "HARM_CATEGORY_HATE_SPEECH",
|
||||
"probability": "NEGLIGIBLE",
|
||||
"probabilityScore": 0.14940722,
|
||||
"severity": "HARM_SEVERITY_NEGLIGIBLE",
|
||||
"severityScore": 0.07477004,
|
||||
},
|
||||
{
|
||||
"category": "HARM_CATEGORY_DANGEROUS_CONTENT",
|
||||
"probability": "NEGLIGIBLE",
|
||||
"probabilityScore": 0.15636235,
|
||||
"severity": "HARM_SEVERITY_NEGLIGIBLE",
|
||||
"severityScore": 0.015967654,
|
||||
},
|
||||
{
|
||||
"category": "HARM_CATEGORY_HARASSMENT",
|
||||
"probability": "NEGLIGIBLE",
|
||||
"probabilityScore": 0.1943678,
|
||||
"severity": "HARM_SEVERITY_NEGLIGIBLE",
|
||||
"severityScore": 0.1284158,
|
||||
},
|
||||
{
|
||||
"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
|
||||
"probability": "NEGLIGIBLE",
|
||||
"probabilityScore": 0.09384396,
|
||||
"severity": "HARM_SEVERITY_NEGLIGIBLE",
|
||||
"severityScore": 0.0726367,
|
||||
},
|
||||
],
|
||||
"groundingMetadata": {
|
||||
"webSearchQueries": ["who won the world cup 2022"],
|
||||
"groundingAttributions": [
|
||||
{
|
||||
"segment": {"endIndex": 38},
|
||||
"confidenceScore": 0.9919262,
|
||||
"web": {
|
||||
"uri": "https://www.careerpower.in/fifa-world-cup-winners-list.html",
|
||||
"title": "FIFA World Cup Winners List from 1930 to 2022, Complete List - Career Power",
|
||||
},
|
||||
},
|
||||
{
|
||||
"segment": {"endIndex": 38},
|
||||
"confidenceScore": 0.9919262,
|
||||
"web": {
|
||||
"uri": "https://www.careerpower.in/fifa-world-cup-winners-list.html",
|
||||
"title": "FIFA World Cup Winners List from 1930 to 2022, Complete List - Career Power",
|
||||
},
|
||||
},
|
||||
{
|
||||
"segment": {"endIndex": 38},
|
||||
"confidenceScore": 0.9919262,
|
||||
"web": {
|
||||
"uri": "https://www.britannica.com/sports/2022-FIFA-World-Cup",
|
||||
"title": "2022 FIFA World Cup | Qatar, Controversy, Stadiums, Winner, & Final - Britannica",
|
||||
},
|
||||
},
|
||||
{
|
||||
"segment": {"endIndex": 38},
|
||||
"confidenceScore": 0.9919262,
|
||||
"web": {
|
||||
"uri": "https://en.wikipedia.org/wiki/2022_FIFA_World_Cup_final",
|
||||
"title": "2022 FIFA World Cup final - Wikipedia",
|
||||
},
|
||||
},
|
||||
{
|
||||
"segment": {"endIndex": 38},
|
||||
"confidenceScore": 0.9919262,
|
||||
"web": {
|
||||
"uri": "https://www.transfermarkt.com/2022-world-cup/erfolge/pokalwettbewerb/WM22",
|
||||
"title": "2022 World Cup - All winners - Transfermarkt",
|
||||
},
|
||||
},
|
||||
{
|
||||
"segment": {"startIndex": 39, "endIndex": 187},
|
||||
"confidenceScore": 0.9919262,
|
||||
"web": {
|
||||
"uri": "https://www.careerpower.in/fifa-world-cup-winners-list.html",
|
||||
"title": "FIFA World Cup Winners List from 1930 to 2022, Complete List - Career Power",
|
||||
},
|
||||
},
|
||||
{
|
||||
"segment": {"startIndex": 39, "endIndex": 187},
|
||||
"confidenceScore": 0.9919262,
|
||||
"web": {
|
||||
"uri": "https://en.wikipedia.org/wiki/2022_FIFA_World_Cup_final",
|
||||
"title": "2022 FIFA World Cup final - Wikipedia",
|
||||
},
|
||||
},
|
||||
],
|
||||
"searchEntryPoint": {
|
||||
"renderedContent": '\u003cstyle\u003e\n.container {\n align-items: center;\n border-radius: 8px;\n display: flex;\n font-family: Google Sans, Roboto, sans-serif;\n font-size: 14px;\n line-height: 20px;\n padding: 8px 12px;\n}\n.chip {\n display: inline-block;\n border: solid 1px;\n border-radius: 16px;\n min-width: 14px;\n padding: 5px 16px;\n text-align: center;\n user-select: none;\n margin: 0 8px;\n -webkit-tap-highlight-color: transparent;\n}\n.carousel {\n overflow: auto;\n scrollbar-width: none;\n white-space: nowrap;\n margin-right: -12px;\n}\n.headline {\n display: flex;\n margin-right: 4px;\n}\n.gradient-container {\n position: relative;\n}\n.gradient {\n position: absolute;\n transform: translate(3px, -9px);\n height: 36px;\n width: 9px;\n}\n@media (prefers-color-scheme: light) {\n .container {\n background-color: #fafafa;\n box-shadow: 0 0 0 1px #0000000f;\n }\n .headline-label {\n color: #1f1f1f;\n }\n .chip {\n background-color: #ffffff;\n border-color: #d2d2d2;\n color: #5e5e5e;\n text-decoration: none;\n }\n .chip:hover {\n background-color: #f2f2f2;\n }\n .chip:focus {\n background-color: #f2f2f2;\n }\n .chip:active {\n background-color: #d8d8d8;\n border-color: #b6b6b6;\n }\n .logo-dark {\n display: none;\n }\n .gradient {\n background: linear-gradient(90deg, #fafafa 15%, #fafafa00 100%);\n }\n}\n@media (prefers-color-scheme: dark) {\n .container {\n background-color: #1f1f1f;\n box-shadow: 0 0 0 1px #ffffff26;\n }\n .headline-label {\n color: #fff;\n }\n .chip {\n background-color: #2c2c2c;\n border-color: #3c4043;\n color: #fff;\n text-decoration: none;\n }\n .chip:hover {\n background-color: #353536;\n }\n .chip:focus {\n background-color: #353536;\n }\n .chip:active {\n background-color: #464849;\n border-color: #53575b;\n }\n .logo-light {\n display: none;\n }\n .gradient {\n background: linear-gradient(90deg, #1f1f1f 15%, #1f1f1f00 100%);\n }\n}\n\u003c/style\u003e\n\u003cdiv class="container"\u003e\n \u003cdiv class="headline"\u003e\n \u003csvg class="logo-light" width="18" height="18" viewBox="9 9 35 35" fill="none" xmlns="http://www.w3.org/2000/svg"\u003e\n \u003cpath fill-rule="evenodd" clip-rule="evenodd" d="M42.8622 27.0064C42.8622 25.7839 42.7525 24.6084 42.5487 23.4799H26.3109V30.1568H35.5897C35.1821 32.3041 33.9596 34.1222 32.1258 35.3448V39.6864H37.7213C40.9814 36.677 42.8622 32.2571 42.8622 27.0064V27.0064Z" fill="#4285F4"/\u003e\n \u003cpath fill-rule="evenodd" clip-rule="evenodd" d="M26.3109 43.8555C30.9659 43.8555 34.8687 42.3195 37.7213 39.6863L32.1258 35.3447C30.5898 36.3792 28.6306 37.0061 26.3109 37.0061C21.8282 37.0061 18.0195 33.9811 16.6559 29.906H10.9194V34.3573C13.7563 39.9841 19.5712 43.8555 26.3109 43.8555V43.8555Z" fill="#34A853"/\u003e\n \u003cpath fill-rule="evenodd" clip-rule="evenodd" d="M16.6559 29.8904C16.3111 28.8559 16.1074 27.7588 16.1074 26.6146C16.1074 25.4704 16.3111 24.3733 16.6559 23.3388V18.8875H10.9194C9.74388 21.2072 9.06992 23.8247 9.06992 26.6146C9.06992 29.4045 9.74388 32.022 10.9194 34.3417L15.3864 30.8621L16.6559 29.8904V29.8904Z" fill="#FBBC05"/\u003e\n \u003cpath fill-rule="evenodd" clip-rule="evenodd" d="M26.3109 16.2386C28.85 16.2386 31.107 17.1164 32.9095 18.8091L37.8466 13.8719C34.853 11.082 30.9659 9.3736 26.3109 9.3736C19.5712 9.3736 13.7563 13.245 10.9194 18.8875L16.6559 23.3388C18.0195 19.2636 21.8282 16.2386 26.3109 16.2386V16.2386Z" fill="#EA4335"/\u003e\n \u003c/svg\u003e\n \u003csvg class="logo-dark" width="18" height="18" viewBox="0 0 48 48" xmlns="http://www.w3.org/2000/svg"\u003e\n \u003ccircle cx="24" cy="23" fill="#FFF" r="22"/\u003e\n \u003cpath d="M33.76 34.26c2.75-2.56 4.49-6.37 4.49-11.26 0-.89-.08-1.84-.29-3H24.01v5.99h8.03c-.4 2.02-1.5 3.56-3.07 4.56v.75l3.91 2.97h.88z" fill="#4285F4"/\u003e\n \u003cpath d="M15.58 25.77A8.845 8.845 0 0 0 24 31.86c1.92 0 3.62-.46 4.97-1.31l4.79 3.71C31.14 36.7 27.65 38 24 38c-5.93 0-11.01-3.4-13.45-8.36l.17-1.01 4.06-2.85h.8z" fill="#34A853"/\u003e\n \u003cpath d="M15.59 20.21a8.864 8.864 0 0 0 0 5.58l-5.03 3.86c-.98-2-1.53-4.25-1.53-6.64 0-2.39.55-4.64 1.53-6.64l1-.22 3.81 2.98.22 1.08z" fill="#FBBC05"/\u003e\n \u003cpath d="M24 14.14c2.11 0 4.02.75 5.52 1.98l4.36-4.36C31.22 9.43 27.81 8 24 8c-5.93 0-11.01 3.4-13.45 8.36l5.03 3.85A8.86 8.86 0 0 1 24 14.14z" fill="#EA4335"/\u003e\n \u003c/svg\u003e\n \u003cdiv class="gradient-container"\u003e\u003cdiv class="gradient"\u003e\u003c/div\u003e\u003c/div\u003e\n \u003c/div\u003e\n \u003cdiv class="carousel"\u003e\n \u003ca class="chip" href="https://www.google.com/search?q=who+won+the+world+cup+2022&client=app-vertex-grounding&safesearch=active"\u003ewho won the world cup 2022\u003c/a\u003e\n \u003c/div\u003e\n\u003c/div\u003e\n'
|
||||
},
|
||||
},
|
||||
}
|
||||
],
|
||||
"usageMetadata": {
|
||||
"promptTokenCount": 6,
|
||||
"candidatesTokenCount": 48,
|
||||
"totalTokenCount": 54,
|
||||
},
|
||||
}
|
||||
|
||||
return mock_response
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value_in_dict", [{}, {"disable_attribution": False}]) #
|
||||
def test_gemini_pro_grounding(value_in_dict):
|
||||
try:
|
||||
|
|
@ -691,16 +813,15 @@ def test_gemini_pro_grounding(value_in_dict):
|
|||
|
||||
client = HTTPHandler()
|
||||
|
||||
with patch.object(client, "post", new=MagicMock()) as mock_call:
|
||||
try:
|
||||
resp = litellm.completion(
|
||||
model="vertex_ai_beta/gemini-1.0-pro-001",
|
||||
messages=[{"role": "user", "content": "Who won the world cup?"}],
|
||||
tools=tools,
|
||||
client=client,
|
||||
)
|
||||
except Exception as e:
|
||||
print("Received Exception - {}".format(str(e)))
|
||||
with patch.object(
|
||||
client, "post", side_effect=vertex_httpx_grounding_post
|
||||
) as mock_call:
|
||||
resp = litellm.completion(
|
||||
model="vertex_ai_beta/gemini-1.0-pro-001",
|
||||
messages=[{"role": "user", "content": "Who won the world cup?"}],
|
||||
tools=tools,
|
||||
client=client,
|
||||
)
|
||||
|
||||
mock_call.assert_called_once()
|
||||
|
||||
|
|
@ -715,6 +836,9 @@ def test_gemini_pro_grounding(value_in_dict):
|
|||
== value_in_dict
|
||||
)
|
||||
|
||||
assert "vertex_ai_grounding_metadata" in resp._hidden_params
|
||||
assert isinstance(resp._hidden_params["vertex_ai_grounding_metadata"], list)
|
||||
|
||||
except litellm.InternalServerError:
|
||||
pass
|
||||
except litellm.RateLimitError:
|
||||
|
|
|
|||
|
|
@ -221,6 +221,7 @@ class SearchEntryPoint(TypedDict, total=False):
|
|||
class GroundingMetadata(TypedDict, total=False):
|
||||
webSearchQueries: List[str]
|
||||
searchEntryPoint: SearchEntryPoint
|
||||
groundingAttributions: List[dict]
|
||||
|
||||
|
||||
class Candidates(TypedDict, total=False):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue