Simplify conditional code. Improve logs to track conversion progress

This commit is contained in:
Debanjum Singh Solanky 2024-05-29 17:50:07 +05:30
parent 80cbaca935
commit 7823ef09dc

View file

@ -1,4 +1,5 @@
from django.core.management.base import BaseCommand
from tqdm import tqdm
from khoj.database.models import Conversation
from khoj.utils.helpers import ImageIntentType, is_none_or_empty
@ -29,7 +30,7 @@ class Command(BaseCommand):
if not options.get("source") or not options.get("destination"):
self.stdout.write(
self.style.ERROR(
"AWS_IMAGE_UPLOAD_BUCKET environment variable or --source, --destination args needs to be set."
"Set --source, --destination args to migrate serving images from source to destination URL."
)
)
return
@ -38,21 +39,21 @@ class Command(BaseCommand):
source = options["destination"] if options["reverse"] else options["source"]
for conversation in Conversation.objects.all():
conversation_updated = False
for chat in conversation.conversation_log.get("chat", []):
for chat in tqdm(conversation.conversation_log.get("chat", []), desc="Processing Conversations"):
if (
chat.get("by", "") == "khoj"
and not is_none_or_empty(chat.get("message"))
and chat.get("message", "").startswith(source)
and chat.get("intent", {}).get("type", "") == ImageIntentType.TEXT_TO_IMAGE2.value
and chat.get("message", "").endswith(".webp")
):
if chat.get("message", "").endswith(".webp"):
# Convert source url to destination url
chat["message"] = chat["message"].replace(source, destination)
conversation_updated = True
updated_count += 1
if conversation_updated:
print("Save the updated conversation")
print(f"Save the updated conversation {conversation.id} to the database.")
conversation.save()
if updated_count > 0: