From 5d7b51ba1dad57df04db48a034d56f4c320108fb Mon Sep 17 00:00:00 2001
From: Webifi <john@webifi.com>
Date: Sun, 28 May 2023 11:57:07 -0500
Subject: [PATCH] Add usage totals

---
 src/lib/Chat.svelte        | 41 +++++++++++++++++++++++++++++---------
 src/lib/EditMessage.svelte | 19 +++++++++++++-----
 src/lib/Stats.svelte       | 17 ----------------
 src/lib/Storage.svelte     | 36 ++++++++++++++++++++++++++++++---
 src/lib/Types.svelte       |  4 +---
 5 files changed, 80 insertions(+), 37 deletions(-)

diff --git a/src/lib/Chat.svelte b/src/lib/Chat.svelte
index 760cbcd..5c4baed 100644
--- a/src/lib/Chat.svelte
+++ b/src/lib/Chat.svelte
@@ -17,7 +17,10 @@
     updateChatSettings,
     resetChatSettings,
     setChatSettingValue,
-    addChatFromJSON
+    addChatFromJSON,
+
+    updateRunningTotal
+
   } from './Storage.svelte'
   import { getChatSettingObjectByKey, getChatSettingList, getRequestSettingList } from './Settings.svelte'
   import {
@@ -59,6 +62,7 @@
   import { v4 as uuidv4 } from 'uuid'
   import { exportChatAsJSON, exportProfileAsJSON } from './Export.svelte'
   import { clickOutside } from 'svelte-use-click-outside'
+    import { getPrice } from './Stats.svelte';
 
   // This makes it possible to override the OpenAI API base URL in the .env file
   const apiBase = import.meta.env.VITE_API_BASE || 'https://api.openai.com'
@@ -86,9 +90,10 @@
   $: chatSettings = chat.settings
   $: globalStore = $globalStorage
 
+  // Make sure chat object is ready to go
+  updateChatSettings(chatId)
+
   onMount(async () => {
-    // Make sure chat object is ready to go
-    updateChatSettings(chatId)
 
     // Focus the input on mount
     focusInput()
@@ -147,7 +152,6 @@
   const sendRequest = async (messages: Message[], doingSummary?:boolean, withSummary?:boolean): Promise<Response> => {
     // Show updating bar
     updating = true
-    updatingMessage = ''
 
     let response: Response
 
@@ -207,6 +211,7 @@
             content: summaryPrompt
           } as Message)
           // Wait for the summary completion
+          updatingMessage = 'Building Summary...'
           const summary = await sendRequest(summarizeReq, true)
           if (summary.error) {
             // Failed to some API issue. let the original caller handle it.
@@ -241,6 +246,7 @@
             saveChatStore()
             // Re-run request with summarized prompts
             // return { error: { message: "End for now" } } as Response
+            updatingMessage = 'Continuing...'
             return await sendRequest(chat.messages, false, true)
           }
         } else if (!summaryPrompt) {
@@ -313,10 +319,7 @@
     updatingMessage = ''
 
     if (!response.error) {
-      // tc.completions++
-      // tc.completionsTokens += response.usage.completion_tokens
-      // chat.totals.push(tc)
-      // console.log('got response:', response)
+      updateRunningTotal(chatId, response.usage, response.model)
     }
 
     return response
@@ -777,7 +780,15 @@
     <button title="Send" class="button is-info" type="submit"><Fa icon={faPaperPlane} /></button>
   </p>
 </form>
-<div class="chat-focus-point" style="height.4em"></div>
+<!-- a target to scroll to -->
+<div class="chat-focus-point running-total-container">
+  {#each Object.entries(chat.usage||{}) as [model, usage]}
+  <p class="is-size-7 running-totals">
+    <em>{model}</em> total <span class="has-text-weight-bold">{usage.total_tokens}</span>
+    tokens ~= <span class="has-text-weight-bold">${getPrice(usage, model).toFixed(6)}</span>
+  </p>
+  {/each}
+</div>
 
 <svelte:window
   on:keydown={(event) => {
@@ -959,3 +970,15 @@
   </div>
 </form>
 <!-- end -->
+
+<style>
+  .running-total-container {
+    min-height:2em;
+    padding-bottom:.6em;
+    padding-left: 1.9em;
+    margin-bottom:-2.6em
+  }
+  .running-totals {
+    opacity: 0.5;
+  }
+</style>
diff --git a/src/lib/EditMessage.svelte b/src/lib/EditMessage.svelte
index 59f5912..5e81e67 100644
--- a/src/lib/EditMessage.svelte
+++ b/src/lib/EditMessage.svelte
@@ -82,6 +82,7 @@
       console.error("Can't find element with message ID", uuid)
     }
   }
+  
 </script>
 
 {#key message.uuid}
@@ -137,19 +138,22 @@
         <div id={'edit-' + message.uuid} class="message-editor" bind:innerText={message.content} contenteditable on:input={update} on:blur={exit} />
       </form>
     {:else}
-      <a href={'#'} class="message-display" on:click|preventDefault={() => {}} on:dblclick|preventDefault={() => edit('edit-' + message.uuid)}>
+      <div 
+        class="message-display" 
+          on:dblclick|preventDefault={() => edit('edit-' + message.uuid)}
+        >
         <SvelteMarkdown 
           source={message.content} 
           options={markedownOptions} 
           renderers={{ code: Code, html: Code }}
         />
-      </a>
+    </div>
     {/if}
     {#if message.role === 'system'}
-      <p class="is-size-7">System Prompt</p>
+      <p class="is-size-7 message-note">System Prompt</p>
     {:else if message.usage}
-      <p class="is-size-7">
-        This message was generated on <em>{message.model || defaultModel}</em> using <span class="has-text-weight-bold">{message.usage.total_tokens}</span>
+      <p class="is-size-7 message-note">
+        <em>{message.model || defaultModel}</em> using <span class="has-text-weight-bold">{message.usage.total_tokens}</span>
         tokens ~= <span class="has-text-weight-bold">${getPrice(message.usage, message.model || defaultModel).toFixed(6)}</span>
       </p>
     {/if}
@@ -158,6 +162,11 @@
 {/key}
 
 <style>
+  .message-note {
+    padding-top: .6em;
+    margin-bottom: -0.6em;
+    opacity: 0.5;
+  }
   .message-edit {
     display: block;
   }
diff --git a/src/lib/Stats.svelte b/src/lib/Stats.svelte
index 16b3af1..6b9422e 100644
--- a/src/lib/Stats.svelte
+++ b/src/lib/Stats.svelte
@@ -29,21 +29,4 @@
     return ((tokens.prompt_tokens * t[0]) + (tokens.completion_tokens * t[1]))
   }
 
-  export const totalUse = (totals: Usage[]): Usage => {
-    const r = {
-      completion_tokens: 0,
-      prompt_tokens: 0,
-      total_tokens: 0,
-      total: 0
-    } as Usage
-
-    (totals || ([] as Usage[])).forEach((t) => {
-      r.total += getPrice(t, t.model as any)
-      r.completion_tokens += t.completion_tokens
-      r.prompt_tokens += t.prompt_tokens
-      r.total_tokens += t.prompt_tokens
-    })
-
-    return r
-  }
 </script>
\ No newline at end of file
diff --git a/src/lib/Storage.svelte b/src/lib/Storage.svelte
index 1430e26..92c5a75 100644
--- a/src/lib/Storage.svelte
+++ b/src/lib/Storage.svelte
@@ -1,7 +1,7 @@
 <script context="module" lang="ts">
   import { persisted } from 'svelte-local-storage-store'
   import { get } from 'svelte/store'
-  import type { Chat, ChatSettings, GlobalSettings, Message, ChatSetting, GlobalSetting } from './Types.svelte'
+  import type { Chat, ChatSettings, GlobalSettings, Message, ChatSetting, GlobalSetting, Usage, Model } from './Types.svelte'
   import { getChatSettingObjectByKey, getGlobalSettingObjectByKey, getChatDefaults, getExcludeFromProfile } from './Settings.svelte'
   import { v4 as uuidv4 } from 'uuid'
   import { applyProfile, getProfile, isStaticProfile } from './Profiles.svelte'
@@ -29,7 +29,8 @@
       id: chatId,
       name: `Chat ${chatId}`,
       settings: {} as ChatSettings,
-      messages: []
+      messages: [],
+      usage:{} as Record<Model,Usage>,
     })
     chatsStorage.set(chats)
     // Apply defaults and prepare it to start
@@ -60,10 +61,12 @@
     // Add a new chat
     chats.push(chat)
     chatsStorage.set(chats)
+    // make sure it's up-to-date
+    updateChatSettings(chatId)
     return chatId
   }
 
-  // Make sure all chat settings are set with current values or defaults
+  // Make sure a chat's settings are set with current values or defaults
   export const updateChatSettings = (chatId) => {
     const chats = get(chatsStorage)
     const chat = chats.find((chat) => chat.id === chatId) as Chat
@@ -78,6 +81,15 @@
     chat.messages.forEach((m) => {
       m.uuid = m.uuid || uuidv4()
     })
+    // Make sure the usage totals object is set
+    // (some earlier versions of this had different structures)
+    const hasUsage = chat.usage && !Array.isArray(chat.usage) 
+      && typeof chat.usage === 'object' 
+      && Object.values(chat.usage).find(v=>'prompt_tokens' in v)
+    if (!hasUsage) {
+      const usageMap:Record<Model,Usage> = {}
+      chat.usage = usageMap
+    }
     chatsStorage.set(chats)
   }
   
@@ -118,6 +130,24 @@
     return (chats.find((chat) => chat.id === chatId) as Chat).settings
   }
 
