Merge pull request #36 from Klimatbyran/feat/save-responses-to-db

Add elastic client

package.json

@@ -13,6 +13,7 @@
   "dependencies": {
     "@bull-board/api": "^5.12.0",
     "@bull-board/express": "^5.12.0",
+    "@elastic/elasticsearch": "^8.12.1",
     "@google/generative-ai": "^0.1.3",
     "bullmq": "^5.1.1",
     "chromadb": "^1.7.3",

src/config/elasticsearch.ts

@@ -0,0 +1,4 @@
+export default {
+    node: process.env.ELASTIC_NODE_URL, 
+    indexName: process.env.ELASTIC_INDEX_NAME,
+}

src/elastic.ts

@@ -0,0 +1,66 @@
+import config from './config/elasticsearch';
+import { Client } from '@elastic/elasticsearch';
+
+class Elastic {
+  client: Client;
+  indexName: string;
+
+  constructor({ node, indexName }) {
+    this.client = new Client({ node });
+    this.indexName = indexName;
+  }
+
+  async setupIndex() {
+    const indexExists = await this.client.indices.exists({ index: this.indexName });
+    if (!indexExists) {
+      await this.client.indices.create({
+        index: this.indexName,
+        body: {
+          mappings: {
+            properties: {
+              url: { type: 'keyword' },
+              pdf: { type: 'binary' },
+              report: { type: 'object' },
+              state: { type: 'keyword' }
+            }
+          }
+        }
+      });
+      console.log(`Index ${this.indexName} created.`);
+    } else {
+      console.log(`Index ${this.indexName} already exists.`);
+    }
+  }
+
+  async indexDocument(documentId: string, pdfContent: string, reportData: object = null) {
+    const docBody = {
+      pdf: pdfContent,
+      report: reportData,
+      state: 'pending'
+    };
+
+    await this.client.index({
+      index: this.indexName,
+      id: documentId,
+      body: docBody
+    });
+
+    console.log(`Document ${documentId} indexed.`);
+  }
+
+  async updateDocumentState(documentId: string, newState: string) {
+    await this.client.update({
+      index: this.indexName,
+      id: documentId,
+      body: {
+        doc: {
+          state: newState
+        }
+      }
+    });
+
+    console.log(`Document ${documentId} state updated to ${newState}.`);
+  }
+}
+
+export default new Elastic(config)

src/index.ts

@@ -7,6 +7,7 @@ import { BullMQAdapter } from '@bull-board/api/bullMQAdapter'
 import { ExpressAdapter } from '@bull-board/express'
 
 import discord from './discord'
+import elastic from './elastic'
 
 // keep this line, otherwise the workers won't be started
 import * as workers from './workers'
@@ -56,6 +57,8 @@ createBullBoard({
 })
 
 const app = express()
+
+elastic.setupIndex()
 discord.login()
 
 app.use('/admin/queues', serverAdapter.getRouter())

src/workers/discordReview.ts

@@ -10,14 +10,19 @@ import {
   ModalActionRowComponentBuilder,
   TextInputStyle,
 } from 'discord.js'
+import { Client } from '@elastic/elasticsearch';
+
+const esClient = new Client({ 
+  node: 'http://elasticsearch.data-pipeline.svc.cluster.local:9200' 
+});
 
 class JobData extends Job {
   data: {
-    url: string
-    json: string
-    channelId: string
-    messageId: string
-  }
+    url: string;
+    json: string;
+    channelId: string;
+    messageId: string;
+  };
 }
 
 const worker = new Worker(

src/workers/downloadPDF.ts

@@ -4,6 +4,11 @@ import pdf from 'pdf-parse'
 import { splitText } from '../queues'
 import discord from '../discord'
 import { TextChannel } from 'discord.js'
+import { Client } from '@elastic/elasticsearch';
+
+const esClient = new Client({ 
+  node: 'http://elasticsearch.data-pipeline.svc.cluster.local:9200' 
+});
 
 class JobData extends Job {
   data: {