Evaluate Orama DB

Author: ggodlewski

package-lock.json

@@ -67,9 +67,9 @@
   ],
   "dependencies": {
     "@nyariv/sandboxjs": "0.8.23",
-    "@opentelemetry/core": "1.28.0",
     "@opentelemetry/api": "1.3.0",
     "@opentelemetry/context-zone": "1.8.0",
+    "@opentelemetry/core": "1.28.0",
     "@opentelemetry/exporter-zipkin": "1.8.0",
     "@opentelemetry/instrumentation": "0.33.0",
     "@opentelemetry/instrumentation-fetch": "0.34.0",
@@ -79,6 +79,9 @@
     "@opentelemetry/sdk-trace-node": "1.8.0",
     "@opentelemetry/sdk-trace-web": "1.8.0",
     "@opentelemetry/semantic-conventions": "1.8.0",
+    "@orama/orama": "3.1.1",
+    "@tensorflow/tfjs-node": "4.22.0",
+    "@tensorflow-models/universal-sentence-encoder": "1.3.3",
     "async": "3.2.3",
     "buffer": "^5.7.1",
     "casual": "1.6.2",
@@ -94,8 +97,8 @@
     "express": "4.21.1",
     "express-jwt": "8.2.1",
     "express-rate-limit": "^6.6.0",
-    "htmlparser2": "9.0.0",
     "file-isignature": "1.0.3",
+    "htmlparser2": "9.0.0",
     "js-yaml": "4.1.0",
     "jsonwebtoken": "9.0.2",
     "jszip": "3.10.1",
@@ -105,6 +108,7 @@
     "marked": "9.0.2",
     "mathjs": "10.5.0",
     "mathml-to-latex": "1.4.0",
+    "@msgpack/msgpack": "3.1.0",
     "minimist": "1.2.6",
     "mitt": "^3.0.0",
     "open": "^7.4.2",

package.json

@@ -18,6 +18,7 @@ import {FolderRegistryContainer} from '../folder_registry/FolderRegistryContaine
 import {ActionRunnerContainer, convertActionYaml} from '../action/ActionRunnerContainer.ts';
 import {getContentFileService} from '../transform/utils.ts';
 import {UploadContainer} from '../google_folder/UploadContainer.ts';
+import {createIndexer} from '../search/Indexer.ts';
 
 const __filename = import.meta.filename;
 
@@ -747,13 +748,15 @@ export class JobManagerContainer extends Container {
       const contentFileService = await getContentFileService(transformedFileSystem, userConfigService);
       const markdownTreeProcessor = new MarkdownTreeProcessor(contentFileService);
 
+      const indexer = await createIndexer();
+
       switch (type) {
         case 'local':
           await gitScanner.resetToLocal({
             privateKeyFile: await userConfigService.getDeployPrivateKeyPath()
           });
 
-          await markdownTreeProcessor.regenerateTree(driveId);
+          await markdownTreeProcessor.regenerateTree(driveId, indexer);
           await markdownTreeProcessor.save();
           break;
         case 'remote':
@@ -762,12 +765,15 @@ export class JobManagerContainer extends Container {
               privateKeyFile: await userConfigService.getDeployPrivateKeyPath()
             });
 
-            await markdownTreeProcessor.regenerateTree(driveId);
+            await markdownTreeProcessor.regenerateTree(driveId, indexer);
             await markdownTreeProcessor.save();
           }
           break;
       }
 
