From 0a7f68f1fce5497a633d87b67779ecf49ccdf12f Mon Sep 17 00:00:00 2001
From: Fraser <tomatoesneverlie@gmail.com>
Date: Sun, 16 Apr 2023 18:25:55 -0400
Subject: [PATCH 1/6] working SSE streaming

---
 api/Procfile            |   2 +-
 api/chat.py             |  93 +++++++++--------
 api/main.py             |  30 ++++--
 web/src/pages/index.tsx | 218 ++++++++++++++++++++++++----------------
 web/src/searchbox.tsx   |   2 +-
 5 files changed, 204 insertions(+), 141 deletions(-)

diff --git a/api/Procfile b/api/Procfile
index 1945201..0d052ab 100644
--- a/api/Procfile
+++ b/api/Procfile
@@ -1 +1 @@
-web: gunicorn main:app
+web: gunicorn main:app --worker-class eventlet --threads 4
diff --git a/api/chat.py b/api/chat.py
index 5d5b30e..9b28faa 100644
--- a/api/chat.py
+++ b/api/chat.py
@@ -24,7 +24,7 @@ CONTEXT_FRACTION = 0.5  # the (approximate) fraction of num_tokens to use for co
 
 ENCODER = tiktoken.get_encoding("cl100k_base")
 
-DEBUG_PRINT = False
+DEBUG_PRINT = True
 
 # --------------------------------- prompt code --------------------------------
 
