From a352940dfd67813eadbf1825956370e20b1ff20c Mon Sep 17 00:00:00 2001
From: Debanjum Singh Solanky <debanjum@gmail.com>
Date: Mon, 15 Apr 2024 17:44:05 +0530
Subject: [PATCH] Use Django management command to update images URL in DB to
 WebP

This provides Khoj server admins more control on migrating their S3
images to WebP format from PNG
---
 src/khoj/database/management/__init__.py      |  0
 .../database/management/commands/__init__.py  |  0
 .../commands/convert_images_png_to_webp.py    | 40 +++++++++++++++++++
 .../migrations/0035_convert_png_to_webp.py    | 10 -----
 4 files changed, 40 insertions(+), 10 deletions(-)
 create mode 100644 src/khoj/database/management/__init__.py
 create mode 100644 src/khoj/database/management/commands/__init__.py
 create mode 100644 src/khoj/database/management/commands/convert_images_png_to_webp.py

diff --git a/src/khoj/database/management/__init__.py b/src/khoj/database/management/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/src/khoj/database/management/commands/__init__.py b/src/khoj/database/management/commands/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/src/khoj/database/management/commands/convert_images_png_to_webp.py b/src/khoj/database/management/commands/convert_images_png_to_webp.py
new file mode 100644
index 00000000..b1ad8615
--- /dev/null
+++ b/src/khoj/database/management/commands/convert_images_png_to_webp.py
@@ -0,0 +1,40 @@
+from django.core.management.base import BaseCommand
+
+from khoj.database.models import Conversation
+from khoj.utils.helpers import ImageIntentType
+
+
+class Command(BaseCommand):
+    help = "Convert all images to WebP format or reverse."
+
+    def add_arguments(self, parser):
+        # Add a new argument 'reverse' to the command
+        parser.add_argument(
+            "--reverse",
+            action="store_true",
+            help="Convert from WebP to PNG instead of PNG to WebP",
+        )
+
+    def handle(self, *args, **options):
+        updated_count = 0
+        for conversation in Conversation.objects.all():
+            conversation_updated = False
+            for chat in conversation.conversation_log["chat"]:
+                if chat["by"] == "khoj" and chat["intent"]["type"] == ImageIntentType.TEXT_TO_IMAGE2.value:
+                    if options["reverse"] and chat["message"].endswith(".webp"):
+                        # Convert WebP url to PNG url
+                        chat["message"] = chat["message"].replace(".webp", ".png")
+                        conversation_updated = True
+                        updated_count += 1
+                    elif chat["message"].endswith(".png"):
+                        # Convert PNG url to WebP url
+                        chat["message"] = chat["message"].replace(".png", ".webp")
+                        conversation_updated = True
+                        updated_count += 1
+            if conversation_updated:
+                conversation.save()
+
+        if updated_count > 0 and options["reverse"]:
+            self.stdout.write(self.style.SUCCESS(f"Successfully converted {updated_count} WebP images to PNG format."))
+        elif updated_count > 0:
+            self.stdout.write(self.style.SUCCESS(f"Successfully converted {updated_count} PNG images to WebP format."))
diff --git a/src/khoj/database/migrations/0035_convert_png_to_webp.py b/src/khoj/database/migrations/0035_convert_png_to_webp.py
index 6ffa024b..35495629 100644
--- a/src/khoj/database/migrations/0035_convert_png_to_webp.py
+++ b/src/khoj/database/migrations/0035_convert_png_to_webp.py
@@ -30,11 +30,6 @@ def convert_png_images_to_webp(apps, schema_editor):
                     chat["intent"]["type"] = ImageIntentType.TEXT_TO_IMAGE_V3.value
                     webp_image_io.close()
 
-            if chat["by"] == "khoj" and chat["intent"]["type"] == ImageIntentType.TEXT_TO_IMAGE2.value:
-                print("❗️ Please MANUALLY update PNG images created by Khoj in your AWS S3 bucket to WebP format.")
-                # Convert PNG url to WebP url
-                chat["message"] = chat["message"].replace(".png", ".webp")
-
         # Save the updated conversation history
         conversation.save()
 
@@ -60,11 +55,6 @@ def convert_webp_images_to_png(apps, schema_editor):
                     chat["intent"]["type"] = ImageIntentType.TEXT_TO_IMAGE.value
                     webp_image_io.close()
 
-            if chat["by"] == "khoj" and chat["intent"]["type"] == ImageIntentType.TEXT_TO_IMAGE2.value:
-                # Convert WebP url to PNG url
-                print("❗️ Please MANUALLY update WebP images created by Khoj in your AWS S3 bucket to PNG format.")
-                chat["message"] = chat["message"].replace(".webp", ".png")
-
         # Save the updated conversation history
         conversation.save()