khoj 1.16.1.dev15__py3-none-any.whl

khoj/database/admin.py

@@ -0,0 +1,290 @@
+import csv
+import json
+
+from apscheduler.job import Job
+from django.contrib import admin, messages
+from django.contrib.auth.admin import UserAdmin
+from django.http import HttpResponse
+from django_apscheduler.admin import DjangoJobAdmin
+from django_apscheduler.jobstores import DjangoJobStore
+from django_apscheduler.models import DjangoJob
+
+from khoj.database.models import (
+    Agent,
+    ChatModelOptions,
+    ClientApplication,
+    Conversation,
+    Entry,
+    GithubConfig,
+    KhojUser,
+    NotionConfig,
+    OpenAIProcessorConversationConfig,
+    ProcessLock,
+    ReflectiveQuestion,
+    SearchModelConfig,
+    ServerChatSettings,
+    SpeechToTextModelOptions,
+    Subscription,
+    TextToImageModelConfig,
+    UserSearchModelConfig,
+    VoiceModelOption,
+)
+from khoj.utils.helpers import ImageIntentType
+
+admin.site.unregister(DjangoJob)
+
+
+class KhojDjangoJobAdmin(DjangoJobAdmin):
+    list_display = (
+        "id",
+        "next_run_time",
+        "job_info",
+    )
+    search_fields = ("id", "next_run_time")
+    ordering = ("-next_run_time",)
+    job_store = DjangoJobStore()
+
+    def job_info(self, obj):
+        job: Job = self.job_store.lookup_job(obj.id)
+        return f"{job.func_ref} {job.args} {job.kwargs}" if job else "None"
+
+    job_info.short_description = "Job Info"  # type: ignore
+
+    def get_search_results(self, request, queryset, search_term):
+        queryset, use_distinct = super().get_search_results(request, queryset, search_term)
+        if search_term:
+            jobs = [job.id for job in self.job_store.get_all_jobs() if search_term in str(job)]
+            queryset |= self.model.objects.filter(id__in=jobs)
+        return queryset, use_distinct
+
+
+admin.site.register(DjangoJob, KhojDjangoJobAdmin)
+
+
+class KhojUserAdmin(UserAdmin):
+    list_display = (
+        "id",
+        "email",
+        "username",
+        "is_active",
+        "is_staff",
+        "is_superuser",
+        "phone_number",
+    )
+    search_fields = ("email", "username", "phone_number", "uuid")
+    filter_horizontal = ("groups", "user_permissions")
+
+    fieldsets = (("Personal info", {"fields": ("phone_number", "email_verification_code")}),) + UserAdmin.fieldsets
+
+    actions = ["get_email_login_url"]
+
+    def get_email_login_url(self, request, queryset):
+        for user in queryset:
+            if user.email:
+                host = request.get_host()
+                unique_id = user.email_verification_code
+                login_url = f"{host}/auth/magic?code={unique_id}"
+                messages.info(request, f"Email login URL for {user.email}: {login_url}")
+
+    get_email_login_url.short_description = "Get email login URL"  # type: ignore
+
+
+admin.site.register(KhojUser, KhojUserAdmin)
+
+admin.site.register(ProcessLock)
+admin.site.register(SpeechToTextModelOptions)
+admin.site.register(SearchModelConfig)
+admin.site.register(ReflectiveQuestion)
+admin.site.register(UserSearchModelConfig)
+admin.site.register(ClientApplication)
+admin.site.register(GithubConfig)
+admin.site.register(NotionConfig)
+admin.site.register(VoiceModelOption)
+
+
+@admin.register(Agent)
+class AgentAdmin(admin.ModelAdmin):
+    list_display = (
+        "id",
+        "name",
+    )
+    search_fields = ("id", "name")
+    ordering = ("-created_at",)
+
+
+@admin.register(Entry)
+class EntryAdmin(admin.ModelAdmin):
+    list_display = (
+        "id",
+        "created_at",
+        "updated_at",
+        "user",
+        "file_source",
+        "file_type",
+        "file_name",
+        "file_path",
+    )
+    search_fields = ("id", "user__email", "user__username", "file_path")
+    list_filter = (
+        "file_type",
+        "user__email",
+    )
+    ordering = ("-created_at",)
+
+
+@admin.register(Subscription)
+class KhojUserSubscription(admin.ModelAdmin):
+    list_display = (
+        "id",
+        "user",
+        "type",
+    )
+
+    search_fields = ("id", "user__email", "user__username", "type")
+    list_filter = ("type",)
+
+
+@admin.register(ChatModelOptions)
+class ChatModelOptionsAdmin(admin.ModelAdmin):
+    list_display = (
+        "id",
+        "chat_model",
+        "model_type",
+        "max_prompt_size",
+    )
+    search_fields = ("id", "chat_model", "model_type")
+
+
+@admin.register(TextToImageModelConfig)
+class TextToImageModelOptionsAdmin(admin.ModelAdmin):
+    list_display = (
+        "id",
+        "model_name",
+        "model_type",
+    )
+    search_fields = ("id", "model_name", "model_type")
+
+
+@admin.register(OpenAIProcessorConversationConfig)
+class OpenAIProcessorConversationConfigAdmin(admin.ModelAdmin):
+    list_display = (
+        "id",
+        "name",
+        "api_key",
+        "api_base_url",
+    )
+    search_fields = ("id", "name", "api_key", "api_base_url")
+
+
+@admin.register(ServerChatSettings)
+class ServerChatSettingsAdmin(admin.ModelAdmin):
+    list_display = (
+        "default_model",
+        "summarizer_model",
+    )
+
+
+@admin.register(Conversation)
+class ConversationAdmin(admin.ModelAdmin):
+    list_display = (
+        "id",
+        "user",
+        "created_at",
+        "updated_at",
+        "client",
+    )
+    search_fields = ("id", "user__email", "user__username", "client__name")
+    list_filter = ("agent",)
+    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"] == ImageIntentType.TEXT_TO_IMAGE.value
+                        or log["intent"]["type"] == ImageIntentType.TEXT_TO_IMAGE_V3.value
+                    )
+                ):
+                    log["message"] = "inline 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"  # type: ignore
+
+    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"] == ImageIntentType.TEXT_TO_IMAGE.value
+                        or log["intent"]["type"] == ImageIntentType.TEXT_TO_IMAGE_V3.value
+                    )
+                ):
+                    updated_log["message"] = "inline 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)"  # type: ignore
+
+    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

