From 149dd13d1fad3d62f10093040c5f15cffff1abbd Mon Sep 17 00:00:00 2001
From: Daniel O'Connell <github@ahiru.pl>
Date: Fri, 29 Sep 2023 00:42:42 +0200
Subject: [PATCH 1/3] Clean up web search

---
 web/src/hooks/useSearch.ts | 220 ++++++++++++++++++++++++++++++++
 web/src/pages/index.tsx    | 254 +++++++------------------------------
 web/src/settings.ts        |   1 +
 3 files changed, 269 insertions(+), 206 deletions(-)
 create mode 100644 web/src/hooks/useSearch.ts

diff --git a/web/src/hooks/useSearch.ts b/web/src/hooks/useSearch.ts
new file mode 100644
index 0000000..f08dff3
--- /dev/null
+++ b/web/src/hooks/useSearch.ts
@@ -0,0 +1,220 @@
+import { API_URL, STAMPY_URL, STAMPY_CONTENT_URL } from "../settings";
+import type {
+  Citation,
+  Entry,
+  AssistantEntry,
+  ErrorMessage,
+  StampyMessage,
+  CurrentSearch,
+  Followup,
+  SearchResult,
+} from "../types";
+
+const MAX_FOLLOWUPS = 4;
+
+type HistoryEntry = {
+  role: "error" | "stampy" | "assistant" | "user";
+  content: string;
+};
+
+export async function* iterateData(res: Response) {
+  const reader = res.body!.getReader();
+  var message = "";
+
+  while (true) {
+    const { done, value } = await reader.read();
+
+    if (done) return;
+
+    const chunk = new TextDecoder("utf-8").decode(value);
+    if (chunk.startsWith("event: close\n")) 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: ")) {
+        message += line.slice(6);
+        // Fixes #43
+      } else if (line !== "") {
+        message += line;
+      } else if (message !== "") {
+        yield JSON.parse(message);
+        message = "";
+      }
+    }
+  }
+}
+
+export const extractAnswer = async (
+  res: Response,
+  baseReferencesIndex: number,
+  setCurrent: (e: CurrentSearch) => void
+): Promise<SearchResult> => {
+  var result: AssistantEntry = {
+    role: "assistant",
+    content: "",
+    citations: [],
+    base_count: baseReferencesIndex,
+  };
+  var followups: Followup[] = [];
+  for await (var data of iterateData(res)) {
+    switch (data.state) {
+      case "loading":
+        // display loading phases, once citations are available toss them
+        // into the current item.
+        result = {
+          ...result,
+          citations: data?.citations || result?.citations || [],
+        };
+        setCurrent({ phase: data.phase, ...result });
+        break;
+
+      case "streaming":
+        // incrementally build up the response
+        result = {
+          role: "assistant",
+          content: (result?.content || "") + data.content,
+          citations: result?.citations || [],
+          base_count: result?.base_count || baseReferencesIndex,
+        };
+        setCurrent({ phase: "streaming", ...result });
+        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 };
+      case "error":
+        throw data.error;
+    }
+  }
+  return { result, followups };
+};
+
+const fetchLLM = async (
+  query: string,
+  mode: string,
+  history: HistoryEntry[]
+): Promise<Response> =>
+  fetch(API_URL + "/chat", {
+    method: "POST",
+    cache: "no-cache",
+    keepalive: true,
+    headers: {
+      "Content-Type": "application/json",
+      Accept: "text/event-stream",
+      "Allow-Control-Allow-Origin": "*",
+    },
+
+    body: JSON.stringify({ query, mode, history }),
+  });
+
+export const queryLLM = async (
+  query: string,
+  mode: string,
+  history: HistoryEntry[],
+  baseReferencesIndex: number,
+  setCurrent: (e?: CurrentSearch) => void
+): Promise<SearchResult> => {
+  // do SSE on a POST request.
+  const res = await fetchLLM(query, mode, history);
+
+  if (!res.ok) {
+    return { result: { role: "error", content: "POST Error: " + res.status } };
+  }
+
+  try {
+    const results = await extractAnswer(res, baseReferencesIndex, setCurrent);
+    setCurrent(undefined);
+    return results;
+  } catch (e) {
+    return {
+      result: { role: "error", content: e ? e.toString() : "unknown error" },
+    };
+  }
+};
+
+export const getStampyContent = async (
+  question_id: string
+): Promise<SearchResult> => {
+  const res = await fetch(`${STAMPY_CONTENT_URL}/${question_id}`, {
+    method: "GET",
+    headers: {
+      "Content-Type": "application/json",
+      Accept: "application/json",
+      "Allow-Control-Allow-Origin": "*",
+    },
+  });
+
+  if (!res.ok) {
+    return { result: { role: "error", content: "POST Error: " + res.status } };
+  }
+
+  const data = (await res.json()).data;
+
+  let result = {
+    role: "stampy",
+    content: data.text,
+    url: `${STAMPY_URL}/?state=${data.pageid}`,
+  } as StampyMessage;
+
+  // re-enable the searchbox, with the question that was just answered
+  // removed from the list of possible followups.
+
+  // create an array of new followup questions from the data
+  const f_new = data.relatedQuestions.map((f: any) => ({
+    pageid: f.pageid!,
+    text: f.title!,
+    score: 0,
+  }));
+
+  const fpids = new Set(f_new.map((f: Followup) => f.pageid));
+  const followups = (f_old: Followup[]): Followup[] => {
+    const f_old_filtered = f_old.filter(
+      (f) => f.pageid !== data.pageid && !fpids.has(f.pageid)
+    );
+    return [...f_new, ...f_old_filtered].slice(0, MAX_FOLLOWUPS); // this is correct, it's N and not N-1 in javascript fsr
+  };
+
+  return { followups, result };
+};
+
+export const runSearch = async (
+  query: string,
+  query_source: "search" | "followups",
+  mode: string,
+  baseReferencesIndex: number,
+  entries: Entry[],
+  setCurrent: (c: CurrentSearch) => void
+) => {
+  if (query_source === "search") {
+    const history = entries
+      .filter((entry) => entry.role !== "error")
+      .map((entry) => ({
+        role: entry.role,
+        content: entry.content.trim(),
+      }));
+
+    return await queryLLM(
+      query,
+      mode,
+      history,
+      baseReferencesIndex,
+      setCurrent
+    );
+  } else {
+    // ----------------- HUMAN AUTHORED CONTENT RETRIEVAL ------------------
+    const [question_id] = query.split("\n", 2);
+    if (question_id) {
+      return await getStampyContent(question_id);
+    }
+    const result = {
+      role: "error",
+      content: "Could not extract Stampy id from " + query,
+    };
+    return { result } as SearchResult;
+  }
+};
diff --git a/web/src/pages/index.tsx b/web/src/pages/index.tsx
index b1765d9..d089953 100644
--- a/web/src/pages/index.tsx
+++ b/web/src/pages/index.tsx
@@ -5,6 +5,7 @@ import Image from 'next/image';
 
 import Page from "../components/page"
 import { API_URL } from "../settings"
