From 9156b7448aef8932517af744ec3bb20c4edc5f82 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 11 May 2024 13:08:16 -0700 Subject: [PATCH] feat - router async batch acompletion --- litellm/router.py | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/litellm/router.py b/litellm/router.py index f0d94908e..7396dab20 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -605,6 +605,33 @@ class Router: self.fail_calls[model_name] += 1 raise e + async def abatch_completion( + self, models: List[str], messages: List[Dict[str, str]], **kwargs + ): + + async def _async_completion_no_exceptions( + model: str, messages: List[Dict[str, str]], **kwargs + ): + """ + Wrapper around self.async_completion that catches exceptions and returns them as a result + """ + try: + return await self.acompletion(model=model, messages=messages, **kwargs) + except Exception as e: + return e + + _tasks = [] + for model in models: + # add each task but if the task fails + _tasks.append( + _async_completion_no_exceptions( + model=model, messages=messages, **kwargs + ) + ) + + response = await asyncio.gather(*_tasks) + return response + def image_generation(self, prompt: str, model: str, **kwargs): try: kwargs["model"] = model