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
|
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,17 +184,19 @@ 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
|
||||||
|
try:
|
||||||
response = completion(
|
response = completion(
|
||||||
model="gemini/gemini-1.5-flash",
|
model="gemini/gemini-1.5-flash",
|
||||||
messages=messages,
|
messages=messages,
|
||||||
|
@ -200,6 +204,8 @@ def test_stream_chunk_builder_litellm_usage_chunks():
|
||||||
complete_response=True,
|
complete_response=True,
|
||||||
stream_options={"include_usage": 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
|
||||||
|
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue