phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

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

Browse source
This commit is contained in:
Krrish Dholakia 2023-12-29 15:27:20 +05:30
parent 1f76b0e721
commit cbdfae1267
1 changed files with 11 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

5
litellm/router.py
View file

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