From 9ceaf97f08b6e71466ad703c7b31e95486133734 Mon Sep 17 00:00:00 2001
From: Joshua Vial <joshua@enspiral.com>
Date: Mon, 11 Dec 2023 22:21:24 +1300
Subject: [PATCH] making image code more robust

---
 aider/commands.py | 11 +++++++----
 aider/sendchat.py |  2 +-
 aider/utils.py    |  2 +-
 3 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/aider/commands.py b/aider/commands.py
index ef9e83f3c..d58198cd1 100644
--- a/aider/commands.py
+++ b/aider/commands.py
@@ -10,7 +10,7 @@ from aider import prompts, voice
 
 from .dump import dump  # noqa: F401
 
-from aider.utils import is_image_file
+from aider.utils import is_image_file, is_gpt4_with_openai_base_url
 
 class Commands:
     voice = None
@@ -171,11 +171,12 @@ class Commands:
         self.io.tool_output("=" * (width + cost_width + 1))
         self.io.tool_output(f"${total_cost:5.2f} {fmt(total)} tokens total")
 
-        # Set image_in_chat to False unless is_gpt4_with_openai_base_url returns True
+        # only switch to image model token count if gpt4 and openai and image in files
         image_in_chat = False
-        if utils.is_gpt4_with_openai_base_url(self.coder.main_model.name, self.coder.client):
+        if is_gpt4_with_openai_base_url(self.coder.main_model.name, self.coder.client):
             image_in_chat = any(is_image_file(relative_fname) for relative_fname in self.coder.get_inchat_relative_files())
         limit = 128000 if image_in_chat else self.coder.main_model.max_context_tokens
+
         remaining = limit - total
         if remaining > 1024:
             self.io.tool_output(f"{cost_pad}{fmt(remaining)} tokens remaining in context window")
@@ -327,7 +328,9 @@ class Commands:
             if abs_file_path in self.coder.abs_fnames:
                 self.io.tool_error(f"{matched_file} is already in the chat")
             else:
-                #TODO put in guard to stop images being added to non openai / gpt-4
+                if is_image_file(matched_file) and not is_gpt4_with_openai_base_url(self.coder.main_model.name, self.coder.client):
+                    self.io.tool_error(f"Cannot add image file {matched_file} as the model does not support image files")
+                    continue
                 content = self.io.read_text(abs_file_path)
                 if content is None:
                     self.io.tool_error(f"Unable to read {matched_file}")
diff --git a/aider/sendchat.py b/aider/sendchat.py
index 18956b83c..64aa9c7b7 100644
--- a/aider/sendchat.py
+++ b/aider/sendchat.py
@@ -8,6 +8,7 @@ import openai
 # from diskcache import Cache
 from openai import APIConnectionError, InternalServerError, RateLimitError
 
+from aider.utils import is_gpt4_with_openai_base_url
 from aider.dump import dump  # noqa: F401
 
 CACHE_PATH = "~/.aider.send.cache.v1"
@@ -41,7 +42,6 @@ def send_with_retries(client, model_name, messages, functions, stream):
     if functions is not None:
         kwargs["functions"] = functions
 
-    from aider.utils import is_gpt4_with_openai_base_url
 
     # Check conditions to switch to gpt-4-vision-preview or strip out image_url messages
     if client and is_gpt4_with_openai_base_url(model_name, client):
diff --git a/aider/utils.py b/aider/utils.py
index 2b02f7bdb..c2053e818 100644
--- a/aider/utils.py
+++ b/aider/utils.py
@@ -1,5 +1,4 @@
 from pathlib import Path
-from openai import OpenAIError
 
 # Set of image file extensions
 IMAGE_EXTENSIONS = {'.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff', '.webp'}
@@ -42,6 +41,7 @@ def show_messages(messages, title=None, functions=None):
 
     if functions:
         dump(functions)
+
 def is_gpt4_with_openai_base_url(model_name, client):
     """
     Check if the model_name starts with 'gpt-4' and the client base URL includes 'api.openai.com'.