feat: adds plugin embeddings

packages/plugin-embeddings/LICENSE.md

@@ -0,0 +1,13 @@
+Copyright 2024 OramaSearch Inc
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.

packages/plugin-embeddings/README.md

@@ -0,0 +1,95 @@
+# Orama Plugin Embeddings
+
+**Orama Plugin Embeddings** allows you to generate fast text embeddings at insert and search time offline, directly on your machine - no OpenAI needed!
+
+## Installation
+
+To get started with **Orama Plugin Embeddings**, just install it with npm:
+
+```sh
+npm i @orama/plugin-embeddings
+```
+
+**Important note**: to use this plugin, you'll also need to install one of the following TensorflowJS backend:
+
+- `@tensorflow/tfjs`
+- `@tensorflow/tfjs-node`
+- `@tensorflow/tfjs-backend-webgl`
+- `@tensorflow/tfjs-backend-cpu`
+- `@tensorflow/tfjs-backend-wasm`
+
+For example, if you're running Orama on the browser, we highly recommend using `@tensorflow/tfjs-backend-webgl`:
+
+```sh
+npm i @tensorflow/tfjs-backend-webgl
+```
+
+If you're using Orama in Node.js, we recommend using `@tensorflow/tfjs-node`:
+
+```sh
+npm i @tensorflow/tfjs-node
+```
+
+## Usage
+
+```js
+import { create } from '@orama/orama'
+import { pluginEmbeddings } from '@orama/plugin-embeddings'
+import '@tensorflow/tfjs-node' // Or any other appropriate TensorflowJS backend
+
+const plugin = await pluginEmbeddings({
+  embeddings: {
+    defaultProperty: 'embeddings', // Property used to store generated embeddings
+    onInsert: {
+      generate: true, // Generate embeddings at insert-time
+      properties: ['description'], // properties to use for generating embeddings at insert time
+      verbose: true,
+    }
+  }
+})
+
+const db = await create({
+  schema: {
+    description: 'string',
+    embeddings: 'vector[512]' // Orama generates 512-dimensions vectors
+  },
+  plugins: [plugin]
+})
+```
+
+Example usage at insert time:
+
+```js
+await insert(db, {
+  description: 'Classroom Headphones Bulk 5 Pack, Student On Ear Color Varieties'
+})
+
+await insert(db, {
+  description: 'Kids Wired Headphones for School Students K-12'
+})
+
+await insert(db, {
+  description: 'Kids Headphones Bulk 5-Pack for K-12 School'
+})
+
+await insert(db, {
+  description: 'Bose QuietComfort Bluetooth Headphones'
+})
+```
+
+Orama will automatically generate text embeddings and store them into the `embeddings` property.
+
+Then, you can use the `vector` or `hybrid` setting to perform hybrid or vector search at runtime:
+
+```js
+await search(db, {
+  term: 'Headphones for 12th grade students',
+  mode: 'vector'
+})
+```
+
+Orama will generate embeddings at search time and perform vector or hybrid search for you.
+
+# License
+
+[Apache 2.0](/LICENSE.md)

packages/plugin-embeddings/package.json

@@ -0,0 +1,55 @@
+{
+  "name": "@orama/plugin-embeddings",
+  "version": "3.0.0",
+  "description": "Orama plugin for generating embeddings locally",
+  "keywords": [
+    "orama",
+    "embeddings",
+    "secure proxy",
+    "vector search"
+  ],
+  "license": "Apache-2.0",
+  "main": "./dist/index.js",
+  "type": "module",
+  "exports": {
+    ".": {
+      "require": "./dist/index.cjs",
+      "import": "./dist/index.js",
+      "types": "./dist/index.d.ts",
+      "browser": "./dist/index.global.js"
+    }
+  },
+  "bugs": {
+    "url": "https://github.com/askorama/orama/issues"
+  },
+  "homepage": "https://github.com/askorama/orama#readme",
+  "repository": {
+    "type": "git",
+    "url": "git+https://github.com/askorama/orama.git"
+  },
+  "sideEffects": false,
+  "types": "./dist/index.d.ts",
+  "files": [
+    "dist"
+  ],
+  "scripts": {
+    "build": "tsup --config tsup.lib.js",
+    "lint": "exit 0",
+    "test": "exit 0"
+  },
+  "publishConfig": {
+    "access": "public"
+  },
+  "devDependencies": {
+    "@orama/orama": "workspace:*",
+    "@types/node": "^20.9.0",
+    "tsup": "^7.2.0",
+    "typescript": "^5.0.0"
+  },
+  "dependencies": {
+    "@tensorflow-models/universal-sentence-encoder": "^1.3.3",
+    "@tensorflow/tfjs": "^4.21.0",
+    "@tensorflow/tfjs-backend-cpu": "^4.21.0",
+    "@tensorflow/tfjs-node": "^4.21.0"
+  }
+}

packages/plugin-embeddings/src/index.ts

@@ -0,0 +1,84 @@
+import type { AnyOrama, SearchParams, TypedDocument, OramaPluginAsync, PartialSchemaDeep } from '@orama/orama'
+import use 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('. ')
+}
+
+export const embeddingsType = 'vector[512]'
+
+export async function pluginEmbeddings(pluginParams: PluginEmbeddingsParams): OramaPluginAsync {
+  const model = await use.load()
+
+  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 = await model.embed(values)
+
+      params[pluginParams.embeddings.defaultProperty] = await embeddings.data()
+    },
+
+    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('Neither "term" nor "vector" parameters were provided')
+      }
+
+      const embeddings = await model.embed(params.term)
+
+      if (!params.vector) {
+        params.vector = {
+          // eslint-disable-next-line
+          // @ts-ignore
+          property: params?.vector?.property ?? pluginParams.embeddings.defaultProperty,
+          value: embeddings
+        }
+      }
+
+      params.vector.value = embeddings
+    }
+  }
+}