From 3f39eac0cbc32cf575cfc59339cc49c752049bdf Mon Sep 17 00:00:00 2001
From: Daniel O'Connell <github@ahiru.pl>
Date: Tue, 3 Oct 2023 21:50:20 +0200
Subject: [PATCH] Add better model validations

---
 api/src/stampy_chat/settings.py | 22 ++++++++++++------
 web/src/pages/playground.tsx    | 41 ++++++++++++++++++++++++---------
 2 files changed, 45 insertions(+), 18 deletions(-)

diff --git a/api/src/stampy_chat/settings.py b/api/src/stampy_chat/settings.py
index 70d6b1a..771ae64 100644
--- a/api/src/stampy_chat/settings.py
+++ b/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):
diff --git a/web/src/pages/playground.tsx b/web/src/pages/playground.tsx
index 4f88a74..29a2d68 100644
--- a/web/src/pages/playground.tsx
+++ b/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 = (obj: {[key: string]: any}, [head, ...rest]: string[], val: any) => {
@@ -157,14 +162,12 @@ type ChatSettingsParams = {
 }
 
 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 updateNum = (field: string) => (num: Parseable) =>
-    updateSettings((prev) => ({...prev, [field]: num}))
+  const updateNum = (field: string) => (num: Parseable) => changeVal(field, num)
 
   return (
     <div
@@ -179,9 +182,24 @@ 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>
@@ -210,6 +228,7 @@ const ChatSettings = ({settings, updateSettings}: ChatSettingsParams) => {
         field="numTokens"
         label="Tokens"
         min="1"
+        max={MODELS[settings.completions as keyof typeof MODELS].numTokens}
         updater={updateNum('numTokens')}
       />
       <NumberInput