From 5239e1c3e9189ad9543092d1f72649d98d5f882d Mon Sep 17 00:00:00 2001 From: Brian Pham Date: Tue, 12 Dec 2023 21:51:05 +0800 Subject: [PATCH] Load default AI Handler from util function --- pr_agent/algo/openai_ai_handler.py | 39 ++++++++++++++++------ pr_agent/algo/utils.py | 6 ++++ pr_agent/tools/pr_add_docs.py | 6 ++-- pr_agent/tools/pr_code_suggestions.py | 6 ++-- pr_agent/tools/pr_description.py | 6 ++-- pr_agent/tools/pr_information_from_user.py | 5 +-- pr_agent/tools/pr_questions.py | 5 +-- pr_agent/tools/pr_reviewer.py | 6 ++-- pr_agent/tools/pr_update_changelog.py | 5 +-- 9 files changed, 55 insertions(+), 29 deletions(-) diff --git a/pr_agent/algo/openai_ai_handler.py b/pr_agent/algo/openai_ai_handler.py index 4c9264d7..c521442d 100644 --- a/pr_agent/algo/openai_ai_handler.py +++ b/pr_agent/algo/openai_ai_handler.py @@ -4,6 +4,7 @@ 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 OPENAI_RETRIES = 5 @@ -37,14 +38,30 @@ class OpenAIHandler(BaseAiHandler): @retry(exceptions=(APIError, Timeout, TryAgain, AttributeError, RateLimitError), tries=OPENAI_RETRIES, delay=2, backoff=2, jitter=(1, 3)) async def chat_completion(self, model: str, system: str, user: str, temperature: float = 0.2): - chat_completion = await openai.ChatCompletion.acreate( - model=model, - messages=[{ - "role": "system", - "content": system - }, { - "role": "user", - "content": user - }], - ) - return chat_completion.choices[0].message.content + try: + deployment_id = self.deployment_id + get_logger().info("System: ", system) + get_logger().info("User: ", user) + messages = [{"role": "system", "content": system}, {"role": "user", "content": user}] + + chat_completion = await openai.ChatCompletion.acreate( + model=model, + deployment_id=deployment_id, + messages=messages, + temperature=temperature, + ) + resp = chat_completion["choices"][0]['message']['content'] + finish_reason = chat_completion["choices"][0]["finish_reason"] + usage = chat_completion.get("usage") + get_logger().info("AI response", response=resp, messages=messages, finish_reason=finish_reason, + model=model, usage=usage) + return resp, finish_reason + except (APIError, Timeout, TryAgain) as e: + get_logger().error("Error during OpenAI inference: ", e) + raise + except (RateLimitError) as e: + get_logger().error("Rate limit error during OpenAI inference: ", e) + raise + except (Exception) as e: + get_logger().error("Unknown error during OpenAI inference: ", e) + raise TryAgain from e \ No newline at end of file diff --git a/pr_agent/algo/utils.py b/pr_agent/algo/utils.py index 4e88b33e..14cde04a 100644 --- a/pr_agent/algo/utils.py +++ b/pr_agent/algo/utils.py @@ -8,6 +8,9 @@ from datetime import datetime from typing import Any, List import yaml +from pr_agent.algo.litellm_ai_handler import LiteLLMAiHandler +from pr_agent.algo.base_ai_handler import BaseAiHandler +from pr_agent.algo.openai_ai_handler import OpenAIHandler from starlette_context import context from pr_agent.config_loader import get_settings, global_settings from pr_agent.log import get_logger @@ -304,3 +307,6 @@ def try_fix_yaml(review_text: str) -> dict: except: pass return data + +def get_ai_handler() -> BaseAiHandler: + return OpenAIHandler() \ 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 3d50afe2..f76baa02 100644 --- a/pr_agent/tools/pr_add_docs.py +++ b/pr_agent/tools/pr_add_docs.py @@ -4,10 +4,10 @@ from typing import Dict 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 load_yaml +from pr_agent.algo.utils import load_yaml, 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 @@ -15,7 +15,7 @@ from pr_agent.log import get_logger class PRAddDocs: - def __init__(self, pr_url: str, cli_mode=False, args: list = None, ai_handler: BaseAiHandler = LiteLLMAiHandler()): + def __init__(self, pr_url: str, cli_mode=False, args: list = None, ai_handler: BaseAiHandler = get_ai_handler()): self.git_provider = get_git_provider()(pr_url) self.main_language = get_main_pr_language( diff --git a/pr_agent/tools/pr_code_suggestions.py b/pr_agent/tools/pr_code_suggestions.py index e55720ae..96bd79fb 100644 --- a/pr_agent/tools/pr_code_suggestions.py +++ b/pr_agent/tools/pr_code_suggestions.py @@ -4,10 +4,10 @@ from typing import Dict, List 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, get_pr_multi_diffs, retry_with_fallback_models from pr_agent.algo.token_handler import TokenHandler -from pr_agent.algo.utils import load_yaml +from pr_agent.algo.utils import load_yaml, 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 @@ -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, ai_handler: BaseAiHandler = LiteLLMAiHandler() ): + def __init__(self, pr_url: str, cli_mode=False, args: list = None, ai_handler: BaseAiHandler = get_ai_handler() ): self.git_provider = get_git_provider()(pr_url) self.main_language = get_main_pr_language( diff --git a/pr_agent/tools/pr_description.py b/pr_agent/tools/pr_description.py index a1b73df2..58de3b4e 100644 --- a/pr_agent/tools/pr_description.py +++ b/pr_agent/tools/pr_description.py @@ -4,10 +4,10 @@ from typing import List, Tuple 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 load_yaml +from pr_agent.algo.utils import load_yaml, 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 @@ -15,7 +15,7 @@ from pr_agent.log import get_logger class PRDescription: - 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()): """ Initialize the PRDescription object with the necessary attributes and objects for generating a PR description using an AI model. diff --git a/pr_agent/tools/pr_information_from_user.py b/pr_agent/tools/pr_information_from_user.py index beb1b0ab..78903490 100644 --- a/pr_agent/tools/pr_information_from_user.py +++ b/pr_agent/tools/pr_information_from_user.py @@ -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() diff --git a/pr_agent/tools/pr_questions.py b/pr_agent/tools/pr_questions.py index b1930924..43c276cb 100644 --- a/pr_agent/tools/pr_questions.py +++ b/pr_agent/tools/pr_questions.py @@ -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 PRQuestions: - def __init__(self, pr_url: str, args=None, ai_handler: BaseAiHandler = LiteLLMAiHandler()): + def __init__(self, pr_url: str, args=None, ai_handler: BaseAiHandler = get_ai_handler()): question_str = self.parse_args(args) self.git_provider = get_git_provider()(pr_url) self.main_pr_language = get_main_pr_language( diff --git a/pr_agent/tools/pr_reviewer.py b/pr_agent/tools/pr_reviewer.py index 4de24b06..d68b893e 100644 --- a/pr_agent/tools/pr_reviewer.py +++ b/pr_agent/tools/pr_reviewer.py @@ -6,10 +6,10 @@ import yaml from jinja2 import Environment, StrictUndefined from yaml import SafeLoader -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 convert_to_markdown, load_yaml, try_fix_yaml +from pr_agent.algo.utils import convert_to_markdown, get_ai_handler, load_yaml, try_fix_yaml 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 IncrementalPR, get_main_pr_language @@ -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, ai_handler: BaseAiHandler = LiteLLMAiHandler()): + def __init__(self, pr_url: str, is_answer: bool = False, is_auto: bool = False, args: list = None, ai_handler: BaseAiHandler = get_ai_handler()): """ Initialize the PRReviewer object with the necessary attributes and objects to review a pull request. diff --git a/pr_agent/tools/pr_update_changelog.py b/pr_agent/tools/pr_update_changelog.py index 7625218b..febe6fec 100644 --- a/pr_agent/tools/pr_update_changelog.py +++ b/pr_agent/tools/pr_update_changelog.py @@ -5,9 +5,10 @@ from typing import Tuple 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 @@ -17,7 +18,7 @@ CHANGELOG_LINES = 50 class PRUpdateChangelog: - def __init__(self, pr_url: str, cli_mode=False, args=None, ai_handler: BaseAiHandler = LiteLLMAiHandler()): + def __init__(self, pr_url: str, cli_mode=False, args=None, ai_handler: BaseAiHandler = get_ai_handler()): self.git_provider = get_git_provider()(pr_url) self.main_language = get_main_pr_language(