khoj 1.24.2.dev2__py3-none-any.whl → 1.24.2.dev16__py3-none-any.whl

khoj/database/adapters/__init__.py

@@ -10,6 +10,7 @@ from enum import Enum
 from typing import Callable, Iterable, List, Optional, Type
 
 import cron_descriptor
+import django
 from apscheduler.job import Job
 from asgiref.sync import sync_to_async
 from django.contrib.sessions.backends.db import SessionStore
@@ -551,26 +552,62 @@ class ClientApplicationAdapters:
 
 class AgentAdapters:
     DEFAULT_AGENT_NAME = "Khoj"
-    DEFAULT_AGENT_AVATAR = "https://assets.khoj.dev/lamp-128.png"
     DEFAULT_AGENT_SLUG = "khoj"
 
+    @staticmethod
+    async def aget_readonly_agent_by_slug(agent_slug: str, user: KhojUser):
+        return await Agent.objects.filter(
+            (Q(slug__iexact=agent_slug.lower()))
+            & (
+                Q(privacy_level=Agent.PrivacyLevel.PUBLIC)
+                | Q(privacy_level=Agent.PrivacyLevel.PROTECTED)
+                | Q(creator=user)
+            )
+        ).afirst()
+
+    @staticmethod
+    async def adelete_agent_by_slug(agent_slug: str, user: KhojUser):
+        agent = await AgentAdapters.aget_agent_by_slug(agent_slug, user)
+        if agent:
+            await agent.adelete()
+            return True
+        return False
+
     @staticmethod
     async def aget_agent_by_slug(agent_slug: str, user: KhojUser):
         return await Agent.objects.filter(
-            (Q(slug__iexact=agent_slug.lower())) & (Q(public=True) | Q(creator=user))
+            (Q(slug__iexact=agent_slug.lower())) & (Q(privacy_level=Agent.PrivacyLevel.PUBLIC) | Q(creator=user))
+        ).afirst()
+
+    @staticmethod
+    async def aget_agent_by_name(agent_name: str, user: KhojUser):
+        return await Agent.objects.filter(
+            (Q(name__iexact=agent_name.lower())) & (Q(privacy_level=Agent.PrivacyLevel.PUBLIC) | Q(creator=user))
         ).afirst()
 
     @staticmethod
     def get_agent_by_slug(slug: str, user: KhojUser = None):
         if user:
-            return Agent.objects.filter((Q(slug__iexact=slug.lower())) & (Q(public=True) | Q(creator=user))).first()
-        return Agent.objects.filter(slug__iexact=slug.lower(), public=True).first()
+            return Agent.objects.filter(
+                (Q(slug__iexact=slug.lower())) & (Q(privacy_level=Agent.PrivacyLevel.PUBLIC) | Q(creator=user))
+            ).first()
+        return Agent.objects.filter(slug__iexact=slug.lower(), privacy_level=Agent.PrivacyLevel.PUBLIC).first()
 
     @staticmethod
     def get_all_accessible_agents(user: KhojUser = None):
+        public_query = Q(privacy_level=Agent.PrivacyLevel.PUBLIC)
         if user:
-            return Agent.objects.filter(Q(public=True) | Q(creator=user)).distinct().order_by("created_at")
-        return Agent.objects.filter(public=True).order_by("created_at")
+            return (
+                Agent.objects.filter(public_query | Q(creator=user))
+                .distinct()
+                .order_by("created_at")
+                .prefetch_related("creator", "chat_model", "fileobject_set")
+            )
+        return (
+            Agent.objects.filter(public_query)
+            .order_by("created_at")
+            .prefetch_related("creator", "chat_model", "fileobject_set")
+        )
 
     @staticmethod
     async def aget_all_accessible_agents(user: KhojUser = None) -> List[Agent]:
@@ -604,17 +641,19 @@ class AgentAdapters:
             agent.chat_model = default_conversation_config
             agent.slug = AgentAdapters.DEFAULT_AGENT_SLUG
             agent.name = AgentAdapters.DEFAULT_AGENT_NAME
