xiaxinhang/pr-agent
1
0
Fork 0
mirror of https://github.com/qodo-ai/pr-agent.git synced 2025-07-04 04:40:38 +08:00
Code Issues Packages Projects Releases Wiki Activity

Refactor AI handler classes

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

6
pr_agent/tools/pr_description.py
View File

@ -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 = {