+import { queryLLM, getStampyContent, runSearch } from "../hooks/useSearch";
 import type { Citation, Entry, UserEntry, AssistantEntry, ErrorMessage, StampyMessage } from "../types";
 import { SearchBox, Followup } from "../components/searchbox";
 import { GlossarySpan } from "../components/glossary";
@@ -32,18 +33,21 @@ type Mode = "rookie" | "concise" | "default";
 // smooth-scroll to the bottom of the window if we're already less than 30% a screen away
 // note: finicky interaction with "smooth" - maybe fix later.
 function scroll30() {
-  if (document.documentElement.scrollHeight - window.scrollY > window.innerHeight * 1.3) return;
-  window.scrollTo({top: document.body.scrollHeight, behavior: "smooth"});
+    if (
+        document.documentElement.scrollHeight - window.scrollY >
+        window.innerHeight * 1.3
+    )
+        return;
+    window.scrollTo({ top: document.body.scrollHeight, behavior: "smooth" });
 }
 
 const Home: NextPage = () => {
-
-  const [ entries, setEntries ] = useState<Entry[]>([]);
-  const [ runningIndex, setRunningIndex ] = useState(0);
-  const [ loadState, setLoadState ] = useState<State>({state: "idle"});
+  const [entries, setEntries] = useState<Entry[]>([]);
+  const [runningIndex, setRunningIndex] = useState(0);
+  const [current, setCurrent] = useState<CurrentSearch>();
 
   // [state, ready to save to localstorage]
-  const [ mode, setMode ] = useState<[Mode, boolean]>(["default", false]);
+  const [mode, setMode] = useState<[Mode, boolean]>(["default", false]);
 
   // store mode in localstorage
   useEffect(() => {
@@ -56,6 +60,12 @@ const Home: NextPage = () => {
     setMode([mode, true]);
   }, []);
 
+  const updateCurrent = (current: CurrentSearch) => {
+      setCurrent(current);
+      if (current?.phase === "streaming") {
+          scroll30();
+      }
+  };
 
   const search = async (
     query: string,
@@ -65,214 +75,46 @@ const Home: NextPage = () => {
   ) => {
 
     // clear the query box, append to entries
+      const userEntry: Entry = {
+          role: "user",
+          content: query_source === "search" ? query : query.split("\n", 2)[1]!,
+      };
+      setEntries((prev) => [...prev, userEntry]);
+      disable();
 
-    const old_entries = entries;
-    const new_entries: Entry[] = [...old_entries, {
-      role: "user",
-      content: query_source === "search" ? query : query.split("\n", 2)[1]!,
-    }];
-    setEntries(new_entries);
-    disable();
+      const { result, followups } = await runSearch(
+          query,
+          query_source,
+          mode[0],
+          runningIndex,
+          entries,
+          updateCurrent
+      );
 
-
-    // ----------------------------- LLM BASED -----------------------------
-    if (query_source === "search") {
-      // 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": "*"
-        },
-
-        body: JSON.stringify({query: query, mode: mode[0], history:
-          old_entries.filter((entry) => entry.role !== "error")
-                     .map((entry) => {
-                       return {
-                         "role" : entry.role,
-                         "content" : entry.content.trim(),
-                       }
-                     })
-        }),
-
-      });
-
-      if (!res.ok) {
-        enable([]);
-        setLoadState({state: "idle"});
-        setEntries([...new_entries, {role: "error", content: "POST Error: " + res.status}]);
-        return;
+      if (query_source === "search") {
+          setRunningIndex(runningIndex + ProcessText(result.content, 0)[1].size);
       }
-
-      // read back the SSE stream
-
-      const reader = res.body!.getReader();
-      var message = "";
-      var followups: Followup[] = [];
-      read: 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;
-
-        // note: this form isn't even remotely close to optimal in terms of
-        // network usage. Lots of json overhead.
-
-        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);
-          // Fixes #43
-          if (!line.startsWith("data: ") && line !== "") message += line;
-          if (line === "") {
-            if (message !== "") {
-              const data = JSON.parse(message);
-
-              switch (data.state) {
-
-                case "loading":
-
-                  // 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: 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
-                    }};
-                  });
-
-                  scroll30();
-                  break;
-
-                case "done":
-
-                  // append the response to the entries, reset to normal
-                  setLoadState((s) => {
-                    if (s.state === "streaming") {
-                      setEntries([...new_entries, s.response]);
-                      setRunningIndex((i) => (i + ProcessText(s.response.content, 0)[1].size));
-                    }
-
-                    return {state: "idle"};
-                  });
-
-                  // add any potential followup questions
-                  var i = 0;
-                  while ('followup_' + i in data) {
-                    followups = [...followups, data['followup_' + i]];
-                    i++;
-                  }
-
-                  break read;
-
-                case "error":
-                  setEntries([...new_entries, {role: "error", content: data.error}]);
-                  setLoadState({state: "idle"});
-                  break read;
-
-              }
-            }
-            message = "";
-          }
-        }
-      }
-
-      enable(followups);
+      setEntries((prev) => [...prev, result]);
+      enable(followups || []);
       scroll30();
