From acb825f4f5689ff34a35d8171c7b159d7a238218 Mon Sep 17 00:00:00 2001 From: sabaimran Date: Sun, 11 Aug 2024 10:41:43 +0530 Subject: [PATCH] Bug fixes for automations - Pass the new conversation id as kwarg for the scheduled_chat function - For edit automations, re-use the original conversation id - Parse images correctly for image automations --- src/khoj/database/adapters/__init__.py | 18 +++++++++++++++--- src/khoj/routers/api.py | 23 +++++++++++++++++++---- src/khoj/routers/email.py | 3 ++- src/khoj/routers/helpers.py | 19 +++++++++++++++++-- 4 files changed, 53 insertions(+), 10 deletions(-) diff --git a/src/khoj/database/adapters/__init__.py b/src/khoj/database/adapters/__init__.py index c25bcb51..b62c2360 100644 --- a/src/khoj/database/adapters/__init__.py +++ b/src/khoj/database/adapters/__init__.py @@ -674,15 +674,27 @@ class ConversationAdapters: @staticmethod async def acreate_conversation_session( - user: KhojUser, client_application: ClientApplication = None, agent_slug: str = None + user: KhojUser, client_application: ClientApplication = None, agent_slug: str = None, title: str = None ): if agent_slug: agent = await AgentAdapters.aget_agent_by_slug(agent_slug, user) if agent is None: raise HTTPException(status_code=400, detail="No such agent currently exists.") - return await Conversation.objects.acreate(user=user, client=client_application, agent=agent) + return await Conversation.objects.acreate(user=user, client=client_application, agent=agent, title=title) agent = await AgentAdapters.aget_default_agent() - return await Conversation.objects.acreate(user=user, client=client_application, agent=agent) + return await Conversation.objects.acreate(user=user, client=client_application, agent=agent, title=title) + + @staticmethod + def create_conversation_session( + user: KhojUser, client_application: ClientApplication = None, agent_slug: str = None, title: str = None + ): + if agent_slug: + agent = AgentAdapters.get_agent_by_slug(agent_slug, user) + if agent is None: + raise HTTPException(status_code=400, detail="No such agent currently exists.") + return Conversation.objects.create(user=user, client=client_application, agent=agent, title=title) + agent = AgentAdapters.get_default_agent() + return Conversation.objects.create(user=user, client=client_application, agent=agent, title=title) @staticmethod async def aget_conversation_by_user( diff --git a/src/khoj/routers/api.py b/src/khoj/routers/api.py index aa519403..1c7697e5 100644 --- a/src/khoj/routers/api.py +++ b/src/khoj/routers/api.py @@ -545,15 +545,19 @@ async def post_automation( if not subject: subject = await acreate_title_from_query(q) - # Create new Conversation Session associated with this new task - conversation = await ConversationAdapters.acreate_conversation_session(user, request.user.client_app) + title = f"Automation: {subject}" - calling_url = request.url.replace(query=f"{request.url.query}&conversation_id={conversation.id}") + # Create new Conversation Session associated with this new task + conversation = await ConversationAdapters.acreate_conversation_session(user, request.user.client_app, title=title) + + calling_url = request.url.replace(query=f"{request.url.query}") # Schedule automation with query_to_run, timezone, subject directly provided by user try: # Use the query to run as the scheduling request if the scheduling request is unset - automation = await schedule_automation(query_to_run, subject, crontime, timezone, q, user, calling_url) + automation = await schedule_automation( + query_to_run, subject, crontime, timezone, q, user, calling_url, conversation.id + ) except Exception as e: logger.error(f"Error creating automation {q} for {user.email}: {e}", exc_info=True) return Response( @@ -649,6 +653,16 @@ def edit_job( automation_metadata["query_to_run"] = query_to_run automation_metadata["subject"] = subject.strip() automation_metadata["crontime"] = crontime + conversation_id = automation_metadata.get("conversation_id") + + if not conversation_id: + title = f"Automation: {subject}" + + # Create new Conversation Session associated with this new task + conversation = ConversationAdapters.create_conversation_session(user, request.user.client_app, title=title) + + conversation_id = conversation.id + automation_metadata["conversation_id"] = conversation_id # Modify automation with updated query, subject automation.modify( @@ -659,6 +673,7 @@ def edit_job( "scheduling_request": q, "user": user, "calling_url": request.url, + "conversation_id": conversation_id, }, ) diff --git a/src/khoj/routers/email.py b/src/khoj/routers/email.py index 96b42edc..5123cca2 100644 --- a/src/khoj/routers/email.py +++ b/src/khoj/routers/email.py @@ -117,7 +117,8 @@ def send_task_email(name, email, query, result, subject, is_image=False): template = env.get_template("task.html") if is_image: - result = f"![{subject}]({result})" + image = result.get("image") + result = f"![{subject}]({image})" html_result = markdown_it.MarkdownIt().render(result) html_content = template.render(name=name, subject=subject, query=query, result=html_result) diff --git a/src/khoj/routers/helpers.py b/src/khoj/routers/helpers.py index 33500fb3..d821ca4b 100644 --- a/src/khoj/routers/helpers.py +++ b/src/khoj/routers/helpers.py @@ -1074,7 +1074,13 @@ def should_notify(original_query: str, executed_query: str, ai_response: str) -> def scheduled_chat( - query_to_run: str, scheduling_request: str, subject: str, user: KhojUser, calling_url: URL, job_id: str = None + query_to_run: str, + scheduling_request: str, + subject: str, + user: KhojUser, + calling_url: URL, + job_id: str = None, + conversation_id: int = None, ): logger.info(f"Processing scheduled_chat: {query_to_run}") if job_id: @@ -1101,6 +1107,10 @@ def scheduled_chat( # Replace the original scheduling query with the scheduled query query_dict["q"] = [query_to_run] + # Replace the original conversation_id with the conversation_id + if conversation_id: + query_dict["conversation_id"] = conversation_id + # Construct the URL to call the chat API with the scheduled query string encoded_query = urlencode(query_dict, doseq=True) url = f"{scheme}://{calling_url.netloc}/api/chat?{encoded_query}" @@ -1130,7 +1140,9 @@ def scheduled_chat( if raw_response.headers.get("Content-Type") == "application/json": response_map = raw_response.json() ai_response = response_map.get("response") or response_map.get("image") - is_image = response_map.get("image") is not None + is_image = False + if type(ai_response) == dict: + is_image = ai_response.get("image") is not None else: ai_response = raw_response.text @@ -1156,6 +1168,7 @@ async def schedule_automation( scheduling_request: str, user: KhojUser, calling_url: URL, + conversation_id: int, ): # Disable minute level automation recurrence minute_value = crontime.split(" ")[0] @@ -1173,6 +1186,7 @@ async def schedule_automation( "scheduling_request": scheduling_request, "subject": subject, "crontime": crontime, + "conversation_id": conversation_id, } ) query_id = hashlib.md5(f"{query_to_run}_{crontime}".encode("utf-8")).hexdigest() @@ -1191,6 +1205,7 @@ async def schedule_automation( "user": user, "calling_url": calling_url, "job_id": job_id, + "conversation_id": conversation_id, }, id=job_id, name=job_metadata,