From 0d15b59461fbfc43082de952254733a294cb2ad3 Mon Sep 17 00:00:00 2001 From: Alexey Rybalchenko Date: Mon, 14 Oct 2024 14:19:15 +0200 Subject: [PATCH] Check available models before access --- src/pearbot.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/src/pearbot.py b/src/pearbot.py index 07bc223..e34678f 100644 --- a/src/pearbot.py +++ b/src/pearbot.py @@ -11,6 +11,7 @@ from storage import get_or_create_session from agents import Agent from review_github import GitHubReviewer +from ollama_utils import is_model_available, get_available_models def extract_commit_info(diff_content): commit_range = None @@ -72,6 +73,17 @@ def analyze_diff(diff_content, initial_review_models, final_review_model): print(f"\n\nImproved feedback:\n{improved_feedback}\n\n") +def validate_models(initial_review_models, final_review_model): + all_models = initial_review_models + [final_review_model] + unavailable_models = [model for model in all_models if not is_model_available(model)] + + if unavailable_models: + print(f"Error: The following models are not available: {', '.join(unavailable_models)}") + print("Available models:") + print(", ".join(get_available_models())) + return False + return True + def main(): parser = argparse.ArgumentParser(description="Pearbot Code Review") parser.add_argument("--server", action="store_true", help="Run as a server") @@ -84,6 +96,11 @@ def main(): initial_review_models = args.initial_review_models.split(',') final_review_model = args.model + print(f"Available models: {', '.join(get_available_models())}") + + if not validate_models(initial_review_models, final_review_model): + sys.exit(1) + if args.server: print("Running as a server...") code_review_agent = Agent(role="code_reviewer", use_post_request=True)