From f0222f6d0828c7c4d6bdfc90b083bfa3596d8878 Mon Sep 17 00:00:00 2001
From: Debanjum Singh Solanky <debanjum@gmail.com>
Date: Mon, 4 Dec 2023 18:54:22 -0500
Subject: [PATCH] Make save_to_conversation_log helper function reusable

- Move it out to conversation.utils from generate_chat_response function
- Log new optional intent_type argument to capture type of response
  expected. This can be type responses by Khoj e.g speech, image. It
  can be used to render responses by Khoj appropriately on clients
- Make user_message_time argument optional, set the time to now by
  default if not passed by calling function
---
 src/khoj/processor/conversation/utils.py | 29 ++++++++++++++++++++++
 src/khoj/routers/helpers.py              | 31 +++---------------------
 2 files changed, 33 insertions(+), 27 deletions(-)

diff --git a/src/khoj/processor/conversation/utils.py b/src/khoj/processor/conversation/utils.py
index e2164719..4606efd9 100644
--- a/src/khoj/processor/conversation/utils.py
+++ b/src/khoj/processor/conversation/utils.py
@@ -4,6 +4,7 @@ from time import perf_counter
 import json
 from datetime import datetime
 import queue
+from typing import Any, Dict, List
 import tiktoken
 
 # External packages
@@ -11,6 +12,8 @@ from langchain.schema import ChatMessage
 from transformers import AutoTokenizer
 
 # Internal Packages
+from khoj.database.adapters import ConversationAdapters
+from khoj.database.models import KhojUser
 from khoj.utils.helpers import merge_dicts
 
 
@@ -89,6 +92,32 @@ def message_to_log(
     return conversation_log
 
 
+def save_to_conversation_log(
+    q: str,
+    chat_response: str,
+    user: KhojUser,
+    meta_log: Dict,
+    user_message_time: str = None,
+    compiled_references: List[str] = [],
+    online_results: Dict[str, Any] = {},
+    inferred_queries: List[str] = [],
+    intent_type: str = "remember",
+):
+    user_message_time = user_message_time or datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+    updated_conversation = message_to_log(
+        user_message=q,
+        chat_response=chat_response,
+        user_message_metadata={"created": user_message_time},
+        khoj_message_metadata={
+            "context": compiled_references,
+            "intent": {"inferred-queries": inferred_queries, "type": intent_type},
+            "onlineContext": online_results,
+        },
+        conversation_log=meta_log.get("chat", []),
+    )
+    ConversationAdapters.save_conversation(user, {"chat": updated_conversation})
+
+
 def generate_chatml_messages_with_context(
     user_message,
     system_message,
diff --git a/src/khoj/routers/helpers.py b/src/khoj/routers/helpers.py
index 618f89ef..3e8ed155 100644
--- a/src/khoj/routers/helpers.py
+++ b/src/khoj/routers/helpers.py
@@ -19,7 +19,7 @@ from khoj.database.models import KhojUser, Subscription
 from khoj.processor.conversation import prompts
 from khoj.processor.conversation.offline.chat_model import converse_offline, send_message_to_model_offline
 from khoj.processor.conversation.openai.gpt import converse, send_message_to_model
-from khoj.processor.conversation.utils import ThreadedGenerator, message_to_log
+from khoj.processor.conversation.utils import ThreadedGenerator, save_to_conversation_log
 
 # Internal Packages
 from khoj.utils import state
@@ -186,30 +186,7 @@ def generate_chat_response(
     conversation_command: ConversationCommand = ConversationCommand.Default,
     user: KhojUser = None,
 ) -> Tuple[Union[ThreadedGenerator, Iterator[str]], Dict[str, str]]:
-    def _save_to_conversation_log(
-        q: str,
-        chat_response: str,
-        user_message_time: str,
-        compiled_references: List[str],
-        online_results: Dict[str, Any],
-        inferred_queries: List[str],
-        meta_log,
-    ):
-        updated_conversation = message_to_log(
-            user_message=q,
-            chat_response=chat_response,
-            user_message_metadata={"created": user_message_time},
-            khoj_message_metadata={
-                "context": compiled_references,
-                "intent": {"inferred-queries": inferred_queries},
-                "onlineContext": online_results,
-            },
-            conversation_log=meta_log.get("chat", []),
-        )
-        ConversationAdapters.save_conversation(user, {"chat": updated_conversation})
-
     # Initialize Variables
-    user_message_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
     chat_response = None
     logger.debug(f"Conversation Type: {conversation_command.name}")
 
@@ -217,13 +194,13 @@ def generate_chat_response(
 
     try:
         partial_completion = partial(
-            _save_to_conversation_log,
+            save_to_conversation_log,
             q,
-            user_message_time=user_message_time,
+            user=user,
+            meta_log=meta_log,
             compiled_references=compiled_references,
             online_results=online_results,
             inferred_queries=inferred_queries,
-            meta_log=meta_log,
         )
 
         conversation_config = ConversationAdapters.get_valid_conversation_config(user)