add llama prefix + add assistant role to openai responses

harmony-one · Oct 18, 2023 · 9ba5791 · 9ba5791
1 parent 1a0c902
commit 9ba5791
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 20 deletions.
diff --git a/src/config.ts b/src/config.ts
@@ -87,7 +87,8 @@ export default {
           : ['d.'],
         newPrefix: process.env.NEW_PREFIX
           ? process.env.NEW_PREFIX.split(',')
-          : ['n.', '..']
+          : ['n.', '..'],
+        llamaPrefix: ['*']
       },
       minimumBalance: parseInt(process.env.MIN_BALANCE ?? '0')
     }

diff --git a/src/modules/llms/helpers.ts b/src/modules/llms/helpers.ts
@@ -39,18 +39,8 @@ export const isMentioned = (
   return false
 }
 
-export const hasChatPrefix = (prompt: string): string => {
-  const prefixList = config.openAi.chatGpt.prefixes.chatPrefix
-  for (let i = 0; i < prefixList.length; i++) {
-    if (prompt.toLocaleLowerCase().startsWith(prefixList[i])) {
-      return prefixList[i]
-    }
-  }
-  return ''
-}
-
-export const hasDallePrefix = (prompt: string): string => {
-  const prefixList = config.openAi.chatGpt.prefixes.dallePrefix
+export const hasLlamaPrefix = (prompt: string): string => {
+  const prefixList = config.openAi.chatGpt.prefixes.llamaPrefix
   for (let i = 0; i < prefixList.length; i++) {
     if (prompt.toLocaleLowerCase().startsWith(prefixList[i])) {
       return prefixList[i]
@@ -200,7 +190,7 @@ export const sendMessage = async (
 
 export const hasPrefix = (prompt: string): string => {
   return (
-    hasBardPrefix(prompt)
+    hasBardPrefix(prompt) || hasLlamaPrefix(prompt)
   )
 }
 

diff --git a/src/modules/llms/index.ts b/src/modules/llms/index.ts
@@ -20,6 +20,7 @@ import {
   addDocToCollection,
   addUrlToCollection,
   hasBardPrefix,
+  hasLlamaPrefix,
   hasPrefix,
   hasUrl,
   isMentioned,
@@ -111,6 +112,11 @@ export class LlmsBot implements PayableBot {
       return
     }
 
+    if (hasLlamaPrefix(ctx.message?.text ?? '') !== '') {
+      await this.onCurrentCollection(ctx)
+      return
+    }
+
     if (ctx.hasCommand(SupportedCommands.pdf.name)) {
       await this.onPdfCommand(ctx)
       return
@@ -174,7 +180,7 @@ export class LlmsBot implements PayableBot {
       if (documentType === 'application/pdf' && ctx.chat?.id && ctx.chat.type === 'private') {
         const url = file.getUrl()
         const fileName = ctx.message?.document?.file_name ?? file.file_id
-        const prompt = ctx.message?.caption ?? 'Summarize this context'
+        const prompt = ctx.message?.caption ?? 'Summarize this context' //  from the PDF file
         await addDocToCollection(ctx, ctx.chat.id, fileName, url, prompt)
         if (!ctx.session.collections.isProcessingQueue) {
           ctx.session.collections.isProcessingQueue = true
@@ -202,7 +208,7 @@ export class LlmsBot implements PayableBot {
       if (ctx.match) {
         prompt = ctx.match as string
       } else {
-        prompt = 'Summarize this context'
+        prompt = 'Summarize this context from the PDF file'
       }
       if (filename !== '' && ctx.chat?.id) {
         const collection = ctx.session.collections.activeCollections.find(c => c.fileName === filename)
@@ -269,8 +275,16 @@ export class LlmsBot implements PayableBot {
   private async onCurrentCollection (ctx: OnMessageContext | OnCallBackQueryData): Promise<void> {
     try {
       let prompt = ''
-      prompt = ctx.match as string
-      // add prefix logic here if prompt == ''
+      if (ctx.match) {
+        prompt = ctx.match as string
+      } else {
+        const prefix = hasLlamaPrefix(ctx.message?.text ?? '')
+        if (prefix && ctx.message?.text) {
+          prompt = ctx.message?.text.slice(prefix.length)
+        } else {
+          prompt = 'Summarize this context'
+        }
+      }
       const collectionName = ctx.session.collections.currentCollection
       const collection = ctx.session.collections.activeCollections.find(c => c.collectionName === collectionName)
       if (collection && collectionName) {
@@ -297,7 +311,7 @@ export class LlmsBot implements PayableBot {
             role: 'user'
           }, {
             content: response.completion,
-            role: 'system'
+            role: 'assistant'
           })
           await ctx.api.editMessageText(ctx.chat?.id ?? '',
             msgId, response.completion,
@@ -331,6 +345,12 @@ export class LlmsBot implements PayableBot {
       const collection = ctx.session.collections.activeCollections.find(c => c.url === url)
       if (collection) {
         const conversation = this.getCollectionConversation(ctx, collection)
+        if (conversation.length === 0) {
+          conversation.push({
+            role: 'system',
+            content: `${collection.collectionType === 'PDF' ? 'The context comes from an URL linked to a PDF file' : 'The context comes from the web crawler of the given URL'}`
+          })
+        }
         const msgId = (
           await ctx.reply('...', {
             message_thread_id:
@@ -356,7 +376,7 @@ export class LlmsBot implements PayableBot {
             role: 'user'
           }, {
             content: response.completion,
-            role: 'system'
+            role: 'assistant'
           })
           await ctx.api.editMessageText(ctx.chat?.id ?? '',
             msgId, response.completion,