Merge pull request #338 from harmony-one/collection-handler

add web crawler processingTime limit + add invalid PDF format respons…
harmony-one · Oct 12, 2023 · f3fcad0 · f3fcad0
2 parents 487837c + d56c166
commit f3fcad0
Show file tree

Hide file tree

Showing 6 changed files with 56 additions and 16 deletions.
diff --git a/src/config.ts b/src/config.ts
@@ -39,7 +39,8 @@ export default {
     minimumBalance: 0,
     isEnabled: Boolean(parseInt(process.env.LLMS_ENABLED ?? '1')),
     prefixes: { bardPrefix: [',', 'b.', 'B.'] },
-    pdfUrl: process.env.PDF_URL ?? ''
+    pdfUrl: process.env.PDF_URL ?? '',
+    processingTime: 300000
   },
   openAi: {
     dalle: {

diff --git a/src/modules/document-handler/index.ts b/src/modules/document-handler/index.ts
@@ -70,7 +70,8 @@ export class DocumentHandler implements PayableBot {
         collectionType: 'PDF',
         fileName,
         url,
-        prompt
+        prompt,
+        processingTime: 0
       })
     } catch (e: any) {
       await this.onError(ctx, e)

diff --git a/src/modules/llms/api/llmApi.ts b/src/modules/llms/api/llmApi.ts
@@ -3,7 +3,7 @@ import config from '../../../config'
 import { type ChatConversation } from '../../types'
 import pino from 'pino'
 
-const API_ENDPOINT = config.llms.apiEndpoint // 'http://localhost:8080' // http://127.0.0.1:5000' // config.llms.apiEndpoint
+const API_ENDPOINT = config.llms.apiEndpoint // config.llms.apiEndpoint // 'http://localhost:8080' // http://127.0.0.1:5000' // config.llms.apiEndpoint
 
 const logger = pino({
   name: 'llmApi',
@@ -42,13 +42,22 @@ export const llmAddUrlDocument = async (args: LlmAddUrlDocument): Promise<string
   return ''
 }
 
-export const llmCheckCollectionStatus = async (name: string): Promise<number> => {
+interface LlmCheckCollectionStatusOutput {
+  price: number
+  status: 'PROCESSING' | 'DONE'
+  error: 'INVALID_COLLECTION' | undefined
+}
+export const llmCheckCollectionStatus = async (name: string): Promise<LlmCheckCollectionStatusOutput> => {
   const endpointUrl = `${API_ENDPOINT}/collections/document/${name}` // ?collectionName=${collectionName}`
   const response = await axios.get(endpointUrl)
   if (response) {
-    return response.data.price
+    return response.data
+  }
+  return {
+    price: -1,
+    status: 'PROCESSING',
+    error: undefined
   }
-  return -1
 }
 
 interface QueryUrlDocumentOutput {

diff --git a/src/modules/llms/helpers.ts b/src/modules/llms/helpers.ts
@@ -255,7 +255,8 @@ export async function addUrlToCollection (ctx: OnMessageContext | OnCallBackQuer
     collectionType: 'URL',
     url: url.toLocaleLowerCase(),
     prompt,
-    msgId
+    msgId,
+    processingTime: 0
   })
 }
 
@@ -276,6 +277,7 @@ export async function addDocToCollection (ctx: OnMessageContext | OnCallBackQuer
     fileName,
     url: url.toLocaleLowerCase(),
     prompt,
-    msgId
+    msgId,
+    processingTime: 0
   })
 }
diff --git a/src/modules/llms/index.ts b/src/modules/llms/index.ts
@@ -349,20 +349,21 @@ export class LlmsBot implements PayableBot {
   }
 
   async onCheckCollectionStatus (ctx: OnMessageContext | OnCallBackQueryData): Promise<void> {
+    const processingTime = config.llms.processingTime
     while (ctx.session.collections.collectionRequestQueue.length > 0) {
       try {
         const collection = ctx.session.collections.collectionRequestQueue.shift()
         if (collection) {
-          const price = await llmCheckCollectionStatus(collection?.collectionName ?? '')
-          if (price > 0) {
+          const result = await llmCheckCollectionStatus(collection?.collectionName ?? '')
+          if (result.price > 0) {
             if (
-              !(await this.payments.pay(ctx as OnMessageContext, price))
+              !(await this.payments.pay(ctx as OnMessageContext, result.price))
             ) {
               await this.onNotBalanceMessage(ctx)
             } else {
               ctx.session.collections.activeCollections.push(collection)
               if (collection.msgId) {
-                const oneFee = await this.payments.getPriceInONE(price)
+                const oneFee = await this.payments.getPriceInONE(result.price)
                 let statusMsg
                 if (collection.collectionType === 'URL') {
                   statusMsg = `${collection.url} processed (${this.payments.toONE(oneFee, false).toFixed(2)} ONE fee)`
@@ -380,12 +381,37 @@ export class LlmsBot implements PayableBot {
               await this.queryUrlCollection(ctx, collection.url ?? '',
                 collection.prompt ?? 'summary')
             }
+          } else if (result.price < 0) {
+            if (collection.msgId) {
+              let statusMsg = ''
+              if (collection.collectionType === 'URL') {
+                statusMsg = `${collection.url} - Invalid URL`
+              } else {
+                statusMsg = `${collection.fileName} - Invalid PDF format`
+              }
+              await ctx.api.editMessageText(ctx.chat?.id ?? '', collection.msgId, statusMsg,
+                { disable_web_page_preview: true })
+            }
           } else {
-            ctx.session.collections.collectionRequestQueue.push(collection)
-            if (ctx.session.collections.collectionRequestQueue.length === 1) {
-              await sleep(5000)
+            if (collection.processingTime && collection.processingTime > processingTime) { // 5 min max
+              if (collection.msgId) {
+                let statusMsg = ''
+                if (collection.collectionType === 'URL') {
+                  statusMsg = `${collection.url} - Processing time limit reached. Please check the file format and try again`
+                } else {
+                  statusMsg = `${collection.fileName} - Processing time limit reached. Please check the file format and try again`
+                }
+                await ctx.api.editMessageText(ctx.chat?.id ?? '', collection.msgId, statusMsg,
+                  { disable_web_page_preview: true })
+              }
             } else {
-              await sleep(2500)
+              const processingTime = collection.processingTime ? collection.processingTime + 5000 : 5000
+              ctx.session.collections.collectionRequestQueue.push({ ...collection, processingTime })
+              if (ctx.session.collections.collectionRequestQueue.length === 1) {
+                await sleep(6000)
+              } else {
+                await sleep(3000)
+              }
             }
           }
         }

diff --git a/src/modules/types.ts b/src/modules/types.ts
@@ -90,6 +90,7 @@ export interface Collection {
   fileName?: string
   prompt?: string
   msgId?: number
+  processingTime?: number // milliseconds
 }
 
 export interface FileDoc {