phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

test: handle vertex ai rate limit errors

Browse source
This commit is contained in:
Krrish Dholakia 2024-03-18 21:33:08 -07:00
parent 2827acc487
commit 2fdc20f549
1 changed files with 40 additions and 25 deletions
Download patch file Download diff file Expand all files Collapse all files

65
litellm/tests/test_amazing_vertex_completion.py
View file

@ -374,7 +374,8 @@ def test_gemini_pro_vision_base64():
print(resp) print(resp)
prompt_tokens = resp.usage.prompt_tokens prompt_tokens = resp.usage.prompt_tokens
except litellm.RateLimitError as e:
pass
except Exception as e: except Exception as e:
if "500 Internal error encountered.'" in str(e): if "500 Internal error encountered.'" in str(e):
pass pass
@ -419,33 +420,43 @@ def test_gemini_pro_function_calling():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_gemini_pro_async_function_calling(): async def test_gemini_pro_async_function_calling():
load_vertex_ai_credentials() load_vertex_ai_credentials()
tools = [ try:
{ tools = [
"type": "function", {
"function": { "type": "function",
"name": "get_current_weather", "function": {
"description": "Get the current weather in a given location", "name": "get_current_weather",
"parameters": { "description": "Get the current weather in a given location",
"type": "object", "parameters": {
"properties": { "type": "object",
"location": { "properties": {
"type": "string", "location": {
"description": "The city and state, e.g. San Francisco, CA", "type": "string",
"description": "The city and state, e.g. San Francisco, CA",
},
"unit": {
"type": "string",
"enum": ["celsius", "fahrenheit"],
},
}, },
"unit": {"type": "string", "enum": ["celsius", "fahrenheit"]}, "required": ["location"],
}, },
"required": ["location"],
}, },
}, }
} ]
] messages = [
messages = [{"role": "user", "content": "What's the weather like in Boston today?"}] {"role": "user", "content": "What's the weather like in Boston today?"}
completion = await litellm.acompletion( ]
model="gemini-pro", messages=messages, tools=tools, tool_choice="auto" completion = await litellm.acompletion(
) model="gemini-pro", messages=messages, tools=tools, tool_choice="auto"
print(f"completion: {completion}") )
assert completion.choices[0].message.content is None print(f"completion: {completion}")
assert len(completion.choices[0].message.tool_calls) == 1 assert completion.choices[0].message.content is None
assert len(completion.choices[0].message.tool_calls) == 1
except litellm.RateLimitError as e:
pass
except Exception as e:
pytest.fail(f"An exception occurred - {str(e)}")
# raise Exception("it worked!") # raise Exception("it worked!")
@ -461,6 +472,8 @@ def test_vertexai_embedding():
input=["good morning from litellm", "this is another item"], input=["good morning from litellm", "this is another item"],
) )
print(f"response:", response) print(f"response:", response)
except litellm.RateLimitError as e:
pass
except Exception as e: except Exception as e:
pytest.fail(f"Error occurred: {e}") pytest.fail(f"Error occurred: {e}")
@ -475,6 +488,8 @@ async def test_vertexai_aembedding():
input=["good morning from litellm", "this is another item"], input=["good morning from litellm", "this is another item"],
) )
print(f"response: {response}") print(f"response: {response}")
except litellm.RateLimitError as e:
pass
except Exception as e: except Exception as e:
pytest.fail(f"Error occurred: {e}") pytest.fail(f"Error occurred: {e}")