refactor command handling on 1country and llms bots

harmony-one · Jan 18, 2024 · 314e514 · 314e514
1 parent 1379b1b
commit 314e514
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 46 deletions.
diff --git a/src/config.ts b/src/config.ts
@@ -38,7 +38,6 @@ export default {
     model: 'chat-bison',
     minimumBalance: 0,
     isEnabled: Boolean(parseInt(process.env.LLMS_ENABLED ?? '1')),
-    prefixes: { bardPrefix: ['b.', 'B.'] },
     pdfUrl: process.env.PDF_URL ?? '',
     processingTime: 300000
   },
@@ -96,10 +95,7 @@ export default {
     defaultRPC: 'https://api.harmony.one',
     restrictedPhrases: process.env.RESTRICTED_PHRASES
       ? process.env.RESTRICTED_PHRASES.split(', ')
-      : ['metamask', 'walletconnect'],
-    registerPrefix: process.env.COUNTRY_PREFIX
-      ? process.env.COUNTRY_PREFIX.split(',')
-      : ['+', '%']
+      : ['metamask', 'walletconnect']
   },
   voiceMemo: {
     isEnabled: Boolean(parseInt(process.env.VOICE_MEMO_ENABLED ?? '1')),

diff --git a/src/modules/1country/index.ts b/src/modules/1country/index.ts
@@ -11,30 +11,21 @@ import { type OnMessageContext, type OnCallBackQueryData, type PayableBot, Reque
 import { type BotPayments } from '../payment'
 import { getCommandNamePrompt, getUrl } from './utils/'
 import { isAdmin } from '../open-ai/utils/context'
-import config from '../../config'
 import { MAX_TRIES, sendMessage } from '../open-ai/helpers'
 import { sleep } from '../sd-images/utils'
 import { isValidUrl } from '../open-ai/utils/web-crawler'
 import { now } from '../../utils/perf'
 
-export const SupportedCommands = {
-  register: { name: 'rent' },
-  visit: { name: 'visit' },
-  check: { name: 'check' },
-  cert: { name: 'cert' },
-  nft: { name: 'nft' },
-  set: { name: 'set' }
+export enum SupportedCommands {
+  register = 'rent',
+  visit = 'visit',
+  check = 'check',
+  cert = 'cert',
+  nft = 'nft',
+  set = 'set'
 }
 
-// enum SupportedCommands {
-//   CHECK = "check",
-//   NFT = "nft",
-//   VISIT = "visit",
-//   CERT = "cert",
-//   RENEW = "renew",
-//   NOTION = "notion",
-//   SUBDOMAIN = "subdomain",
-// }
+const COUNTRY_PREFIX_LIST = ['+', '%']
 
 export class OneCountryBot implements PayableBot {
   public readonly module = 'OneCountryBot'
@@ -58,7 +49,7 @@ export class OneCountryBot implements PayableBot {
     ctx: OnMessageContext | OnCallBackQueryData
   ): boolean {
     const hasCommand = ctx.hasCommand(
-      Object.values(SupportedCommands).map((command) => command.name)
+      Object.values(SupportedCommands).map((command) => command)
     )
     const hasPrefix = this.hasPrefix(ctx.message?.text ?? '')
     if (hasPrefix && ctx.session.oneCountry.lastDomain) {
@@ -68,7 +59,7 @@ export class OneCountryBot implements PayableBot {
   }
 
   private hasPrefix (prompt: string): boolean {
-    const prefixList = config.country.registerPrefix
+    const prefixList = COUNTRY_PREFIX_LIST
     for (let i = 0; i < prefixList.length; i++) {
       if (prompt.toLocaleLowerCase().startsWith(prefixList[i])) {
         return true
@@ -88,17 +79,17 @@ export class OneCountryBot implements PayableBot {
       return
     }
 
-    if (ctx.hasCommand(SupportedCommands.visit.name)) {
+    if (ctx.hasCommand(SupportedCommands.visit)) {
       await this.onVistitCmd(ctx)
       return
     }
 
-    if (ctx.hasCommand(SupportedCommands.check.name)) {
+    if (ctx.hasCommand(SupportedCommands.check)) {
       await this.onCheckCmd(ctx)
       return
     }
 
-    if (ctx.hasCommand(SupportedCommands.register.name)) {
+    if (ctx.hasCommand(SupportedCommands.register)) {
       await this.onRegister(ctx)
       return
     }
@@ -108,17 +99,17 @@ export class OneCountryBot implements PayableBot {
       return
     }
 
-    if (ctx.hasCommand(SupportedCommands.nft.name)) {
+    if (ctx.hasCommand(SupportedCommands.nft)) {
       await this.onNftCmd(ctx)
       return
     }
 
-    if (ctx.hasCommand(SupportedCommands.cert.name)) {
+    if (ctx.hasCommand(SupportedCommands.cert)) {
       await this.onCertCmd(ctx)
       return
     }
 
-    if (ctx.hasCommand(SupportedCommands.set.name)) {
+    if (ctx.hasCommand(SupportedCommands.set)) {
       await this.onSet(ctx)
       return
     }

diff --git a/src/modules/llms/helpers.ts b/src/modules/llms/helpers.ts
@@ -1,4 +1,3 @@
-import config from '../../config'
 import {
   type OnMessageContext,
   type OnCallBackQueryData,
@@ -7,22 +6,22 @@ import {
   type ChatPayload
 } from '../types'
 import { type ParseMode } from 'grammy/types'
-// import { getChatModel, getChatModelPrice, getTokenNumber } from "./api/openAi";
 import { LlmsModelsEnum } from './types'
 import { type Message } from 'grammy/out/types'
 import { llmAddUrlDocument } from './api/llmApi'
 
-export const SupportedCommands = {
-  bardF: { name: 'bard' },
-  bard: { name: 'b' },
-  j2Ultra: { name: 'j2-ultra' },
-  sum: { name: 'sum' },
-  ctx: { name: 'ctx' },
-  pdf: { name: 'pdf' }
+export enum SupportedCommands {
+  bardF = 'bard',
+  bard = 'b',
+  j2Ultra = 'j2-ultra',
+  sum = 'sum',
+  ctx = 'ctx',
+  pdf = 'pdf'
 }
 
 export const MAX_TRIES = 3
 const LLAMA_PREFIX_LIST = ['*']
+const BARD_PREFIX_LIST = ['b.', 'B.']
 
 export const isMentioned = (
   ctx: OnMessageContext | OnCallBackQueryData
@@ -51,7 +50,7 @@ export const hasLlamaPrefix = (prompt: string): string => {
 }
 
 export const hasBardPrefix = (prompt: string): string => {
-  const prefixList = config.llms.prefixes.bardPrefix
+  const prefixList = BARD_PREFIX_LIST
   for (let i = 0; i < prefixList.length; i++) {
     if (prompt.toLocaleLowerCase().startsWith(prefixList[i])) {
       return prefixList[i]

diff --git a/src/modules/llms/index.ts b/src/modules/llms/index.ts
@@ -63,7 +63,7 @@ export class LlmsBot implements PayableBot {
     ctx: OnMessageContext | OnCallBackQueryData
   ): boolean {
     const hasCommand = ctx.hasCommand(
-      Object.values(SupportedCommands).map((command) => command.name)
+      Object.values(SupportedCommands).map((command) => command)
     )
     if (isMentioned(ctx)) {
       return true
@@ -117,12 +117,12 @@ export class LlmsBot implements PayableBot {
       return
     }
 
-    if (ctx.hasCommand(SupportedCommands.pdf.name)) {
+    if (ctx.hasCommand(SupportedCommands.pdf)) {
       await this.onPdfCommand(ctx)
       return
     }
 
-    if (ctx.hasCommand(SupportedCommands.bard.name) || ctx.hasCommand(SupportedCommands.bardF.name)) {
+    if (ctx.hasCommand(SupportedCommands.bard) || ctx.hasCommand(SupportedCommands.bardF)) {
       await this.onChat(ctx, LlmsModelsEnum.BISON)
       return
     }
@@ -142,17 +142,17 @@ export class LlmsBot implements PayableBot {
       return
     }
 
-    if (ctx.hasCommand(SupportedCommands.j2Ultra.name)) {
+    if (ctx.hasCommand(SupportedCommands.j2Ultra)) {
       await this.onChat(ctx, LlmsModelsEnum.J2_ULTRA)
       return
     }
 
-    if (ctx.hasCommand(SupportedCommands.ctx.name)) {
+    if (ctx.hasCommand(SupportedCommands.ctx)) {
       await this.onCurrentCollection(ctx)
       return
     }
 
-    if (ctx.hasCommand(SupportedCommands.sum.name) ||
+    if (ctx.hasCommand(SupportedCommands.sum) ||
       (ctx.message?.text?.startsWith('sum ') && ctx.chat?.type === 'private')
     ) {
       await this.onSum(ctx)