From 25a28e6f9c37006328f4691adde2fc61a05654a7 Mon Sep 17 00:00:00 2001
From: Paul Gauthier <aider@paulg.org>
Date: Tue, 13 Jun 2023 12:19:56 -0700
Subject: [PATCH] Improved support for gpt-3.5-turbo-16k

---
 README.md         |  4 ++--
 aider/coder.py    | 20 ++++++++++----------
 aider/commands.py |  2 +-
 aider/main.py     |  2 +-
 aider/models.py   | 11 +++++------
 5 files changed, 19 insertions(+), 20 deletions(-)

diff --git a/README.md b/README.md
index 47dfa1758..f9c33d8bf 100644
--- a/README.md
+++ b/README.md
@@ -59,8 +59,8 @@ You can find more chat transcripts on the [examples page](https://aider.chat/exa
 ## GPT-4 vs GPT-3.5
 
 Aider now supports the brand new `gpt-3.5-turbo-16k` model.
-You'll need to pull the latest build from GitHub for that to work.
-For now, it just treats it like 3.5 with a bigger context window
+To use it, pull the latest build from GitHub and run `aider -3`.
+For now, aider just treats it like 3.5 with a bigger context window
 (repo maps are disabled and it does not try and use a diff based output format).
 
 Aider supports `gpt-4`, `gpt-3.5-turbo` and `gpt-4-32k`.
diff --git a/aider/coder.py b/aider/coder.py
index 824af5801..f93a4c79e 100755
--- a/aider/coder.py
+++ b/aider/coder.py
@@ -79,14 +79,14 @@ class Coder:
             self.console = Console(force_terminal=True, no_color=True)
 
         main_model = models.get_model(main_model)
-        if main_model != models.GPT35:
+        if main_model not in models.GPT35_models:
             if not self.check_model_availability(main_model):
                 if main_model != models.GPT4:
                     self.io.tool_error(f"API key does not support {main_model.name}.")
-                main_model = models.GPT35
+                main_model = models.GPT35_16k
 
         self.main_model = main_model
-        if main_model == models.GPT35:
+        if main_model in models.GPT35_models:
             self.io.tool_output(
                 f"Using {main_model.name} (experimental): disabling ctags/repo-maps.",
             )
@@ -107,7 +107,7 @@ class Coder:
             self.io.tool_output("Not using git.")
             self.find_common_root()
 
-        if main_model != models.GPT35:
+        if main_model in models.GPT4_models:
             rm_io = io if self.verbose else None
             self.repo_map = RepoMap(
                 map_tokens,
@@ -314,7 +314,7 @@ class Coder:
         ]
 
         main_sys = self.gpt_prompts.main_system
-        if self.main_model != models.GPT35:
+        if self.main_model in models.GPT4_models:
             main_sys += "\n" + self.gpt_prompts.system_reminder
 
         messages = [
@@ -343,8 +343,8 @@ class Coder:
         if edit_error:
             return edit_error
 
-        if self.main_model != models.GPT35 or not edited:
-            # Don't add assistant messages to the history if they contain "edits"
+        if self.main_model in models.GPT4_models or not edited:
+            # Don't add 3.5 assistant messages to the history if they contain "edits"
             # Because those edits are actually fully copies of the file!
             # That wastes too much context window.
             self.cur_messages += [dict(role="assistant", content=content)]
@@ -478,7 +478,7 @@ class Coder:
 
                 if self.pretty:
                     show_resp = self.resp
-                    if self.main_model == models.GPT35:
+                    if self.main_model in models.GPT35_models:
                         try:
                             show_resp = self.update_files_gpt35(self.resp, mode="diff")
                         except ValueError:
@@ -782,9 +782,9 @@ class Coder:
         return set(self.get_all_relative_files()) - set(self.get_inchat_relative_files())
 
     def apply_updates(self, content):
-        if self.main_model in (models.GPT4, models.GPT4_32k):
+        if self.main_model in models.GPT4_models:
             method = self.update_files_gpt4
-        elif self.main_model == models.GPT35:
+        elif self.main_model in models.GPT35_models:
             method = self.update_files_gpt35
         else:
             raise ValueError(f"apply_updates() doesn't support {self.main_model.name}")
diff --git a/aider/commands.py b/aider/commands.py
index dd42c2e85..2928f0f9a 100644
--- a/aider/commands.py
+++ b/aider/commands.py
@@ -183,7 +183,7 @@ class Commands:
             "was reset and removed from git.\n"
         )
 
-        if self.coder.main_model != models.GPT35:
+        if self.coder.main_model in models.GPT4_models:
             return prompts.undo_command_reply
 
     def cmd_diff(self, args):
diff --git a/aider/main.py b/aider/main.py
index 8cc34ac12..da2a1ec4b 100644
--- a/aider/main.py
+++ b/aider/main.py
@@ -83,7 +83,7 @@ def main(args=None, input=None, output=None):
         "-3",
         action="store_const",
         dest="model",
-        const=models.GPT35.name,
+        const=models.GPT35_16k.name,
         help=f"Use {models.GPT35.name} model for the main chat (not advised)",
     )
     parser.add_argument(
diff --git a/aider/models.py b/aider/models.py
index dc3057a06..36d6c10c9 100644
--- a/aider/models.py
+++ b/aider/models.py
@@ -27,15 +27,14 @@ class Model_GPT35_16k:
     max_context_tokens = 16 * 1024
 
 
-GPT35 = Model_GPT35_16k()
+GPT35_16k = Model_GPT35_16k()
+
+GPT35_models = [GPT35, GPT35_16k]
+GPT4_models = [GPT4, GPT4_32k]
 
 
 def get_model(name):
-    models = [
-        GPT4_32k,
-        GPT4,
-        GPT35,
-    ]
+    models = GPT35_models + GPT4_models
 
     for model in models:
         if model.name == name: