From 6863b47e6a60ce9295f1df23f413500e94328c68 Mon Sep 17 00:00:00 2001
From: Paul Gauthier <aider@paulg.org>
Date: Sat, 20 Apr 2024 08:09:17 -0700
Subject: [PATCH] Check the environment inside Model

---
 aider/main.py   | 18 +-----------------
 aider/models.py | 23 +++++++++++++++++++++++
 2 files changed, 24 insertions(+), 17 deletions(-)

diff --git a/aider/main.py b/aider/main.py
index bab2a71fc..21b4e356a 100644
--- a/aider/main.py
+++ b/aider/main.py
@@ -602,28 +602,12 @@ def main(argv=None, input=None, output=None, force_git_root=None):
     if args.openai_organization_id:
         os.environ["OPENAI_ORGANIZATION"] = args.openai_organization_id
 
-    # Is the model known and are all needed keys/params available?
-    res = litellm.validate_environment(args.model)
-    missing_keys = res.get("missing_keys")
-    keys_in_environment = res.get("keys_in_environment")
-
-    if not keys_in_environment and not missing_keys:
-        io.tool_error(f"Unable to check environment variables for model {args.model}")
-    elif missing_keys:
-        io.tool_error(f"To use model {args.model}, please set these environment variables:")
-        for key in missing_keys:
-            io.tool_error(f"- {key}")
-        return 1
-    elif not keys_in_environment and args.require_model_info:
-        io.tool_error(models.check_model_name(args.model))
-        return 1
-
     # Check in advance that we have model metadata
     try:
         main_model = models.Model(
             args.model, weak_model=args.weak_model, require_model_info=args.require_model_info
         )
-    except models.NoModelInfo as err:
+    except (models.NoModelInfo, models.ModelEnvironmentError) as err:
         io.tool_error(str(err))
         return 1
 
diff --git a/aider/models.py b/aider/models.py
index fa3ca5224..026ea35f1 100644
--- a/aider/models.py
+++ b/aider/models.py
@@ -22,6 +22,15 @@ class NoModelInfo(Exception):
         super().__init__(check_model_name(model))
 
 
+class ModelEnvironmentError(Exception):
+    """
+    Exception raised when the environment isn't setup for the model
+    """
+
+    def __init__(self, message):
+        super().__init__(message)
+
+
 @dataclass
 class ModelSettings:
     name: str
@@ -126,6 +135,20 @@ class Model:
     def __init__(self, model, weak_model=None, require_model_info=True):
         self.name = model
 
+        # Are all needed keys/params available?
+        res = litellm.validate_environment(model)
+        missing_keys = res.get("missing_keys")
+        keys_in_environment = res.get("keys_in_environment")
+
+        if missing_keys:
+            res = f"To use model {model}, please set these environment variables:"
+            for key in missing_keys:
+                res += f"- {key}"
+            raise ModelEnvironmentError(res)
+        elif not keys_in_environment:
+            print(f"Unable to check environment variables for model {model}")
+
+        # Do we have the model_info?
         try:
             self.info = litellm.get_model_info(model)
         except Exception: