test: update test to handle model overloaded error

This commit is contained in:
Krrish Dholakia 2024-11-20 22:02:33 +05:30
parent f6ebba7538
commit 361d010068

View file

@ -4490,19 +4490,22 @@ async def test_dynamic_azure_params(stream, sync_mode):
@pytest.mark.flaky(retries=3, delay=1) @pytest.mark.flaky(retries=3, delay=1)
async def test_completion_ai21_chat(): async def test_completion_ai21_chat():
litellm.set_verbose = True litellm.set_verbose = True
response = await litellm.acompletion( try:
model="jamba-1.5-large", response = await litellm.acompletion(
user="ishaan", model="jamba-1.5-large",
tool_choice="auto", user="ishaan",
seed=123, tool_choice="auto",
messages=[{"role": "user", "content": "what does the document say"}], seed=123,
documents=[ messages=[{"role": "user", "content": "what does the document say"}],
{ documents=[
"content": "hello world", {
"metadata": {"source": "google", "author": "ishaan"}, "content": "hello world",
} "metadata": {"source": "google", "author": "ishaan"},
], }
) ],
)
except litellm.InternalServerError:
pytest.skip("Model is overloaded")
@pytest.mark.parametrize( @pytest.mark.parametrize(