Merge pull request #342 from harmony-one/group-whitelist

Group whitelist

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')
     }

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)
   )
 }
 

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,

src/modules/payment/index.ts

@@ -280,11 +280,21 @@ export class BotPayments {
     )
   }
 
-  public isGroupInWhitelist (chatId: number | string): boolean {
-    const { groupWhitelist } = config.payment
-    return (
-      groupWhitelist.includes(chatId.toString())
-    )
+  public async isGroupInWhitelist (ctx: OnMessageContext): Promise<boolean> {
+    if (ctx.chat.id && ctx.chat.type !== 'private') {
+      const { whitelist } = config.payment
+      const admins = await ctx.getChatAdministrators()
+      for (let i = 0; i < admins.length; i++) {
+        const username = admins[i].user.username ?? ''
+        if (whitelist.includes(admins[i].user.id.toString()) ||
+        (username && whitelist.includes(username.toString().toLowerCase()))) {
+          return true
+        }
+      }
+      return false
+    } else {
+      return true
+    }
   }
 
   public isPaymentsEnabled (): boolean {
@@ -295,22 +305,31 @@ export class BotPayments {
     )
   }
 
-  private skipPayment (ctx: OnMessageContext, amountUSD: number): boolean {
+  private async skipPayment (ctx: OnMessageContext, amountUSD: number): Promise<boolean> {
     const { id: userId, username = '' } = ctx.update.message.from
 
     if (!this.isPaymentsEnabled()) {
       return true
     }
+
     if (amountUSD === 0) {
       return true
     }
+
     if (this.isUserInWhitelist(userId, username)) {
       this.logger.info(
         `@${username} (${userId}) is in the whitelist, skip payment`
       )
       return true
     }
 
+    if (await this.isGroupInWhitelist(ctx)) {
+      this.logger.info(
+        `The chat ${ctx.chat.id} is in the whitelist, skip payment`
+      )
+      return true
+    }
+
     if (this.ONERate === 0) {
       this.logger.error('ONE token rate is 0, skip payment')
       return true
@@ -384,7 +403,7 @@ export class BotPayments {
       `Payment requested @${from.username}(${from.id}) in chat ${chat.id} (${chat.type}), accountId: ${accountId}, account address: ${userAccount.address}`
     )
 
-    if (this.skipPayment(ctx, amountUSD)) {
+    if (await this.skipPayment(ctx, amountUSD)) {
       await this.writePaymentLog(ctx, BigNumber(0))
       return true
     }