diff --git a/pr_agent/algo/ai_handler.py b/pr_agent/algo/ai_handler.py index c3989563..3e4dc18a 100644 --- a/pr_agent/algo/ai_handler.py +++ b/pr_agent/algo/ai_handler.py @@ -7,11 +7,11 @@ from openai.error import APIError, RateLimitError, Timeout, TryAgain from retry import retry from pr_agent.config_loader import get_settings from pr_agent.log import get_logger - +from pr_agent.algo.base_ai_handler import BaseAiHandler OPENAI_RETRIES = 5 -class AiHandler: +class AiHandler(BaseAiHandler): """ This class handles interactions with the OpenAI API for chat completions. It initializes the API key and other settings from a configuration file, diff --git a/pr_agent/algo/base_ai_handler.py b/pr_agent/algo/base_ai_handler.py new file mode 100644 index 00000000..9cb40b5c --- /dev/null +++ b/pr_agent/algo/base_ai_handler.py @@ -0,0 +1,48 @@ +from abc import ABC, abstractmethod + +class BaseAiHandler(ABC): + """ + This class defines the interface for an AI handler. + """ + + @abstractmethod + def __init__(self): + pass + + @property + @abstractmethod + def deployment_id(self): + pass + + @abstractmethod + async def chat_completion(self, model: str, system: str, user: str, temperature: float = 0.2): + pass + + +class AiHandler(BaseAiHandler): + """ + This class handles interactions with the OpenAI API for chat completions. + It initializes the API key and other settings from a configuration file, + and provides a method for performing chat completions using the OpenAI ChatCompletion API. + """ + + # ... rest of your code ... + + +class CustomAiHandler(BaseAiHandler): + """ + This class is your custom AI handler that uses a different LLM library. + """ + + def __init__(self): + # Initialize your custom AI handler + pass + + @property + def deployment_id(self): + # Return the deployment ID for your custom AI handler + pass + + async def chat_completion(self, model: str, system: str, user: str, temperature: float = 0.2): + # Implement the chat completion method for your custom AI handler + pass \ No newline at end of file diff --git a/pr_agent/tools/pr_add_docs.py b/pr_agent/tools/pr_add_docs.py index eec75b9c..04cc53f4 100644 --- a/pr_agent/tools/pr_add_docs.py +++ b/pr_agent/tools/pr_add_docs.py @@ -4,7 +4,7 @@ from typing import Dict 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.algo.utils import load_yaml @@ -15,14 +15,14 @@ from pr_agent.log import get_logger class PRAddDocs: - 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( self.git_provider.get_languages(), self.git_provider.get_files() ) - self.ai_handler = AiHandler() + self.ai_handler = ai_handler self.patches_diff = None self.prediction = None self.cli_mode = cli_mode diff --git a/pr_agent/tools/pr_code_suggestions.py b/pr_agent/tools/pr_code_suggestions.py index 9e8d7f15..00e58a21 100644 --- a/pr_agent/tools/pr_code_suggestions.py +++ b/pr_agent/tools/pr_code_suggestions.py @@ -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 diff --git a/pr_agent/tools/pr_description.py b/pr_agent/tools/pr_description.py index c1bd03fd..9b02cd35 100644 --- a/pr_agent/tools/pr_description.py +++ b/pr_agent/tools/pr_description.py @@ -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_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.algo.utils import load_yaml @@ -15,7 +15,7 @@ from pr_agent.log import get_logger class PRDescription: - def __init__(self, pr_url: str, args: list = None): + def __init__(self, pr_url: str, args: list = None, ai_handler: BaseAiHandler = AiHandler()): """ Initialize the PRDescription object with the necessary attributes and objects for generating a PR description using an AI model. @@ -31,7 +31,7 @@ class PRDescription: 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 = { diff --git a/pr_agent/tools/pr_information_from_user.py b/pr_agent/tools/pr_information_from_user.py index 059966e1..27c77180 100644 --- a/pr_agent/tools/pr_information_from_user.py +++ b/pr_agent/tools/pr_information_from_user.py @@ -2,7 +2,7 @@ import copy 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 @@ -12,12 +12,12 @@ from pr_agent.log import get_logger class PRInformationFromUser: - def __init__(self, pr_url: str, args: list = None): + def __init__(self, pr_url: str, args: list = None, ai_handler: BaseAiHandler = AiHandler()): 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 = AiHandler() + self.ai_handler = ai_handler self.vars = { "title": self.git_provider.pr.title, "branch": self.git_provider.get_pr_branch(), diff --git a/pr_agent/tools/pr_questions.py b/pr_agent/tools/pr_questions.py index 7740fd4a..4aec3edf 100644 --- a/pr_agent/tools/pr_questions.py +++ b/pr_agent/tools/pr_questions.py @@ -2,7 +2,7 @@ import copy 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 @@ -12,13 +12,13 @@ from pr_agent.log import get_logger class PRQuestions: - def __init__(self, pr_url: str, args=None): + def __init__(self, pr_url: str, args=None, ai_handler: BaseAiHandler = AiHandler()): 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 = AiHandler() + self.ai_handler = ai_handler self.question_str = question_str self.vars = { "title": self.git_provider.pr.title, diff --git a/pr_agent/tools/pr_reviewer.py b/pr_agent/tools/pr_reviewer.py index ed99ddf6..c78a11e8 100644 --- a/pr_agent/tools/pr_reviewer.py +++ b/pr_agent/tools/pr_reviewer.py @@ -6,7 +6,7 @@ import yaml from jinja2 import Environment, StrictUndefined from yaml import SafeLoader -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.algo.utils import convert_to_markdown, load_yaml, try_fix_yaml @@ -21,7 +21,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): + def __init__(self, pr_url: str, is_answer: bool = False, is_auto: bool = False, args: list = None, ai_handler: BaseAiHandler = AiHandler()): """ Initialize the PRReviewer object with the necessary attributes and objects to review a pull request. @@ -42,7 +42,7 @@ class PRReviewer: 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 = AiHandler() + self.ai_handler = ai_handler self.patches_diff = None self.prediction = None diff --git a/pr_agent/tools/pr_update_changelog.py b/pr_agent/tools/pr_update_changelog.py index a5f24e0d..f8a84960 100644 --- a/pr_agent/tools/pr_update_changelog.py +++ b/pr_agent/tools/pr_update_changelog.py @@ -5,7 +5,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