From 8f751f7371569a503948ec6188b16285e0dc65bb Mon Sep 17 00:00:00 2001
From: Ori Kotek <ori.k@codium.ai>
Date: Mon, 7 Aug 2023 13:26:28 +0300
Subject: [PATCH] Default timeout for AI is now 180s, configurable

---
 pr_agent/algo/ai_handler.py          | 31 +++++++++++++++-------------
 pr_agent/settings/configuration.toml |  1 +
 2 files changed, 18 insertions(+), 14 deletions(-)

diff --git a/pr_agent/algo/ai_handler.py b/pr_agent/algo/ai_handler.py
index 57221518..27e9533a 100644
--- a/pr_agent/algo/ai_handler.py
+++ b/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
\ No newline at end of file
+        return resp, finish_reason
diff --git a/pr_agent/settings/configuration.toml b/pr_agent/settings/configuration.toml
index cb9af775..8334049d 100644
--- a/pr_agent/settings/configuration.toml
+++ b/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