-
-    } else {
-    // ----------------- HUMAN AUTHORED CONTENT RETRIEVAL ------------------
-      const query_id = query.split("\n", 2)[0];
-
-      const res = await fetch(API_URL + "/human/" + query_id, {
-        method: "GET",
-        headers: {
-          "Content-Type": "application/json",
-          "Accept": "application/json",
-          "Allow-Control-Allow-Origin": "*"
-        },
-      });
-
-      if (!res.ok) {
-        enable([]);
-        setLoadState({state: "idle"});
-        setEntries([...new_entries, {role: "error", content: "POST Error: " + res.status}]);
-        return;
-      }
-
-      const data = (await res.json()).data;
-
-      setEntries([...new_entries, {
-        role: "stampy",
-        content: data.text,
-        url: "https://aisafety.info/?state=" + data.pageid,
-      }]);
-
-      // re-enable the searchbox, with the question that was just answered
-      // removed from the list of possible followups.
-
-      // create an array of new followup questions from the data
-      const f_new = data.relatedQuestions.map((f: any) => { return {
-        pageid: f.pageid!,
-        text: f.title!,
-        score: 0
-      };});
-
-      const fpids = new Set(f_new.map((f: Followup) => f.pageid));
-
-      enable((f_old: Followup[]) => {
-        const f_old_filtered = f_old.filter((f) => f.pageid !== data.pageid && !fpids.has(f.pageid));
-        return [...f_new, ...f_old_filtered].slice(0, MAX_FOLLOWUPS); // this is correct, it's N and not N-1 in javascript fsr
-      });
-
-      scroll30();
-    }
   };
 
   var last_entry = <></>;