+      await transformedFileSystem.mkdir('/.private');
+      await transformedFileSystem.writeBuffer('/.private/' + indexer.getFileName(), await indexer.getData());
+
       await this.schedule(driveId, {
         ...initJob(),
         type: 'run_action',

src/containers/job/JobManagerContainer.ts

@@ -0,0 +1,25 @@
+// import {LunrIndexer} from './LunrIndexer.ts';
+import {OramaIndexer} from './OramaIndexer.ts';
+
+export interface SearchResults {
+  result: any[];
+}
+
+export interface PageToIndex {
+  content: string;
+  id: string;
+  title: string;
+  path: string;
+}
+
+export interface Indexer {
+  addPage(page: PageToIndex): Promise<void>;
+  getData(): Promise<Uint8Array>;
+  setData(data: Uint8Array): Promise<void>;
+  search(term: string): Promise<SearchResults>;
+  getFileName(): string;
+}
+
+export function createIndexer() {
+  return OramaIndexer.init();
+}

src/containers/search/Indexer.ts

@@ -1,26 +1,34 @@
 import lunr from 'lunr';
 import stemmerSupport from 'lunr-languages/lunr.stemmer.support.js';
+import { Indexer, PageToIndex } from './Indexer.ts';
 
 stemmerSupport(lunr);
 
 lunr.tokenizer.separator = /[\s-_@#,.]+/;
 
-export class LunrIndexer {
+export class LunrIndexer implements Indexer {
   private lunrBuilder: lunr.Builder;
-  private store;
+  private store: Record<string, any>;
+  lunrIndex: lunr.Index;
 
   constructor() {
     this.lunrBuilder = new lunr.Builder();
     this.lunrBuilder.ref('path');
-    this.lunrBuilder.field('id2', { extractor: (doc) => {
+    this.lunrBuilder.field('id2', {
+      extractor: (doc) => {
         return doc['id'] ? doc['id'].replace(/[_-]*/g, '') : undefined;
-    }});
+      }
+    });
     this.lunrBuilder.field('id');
     this.lunrBuilder.field('title');
     this.store = {};
   }
 
-  async addPage(page) {
+  static async init() {
+    return new LunrIndexer();
+  }
+
+  async addPage(page: PageToIndex) {
     this.lunrBuilder.add({
       path: page.path,
       title: page.title,
@@ -33,9 +41,56 @@ export class LunrIndexer {
     };
   }
 
-  getJson() {
+  async getData() {
     const lunrIndex = this.lunrBuilder.build();
-    return { index: lunrIndex.toJSON(), store: this.store };
+    const str = JSON.stringify({ index: lunrIndex.toJSON(), store: this.store });
+    return new TextEncoder().encode(str);
+  }
+
+  async setData(data: Uint8Array) {
+    this.lunrIndex = undefined;
+    this.store = {};
+
+    try {
+      const lunrData = JSON.parse(new TextDecoder().decode(data));
+      this.store = lunrData.store || {};
+      if (lunrData?.index) {
+        this.lunrIndex = lunr.Index.load(lunrData.index);
+      }
+    // deno-lint-ignore no-unused-vars
+    } catch (err) {
+      this.store = {};
+    }
   }
 
+  async search(queryParam: string) {
+    if (!this.lunrIndex) {
+      return {
+        result: []
+      };
+    }
+
+    queryParam = (queryParam || '').trim().replace(/:/g, ' ');
+
+    let result = this.lunrIndex.search(queryParam);
+    if (result.length === 0 && queryParam.indexOf('*') === -1) {
+      result = this.lunrIndex.search(queryParam.split(/\s+/g).map(w => w.length > 2 ? w + '*' : w).join(' '));
+    }
+    if (result.length === 0 && queryParam.replace(/[_-]*/g, '').length > 10) {
+      result = this.lunrIndex.search(queryParam.replace(/[_-]*/g, ''));
+    }
+
+    return {
+      result: result.map((doc) => ({
+        path: doc.ref,
+        score: doc.score,
+        matchData: doc.matchData,
+        ...this.store[doc.ref]
+      }))
+    };
+  }
+
+  getFileName() {
+    return 'lonr.json';
+  }
 }

src/containers/search/LunrIndexer.ts

@@ -0,0 +1,75 @@
+import { create, load, save, insert, search } from '@orama/orama';
+import type { Orama, RawData } from '@orama/orama';
+import msgpack from '@msgpack/msgpack';
+import '@tensorflow/tfjs-node'; // Or any other appropriate TensorflowJS backend, like @tensorflow/tfjs-backend-webgl
+
+import { pluginEmbeddings } from './pluginEmbeddings.ts';
+
+import {Indexer, SearchResults, PageToIndex} from './Indexer.ts';
+
+const plugin = await pluginEmbeddings({
+  embeddings: {
+    defaultProperty: 'embeddings',
+    onInsert: {
+      generate: true,
+      properties: ['title', 'content']
+    }
+  }
+});
+
+export class OramaIndexer implements Indexer {
+  private constructor(private db: Orama<unknown, unknown, unknown, unknown>) {
+  }
+
+  static async init() {
+    const db = create({
+      schema: {
+        title: 'string',
+        id: 'string',
+        path: 'string',
+        embeddings: 'vector[512]',
+      },
+      plugins: [plugin]
+    });
+
+    return new OramaIndexer(db);
+  }
+
+  async addPage(page: PageToIndex): Promise<void> {
+    insert(this.db, {
+      title: page.title,
+      id: page.id,
+      path: page.path,
+      content: page.content
+    });
+  }
+
+  async search(term: string): Promise<SearchResults> {
+    const results = await search(this.db, {
+      term,
+      mode: 'vector',
+      similarity: 0.5,
+    });
+
+    return {
+      result: results.hits.map(h => ({
+        score: h.score,
+        id: h?.document?.id,
+        path: h?.document?.path,
+        title: h?.document?.title,
+      }))
+    };
+  }
+
+  async setData(data: Uint8Array): Promise<void> {
+    load(this.db, <RawData>msgpack.decode(data));
+  }
+
+  async getData(): Promise<Uint8Array> {
+    return msgpack.encode(save(this.db));
+  }
+
+  getFileName() {
+    return 'orama.msgpack';
+  }
+}

src/containers/search/OramaIndexer.ts

@@ -0,0 +1,89 @@
+import type { AnyOrama, SearchParams, TypedDocument, OramaPluginAsync, PartialSchemaDeep } from '@orama/orama'
+import { load as loadModel } from '@tensorflow-models/universal-sentence-encoder'
+
+export type PluginEmbeddingsParams = {
+  embeddings: {
+    defaultProperty: string
+    onInsert?: {
+      generate: boolean
+      properties: string[]
+      verbose?: boolean
+    }
+  }
+}
+
+function getPropertyValue (obj: object, path: string) {
+  return path.split('.').reduce((current, key) =>
+    current && current[key] !== undefined ? current[key] : undefined, obj
+  );
+}
+
+function getPropertiesValues(schema: object, properties: string[]) {
+  return properties
+    .map(prop => getPropertyValue(schema, prop))
+    .filter(value => value !== undefined)
+    .join('. ');
+}
+
+function normalizeVector(v: number[]): number[] {
+  const norm = Math.sqrt(v.reduce((sum, val) => sum + val * val, 0));
+  return v.map(val => val / norm);
+}
+
+export const embeddingsType = 'vector[512]';
+
+export async function pluginEmbeddings(pluginParams: PluginEmbeddingsParams): Promise<OramaPluginAsync> {
+  const model = await loadModel();
+
+  return {
+    name: 'orama-plugin-embeddings',
+
+    async beforeInsert<T extends TypedDocument<any>>(_db: AnyOrama, _id: string, params: PartialSchemaDeep<T>) {
+      if (!pluginParams.embeddings?.onInsert?.generate) {
+        return;
+      }
+
+      if (!pluginParams.embeddings?.onInsert?.properties) {
+        throw new Error('Missing "embeddingsConfig.properties" parameter for plugin-secure-proxy');
+      }
+
+      const properties = pluginParams.embeddings.onInsert.properties;
+      const values = getPropertiesValues(params, properties);
+
+      if (pluginParams.embeddings.onInsert.verbose) {
+        console.log(`Generating embeddings for properties "${properties.join(', ')}": "${values}"`);
+      }
+
+      const embeddings = Array.from(await (await model.embed(values)).data());
+
+      params[pluginParams.embeddings.defaultProperty] = normalizeVector(embeddings);
+    },
+
+    async beforeSearch<T extends AnyOrama>(_db: AnyOrama, params: SearchParams<T, TypedDocument<any>>) {
+      if (params.mode !== 'vector' && params.mode !== 'hybrid') {
+        return;
+      }
+
+      if (params?.vector?.value) {
+        return;
+      }
+
+      if (!params.term) {
+        throw new Error('No "term" or "vector" parameters were provided');
+      }
+
+      const embeddings = Array.from(await (await model.embed(params.term)).data()) as unknown as number[];
+
+      if (!params.vector) {
+        params.vector = {
+          // eslint-disable-next-line
+          // @ts-ignore
+          property: params?.vector?.property ?? pluginParams.embeddings.defaultProperty,
+          value: normalizeVector(embeddings)
+        };
+      }
+
+      params.vector.value = normalizeVector(embeddings);
+    }
+  };
+}