khoj/database/apps.py

@@ -0,0 +1,6 @@
+from django.apps import AppConfig
+
+
+class DatabaseConfig(AppConfig):
+    default_auto_field = "django.db.models.BigAutoField"
+    name = "khoj.database"

khoj/database/management/commands/change_generated_images_url.py

@@ -0,0 +1,61 @@
+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
+
+
+class Command(BaseCommand):
+    help = "Serve Khoj generated images from a different URL."
+
+    def add_arguments(self, parser):
+        # Pass Source URL
+        parser.add_argument(
+            "--source",
+            action="store",
+            help="URL from which generated images are currently served.",
+        )
+        # Pass Destination URL
+        parser.add_argument("--destination", action="store", help="URL to serve generated image from going forward.")
+
+        # Add a new argument 'reverse' to the command
+        parser.add_argument(
+            "--reverse",
+            action="store_true",
+            help="Revert to serve generated images from source instead of destination URL.",
+        )
+
+    def handle(self, *args, **options):
+        updated_count = 0
+        if not options.get("source") or not options.get("destination"):
+            self.stdout.write(
+                self.style.ERROR(
+                    "Set --source, --destination args to migrate serving images from source to destination URL."
+                )
+            )
+            return
+
+        destination = options["source"] if options["reverse"] else options["destination"]
+        source = options["destination"] if options["reverse"] else options["source"]
+        for conversation in Conversation.objects.all():
+            conversation_updated = False
+            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")
+                ):
+                    # Convert source url to destination url
+                    chat["message"] = chat["message"].replace(source, destination)
+                    conversation_updated = True
+                    updated_count += 1
+
+            if conversation_updated:
+                print(f"Save the updated conversation {conversation.id} to the database.")
+                conversation.save()
+
+        if updated_count > 0:
+            success = f"Successfully converted {updated_count} image URLs from {source} to {destination}.".strip()
+            self.stdout.write(self.style.SUCCESS(success))

