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 django.core.management.base import BaseCommand
from tqdm import tqdm
from khoj.database.models import Conversation from khoj.database.models import Conversation
from khoj.utils.helpers import ImageIntentType, is_none_or_empty 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"): if not options.get("source") or not options.get("destination"):
self.stdout.write( self.stdout.write(
self.style.ERROR( 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 return
@ -38,21 +39,21 @@ class Command(BaseCommand):
source = options["destination"] if options["reverse"] else options["source"] source = options["destination"] if options["reverse"] else options["source"]
for conversation in Conversation.objects.all(): for conversation in Conversation.objects.all():
conversation_updated = False conversation_updated = False
for chat in conversation.conversation_log.get("chat", []): for chat in tqdm(conversation.conversation_log.get("chat", []), desc="Processing Conversations"):
if ( if (
chat.get("by", "") == "khoj" chat.get("by", "") == "khoj"
and not is_none_or_empty(chat.get("message")) and not is_none_or_empty(chat.get("message"))
and chat.get("message", "").startswith(source) and chat.get("message", "").startswith(source)
and chat.get("intent", {}).get("type", "") == ImageIntentType.TEXT_TO_IMAGE2.value 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
# Convert source url to destination url chat["message"] = chat["message"].replace(source, destination)
chat["message"] = chat["message"].replace(source, destination) conversation_updated = True
conversation_updated = True updated_count += 1
updated_count += 1
if conversation_updated: if conversation_updated:
print("Save the updated conversation") print(f"Save the updated conversation {conversation.id} to the database.")
conversation.save() conversation.save()
if updated_count > 0: if updated_count > 0: