Abstract AiHandler to BaseAiHandler

This commit is contained in:
Brian Pham
2023-12-09 16:39:25 +00:00
parent 8d075b76ae
commit f2abe5c73e
8 changed files with 69 additions and 21 deletions

View File

@ -6,7 +6,7 @@ from typing import Tuple
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,7 +17,7 @@ CHANGELOG_LINES = 50
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(
@ -25,7 +25,7 @@ class PRUpdateChangelog:
)
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