From 884fe42602f6aa196098d3f3a5524d0f596e14ce Mon Sep 17 00:00:00 2001
From: Debanjum Singh Solanky <debanjum@gmail.com>
Date: Thu, 17 Oct 2024 11:56:43 -0700
Subject: [PATCH] Allow automation as an output mode supported by custom agents

---
 .../0068_alter_agent_output_modes.py          | 24 +++++++++++++++++++
 src/khoj/database/models/__init__.py          |  1 +
 src/khoj/utils/helpers.py                     |  3 ++-
 3 files changed, 27 insertions(+), 1 deletion(-)
 create mode 100644 src/khoj/database/migrations/0068_alter_agent_output_modes.py

diff --git a/src/khoj/database/migrations/0068_alter_agent_output_modes.py b/src/khoj/database/migrations/0068_alter_agent_output_modes.py
new file mode 100644
index 00000000..31c01bd4
--- /dev/null
+++ b/src/khoj/database/migrations/0068_alter_agent_output_modes.py
@@ -0,0 +1,24 @@
+# Generated by Django 5.0.8 on 2024-10-17 18:13
+
+import django.contrib.postgres.fields
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("database", "0067_alter_agent_style_icon"),
+    ]
+
+    operations = [
+        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
+                ),
+                default=list,
+                size=None,
+            ),
+        ),
+    ]
diff --git a/src/khoj/database/models/__init__.py b/src/khoj/database/models/__init__.py
index dfe91b14..ec4b61d1 100644
--- a/src/khoj/database/models/__init__.py
+++ b/src/khoj/database/models/__init__.py
@@ -174,6 +174,7 @@ class Agent(BaseModel):
         # These map to various ConversationCommand types
         TEXT = "text"
         IMAGE = "image"
+        AUTOMATION = "automation"
 
     creator = models.ForeignKey(
         KhojUser, on_delete=models.CASCADE, default=None, null=True, blank=True
diff --git a/src/khoj/utils/helpers.py b/src/khoj/utils/helpers.py
index 4c7bf985..e0908e51 100644
--- a/src/khoj/utils/helpers.py
+++ b/src/khoj/utils/helpers.py
@@ -347,12 +347,13 @@ tool_descriptions_for_llm = {
 
 mode_descriptions_for_llm = {
     ConversationCommand.Image: "Use this if the user is requesting you to generate a picture based on their description.",
-    ConversationCommand.Automation: "Use this if the user is requesting a response at a scheduled date or time.",
+    ConversationCommand.Automation: "Use this if you are confident the user is requesting a response at a scheduled date, time and frequency",
     ConversationCommand.Text: "Use this if the other response modes don't seem to fit the query.",
 }
 
 mode_descriptions_for_agent = {
     ConversationCommand.Image: "Agent can generate image in response.",
+    ConversationCommand.Automation: "Agent can schedule a task to run at a scheduled date, time and frequency in response.",
     ConversationCommand.Text: "Agent can generate text in response.",
 }