Default timeout for AI is now 180s, configurable

pr_agent/algo/ai_handler.py

@@ -1,13 +1,15 @@
 import logging
 
+import litellm
 import openai
+from litellm import acompletion
 from openai.error import APIError, RateLimitError, Timeout, TryAgain
 from retry import retry
-import litellm
-from litellm import acompletion
+
 from pr_agent.config_loader import get_settings
-import traceback
-OPENAI_RETRIES=5
+
+OPENAI_RETRIES = 5
+
 
 class AiHandler:
     """
@@ -69,15 +71,16 @@ class AiHandler:
         """
         try:
             response = await acompletion(
-                            model=model,
-                            deployment_id=self.deployment_id,
-                            messages=[
-                                {"role": "system", "content": system},
-                                {"role": "user", "content": user}
-                            ],
-                            temperature=temperature,
-                            azure=self.azure
-                        )
+                model=model,
+                deployment_id=self.deployment_id,
+                messages=[
+                    {"role": "system", "content": system},
+                    {"role": "user", "content": user}
+                ],
+                temperature=temperature,
+                azure=self.azure,
+                force_timeout=get_settings().config.ai_timeout
+            )
         except (APIError, Timeout, TryAgain) as e:
             logging.error("Error during OpenAI inference: ", e)
             raise
@@ -92,4 +95,4 @@ class AiHandler:
         resp = response["choices"][0]['message']['content']
         finish_reason = response["choices"][0]["finish_reason"]
         print(resp, finish_reason)
-        return resp, finish_reason
+        return resp, finish_reason

pr_agent/settings/configuration.toml

@@ -7,6 +7,7 @@ publish_output_progress=true
 verbosity_level=0 # 0,1,2
 use_extra_bad_extensions=false
 use_repo_settings_file=true
+ai_timeout=180
 
 [pr_reviewer] # /review #
 require_focused_review=true