Abstract AiHandler to BaseAiHandler

pr_agent/algo/ai_handler.py

@@ -7,12 +7,13 @@
 from openai.error import APIError, RateLimitError, Timeout, TryAgain
 from retry import retry
 from pr_agent.config_loader import get_settings
+from pr_agent.algo.base_ai_handler import BaseAiHandler
 from pr_agent.log import get_logger
 
 OPENAI_RETRIES = 5
 
 
-class AiHandler:
+class AiHandler(BaseAiHandler):
  """
  This class handles interactions with the OpenAI API for chat completions.
  It initializes the API key and other settings from a configuration file,

pr_agent/algo/base_ai_handler.py

@@ -0,0 +1,20 @@
+from abc import ABC, abstractmethod
+
+class BaseAiHandler(ABC):
+ """
+ This class defines the interface for an AI handler.
+ """
+
+ @abstractmethod
+ def __init__(self):
+ pass
+
+ @property
+ @abstractmethod
+ def deployment_id(self):
+ pass
+
+ @abstractmethod
+ async def chat_completion(self, model: str, system: str, user: str, temperature: float = 0.2):
+ pass
+

pr_agent/tools/pr_code_suggestions.py

@@ -3,7 +3,7 @@
 from typing import Dict, List
 from jinja2 import Environment, StrictUndefined
 
-from pr_agent.algo.ai_handler import AiHandler
+from pr_agent.algo.ai_handler import BaseAiHandler, AiHandler
 from pr_agent.algo.pr_processing import get_pr_diff, get_pr_multi_diffs, retry_with_fallback_models
 from pr_agent.algo.token_handler import TokenHandler
 from pr_agent.algo.utils import load_yaml
@@ -14,7 +14,7 @@
 
 
 class PRCodeSuggestions:
- def __init__(self, pr_url: str, cli_mode=False, args: list = None):
+ def __init__(self, pr_url: str, cli_mode=False, args: list = None, ai_handler: BaseAiHandler = AiHandler()):
 
  self.git_provider = get_git_provider()(pr_url)
  self.main_language = get_main_pr_language(
@@ -31,7 +31,7 @@ def __init__(self, pr_url: str, cli_mode=False, args: list = None):
  else:
  num_code_suggestions = get_settings().pr_code_suggestions.num_code_suggestions
 
- self.ai_handler = AiHandler()
+ self.ai_handler = ai_handler
  self.patches_diff = None
  self.prediction = None
  self.cli_mode = cli_mode

pr_agent/tools/pr_description.py

@@ -4,7 +4,7 @@
 
 from jinja2 import Environment, StrictUndefined
 
-from pr_agent.algo.ai_handler import AiHandler
+from pr_agent.algo.ai_handler import BaseAiHandler, AiHandler
 from pr_agent.algo.pr_processing import get_pr_diff, retry_with_fallback_models
 from pr_agent.algo.token_handler import TokenHandler
 from pr_agent.algo.utils import load_yaml, set_custom_labels, get_user_labels
@@ -15,7 +15,7 @@
 
 
 class PRDescription:
- def __init__(self, pr_url: str, args: list = None):
+ def __init__(self, pr_url: str, args: list = None, ai_handler: BaseAiHandler = AiHandler()):
  """
  Initialize the PRDescription object with the necessary attributes and objects for generating a PR description
  using an AI model.
@@ -36,7 +36,7 @@ def __init__(self, pr_url: str, args: list = None):
  get_settings().pr_description.enable_semantic_files_types = False
 
  # Initialize the AI handler
