Load default AI Handler from util function

This commit is contained in:
Brian Pham
2023-12-12 21:51:05 +08:00
parent ebf7027aab
commit 5239e1c3e9
9 changed files with 55 additions and 29 deletions

View File

@ -2,9 +2,10 @@ import copy
from jinja2 import Environment, StrictUndefined
from pr_agent.algo.litellm_ai_handler import BaseAiHandler, LiteLLMAiHandler
from pr_agent.algo.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 get_ai_handler
from pr_agent.config_loader import get_settings
from pr_agent.git_providers import get_git_provider
from pr_agent.git_providers.git_provider import get_main_pr_language
@ -12,7 +13,7 @@ from pr_agent.log import get_logger
class PRInformationFromUser:
def __init__(self, pr_url: str, args: list = None, ai_handler: BaseAiHandler = LiteLLMAiHandler()):
def __init__(self, pr_url: str, args: list = None, ai_handler: BaseAiHandler = get_ai_handler()):
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()