improve --extend

pr_agent/tools/pr_code_suggestions.py

@@ -2,11 +2,13 @@ import copy
 import json
 import logging
 import textwrap
+from typing import List
 
+import yaml
 from jinja2 import Environment, StrictUndefined
 
 from pr_agent.algo.ai_handler import AiHandler
-from pr_agent.algo.pr_processing import get_pr_diff, retry_with_fallback_models
+from pr_agent.algo.pr_processing import get_pr_diff, retry_with_fallback_models, get_pr_multi_diffs
 from pr_agent.algo.token_handler import TokenHandler
 from pr_agent.algo.utils import try_fix_json
 from pr_agent.config_loader import get_settings
@@ -22,6 +24,13 @@ class PRCodeSuggestions:
             self.git_provider.get_languages(), self.git_provider.get_files()
         )
 
+        # extended mode
+        self.is_extended = any(["extended" in arg for arg in args])
+        if self.is_extended:
+            num_code_suggestions = get_settings().pr_code_suggestions.num_code_suggestions_per_chunk
+        else:
+            num_code_suggestions = get_settings().pr_code_suggestions.num_code_suggestions
+
         self.ai_handler = AiHandler()
         self.patches_diff = None
         self.prediction = None
@@ -32,7 +41,7 @@ class PRCodeSuggestions:
             "description": self.git_provider.get_pr_description(),
             "language": self.main_language,
             "diff": "",  # empty diff for initial calculation
-            "num_code_suggestions": get_settings().pr_code_suggestions.num_code_suggestions,
+            "num_code_suggestions": num_code_suggestions,
             "extra_instructions": get_settings().pr_code_suggestions.extra_instructions,
             "commit_messages_str": self.git_provider.get_commit_messages(),
         }
@@ -42,14 +51,22 @@ class PRCodeSuggestions:
                                           get_settings().pr_code_suggestions_prompt.user)
 
     async def run(self):
-        # assert type(self.git_provider) != BitbucketProvider, "Bitbucket is not supported for now"
-
         logging.info('Generating code suggestions for PR...')
         if get_settings().config.publish_output:
             self.git_provider.publish_comment("Preparing review...", is_temporary=True)
-        await retry_with_fallback_models(self._prepare_prediction)
+
         logging.info('Preparing PR review...')
-        data = self._prepare_pr_code_suggestions()
+        if not self.is_extended:
+            await retry_with_fallback_models(self._prepare_prediction)
+            data = self._prepare_pr_code_suggestions()
+        else:
+            data = await retry_with_fallback_models(self._prepare_prediction_extended)
+
+        if (not self.is_extended and get_settings().pr_code_suggestions.rank_suggestions) or \
+                (self.is_extended and get_settings().pr_code_suggestions.rank_extended_suggestions):
+            logging.info('Ranking Suggestions...')
+            data['Code suggestions'] = await self.rank_suggestions(data['Code suggestions'])
+
         if get_settings().config.publish_output:
             logging.info('Pushing PR review...')
             self.git_provider.remove_initial_comment()
@@ -145,3 +162,81 @@ class PRCodeSuggestions:
 
         return new_code_snippet
 
+    async def _prepare_prediction_extended(self, model: str) -> dict:
+        logging.info('Getting PR diff...')
+        patches_diff_list = get_pr_multi_diffs(self.git_provider, self.token_handler, model,
+                                               max_calls=get_settings().pr_code_suggestions.max_number_of_calls)
+
+        logging.info('Getting multi AI predictions...')
+        prediction_list = []
+        for i, patches_diff in enumerate(patches_diff_list):
+            logging.info(f"Processing chunk {i + 1} of {len(patches_diff_list)}")
+            self.patches_diff = patches_diff
+            prediction = await self._get_prediction(model)
+            prediction_list.append(prediction)
+        self.prediction_list = prediction_list
+
+        data = {}
+        for prediction in prediction_list:
+            self.prediction = prediction
+            data_per_chunk = self._prepare_pr_code_suggestions()
+            if "Code suggestions" in data:
+                data["Code suggestions"].extend(data_per_chunk["Code suggestions"])
+            else:
+                data.update(data_per_chunk)
+        self.data = data
+        return data
+
+    async def rank_suggestions(self, data: List) -> List:
+        """
+        Call a model to rank (sort) code suggestions based on their importance order.
+
+        Args:
+            data (List): A list of code suggestions to be ranked.
+
+        Returns:
+            List: The ranked list of code suggestions.
+        """
+
+        suggestion_list = []
+        # remove invalid suggestions
+        for i, suggestion in enumerate(data):
+            if suggestion['existing code'] != suggestion['improved code']:
+                suggestion_list.append(suggestion)
+
+        data_sorted = [[]] * len(suggestion_list)
+
+        try:
+            suggestion_str = ""
+            for i, suggestion in enumerate(suggestion_list):
+                suggestion_str += f"suggestion {i + 1}: " + str(suggestion) + '\n\n'
+
+            variables = {'suggestion_list': suggestion_list, 'suggestion_str': suggestion_str}
+            model = get_settings().config.model
+            environment = Environment(undefined=StrictUndefined)
+            system_prompt = environment.from_string(get_settings().pr_sort_code_suggestions_prompt.system).render(
+                variables)
+            user_prompt = environment.from_string(get_settings().pr_sort_code_suggestions_prompt.user).render(variables)
+            if get_settings().config.verbosity_level >= 2:
+                logging.info(f"\nSystem prompt:\n{system_prompt}")
+                logging.info(f"\nUser prompt:\n{user_prompt}")
+            response, finish_reason = await self.ai_handler.chat_completion(model=model, system=system_prompt,
+                                                                            user=user_prompt)
+
+            sort_order = yaml.safe_load(response)
+            for s in sort_order['Sort Order']:
+                suggestion_number = s['suggestion number']
+                importance_order = s['importance order']
+                data_sorted[importance_order - 1] = suggestion_list[suggestion_number - 1]
+
+            if get_settings().pr_extendeted_code_suggestions.final_clip_factor != 1:
+                new_len = int(0.5 + len(data_sorted) * get_settings().pr_extendeted_code_suggestions.final_clip_factor)
+                data_sorted = data_sorted[:new_len]
+        except Exception as e:
+            if get_settings().config.verbosity_level >= 1:
+                logging.info(f"Could not sort suggestions, error: {e}")
+            data_sorted = suggestion_list
+
+        return data_sorted
+
+