list citations in obvious way

api/chat.py

@@ -84,7 +84,7 @@ def normal_completion(prompt: List[Dict[str, str]]) -> str:
         print(e)
         return "I'm sorry, I failed to process your query. Please try again. If the problem persists, please contact the administrator."
 
-def talk_to_robot(query: str, history: List[Dict[str, str]] = [], k: int = 10) -> str:
+def talk_to_robot(query: str, history: List[Dict[str, str]] = [], k: int = 10):
 
     # 1. Find the most relevant blocks from the Alignment Research Dataset
     top_k_blocks: List[Block] = get_top_k_blocks(query, k)
@@ -93,4 +93,4 @@ def talk_to_robot(query: str, history: List[Dict[str, str]] = [], k: int = 10) -
     prompt: List[Dict[str, str]] = construct_prompt(query, history, top_k_blocks)
     
     # 3. Answer the user query
-    return normal_completion(prompt)
+    return (normal_completion(prompt), top_k_blocks)

api/main.py

@@ -35,10 +35,8 @@ def semantic():
 @cross_origin()
 def chat():
     query = request.json['query']
-    return talk_to_robot(query)
-
-
-
+    response, context = talk_to_robot(query)
+    return jsonify({'response': response, 'citations': [{'title': block.title, 'author': block.author, 'date': block.date, 'url': block.url} for block in context]})
 
 # ------------------------------------------------------------------------------
 

web/src/pages/index.tsx

@@ -8,21 +8,76 @@ import { useState } from "react";
 import Header from "../header";
 import SearchBox from "../searchbox";
 
-type Entry = {
-    role: "user" | "assistant";
+
+type UserEntry = {
+    role: "user";
     content: string;
+}
+
+type AssistantEntry = {
+    role: "assistant";
+    content: string;
+    citations: Citation[];
+}
+
+type Entry = UserEntry | AssistantEntry;
+
+type Citation = {
+    title: string;
+    author: string;
+    date: string;
+    url: string;
+}
+
+// const Colours = ["blue", "cyan", "teal", "green", "amber"].map(colour => `bg-${colour}-100 border-${colour}-300 text-${colour}-800`);
+// this would be nice, but Tailwind needs te actual string of the class to be in
+// the source file for it to be included in the build
+
+const Colours = [
+    // "bg-teal-100   border-teal-300   text-teal-800",
+    "bg-red-100    border-red-300    text-red-800",
+    "bg-amber-100  border-amber-300  text-amber-800",
+    "bg-orange-100 border-orange-300 text-orange-800",
+    "bg-lime-100   border-lime-300   text-lime-800",
+    "bg-green-100  border-green-300  text-green-800",
+    "bg-cyan-100  border-cyan-300  text-cyan-800",
+    "bg-blue-100  border-blue-300  text-blue-800",
+];
+
+const ShowCitation: React.FC<{citation: Citation, i: number}> = ({citation, i}) => {
+    return (
+        <a className={Colours[i % Colours.length] + " border-2 flex items-center rounded my-2 text-sm no-underline w-fit"}
+            href={citation.url} target="_blank" rel="noreferrer">
+            <span className="mx-1"> [{i + 1}] </span>
+            <p className="mx-1 my-0">
+                {citation.title + " - " + citation.author + " - " + citation.date}
+            </p>
+        </a>
+    );
 };
 
 const ShowEntry: React.FC<{entry: Entry}> = ({entry}) => {
+
+    // user message
     if (entry.role === "user") {
         return ( <p className="border border-gray-300 px-1"> {entry.content} </p>);
     }
 
+    // system reply
     return (
         <div className="my-3">
             {   // split into paragraphs
-                entry.content.split("\n").map((paragraph, i) => ( <p key={i}> {paragraph} </p>))
+                entry.content.split("\n").map(paragraph => ( <p> {paragraph} </p>))
             }
+            <ul>
+                {   // show citations
+                    entry.citations.map((citation, i) => (
+                        <li key={i}>
+                            <ShowCitation citation={citation} i={i} />
+                        </li>
+                    ))
+                }
+            </ul>
         </div>
     );
 };
@@ -58,11 +113,9 @@ const Home: NextPage = () => {
             return;
         }
 
-        const response = res.body!.getReader().read().then(({value}) => {
-            return new TextDecoder("utf-8").decode(value);
-        });
+        const data = await res.json();
 
-        setEntries([...new_entries, {role: "assistant", content: await response}]);
+        setEntries([...new_entries, {role: "assistant", content: await data.response, citations: await data.citations}]);
 
         setLoading(false);