From 820d70a3abf407efb6951e33e4fce1ea9a09e272 Mon Sep 17 00:00:00 2001
From: Fraser <tomatoesneverlie@gmail.com>
Date: Thu, 1 Jun 2023 23:31:07 -0400
Subject: [PATCH] very basic followups check in backend

---
 api/chat.py      |  3 +++
 api/followups.py | 45 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 48 insertions(+)
 create mode 100644 api/followups.py

diff --git a/api/chat.py b/api/chat.py
index d08d2a8..8af6567 100644
--- a/api/chat.py
+++ b/api/chat.py
@@ -1,4 +1,5 @@
 # ------------------------------- env, constants -------------------------------
+from followups import search_authored
 
 from get_blocks import get_top_k_blocks, Block
 
@@ -170,6 +171,8 @@ def talk_to_robot(index, query: str, history: List[Dict[str, str]], k: int = STA
         log(query)
         log(response)
 
+        search_authored(response, DEBUG_PRINT)
+
         yield json.dumps({"state": "done"})
 
     except Exception as e:
diff --git a/api/followups.py b/api/followups.py
new file mode 100644
index 0000000..87e4cef
--- /dev/null
+++ b/api/followups.py
@@ -0,0 +1,45 @@
+from dataclasses import dataclass
+from urllib.parse import quote
+import requests
+
+SIMILARITY_THRESHOLD = 0.5 # total shot in the dark - play with this later
+MAX_FOLLOWUPS = 3
+
+@dataclass
+class Followup:
+    text: str
+    pageid: str
+    score: float
+
+# do a search like this:
+# https://nlp.stampy.ai/api/search?query=what%20is%20agi
+
+def search_authored(query: str, DEBUG_PRINT: bool = False):
+    url = 'https://nlp.stampy.ai/api/search?query=' + quote(query)
+    response = requests.get(url).json()
+    followups = [ Followup(entry['title'], entry['pageid'], entry['score']) for entry in response ]
+
+    # (note: api presently returns followups pre-sorted, but idk if that's 
+    # guaranteed to stay the case. Re-sorting should be cheap anyway).
+
+    followups.sort(key=lambda f: f.score, reverse=True)
+
+    followups = followups[:MAX_FOLLOWUPS]
+
+    if DEBUG_PRINT:
+        print(" ------------------------------ suggested followups: -----------------------------")
+        for followup in followups:
+            if followup.score > SIMILARITY_THRESHOLD:
+                print(f'{followup.score:.2f} - suggested to user')
+            else:
+                print(f'{followup.score:.2f} - not suggested')
+
+            print(followup.text)
+            print(followup.pageid)
+            print()
+
+    followups = [ f for f in followups if f.score > SIMILARITY_THRESHOLD ]
+
+    return followups
+
+