From cbc8a021795b2199aaf190aaa11f481c8ac4fb57 Mon Sep 17 00:00:00 2001
From: Debanjum Singh Solanky <debanjum@gmail.com>
Date: Tue, 30 Apr 2024 18:23:18 +0530
Subject: [PATCH] Make, use func for constructing the automation created
 response

- Dedupe logic across http, ws chat API endpoints
- Reduces size of already too long http, ws chat API endpoint funcs
---
 src/khoj/routers/api_chat.py | 41 +++++-------------------------------
 src/khoj/routers/helpers.py  | 22 +++++++++++++++++++
 2 files changed, 27 insertions(+), 36 deletions(-)

diff --git a/src/khoj/routers/api_chat.py b/src/khoj/routers/api_chat.py
index ae3f7cd9..ce9d4c01 100644
--- a/src/khoj/routers/api_chat.py
+++ b/src/khoj/routers/api_chat.py
@@ -1,12 +1,10 @@
 import json
 import logging
 import math
-import re
 from datetime import datetime
 from typing import Dict, Optional
 from urllib.parse import unquote
 
-import cron_descriptor
 from asgiref.sync import sync_to_async
 from fastapi import APIRouter, Depends, HTTPException, Request, WebSocket
 from fastapi.requests import Request
@@ -37,6 +35,7 @@ from khoj.routers.helpers import (
     agenerate_chat_response,
     aget_relevant_information_sources,
     aget_relevant_output_modes,
+    construct_automation_created_message,
     create_automation,
     get_conversation_command,
     is_ready_to_chat,
@@ -401,24 +400,10 @@ async def websocket_endpoint(
                     f"Unable to create automation. Ensure the automation doesn't already exist."
                 )
                 continue
-            # Display next run time in user timezone instead of UTC
-            schedule = f'{cron_descriptor.get_description(crontime)} {automation.next_run_time.strftime("%Z")}'
-            next_run_time = automation.next_run_time.strftime("%Y-%m-%d %I:%M %p %Z")
-            # Remove /automated_task prefix from inferred_query
-            unprefixed_query_to_run = re.sub(r"^\/automated_task\s*", "", query_to_run)
-            # Create the automation response
-            scheme = "http" if not websocket.url.is_secure else "https"
-            automation_icon_url = f"{scheme}://{websocket.url.netloc}/static/assets/icons/automation.svg"
-            llm_response = f"""
-            ### ![]({automation_icon_url}) Created Automation
-- Subject: **{subject}**
-- Query to Run: "{unprefixed_query_to_run}"
-- Schedule: `{schedule}`
-- Next Run At: {next_run_time}
-
-Manage your tasks [here](/config#automations).
-            """.strip()
 
+            llm_response = construct_automation_created_message(
+                automation, crontime, query_to_run, subject, websocket.url
+            )
             await sync_to_async(save_to_conversation_log)(
                 q,
                 llm_response,
@@ -657,24 +642,8 @@ async def chat(
                 media_type="text/plain",
                 status_code=500,
             )
-        # Display next run time in user timezone instead of UTC
-        schedule = f'{cron_descriptor.get_description(crontime)} {automation.next_run_time.strftime("%Z")}'
-        next_run_time = automation.next_run_time.strftime("%Y-%m-%d %I:%M %p %Z")
-        # Remove /automated_task prefix from inferred_query
-        unprefixed_query_to_run = re.sub(r"^\/automated_task\s*", "", query_to_run)
-        # Create the Automation response
-        scheme = "http" if not request.url.is_secure else "https"
-        automation_icon_url = f"{scheme}://{request.url.netloc}/static/assets/icons/automation.svg"
-        llm_response = f"""
-        ### ![]({automation_icon_url}) Created Automation
-- Subject: **{subject}**
-- Query to Run: "{unprefixed_query_to_run}"
-- Schedule: `{schedule}`
-- Next Run At: {next_run_time}
-
-Manage your automations [here](/config#automations).
-        """.strip()
 
+        llm_response = construct_automation_created_message(automation, crontime, query_to_run, subject, request.url)
         await sync_to_async(save_to_conversation_log)(
             q,
             llm_response,
diff --git a/src/khoj/routers/helpers.py b/src/khoj/routers/helpers.py
index 53808c21..fda45469 100644
--- a/src/khoj/routers/helpers.py
+++ b/src/khoj/routers/helpers.py
@@ -21,9 +21,11 @@ from typing import (
 )
 from urllib.parse import parse_qs, urlencode
 
+import cron_descriptor
 import openai
 import pytz
 import requests
+from apscheduler.job import Job
 from apscheduler.triggers.cron import CronTrigger
 from asgiref.sync import sync_to_async
 from fastapi import Depends, Header, HTTPException, Request, UploadFile
@@ -948,3 +950,23 @@ async def create_automation(
         jitter=30,
     )
     return job, crontime_string, query_to_run, subject
+
+
+def construct_automation_created_message(automation: Job, crontime: str, query_to_run: str, subject: str, url: URL):
+    # Display next run time in user timezone instead of UTC
+    schedule = f'{cron_descriptor.get_description(crontime)} {automation.next_run_time.strftime("%Z")}'
+    next_run_time = automation.next_run_time.strftime("%Y-%m-%d %I:%M %p %Z")
+    # Remove /automated_task prefix from inferred_query
+    unprefixed_query_to_run = re.sub(r"^\/automated_task\s*", "", query_to_run)
+    # Create the automation response
+    scheme = "http" if not url.is_secure else "https"
+    automation_icon_url = f"{scheme}://{url.netloc}/static/assets/icons/automation.svg"
+    return f"""
+    ### ![]({automation_icon_url}) Created Automation
+- Subject: **{subject}**
+- Query to Run: "{unprefixed_query_to_run}"
+- Schedule: `{schedule}`
+- Next Run At: {next_run_time}
+
+Manage your tasks [here](/config#automations).
+    """.strip()