From a0b00ce4a13c2f27141ca5c41d1f3034b69ce265 Mon Sep 17 00:00:00 2001
From: sabaimran <narmiabas@gmail.com>
Date: Fri, 29 Nov 2024 18:10:14 -0800
Subject: [PATCH] Don't include null attributes when filling in stored
 conversation metadata

- Prompt adjustments to indicate to LLM what context it has
---
 src/khoj/database/models/__init__.py       |  2 +-
 src/khoj/processor/conversation/prompts.py |  2 +-
 src/khoj/processor/conversation/utils.py   | 41 ++++++++++++----------
 3 files changed, 25 insertions(+), 20 deletions(-)

diff --git a/src/khoj/database/models/__init__.py b/src/khoj/database/models/__init__.py
index ae6a729d..88b15404 100644
--- a/src/khoj/database/models/__init__.py
+++ b/src/khoj/database/models/__init__.py
@@ -108,7 +108,7 @@ class ChatMessage(PydanticBaseModel):
     created: str
     images: Optional[List[str]] = None
     queryFiles: Optional[List[Dict]] = None
-    excalidrawDiagram: Optional[str] = None
+    excalidrawDiagram: Optional[List[Dict]] = None
     by: str
     turnId: Optional[str]
     intent: Optional[Intent] = None
diff --git a/src/khoj/processor/conversation/prompts.py b/src/khoj/processor/conversation/prompts.py
index a943cbd0..5e28a912 100644
--- a/src/khoj/processor/conversation/prompts.py
+++ b/src/khoj/processor/conversation/prompts.py
@@ -186,7 +186,7 @@ Here is the image you generated based on my query. You can follow-up with a gene
 
 generated_diagram_attachment = PromptTemplate.from_template(
     f"""
-The AI has successfully created a diagram based on the user's query and handled the request. Good job!
+The AI has successfully created a diagram based on the user's query and handled the request. Good job! This will be shared with the user.
 
 AI can follow-up with a general response or summary. Limit to 1-2 sentences.
 """.strip()
diff --git a/src/khoj/processor/conversation/utils.py b/src/khoj/processor/conversation/utils.py
index 4d238b3b..86c2d87e 100644
--- a/src/khoj/processor/conversation/utils.py
+++ b/src/khoj/processor/conversation/utils.py
@@ -276,27 +276,32 @@ def save_to_conversation_log(
 ):
     user_message_time = user_message_time or datetime.now().strftime("%Y-%m-%d %H:%M:%S")
     turn_id = tracer.get("mid") or str(uuid.uuid4())
+
+    user_message_metadata = {"created": user_message_time, "images": query_images, "turnId": turn_id}
+
+    if raw_query_files and len(raw_query_files) > 0:
+        user_message_metadata["queryFiles"] = [file.model_dump(mode="json") for file in raw_query_files]
+
+    khoj_message_metadata = {
+        "context": compiled_references,
+        "intent": {"inferred-queries": inferred_queries, "type": intent_type},
+        "onlineContext": online_results,
+        "codeContext": code_results,
+        "automationId": automation_id,
+        "trainOfThought": train_of_thought,
+        "turnId": turn_id,
+        "images": generated_images,
+        "queryFiles": [file.model_dump(mode="json") for file in raw_generated_files],
+    }
+
+    if generated_excalidraw_diagram:
+        khoj_message_metadata["excalidrawDiagram"] = generated_excalidraw_diagram
+
     updated_conversation = message_to_log(
         user_message=q,
         chat_response=chat_response,
-        user_message_metadata={
-            "created": user_message_time,
-            "images": query_images,
-            "turnId": turn_id,
-            "queryFiles": [file.model_dump(mode="json") for file in raw_query_files],
-        },
-        khoj_message_metadata={
-            "context": compiled_references,
-            "intent": {"inferred-queries": inferred_queries, "type": intent_type},
-            "onlineContext": online_results,
-            "codeContext": code_results,
-            "automationId": automation_id,
-            "trainOfThought": train_of_thought,
-            "turnId": turn_id,
-            "images": generated_images,
-            "queryFiles": [file.model_dump(mode="json") for file in raw_generated_files],
-            "excalidrawDiagram": str(generated_excalidraw_diagram) if generated_excalidraw_diagram else None,
-        },
+        user_message_metadata=user_message_metadata,
+        khoj_message_metadata=khoj_message_metadata,
         conversation_log=meta_log.get("chat", []),
     )
     ConversationAdapters.save_conversation(