forked from phoenix/litellm-mirror
test: handle gemini error
This commit is contained in:
parent
8b4c33c398
commit
7d4829f908
1 changed files with 19 additions and 13 deletions
|
@ -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
|
||||
"""
|
||||
from litellm.types.utils import Usage
|
||||
|
||||
messages = [
|
||||
{"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": "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
|
||||
|
||||
# make a streaming gemini call
|
||||
response = completion(
|
||||
model="gemini/gemini-1.5-flash",
|
||||
messages=messages,
|
||||
stream=True,
|
||||
complete_response=True,
|
||||
stream_options={"include_usage": True},
|
||||
)
|
||||
try:
|
||||
response = completion(
|
||||
model="gemini/gemini-1.5-flash",
|
||||
messages=messages,
|
||||
stream=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
|
||||
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue