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

19
litellm/router.py
View file

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