+            agent.privacy_level = Agent.PrivacyLevel.PUBLIC
+            agent.managed_by_admin = True
+            agent.input_tools = []
+            agent.output_modes = []
             agent.save()
         else:
             # The default agent is public and managed by the admin. It's handled a little differently than other agents.
             agent = Agent.objects.create(
                 name=AgentAdapters.DEFAULT_AGENT_NAME,
-                public=True,
+                privacy_level=Agent.PrivacyLevel.PUBLIC,
                 managed_by_admin=True,
                 chat_model=default_conversation_config,
                 personality=default_personality,
-                tools=["*"],
-                avatar=AgentAdapters.DEFAULT_AGENT_AVATAR,
                 slug=AgentAdapters.DEFAULT_AGENT_SLUG,
             )
             Conversation.objects.filter(agent=None).update(agent=agent)
@@ -625,6 +664,68 @@ class AgentAdapters:
     async def aget_default_agent():
         return await Agent.objects.filter(name=AgentAdapters.DEFAULT_AGENT_NAME).afirst()
 
+    @staticmethod
+    async def aupdate_agent(
+        user: KhojUser,
+        name: str,
+        personality: str,
+        privacy_level: str,
+        icon: str,
+        color: str,
+        chat_model: str,
+        files: List[str],
+        input_tools: List[str],
+        output_modes: List[str],
+    ):
+        chat_model_option = await ChatModelOptions.objects.filter(chat_model=chat_model).afirst()
+
+        agent, created = await Agent.objects.filter(name=name, creator=user).aupdate_or_create(
+            defaults={
+                "name": name,
+                "creator": user,
+                "personality": personality,
+                "privacy_level": privacy_level,
+                "style_icon": icon,
+                "style_color": color,
+                "chat_model": chat_model_option,
+                "input_tools": input_tools,
+                "output_modes": output_modes,
+            }
+        )
+
+        # Delete all existing files and entries
+        await FileObject.objects.filter(agent=agent).adelete()
+        await Entry.objects.filter(agent=agent).adelete()
+
+        for file in files:
+            reference_file = await FileObject.objects.filter(file_name=file, user=agent.creator).afirst()
+            if reference_file:
+                await FileObject.objects.acreate(file_name=file, agent=agent, raw_text=reference_file.raw_text)
+
+                # Duplicate all entries associated with the file
+                entries: List[Entry] = []
+                async for entry in Entry.objects.filter(file_path=file, user=agent.creator).aiterator():
+                    entries.append(
+                        Entry(
+                            agent=agent,
+                            embeddings=entry.embeddings,
+                            raw=entry.raw,
+                            compiled=entry.compiled,
+                            heading=entry.heading,
+                            file_source=entry.file_source,
+                            file_type=entry.file_type,
+                            file_path=entry.file_path,
+                            file_name=entry.file_name,
+                            url=entry.url,
+                            hashed_value=entry.hashed_value,
+                        )
+                    )
+
+                # Bulk create entries
+                await Entry.objects.abulk_create(entries)
+
+        return agent
+
 
 class PublicConversationAdapters:
     @staticmethod
@@ -1113,8 +1214,8 @@ class FileObjectAdapters:
         return await FileObject.objects.acreate(user=user, file_name=file_name, raw_text=raw_text)
 
     @staticmethod
-    async def async_get_file_objects_by_name(user: KhojUser, file_name: str):
-        return await sync_to_async(list)(FileObject.objects.filter(user=user, file_name=file_name))
+    async def async_get_file_objects_by_name(user: KhojUser, file_name: str, agent: Agent = None):
+        return await sync_to_async(list)(FileObject.objects.filter(user=user, file_name=file_name, agent=agent))
 
     @staticmethod
     async def async_get_all_file_objects(user: KhojUser):
@@ -1196,10 +1297,18 @@ class EntryAdapters:
     def user_has_entries(user: KhojUser):
         return Entry.objects.filter(user=user).exists()
 
