From 266936a4b508ac57ddd9f7b0ff84142bc642552b Mon Sep 17 00:00:00 2001
From: Daniel O'Connell <github@ahiru.pl>
Date: Sat, 30 Sep 2023 18:02:35 +0200
Subject: [PATCH] Notify when fetching followups

---
 api/src/stampy_chat/chat.py        | 11 +++++++----
 api/tests/stampy_chat/test_chat.py | 17 +++++++++--------
 web/src/hooks/useSearch.ts         | 18 +++++++++---------
 web/src/pages/index.tsx            |  6 ++++++
 4 files changed, 31 insertions(+), 21 deletions(-)

diff --git a/api/src/stampy_chat/chat.py b/api/src/stampy_chat/chat.py
index 3a417e6..96e67db 100644
--- a/api/src/stampy_chat/chat.py
+++ b/api/src/stampy_chat/chat.py
@@ -226,11 +226,14 @@ def talk_to_robot_internal(index, query: str, mode: str, history: Prompt, sessio
 
         logger.interaction(session_id, query, response, history, prompt, top_k_blocks)
 
-        # yield done state, possibly with followup questions
-        fin_json = {'state': 'done'}
+        yield {"state": "loading", "phase": "followups"}
+        # yield optional followups
         followups = multisearch_authored([query, response])
-        for i, followup in enumerate(followups):
-            fin_json[f'followup_{i}'] = asdict(followup)
+        if followups:
+            yield {'state': 'followups', 'followups': list(map(asdict, followups))}
+
+        # yield done state
+        fin_json = {'state': 'done'}
         yield fin_json
 
     except Exception as e:
diff --git a/api/tests/stampy_chat/test_chat.py b/api/tests/stampy_chat/test_chat.py
index 38f0a32..b1b64f5 100644
--- a/api/tests/stampy_chat/test_chat.py
+++ b/api/tests/stampy_chat/test_chat.py
@@ -268,7 +268,7 @@ def test_talk_to_robot_internal(history, context):
     with patch('stampy_chat.chat.get_top_k_blocks', return_value=context):
         with patch('stampy_chat.chat.multisearch_authored', return_value=followups):
             with patch('openai.ChatCompletion.create', return_value=chunks):
-                assert list(talk_to_robot_internal("index", "what is this about?", "default", history)) == [
+                assert list(talk_to_robot_internal("index", "what is this about?", "default", history, 'session id')) == [
                     {'phase': 'semantic', 'state': 'loading'},
                     {'citations': [], 'phase': 'semantic', 'state': 'loading'},
                     {'phase': 'prompt', 'state': 'loading'},
@@ -277,12 +277,13 @@ def test_talk_to_robot_internal(history, context):
                     {'content': 'response 2', 'state': 'streaming'},
                     {'content': 'response 3', 'state': 'streaming'},
                     {'content': 'response 4', 'state': 'streaming'},
-                    {
-                        'followup_0': {'pageid': '1', 'score': 0.231, 'text': 'followup 1'},
-                        'followup_1': {'pageid': '2', 'score': 0.231, 'text': 'followup 2'},
-                        'followup_2': {'pageid': '3', 'score': 0.231, 'text': 'followup 3'},
-                        'state': 'done'
-                    },
+                    {'state': 'loading', 'phase': 'followups'},
+                    {'state': 'followups', 'followups': [
+                        {'pageid': '1', 'score': 0.231, 'text': 'followup 1'},
+                        {'pageid': '2', 'score': 0.231, 'text': 'followup 2'},
+                        {'pageid': '3', 'score': 0.231, 'text': 'followup 3'},
+                    ]},
+                    {'state': 'done'},
                 ]
 
 
@@ -297,7 +298,7 @@ def test_talk_to_robot_internal_error(history, context):
     ]
     with patch('stampy_chat.chat.get_top_k_blocks', return_value=context):
         with patch('openai.ChatCompletion.create', return_value=chunks):
-            assert list(talk_to_robot_internal("index", "what is this about?", "default", history)) == [
+            assert list(talk_to_robot_internal("index", "what is this about?", "default", history, 'session id')) == [
                 {'phase': 'semantic', 'state': 'loading'},
                 {'citations': [], 'phase': 'semantic', 'state': 'loading'},
                 {'phase': 'prompt', 'state': 'loading'},
diff --git a/web/src/hooks/useSearch.ts b/web/src/hooks/useSearch.ts
index 44e48ee..919d62d 100644
--- a/web/src/hooks/useSearch.ts
+++ b/web/src/hooks/useSearch.ts
@@ -12,7 +12,7 @@ import type {
 
 const MAX_FOLLOWUPS = 4;
 const DATA_HEADER = "data: "
-const EVENT_END_HEADER = "event: close\n"
+const EVENT_END_HEADER = "event: close"
 
 type HistoryEntry = {
   role: "error" | "stampy" | "assistant" | "user";
@@ -29,13 +29,13 @@ export async function* iterateData(res: Response) {
     if (done) return;
 
     const chunk = new TextDecoder("utf-8").decode(value);
-    if (chunk.startsWith(EVENT_END_HEADER)) return;
-
     for (const line of chunk.split("\n")) {
       // Most times, it seems that a single read() call will be one SSE "message",
       // but I'll do the proper aggregation spec thing in case that's not always true.
 
-      if (line.startsWith(DATA_HEADER)) {
+      if (line.startsWith(EVENT_END_HEADER)) {
+        return;
+      } else if (line.startsWith(DATA_HEADER)) {
         message += line.slice(DATA_HEADER.length);
         // Fixes #43
       } else if (line !== "") {
@@ -83,12 +83,12 @@ export const extractAnswer = async (
         setCurrent({ phase: "streaming", ...result });
         break;
 
+      case "followups":
+         // add any potential followup questions
+         followups = data.followups.map((value) => value as Followup);
+         break;
       case "done":
-        // add any potential followup questions
-        const followups = Object.entries(data)
-          .filter(([key]) => key.startsWith("followup_"))
-          .map(([k, value]) => value as Followup);
-        return { result, followups };
+        break;
       case "error":
         throw data.error;
     }
diff --git a/web/src/pages/index.tsx b/web/src/pages/index.tsx
index 5cb65ae..0a8c79b 100644
--- a/web/src/pages/index.tsx
+++ b/web/src/pages/index.tsx
@@ -118,6 +118,12 @@ const Home: NextPage = () => {
     case "streaming":
       last_entry = <ShowAssistantEntry entry={current} />;
       break;
+    case "followups":
+      last_entry = <>
+          <ShowAssistantEntry entry={current} />
+          <p>Loading: Checking for followups...</p>
+      </>;
+      break;
   }
 
   return (