khoj/database/management/commands/convert_images_png_to_webp.py

@@ -0,0 +1,99 @@
+import base64
+import io
+
+from django.core.management.base import BaseCommand
+from PIL import Image
+
+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.get("chat", []):
+                if (
+                    chat.get("by", "") == "khoj"
+                    and chat.get("intent", {}).get("type", "") == ImageIntentType.TEXT_TO_IMAGE.value
+                    and not options["reverse"]
+                ):
+                    # Decode the base64 encoded PNG image
+                    print("Decode the base64 encoded PNG image")
+                    decoded_image = base64.b64decode(chat["message"])
+
+                    # Convert images from PNG to WebP format
+                    print("Convert images from PNG to WebP format")
+                    image_io = io.BytesIO(decoded_image)
+                    with Image.open(image_io) as png_image:
+                        webp_image_io = io.BytesIO()
+                        png_image.save(webp_image_io, "WEBP")
+
+                        # Encode the WebP image back to base64
+                        webp_image_bytes = webp_image_io.getvalue()
+                        chat["message"] = base64.b64encode(webp_image_bytes).decode()
+                        chat["intent"]["type"] = ImageIntentType.TEXT_TO_IMAGE_V3.value
+                        webp_image_io.close()
+                    conversation_updated = True
+                    updated_count += 1
+
+                elif (
+                    chat.get("by", "") == "khoj"
+                    and chat.get("intent", {}).get("type", "") == ImageIntentType.TEXT_TO_IMAGE_V3.value
+                    and options["reverse"]
+                ):
+                    # Decode the base64 encoded WebP image
+                    print("Decode the base64 encoded WebP image")
+                    decoded_image = base64.b64decode(chat["message"])
+
+                    # Convert images from WebP to PNG format
+                    print("Convert images from WebP to PNG format")
+                    image_io = io.BytesIO(decoded_image)
+                    with Image.open(image_io) as png_image:
+                        webp_image_io = io.BytesIO()
+                        png_image.save(webp_image_io, "PNG")
+
+                        # Encode the WebP image back to base64
+                        webp_image_bytes = webp_image_io.getvalue()
+                        chat["message"] = base64.b64encode(webp_image_bytes).decode()
+                        chat["intent"]["type"] = ImageIntentType.TEXT_TO_IMAGE.value
+                        webp_image_io.close()
+                    conversation_updated = True
+                    updated_count += 1
+
+                elif (
+                    chat.get("by", "") == "khoj"
+                    and chat.get("intent", {}).get("type", "") == ImageIntentType.TEXT_TO_IMAGE2.value
+                ):
+                    if options["reverse"] and chat.get("message", "").endswith(".webp"):
+                        # Convert WebP url to PNG url
+                        print("Convert WebP url to PNG url")
+                        chat["message"] = chat["message"].replace(".webp", ".png")
+                        conversation_updated = True
+                        updated_count += 1
+                    elif chat.get("message", "").endswith(".png"):
+                        # Convert PNG url to WebP url
+                        print("Convert PNG url to WebP url")
+                        chat["message"] = chat["message"].replace(".png", ".webp")
+                        conversation_updated = True
+                        updated_count += 1
+
+            if conversation_updated:
+                print("Save the updated conversation")
+                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."))

khoj/database/migrations/0001_khojuser.py