@@ -116,48 +116,59 @@ def construct_prompt(query: str, history: List[Dict[str, str]], context: List[Bl
     return prompt
 
 # ------------------------------- completion code -------------------------------
+import time
+import json
 
 # returns either (True, reply string, top_k_blocks)) or (False, error message string, None)
 def talk_to_robot(index, query: str, history: List[Dict[str, str]], k: int = STANDARD_K):
+    yield json.dumps({"state": "loading", "phase": "semantic"})
+    time.sleep(1)
+    yield json.dumps({"state": "loading", "phase": "prompt"})
+    time.sleep(1)
+    yield json.dumps({"state": "loading", "phase": "llm"})
+    time.sleep(1)
+    for c in "Hi. I'm a big dumb LLM. Hubris will be the end of us all.":
+        yield json.dumps({"state": "streaming", "response": c})
+        time.sleep(0.1)
 
-    try:
-        # 1. Find the most relevant blocks from the Alignment Research Dataset
-        top_k_blocks = get_top_k_blocks(index, query, k)
-
-
-        # 2. Generate a prompt
-        prompt = construct_prompt(query, history, top_k_blocks)
-
-
-        # 3. Count number of tokens left for completion (-50 for a buffer)
-        max_tokens_completion = NUM_TOKENS - sum([len(ENCODER.encode(message["content"]) + ENCODER.encode(message["role"])) for message in prompt]) - 50
-
-        # 4. Answer the user query
-        t1 = time.time()
-        response = openai.ChatCompletion.create(
-            model=COMPLETIONS_MODEL,
-            messages=prompt,
-            max_tokens=max_tokens_completion
-        )["choices"][0]["message"]["content"]
-        t2 = time.time()
-        print("Time to get response: ", t2 - t1)
-        
-
-        if DEBUG_PRINT:
-            print('\n' * 10)
-            print(" ------------------------------ prompt: -----------------------------")
-            for message in prompt:
-                print(f"----------- {message['role']}: ------------------")
-                print(message['content'])
-
-            print('\n' * 10)
-
-            print(" ------------------------------ response: -----------------------------")
-            print(response)
-
-        return (True, response, top_k_blocks)
-
-    except Exception as e:
-        print(e)
-        return (False, "Error: " + str(e), None)
+    # try:
+    #     # 1. Find the most relevant blocks from the Alignment Research Dataset
+    #     top_k_blocks = get_top_k_blocks(index, query, k)
+    #
+    #
+    #     # 2. Generate a prompt
+    #     prompt = construct_prompt(query, history, top_k_blocks)
+    #
+    #
+    #     # 3. Count number of tokens left for completion (-50 for a buffer)
+    #     max_tokens_completion = NUM_TOKENS - sum([len(ENCODER.encode(message["content"]) + ENCODER.encode(message["role"])) for message in prompt]) - 50
+    #
+    #     # 4. Answer the user query
+    #     t1 = time.time()
+    #     response = openai.ChatCompletion.create(
+    #         model=COMPLETIONS_MODEL,
+    #         messages=prompt,
+    #         max_tokens=max_tokens_completion
+    #     )["choices"][0]["message"]["content"]
+    #     t2 = time.time()
+    #     print("Time to get response: ", t2 - t1)
+    #     
+    #
+    #     if DEBUG_PRINT:
+    #         print('\n' * 10)
+    #         print(" ------------------------------ prompt: -----------------------------")
+    #         for message in prompt:
+    #             print(f"----------- {message['role']}: ------------------")
+    #             print(message['content'])
+    #
+    #         print('\n' * 10)
+    #
+    #         print(" ------------------------------ response: -----------------------------")
+    #         print(response)
+    #
+    #     return (True, response, top_k_blocks)
+    #
+    # except Exception as e:
+    #     print(e)
+    #     return (False, "Error: " + str(e), None)
 
diff --git a/api/main.py b/api/main.py
index 099c2ca..96b8e04 100644
--- a/api/main.py
+++ b/api/main.py
@@ -1,4 +1,4 @@
-from flask import Flask, jsonify, request
+from flask import Flask, jsonify, request, Response
 from flask_cors import CORS, cross_origin
 from get_blocks import get_top_k_blocks
 from chat import talk_to_robot
@@ -34,6 +34,12 @@ app = Flask(__name__)
 cors = CORS(app)
 app.config['CORS_HEADERS'] = 'Content-Type'
 
+# ---------------------------------- sse stuff ---------------------------------
+
+def stream(src):
+    yield from ('data: ' + '\ndata: '.join(message.splitlines()) + '\n\n' for message in src)
+    yield 'data: close\n\n'
+
 # ------------------------------- semantic search ------------------------------
 
 
@@ -50,20 +56,22 @@ def semantic():
 @app.route('/chat', methods=['POST'])
 @cross_origin()
 def chat():
+    # 
+    # query = request.json['query']
+    # history = request.json['history']
+
+    return Response(stream(talk_to_robot(None, None, None)), mimetype='text/event-stream')
+
+    # is_valid, response, context = talk_to_robot(index, query, history)
     
-    query = request.json['query']
-    history = request.json['history']
-
-    is_valid, response, context = talk_to_robot(index, query, history)
-
-    if is_valid:
-        return jsonify({'response': response, 'citations': [{'title': block.title, 'author': block.author, 'date': block.date, 'url': block.url} for block in context]})
-    else:
-        return jsonify({'error': response})
+    # if is_valid:
+    #     return jsonify({'response': response, 'citations': [{'title': block.title, 'author': block.author, 'date': block.date, 'url': block.url} for block in context]})
+    # else:
+    #     return jsonify({'error': response})
 
 
 # ------------------------------------------------------------------------------
 
 
 if __name__ == '__main__':
-    app.run(debug=True, port=3000)
\ No newline at end of file
+    app.run(debug=True, port=3000)
diff --git a/web/src/pages/index.tsx b/web/src/pages/index.tsx
index 6f6e397..cd0e695 100644
--- a/web/src/pages/index.tsx
+++ b/web/src/pages/index.tsx
@@ -121,10 +121,22 @@ const ShowEntry: React.FC<{entry: Entry}> = ({entry}) => {
     );
 };
 
+type State = {
+    state: "idle";
+} | {
+    state: "loading";
+    phase: "semantic" | "prompt" | "llm";
+} | {
+    state: "streaming";
+    response: string;
+};
+
+
 const Home: NextPage = () => {
 
     const [ entries, setEntries ] = useState<Entry[]>([]);
     const [ runningIndex, setRunningIndex ] = useState(0);
+    const [ loadState, setLoadState ] = useState<State>({state: "idle"});
 
     const search = async (
         query: string,
@@ -140,6 +152,37 @@ const Home: NextPage = () => {
 
         setLoading(true);
 
+
+
+        // sse connection
+        const eventSource = new EventSource(API_URL + "/chat");
+                            
+        eventSource.onmessage = (event) => {
+
+            if (event.data === "close") {
+                eventSource.close();
+                setLoading(false);
+                setLoadState({state: "idle"});
+                return;
+            }
+
+            console.log(event.data);
+            const data = JSON.parse(event.data);
+
+            if (data.state === "loading") {
+                setLoadState({state: "loading", phase: data.phase});
+                return;
+            }
+
+            if (data.state === "streaming") {
+                setLoadState((s) => {
+                    const response = s.state === "streaming" ? s.response : "";
+                    return {state: "streaming", response: response + data.response};
+                });
+                return;
+            }
+        }
+
         const res = await fetch(API_URL + "/chat", {
             method: "POST",
             headers: { "Content-Type": "application/json", "Allow-Control-Allow-Origin": "*" },
@@ -154,93 +197,93 @@ const Home: NextPage = () => {
                })
         })
 
-        if (!res.ok) {
-            setLoading(false);
-            console.log("load failure: " + res.status);
-            return;
-        }
-
-        const data = await res.json();
-
-        // -------------------------- error checking ---------------------------
-
-        if (data.error) {
-            setEntries([...new_entries, {role: "error", content: data.error}]);
-            setLoading(false);
-            return;
-        }
-
-        // ---------------------- normalize citation form ----------------------
-
-        // transform all things that look like [a, b, c] into [a][b][c]
-        let response = data.response.replace(
-
-                    /\[((?:[a-z]+,\s*)*[a-z]+)\]/g, // identify groups of this form
-
-                    (block: string) => block.split(',')
-                                            .map((x) => x.trim())
-                                            .join("][")
-        )
-
-        // transform all things that look like [(a), (b), (c)] into [(a)][(b)][(c)]
-        response = response.replace(
-            
-                /\[((?:\([a-z]+\),\s*)*\([a-z]+\))\]/g, // identify groups of this form
-
-                (block: string) => block.split(',')
-                                        .map((x) => x.trim())
-                                        .join("][")
-        )
-
-        // transform all things that look like [(a)] into [a]
-        response = response.replace(
-            /\[\(([a-z]+)\)\]/g,
-            (_match: string, x: string) => `[${x}]`
-        )
-
-        // transform all things that look like [ a ] into [a]
-        response = response.replace(
-            /\[\s*([a-z]+)\s*\]/g,
-            (_match: string, x: string) => `[${x}]`
-        )
-
-        // -------------- map citations from strings into numbers --------------
-
-        // figure out what citations are in the response, and map them appropriately
-        const cite_map = new Map<string, number>();
-        let cite_count = runningIndex;
-
-        // scan a regex for [x] over the response. If x isn't in the map, add it.
-        const regex = /\[([a-z]+)\]/g;
-        let match;
-        let response_copy = ""
-        while ((match = regex.exec(response)) !== null) {
-            if (!cite_map.has(match[1]!)) {
-                cite_map.set(match[1]!, cite_count++);
-            }
-            // replace [x] with [i]
-            response_copy += response.slice(response_copy.length, match.index) + `[${cite_map.get(match[1]!)! + 1}]`;
-        }
-
-        setRunningIndex(cite_count);
-
-        response = response_copy + response.slice(response_copy.length);
-
-        // ----------------- create the ordered citation array -----------------
-
-        const citations = new Map<number, Citation>();
-        cite_map.forEach((value, key) => {
-            const index = key.charCodeAt(0) - 'a'.charCodeAt(0);
-            if (index >= data.citations.length) {
-                console.log("invalid citation index: " + index);
-            } else {
-                citations.set(value, data.citations[index]);
-            }
-        });
-
-        setEntries([...new_entries, {role: "assistant", 
-                                     content: response,
-                                     citations: citations}]);
+        // if (!res.ok) {
+        //     setLoading(false);
+        //     console.log("load failure: " + res.status);
+        //     return;
+        // }
+        //
+        // const data = await res.json();
+        //
+        // // -------------------------- error checking ---------------------------
+        //
+        // if (data.error) {
+        //     setEntries([...new_entries, {role: "error", content: data.error}]);
+        //     setLoading(false);
+        //     return;
+        // }
+        //
+        // // ---------------------- normalize citation form ----------------------
+        //
+        // // transform all things that look like [a, b, c] into [a][b][c]
+        // let response = data.response.replace(
+        //
+        //             /\[((?:[a-z]+,\s*)*[a-z]+)\]/g, // identify groups of this form
+        //
+        //             (block: string) => block.split(',')
+        //                                     .map((x) => x.trim())
+        //                                     .join("][")
+        // )
+        //
+        // // transform all things that look like [(a), (b), (c)] into [(a)][(b)][(c)]
+        // response = response.replace(
+        //     
+        //         /\[((?:\([a-z]+\),\s*)*\([a-z]+\))\]/g, // identify groups of this form
+        //
+        //         (block: string) => block.split(',')
+        //                                 .map((x) => x.trim())
+        //                                 .join("][")
+        // )
+        //
+        // // transform all things that look like [(a)] into [a]
+        // response = response.replace(
+        //     /\[\(([a-z]+)\)\]/g,
+        //     (_match: string, x: string) => `[${x}]`
+        // )
+        //
+        // // transform all things that look like [ a ] into [a]
+        // response = response.replace(
+        //     /\[\s*([a-z]+)\s*\]/g,
+        //     (_match: string, x: string) => `[${x}]`
+        // )
+        //
+        // // -------------- map citations from strings into numbers --------------
+        //
+        // // figure out what citations are in the response, and map them appropriately
+        // const cite_map = new Map<string, number>();
+        // let cite_count = runningIndex;
+        //
+        // // scan a regex for [x] over the response. If x isn't in the map, add it.
+        // const regex = /\[([a-z]+)\]/g;
+        // let match;
+        // let response_copy = ""
+        // while ((match = regex.exec(response)) !== null) {
+        //     if (!cite_map.has(match[1]!)) {
+        //         cite_map.set(match[1]!, cite_count++);
+        //     }
+        //     // replace [x] with [i]
+        //     response_copy += response.slice(response_copy.length, match.index) + `[${cite_map.get(match[1]!)! + 1}]`;
+        // }
+        //
+        // setRunningIndex(cite_count);
+        //
+        // response = response_copy + response.slice(response_copy.length);
+        //
+        // // ----------------- create the ordered citation array -----------------
+        //
+        // const citations = new Map<number, Citation>();
+        // cite_map.forEach((value, key) => {
+        //     const index = key.charCodeAt(0) - 'a'.charCodeAt(0);
+        //     if (index >= data.citations.length) {
+        //         console.log("invalid citation index: " + index);
+        //     } else {
+        //         citations.set(value, data.citations[index]);
+        //     }
+        // });
+        //
+        // setEntries([...new_entries, {role: "assistant", 
+        //                              content: response,
+        //                              citations: citations}]);
 
         setLoading(false);
 
@@ -261,6 +304,7 @@ const Home: NextPage = () => {
                     ))}
                 </ul>
                 <SearchBox search={search} />
+                <p>{loadState.state === "loading" ? loadState.phase : ( loadState.state === "streaming" ? loadState.response : "" )}</p>
             </main>
         </>
     );
