From a6f6e4f418c328f9d9cf74fb244a10ab40d2428d Mon Sep 17 00:00:00 2001
From: sabaimran <narmiabas@gmail.com>
Date: Wed, 9 Oct 2024 20:34:20 -0700
Subject: [PATCH] Fix notes references and passage of user query in the chat
 flow

---
 src/khoj/routers/api_chat.py | 14 ++++++--------
 src/khoj/routers/helpers.py  |  2 +-
 src/khoj/routers/research.py |  4 ++--
 3 files changed, 9 insertions(+), 11 deletions(-)

diff --git a/src/khoj/routers/api_chat.py b/src/khoj/routers/api_chat.py
index 89b56a29..b2689a7f 100644
--- a/src/khoj/routers/api_chat.py
+++ b/src/khoj/routers/api_chat.py
@@ -720,13 +720,11 @@ async def chat(
                 if type(research_result) == InformationCollectionIteration:
                     if research_result.summarizedResult:
                         pending_research = False
-                        # if research_result.onlineContext:
-                        #     researched_results += str(research_result.onlineContext)
-                        #     online_results.update(research_result.onlineContext)
+                        if research_result.onlineContext:
+                            online_results.update(research_result.onlineContext)
 
-                        # if research_result.context:
-                        #     researched_results += str(research_result.context)
-                        #     compiled_references.extend(research_result.context)
+                        if research_result.context:
+                            compiled_references.extend(research_result.context)
 
                         researched_results += research_result.summarizedResult
 
@@ -1021,10 +1019,9 @@ async def chat(
         async for result in send_event(ChatEvent.STATUS, f"**Generating a well-informed response**"):
             yield result
         llm_response, chat_metadata = await agenerate_chat_response(
-            defiltered_query,
+            q,
             meta_log,
             conversation,
-            researched_results,
             compiled_references,
             online_results,
             inferred_queries,
@@ -1035,6 +1032,7 @@ async def chat(
             location,
             user_name,
             uploaded_image_url,
+            researched_results,
         )
 
         # Send Response
diff --git a/src/khoj/routers/helpers.py b/src/khoj/routers/helpers.py
index 279ad85e..e4ebdb51 100644
--- a/src/khoj/routers/helpers.py
+++ b/src/khoj/routers/helpers.py
@@ -947,7 +947,6 @@ def generate_chat_response(
     q: str,
     meta_log: dict,
     conversation: Conversation,
-    meta_research: str = "",
     compiled_references: List[Dict] = [],
     online_results: Dict[str, Dict] = {},
     inferred_queries: List[str] = [],
@@ -958,6 +957,7 @@ def generate_chat_response(
     location_data: LocationData = None,
     user_name: Optional[str] = None,
     uploaded_image_url: Optional[str] = None,
+    meta_research: str = "",
 ) -> Tuple[Union[ThreadedGenerator, Iterator[str]], Dict[str, str]]:
     # Initialize Variables
     chat_response = None
diff --git a/src/khoj/routers/research.py b/src/khoj/routers/research.py
index 028fffb7..5960c88c 100644
--- a/src/khoj/routers/research.py
+++ b/src/khoj/routers/research.py
@@ -33,7 +33,7 @@ class InformationCollectionIteration:
         self,
         data_source: str,
         query: str,
-        context: str = None,
+        context: Dict[str, Dict] = None,
         onlineContext: dict = None,
         summarizedResult: str = None,
     ):
@@ -187,7 +187,7 @@ async def execute_information_collection(
                     compiled_references.extend(result[0])
                     inferred_queries.extend(result[1])
                     defiltered_query = result[2]
-                    this_iteration.context = str(compiled_references)
+                    this_iteration.context = compiled_references
 
         elif this_iteration.data_source == ConversationCommand.Online:
             async for result in search_online(