From 19ab0614c4ba2c5353695bdb6a0a575c728f598b Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 30 Jul 2024 22:49:33 -0700 Subject: [PATCH] fix predibase tests --- litellm/tests/test_streaming.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/litellm/tests/test_streaming.py b/litellm/tests/test_streaming.py index d708c7155a3..f5b2f363b26 100644 --- a/litellm/tests/test_streaming.py +++ b/litellm/tests/test_streaming.py @@ -510,10 +510,10 @@ def test_completion_azure_stream(): async def test_completion_predibase_streaming(sync_mode): try: litellm.set_verbose = True - if sync_mode: response = completion( model="predibase/llama-3-8b-instruct", + timeout=5, tenant_id="c4768f95", max_tokens=10, api_base="https://serving.app.predibase.com", @@ -540,6 +540,7 @@ async def test_completion_predibase_streaming(sync_mode): response = await litellm.acompletion( model="predibase/llama-3-8b-instruct", tenant_id="c4768f95", + timeout=5, max_tokens=10, api_base="https://serving.app.predibase.com", api_key=os.getenv("PREDIBASE_API_KEY"),