@instant.dev/vectors 0.1.0

package/README.md

@@ -0,0 +1,117 @@
+# Simple vector creation with automatic batching
+![npm version](https://img.shields.io/npm/v/@instant.dev/vectors?label=) ![Build Status](https://app.travis-ci.com/instant-dev/vectors.svg?branch=main)
+
+## Batch create vectors without thinking about it
+
+When you're creating a lot of vectors - for example, indexing a bunch of documents
+at once using [OpenAI embeddings](https://platform.openai.com/docs/guides/embeddings) -
+you quickly run into IO-related performance issues. Your web requests will be throttled
+if you make too many parallel API requests, so OpenAI allows for batched requests
+via the [OpenAI embeddings API](https://platform.openai.com/docs/api-reference/embeddings).
+However, this API only allows for a maximum of 8,191 tokens per request:
+about 32,764 characters.
+
+**Solution:** `@instant.dev/vectors` provides a simple `VectorManager` utility that performs
+automatic, efficient batch creation of vectors via APIs. It will automatically collect
+vector creation requests over a 100ms (configurable) timeframe and batch them to minimize
+vector creation requests.
+
+It is most useful in web server contexts where multiple user requests may be
+creating vectors at the same time. If you rely on the same `VectorManager` instance
+all of these disparate requests will be efficiently batched.
+
+## Installation and Importing
+
+To use this library you'll need to also work with a vector creation tool, like OpenAI.
+
+```shell
+npm i @instant.dev/vectors --save # vector management
+npm i openai --save # openai for the engine
+```
+
+CommonJS:
+
+```javascript
+const { VectorManager } = require('@instant.dev/vectors');
+const OpenAI = require('openai');
+
+const openai = new OpenAI({apiKey: process.env.OPENAI_API_KEY});
+const Vectors = new VectorManager();
+```
+
+ESM:
+
+```javascript
+import { VectorManager } from '@instant.dev/vectors';
+import { Configuration, OpenAIApi } from "openai";
+const configuration = new Configuration({
+    organization: "YOUR_ORG_ID",
+    apiKey: process.env.OPENAI_API_KEY,
+});
+
+const openai = new OpenAIApi(configuration);
+const Vectors = new VectorManager();
+```
+
+## Usage
+
+Once you've imported and instantiated the package, you can use it like so.
+
+Set a batch engine:
+
+```javascript
+// values will automatically be batched appropriately
+Vectors.setEngine(async (values) => {
+  const embeddingResult = await openai.embeddings.create({
+    model: 'text-embedding-ada-002',
+    input: values,
+  });
+  return embeddingResult.data.map(entry => entry.embedding);
+});
+```
+
+Create a single vector:
+
+```javascript
+let vector = await Vectors.create(`Something to vectorize!`);
+```
+
+Create multiple vectors, will automatically batch:
+
+```javascript
+const myStrings = [
+  `Some string!`,
+  `Can also be a lot longer`,
+  `W`.repeat(1000),
+  // ...
+];
+
+let vectors = await Promise.all(myStrings.map(str => Vectors.create(str)));
+```
+
+Create multiple vectors with `batchCreate` utility:
+
+```javascript
+const myStrings = [
+  `Some string!`,
+  `Can also be a lot longer`,
+  `W`.repeat(1000),
+  // ...
+];
+
+let vectors = await Vectors.batchCreate(myStrings);
+```
+
+# Acknowledgements
+
+Special thank you to [Scott Gamble](https://x.com/threesided) who helps run all
+of the front-of-house work for instant.dev 💜!
+
+| Destination | Link |
+| ----------- | ---- |
+| Home | [instant.dev](https://instant.dev) |
+| GitHub | [github.com/instant-dev](https://github.com/instant-dev) |
+| Discord | [discord.gg/puVYgA7ZMh](https://discord.gg/puVYgA7ZMh) |
+| X / instant.dev | [x.com/instantdevs](https://x.com/instantdevs) |
+| X / Keith Horwood | [x.com/keithwhor](https://x.com/keithwhor) |
+| X / Scott Gamble | [x.com/threesided](https://x.com/threesided) |

package/core/vector_manager.js

@@ -0,0 +1,187 @@
+/**
+ * Convert strings and other objects to vectors, while automatically batching
+ * requests to minimize network and / or processor IO
+ * Defaults are set to work best with OpenAI's text-embedding-ada-002 model
+ */
+class VectorManager {
+
+  constructor () {
+    this.maximumBatchSize = 7168 * 4; // 4 tokens per word, estimated
+    this.maximumParallelRequests = 10; // 10 requests simultaneously max
+    this.fastQueueTime = 10; // time to wait if no other entries are added
+    this.waitQueueTime = 100; // time to wait to collect entries if 1+ entries are added
+    /**
+     * @private
+     */
+    this._vectorize = null;
+    /**
+     * @private
+     */
+    this._queue = [];
+    /**
+     * @private
+     */
+    this._results = new WeakMap();
+    /**
+     * @private
+     */
+    this._timeout = null;
+    /**
+     * @private
+     */
+    this._processing = false;
+  }
+
+  /**
+   * @private
+   */
+  async __sleep__ (t) {
+    return new Promise(r => setTimeout(() => r(true), t));
+  }
+
+  /**
+   * @private
+   */
+  async __dequeue__ () {
+    this._processing = true;
+    clearTimeout(this._timeout);
+    const queue = this._queue.slice();
+    this._timeout = null;
+    this._queue = [];
+    await this.batchVectorize(queue);
+    this._processing = false;
+    return true;
+  }
+
+  /**
+   * @private
+   */
+  async vectorizeValues (values) {
+    const vectors = await this._vectorize(values);
+    return vectors;
+  }
+
+  /**
+   * @private
+   */
+  async batchVectorize (queue) {
+    const strValues = queue.map(item => this.convertToString(item.value));
+    const batches = [[]];
+    let curBatchSize = 0;
+    while (strValues.length) {
+      const str = strValues.shift().slice(0, this.maximumBatchSize);
+      let n = batches.length - 1;
+      curBatchSize += str.length;
+      if (curBatchSize > this.maximumBatchSize) {
+        batches.push([str]);
+        n = batches.length - 1;
+        curBatchSize = str.length;
+      } else {
+        batches[n].push(str);
+      }
+    }
+    let i = 0;
+    while (batches.length) {
+      const parallelBatches = batches.splice(0, this.maximumParallelRequests);
+      const parallelVectors = await Promise.all(parallelBatches.map(strValues => this.vectorizeValues(strValues)));
+      parallelVectors.forEach((vectors, j) => {
+        vectors = Array.isArray(vectors)
+          ? vectors
+          : [];
+        parallelBatches[j].forEach((str, k) => {
+          if (vectors[k]) {
+            this._results.set(queue[i++], vectors[k]);
+          } else {
+            this._results.set(queue[i++], -1);
+          }
+        });
+      });
+    }
+    return true;
+  }
+
+  /**
+   * Converts the provided value to a string for easy vectorization
+   * @param {any} value 
+   * @returns {string}
+   */
+  convertToString (value) {
+    return (value === null || value === void 0)
+      ? ''
+      : typeof value === 'object'
+        ? JSON.stringify(value)
+        : (value + '');
+  }
+
+  /**
+   * Sets the vector engine: takes in an array of values with a maximum string length of this.maximumBatchSize
+   * @param {function} fnVectorize Expects a single argument, values, that is an array
+   * @returns {boolean}
+   */
+  setEngine (fnVectorize) {
+    if (typeof fnVectorize !== 'function') {
+      throw new Error(`.setEngine(fn) expects a valid function`);
+    } else if (fnVectorize.constructor.name !== 'AsyncFunction') {
+      throw new Error(`.setEngine(fn) expects an Asynchronous function`);
+    }
+    this._vectorize = fnVectorize;
+    return true;
+  }
+
+  async createTimeout () {
+    this._timeout = setTimeout(() => {
+      if (this._processing) {
+        this.createTimeout();
+      } else if (this._queue.length <= 1) {
+        this.__dequeue__();
+      } else {
+        this._timeout = setTimeout(
+          () => this.__dequeue__(),
+          this.waitQueueTime - this.fastQueueTime
+        );
+      }
+    }, this.fastQueueTime);
+  }
+
+  /**
+   * Creates a vector from a value
+   * @param {any} value Any value. null and undefined are converted to empty strings, non-string values are JSONified
+   * @returns 
+   */
+  async create (value) {
+    if (!this._vectorize) {
+      throw new Error(
+        `Could not vectorize: no vector engine has been set.\n` +
+        `Use Instant.Vectors.setEngine(fn) to enable.`
+      );
+    }
+    const item = {value};
+    this._queue.push(item);
+    if (!this._timeout) {
+      this.createTimeout();
+    }
+    let result = null;
+    while (!(result = this._results.get(item))) {
+      await this.__sleep__(10);
+    }
+    this._results.delete(item);
+    if (!Array.isArray(result)) {
+      throw new Error(
+        `Could not vectorize: vector engine did not return a valid vector for input "${value}"`
+      );
+    }
+    return result;
+  }
+
+  /**
+   * Batch creates vectors from an array
+   * @param {array} values Any values. null and undefined are converted to empty strings, non-string values are JSONified
+   * @returns 
+   */
+  async batchCreate (values) {
+    return Promise.all(values.map(value => this.create(value)));
+  }
+
+};
+
+module.exports = VectorManager;

package/index.js

@@ -0,0 +1,5 @@
+const VectorManager = require('./core/vector_manager.js');
+
+module.exports = {
+  VectorManager
+};

package/package.json

@@ -0,0 +1,25 @@
+{
+  "name": "@instant.dev/vectors",
+  "version": "0.1.0",
+  "description": "Utility for batch creating vectors via API",
+  "main": "index.js",
+  "scripts": {
+    "test": "mocha ./test/runner.js"
+  },
+  "keywords": [
+    "vectors",
+    "vector",
+    "management",
+    "openai",
+    "pinecone",
+    "pgvector"
+  ],
+  "author": "Keith Horwood",
+  "license": "MIT",
+  "devDependencies": {
+    "chai": "^4.3.10",
+    "dotenv": "^16.3.1",
+    "mocha": "^10.2.0",
+    "openai": "^4.11.0"
+  }
+}

package/test/runner.js

@@ -0,0 +1,19 @@
+require('dotenv').config({path: '.test.env'});
+const child_process = require('child_process');
+const os = require('os');
+const fs = require('fs');
+
+let args = [];
+try {
+  args = JSON.parse(process.env.npm_argv);
+  args = args.slice(3);
+} catch (e) {
+  args = [];
+}
+
+describe('Test Suite', function() {
+
+  let testFilenames = fs.readdirSync('./test/tests');
+  testFilenames.forEach(filename => require(`./tests/${filename}`)());
+
+});

package/test/tests/vectors.js

@@ -0,0 +1,147 @@
+const OpenAI = require('openai');
+const openai = new OpenAI({apiKey: process.env.OPENAI_API_KEY});
+
+const { VectorManager } = require('../../index.js');
+
+module.exports = (InstantORM, Databases) => {
+
+  const expect = require('chai').expect;
+
+  const Vectors = new VectorManager();
+
+  describe('VectorManager', async () => {
+
+    it('Should fail to vectorize without vector engine set', async () => {
+
+      const testPhrase = `I am extremely happy`;
+
+      let vector;
+      let error;
+
+      try {
+        vector = await Vectors.create(testPhrase);
+      } catch (e) {
+        error = e;
+      }
+
+      expect(error).to.exist;
+      expect(error.message).to.contain('Could not vectorize: no vector engine has been set');
+
+    });
+
+    it('Should fail to vectorize with a bad vector engine', async () => {
+
+      const testPhrase = `I am extremely happy`;
+
+      Vectors.setEngine(async (values) => {
+        // do nothing
+      });
+      
+      let vector;
+
+      try {
+        vector = await Vectors.create(testPhrase);
+      } catch (e) {
+        error = e;
+      }
+
+      expect(error).to.exist;
+      expect(error.message).to.contain('Could not vectorize: vector engine did not return a valid vector for input "I am extremely happy"');
+
+    });
+
+    it('Should succeed at vectorizing when vector engine is set properly', async () => {
+
+      const testPhrase = `I am extremely happy`;
+      let testVector;
+
+      Vectors.setEngine(async (values) => {
+        const embedding = await openai.embeddings.create({
+          model: 'text-embedding-ada-002',
+          input: values,
+        });
+        const vectors = embedding.data.map((entry, i) => {
+          let embedding = entry.embedding;
+          if (values[i] === testPhrase) {
+            testVector = embedding;
+          }
+          return embedding;
+        });
+        return vectors;
+      });
+      
+      const vector = await Vectors.create(testPhrase);
+
+      expect(vector).to.deep.equal(testVector);
+
+    });
+
+    it('Should create more vectors', async () => {
+
+      const vectorMap = {};
+
+      Vectors.setEngine(async (values) => {
+        const embedding = await openai.embeddings.create({
+          model: 'text-embedding-ada-002',
+          input: values,
+        });
+        return embedding.data.map((entry, i) => {
+          vectorMap[values[i]] = entry.embedding;
+          return vectorMap[values[i]];
+        });
+      });
+
+      const strings = [
+        `I am feeling awful`,
+        `I am in extreme distress`,
+        `I am feeling pretty good`,
+        `I am so-so`
+      ]
+
+      const vectors = await Promise.all(strings.map(str => Vectors.create(str)));
+
+      expect(vectors).to.exist;
+      expect(vectors.length).to.equal(4);
+      vectors.forEach((vector, i) => {
+        expect(vector).to.deep.equal(vectorMap[strings[i]]);
+      });
+
+    });
+
+    it('Should create many more vectors (50 vectors, ~4 per batch)', async function () {
+
+      this.timeout(10000);
+
+      const vectorMap = {};
+
+      Vectors.maximumBatchSize = 1000;
+      Vectors.setEngine(async (values) => {
+        const embedding = await openai.embeddings.create({
+          model: 'text-embedding-ada-002',
+          input: values,
+        });
+        return embedding.data.map((entry, i) => {
+          vectorMap[values[i]] = entry.embedding;
+          return vectorMap[values[i]];
+        });
+      });
+
+      let strings = Array(50).fill(0).map((_, i) => {
+        return i + '_ ' + Array(50).fill(0).map(() => {
+          return ['alpha', 'beta', 'gamma'][(Math.random() * 3) | 0]
+        }).join(' ');
+      });
+
+      const vectors = await Promise.all(strings.map(str => Vectors.create(str)));
+
+      expect(vectors).to.exist;
+      expect(vectors.length).to.equal(50);
+      vectors.forEach((vector, i) => {
+        expect(vector).to.deep.equal(vectorMap[strings[i]]);
+      });
+
+    });
+
+  });
+
+};