From 52585d7991a7f7fe8c4f23d7ff437dbd1bbb72b3 Mon Sep 17 00:00:00 2001
From: Fraser <tomatoesneverlie@gmail.com>
Date: Fri, 31 Mar 2023 18:32:52 -0400
Subject: [PATCH] add error passing mechanism

---
 api/chat.py             |  8 ++++++--
 api/main.py             |  9 +++++++--
 web/src/pages/index.tsx | 35 ++++++++++++++++++++++++++++-------
 3 files changed, 41 insertions(+), 11 deletions(-)

diff --git a/api/chat.py b/api/chat.py
index f2e685f..2254223 100644
--- a/api/chat.py
+++ b/api/chat.py
@@ -73,7 +73,6 @@ def construct_prompt(query: str, history: List[Dict[str, str]], context: List[Bl
 
 # ------------------------------------------------------------------------------
 
-        
 def normal_completion(prompt: List[Dict[str, str]]) -> str:
     try:
         return openai.ChatCompletion.create(
@@ -84,6 +83,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."
 
+# returns either (True, reply string, embeddings) or (False, error message string, None)
 def talk_to_robot(dataset_dict, query: str, history: List[Dict[str, str]] = [], k: int = 10):
 
     # 1. Find the most relevant blocks from the Alignment Research Dataset
@@ -91,6 +91,10 @@ def talk_to_robot(dataset_dict, query: str, history: List[Dict[str, str]] = [],
     
     # 2. Generate a prompt for the ChatCompletions API
     prompt: List[Dict[str, str]] = construct_prompt(query, history, top_k_blocks)
+
+    # if we were to error out, return something like this
+    # return (False, "Example error message", None)
     
     # 3. Answer the user query
-    return (normal_completion(prompt), top_k_blocks)
+    return (True, normal_completion(prompt), top_k_blocks)
+    
diff --git a/api/main.py b/api/main.py
index e076a32..00799a1 100644
--- a/api/main.py
+++ b/api/main.py
@@ -55,8 +55,13 @@ def semantic():
 @cross_origin()
 def chat():
     query = request.json['query']
-    response, context = talk_to_robot(dataset_dict, query)
-    return jsonify({'response': response, 'citations': [{'title': block.title, 'author': block.author, 'date': block.date, 'url': block.url} for block in context]})
+
+    is_valid, response, context = talk_to_robot(dataset_dict, query)
+
+    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 fe718e3..5a76ec2 100644
--- a/web/src/pages/index.tsx
+++ b/web/src/pages/index.tsx
@@ -27,7 +27,12 @@ type AssistantEntry = {
     citations: Map<number, Citation>;
 }
 
-type Entry = UserEntry | AssistantEntry;
+type ErrorMessage = {
+    role: "error";
+    content: string;
+}
+
+type Entry = UserEntry | AssistantEntry | ErrorMessage;
 
 // 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
@@ -72,6 +77,11 @@ 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>);
+    }
 
     const in_text_citation_regex = /\[([0-9]+)\]/g;
 
@@ -128,12 +138,15 @@ 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.map((entry) => {
-                return {
-                    "role" : entry.role,
-                    "content" : entry.content
-                }
-            })})
+            body: JSON.stringify({query: query, history: 
+                old_entries.filter((entry) => entry.role !== "error")
+                           .map((entry) => {
+                               return {
+                                   "role" : entry.role,
+                                   "content" : entry.content
+                               }
+                           })
+               })
         })
 
         if (!res.ok) {
@@ -144,6 +157,14 @@ const Home: NextPage = () => {
 
         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]