iframe-proxy | Sunbelt Computer Software

144 lines (123 loc) · 4.19 KB
import assert from 'node:assert';
import test from 'node:test';
import { Sequelize, DataTypes } from 'sequelize';
import 'pgvector/sequelize';
import { l2Distance, maxInnerProduct, cosineDistance, l1Distance, hammingDistance, jaccardDistance } from 'pgvector/sequelize';
import { SparseVector } from 'pgvector';
/** @import { Model, ModelCtor, Optional } from 'sequelize' */
 * @typedef {object} ItemAttributes
 * @property {number} id
 * @property {?number[]} embedding
 * @property {?number[]} half_embedding
 * @property {?string} binary_embedding
 * @property {?SparseVector} sparse_embedding
test('sequelize example', async () => {
  let sequelize = new Sequelize('postgres://localhost/pgvector_node_test', {
    logging: false
  await sequelize.query('CREATE EXTENSION IF NOT EXISTS vector');
  // need to reconnect after the vector extension has been created
  sequelize.close();
  sequelize = new Sequelize('postgres://localhost/pgvector_node_test', {
    logging: false
  /** @type {ModelCtor<Model<ItemAttributes, Optional<ItemAttributes, 'id'>> & ItemAttributes>} */
  const Item = sequelize.define('Item', {
    embedding: {
      type: DataTypes.VECTOR(3)
    half_embedding: {
      type: DataTypes.HALFVEC(3)
    binary_embedding: {
      type: 'BIT(3)'
    sparse_embedding: {
      type: DataTypes.SPARSEVEC(3)
    modelName: 'Item',
    tableName: 'sequelize_items',
    indexes: [
        fields: ['embedding'],
        using: 'hnsw',
        operator: 'vector_l2_ops'
  await Item.sync({force: true});
  await Item.create({embedding: [1, 1, 1], half_embedding: [1, 1, 1], binary_embedding: '000', sparse_embedding: new SparseVector([1, 1, 1])});
  await Item.create({embedding: [2, 2, 2], half_embedding: [2, 2, 2], binary_embedding: '101', sparse_embedding: new SparseVector([2, 2, 2])});
  await Item.create({embedding: [1, 1, 2], half_embedding: [1, 1, 2], binary_embedding: '111', sparse_embedding: new SparseVector([1, 1, 2])});
  // L2 distance
  let items = await Item.findAll({
    order: l2Distance('embedding', [1, 1, 1], sequelize),
    limit: 5
  assert.deepEqual(items.map(v => v.id), [1, 3, 2]);
  assert.deepEqual(items[0].embedding, [1, 1, 1]);
  assert.deepEqual(items[1].embedding, [1, 1, 2]);
  assert.deepEqual(items[2].embedding, [2, 2, 2]);
  // L2 distance - halfvec
  items = await Item.findAll({
    order: l2Distance('half_embedding', [1, 1, 1], sequelize),
    limit: 5
  assert.deepEqual(items.map(v => v.id), [1, 3, 2]);
  // L2 distance - sparsevec
  items = await Item.findAll({
    order: l2Distance('sparse_embedding', new SparseVector([1, 1, 1]), sequelize),
    limit: 5
  assert.deepEqual(items.map(v => v.id), [1, 3, 2]);
  await Item.create({});
  // max inner product
  items = await Item.findAll({
    order: maxInnerProduct(sequelize.literal('"Item".embedding'), [1, 1, 1], sequelize),
    limit: 5
  assert.deepEqual(items.map(v => v.id), [2, 3, 1, 4]);
  // cosine distance
  items = await Item.findAll({
    order: cosineDistance('embedding', [1, 1, 1], sequelize),
    limit: 5
  assert.deepEqual(items.map(v => v.id).slice(2), [3, 4]);
  // L1 distance
  items = await Item.findAll({
    order: l1Distance('embedding', [1, 1, 1], sequelize),
    limit: 5
  assert.deepEqual(items.map(v => v.id), [1, 3, 2, 4]);
  // Hamming distance
  items = await Item.findAll({
    order: hammingDistance('binary_embedding', '101', sequelize),
    limit: 5
  assert.deepEqual(items.map(v => v.id), [2, 3, 1, 4]);
  // Jaccard distance
  items = await Item.findAll({
    order: jaccardDistance('binary_embedding', '101', sequelize),
    limit: 5
  assert.deepEqual(items.map(v => v.id), [2, 3, 1, 4]);
  // bad value
  // @ts-ignore
  await assert.rejects(Item.create({embedding: 'bad'}), {message: /invalid input syntax for type vector/})
  sequelize.close();
test('dimensions', () => {
  assert.equal(DataTypes.VECTOR(3).toSql(), 'VECTOR(3)');
test('no dimensions', () => {
  assert.equal(DataTypes.VECTOR().toSql(), 'VECTOR');
test('bad dimensions', () => {
  assert.throws(() => {
    // @ts-ignore
    DataTypes.VECTOR('bad').toSql();
  }, {message: 'expected integer'});
Sunbelt Computer Software

PL/B Language Development and Support

Navigation Menu

Uh oh!

Sunbelt Computer Software

PL/B Language Development and Support

Uh oh!

FilesExpand file tree

sequelize.test.js

Latest commit

History

sequelize.test.js

File metadata and controls
