Refactor AI handler classes

This commit is contained in:
Brian Pham
2023-12-10 00:25:25 +08:00
parent b6cabda586
commit 7e47baa9db
9 changed files with 71 additions and 23 deletions

View File

@ -4,7 +4,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
@ -15,7 +15,7 @@ from pr_agent.log import get_logger
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(
@ -32,7 +32,7 @@ class PRCodeSuggestions:
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