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

test: handle gemini error

Browse source
This commit is contained in:
Krrish Dholakia 2024-11-14 22:26:06 +05:30
parent 8b4c33c398
commit 7d4829f908
1 changed files with 19 additions and 13 deletions
Download patch file Download diff file Expand all files Collapse all files

32
tests/local_testing/test_stream_chunk_builder.py
View file

@ -172,6 +172,8 @@ def test_stream_chunk_builder_litellm_usage_chunks():
""" """
Checks if stream_chunk_builder is able to correctly rebuild with given metadata from streaming chunks Checks if stream_chunk_builder is able to correctly rebuild with given metadata from streaming chunks
""" """
from litellm.types.utils import Usage
messages = [ messages = [
{"role": "user", "content": "Tell me the funniest joke you know."}, {"role": "user", "content": "Tell me the funniest joke you know."},
{ {
@ -182,24 +184,28 @@ def test_stream_chunk_builder_litellm_usage_chunks():
{"role": "assistant", "content": "uhhhh\n\n\nhmmmm.....\nthinking....\n"}, {"role": "assistant", "content": "uhhhh\n\n\nhmmmm.....\nthinking....\n"},
{"role": "user", "content": "\nI am waiting...\n\n...\n"}, {"role": "user", "content": "\nI am waiting...\n\n...\n"},
] ]
# make a regular gemini call
response = completion(
model="gemini/gemini-1.5-flash",
messages=messages,
)
usage: litellm.Usage = response.usage usage: litellm.Usage = Usage(
completion_tokens=27,
prompt_tokens=55,
total_tokens=82,
completion_tokens_details=None,
prompt_tokens_details=None,
)
gemini_pt = usage.prompt_tokens gemini_pt = usage.prompt_tokens
# make a streaming gemini call # make a streaming gemini call
response = completion( try:
model="gemini/gemini-1.5-flash", response = completion(
messages=messages, model="gemini/gemini-1.5-flash",
stream=True, messages=messages,
complete_response=True, stream=True,
stream_options={"include_usage": True}, complete_response=True,
) stream_options={"include_usage": True},
)
except litellm.InternalServerError as e:
pytest.skip(f"Skipping test due to internal server error - {str(e)}")
usage: litellm.Usage = response.usage usage: litellm.Usage = response.usage