@@ -0,0 +1,98 @@
+# Generated by Django 4.2.5 on 2023-09-14 19:00
+
+import django.contrib.auth.models
+import django.contrib.auth.validators
+import django.utils.timezone
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    initial = True
+
+    dependencies = [
+        ("auth", "0012_alter_user_first_name_max_length"),
+    ]
+
+    run_before = [
+        ("admin", "0001_initial"),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name="KhojUser",
+            fields=[
+                ("id", models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name="ID")),
+                ("password", models.CharField(max_length=128, verbose_name="password")),
+                ("last_login", models.DateTimeField(blank=True, null=True, verbose_name="last login")),
+                (
+                    "is_superuser",
+                    models.BooleanField(
+                        default=False,
+                        help_text="Designates that this user has all permissions without explicitly assigning them.",
+                        verbose_name="superuser status",
+                    ),
+                ),
+                (
+                    "username",
+                    models.CharField(
+                        error_messages={"unique": "A user with that username already exists."},
+                        help_text="Required. 150 characters or fewer. Letters, digits and @/./+/-/_ only.",
+                        max_length=150,
+                        unique=True,
+                        validators=[django.contrib.auth.validators.UnicodeUsernameValidator()],
+                        verbose_name="username",
+                    ),
+                ),
+                ("first_name", models.CharField(blank=True, max_length=150, verbose_name="first name")),
+                ("last_name", models.CharField(blank=True, max_length=150, verbose_name="last name")),
+                ("email", models.EmailField(blank=True, max_length=254, verbose_name="email address")),
+                (
+                    "is_staff",
+                    models.BooleanField(
+                        default=False,
+                        help_text="Designates whether the user can log into this admin site.",
+                        verbose_name="staff status",
+                    ),
+                ),
+                (
+                    "is_active",
+                    models.BooleanField(
+                        default=True,
+                        help_text="Designates whether this user should be treated as active. Unselect this instead of deleting accounts.",
+                        verbose_name="active",
+                    ),
+                ),
+                ("date_joined", models.DateTimeField(default=django.utils.timezone.now, verbose_name="date joined")),
+                (
+                    "groups",
+                    models.ManyToManyField(
+                        blank=True,
+                        help_text="The groups this user belongs to. A user will get all permissions granted to each of their groups.",
+                        related_name="user_set",
+                        related_query_name="user",
+                        to="auth.group",
+                        verbose_name="groups",
+                    ),
+                ),
+                (
+                    "user_permissions",
+                    models.ManyToManyField(
+                        blank=True,
+                        help_text="Specific permissions for this user.",
+                        related_name="user_set",
+                        related_query_name="user",
+                        to="auth.permission",
+                        verbose_name="user permissions",
+                    ),
+                ),
+            ],
+            options={
+                "verbose_name": "user",
+                "verbose_name_plural": "users",
+                "abstract": False,
+            },
+            managers=[
+                ("objects", django.contrib.auth.models.UserManager()),
+            ],
+        ),
+    ]

khoj/database/migrations/0002_googleuser.py

@@ -0,0 +1,32 @@
+# Generated by Django 4.2.4 on 2023-09-18 23:24
+
+import django.db.models.deletion
+from django.conf import settings
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("database", "0001_khojuser"),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name="GoogleUser",
+            fields=[
+                ("id", models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name="ID")),
+                ("sub", models.CharField(max_length=200)),
+                ("azp", models.CharField(max_length=200)),
+                ("email", models.CharField(max_length=200)),
+                ("name", models.CharField(max_length=200)),
+                ("given_name", models.CharField(max_length=200)),
+                ("family_name", models.CharField(max_length=200)),
+                ("picture", models.CharField(max_length=200)),
+                ("locale", models.CharField(max_length=200)),
+                (
+                    "user",
+                    models.OneToOneField(on_delete=django.db.models.deletion.CASCADE, to=settings.AUTH_USER_MODEL),
+                ),
+            ],
+        ),
+    ]

khoj/database/migrations/0003_vector_extension.py

@@ -0,0 +1,10 @@
+from django.db import migrations
+from pgvector.django import VectorExtension
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("database", "0002_googleuser"),
+    ]
+
+    operations = [VectorExtension()]