From 9a4d2f6f39ce6726fa5e2a6863527e0492d4dc5a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Erik=20Bj=C3=A4reholt?= <erik@bjareho.lt>
Date: Wed, 9 Oct 2024 20:54:56 +0200
Subject: [PATCH] fix: limit image size, fixes exception when viewing large
 images (#185)

---
 gptme/message.py | 28 ++++++++++++++++++++++++----
 1 file changed, 24 insertions(+), 4 deletions(-)

diff --git a/gptme/message.py b/gptme/message.py
index 75f8d4f6..1810fdc8 100644
--- a/gptme/message.py
+++ b/gptme/message.py
@@ -77,6 +77,10 @@ def replace(self, **kwargs) -> Self:
     def _content_files_list(
         self, openai: bool = False, anthropic: bool = False
     ) -> list[dict[str, Any]]:
+        # only these providers support files in the content
+        if not openai and not anthropic:
+            raise ValueError("Provider does not support files in the content")
+
         # combines a content message with a list of files
         content: list[dict[str, Any]] = (
             self.content
@@ -99,6 +103,24 @@ def _content_files_list(
                     "text": f"![{f.name}]({f.name}):",
                 }
             )
+
+            # read file
+            data_bytes = f.read_bytes()
+            data = base64.b64encode(data_bytes).decode("utf-8")
+
+            # check that the file is not too large
+            # anthropic limit is 5MB
+            # TODO: use compression to reduce file size
+            # print(f"{len(data)=}")
+            if len(data) > 5_000_000:
+                content.append(
+                    {
+                        "type": "text",
+                        "text": "Image size exceeds 5MB. Please upload a smaller image.",
+                    }
+                )
+                continue
+
             if anthropic:
                 content.append(
                     {
@@ -106,7 +128,7 @@ def _content_files_list(
                         "source": {
                             "type": "base64",
                             "media_type": media_type,
-                            "data": base64.b64encode(f.read_bytes()).decode("utf-8"),
+                            "data": data,
                         },
                     }
                 )
@@ -115,9 +137,7 @@ def _content_files_list(
                 content.append(
                     {
                         "type": "image_url",
-                        "image_url": {
-                            "url": f"data:{media_type};base64,{base64.b64encode(f.read_bytes()).decode('utf-8')}"
-                        },
+                        "image_url": {"url": f"data:{media_type};base64,{data}"},
                     }
                 )
             else: