refactor: getConversationPairs and promptBase

josStorer · Apr 5, 2023 · abccc6c · abccc6c
1 parent ab95a3e
commit abccc6c
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 34 deletions.
diff --git a/src/background/apis/custom-api.mjs b/src/background/apis/custom-api.mjs
@@ -9,11 +9,7 @@ import { getUserConfig, maxResponseTokenLength } from '../../config/index.mjs'
 import { fetchSSE } from '../../utils/fetch-sse'
 import { getConversationPairs } from '../../utils/get-conversation-pairs'
 import { isEmpty } from 'lodash-es'
-import { pushRecord, setAbortController } from './shared.mjs'
-
-const getCustomApiPromptBase = async () => {
-  return `I am a helpful, creative, clever, and very friendly assistant. I am familiar with various languages in the world.`
-}
+import { getCustomApiPromptBase, pushRecord, setAbortController } from './shared.mjs'
 
 /**
  * @param {Browser.Runtime.Port} port
@@ -25,7 +21,7 @@ const getCustomApiPromptBase = async () => {
 export async function generateAnswersWithCustomApi(port, question, session, apiKey, modelName) {
   const { controller, messageListener } = setAbortController(port)
 
-  const prompt = getConversationPairs(session.conversationRecords, true)
+  const prompt = getConversationPairs(session.conversationRecords, false)
   prompt.unshift({ role: 'system', content: await getCustomApiPromptBase() })
   prompt.push({ role: 'user', content: question })
   const apiUrl = (await getUserConfig()).customModelApiUrl

diff --git a/src/background/apis/openai-api.mjs b/src/background/apis/openai-api.mjs
@@ -4,22 +4,12 @@ import { maxResponseTokenLength, Models, getUserConfig } from '../../config/inde
 import { fetchSSE } from '../../utils/fetch-sse'
 import { getConversationPairs } from '../../utils/get-conversation-pairs'
 import { isEmpty } from 'lodash-es'
-import { pushRecord, setAbortController } from './shared.mjs'
-
-const getChatgptPromptBase = async () => {
-  return `You are a helpful, creative, clever, and very friendly assistant. You are familiar with various languages in the world.`
-}
-
-const getGptPromptBase = async () => {
-  return (
-    `The following is a conversation with an AI assistant.` +
-    `The assistant is helpful, creative, clever, and very friendly. The assistant is familiar with various languages in the world.\n\n` +
-    `Human: Hello, who are you?\n` +
-    `AI: I am an AI created by OpenAI. How can I help you today?\n` +
-    `Human: 谢谢\n` +
-    `AI: 不客气\n`
-  )
-}
+import {
+  getChatSystemPromptBase,
+  getCompletionPromptBase,
+  pushRecord,
+  setAbortController,
+} from './shared.mjs'
 
 /**
  * @param {Browser.Runtime.Port} port
@@ -38,9 +28,9 @@ export async function generateAnswersWithGptCompletionApi(
   const { controller, messageListener } = setAbortController(port)
 
   const prompt =
-    (await getGptPromptBase()) +
-    getConversationPairs(session.conversationRecords, false) +
-    `Human:${question}\nAI:`
+    (await getCompletionPromptBase()) +
+    getConversationPairs(session.conversationRecords, true) +
+    `Human: ${question}\nAI: `
   const apiUrl = (await getUserConfig()).customOpenAiApiUrl
 
   let answer = ''
@@ -101,8 +91,8 @@ export async function generateAnswersWithGptCompletionApi(
 export async function generateAnswersWithChatgptApi(port, question, session, apiKey, modelName) {
   const { controller, messageListener } = setAbortController(port)
 
-  const prompt = getConversationPairs(session.conversationRecords, true)
-  prompt.unshift({ role: 'system', content: await getChatgptPromptBase() })
+  const prompt = getConversationPairs(session.conversationRecords, false)
+  prompt.unshift({ role: 'system', content: await getChatSystemPromptBase() })
   prompt.push({ role: 'user', content: question })
   const apiUrl = (await getUserConfig()).customOpenAiApiUrl
 

diff --git a/src/background/apis/shared.mjs b/src/background/apis/shared.mjs
@@ -1,3 +1,22 @@
+export const getChatSystemPromptBase = async () => {
+  return `You are a helpful, creative, clever, and very friendly assistant. You are familiar with various languages in the world.`
+}
+
+export const getCompletionPromptBase = async () => {
+  return (
+    `The following is a conversation with an AI assistant.` +
+    `The assistant is helpful, creative, clever, and very friendly. The assistant is familiar with various languages in the world.\n\n` +
+    `Human: Hello, who are you?\n` +
+    `AI: I am an AI assistant. How can I help you today?\n` +
+    `Human: 没什么\n` +
+    `AI: 好的, 如果有什么需要, 随时告诉我\n`
+  )
+}
+
+export const getCustomApiPromptBase = async () => {
+  return `I am a helpful, creative, clever, and very friendly assistant. I am familiar with various languages in the world.`
+}
+
 export function setAbortController(port, onStop, onDisconnect) {
   const controller = new AbortController()
   const messageListener = (msg) => {

diff --git a/src/utils/get-conversation-pairs.mjs b/src/utils/get-conversation-pairs.mjs
@@ -1,15 +1,15 @@
-export function getConversationPairs(records, isChatgpt) {
+export function getConversationPairs(records, isCompletion) {
   let pairs
-  if (isChatgpt) {
-    pairs = []
+  if (isCompletion) {
+    pairs = ''
     for (const record of records) {
-      pairs.push({ role: 'user', content: record['question'] })
-      pairs.push({ role: 'assistant', content: record['answer'] })
+      pairs += 'Human: ' + record.question + '\nAI: ' + record.answer + '\n'
     }
   } else {
-    pairs = ''
+    pairs = []
     for (const record of records) {
-      pairs += 'Human:' + record.question + '\nAI:' + record.answer + '\n'
+      pairs.push({ role: 'user', content: record['question'] })
+      pairs.push({ role: 'assistant', content: record['answer'] })
     }
   }