From eee29179db0a1bdb98eccd5703a04d94f410ba97 Mon Sep 17 00:00:00 2001
From: Fraser <tomatoesneverlie@gmail.com>
Date: Fri, 2 Jun 2023 00:02:39 -0400
Subject: [PATCH] pass followups to front-end on finish response

---
 api/chat.py             | 13 ++++++++-----
 web/src/pages/index.tsx | 27 +++++++++++++++++++++------
 2 files changed, 29 insertions(+), 11 deletions(-)

diff --git a/api/chat.py b/api/chat.py
index e266206..8c7b230 100644
--- a/api/chat.py
+++ b/api/chat.py
@@ -1,13 +1,13 @@
 # ------------------------------- env, constants -------------------------------
 from followups import multisearch_authored
-
 from get_blocks import get_top_k_blocks, Block
 
+from dataclasses import asdict
 from typing import List, Dict, Callable
 import openai
+import re
 import tiktoken
 import time
-import re
 
 # OpenAI models
 EMBEDDING_MODEL = "text-embedding-ada-002"
@@ -171,9 +171,12 @@ def talk_to_robot(index, query: str, history: List[Dict[str, str]], k: int = STA
         log(query)
         log(response)
 
-        multisearch_authored([query, response], DEBUG_PRINT)
-
-        yield json.dumps({"state": "done"})
+        # yield done state, possibly with followup questions
+        fin_json = {"state": "done"}
+        followups = multisearch_authored([query, response], DEBUG_PRINT)
+        for i, followup in enumerate(followups):
+            fin_json[f"followup_{i}"] = asdict(followup)
+        yield json.dumps(fin_json)
 
     except Exception as e:
         print(e)
diff --git a/web/src/pages/index.tsx b/web/src/pages/index.tsx
index 1ea57c0..687d66a 100644
--- a/web/src/pages/index.tsx
+++ b/web/src/pages/index.tsx
@@ -101,7 +101,7 @@ const A: React.FC<{href: string, className?: string, children: React.ReactNode}>
 const ProcessText: (text: string, base_count: number) => [string, Map<string, number>] = (text, base_count) => {
 
     // ---------------------- normalize citation form ----------------------
-    // the general plan here is just to add parsing cases until we can respond 
+    // the general plan here is just to add parsing cases until we can respond
     // well to almost everything the LLM emits. We won't ever reach five nines,
     // but the domain is one where occasionally failing isn't catastrophic.
 
@@ -214,8 +214,15 @@ const ShowAssistantEntry: React.FC<{entry: AssistantEntry}> = ({entry}) => {
 
 
 
+type Followup = {
+    text: string;
+    pageid: string;
+    score: number;
+}
+
 type State = {
     state: "idle";
+    followups: Followup[];
 } | {
     state: "loading";
     phase: "semantic" | "prompt" | "llm";
@@ -237,7 +244,7 @@ const Home: NextPage = () => {
 
     const [ entries, setEntries ] = useState<Entry[]>([]);
     const [ runningIndex, setRunningIndex ] = useState(0);
-    const [ loadState, setLoadState ] = useState<State>({state: "idle"});
+    const [ loadState, setLoadState ] = useState<State>({state: "idle", followups: []});
 
     const search = async (
         query: string,
@@ -279,7 +286,7 @@ const Home: NextPage = () => {
 
         if (!res.ok) {
             setLoading(false);
-            setLoadState({state: "idle"});
+            setLoadState({state: "idle", followups: []});
             setEntries([...new_entries, {role: "error", content: "POST Error: " + res.status}]);
             return;
         }
@@ -351,15 +358,24 @@ const Home: NextPage = () => {
 
                             case "done":
 
-                                // append the response to the entries, reset to normal
+                                // append the response to the entries, add any potential followup questions, reset to normal
+                                var followups: Followup[] = [];
+                                var i = 0;
+                                while ('followup_' + i in data) {
+                                    followups = [...followups, data['followup_' + i]];
+                                    i++;
+                                }
 
                                 setLoadState((s) => {
                                     if (s.state === "streaming") {
                                         setEntries([...new_entries, s.response]);
                                         setRunningIndex((i) => (i + ProcessText(s.response.content, 0)[1].size));
                                     }
-                                    return {state: "idle"};
+
+
+                                    return {state: "idle", followups: followups};
                                 });
+
                                 break read;
 
                             case "error":
@@ -374,7 +390,6 @@ const Home: NextPage = () => {
         }
 
         setLoading(false);
-        setLoadState({state: "idle"});
         scroll30();
     };