From 2d8d7e35696756b581db914d1cb69a720ef35a46 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 6 Jan 2024 17:05:55 +0530 Subject: [PATCH] perf(router.py): don't use asyncio.wait for - just pass it to the completion call for timeouts --- litellm/router.py | 16 ++++++---------- 1 file changed, 6 insertions(+), 10 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 770098df0c..39c0c4b56f 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -352,18 +352,16 @@ class Router: else: model_client = potential_model_client self.total_calls[model_name] += 1 - response = await asyncio.wait_for( - litellm.acompletion( + response = await litellm.acompletion( **{ **data, "messages": messages, "caching": self.cache_responses, "client": model_client, + "timeout": self.timeout, **kwargs, } - ), - timeout=self.timeout, - ) + ) self.success_calls[model_name] += 1 return response except Exception as e: @@ -614,18 +612,16 @@ class Router: else: model_client = potential_model_client self.total_calls[model_name] += 1 - response = await asyncio.wait_for( - litellm.atext_completion( + response = await litellm.atext_completion( **{ **data, "prompt": prompt, "caching": self.cache_responses, "client": model_client, + "timeout": self.timeout, **kwargs, } - ), - timeout=self.timeout, - ) + ) self.success_calls[model_name] += 1 return response except Exception as e: