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 =
Loading: Checking for followups...
+ >; + break; } return (