From 4fbaef10e93e2b89cabcee09a188cec392b31a74 Mon Sep 17 00:00:00 2001
From: sabaimran <narmiabas@gmail.com>
Date: Wed, 9 Oct 2024 15:58:05 -0700
Subject: [PATCH] Correct usage of the summarize function

---
 src/khoj/routers/api_chat.py |  6 ++++--
 src/khoj/routers/research.py | 15 +++++++++++++--
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/src/khoj/routers/api_chat.py b/src/khoj/routers/api_chat.py
index 7b0241b5..64c50a1c 100644
--- a/src/khoj/routers/api_chat.py
+++ b/src/khoj/routers/api_chat.py
@@ -783,7 +783,7 @@ async def chat(
                 async for result in send_llm_response(response_log):
                     yield result
             else:
-                response_log = await generate_summary_from_files(
+                async for response in generate_summary_from_files(
                     q=q,
                     user=user,
                     file_filters=file_filters,
@@ -791,7 +791,9 @@ async def chat(
                     subscribed=subscribed,
                     send_status_func=partial(send_event, ChatEvent.STATUS),
                     send_response_func=partial(send_llm_response),
-                )
+                ):
+                    yield response
+
             await sync_to_async(save_to_conversation_log)(
                 q,
                 response_log,
diff --git a/src/khoj/routers/research.py b/src/khoj/routers/research.py
index 65c3f42d..4211b072 100644
--- a/src/khoj/routers/research.py
+++ b/src/khoj/routers/research.py
@@ -245,14 +245,25 @@ async def execute_information_collection(
                     )
                 )
             else:
-                response_log = await generate_summary_from_files(
+                async for response in generate_summary_from_files(
                     q=query,
                     user=user,
                     file_filters=file_filters,
                     meta_log=conversation_history,
                     subscribed=subscribed,
                     send_status_func=send_status_func,
-                )
+                ):
+                    if isinstance(response, dict) and ChatEvent.STATUS in response:
+                        yield response[ChatEvent.STATUS]
+                    else:
+                        response_log = response
+                        previous_iterations.append(
+                            InformationCollectionIteration(
+                                data_source=this_iteration.data_source,
+                                query=this_iteration.query,
+                                context=response_log,
+                            )
+                        )
         else:
             iteration = MAX_ITERATIONS