Implement Lazy Initialization for AI Handlers in PR Tools

pr_agent/agent/pr_agent.py

@@ -1,4 +1,6 @@
 import shlex
+from functools import partial
+
 from pr_agent.algo.ai_handlers.base_ai_handler import BaseAiHandler
 from pr_agent.algo.ai_handlers.litellm_ai_handler import LiteLLMAIHandler
 
@@ -41,8 +43,8 @@
 commands = list(command2class.keys())
 
 class PRAgent:
- def __init__(self, ai_handler: BaseAiHandler = LiteLLMAIHandler()):
- self.ai_handler = ai_handler
+ def __init__(self, ai_handler: partial[BaseAiHandler,] = LiteLLMAIHandler):
+ self.ai_handler = ai_handler # will be initialized in run_action
 
  async def handle_request(self, pr_url, request, notify=None) -> bool:
  # First, apply repo specific settings if exists

pr_agent/tools/pr_add_docs.py

@@ -1,5 +1,6 @@
 import copy
 import textwrap
+from functools import partial
 from typing import Dict
 
 from jinja2 import Environment, StrictUndefined
@@ -17,14 +18,14 @@
 
 class PRAddDocs:
  def __init__(self, pr_url: str, cli_mode=False, args: list = None,
- ai_handler: BaseAiHandler = LiteLLMAIHandler()):
+ ai_handler: partial[BaseAiHandler,] = LiteLLMAIHandler):
 
  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.ai_handler = ai_handler
+ self.ai_handler = ai_handler()
  self.patches_diff = None
  self.prediction = None
  self.cli_mode = cli_mode

pr_agent/tools/pr_code_suggestions.py

@@ -1,5 +1,6 @@
 import copy
 import textwrap
+from functools import partial
 from typing import Dict, List
 from jinja2 import Environment, StrictUndefined
 
@@ -16,7 +17,7 @@
 
 class PRCodeSuggestions:
  def __init__(self, pr_url: str, cli_mode=False, args: list = None,
- ai_handler: BaseAiHandler = LiteLLMAIHandler()):
+ ai_handler: partial[BaseAiHandler,] = LiteLLMAIHandler):
 
  self.git_provider = get_git_provider()(pr_url)
  self.main_language = get_main_pr_language(
@@ -33,7 +34,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 = ai_handler
+ self.ai_handler = ai_handler()
  self.patches_diff = None
  self.prediction = None
  self.cli_mode = cli_mode

pr_agent/tools/pr_description.py

@@ -1,5 +1,6 @@
 import copy
 import re
+from functools import partial
 from typing import List, Tuple
 
 from jinja2 import Environment, StrictUndefined
@@ -17,7 +18,7 @@
 
 class PRDescription:
  def __init__(self, pr_url: str, args: list = None,
- ai_handler: BaseAiHandler = LiteLLMAIHandler()):
+ ai_handler: partial[BaseAiHandler,] = LiteLLMAIHandler):
  """
  Initialize the PRDescription object with the necessary attributes and objects for generating a PR description
  using an AI model.
@@ -38,7 +39,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 = ai_handler
+ self.ai_handler = ai_handler()
 
  # Initialize the variables dictionary
  self.vars = {

pr_agent/tools/pr_generate_labels.py

@@ -1,5 +1,6 @@
 import copy
 import re
+from functools import partial
 from typing import List, Tuple
 
 from jinja2 import Environment, StrictUndefined
@@ -17,7 +18,7 @@
 
 class PRGenerateLabels:
  def __init__(self, pr_url: str, args: list = None,
- ai_handler: BaseAiHandler = LiteLLMAIHandler()):
+ ai_handler: partial[BaseAiHandler,] = LiteLLMAIHandler):
  """
  Initialize the PRGenerateLabels object with the necessary attributes and objects for generating labels
  corresponding to the PR using an AI model.
@@ -33,7 +34,7 @@ def __init__(self, pr_url: str, args: list = None,
  self.pr_id = self.git_provider.get_pr_id()
 
  # Initialize the AI handler
- self.ai_handler = ai_handler
+ self.ai_handler = ai_handler()
 
  # Initialize the variables dictionary
  self.vars = {

pr_agent/tools/pr_information_from_user.py

@@ -1,4 +1,5 @@
 import copy
+from functools import partial
 
 from jinja2 import Environment, StrictUndefined
 
@@ -14,12 +15,12 @@
 
 class PRInformationFromUser:
  def __init__(self, pr_url: str, args: list = None,
- ai_handler: BaseAiHandler = LiteLLMAIHandler()):
+ ai_handler: partial[BaseAiHandler,] = LiteLLMAIHandler):
  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 = ai_handler
+ 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

@@ -1,4 +1,5 @@
 import copy
+from functools import partial
 
 from jinja2 import Environment, StrictUndefined
 
@@ -13,13 +14,13 @@
 
 
 class PRQuestions:
- def __init__(self, pr_url: str, args=None, ai_handler: BaseAiHandler = LiteLLMAIHandler()):
+ def __init__(self, pr_url: str, args=None, ai_handler: partial[BaseAiHandler,] = LiteLLMAIHandler):
  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 = ai_handler
+ self.ai_handler = ai_handler()
  self.question_str = question_str
  self.vars = {
  "title": self.git_provider.pr.title,

pr_agent/tools/pr_reviewer.py

@@ -1,6 +1,7 @@
 import copy
 import datetime
 from collections import OrderedDict
+from functools import partial
 from typing import List, Tuple
 
 import yaml
@@ -24,7 +25,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,
- ai_handler: BaseAiHandler = LiteLLMAIHandler()):
+ ai_handler: partial[BaseAiHandler,] = LiteLLMAIHandler):
  """
  Initialize the PRReviewer object with the necessary attributes and objects to review a pull request.
 
@@ -47,7 +48,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 = ai_handler
+ self.ai_handler = ai_handler()
  self.patches_diff = None
  self.prediction = None
 

pr_agent/tools/pr_update_changelog.py

@@ -1,5 +1,6 @@
 import copy
 from datetime import date
+from functools import partial
 from time import sleep
 from typing import Tuple
 
@@ -18,15 +19,15 @@
 
 
 class PRUpdateChangelog:
- def __init__(self, pr_url: str, cli_mode=False, args=None, ai_handler: BaseAiHandler = LiteLLMAIHandler()):
+ def __init__(self, pr_url: str, cli_mode=False, args=None, ai_handler: partial[BaseAiHandler,] = LiteLLMAIHandler):
 
  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 = ai_handler
+ self.ai_handler = ai_handler()
  self.patches_diff = None
  self.prediction = None
  self.cli_mode = cli_mode