From 072ecba4c5972d4bd419ed5f1959973a81ad422e Mon Sep 17 00:00:00 2001
From: "Carles Sala (aider)" <carles@pythiac.com>
Date: Wed, 26 Feb 2025 10:50:16 +0100
Subject: [PATCH] feat: Add commands to switch main, editor, and weak models

---
 aider/commands.py            | 32 ++++++++++++++++++++--
 tests/basic/test_commands.py | 52 ++++++++++++++++++++++++++++++++++++
 2 files changed, 82 insertions(+), 2 deletions(-)

diff --git a/aider/commands.py b/aider/commands.py
index e1bb130d2..7b54973b6 100644
--- a/aider/commands.py
+++ b/aider/commands.py
@@ -83,10 +83,38 @@ class Commands:
         self.original_read_only_fnames = set(original_read_only_fnames or [])
 
     def cmd_model(self, args):
-        "Switch to a new LLM"
+        "Switch the Main Model to a new LLM"
 
         model_name = args.strip()
-        model = models.Model(model_name, weak_model=self.coder.main_model.weak_model.name)
+        model = models.Model(
+            model_name,
+            editor_model=self.coder.main_model.editor_model.name,
+            weak_model=self.coder.main_model.weak_model.name,
+        )
+        models.sanity_check_models(self.io, model)
+        raise SwitchCoder(main_model=model)
+
+    def cmd_editor_model(self, args):
+        "Switch the Editor Model to a new LLM"
+
+        model_name = args.strip()
+        model = models.Model(
+            self.coder.main_model.name,
+            editor_model=model_name,
+            weak_model=self.coder.main_model.weak_model.name,
+        )
+        models.sanity_check_models(self.io, model)
+        raise SwitchCoder(main_model=model)
+
+    def cmd_weak_model(self, args):
+        "Switch the Weak Model to a new LLM"
+
+        model_name = args.strip()
+        model = models.Model(
+            self.coder.main_model.name,
+            editor_model=self.coder.main_model.editor_model.name,
+            weak_model=model_name,
+        )
         models.sanity_check_models(self.io, model)
         raise SwitchCoder(main_model=model)
 
diff --git a/tests/basic/test_commands.py b/tests/basic/test_commands.py
index cc37495a4..154b32885 100644
--- a/tests/basic/test_commands.py
+++ b/tests/basic/test_commands.py
@@ -1664,6 +1664,58 @@ class TestCommands(TestCase):
                 self.assertIn("-Further modified content", diff_output)
                 self.assertIn("+Final modified content", diff_output)
 
+    def test_cmd_model(self):
+        io = InputOutput(pretty=False, fancy_input=False, yes=True)
+        coder = Coder.create(self.GPT35, None, io)
+        commands = Commands(io, coder)
+
+        # Test switching the main model
+        with self.assertRaises(SwitchCoder) as context:
+            commands.cmd_model("gpt-4")
+
+        # Check that the SwitchCoder exception contains the correct model configuration
+        self.assertEqual(context.exception.kwargs.get("main_model").name, "gpt-4")
+        self.assertEqual(
+            context.exception.kwargs.get("main_model").editor_model.name,
+            self.GPT35.editor_model.name,
+        )
+        self.assertEqual(
+            context.exception.kwargs.get("main_model").weak_model.name, self.GPT35.weak_model.name
+        )
+
+    def test_cmd_editor_model(self):
+        io = InputOutput(pretty=False, fancy_input=False, yes=True)
+        coder = Coder.create(self.GPT35, None, io)
+        commands = Commands(io, coder)
+
+        # Test switching the editor model
+        with self.assertRaises(SwitchCoder) as context:
+            commands.cmd_editor_model("gpt-4")
+
+        # Check that the SwitchCoder exception contains the correct model configuration
+        self.assertEqual(context.exception.kwargs.get("main_model").name, self.GPT35.name)
+        self.assertEqual(context.exception.kwargs.get("main_model").editor_model.name, "gpt-4")
+        self.assertEqual(
+            context.exception.kwargs.get("main_model").weak_model.name, self.GPT35.weak_model.name
+        )
+
+    def test_cmd_weak_model(self):
+        io = InputOutput(pretty=False, fancy_input=False, yes=True)
+        coder = Coder.create(self.GPT35, None, io)
+        commands = Commands(io, coder)
+
+        # Test switching the weak model
+        with self.assertRaises(SwitchCoder) as context:
+            commands.cmd_weak_model("gpt-4")
+
+        # Check that the SwitchCoder exception contains the correct model configuration
+        self.assertEqual(context.exception.kwargs.get("main_model").name, self.GPT35.name)
+        self.assertEqual(
+            context.exception.kwargs.get("main_model").editor_model.name,
+            self.GPT35.editor_model.name,
+        )
+        self.assertEqual(context.exception.kwargs.get("main_model").weak_model.name, "gpt-4")
+
     def test_cmd_ask(self):
         io = InputOutput(pretty=False, fancy_input=False, yes=True)
         coder = Coder.create(self.GPT35, None, io)