Refactor AI handler instantiation in PRAgent and related classes

This commit is contained in:
mrT23
2023-12-14 08:53:22 +02:00
parent e37598fdca
commit 3531016a2c
11 changed files with 15 additions and 51 deletions

View File

@ -14,7 +14,7 @@ from pr_agent.log import get_logger
class PRCodeSuggestions:
def __init__(self, pr_url: str, cli_mode=False, args: list = None, ai_handler: BaseAiHandler = None ):
def __init__(self, pr_url: str, cli_mode=False, args: list = None, ai_handler: BaseAiHandler = None):
self.git_provider = get_git_provider()(pr_url)
self.main_language = get_main_pr_language(

View File

@ -4,7 +4,7 @@ from typing import List, Tuple
from jinja2 import Environment, StrictUndefined
from pr_agent.algo.ai_handler import AiHandler
from pr_agent.algo.ai_handlers.base_ai_handler import BaseAiHandler
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 @@ from pr_agent.log import get_logger
class PRGenerateLabels:
def __init__(self, pr_url: str, args: list = None):
def __init__(self, pr_url: str, args: list = None, ai_handler: BaseAiHandler = None):
"""
Initialize the PRGenerateLabels object with the necessary attributes and objects for generating labels
corresponding to the PR using an AI model.
@ -31,7 +31,7 @@ class PRGenerateLabels:
self.pr_id = self.git_provider.get_pr_id()
# Initialize the AI handler
self.ai_handler = AiHandler()
self.ai_handler = ai_handler
# Initialize the variables dictionary
self.vars = {