Merge pull request #114 from StampyAI/chat-options

Chat options

api/src/stampy_chat/settings.py

@@ -1,8 +1,12 @@
+from collections import namedtuple
 import tiktoken
 
 from stampy_chat.env import COMPLETIONS_MODEL
 
 
+Model = namedtuple('Model', ['maxTokens', 'topKBlocks'])
+
+
 SOURCE_PROMPT = (
     "You are a helpful assistant knowledgeable about AI Alignment and Safety. "
     "Please give a clear and coherent answer to the user's questions.(written after \"Q:\") "
@@ -44,6 +48,12 @@ DEFAULT_PROMPTS = {
     'question': QUESTION_PROMPT,
     'modes': PROMPT_MODES,
 }
+MODELS = {
+    'gpt-3.5-turbo': Model(4097, 10),
+    'gpt-3.5-turbo-16k': Model(16385, 30),
+    'gpt-4': Model(8192, 20),
+    # 'gpt-4-32k': Model(32768, 30),
+}
 
 
 class Settings:
@@ -99,23 +109,21 @@ class Settings:
             self.encoders[value] = tiktoken.get_encoding(value)
 
     def set_completions(self, completions, numTokens=None, topKBlocks=None):
+        if completions not in MODELS:
+            raise ValueError(f'Unknown model: {completions}')
         self.completions = completions
 
-        # Set the max number of tokens sent in the prompt
+        # Set the max number of tokens sent in the prompt - see https://platform.openai.com/docs/models/gpt-4
         if numTokens is not None:
             self.numTokens = numTokens
-        elif completions == 'gtp-4':
-            self.numTokens = 8191
         else:
-            self.numTokens = 4095
+            self.numTokens = MODELS[completions].maxTokens
 
         # Set the max number of blocks used as citations
         if topKBlocks is not None:
             self.topKBlocks = topKBlocks
-        elif completions == 'gtp-4':
-            self.topKBlocks = 20
         else:
-            self.topKBlocks = 10
+            self.topKBlocks = MODELS[completions].topKBlocks
 
     @property
     def prompt_modes(self):

api/tests/stampy_chat/test_chat.py

@@ -255,14 +255,14 @@ def test_check_openai_moderation_not_flagged():
 
 
 @pytest.mark.parametrize('prompt, remaining', (
-    ([{'role': 'system', 'content': 'bla'}], 4043),
+    ([{'role': 'system', 'content': 'bla'}], 4045),
     (
         [
             {'role': 'system', 'content': 'bla'},
             {'role': 'user', 'content': 'message 1'},
             {'role': 'assistant', 'content': 'response 1'},
         ],
-        4035
+        4037
     ),
     (
         [

web/src/components/chat.tsx

@@ -103,16 +103,14 @@ const Chat = ({ sessionId, settings, onQuery, onNewEntry }: ChatParams) => {
   const search = async (
     query: string,
     query_source: "search" | "followups",
-    disable: () => void,
-    enable: (f_set: Followup[] | ((fs: Followup[]) => Followup[])) => void
+    enable: (f_set: Followup[] | ((fs: Followup[]) => Followup[])) => void,
+    controller: AbortController
   ) => {
     // clear the query box, append to entries
     const userEntry: Entry = {
       role: "user",
       content: query_source === "search" ? query : query.split("\n", 2)[1]!,
     };
-    addEntry(userEntry);
-    disable();
 
     const { result, followups } = await runSearch(
       query,
@@ -120,13 +118,18 @@ const Chat = ({ sessionId, settings, onQuery, onNewEntry }: ChatParams) => {
       settings,
       entries,
       updateCurrent,
-      sessionId
+      sessionId,
+      controller
     );
+    if (result.content !== "aborted") {
+      addEntry(userEntry);
+      addEntry(result);
+      enable(followups || []);
+      scroll30();
+    } else {
+      enable([]);
+    }
     setCurrent(undefined);
-
-    addEntry(result);
-    enable(followups || []);
-    scroll30();
   };
 
   var last_entry = <></>;

web/src/components/searchbox.tsx

@@ -32,8 +32,8 @@ const SearchBoxInternal: React.FC<{
   search: (
     query: string,
     query_source: "search" | "followups",
-    disable: () => void,
-    enable: (f_set: Followup[] | ((fs: Followup[]) => Followup[])) => void
+    enable: (f_set: Followup[] | ((fs: Followup[]) => Followup[])) => void,
+    controller: AbortController
   ) => void;
   onQuery?: (q: string) => any;
 }> = ({ search, onQuery }) => {
@@ -43,20 +43,21 @@ const SearchBoxInternal: React.FC<{
   const [query, setQuery] = useState(initial_query);
   const [loading, setLoading] = useState(false);
   const [followups, setFollowups] = useState<Followup[]>([]);
+  const [controller, setController] = useState(new AbortController());
 
   const inputRef = React.useRef<HTMLTextAreaElement>(null);
 
   // because everything is async, I can't just manually set state at the
   // point we do a search. Instead it needs to be passed into the search
   // method, for some reason.
-  const enable = (f_set: Followup[] | ((fs: Followup[]) => Followup[])) => {
-    setLoading(false);
-    setFollowups(f_set);
-  };
-  const disable = () => {
-    setLoading(true);
-    setQuery("");
-  };
+  const enable =
+    (controller: AbortController) =>
+    (f_set: Followup[] | ((fs: Followup[]) => Followup[])) => {
+      if (!controller.signal.aborted) setQuery("");
+
+      setLoading(false);
+      setFollowups(f_set);
+    };
 
   useEffect(() => {
     // set focus on the input box
@@ -71,7 +72,17 @@ const SearchBoxInternal: React.FC<{
     inputRef.current.selectionEnd = inputRef.current.textLength;
   }, []);
 
-  if (loading) return <></>;
+  const runSearch =
+    (query: string, searchtype: "search" | "followups") => () => {
+      if (loading || query.trim() === "") return;
+
+      setLoading(true);
+      const controller = new AbortController();
+      setController(controller);
+      search(query, searchtype, enable(controller), controller);
+    };
+  const cancelSearch = () => controller.abort();
+
   return (
     <>
       <div className="mt-1 flex flex-col items-end">
@@ -81,14 +92,10 @@ const SearchBoxInternal: React.FC<{
             <li key={i}>
               <button
                 className="my-1 border border-gray-300 px-1"
-                onClick={() => {
-                  search(
-                    followup.pageid + "\n" + followup.text,
-                    "followups",
-                    disable,
-                    enable
-                  );
-                }}
+                onClick={runSearch(
+                  followup.pageid + "\n" + followup.text,
+                  "followups"
+                )}
               >
                 <span> {followup.text} </span>
               </button>
@@ -97,13 +104,7 @@ const SearchBoxInternal: React.FC<{
         })}
       </div>
 
-      <form
-        className="mt-1 mb-2 flex"
-        onSubmit={(e) => {
-          e.preventDefault();
-          search(query, "search", disable, enable);
-        }}
-      >
+      <div className="mt-1 mb-2 flex">
         <TextareaAutosize
           className="flex-1 resize-none border border-gray-300 px-1"
           ref={inputRef}
@@ -118,14 +119,18 @@ const SearchBoxInternal: React.FC<{
             // if <enter> without <shift>, submit the form (if it's not empty)
             if (e.key === "Enter" && !e.shiftKey) {
               e.preventDefault();
-              if (query.trim() !== "") search(query, "search", disable, enable);
+              runSearch(query, "search")();
             }
           }}
         />
-        <button className="ml-2" type="submit" disabled={loading}>
-          {loading ? "Loading..." : "Search"}
+        <button
+          className="ml-2"
+          type="button"
+          onClick={loading ? cancelSearch : runSearch(query, "search")}
+        >
+          {loading ? "Cancel" : "Search"}
         </button>
-      </form>
+      </div>
     </>
   );
 };

web/src/hooks/useSearch.ts

@@ -21,6 +21,12 @@ type HistoryEntry = {
   content: string;
 };
 
+const ignoreAbort = (error: Error) => {
+  if (error.name !== "AbortError") {
+    throw error;
+  }
+};
+
 export async function* iterateData(res: Response) {
   const reader = res.body!.getReader();
   var message = "";
@@ -104,9 +110,11 @@ const fetchLLM = async (
   sessionId: string,
   query: string,
   settings: LLMSettings,
-  history: HistoryEntry[]
-): Promise<Response> =>
+  history: HistoryEntry[],
+  controller: AbortController
+): Promise<Response | void> =>
   fetch(API_URL + "/chat", {
+    signal: controller.signal,
     method: "POST",
     cache: "no-cache",
     keepalive: true,
@@ -116,25 +124,31 @@ const fetchLLM = async (
     },
 
     body: JSON.stringify({ sessionId, query, history, settings }),
-  });
+  }).catch(ignoreAbort);
 
 export const queryLLM = async (
   query: string,
   settings: LLMSettings,
   history: HistoryEntry[],
   setCurrent: (e?: CurrentSearch) => void,
-  sessionId: string
+  sessionId: string,
+  controller: AbortController
 ): Promise<SearchResult> => {
   // do SSE on a POST request.
-  const res = await fetchLLM(sessionId, query, settings, history);
+  const res = await fetchLLM(sessionId, query, settings, history, controller);
 
-  if (!res.ok) {
+  if (!res) {
+    return { result: { role: "error", content: "No response from server" } };
+  } else if (!res.ok) {
     return { result: { role: "error", content: "POST Error: " + res.status } };
   }
 
   try {
     return await extractAnswer(res, setCurrent);
   } catch (e) {
+    if ((e as Error)?.name === "AbortError") {
+      return { result: { role: "error", content: "aborted" } };
+    }
     return {
       result: { role: "error", content: e ? e.toString() : "unknown error" },
     };
@@ -149,17 +163,21 @@ const cleanStampyContent = (contents: string) =>
   );
 
 export const getStampyContent = async (
-  questionId: string
+  questionId: string,
+  controller: AbortController
 ): Promise<SearchResult> => {
   const res = await fetch(`${STAMPY_CONTENT_URL}/${questionId}`, {
     method: "GET",
+    signal: controller.signal,
     headers: {
       "Content-Type": "application/json",
       Accept: "application/json",
     },
-  });
+  }).catch(ignoreAbort);
 
-  if (!res.ok) {
+  if (!res) {
+    return { result: { role: "error", content: "No response from server" } };
+  } else if (!res.ok) {
     return { result: { role: "error", content: "POST Error: " + res.status } };
   }
 
@@ -198,7 +216,8 @@ export const runSearch = async (
   settings: LLMSettings,
   entries: Entry[],
   setCurrent: (c: CurrentSearch) => void,
-  sessionId: string
+  sessionId: string,
+  controller: AbortController
 ): Promise<SearchResult> => {
   if (query_source === "search") {
     const history = entries
@@ -208,12 +227,19 @@ export const runSearch = async (
         content: entry.content.trim(),
       }));
 
-    return await queryLLM(query, settings, history, setCurrent, sessionId);
+    return await queryLLM(
+      query,
+      settings,
+      history,
+      setCurrent,
+      sessionId,
+      controller
+    );
   } else {
     // ----------------- HUMAN AUTHORED CONTENT RETRIEVAL ------------------
     const [questionId] = query.split("\n", 2);
     if (questionId) {
-      return await getStampyContent(questionId);
+      return await getStampyContent(questionId, controller);
     }
     const result = {
       role: "error",

web/src/pages/playground.tsx

@@ -42,19 +42,24 @@ const DEFAULT_PROMPTS = {
       "rather than just giving a formal definition.\n\n",
   },
 };
+const MODELS = {
+  "gpt-3.5-turbo": { numTokens: 4095, topKBlocks: 10 },
+  "gpt-3.5-turbo-16k": { numTokens: 16385, topKBlocks: 30 },
+  "gpt-4": { numTokens: 8192, topKBlocks: 20 },
+  /* 'gpt-4-32k': {numTokens: 32768, topKBlocks: 30}, */
+};
 const DEFAULT_SETTINGS = {
   prompts: DEFAULT_PROMPTS,
   mode: "default" as Mode,
   completions: "gpt-3.5-turbo",
   encoder: "cl100k_base",
-  topKBlocks: 10, //  the number of blocks to use as citations
-  numTokens: 4095,
+  topKBlocks: MODELS["gpt-3.5-turbo"].topKBlocks, //  the number of blocks to use as citations
+  numTokens: MODELS["gpt-3.5-turbo"].numTokens,
   tokensBuffer: 50, //  the number of tokens to leave as a buffer when calculating remaining tokens
   maxHistory: 10, //  the max number of previous items to use as history
   historyFraction: 0.25, //  the (approximate) fraction of num_tokens to use for history text before truncating
   contextFraction: 0.5, //  the (approximate) fraction of num_tokens to use for context text before truncating
 };
-const COMPLETION_MODELS = ["gpt-3.5-turbo", "gpt-4"];
 const ENCODERS = ["cl100k_base"];
 
 const updateIn = (
@@ -72,101 +77,107 @@ const updateIn = (
   return obj;
 };
 
-type ChatSettingsParams = {
-  settings: LLMSettings;
-  updateSettings: (updater: (settings: LLMSettings) => LLMSettings) => void;
-};
-type NumberParser = (v: any) => number;
+type Parseable = string | number | undefined;
+type NumberParser = (v: Parseable) => number;
 type InputFields = {
   field: string;
   label: string;
+  value?: Parseable;
   min?: string | number;
   max?: string | number;
   step?: string | number;
   parser?: NumberParser;
+  updater: (v: any) => any;
+};
+
+const between =
+  (
+    min: Parseable,
+    max: Parseable,
+    parser: NumberParser,
+    updater: (v: any) => any
+  ) =>
+  (event: ChangeEvent) => {
+    let num = parser((event.target as HTMLInputElement).value);
+    if (isNaN(num)) {
+      return;
+    } else if (min !== undefined && num < parser(min)) {
+      num = parser(min);
+    } else if (max !== undefined && num > parser(max)) {
+      num = parser(max);
+    }
+    updater(num);
+  };
+
+const SectionHeader = ({ text }: { text: string }) => (
+  <h4 className="col-span-4 text-lg font-semibold">{text}</h4>
+);
+
+const NumberInput = ({
+  field,
+  value,
+  label,
+  min,
+  max,
+  updater,
+  // this cast is just to satisfy typescript - it can handle numbers, strings and undefined just fine
+  parser = (v) => parseInt(v as string, 10),
+}: InputFields) => (
+  <>
+    <label htmlFor={field} className="col-span-3 inline-block">
+      {label}:{" "}
+    </label>
+    <input
+      name={field}
+      value={value}
+      className="w-20"
+      onChange={between(min, max, parser, updater)}
+      type="number"
+    />
+  </>
+);
+
+const Slider = ({
+  field,
+  value,
+  label,
+  min = 0,
+  max = 1,
+  step = 0.01,
+  // this cast is just to satisfy typescript - it can handle numbers, strings and undefined just fine
+  parser = (v) => parseFloat(v as string),
+  updater,
+}: InputFields) => (
+  <>
+    <label htmlFor={field} className="col-span-2">
+      {label}:
+    </label>
+    <input
+      name={field}
+      className="col-span-2"
+      value={value}
+      onChange={between(min, max, parser, updater)}
+      type="range"
+      min={min}
+      max={max}
+      step={step}
+    />
+  </>
+);
+
+type ChatSettingsParams = {
+  settings: LLMSettings;
+  updateSettings: (updater: (settings: LLMSettings) => LLMSettings) => void;
 };
 
 const ChatSettings = ({ settings, updateSettings }: ChatSettingsParams) => {
+  const changeVal = (field: string, value: any) =>
+    updateSettings((prev) => ({ ...prev, [field]: value }));
   const update = (setting: string) => (event: ChangeEvent) => {
-    updateSettings((prev) => ({
-      ...prev,
-      [setting]: (event.target as HTMLInputElement).value,
-    }));
+    changeVal(setting, (event.target as HTMLInputElement).value);
   };
-  const between =
-    (
-      setting: string,
-      min: number | undefined,
-      max: number | undefined,
-      parser: NumberParser
-    ) =>
-    (event: ChangeEvent) => {
-      let num = parser((event.target as HTMLInputElement).value);
-      if (isNaN(num)) {
-        return;
-      } else if (min !== undefined && num < min) {
-        num = min;
-      } else if (max !== undefined && num > max) {
-        num = max;
-      }
-      updateSettings((prev) => ({ ...prev, [setting]: num }));
-    };
-  const floatBetween = (setting: string, min?: number, max?: number) =>
-    between(setting, min, max, parseFloat);
-
-  const SectionHeader = ({ text }: { text: string }) => (
-    <h4 className="col-span-4 text-lg font-semibold">{text}</h4>
-  );
-
-  const NumberInput = ({
-    field,
-    label,
-    min,
-    max,
-    parser = (v) => parseInt(v, 10),
-  }: InputFields) => (
-    <>
-      <label htmlFor={field} className="col-span-3 inline-block">
-        {label}:{" "}
-      </label>
-      <input
-        name={field}
-        value={settings[field]}
-        className="w-20"
-        onChange={between(
-          field,
-          min ? parser(min) : undefined,
-          max ? parser(max) : undefined,
-          parser
-        )}
-        type="number"
-      />
-    </>
-  );
-  const Slider = ({
-    field,
-    label,
-    min = 0,
-    max = 1,
-    step = 0.01,
-    parser = parseFloat,
-  }: InputFields) => (
-    <>
-      <label htmlFor={field} className="col-span-2">
-        {label}:
-      </label>
-      <input
-        name={field}
-        className="col-span-2"
-        value={settings[field]}
-        onChange={floatBetween(field, parser(min), parser(max))}
-        type="range"
-        min={min}
-        max={max}
-        step={step}
-      />
-    </>
-  );
+  const updateNum = (field: string) => (num: Parseable) =>
+    changeVal(field, num);
 
   return (
     <div
@@ -181,9 +192,30 @@ const ChatSettings = ({ settings, updateSettings }: ChatSettingsParams) => {
         name="completions-model"
         className="col-span-2"
         value={settings.completions}
-        onChange={update("completions")}
+        onChange={(event: ChangeEvent) => {
+          const value = (event.target as HTMLInputElement).value;
+          const { numTokens, topKBlocks } =
+            MODELS[value as keyof typeof MODELS];
+          const prevNumTokens =
+            MODELS[settings.completions as keyof typeof MODELS].numTokens;
+          const prevTopKBlocks =
+            MODELS[settings.completions as keyof typeof MODELS].topKBlocks;
+
+          if (settings.numTokens === prevNumTokens) {
+            changeVal("numTokens", numTokens);
+          } else {
+            changeVal(
+              "numTokens",
+              Math.min(settings.numTokens || 0, numTokens)
+            );
+          }
+          if (settings.topKBlocks === prevTopKBlocks) {
+            changeVal("topKBlocks", topKBlocks);
+          }
+          changeVal("completions", value);
+        }}
       >
-        {COMPLETION_MODELS.map((name) => (
+        {Object.keys(MODELS).map((name) => (
           <option value={name} key={name}>
             {name}
           </option>
@@ -207,33 +239,50 @@ const ChatSettings = ({ settings, updateSettings }: ChatSettingsParams) => {
       </select>
 
       <SectionHeader text="Token options" />
-      <NumberInput field="numTokens" label="Tokens" min="1" />
+      <NumberInput
+        value={settings.numTokens}
+        field="numTokens"
+        label="Tokens"
+        min="1"
+        max={MODELS[settings.completions as keyof typeof MODELS].numTokens}
+        updater={updateNum("numTokens")}
+      />
       <NumberInput
         field="tokensBuffer"
+        value={settings.tokensBuffer}
         label="Number of tokens to leave as a buffer when calculating remaining tokens"
         min="0"
-        max={settings.tokensBuffer}
+        max={settings.numTokens}
+        updater={updateNum("tokensBuffer")}
       />
 
       <SectionHeader text="Prompt options" />
       <NumberInput
+        value={settings.topKBlocks}
         field="topKBlocks"
         label="Number of blocks to use as citations"
         min="1"
+        updater={updateNum("topKBlocks")}
       />
       <NumberInput
+        value={settings.maxHistory}
         field="maxHistory"
         label="The max number of previous interactions to use"
         min="0"
+        updater={updateNum("maxHistory")}
       />
 
       <Slider
+        value={settings.contextFraction}
         field="contextFraction"
         label="Approximate fraction of num_tokens to use for citations text before truncating"
+        updater={updateNum("contextFraction")}
       />
       <Slider
+        value={settings.historyFraction}
         field="historyFraction"
         label="Approximate fraction of num_tokens to use for history text before truncating"
+        updater={updateNum("historyFraction")}
       />
     </div>
   );

web/src/pages/semantic.tsx

@@ -5,35 +5,42 @@ import type { Followup } from "../types";
 import Page from "../components/page";
 import { SearchBox } from "../components/searchbox";
 
+const ignoreAbort = (error: Error) => {
+  if (error.name !== "AbortError") {
+    throw error;
+  }
+};
+
 const Semantic: NextPage = () => {
   const [results, setResults] = useState<SemanticEntry[]>([]);
 
   const semantic_search = async (
     query: string,
     _query_source: "search" | "followups",
-    disable: () => void,
-    enable: (f_set: Followup[]) => void
+    enable: (f_set: Followup[]) => void,
+    controller: AbortController
   ) => {
-    disable();
-
     const res = await fetch(API_URL + "/semantic", {
       method: "POST",
+      signal: controller.signal,
       headers: {
         "Content-Type": "application/json",
         "Access-Control-Allow-Origin": "*",
       },
       body: JSON.stringify({ query: query }),
-    });
+    }).catch(ignoreAbort);
 
-    if (!res.ok) {
+    if (!res) {
       enable([]);
+      return;
+    } else if (!res.ok) {
       console.error("load failure: " + res.status);
     }
+    enable([]);
 
     const data = await res.json();
 
     setResults(data);
-    enable([]);
   };
 
   return (