From a2893c5c55e4c0bcbd016d6270bef92d6f3fc6d9 Mon Sep 17 00:00:00 2001
From: smathot <s.mathot@cogsci.nl>
Date: Wed, 6 Mar 2024 15:25:56 +0100
Subject: [PATCH] Allow selection of model through menu

---
 heymans/routes/api.py                  |  1 +
 heymans/routes/app.py                  | 17 +++++++++++++++++
 heymans/templates/main.js              |  7 ++++---
 heymans/templates/menu.html            | 24 ++++++++++++++++++++++--
 heymans/templates/stylesheet.css.jinja | 11 +++++++----
 5 files changed, 51 insertions(+), 9 deletions(-)

diff --git a/heymans/routes/api.py b/heymans/routes/api.py
index 7cbef8e..871b12c 100644
--- a/heymans/routes/api.py
+++ b/heymans/routes/api.py
@@ -19,6 +19,7 @@ def api_chat_start():
     data = request.json
     session['user_message'] = data.get('message', '')
     session['search_first'] = data.get('search_first', True)
+    session['model'] = data.get('model', None)
     session['message_id'] = data.get('message_id', None)
     heymans = get_heymans()
     redis_client.delete(f'stream_cancel_{heymans.user_id}')
diff --git a/heymans/routes/app.py b/heymans/routes/app.py
index f16585a..7533c87 100644
--- a/heymans/routes/app.py
+++ b/heymans/routes/app.py
@@ -23,6 +23,23 @@ def __init__(self, username):
 
 
 def get_heymans():
+    model = session.get('model', None)
+    logger.info(f'model: {model}')
+    if model is not None:
+        if model == 'gpt4':
+            config.search_model = 'gpt-3.5'
+            config.condense_model = 'gpt-3.5'
+            config.answer_model = 'gpt-4'
+        elif model == 'claude':
+            config.search_model = 'claude-3-sonnet'
+            config.condense_model = 'claude-3-sonnet'
+            config.answer_model = 'claude-3-opus'
+        elif model == 'mistral':
+            config.search_model = 'mistral-small'
+            config.condense_model = 'mistral-small'
+            config.answer_model = 'mistral-medium'
+        else:
+            logger.warning('unknown model, using default')
     return Heymans(user_id=current_user.get_id(), persistent=True,
                    encryption_key=session['encryption_key'],
                    search_first=session.get('search_first', True))
diff --git a/heymans/templates/main.js b/heymans/templates/main.js
index a70113c..b5bb3aa 100644
--- a/heymans/templates/main.js
+++ b/heymans/templates/main.js
@@ -119,7 +119,7 @@ async function sendMessage(message) {
     await fetchWithRetry('{{ server_url }}/api/chat/start', {
         method: 'POST',
         headers: {'Content-Type': 'application/json'},
-        body: requestBody(message, searchFirst, user_message_id)
+        body: requestBody(message, searchFirst, selectedModel, user_message_id)
     }).catch(e => {
         console.error('Failed to start chat session:', e);
     });
@@ -244,11 +244,12 @@ function deleteMessage(messageId) {
 
 
 
-function requestBody(message, searchFirst, user_message_id) {
+function requestBody(message, searchFirst, selectedModel, user_message_id) {
     return JSON.stringify({
         message: message,
         search_first: searchFirst,
-        message_id:user_message_id
+        model: selectedModel,
+        message_id: user_message_id
     })
 }
 
diff --git a/heymans/templates/menu.html b/heymans/templates/menu.html
index ef9e08f..90e49d9 100644
--- a/heymans/templates/menu.html
+++ b/heymans/templates/menu.html
@@ -149,6 +149,16 @@
     <hr class="menu-divider">
     <div @click="toggleDocs"><i :class="toggleDocsIcon"></i>&nbsp; {{ toggleDocsLabel }}</div>
     <hr class="menu-divider">
+    <!-- Model selection -->
+    <div>
+      <label for="model-selection">Model: </label>
+      <select id="model-selection" v-model="selectedModel" @change="onModelChange">
+        <option value="gpt4">OpenAI (GPT4)</option>
+        <option value="claude">Anthropic (Claude 3)</option>
+        <option value="mistral">Mistral</option>
+      </select>
+    </div>
+    <hr class="menu-divider">
     <div @click="navigateTo('/about')">About</div>
     <div @click="navigateTo('/terms')">Terms and conditions</div>
     <div @click="navigateTo('/subscribe/customer-portal')" v-if="showManageSubscriptions">Manage subscription</div>
@@ -210,6 +220,10 @@
 if (window.searchFirst === null) {
   window.searchFirst = true;
 }
+window.selectedModel = JSON.parse(localStorage.getItem('selectedModel'));
+if (window.selectedModel == null) {
+  window.selectedModel = 'gpt4';
+}
 new Vue({
   el: '#app',
   data: {
@@ -221,7 +235,8 @@
     toggleDocsLabel: "{{ search_first_menu_label }}",
     attachments: {},
     isUploading: false,
-    isDocsToggled: window.searchFirst
+    isDocsToggled: window.searchFirst,
+    selectedModel: window.selectedModel
   },
   computed: {
     sortedConversations() {
@@ -234,7 +249,7 @@
     },
   },
   methods: {
-  getNewConversationUrl() {
+    getNewConversationUrl() {
       return '/api/conversation/new/' + Number(searchFirst);
     },
     toggleDocs: function() {
@@ -243,6 +258,11 @@
       console.log('documentation search enabled: ' + this.isDocsToggled);
       localStorage.setItem('searchFirst', JSON.stringify(window.searchFirst));
     },
+    onModelChange() {
+      console.log('model changed:', this.selectedModel);
+      window.selectedModel = this.selectedModel;
+      localStorage.setItem('selectedModel', JSON.stringify(window.selectedModel));
+    },
     toggleMenu: function() {
       this.menuVisible = !this.menuVisible;
     },
diff --git a/heymans/templates/stylesheet.css.jinja b/heymans/templates/stylesheet.css.jinja
index 94144b2..77f14ca 100644
--- a/heymans/templates/stylesheet.css.jinja
+++ b/heymans/templates/stylesheet.css.jinja
@@ -45,12 +45,19 @@ a:visited {
     cursor: pointer;
 }
 
+.menu select,
 #container select,
 #container button,
 #container input {
     font-family: "Roboto Light", sans-serif;
 }
 
+.menu select {
+    font-size: 1em;
+    background: none;
+    border: none;
+}
+
 .codehilite pre {
     overflow: auto;
     line-height: 150%;
@@ -306,10 +313,6 @@ button#start {
     color: #90A4AE;
 }
 
-.message-answer-model {
-    display: none;
-}
-
 .message-ai,
 .message-user {
     border-radius: 4px;