fix(router.py): support wait_for for async completion calls

This commit is contained in:
Krrish Dholakia 2023-12-29 15:27:20 +05:30
parent 1f76b0e721
commit cbdfae1267

View file

@ -332,7 +332,8 @@ class Router:
else:
model_client = potential_model_client
self.total_calls[model_name] += 1
response = await litellm.acompletion(
response = await asyncio.wait_for(
litellm.acompletion(
**{
**data,
"messages": messages,
@ -340,6 +341,8 @@ class Router:
"client": model_client,
**kwargs,
}
),
timeout=self.timeout,
)
self.success_calls[model_name] += 1
return response