pgvector-node/tests/slonik.test.js at master · pgvector/pgvector-node · GitHub

History

39 lines (33 loc) · 2.2 KB

Raw

import assert from 'node:assert';

import test from 'node:test';

import pgvector from 'pgvector';

import { SparseVector } from 'pgvector';

import { createPool, sql } from 'slonik';

test('slonik example', async () => {

const pool = await createPool('postgres://localhost/pgvector_node_test');

await pool.query(sql.unsafe`CREATE EXTENSION IF NOT EXISTS vector`);

await pool.query(sql.unsafe`DROP TABLE IF EXISTS slonik_items`);

await pool.query(sql.unsafe`CREATE TABLE slonik_items (id serial PRIMARY KEY, embedding vector(3), half_embedding halfvec(3), binary_embedding bit(3), sparse_embedding sparsevec(3))`);

const embedding1 = pgvector.toSql([1, 1, 1]);

const embedding2 = pgvector.toSql([2, 2, 2]);

const embedding3 = pgvector.toSql([1, 1, 2]);

const halfEmbedding1 = pgvector.toSql([1, 1, 1]);

const halfEmbedding2 = pgvector.toSql([2, 2, 2]);

const halfEmbedding3 = pgvector.toSql([1, 1, 2]);

const binaryEmbedding1 = '000';

const binaryEmbedding2 = '101';

const binaryEmbedding3 = '111';

const sparseEmbedding1 = pgvector.toSql(new SparseVector([1, 1, 1]));

const sparseEmbedding2 = pgvector.toSql(new SparseVector([2, 2, 2]));

const sparseEmbedding3 = pgvector.toSql(new SparseVector([1, 1, 2]));

await pool.query(sql.unsafe`INSERT INTO slonik_items (embedding, half_embedding, binary_embedding, sparse_embedding) VALUES (${embedding1}, ${halfEmbedding1}, ${binaryEmbedding1}, ${sparseEmbedding1}), (${embedding2}, ${halfEmbedding2}, ${binaryEmbedding2}, ${sparseEmbedding2}), (${embedding3}, ${halfEmbedding3}, ${binaryEmbedding3}, ${sparseEmbedding3})`);

const embedding = pgvector.toSql([1, 1, 1]);

const items = await pool.query(sql.unsafe`SELECT * FROM slonik_items ORDER BY embedding <-> ${embedding} LIMIT 5`);

assert.deepEqual(items.rows.map(v => v.id), [1, 3, 2]);

assert.deepEqual(pgvector.fromSql(items.rows[0].embedding), [1, 1, 1]);

assert.deepEqual(pgvector.fromSql(items.rows[0].half_embedding), [1, 1, 1]);

assert.equal(items.rows[0].binary_embedding, '000');

assert.deepEqual(pgvector.fromSql(items.rows[0].sparse_embedding), new SparseVector([1, 1, 1]));

await pool.query(sql.unsafe`CREATE INDEX ON slonik_items USING hnsw (embedding vector_l2_ops)`);

await pool.end();

});

pgvector-node/tests/slonik.test.js at master · pgvector/pgvector-node · GitHub

Uh oh!

Expand file tree

Breadcrumbs

slonik.test.js

Latest commit

History

Breadcrumbs

slonik.test.js

File metadata and controls