+  export const updateRunningTotal = (chatId: number, usage: Usage, model:Model) => {
+    const chats = get(chatsStorage)
+    const chat = chats.find((chat) => chat.id === chatId) as Chat
+    let total:Usage = chat.usage[model]
+    if (!total) {
+      total = {
+        prompt_tokens: 0,
+        completion_tokens: 0,
+        total_tokens: 0,
+      } 
+      chat.usage[model] = total
+    }
+    total.completion_tokens += usage.completion_tokens
+    total.prompt_tokens += usage.prompt_tokens
+    total.total_tokens += usage.total_tokens
+    chatsStorage.set(chats)
+  }
+
   export const addMessage = (chatId: number, message: Message) => {
     const chats = get(chatsStorage)
     const chat = chats.find((chat) => chat.id === chatId) as Chat
diff --git a/src/lib/Types.svelte b/src/lib/Types.svelte
index 919988e..dfeaa0e 100644
--- a/src/lib/Types.svelte
+++ b/src/lib/Types.svelte
@@ -15,8 +15,6 @@
     completion_tokens: number;
     prompt_tokens: number;
     total_tokens: number;
-    total: number;
-    model?: Model;
   };
 
   export type Message = {
@@ -69,7 +67,7 @@
     id: number;
     name: string;
     messages: Message[];
-    usage?: Usage[];
+    usage: Record<Model, Usage>;
     settings: ChatSettings;
   };