diff --git a/web/src/searchbox.tsx b/web/src/searchbox.tsx
index d21dedd..9b4e7dd 100644
--- a/web/src/searchbox.tsx
+++ b/web/src/searchbox.tsx
@@ -20,7 +20,7 @@ const SearchBox: React.FC<{search: (
         if (!loading) inputRef.current?.focus();
     }, [loading]);
 
-    if (loading) return <p>loading...</p>;
+    if (loading) return <></>;
     return (<>
         <form className="flex mb-2" onSubmit={async (e) => {
             e.preventDefault();

From 421b755a753c9654575f03ea469f7fea8468c02a Mon Sep 17 00:00:00 2001
From: Fraser <tomatoesneverlie@gmail.com>
Date: Sun, 16 Apr 2023 18:52:58 -0400
Subject: [PATCH 2/6] somewhat more complete integration

---
 api/chat.py             |  27 +++++-----
 api/main.py             |  12 ++---
 web/src/pages/index.tsx | 107 ++++++++++++++++++++++++++--------------
 3 files changed, 91 insertions(+), 55 deletions(-)

diff --git a/api/chat.py b/api/chat.py
index 9b28faa..8360857 100644
--- a/api/chat.py
+++ b/api/chat.py
@@ -121,27 +121,28 @@ import json
 
 # returns either (True, reply string, top_k_blocks)) or (False, error message string, None)
 def talk_to_robot(index, query: str, history: List[Dict[str, str]], k: int = STANDARD_K):
+    # 1. Find the most relevant blocks from the Alignment Research Dataset
     yield json.dumps({"state": "loading", "phase": "semantic"})
-    time.sleep(1)
+    top_k_blocks = get_top_k_blocks(index, query, k)
+
+    # 2. Generate a prompt
+    prompt = construct_prompt(query, history, top_k_blocks)
     yield json.dumps({"state": "loading", "phase": "prompt"})
-    time.sleep(1)
+
+    # 3. Count number of tokens left for completion (-50 for a buffer)
+    max_tokens_completion = NUM_TOKENS - sum([len(ENCODER.encode(message["content"]) + ENCODER.encode(message["role"])) for message in prompt]) - 50
+
+
+
+
+
     yield json.dumps({"state": "loading", "phase": "llm"})
     time.sleep(1)
-    for c in "Hi. I'm a big dumb LLM. Hubris will be the end of us all.":
+    for c in "Hi. I'm a big dumb LLM.\nHubris will be the end of us all.":
         yield json.dumps({"state": "streaming", "response": c})
         time.sleep(0.1)
 
     # try:
-    #     # 1. Find the most relevant blocks from the Alignment Research Dataset
-    #     top_k_blocks = get_top_k_blocks(index, query, k)
-    #
-    #
-    #     # 2. Generate a prompt
-    #     prompt = construct_prompt(query, history, top_k_blocks)
-    #
-    #
-    #     # 3. Count number of tokens left for completion (-50 for a buffer)
-    #     max_tokens_completion = NUM_TOKENS - sum([len(ENCODER.encode(message["content"]) + ENCODER.encode(message["role"])) for message in prompt]) - 50
     #
     #     # 4. Answer the user query
     #     t1 = time.time()
diff --git a/api/main.py b/api/main.py
index 96b8e04..8389e2b 100644
--- a/api/main.py
+++ b/api/main.py
@@ -38,7 +38,7 @@ app.config['CORS_HEADERS'] = 'Content-Type'
 
 def stream(src):
     yield from ('data: ' + '\ndata: '.join(message.splitlines()) + '\n\n' for message in src)
-    yield 'data: close\n\n'
+    yield 'event: close\n\n'
 
 # ------------------------------- semantic search ------------------------------
 
@@ -56,14 +56,14 @@ def semantic():
 @app.route('/chat', methods=['POST'])
 @cross_origin()
 def chat():
-    # 
-    # query = request.json['query']
-    # history = request.json['history']
 
-    return Response(stream(talk_to_robot(None, None, None)), mimetype='text/event-stream')
+    query = request.json['query']
+    history = request.json['history']
+
+    return Response(stream(talk_to_robot(index, query, history)), mimetype='text/event-stream')
 
     # is_valid, response, context = talk_to_robot(index, query, history)
-    
+
     # if is_valid:
     #     return jsonify({'response': response, 'citations': [{'title': block.title, 'author': block.author, 'date': block.date, 'url': block.url} for block in context]})
     # else:
diff --git a/web/src/pages/index.tsx b/web/src/pages/index.tsx
index cd0e695..4fe860c 100644
--- a/web/src/pages/index.tsx
+++ b/web/src/pages/index.tsx
@@ -82,7 +82,7 @@ const ShowEntry: React.FC<{entry: Entry}> = ({entry}) => {
     if (entry.role === "user") {
         return ( <p className="border border-gray-300 px-1 text-right"> {entry.content} </p>);
     }
-    
+
     // error message
     if (entry.role === "error") {
         return ( <p className="border bg-red-100 border-red-500 text-red-800 px-1"> {entry.content} </p>);
@@ -143,31 +143,48 @@ const Home: NextPage = () => {
         setQuery: (query: string) => void,
         setLoading: (loading: boolean) => void
     ) => {
-        
+
         // clear the query box, append to entries
+
         const old_entries = entries;
         const new_entries: Entry[] = [...old_entries, {role: "user", content: query}];
         setEntries(new_entries);
         setQuery("");
-
         setLoading(true);
 
+        // do SSE on a POST request.
 
+        const res = await fetch(API_URL + "/chat", {
+            method: "POST",
+            cache: "no-cache",
+            keepalive: true,
+            headers: {
+                "Content-Type": "application/json",
+                "Accept": "text/event-stream",
+                "Allow-Control-Allow-Origin": "*"
+            },
 
-        // sse connection
-        const eventSource = new EventSource(API_URL + "/chat");
-                            
-        eventSource.onmessage = (event) => {
+            body: JSON.stringify({query: query, history:
+                old_entries.filter((entry) => entry.role !== "error")
+                           .map((entry) => {
+                               return {
+                                   "role" : entry.role,
+                                   "content" : entry.content.trim(),
+                               }
+                           })
+            }),
 
-            if (event.data === "close") {
-                eventSource.close();
-                setLoading(false);
-                setLoadState({state: "idle"});
-                return;
-            }
+        });
 
-            console.log(event.data);
-            const data = JSON.parse(event.data);
+        if (!res.ok) {
+            setLoading(false);
+            console.log("load failure: " + res.status);
+            return;
+        }
+
+        const process = (message: string) => {
+            console.log(message);
+            const data = JSON.parse(message);
 
             if (data.state === "loading") {
                 setLoadState({state: "loading", phase: data.phase});
@@ -183,25 +200,30 @@ const Home: NextPage = () => {
             }
         }
 
-        const res = await fetch(API_URL + "/chat", {
-            method: "POST",
-            headers: { "Content-Type": "application/json", "Allow-Control-Allow-Origin": "*" },
-            body: JSON.stringify({query: query, history: 
-                old_entries.filter((entry) => entry.role !== "error")
-                           .map((entry) => {
-                               return {
-                                   "role" : entry.role,
-                                   "content" : entry.content.trim(),
-                               }
-                           })
-               })
-        })
+        // read back sse stream
+
+        const reader = res.body!.getReader();
+        var message = "";
+        while (true) {
+
+            const {done, value} = await reader.read();
+
+            if (done) break;
+            const chunk = new TextDecoder("utf-8").decode(value);
+            if (chunk.startsWith("event: close\n")) break;
+
+            for (const line of chunk.split('\n')) {
+                if (line.startsWith("data: ")) message += line.slice(6);
+                if (line === "") {
+                    if (message !== "") process(message);
+                    message = "";
+                }
+            }
+        }
+
+        setLoading(false);
+        setLoadState({state: "idle"});
 
-        // if (!res.ok) {
-        //     setLoading(false);
-        //     console.log("load failure: " + res.status);
-        //     return;
-        // }
         //
         // const data = await res.json();
         //
@@ -227,7 +249,7 @@ const Home: NextPage = () => {
         //
         // // transform all things that look like [(a), (b), (c)] into [(a)][(b)][(c)]
         // response = response.replace(
-        //     
+        //
         //         /\[((?:\([a-z]+\),\s*)*\([a-z]+\))\]/g, // identify groups of this form
         //
         //         (block: string) => block.split(',')
@@ -281,7 +303,7 @@ const Home: NextPage = () => {
         //     }
         // });
         //
-        // setEntries([...new_entries, {role: "assistant", 
+        // setEntries([...new_entries, {role: "assistant",
         //                              content: response,
         //                              citations: citations}]);
 
@@ -304,7 +326,20 @@ const Home: NextPage = () => {
                     ))}
                 </ul>
                 <SearchBox search={search} />
-                <p>{loadState.state === "loading" ? loadState.phase : ( loadState.state === "streaming" ? loadState.response : "" )}</p>
+                {(() => {
+                    if (loadState.state === "loading") {
+                        switch (loadState.phase) {
+                            case "semantic": return <p>Loading: Performing semantic search...</p>;
+                            case "prompt": return <p>Loading: Creating prompt...</p>;
+                            case "llm": return <p>Loading: Waiting for LLM...</p>;
+                        }
+
+                    } else if (loadState.state === "streaming") {
+                        return <p>{loadState.response}</p>;
+                    }
+                    return <></>;
+                })()}
+
             </main>
         </>
     );

From 10c3068aa52d4e6113499251f9eb75fbade16e03 Mon Sep 17 00:00:00 2001
From: Fraser <tomatoesneverlie@gmail.com>
Date: Sun, 16 Apr 2023 19:01:12 -0400
Subject: [PATCH 3/6] error passing and logging in SSE

---
 api/chat.py             | 33 ++++++++++++++++++---------------
 web/src/pages/index.tsx | 41 +++++++++++++++++++----------------------
 2 files changed, 37 insertions(+), 37 deletions(-)

diff --git a/api/chat.py b/api/chat.py
index 8360857..c368fcc 100644
--- a/api/chat.py
+++ b/api/chat.py
@@ -121,26 +121,29 @@ import json
 
 # returns either (True, reply string, top_k_blocks)) or (False, error message string, None)
 def talk_to_robot(index, query: str, history: List[Dict[str, str]], k: int = STANDARD_K):
-    # 1. Find the most relevant blocks from the Alignment Research Dataset
-    yield json.dumps({"state": "loading", "phase": "semantic"})
-    top_k_blocks = get_top_k_blocks(index, query, k)
+    try:
+        # 1. Find the most relevant blocks from the Alignment Research Dataset
+        yield json.dumps({"state": "loading", "phase": "semantic"})
+        top_k_blocks = get_top_k_blocks(index, query, k)
 
-    # 2. Generate a prompt
-    prompt = construct_prompt(query, history, top_k_blocks)
-    yield json.dumps({"state": "loading", "phase": "prompt"})
+        # 2. Generate a prompt
+        prompt = construct_prompt(query, history, top_k_blocks)
+        yield json.dumps({"state": "loading", "phase": "prompt"})
 
-    # 3. Count number of tokens left for completion (-50 for a buffer)
-    max_tokens_completion = NUM_TOKENS - sum([len(ENCODER.encode(message["content"]) + ENCODER.encode(message["role"])) for message in prompt]) - 50
+        # 3. Count number of tokens left for completion (-50 for a buffer)
+        max_tokens_completion = NUM_TOKENS - sum([len(ENCODER.encode(message["content"]) + ENCODER.encode(message["role"])) for message in prompt]) - 50
 
+        x = int("non-int")
 
+        yield json.dumps({"state": "loading", "phase": "llm"})
+        time.sleep(1)
+        for c in "Hi. I'm a big dumb LLM.\nHubris will be the end of us all.":
+            yield json.dumps({"state": "streaming", "response": c})
+            time.sleep(0.1)
 
-
-
-    yield json.dumps({"state": "loading", "phase": "llm"})
-    time.sleep(1)
-    for c in "Hi. I'm a big dumb LLM.\nHubris will be the end of us all.":
-        yield json.dumps({"state": "streaming", "response": c})
-        time.sleep(0.1)
+    except Exception as e:
+        print(e)
+        yield json.dumps({"state": "error", "error": str(e)})
 
     # try:
     #
diff --git a/web/src/pages/index.tsx b/web/src/pages/index.tsx
index 4fe860c..0d59d83 100644
--- a/web/src/pages/index.tsx
+++ b/web/src/pages/index.tsx
@@ -182,29 +182,11 @@ const Home: NextPage = () => {
             return;
         }
 
-        const process = (message: string) => {
-            console.log(message);
-            const data = JSON.parse(message);
-
-            if (data.state === "loading") {
-                setLoadState({state: "loading", phase: data.phase});
-                return;
-            }
-
-            if (data.state === "streaming") {
-                setLoadState((s) => {
-                    const response = s.state === "streaming" ? s.response : "";
-                    return {state: "streaming", response: response + data.response};
-                });
-                return;
-            }
-        }
-
         // read back sse stream
 
         const reader = res.body!.getReader();
         var message = "";
-        while (true) {
+        read: while (true) {
 
             const {done, value} = await reader.read();
 
@@ -215,7 +197,24 @@ const Home: NextPage = () => {
             for (const line of chunk.split('\n')) {
                 if (line.startsWith("data: ")) message += line.slice(6);
                 if (line === "") {
-                    if (message !== "") process(message);
+                    if (message !== "") {
+                        const data = JSON.parse(message);
+
+                        switch (data.state) {
+                            case "loading":
+                                setLoadState({state: "loading", phase: data.phase});
+                                break;
+                            case "streaming":
+                                setLoadState((s) => {
+                                    const response = s.state === "streaming" ? s.response : "";
+                                    return {state: "streaming", response: response + data.response};
+                                });
+                                break;
+                            case "error":
+                                setEntries([...new_entries, {role: "error", content: data.error}]);
+                                break read;
+                        }
+                    }
                     message = "";
                 }
             }
@@ -333,13 +332,11 @@ const Home: NextPage = () => {
                             case "prompt": return <p>Loading: Creating prompt...</p>;
                             case "llm": return <p>Loading: Waiting for LLM...</p>;
                         }
-
                     } else if (loadState.state === "streaming") {
                         return <p>{loadState.response}</p>;
                     }
                     return <></>;
                 })()}
-
             </main>
         </>
     );

From 6d91e423639c0f6c0f2370af17e70c92193370ad Mon Sep 17 00:00:00 2001
From: Fraser <tomatoesneverlie@gmail.com>
Date: Sun, 16 Apr 2023 19:37:57 -0400
Subject: [PATCH 4/6] stream answer

---
 api/chat.py             | 71 +++++++++++++++++++----------------------
 api/main.py             |  7 ----
 web/src/pages/index.tsx | 25 ++++++++++++---
 3 files changed, 53 insertions(+), 50 deletions(-)

diff --git a/api/chat.py b/api/chat.py
index c368fcc..25ebc95 100644
--- a/api/chat.py
+++ b/api/chat.py
@@ -127,52 +127,47 @@ def talk_to_robot(index, query: str, history: List[Dict[str, str]], k: int = STA
         top_k_blocks = get_top_k_blocks(index, query, k)
 
         # 2. Generate a prompt
-        prompt = construct_prompt(query, history, top_k_blocks)
         yield json.dumps({"state": "loading", "phase": "prompt"})
+        prompt = construct_prompt(query, history, top_k_blocks)
 
         # 3. Count number of tokens left for completion (-50 for a buffer)
         max_tokens_completion = NUM_TOKENS - sum([len(ENCODER.encode(message["content"]) + ENCODER.encode(message["role"])) for message in prompt]) - 50
 
-        x = int("non-int")
-
+        # 4. Answer the user query
         yield json.dumps({"state": "loading", "phase": "llm"})
-        time.sleep(1)
-        for c in "Hi. I'm a big dumb LLM.\nHubris will be the end of us all.":
-            yield json.dumps({"state": "streaming", "response": c})
-            time.sleep(0.1)
+        t1 = time.time()
+        response = ''
+
+        for chunk in openai.ChatCompletion.create(
+            model=COMPLETIONS_MODEL,
+            messages=prompt,
+            max_tokens=max_tokens_completion,
+            stream=True
+        ):
+            res = chunk["choices"][0]["delta"]
+            if res is not None and res.get("content") is not None:
+                response += res["content"]
+                yield json.dumps({"state": "streaming", "content": res["content"]})
+
+
+        t2 = time.time()
+        print("Time to get response: ", t2 - t1)
+
+        if DEBUG_PRINT:
+            print('\n' * 10)
+            print(" ------------------------------ prompt: -----------------------------")
+            for message in prompt:
+                print(f"----------- {message['role']}: ------------------")
+                print(message['content'])
+
+            print('\n' * 10)
+
+            print(" ------------------------------ response: -----------------------------")
+            print(response)
+
+        yield json.dumps({"state": "done"})
 
     except Exception as e:
         print(e)
         yield json.dumps({"state": "error", "error": str(e)})
 
-    # try:
-    #
-    #     # 4. Answer the user query
-    #     t1 = time.time()
-    #     response = openai.ChatCompletion.create(
-    #         model=COMPLETIONS_MODEL,
-    #         messages=prompt,
-    #         max_tokens=max_tokens_completion
-    #     )["choices"][0]["message"]["content"]
-    #     t2 = time.time()
-    #     print("Time to get response: ", t2 - t1)
-    #     
-    #
-    #     if DEBUG_PRINT:
-    #         print('\n' * 10)
-    #         print(" ------------------------------ prompt: -----------------------------")
-    #         for message in prompt:
-    #             print(f"----------- {message['role']}: ------------------")
-    #             print(message['content'])
-    #
-    #         print('\n' * 10)
-    #
-    #         print(" ------------------------------ response: -----------------------------")
-    #         print(response)
-    #
-    #     return (True, response, top_k_blocks)
-    #
-    # except Exception as e:
-    #     print(e)
-    #     return (False, "Error: " + str(e), None)
-
diff --git a/api/main.py b/api/main.py
index 8389e2b..45f04c9 100644
--- a/api/main.py
+++ b/api/main.py
@@ -62,13 +62,6 @@ def chat():
 
     return Response(stream(talk_to_robot(index, query, history)), mimetype='text/event-stream')
 
-    # is_valid, response, context = talk_to_robot(index, query, history)
-
-    # if is_valid:
-    #     return jsonify({'response': response, 'citations': [{'title': block.title, 'author': block.author, 'date': block.date, 'url': block.url} for block in context]})
-    # else:
-    #     return jsonify({'error': response})
-
 
 # ------------------------------------------------------------------------------
 
diff --git a/web/src/pages/index.tsx b/web/src/pages/index.tsx
index 0d59d83..5ae72a7 100644
--- a/web/src/pages/index.tsx
+++ b/web/src/pages/index.tsx
@@ -128,7 +128,7 @@ type State = {
     phase: "semantic" | "prompt" | "llm";
 } | {
     state: "streaming";
-    response: string;
+    response: AssistantEntry;
 };
 
 
@@ -201,18 +201,33 @@ const Home: NextPage = () => {
                         const data = JSON.parse(message);
 
                         switch (data.state) {
+
                             case "loading":
                                 setLoadState({state: "loading", phase: data.phase});
                                 break;
+
                             case "streaming":
                                 setLoadState((s) => {
-                                    const response = s.state === "streaming" ? s.response : "";
-                                    return {state: "streaming", response: response + data.response};
+                                    const response = s.state === "streaming" ? s.response : {role: "assistant", content: "", citations: new Map()};
+                                    return {state: "streaming", response: {
+                                        role: "assistant",
+                                        content: response.content + data.content,
+                                        citations: response.citations,
+                                    }};
                                 });
                                 break;
+
+                            case "done":
+                                setLoadState((s) => {
+                                    if (s.state === "streaming") setEntries([...new_entries, s.response]);
+                                    return {state: "idle"};
+                                });
+                                break read;
+
                             case "error":
                                 setEntries([...new_entries, {role: "error", content: data.error}]);
                                 break read;
+
                         }
                     }
                     message = "";
@@ -323,7 +338,6 @@ const Home: NextPage = () => {
                             <ShowEntry entry={entry} />
                         </li>
                     ))}
-                </ul>
                 <SearchBox search={search} />
                 {(() => {
                     if (loadState.state === "loading") {
@@ -333,10 +347,11 @@ const Home: NextPage = () => {
                             case "llm": return <p>Loading: Waiting for LLM...</p>;
                         }
                     } else if (loadState.state === "streaming") {
-                        return <p>{loadState.response}</p>;
+                        return <ShowEntry entry={loadState.response} />;
                     }
                     return <></>;
                 })()}
+                </ul>
             </main>
         </>
     );

From a1fd3a852d3e80461c0ee7a114fc2e1b8e7a3d3d Mon Sep 17 00:00:00 2001
From: Fraser <tomatoesneverlie@gmail.com>
Date: Sun, 16 Apr 2023 19:52:43 -0400
Subject: [PATCH 5/6] move parsing towards point of use.

---
 web/src/pages/index.tsx | 161 +++++++++++++++++++++-------------------
 1 file changed, 83 insertions(+), 78 deletions(-)

diff --git a/web/src/pages/index.tsx b/web/src/pages/index.tsx
index 5ae72a7..9870807 100644
--- a/web/src/pages/index.tsx
+++ b/web/src/pages/index.tsx
@@ -3,7 +3,7 @@ const API_URL = process.env.NEXT_PUBLIC_API_URL || "http://127.0.0.1:3000";
 import Head from "next/head";
 import React from "react";
 import { type NextPage } from "next";
-import { useState } from "react";
+import { useState, useMemo } from "react";
 
 import Header from "../header";
 import SearchBox from "../searchbox";
@@ -23,7 +23,7 @@ type UserEntry = {
 type AssistantEntry = {
     role: "assistant";
     content: string;
-    citations: Map<number, Citation>;
+    citations: Citation[];
 }
 
 type ErrorMessage = {
@@ -88,9 +88,84 @@ const ShowEntry: React.FC<{entry: Entry}> = ({entry}) => {
         return ( <p className="border bg-red-100 border-red-500 text-red-800 px-1"> {entry.content} </p>);
     }
 
+    // robot message
+    const res = useMemo(() => ShowAssistantEntry(entry), [entry]);
+    return res;
+};
+
+const ShowAssistantEntry = (entry: AssistantEntry) => {
     const in_text_citation_regex = /\[([0-9]+)\]/g;
 
-    // system reply
+    // ---------------------- normalize citation form ----------------------
+   
+    // transform all things that look like [a, b, c] into [a][b][c]
+    let response = entry.content.replace(
+   
+                /\[((?:[a-z]+,\s*)*[a-z]+)\]/g, // identify groups of this form
+   
+                (block: string) => block.split(',')
+                                        .map((x) => x.trim())
+                                        .join("][")
+    )
+   
+    // transform all things that look like [(a), (b), (c)] into [(a)][(b)][(c)]
+    response = response.replace(
+   
+            /\[((?:\([a-z]+\),\s*)*\([a-z]+\))\]/g, // identify groups of this form
+   
+            (block: string) => block.split(',')
+                                    .map((x) => x.trim())
+                                    .join("][")
+    )
+   
+    // transform all things that look like [(a)] into [a]
+    response = response.replace(
+        /\[\(([a-z]+)\)\]/g,
+        (_match: string, x: string) => `[${x}]`
+    )
+   
+    // transform all things that look like [ a ] into [a]
+    response = response.replace(
+        /\[\s*([a-z]+)\s*\]/g,
+        (_match: string, x: string) => `[${x}]`
+    )
+   
+    // -------------- map citations from strings into numbers --------------
+   
+    // figure out what citations are in the response, and map them appropriately
+    const cite_map = new Map<string, number>();
+    // let cite_count = runningIndex;
+    let cite_count = 0;
+   
+    // scan a regex for [x] over the response. If x isn't in the map, add it.
+    const regex = /\[([a-z]+)\]/g;
+    let match;
+    let response_copy = ""
+    while ((match = regex.exec(response)) !== null) {
+        if (!cite_map.has(match[1]!)) {
+            cite_map.set(match[1]!, cite_count++);
+        }
+        // replace [x] with [i]
+        response_copy += response.slice(response_copy.length, match.index) + `[${cite_map.get(match[1]!)! + 1}]`;
+    }
+   
+    // setRunningIndex(cite_count);
+    // TODO
+   
+    response = response_copy + response.slice(response_copy.length);
+   
+    // ----------------- create the ordered citation array -----------------
+   
+    const citations = new Map<number, Citation>();
+    cite_map.forEach((value, key) => {
+        const index = key.charCodeAt(0) - 'a'.charCodeAt(0);
+        if (index >= entry.citations.length) {
+            console.log("invalid citation index: " + index);
+        } else {
+            citations.set(value, entry.citations[index]!);
+        }
+    });
+   
     return (
         <div className="mt-3 mb-8">
             {   // split into paragraphs
@@ -100,8 +175,8 @@ const ShowEntry: React.FC<{entry: Entry}> = ({entry}) => {
                             return text.trim();
                         }
                         i = parseInt(text) - 1;
-                        if (!entry.citations.has(i)) return `[${text}]`;
-                        const citation = entry.citations.get(i)!;
+                        if (!citations.has(i)) return `[${text}]`;
+                        const citation = citations.get(i)!;
                         return (
                             <ShowInTextCitation citation={citation} i={i} />
                         );
@@ -110,7 +185,7 @@ const ShowEntry: React.FC<{entry: Entry}> = ({entry}) => {
             }
             <ul>
                 {   // show citations
-                    Array.from(entry.citations.entries()).map(([i, citation]) => (
+                    Array.from(citations.entries()).map(([i, citation]) => (
                         <li key={i}>
                             <ShowCitation citation={citation} i={i} />
                         </li>
@@ -120,6 +195,7 @@ const ShowEntry: React.FC<{entry: Entry}> = ({entry}) => {
         </div>
     );
 };
+    
 
 type State = {
     state: "idle";
@@ -208,7 +284,7 @@ const Home: NextPage = () => {
 
                             case "streaming":
                                 setLoadState((s) => {
-                                    const response = s.state === "streaming" ? s.response : {role: "assistant", content: "", citations: new Map()};
+                                    const response = s.state === "streaming" ? s.response : {role: "assistant", content: "", citations: []};
                                     return {state: "streaming", response: {
                                         role: "assistant",
                                         content: response.content + data.content,
@@ -249,77 +325,6 @@ const Home: NextPage = () => {
         //     return;
         // }
         //
-        // // ---------------------- normalize citation form ----------------------
-        //
-        // // transform all things that look like [a, b, c] into [a][b][c]
-        // let response = data.response.replace(
-        //
-        //             /\[((?:[a-z]+,\s*)*[a-z]+)\]/g, // identify groups of this form
-        //
-        //             (block: string) => block.split(',')
-        //                                     .map((x) => x.trim())
-        //                                     .join("][")
-        // )
-        //
-        // // transform all things that look like [(a), (b), (c)] into [(a)][(b)][(c)]
-        // response = response.replace(
-        //
-        //         /\[((?:\([a-z]+\),\s*)*\([a-z]+\))\]/g, // identify groups of this form
-        //
-        //         (block: string) => block.split(',')
-        //                                 .map((x) => x.trim())
-        //                                 .join("][")
-        // )
-        //
-        // // transform all things that look like [(a)] into [a]
-        // response = response.replace(
-        //     /\[\(([a-z]+)\)\]/g,
-        //     (_match: string, x: string) => `[${x}]`
-        // )
-        //
-        // // transform all things that look like [ a ] into [a]
-        // response = response.replace(
-        //     /\[\s*([a-z]+)\s*\]/g,
-        //     (_match: string, x: string) => `[${x}]`
-        // )
-        //
-        // // -------------- map citations from strings into numbers --------------
-        //
-        // // figure out what citations are in the response, and map them appropriately
-        // const cite_map = new Map<string, number>();
-        // let cite_count = runningIndex;
-        //
-        // // scan a regex for [x] over the response. If x isn't in the map, add it.
-        // const regex = /\[([a-z]+)\]/g;
-        // let match;
-        // let response_copy = ""
-        // while ((match = regex.exec(response)) !== null) {
-        //     if (!cite_map.has(match[1]!)) {
-        //         cite_map.set(match[1]!, cite_count++);
-        //     }
-        //     // replace [x] with [i]
-        //     response_copy += response.slice(response_copy.length, match.index) + `[${cite_map.get(match[1]!)! + 1}]`;
-        // }
-        //
-        // setRunningIndex(cite_count);
-        //
-        // response = response_copy + response.slice(response_copy.length);
-        //
-        // // ----------------- create the ordered citation array -----------------
-        //
-        // const citations = new Map<number, Citation>();
-        // cite_map.forEach((value, key) => {
-        //     const index = key.charCodeAt(0) - 'a'.charCodeAt(0);
-        //     if (index >= data.citations.length) {
-        //         console.log("invalid citation index: " + index);
-        //     } else {
-        //         citations.set(value, data.citations[index]);
-        //     }
-        // });
-        //
-        // setEntries([...new_entries, {role: "assistant",
-        //                              content: response,
-        //                              citations: citations}]);
 
         setLoading(false);
 

From dee78258d6a8972aba9f29237f06e22bf3085a78 Mon Sep 17 00:00:00 2001
From: Fraser <tomatoesneverlie@gmail.com>
Date: Sun, 16 Apr 2023 21:12:25 -0400
Subject: [PATCH 6/6] fully working streaming

---
 api/chat.py             |   2 +
 web/src/pages/index.tsx | 154 +++++++++++++++++++++++++---------------
 2 files changed, 99 insertions(+), 57 deletions(-)

diff --git a/api/chat.py b/api/chat.py
index 25ebc95..10fbf1d 100644
--- a/api/chat.py
+++ b/api/chat.py
@@ -126,6 +126,8 @@ def talk_to_robot(index, query: str, history: List[Dict[str, str]], k: int = STA
         yield json.dumps({"state": "loading", "phase": "semantic"})
         top_k_blocks = get_top_k_blocks(index, query, k)
 
+        yield json.dumps({"state": "loading", "phase": "semantic", 'citations': [{'title': block.title, 'author': block.author, 'date': block.date, 'url': block.url} for block in top_k_blocks]})
+
         # 2. Generate a prompt
         yield json.dumps({"state": "loading", "phase": "prompt"})
         prompt = construct_prompt(query, history, top_k_blocks)
diff --git a/web/src/pages/index.tsx b/web/src/pages/index.tsx
index 9870807..a770c5d 100644
--- a/web/src/pages/index.tsx
+++ b/web/src/pages/index.tsx
@@ -3,7 +3,7 @@ const API_URL = process.env.NEXT_PUBLIC_API_URL || "http://127.0.0.1:3000";
 import Head from "next/head";
 import React from "react";
 import { type NextPage } from "next";
-import { useState, useMemo } from "react";
+import { useState } from "react";
 
 import Header from "../header";
 import SearchBox from "../searchbox";
@@ -24,6 +24,7 @@ type AssistantEntry = {
     role: "assistant";
     content: string;
     citations: Citation[];
+    base_count: number; // the number to start counting citations at
 }
 
 type ErrorMessage = {
@@ -83,23 +84,17 @@ const ShowEntry: React.FC<{entry: Entry}> = ({entry}) => {
         return ( <p className="border border-gray-300 px-1 text-right"> {entry.content} </p>);
     }
 
-    // error message
-    if (entry.role === "error") {
-        return ( <p className="border bg-red-100 border-red-500 text-red-800 px-1"> {entry.content} </p>);
-    }
+        
+    
 
-    // robot message
-    const res = useMemo(() => ShowAssistantEntry(entry), [entry]);
-    return res;
-};
 
-const ShowAssistantEntry = (entry: AssistantEntry) => {
-    const in_text_citation_regex = /\[([0-9]+)\]/g;
+// todo: memoize this if too slow.
+const ProcessText: (text: string, base_count: number) => [string, Map<string, number>] = (text, base_count) => {
 
     // ---------------------- normalize citation form ----------------------
    
     // transform all things that look like [a, b, c] into [a][b][c]
-    let response = entry.content.replace(
+    let response = text.replace(
    
                 /\[((?:[a-z]+,\s*)*[a-z]+)\]/g, // identify groups of this form
    
@@ -134,7 +129,6 @@ const ShowAssistantEntry = (entry: AssistantEntry) => {
    
     // figure out what citations are in the response, and map them appropriately
     const cite_map = new Map<string, number>();
-    // let cite_count = runningIndex;
     let cite_count = 0;
    
     // scan a regex for [x] over the response. If x isn't in the map, add it.
@@ -143,16 +137,22 @@ const ShowAssistantEntry = (entry: AssistantEntry) => {
     let response_copy = ""
     while ((match = regex.exec(response)) !== null) {
         if (!cite_map.has(match[1]!)) {
-            cite_map.set(match[1]!, cite_count++);
+            cite_map.set(match[1]!, base_count + cite_count++);
         }
         // replace [x] with [i]
         response_copy += response.slice(response_copy.length, match.index) + `[${cite_map.get(match[1]!)! + 1}]`;
     }
    
-    // setRunningIndex(cite_count);
-    // TODO
-   
     response = response_copy + response.slice(response_copy.length);
+
+    return [response, cite_map]
+}
+
+
+const ShowAssistantEntry: React.FC<{entry: AssistantEntry}> = ({entry}) => {
+    const in_text_citation_regex = /\[([0-9]+)\]/g;
+
+    let [response, cite_map] = ProcessText(entry.content, entry.base_count);
    
     // ----------------- create the ordered citation array -----------------
    
@@ -165,11 +165,11 @@ const ShowAssistantEntry = (entry: AssistantEntry) => {
             citations.set(value, entry.citations[index]!);
         }
     });
-   
+
     return (
         <div className="mt-3 mb-8">
             {   // split into paragraphs
-                entry.content.split("\n").map(paragraph => ( <p> {
+                response.split("\n").map(paragraph => ( <p> {
                     paragraph.split(in_text_citation_regex).map((text, i) => {
                         if (i % 2 === 0) {
                             return text.trim();
@@ -183,7 +183,7 @@ const ShowAssistantEntry = (entry: AssistantEntry) => {
                     })
                 } </p>))
             }
-            <ul>
+            <ul className="mt-5">
                 {   // show citations
                     Array.from(citations.entries()).map(([i, citation]) => (
                         <li key={i}>
@@ -195,13 +195,17 @@ const ShowAssistantEntry = (entry: AssistantEntry) => {
         </div>
     );
 };
-    
+
+
+
+
 
 type State = {
     state: "idle";
 } | {
     state: "loading";
     phase: "semantic" | "prompt" | "llm";
+    citations: Citation[];
 } | {
     state: "streaming";
     response: AssistantEntry;
@@ -254,11 +258,12 @@ const Home: NextPage = () => {
 
         if (!res.ok) {
             setLoading(false);
-            console.log("load failure: " + res.status);
+            setLoadState({state: "idle"});
+            setEntries([...new_entries, {role: "error", content: "POST Error: " + res.status}]);
             return;
         }
 
-        // read back sse stream
+        // read back the SSE stream
 
         const reader = res.body!.getReader();
         var message = "";
@@ -270,7 +275,13 @@ const Home: NextPage = () => {
             const chunk = new TextDecoder("utf-8").decode(value);
             if (chunk.startsWith("event: close\n")) break;
 
+            // note: this form isn't even remotely close to optimal in terms of network usage.
+
             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: ")) message += line.slice(6);
                 if (line === "") {
                     if (message !== "") {
@@ -279,23 +290,50 @@ const Home: NextPage = () => {
                         switch (data.state) {
 
                             case "loading":
-                                setLoadState({state: "loading", phase: data.phase});
+
+                                // display loading phases, once citations are available toss them
+                                // into the loading state.
+
+                                setLoadState((s) => {
+                                    var citations = s.state === "loading" ? s.citations : [];
+                                    if (data.citations !== undefined) {
+                                        citations = data.citations;
+                                    }
+                                    return {state: "loading", phase: data.phase, citations: citations};
+                                });
+
                                 break;
 
                             case "streaming":
+
+                                // incrementally build up the response
+
                                 setLoadState((s) => {
-                                    const response = s.state === "streaming" ? s.response : {role: "assistant", content: "", citations: []};
+                                    const response = s.state === "streaming" ? s.response : 
+                                                {role: "assistant", 
+                                                 content: "", 
+                                                 citations: s.state === "loading" ? s.citations : [], 
+                                                 base_count: runningIndex
+                                                };
+
                                     return {state: "streaming", response: {
                                         role: "assistant",
                                         content: response.content + data.content,
                                         citations: response.citations,
+                                        base_count: response.base_count
                                     }};
                                 });
                                 break;
 
                             case "done":
+
+                                // append the response to the entries, reset to normal
+
                                 setLoadState((s) => {
-                                    if (s.state === "streaming") setEntries([...new_entries, s.response]);
+                                    if (s.state === "streaming") {
+                                        setEntries([...new_entries, s.response]);
+                                        setRunningIndex((i) => (i + ProcessText(s.response.content, 0)[1].size));
+                                    }
                                     return {state: "idle"};
                                 });
                                 break read;
@@ -314,20 +352,6 @@ const Home: NextPage = () => {
         setLoading(false);
         setLoadState({state: "idle"});
 
-        //
-        // const data = await res.json();
-        //
-        // // -------------------------- error checking ---------------------------
-        //
-        // if (data.error) {
-        //     setEntries([...new_entries, {role: "error", content: data.error}]);
-        //     setLoading(false);
-        //     return;
-        // }
-        //
-
-        setLoading(false);
-
     };
 
     return (
@@ -338,24 +362,40 @@ const Home: NextPage = () => {
             <main>
                 <Header page="index" />
                 <ul>
-                    {entries.map((entry, i) => (
-                        <li key={i}>
-                            <ShowEntry entry={entry} />
-                        </li>
-                    ))}
-                <SearchBox search={search} />
-                {(() => {
-                    if (loadState.state === "loading") {
-                        switch (loadState.phase) {
-                            case "semantic": return <p>Loading: Performing semantic search...</p>;
-                            case "prompt": return <p>Loading: Creating prompt...</p>;
-                            case "llm": return <p>Loading: Waiting for LLM...</p>;
+                    {entries.map((entry, i) => {
+                        if (entry.role === "user") {
+                            return <li key={i}>
+                                <p className="border border-gray-300 px-1 text-right"> {entry.content} </p>
+                            </li>
                         }
-                    } else if (loadState.state === "streaming") {
-                        return <ShowEntry entry={loadState.response} />;
-                    }
-                    return <></>;
-                })()}
+                        if (entry.role === "error") {
+                            return <li key={i}>
+                                <p className="border bg-red-100 border-red-500 text-red-800 px-1"> {entry.content} </p>
+                            </li>
+                        }
+                        if (entry.role === "assistant") {
+                            return <li key={i}>
+                                <ShowAssistantEntry entry={entry}/>
+                            </li>
+                        }
+                        return <></>
+                    })}
+
+                    <SearchBox search={search} />
+
+                    {(() => {
+                        if (loadState.state === "loading") {
+                            switch (loadState.phase) {
+                                case "semantic": return <p>Loading: Performing semantic search...</p>;
+                                case "prompt": return <p>Loading: Creating prompt...</p>;
+                                case "llm": return <p>Loading: Waiting for LLM...</p>;
+                            }
+                        } else if (loadState.state === "streaming") {
+                            return <ShowAssistantEntry entry={loadState.response}/>;
+                        }
+                        return <></>;
+                    })()}
+
                 </ul>
             </main>
         </>