diff --git a/src/khoj/database/admin.py b/src/khoj/database/admin.py index f64a02df..f92b670d 100644 --- a/src/khoj/database/admin.py +++ b/src/khoj/database/admin.py @@ -1,5 +1,9 @@ +import csv +import json + from django.contrib import admin from django.contrib.auth.admin import UserAdmin +from django.http import HttpResponse # Register your models here. @@ -13,6 +17,7 @@ from khoj.database.models import ( Subscription, ReflectiveQuestion, TextToImageModelConfig, + Conversation, ) admin.site.register(KhojUser, UserAdmin) @@ -25,3 +30,101 @@ admin.site.register(SearchModelConfig) admin.site.register(Subscription) admin.site.register(ReflectiveQuestion) admin.site.register(TextToImageModelConfig) + + +@admin.register(Conversation) +class ConversationAdmin(admin.ModelAdmin): + list_display = ( + "id", + "user", + "created_at", + "updated_at", + ) + search_fields = ("conversation_id",) + ordering = ("-created_at",) + + actions = ["export_selected_objects", "export_selected_minimal_objects"] + + def export_selected_objects(self, request, queryset): + response = HttpResponse(content_type="text/csv") + response["Content-Disposition"] = 'attachment; filename="conversations.csv"' + + writer = csv.writer(response) + writer.writerow(["id", "user", "created_at", "updated_at", "conversation_log"]) + + for conversation in queryset: + modified_log = conversation.conversation_log + chat_log = modified_log.get("chat", []) + for idx, log in enumerate(chat_log): + if ( + log["by"] == "khoj" + and log["intent"] + and log["intent"]["type"] + and log["intent"]["type"] == "text-to-image" + ): + log["message"] = "image redacted for space" + chat_log[idx] = log + modified_log["chat"] = chat_log + + writer.writerow( + [ + conversation.id, + conversation.user, + conversation.created_at, + conversation.updated_at, + json.dumps(modified_log), + ] + ) + + return response + + export_selected_objects.short_description = "Export selected conversations" + + def export_selected_minimal_objects(self, request, queryset): + response = HttpResponse(content_type="text/csv") + response["Content-Disposition"] = 'attachment; filename="conversations.csv"' + + writer = csv.writer(response) + writer.writerow(["id", "user", "created_at", "updated_at", "conversation_log"]) + + fields_to_keep = set(["message", "by", "created"]) + + for conversation in queryset: + return_log = dict() + chat_log = conversation.conversation_log.get("chat", []) + for idx, log in enumerate(chat_log): + updated_log = {} + for key in fields_to_keep: + updated_log[key] = log[key] + if ( + log["by"] == "khoj" + and log["intent"] + and log["intent"]["type"] + and log["intent"]["type"] == "text-to-image" + ): + updated_log["message"] = "image redacted for space" + chat_log[idx] = updated_log + return_log["chat"] = chat_log + + writer.writerow( + [ + conversation.id, + conversation.user, + conversation.created_at, + conversation.updated_at, + json.dumps(return_log), + ] + ) + + return response + + export_selected_minimal_objects.short_description = "Export selected conversations (minimal)" + + def get_actions(self, request): + actions = super().get_actions(request) + if not request.user.is_superuser: + if "export_selected_objects" in actions: + del actions["export_selected_objects"] + if "export_selected_minimal_objects" in actions: + del actions["export_selected_minimal_objects"] + return actions