diff --git a/src/config.ts b/src/config.ts index 85868db..e6caf37 100644 --- 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 index e4e80bc..cfd8617 100644 --- 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 index 06682a3..81d2540 100644 --- a/src/modules/llms/api/llmApi.ts +++ b/src/modules/llms/api/llmApi.ts @@ -42,13 +42,22 @@ export const llmAddUrlDocument = async (args: LlmAddUrlDocument): Promise => { +interface LlmCheckCollectionStatusOutput { + price: number + status: 'PROCESSING' | 'DONE' + error: 'INVALID_PDF' | undefined +} +export const llmCheckCollectionStatus = async (name: string): Promise => { 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 index 1d4a6b4..85d829c 100644 --- 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 index e08b92e..357a586 100644 --- 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 { + 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 index 079a4c2..e9037a4 100644 --- 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 {