Merge pull request #41 from Klimatbyran/feat/elasticsearch-client

Feat: Elasticsearch client

k8s/deployment.yaml

@@ -45,6 +45,16 @@ spec:
                 secretKeyRef:
                   name: env
                   key: OPENAI_ORG_ID
+            - name: ELASTIC_NODE_URL
+              valueFrom:
+                secretKeyRef:
+                  name: env
+                  key: ELASTIC_NODE_URL
+            - name: ELASTIC_INDEX_NAME
+              valueFrom:
+                secretKeyRef:
+                  name: env
+                  key: ELASTIC_INDEX_NAME
             # - name: CHROMA_TOKEN
             #   valueFrom:
             #     secretKeyRef:

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,124 @@
+import config from './config/elasticsearch';
+import { Client } from '@elastic/elasticsearch';
+import * as crypto from 'crypto';
+
+class Elastic {
+  client: Client;
+  indexName: string;
+
+  constructor({ node, indexName }) {
+    try  {
+      this.client = new Client({ node });
+      this.indexName = indexName;
+    } catch (error) {
+      console.error('Elasticsearch constructor error:', error);
+    }
+  }
+
+  // Hash the URL to use as the document ID
+  private hashUrl(url: string): string {
+    return crypto.createHash('sha256').update(url).digest('hex');
+  }
+
+  async setupIndex() {
+    try {
+      console.log(`Checking if index ${this.indexName} exists...`);
+      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: 'text' },
+                state: { type: 'keyword' }
+              }
+            }
+          }
+      });
+      console.log(`Index ${this.indexName} created.`);
+    } else {
+      console.log(`Index ${this.indexName} already exists.`);
+    }
+    } catch (error) {
+      console.error('Elasticsearch setupIndex error:', error);
+    }
+  }
+
+  async createEntryWithUrl(url: string) {
+    const documentId = this.hashUrl(url);
+    const docBody = {
+      url: url,
+      state: 'pending'
+    };
+
+    try {
+      await this.client.index({
+        index: this.indexName,
+        id: documentId,
+        body: docBody
+      });
+      console.log(`Entry created with PDF URL. Document ID: ${documentId}`);
+    } catch (error) {
+      console.error(`Error creating entry with PDF URL for Document ID ${documentId}:`, error);
+    }
+  }
+
+  async addPDFContent(url: string, pdfContent: string) {
+    const documentId = this.hashUrl(url);
+    try {
+      await this.client.update({
+        index: this.indexName,
+        id: documentId,
+        body: {
+          doc: {
+            pdf: pdfContent
+          }
+        }
+      });
+      console.log(`PDF content added. Document ID: ${documentId}`);
+    } catch (error) {
+      console.error(`Error adding PDF content for Document ID ${documentId}:`, error);
+    }
+  }
+
+  async addReportData(url: string, reportData: object) {
+    const documentId = this.hashUrl(url);
+    try {
+      await this.client.update({
+        index: this.indexName,
+        id: documentId,
+        body: {
+          doc: {
+            report: reportData
+          }
+        }
+      });
+      console.log(`Report data added. Document ID: ${documentId}`);
+    } catch (error) {
+      console.error(`Error adding report data for Document ID ${documentId}:`, error);
+    }
+  }
+
+  async updateDocumentState(url: string, newState: string) {
+    const documentId = this.hashUrl(url);
+    try {
+      await this.client.update({
+        index: this.indexName,
+        id: documentId,
+        body: {
+          doc: {
+            state: newState
+          }
+        }
+      });
+      console.log(`Document ${documentId} state updated to ${newState}.`);
+    } catch (error) {
+      console.error(`Error updating document state for Document ID ${documentId}:`, error);
+    }
+  }
+}
+
+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'
@@ -57,6 +58,7 @@ createBullBoard({
 
 const app = express()
 discord.login()
+elastic.setupIndex()
 
 app.use('/admin/queues', serverAdapter.getRouter())
 app.listen(3000, () => {