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

2026-08-13 06:23:32 +00:00 · 2023-12-29 15:27:20 +05:30
parent 1f76b0e721
commit cbdfae1267
1 changed files with 11 additions and 8 deletions
@@ -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