- self.ai_handler = AiHandler()
+ self.ai_handler = ai_handler
 
  # Initialize the variables dictionary
  self.vars = {

pr_agent/tools/pr_information_from_user.py

@@ -2,7 +2,7 @@
 
 from jinja2 import Environment, StrictUndefined
 
-from pr_agent.algo.ai_handler import AiHandler
+from pr_agent.algo.ai_handler import BaseAiHandler, AiHandler
 from pr_agent.algo.pr_processing import get_pr_diff, retry_with_fallback_models
 from pr_agent.algo.token_handler import TokenHandler
 from pr_agent.config_loader import get_settings
@@ -12,12 +12,12 @@
 
 
 class PRInformationFromUser:
- def __init__(self, pr_url: str, args: list = None):
+ def __init__(self, pr_url: str, args: list = None, ai_handler: BaseAiHandler = AiHandler()):
  self.git_provider = get_git_provider()(pr_url)
  self.main_pr_language = get_main_pr_language(
  self.git_provider.get_languages(), self.git_provider.get_files()
  )
- self.ai_handler = AiHandler()
+ self.ai_handler = ai_handler
  self.vars = {
  "title": self.git_provider.pr.title,
  "branch": self.git_provider.get_pr_branch(),

pr_agent/tools/pr_questions.py

@@ -2,7 +2,7 @@
 
 from jinja2 import Environment, StrictUndefined
 
-from pr_agent.algo.ai_handler import AiHandler
+from pr_agent.algo.ai_handler import BaseAiHandler, AiHandler
 from pr_agent.algo.pr_processing import get_pr_diff, retry_with_fallback_models
 from pr_agent.algo.token_handler import TokenHandler
 from pr_agent.config_loader import get_settings
@@ -12,13 +12,13 @@
 
 
 class PRQuestions:
- def __init__(self, pr_url: str, args=None):
+ def __init__(self, pr_url: str, args=None, ai_handler: BaseAiHandler = AiHandler()):
  question_str = self.parse_args(args)
  self.git_provider = get_git_provider()(pr_url)
  self.main_pr_language = get_main_pr_language(
  self.git_provider.get_languages(), self.git_provider.get_files()
  )
- self.ai_handler = AiHandler()
+ self.ai_handler = ai_handler
  self.question_str = question_str
  self.vars = {
  "title": self.git_provider.pr.title,

pr_agent/tools/pr_reviewer.py

@@ -7,7 +7,7 @@
 from jinja2 import Environment, StrictUndefined
 from yaml import SafeLoader
 
-from pr_agent.algo.ai_handler import AiHandler
+from pr_agent.algo.ai_handler import BaseAiHandler, AiHandler
 from pr_agent.algo.pr_processing import get_pr_diff, retry_with_fallback_models
 from pr_agent.algo.token_handler import TokenHandler
 from pr_agent.algo.utils import convert_to_markdown, load_yaml, try_fix_yaml, set_custom_labels, get_user_labels
@@ -22,7 +22,7 @@ class PRReviewer:
  """
  The PRReviewer class is responsible for reviewing a pull request and generating feedback using an AI model.
  """
- def __init__(self, pr_url: str, is_answer: bool = False, is_auto: bool = False, args: list = None):
+ def __init__(self, pr_url: str, is_answer: bool = False, is_auto: bool = False, args: list = None, ai_handler: BaseAiHandler = AiHandler()):
  """
  Initialize the PRReviewer object with the necessary attributes and objects to review a pull request.
 
@@ -43,7 +43,7 @@ def __init__(self, pr_url: str, is_answer: bool = False, is_auto: bool = False,
 
  if self.is_answer and not self.git_provider.is_supported("get_issue_comments"):
  raise Exception(f"Answer mode is not supported for {get_settings().config.git_provider} for now")
- self.ai_handler = AiHandler()
+ self.ai_handler = ai_handler
  self.patches_diff = None
  self.prediction = None
 

pr_agent/tools/pr_update_changelog.py

@@ -5,7 +5,7 @@
 
 from jinja2 import Environment, StrictUndefined
 
-from pr_agent.algo.ai_handler import AiHandler
+from pr_agent.algo.ai_handler import BaseAiHandler, AiHandler
 from pr_agent.algo.pr_processing import get_pr_diff, retry_with_fallback_models
 from pr_agent.algo.token_handler import TokenHandler
 from pr_agent.config_loader import get_settings
@@ -17,15 +17,15 @@
 
 
 class PRUpdateChangelog:
- def __init__(self, pr_url: str, cli_mode=False, args=None):
+ def __init__(self, pr_url: str, cli_mode=False, args=None, ai_handler: BaseAiHandler = AiHandler()):
 
  self.git_provider = get_git_provider()(pr_url)
  self.main_language = get_main_pr_language(
  self.git_provider.get_languages(), self.git_provider.get_files()
  )
  self.commit_changelog = get_settings().pr_update_changelog.push_changelog_changes
  self._get_changlog_file() # self.changelog_file_str
- self.ai_handler = AiHandler()
+ self.ai_handler = ai_handler
  self.patches_diff = None
  self.prediction = None
  self.cli_mode = cli_mode