+    @staticmethod
+    def agent_has_entries(agent: Agent):
+        return Entry.objects.filter(agent=agent).exists()
+
     @staticmethod
     async def auser_has_entries(user: KhojUser):
         return await Entry.objects.filter(user=user).aexists()
 
+    @staticmethod
+    async def aagent_has_entries(agent: Agent):
+        return await Entry.objects.filter(agent=agent).aexists()
+
     @staticmethod
     async def adelete_entry_by_file(user: KhojUser, file_path: str):
         return await Entry.objects.filter(user=user, file_path=file_path).adelete()
@@ -1214,6 +1323,10 @@ class EntryAdapters:
 
         return deleted_count
 
+    @staticmethod
+    async def aget_agent_entry_filepaths(agent: Agent):
+        return await sync_to_async(list)(Entry.objects.filter(agent=agent).values_list("file_path", flat=True))
+
     @staticmethod
     def get_all_filenames_by_source(user: KhojUser, file_source: str):
         return (
@@ -1229,15 +1342,19 @@ class EntryAdapters:
         return total_size / 1024 / 1024
 
     @staticmethod
-    def apply_filters(user: KhojUser, query: str, file_type_filter: str = None):
+    def apply_filters(user: KhojUser, query: str, file_type_filter: str = None, agent: Agent = None):
         q_filter_terms = Q()
 
         word_filters = EntryAdapters.word_filter.get_filter_terms(query)
         file_filters = EntryAdapters.file_filter.get_filter_terms(query)
         date_filters = EntryAdapters.date_filter.get_query_date_range(query)
 
+        user_or_agent = Q(user=user)
+        if agent != None:
+            user_or_agent |= Q(agent=agent)
+
         if len(word_filters) == 0 and len(file_filters) == 0 and len(date_filters) == 0:
-            return Entry.objects.filter(user=user)
+            return Entry.objects.filter(user_or_agent)
 
         for term in word_filters:
             if term.startswith("+"):
@@ -1273,7 +1390,7 @@ class EntryAdapters:
                 formatted_max_date = date.fromtimestamp(max_date).strftime("%Y-%m-%d")
                 q_filter_terms &= Q(embeddings_dates__date__lte=formatted_max_date)
 
-        relevant_entries = Entry.objects.filter(user=user).filter(q_filter_terms)
+        relevant_entries = Entry.objects.filter(user_or_agent).filter(q_filter_terms)
         if file_type_filter:
             relevant_entries = relevant_entries.filter(file_type=file_type_filter)
         return relevant_entries
@@ -1286,9 +1403,15 @@ class EntryAdapters:
         file_type_filter: str = None,
         raw_query: str = None,
         max_distance: float = math.inf,
+        agent: Agent = None,
     ):
