feat: Add commands to switch main, editor, and weak models

This commit is contained in:
Carles Sala (aider) 2025-02-26 10:50:16 +01:00 committed by Paul Gauthier
parent 3ad5d75bee
commit 072ecba4c5
2 changed files with 82 additions and 2 deletions

View file

@ -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)

View file

@ -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)