From a11b5293fbe535fb87019fac6de087550b3de511 Mon Sep 17 00:00:00 2001
From: sabaimran <narmiabas@gmail.com>
Date: Wed, 23 Oct 2024 21:45:17 -0700
Subject: [PATCH] Add uploaded images to research mode, code slash command,
 include code references

---
 .../components/chatMessage/chatMessage.tsx    |  3 +++
 src/khoj/routers/api_chat.py                  | 25 ++++++++++---------
 src/khoj/routers/helpers.py                   |  2 ++
 3 files changed, 18 insertions(+), 12 deletions(-)

diff --git a/src/interface/web/app/components/chatMessage/chatMessage.tsx b/src/interface/web/app/components/chatMessage/chatMessage.tsx
index 85bcc517..1c7757b4 100644
--- a/src/interface/web/app/components/chatMessage/chatMessage.tsx
+++ b/src/interface/web/app/components/chatMessage/chatMessage.tsx
@@ -301,6 +301,9 @@ export function TrainOfThought(props: TrainOfThoughtProps) {
     const iconColor = props.primary ? convertColorToTextClass(props.agentColor) : "text-gray-500";
     const icon = chooseIconFromHeader(header, iconColor);
     let markdownRendered = DOMPurify.sanitize(md.render(props.message));
+
+    // Remove any header tags from markdownRendered
+    markdownRendered = markdownRendered.replace(/<h[1-6].*?<\/h[1-6]>/g, "");
     return (
         <div
             className={`${styles.trainOfThoughtElement} break-all items-center ${props.primary ? "text-gray-400" : "text-gray-300"} ${styles.trainOfThought} ${props.primary ? styles.primary : ""}`}
diff --git a/src/khoj/routers/api_chat.py b/src/khoj/routers/api_chat.py
index 32513b55..5d811fe7 100644
--- a/src/khoj/routers/api_chat.py
+++ b/src/khoj/routers/api_chat.py
@@ -690,14 +690,14 @@ async def chat(
         inferred_queries: List[Any] = []
         defiltered_query = defilter_query(q)
 
-        if conversation_commands == [ConversationCommand.Default] or is_automated_task:
+        if conversation_commands == [ConversationCommand.Default]:
             async for research_result in execute_information_collection(
                 request=request,
                 user=user,
                 query=defiltered_query,
                 conversation_id=conversation_id,
                 conversation_history=meta_log,
-                query_images=raw_images,
+                query_images=uploaded_images,
                 agent=agent,
                 send_status_func=partial(send_event, ChatEvent.STATUS),
                 user_name=user_name,
@@ -956,16 +956,17 @@ async def chat(
                     ):
                         yield result
 
-            ## Send Gathered References
-            async for result in send_event(
-                ChatEvent.REFERENCES,
-                {
-                    "inferredQueries": inferred_queries,
-                    "context": compiled_references,
-                    "onlineContext": online_results,
-                },
-            ):
-                yield result
+        ## Send Gathered References
+        async for result in send_event(
+            ChatEvent.REFERENCES,
+            {
+                "inferredQueries": inferred_queries,
+                "context": compiled_references,
+                "onlineContext": online_results,
+                "codeContext": code_results,
+            },
+        ):
+            yield result
 
         if pending_research:
             ## Gather Code Results
diff --git a/src/khoj/routers/helpers.py b/src/khoj/routers/helpers.py
index bfe25fe3..083c38a4 100644
--- a/src/khoj/routers/helpers.py
+++ b/src/khoj/routers/helpers.py
@@ -232,6 +232,8 @@ def get_conversation_command(query: str, any_references: bool = False) -> Conver
         return ConversationCommand.Summarize
     elif query.startswith("/diagram"):
         return ConversationCommand.Diagram
+    elif query.startswith("/code"):
+        return ConversationCommand.Code
     # If no relevant notes found for the given query
     elif not any_references:
         return ConversationCommand.General