-        relevant_entries = EntryAdapters.apply_filters(user, raw_query, file_type_filter)
-        relevant_entries = relevant_entries.filter(user=user).annotate(
+        user_or_agent = Q(user=user)
+
+        if agent != None:
+            user_or_agent |= Q(agent=agent)
+
+        relevant_entries = EntryAdapters.apply_filters(user, raw_query, file_type_filter, agent)
+        relevant_entries = relevant_entries.filter(user_or_agent).annotate(
             distance=CosineDistance("embeddings", embeddings)
         )
         relevant_entries = relevant_entries.filter(distance__lte=max_distance)

khoj/database/admin.py

@@ -27,6 +27,7 @@ from khoj.database.models import (
     Subscription,
     TextToImageModelConfig,
     UserConversationConfig,
+    UserRequests,
     UserSearchModelConfig,
     UserVoiceModelConfig,
     VoiceModelOption,
@@ -103,6 +104,7 @@ admin.site.register(NotionConfig)
 admin.site.register(UserVoiceModelConfig)
 admin.site.register(VoiceModelOption)
 admin.site.register(UserConversationConfig)
+admin.site.register(UserRequests)
 
 
 @admin.register(Agent)

khoj/database/migrations/0065_remove_agent_avatar_remove_agent_public_and_more.py

@@ -0,0 +1,49 @@
+# Generated by Django 5.0.8 on 2024-09-18 02:54
+
+import django.db.models.deletion
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("database", "0064_remove_conversation_temp_id_alter_conversation_id"),
+    ]
+
+    operations = [
+        migrations.RemoveField(
+            model_name="agent",
+            name="avatar",
+        ),
+        migrations.RemoveField(
+            model_name="agent",
+            name="public",
+        ),
+        migrations.AddField(
+            model_name="agent",
+            name="privacy_level",
+            field=models.CharField(
+                choices=[("public", "Public"), ("private", "Private"), ("protected", "Protected")],
+                default="private",
+                max_length=30,
+            ),
+        ),
+        migrations.AddField(
+            model_name="entry",
+            name="agent",
+            field=models.ForeignKey(
+                blank=True, default=None, null=True, on_delete=django.db.models.deletion.CASCADE, to="database.agent"
+            ),
+        ),
+        migrations.AddField(
+            model_name="fileobject",
+            name="agent",
+            field=models.ForeignKey(
+                blank=True, default=None, null=True, on_delete=django.db.models.deletion.CASCADE, to="database.agent"
+            ),
+        ),
+        migrations.AlterField(
+            model_name="agent",
+            name="slug",
+            field=models.CharField(max_length=200, unique=True),
+        ),
+    ]

khoj/database/migrations/0066_remove_agent_tools_agent_input_tools_and_more.py

@@ -0,0 +1,69 @@
+# Generated by Django 5.0.8 on 2024-10-01 00:42
+
+import django.contrib.postgres.fields
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("database", "0065_remove_agent_avatar_remove_agent_public_and_more"),
+    ]
+
+    operations = [
+        migrations.RemoveField(
+            model_name="agent",
+            name="tools",
+        ),
+        migrations.AddField(
+            model_name="agent",
+            name="input_tools",
+            field=django.contrib.postgres.fields.ArrayField(
+                base_field=models.CharField(
+                    choices=[
+                        ("general", "General"),
+                        ("online", "Online"),
+                        ("notes", "Notes"),
+                        ("summarize", "Summarize"),
+                        ("webpage", "Webpage"),
+                    ],
+                    max_length=200,
+                ),
+                default=list,
+                size=None,
+            ),
+        ),
+        migrations.AddField(
+            model_name="agent",
+            name="output_modes",
+            field=django.contrib.postgres.fields.ArrayField(
+                base_field=models.CharField(choices=[("text", "Text"), ("image", "Image")], max_length=200),
+                default=list,
+                size=None,
+            ),
+        ),
+        migrations.AlterField(
+            model_name="agent",
+            name="style_icon",
+            field=models.CharField(
+                choices=[
+                    ("Lightbulb", "Lightbulb"),
+                    ("Health", "Health"),
+                    ("Robot", "Robot"),
+                    ("Aperture", "Aperture"),
+                    ("GraduationCap", "Graduation Cap"),
+                    ("Jeep", "Jeep"),
+                    ("Island", "Island"),
+                    ("MathOperations", "Math Operations"),
+                    ("Asclepius", "Asclepius"),
+                    ("Couch", "Couch"),
+                    ("Code", "Code"),
+                    ("Atom", "Atom"),
+                    ("ClockCounterClockwise", "Clock Counter Clockwise"),
+                    ("PencilLine", "Pencil Line"),
+                    ("Chalkboard", "Chalkboard"),
+                ],
+                default="Lightbulb",
+                max_length=200,
+            ),
+        ),
+    ]

khoj/database/migrations/0067_alter_agent_style_icon.py

