Load default AI Handler from util function

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):
+        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,
-            messages=[{
+                deployment_id=deployment_id,
-                "role": "system",
+                messages=messages,
-                "content": system
+                temperature=temperature,
-            }, {
-                "role": "user",
-                "content": user            
-            }],
             )
-        return chat_completion.choices[0].message.content
+            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        

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()

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(

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(

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.

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()

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(

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.
 

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(