-  if (loadState.state === "loading") {
-    switch (loadState.phase) {
-      case "semantic": last_entry = <p>Loading: Performing semantic search...</p>; break;
-      case "prompt": last_entry = <p>Loading: Creating prompt...</p>; break;
-      case "llm": last_entry = <p>Loading: Waiting for LLM...</p>; break;
-    }
-  } else if (loadState.state === "streaming") {
-    last_entry = <ShowAssistantEntry entry={loadState.response}/>;
+  switch (current?.phase) {
+    case "semantic":
+      last_entry = <p>Loading: Performing semantic search...</p>;
+      break;
+    case "prompt":
+      last_entry = <p>Loading: Creating prompt...</p>;
+      break;
+    case "llm":
+      last_entry = <p>Loading: Waiting for LLM...</p>;
+      break;
+    case "streaming":
+      last_entry = <ShowAssistantEntry entry={current} />;
+      break;
   }
 
-
-
   return (
     <Page page="index">
       <Controls mode={mode} setMode={setMode} />
diff --git a/web/src/settings.ts b/web/src/settings.ts
index cf42900..abb4211 100644
--- a/web/src/settings.ts
+++ b/web/src/settings.ts
@@ -1,2 +1,3 @@
 export const API_URL = process.env.NEXT_PUBLIC_API_URL || "http://127.0.0.1:3001";
 export const STAMPY_URL = process.env.STAMPY_URL || "https://aisafety.info"
+export const STAMPY_CONTENT_URL = process.env.STAMPY_CONTENT_URL || `${API_URL}/human`

From 025e6881f0cdce46c4a8812b44f0b15797538d0e Mon Sep 17 00:00:00 2001
From: Daniel O'Connell <github@ahiru.pl>
Date: Fri, 29 Sep 2023 00:43:18 +0200
Subject: [PATCH 2/3] move Stampy URL cleaner to JS

---
 web/src/hooks/useSearch.ts | 13 ++++++++-----
 web/src/pages/index.tsx    |  1 +
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/web/src/hooks/useSearch.ts b/web/src/hooks/useSearch.ts
index f08dff3..6c1596e 100644
--- a/web/src/hooks/useSearch.ts
+++ b/web/src/hooks/useSearch.ts
@@ -127,9 +127,7 @@ export const queryLLM = async (
   }
 
   try {
-    const results = await extractAnswer(res, baseReferencesIndex, setCurrent);
-    setCurrent(undefined);
-    return results;
+    return await extractAnswer(res, baseReferencesIndex, setCurrent);
   } catch (e) {
     return {
       result: { role: "error", content: e ? e.toString() : "unknown error" },
@@ -137,6 +135,11 @@ export const queryLLM = async (
   }
 };
 
+const cleanStampyContent = (contents: string) => contents.replace(
+    /<a(.*?)href="\/\?state=([a-zA-Z0-9]+.*?)"(.*?)<\/a>/g,
+    (_, pre, linkParts, post) => `<a${pre}href="${STAMPY_URL}/?state=${linkParts}"${post}</a>`
+);
+
 export const getStampyContent = async (
   question_id: string
 ): Promise<SearchResult> => {
@@ -157,7 +160,7 @@ export const getStampyContent = async (
 
   let result = {
     role: "stampy",
-    content: data.text,
+    content: cleanStampyContent(data.text),
     url: `${STAMPY_URL}/?state=${data.pageid}`,
   } as StampyMessage;
 
@@ -189,7 +192,7 @@ export const runSearch = async (
   baseReferencesIndex: number,
   entries: Entry[],
   setCurrent: (c: CurrentSearch) => void
-) => {
+): SearchResult => {
   if (query_source === "search") {
     const history = entries
       .filter((entry) => entry.role !== "error")
diff --git a/web/src/pages/index.tsx b/web/src/pages/index.tsx
index d089953..0c68982 100644
--- a/web/src/pages/index.tsx
+++ b/web/src/pages/index.tsx
@@ -90,6 +90,7 @@ const Home: NextPage = () => {
           entries,
           updateCurrent
       );
+      setCurrent(undefined);
 
       if (query_source === "search") {
           setRunningIndex(runningIndex + ProcessText(result.content, 0)[1].size);

From 04c0d743f484730d1ab7ee460b3b29eb710065a3 Mon Sep 17 00:00:00 2001
From: Daniel O'Connell <github@ahiru.pl>
Date: Fri, 29 Sep 2023 13:23:11 +0200
Subject: [PATCH 3/3] PR comments

---
 web/src/hooks/useSearch.ts | 22 +++++++++++-----------
 1 file changed, 11 insertions(+), 11 deletions(-)

diff --git a/web/src/hooks/useSearch.ts b/web/src/hooks/useSearch.ts
index 6c1596e..8cceb0f 100644
--- a/web/src/hooks/useSearch.ts
+++ b/web/src/hooks/useSearch.ts
@@ -11,6 +11,8 @@ import type {
 } from "../types";
 
 const MAX_FOLLOWUPS = 4;
+const DATA_HEADER = "data: "
+const EVENT_END_HEADER = "event: close\n"
 
 type HistoryEntry = {
   role: "error" | "stampy" | "assistant" | "user";
@@ -27,14 +29,14 @@ export async function* iterateData(res: Response) {
     if (done) return;
 
     const chunk = new TextDecoder("utf-8").decode(value);
-    if (chunk.startsWith("event: close\n")) return;
+    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: ")) {
-        message += line.slice(6);
+      if (line.startsWith(DATA_HEADER)) {
+        message += line.slice(DATA_HEADER.length);
         // Fixes #43
       } else if (line !== "") {
         message += line;
@@ -106,7 +108,6 @@ const fetchLLM = async (
     headers: {
       "Content-Type": "application/json",
       Accept: "text/event-stream",
-      "Allow-Control-Allow-Origin": "*",
     },
 
     body: JSON.stringify({ query, mode, history }),
@@ -141,14 +142,13 @@ const cleanStampyContent = (contents: string) => contents.replace(
 );
 
 export const getStampyContent = async (
-  question_id: string
+  questionId: string
 ): Promise<SearchResult> => {
-  const res = await fetch(`${STAMPY_CONTENT_URL}/${question_id}`, {
+  const res = await fetch(`${STAMPY_CONTENT_URL}/${questionId}`, {
     method: "GET",
     headers: {
       "Content-Type": "application/json",
       Accept: "application/json",
-      "Allow-Control-Allow-Origin": "*",
     },
   });
 
@@ -179,7 +179,7 @@ export const getStampyContent = async (
     const f_old_filtered = f_old.filter(
       (f) => f.pageid !== data.pageid && !fpids.has(f.pageid)
     );
-    return [...f_new, ...f_old_filtered].slice(0, MAX_FOLLOWUPS); // this is correct, it's N and not N-1 in javascript fsr
+    return [...f_new, ...f_old_filtered].slice(0, MAX_FOLLOWUPS);
   };
 
   return { followups, result };
@@ -210,9 +210,9 @@ export const runSearch = async (
     );
   } else {
     // ----------------- HUMAN AUTHORED CONTENT RETRIEVAL ------------------
-    const [question_id] = query.split("\n", 2);
-    if (question_id) {
-      return await getStampyContent(question_id);
+    const [questionId] = query.split("\n", 2);
+    if (questionId) {
+      return await getStampyContent(questionId);
     }
     const result = {
       role: "error",