Move the bos_token and eos_token to the top level tokens field in the…

… manifest
axonzeta · Nov 6, 2024 · fd82494 · fd82494
1 parent ddbd5c1
commit fd82494
Show file tree

Hide file tree

Showing 5 changed files with 10 additions and 10 deletions.
diff --git a/src/extension/background/APIHandler/AILanguageModelHandler.ts b/src/extension/background/APIHandler/AILanguageModelHandler.ts
@@ -86,8 +86,8 @@ class AILanguageModelHandler {
       const template = new Template(promptConfig.template)
       const prompt = template.render({
         messages: history,
-        bos_token: promptConfig.bosToken,
-        eos_token: promptConfig.eosToken
+        bos_token: manifest.tokens.bosToken,
+        eos_token: manifest.tokens.eosToken
       })
 
       const tokenCount = await AIPrompter.countTokens(prompt, props, {})

diff --git a/src/extension/background/APIHandler/AIRewriterHandler.ts b/src/extension/background/APIHandler/AIRewriterHandler.ts
@@ -134,8 +134,8 @@ class AIRewriterHandler {
       shared_context: sharedContext,
       context,
       input,
-      bos_token: config.bosToken,
-      eos_token: config.eosToken
+      bos_token: manifest.tokens.bosToken,
+      eos_token: manifest.tokens.eosToken
     })
   }
 }

diff --git a/src/extension/background/APIHandler/AISummarizerHandler.ts b/src/extension/background/APIHandler/AISummarizerHandler.ts
@@ -133,8 +133,8 @@ class AISummarizerHandler {
       shared_context: sharedContext,
       context,
       input,
-      bos_token: config.bosToken,
-      eos_token: config.eosToken
+      bos_token: manifest.tokens.bosToken,
+      eos_token: manifest.tokens.eosToken
     })
   }
 }

diff --git a/src/extension/background/APIHandler/AIWriterHandler.ts b/src/extension/background/APIHandler/AIWriterHandler.ts
@@ -133,8 +133,8 @@ class AIWriterHandler {
       sharedContext,
       context,
       input,
-      bos_token: config.bosToken,
-      eos_token: config.eosToken
+      bos_token: manifest.tokens.bosToken,
+      eos_token: manifest.tokens.eosToken
     })
   }
 }

diff --git a/src/shared/AIModelManifest.ts b/src/shared/AIModelManifest.ts
@@ -22,8 +22,6 @@ export type AIModelManifestConfig = {
 
 export type AIModelPromptConfig = {
   template: string
-  bosToken?: string
-  eosToken?: string
 }
 
 export type AIModelManifest = {
@@ -37,6 +35,8 @@ export type AIModelManifest = {
     max: number
     default: number
     stop?: string[]
+    bosToken?: string
+    eosToken?: string
   }
   prompts?: {
     languageModel?: AIModelPromptConfig