Merge remote-tracking branch 'origin/main' into tr/user_description

pr_agent/algo/ai_handlers/langchain_ai_handler.py

@@ -1,5 +1,5 @@
 try:
-    from langchain.chat_models import ChatOpenAI
+    from langchain.chat_models import ChatOpenAI, AzureChatOpenAI
     from langchain.schema import SystemMessage, HumanMessage
 except: # we don't enforce langchain as a dependency, so if it's not installed, just move on
     pass
@@ -9,23 +9,41 @@ from pr_agent.config_loader import get_settings
 from pr_agent.log import get_logger
 
 from openai.error import APIError, RateLimitError, Timeout, TryAgain
-from retry import retry    
+from retry import retry
+import functools
 
 OPENAI_RETRIES = 5
 
 class LangChainOpenAIHandler(BaseAiHandler):
     def __init__(self):
         # Initialize OpenAIHandler specific attributes here
+        super().__init__()
+        self.azure = get_settings().get("OPENAI.API_TYPE", "").lower() == "azure"
         try:
-            super().__init__()            
-            self._chat = ChatOpenAI(openai_api_key=get_settings().openai.key)
-
+            if self.azure:
+                # using a partial function so we can set the deployment_id later to support fallback_deployments
+                # but still need to access the other settings now so we can raise a proper exception if they're missing
+                self._chat = functools.partial(
+                    lambda **kwargs: AzureChatOpenAI(**kwargs),
+                    openai_api_key=get_settings().openai.key,
+                    openai_api_base=get_settings().openai.api_base,
+                    openai_api_version=get_settings().openai.api_version,
+                )
+            else:
+                self._chat = ChatOpenAI(openai_api_key=get_settings().openai.key)
         except AttributeError as e:
-            raise ValueError("OpenAI key is required") from e
+            if getattr(e, "name"):
+                raise ValueError(f"OpenAI {e.name} is required") from e
+            else:
+                raise e
     
     @property
     def chat(self):
-        return self._chat
+        if self.azure:
+            # we must set the deployment_id only here (instead of the __init__ method) to support fallback_deployments
+            return self._chat(deployment_name=self.deployment_id)
+        else:
+            return self._chat
 
     @property
     def deployment_id(self):

pr_agent/algo/ai_handlers/litellm_ai_handler.py

@@ -101,11 +101,6 @@ class LiteLLMAIHandler(BaseAiHandler):
         """
         try:
             deployment_id = self.deployment_id
-            if get_settings().config.verbosity_level >= 2:
-                get_logger().debug(
-                    f"Generating completion with {model}"
-                    f"{(' from deployment ' + deployment_id) if deployment_id else ''}"
-                )
             if self.azure:
                 model = 'azure/' + model
             messages = [{"role": "system", "content": system}, {"role": "user", "content": user}]

pr_agent/algo/pr_processing.py

@@ -226,6 +226,11 @@ async def retry_with_fallback_models(f: Callable):
     # try each (model, deployment_id) pair until one is successful, otherwise raise exception
     for i, (model, deployment_id) in enumerate(zip(all_models, all_deployments)):
         try:
+            if get_settings().config.verbosity_level >= 2:
+                get_logger().debug(
+                    f"Generating prediction with {model}"
+                    f"{(' from deployment ' + deployment_id) if deployment_id else ''}"
+                )
             get_settings().set("openai.deployment_id", deployment_id)
             return await f(model)
         except Exception as e:

pr_agent/settings/configuration.toml

@@ -70,6 +70,7 @@ include_improved_code = true
 extra_instructions = ""
 rank_suggestions = false
 # params for '/improve --extended' mode
+auto_extended_mode=false
 num_code_suggestions_per_chunk=8
 rank_extended_suggestions = true
 max_number_of_calls = 5

pr_agent/tools/pr_code_suggestions.py

@@ -26,7 +26,7 @@ class PRCodeSuggestions:
 
         # extended mode
         try:
-            self.is_extended = any(["extended" in arg for arg in args])
+            self.is_extended = self._get_is_extended(args or [])
         except:
             self.is_extended = False
         if self.is_extended:
@@ -206,6 +206,16 @@ class PRCodeSuggestions:
 
         return new_code_snippet
 
+    def _get_is_extended(self, args: list[str]) -> bool:
+        """Check if extended mode should be enabled by the `--extended` flag or automatically according to the configuration"""
+        if any(["extended" in arg for arg in args]):
+            get_logger().info("Extended mode is enabled by the `--extended` flag")
+            return True
+        if get_settings().pr_code_suggestions.auto_extended_mode:
+            get_logger().info("Extended mode is enabled automatically based on the configuration toggle")
+            return True
+        return False
+
     async def _prepare_prediction_extended(self, model: str) -> dict:
         get_logger().info('Getting PR diff...')
         patches_diff_list = get_pr_multi_diffs(self.git_provider, self.token_handler, model,
@@ -271,8 +281,14 @@ class PRCodeSuggestions:
                 data_sorted[importance_order - 1] = suggestion_list[suggestion_number - 1]
 
             if get_settings().pr_code_suggestions.final_clip_factor != 1:
-                new_len = int(0.5 + len(data_sorted) * get_settings().pr_code_suggestions.final_clip_factor)
-                data_sorted = data_sorted[:new_len]
+                max_len = max(
+                    len(data_sorted),
+                    get_settings().pr_code_suggestions.num_code_suggestions,
+                    get_settings().pr_code_suggestions.num_code_suggestions_per_chunk,
+                )
+                new_len = int(0.5 + max_len * get_settings().pr_code_suggestions.final_clip_factor)
+                if new_len < len(data_sorted):
+                    data_sorted = data_sorted[:new_len]
         except Exception as e:
             if get_settings().config.verbosity_level >= 1:
                 get_logger().info(f"Could not sort suggestions, error: {e}")

pr_agent/tools/pr_config.py

@@ -7,7 +7,7 @@ class PRConfig:
     """
     The PRConfig class is responsible for listing all configuration options available for the user.
     """
-    def __init__(self, pr_url: str, args=None):
+    def __init__(self, pr_url: str, args=None, ai_handler=None):
         """
         Initialize the PRConfig object with the necessary attributes and objects to comment on a pull request.