@@ -0,0 +1,50 @@
+# Generated by Django 5.0.8 on 2024-10-01 18:42
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("database", "0066_remove_agent_tools_agent_input_tools_and_more"),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name="agent",
+            name="style_icon",
+            field=models.CharField(
+                choices=[
+                    ("Lightbulb", "Lightbulb"),
+                    ("Health", "Health"),
+                    ("Robot", "Robot"),
+                    ("Aperture", "Aperture"),
+                    ("GraduationCap", "Graduation Cap"),
+                    ("Jeep", "Jeep"),
+                    ("Island", "Island"),
+                    ("MathOperations", "Math Operations"),
+                    ("Asclepius", "Asclepius"),
+                    ("Couch", "Couch"),
+                    ("Code", "Code"),
+                    ("Atom", "Atom"),
+                    ("ClockCounterClockwise", "Clock Counter Clockwise"),
+                    ("PencilLine", "Pencil Line"),
+                    ("Chalkboard", "Chalkboard"),
+                    ("Cigarette", "Cigarette"),
+                    ("CraneTower", "Crane Tower"),
+                    ("Heart", "Heart"),
+                    ("Leaf", "Leaf"),
+                    ("NewspaperClipping", "Newspaper Clipping"),
+                    ("OrangeSlice", "Orange Slice"),
+                    ("SmileyMelting", "Smiley Melting"),
+                    ("YinYang", "Yin Yang"),
+                    ("SneakerMove", "Sneaker Move"),
+                    ("Student", "Student"),
+                    ("Oven", "Oven"),
+                    ("Gavel", "Gavel"),
+                    ("Broadcast", "Broadcast"),
+                ],
+                default="Lightbulb",
+                max_length=200,
+            ),
+        ),
+    ]

khoj/database/models/__init__.py

@@ -3,6 +3,7 @@ import uuid
 from random import choice
 
 from django.contrib.auth.models import AbstractUser
+from django.contrib.postgres.fields import ArrayField
 from django.core.exceptions import ValidationError
 from django.db import models
 from django.db.models.signals import pre_save
@@ -10,6 +11,8 @@ from django.dispatch import receiver
 from pgvector.django import VectorField
 from phonenumber_field.modelfields import PhoneNumberField
 
+from khoj.utils.helpers import ConversationCommand
+
 
 class BaseModel(models.Model):
     created_at = models.DateTimeField(auto_now_add=True)
@@ -125,7 +128,7 @@ class Agent(BaseModel):
         EMERALD = "emerald"
 
     class StyleIconTypes(models.TextChoices):
-        LIGHBULB = "Lightbulb"
+        LIGHTBULB = "Lightbulb"
         HEALTH = "Health"
         ROBOT = "Robot"
         APERTURE = "Aperture"
@@ -140,20 +143,64 @@ class Agent(BaseModel):
         CLOCK_COUNTER_CLOCKWISE = "ClockCounterClockwise"
         PENCIL_LINE = "PencilLine"
         CHALKBOARD = "Chalkboard"
+        CIGARETTE = "Cigarette"
+        CRANE_TOWER = "CraneTower"
+        HEART = "Heart"
+        LEAF = "Leaf"
+        NEWSPAPER_CLIPPING = "NewspaperClipping"
+        ORANGE_SLICE = "OrangeSlice"
+        SMILEY_MELTING = "SmileyMelting"
+        YIN_YANG = "YinYang"
+        SNEAKER_MOVE = "SneakerMove"
+        STUDENT = "Student"
+        OVEN = "Oven"
+        GAVEL = "Gavel"
+        BROADCAST = "Broadcast"
+
+    class PrivacyLevel(models.TextChoices):
+        PUBLIC = "public"
+        PRIVATE = "private"
+        PROTECTED = "protected"
+
+    class InputToolOptions(models.TextChoices):
+        # These map to various ConversationCommand types
+        GENERAL = "general"
+        ONLINE = "online"
+        NOTES = "notes"
+        SUMMARIZE = "summarize"
+        WEBPAGE = "webpage"
+
+    class OutputModeOptions(models.TextChoices):
+        # These map to various ConversationCommand types
+        TEXT = "text"
+        IMAGE = "image"
 
     creator = models.ForeignKey(
         KhojUser, on_delete=models.CASCADE, default=None, null=True, blank=True
     )  # Creator will only be null when the agents are managed by admin
     name = models.CharField(max_length=200)
     personality = models.TextField()
