diff --git a/src/pearbot.py b/src/pearbot.py index d19663d..e071767 100644 --- a/src/pearbot.py +++ b/src/pearbot.py @@ -31,9 +31,7 @@ def main(): if args.server: print("Running as a server...") - github_reviewer = GitHubReviewer(code_review_agent, feedback_improver_agent) - github_reviewer.app.config['INITIAL_REVIEW_MODELS'] = initial_review_models - github_reviewer.app.config['FINAL_REVIEW_MODEL'] = final_review_model + github_reviewer = GitHubReviewer(code_review_agent, feedback_improver_agent, initial_review_models, final_review_model) github_reviewer.run_server() elif args.diff or not sys.stdin.isatty(): if args.diff == '-' or not sys.stdin.isatty(): diff --git a/src/review_github.py b/src/review_github.py index 86c59b5..c541aa0 100644 --- a/src/review_github.py +++ b/src/review_github.py @@ -13,7 +13,7 @@ from ollama_utils import validate_models class GitHubReviewer: - def __init__(self, code_review_agent, feedback_improver_agent): + def __init__(self, code_review_agent, feedback_improver_agent, initial_review_models, final_review_model): try: self.GITHUB_APP_ID = os.getenv("GITHUB_APP_ID") self.GITHUB_PRIVATE_KEY = os.getenv("GITHUB_PRIVATE_KEY") @@ -30,6 +30,8 @@ def __init__(self, code_review_agent, feedback_improver_agent): self.code_review_agent = code_review_agent self.feedback_improver_agent = feedback_improver_agent + self.initial_review_models = initial_review_models + self.final_review_model = final_review_model self.app = Flask(__name__) self.setup_routes() @@ -108,7 +110,7 @@ def handle_issue_comment(self, payload): print(f"Review condition not found") def perform_review(self, pr_number, repo_full_name, installation_id): - if not validate_models(self.app.config['INITIAL_REVIEW_MODELS'] + [self.app.config['FINAL_REVIEW_MODEL']]): + if not validate_models(self.initial_review_models + [self.final_review_model]): sys.exit(1) access_token = self.get_installation_access_token(installation_id) @@ -125,7 +127,7 @@ def perform_review(self, pr_number, repo_full_name, installation_id): } initial_reviews = [] - for model in self.app.config['INITIAL_REVIEW_MODELS']: + for model in self.initial_review_models: print(f"\n\n >>> Requesting initial review with {model}...") _, initial_review = self.code_review_agent.analyze(pr_data, model) initial_reviews.append(initial_review) @@ -134,8 +136,8 @@ def perform_review(self, pr_number, repo_full_name, installation_id): "pr_data": pr_data, "initial_reviews": initial_reviews } - print(f"\n\n >>> Requesting improved review (with {self.app.config['FINAL_REVIEW_MODEL']})...\n\n") - _, improved_feedback = self.feedback_improver_agent.analyze(improvement_data, self.app.config['FINAL_REVIEW_MODEL']) + print(f"\n\n >>> Requesting improved review (with {self.final_review_model})...\n\n") + _, improved_feedback = self.feedback_improver_agent.analyze(improvement_data, self.final_review_model) try: print(f"\n\nPosting improved feedback:\n{improved_feedback}\n\n")