Merge pull request #323 from harmony-one/vertex-bard-integration

collection session object + pdf collection upload

src/bot.ts

@@ -187,6 +187,11 @@ function createInitialSessionData (): BotSessionData {
       languages: [],
       enable: false
     },
+    collections: {
+      activeCollections: [],
+      collectionRequestQueue: [],
+      isProcessingQueue: false
+    },
     llms: {
       model: config.llms.model,
       isEnabled: config.llms.isEnabled,

src/modules/document-handler/index.ts

@@ -1,6 +1,7 @@
 import { type OnMessageContext, type PayableBot, type RefundCallback, RequestState } from '../types'
 import * as Sentry from '@sentry/node'
 import { now } from '../../utils/perf'
+import { llmAddUrlDocument } from '../llms/api/llmApi'
 
 const SupportedDocuments = { PDF: 'application/pdf' }
 
@@ -14,12 +15,16 @@ export class DocumentHandler implements PayableBot {
     ctx.transient.analytics.module = this.module
     try {
       const file = await ctx.getFile()
+      const documentType = ctx.message.document?.mime_type
+      if (documentType === 'application/pdf' && ctx.chat.id) {
+        const pdfUrl = file.getUrl()
+        const fileName = ctx.message.document?.file_name ?? file.file_id
+        await this.addDocToCollection(ctx, ctx.chat.id, fileName, pdfUrl)
+      }
       console.log(file)
-      await ctx.reply('you did it kid')
       ctx.transient.analytics.sessionState = RequestState.Success
     } catch (ex) {
       Sentry.captureException(ex)
-      await ctx.reply('you failed kid')
       ctx.transient.analytics.sessionState = RequestState.Error
     } finally {
       ctx.transient.analytics.actualResponseTime = now()
@@ -32,7 +37,21 @@ export class DocumentHandler implements PayableBot {
     if (documentType !== undefined) {
       return Object.values(SupportedDocuments).includes(documentType)
     }
-
     return false
   }
+
+  private async addDocToCollection (ctx: OnMessageContext, chatId: number, fileName: string, pdfUrl: string): Promise<void> {
+    const collectionName = await llmAddUrlDocument({
+      chatId,
+      pdfUrl,
+      fileName
+    })
+    ctx.session.collections.collectionRequestQueue.push({
+      collectionName,
+      collectionType: 'PDF',
+      fileName,
+      url: pdfUrl
+    })
+    ctx.session.collections.isProcessingQueue = true
+  }
 }