-    avatar = models.URLField(max_length=400, default=None, null=True, blank=True)
-    tools = models.JSONField(default=list)  # List of tools the agent has access to, like online search or notes search
-    public = models.BooleanField(default=False)
+    input_tools = ArrayField(models.CharField(max_length=200, choices=InputToolOptions.choices), default=list)
+    output_modes = ArrayField(models.CharField(max_length=200, choices=OutputModeOptions.choices), default=list)
     managed_by_admin = models.BooleanField(default=False)
     chat_model = models.ForeignKey(ChatModelOptions, on_delete=models.CASCADE)
-    slug = models.CharField(max_length=200)
+    slug = models.CharField(max_length=200, unique=True)
     style_color = models.CharField(max_length=200, choices=StyleColorTypes.choices, default=StyleColorTypes.BLUE)
-    style_icon = models.CharField(max_length=200, choices=StyleIconTypes.choices, default=StyleIconTypes.LIGHBULB)
+    style_icon = models.CharField(max_length=200, choices=StyleIconTypes.choices, default=StyleIconTypes.LIGHTBULB)
+    privacy_level = models.CharField(max_length=30, choices=PrivacyLevel.choices, default=PrivacyLevel.PRIVATE)
+
+    def save(self, *args, **kwargs):
+        is_new = self._state.adding
+
+        if self.creator is None:
+            self.managed_by_admin = True
+
+        if is_new:
+            random_sequence = "".join(choice("0123456789") for i in range(6))
+            slug = f"{self.name.lower().replace(' ', '-')}-{random_sequence}"
+            self.slug = slug
+
+        super().save(*args, **kwargs)
 
 
 class ProcessLock(BaseModel):
@@ -173,22 +220,11 @@ class ProcessLock(BaseModel):
 def verify_agent(sender, instance, **kwargs):
     # check if this is a new instance
     if instance._state.adding:
-        if Agent.objects.filter(name=instance.name, public=True).exists():
+        if Agent.objects.filter(name=instance.name, privacy_level=Agent.PrivacyLevel.PUBLIC).exists():
             raise ValidationError(f"A public Agent with the name {instance.name} already exists.")
         if Agent.objects.filter(name=instance.name, creator=instance.creator).exists():
             raise ValidationError(f"A private Agent with the name {instance.name} already exists.")
 
-        slug = instance.name.lower().replace(" ", "-")
-        observed_random_numbers = set()
-        while Agent.objects.filter(slug=slug).exists():
-            try:
-                random_number = choice([i for i in range(0, 1000) if i not in observed_random_numbers])
-            except IndexError:
-                raise ValidationError("Unable to generate a unique slug for the Agent. Please try again later.")
-            observed_random_numbers.add(random_number)
-            slug = f"{slug}-{random_number}"
-        instance.slug = slug
-
 
 class NotionConfig(BaseModel):
     token = models.CharField(max_length=200)
@@ -406,6 +442,7 @@ class Entry(BaseModel):
         GITHUB = "github"
 
     user = models.ForeignKey(KhojUser, on_delete=models.CASCADE, default=None, null=True, blank=True)
+    agent = models.ForeignKey(Agent, on_delete=models.CASCADE, default=None, null=True, blank=True)
     embeddings = VectorField(dimensions=None)
     raw = models.TextField()
     compiled = models.TextField()
@@ -418,12 +455,17 @@ class Entry(BaseModel):
     hashed_value = models.CharField(max_length=100)
     corpus_id = models.UUIDField(default=uuid.uuid4, editable=False)
 
+    def save(self, *args, **kwargs):
+        if self.user and self.agent:
+            raise ValidationError("An Entry cannot be associated with both a user and an agent.")
+
 
 class FileObject(BaseModel):
     # Same as Entry but raw will be a much larger string
     file_name = models.CharField(max_length=400, default=None, null=True, blank=True)
     raw_text = models.TextField()
     user = models.ForeignKey(KhojUser, on_delete=models.CASCADE, default=None, null=True, blank=True)
+    agent = models.ForeignKey(Agent, on_delete=models.CASCADE, default=None, null=True, blank=True)
 
 
 class EntryDates(BaseModel):