From 59fec37943fd3105d0a52773581a65eccb191cab Mon Sep 17 00:00:00 2001
From: sabaimran <narmiabas@gmail.com>
Date: Sun, 20 Oct 2024 22:24:51 -0700
Subject: [PATCH] Improve agents management, and limit agents view to private
 and official agents

- Default to None for the input_tools and output_modes so that they can be managed in the admin panel
- Hold off on showing off all Public Agents until we have a better experience for user profiles etc.
---
 src/khoj/database/adapters/__init__.py        |  2 +
 ...nt_input_tools_alter_agent_output_modes.py | 46 +++++++++++++++++++
 src/khoj/database/models/__init__.py          |  8 +++-
 3 files changed, 54 insertions(+), 2 deletions(-)
 create mode 100644 src/khoj/database/migrations/0070_alter_agent_input_tools_alter_agent_output_modes.py

diff --git a/src/khoj/database/adapters/__init__.py b/src/khoj/database/adapters/__init__.py
index d025459e..eb914ed8 100644
--- a/src/khoj/database/adapters/__init__.py
+++ b/src/khoj/database/adapters/__init__.py
@@ -622,6 +622,8 @@ class AgentAdapters:
     @staticmethod
     def get_all_accessible_agents(user: KhojUser = None):
         public_query = Q(privacy_level=Agent.PrivacyLevel.PUBLIC)
+        # TODO Update this to allow any public agent that's officially approved once that experience is launched
+        public_query &= Q(managed_by_admin=True)
         if user:
             return (
                 Agent.objects.filter(public_query | Q(creator=user))
diff --git a/src/khoj/database/migrations/0070_alter_agent_input_tools_alter_agent_output_modes.py b/src/khoj/database/migrations/0070_alter_agent_input_tools_alter_agent_output_modes.py
new file mode 100644
index 00000000..74dfc229
--- /dev/null
+++ b/src/khoj/database/migrations/0070_alter_agent_input_tools_alter_agent_output_modes.py
@@ -0,0 +1,46 @@
+# Generated by Django 5.0.8 on 2024-10-21 05:16
+
+import django.contrib.postgres.fields
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("database", "0069_webscraper_serverchatsettings_web_scraper"),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            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,
+                ),
+                blank=True,
+                default=list,
+                null=True,
+                size=None,
+            ),
+        ),
+        migrations.AlterField(
+            model_name="agent",
+            name="output_modes",
+            field=django.contrib.postgres.fields.ArrayField(
+                base_field=models.CharField(
+                    choices=[("text", "Text"), ("image", "Image"), ("automation", "Automation")], max_length=200
+                ),
+                blank=True,
+                default=list,
+                null=True,
+                size=None,
+            ),
+        ),
+    ]
diff --git a/src/khoj/database/models/__init__.py b/src/khoj/database/models/__init__.py
index 2b2fde2d..6b122dac 100644
--- a/src/khoj/database/models/__init__.py
+++ b/src/khoj/database/models/__init__.py
@@ -180,8 +180,12 @@ class Agent(BaseModel):
     )  # Creator will only be null when the agents are managed by admin
     name = models.CharField(max_length=200)
     personality = models.TextField()
-    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)
+    input_tools = ArrayField(
+        models.CharField(max_length=200, choices=InputToolOptions.choices), default=list, null=True, blank=True
+    )
+    output_modes = ArrayField(
+        models.CharField(max_length=200, choices=OutputModeOptions.choices), default=list, null=True, blank=True
+    )
     managed_by_admin = models.BooleanField(default=False)
     chat_model = models.ForeignKey(ChatModelOptions, on_delete=models.CASCADE)
     slug = models.CharField(max_length=200, unique=True)