Added Vector class

README.md

@@ -220,11 +220,11 @@ await db.schema.createTable('items')
 Insert vectors
 
 ```javascript
-import pgvector from 'pgvector/kysely';
+import { vector } from 'pgvector/kysely';
 
 const newItems = [
-  {embedding: pgvector.toSql([1, 2, 3])},
-  {embedding: pgvector.toSql([4, 5, 6])}
+  {embedding: vector([1, 2, 3])},
+  {embedding: vector([4, 5, 6])}
 ];
 await db.insertInto('items').values(newItems).execute();
 ```

package.json

@@ -12,6 +12,10 @@
     "url": "https://github.com/pgvector/pgvector-node"
   },
   "exports": {
+    ".": {
+      "types": "./types/index.d.ts",
+      "default": "./src/index.js"
+    },
     "./drizzle-orm": {
       "types": "./types/drizzle-orm/index.d.ts",
       "default": "./src/drizzle-orm/index.js"
@@ -48,6 +52,9 @@
   },
   "typesVersions": {
     "*": {
+      "*": [
+        "types/index.d.ts"
+      ],
       "drizzle-orm": [
         "types/drizzle-orm/index.d.ts"
       ],

src/index.js

@@ -0,0 +1,29 @@
+const utils = require('./utils');
+
+class Vector {
+  constructor(vec) {
+    if (Array.isArray(vec)) {
+      this.vec = vec;
+    } else {
+      this.vec = utils.fromSql(vec);
+    }
+  }
+
+  toArray() {
+    return this.vec;
+  }
+
+  toPostgres() {
+    return this.toString();
+  }
+
+  toString() {
+    return utils.toSql(this.vec);
+  }
+}
+
+function vector(vec) {
+  return new Vector(vec);
+}
+
+module.exports = {vector};

src/kysely/index.js

@@ -1,16 +1,16 @@
 const { sql } = require('kysely');
-const { fromSql, toSql } = require('../utils');
+const { vector } = require('..');
 
 function l2Distance(column, value) {
-  return sql`${sql.ref(column)} <-> ${toSql(value)}`;
+  return sql`${sql.ref(column)} <-> ${vector(value)}`;
 }
 
 function maxInnerProduct(column, value) {
-  return sql`${sql.ref(column)} <#> ${toSql(value)}`;
+  return sql`${sql.ref(column)} <#> ${vector(value)}`;
 }
 
 function cosineDistance(column, value) {
-  return sql`${sql.ref(column)} <=> ${toSql(value)}`;
+  return sql`${sql.ref(column)} <=> ${vector(value)}`;
 }
 
-module.exports = {fromSql, toSql, l2Distance, maxInnerProduct, cosineDistance};
+module.exports = {vector, l2Distance, maxInnerProduct, cosineDistance};

tests/kysely/index.test.mjs

@@ -1,7 +1,6 @@
 import pg from 'pg';
 import { Kysely, PostgresDialect, sql } from 'kysely';
-import pgvector from 'pgvector/kysely';
-import { l2Distance, maxInnerProduct, cosineDistance } from 'pgvector/kysely';
+import { vector, l2Distance, maxInnerProduct, cosineDistance } from 'pgvector/kysely';
 
 test('example', async () => {
   const dialect = new PostgresDialect({
@@ -26,9 +25,9 @@ test('example', async () => {
     .execute();
 
   const newItems = [
-    {embedding: pgvector.toSql([1, 1, 1])},
-    {embedding: pgvector.toSql([2, 2, 2])},
-    {embedding: pgvector.toSql([1, 1, 2])}
+    {embedding: vector([1, 1, 1])},
+    {embedding: vector([2, 2, 2])},
+    {embedding: vector([1, 1, 2])}
   ];
   await db.insertInto('kysely_items')
     .values(newItems)
@@ -40,9 +39,9 @@ test('example', async () => {
     .limit(5)
     .execute();
   expect(items.map(v => v.id)).toStrictEqual([1, 3, 2]);
-  expect(pgvector.fromSql(items[0].embedding)).toStrictEqual([1, 1, 1]);
-  expect(pgvector.fromSql(items[1].embedding)).toStrictEqual([1, 1, 2]);
-  expect(pgvector.fromSql(items[2].embedding)).toStrictEqual([2, 2, 2]);
+  expect(vector(items[0].embedding).toArray()).toStrictEqual([1, 1, 1]);
+  expect(vector(items[1].embedding).toArray()).toStrictEqual([1, 1, 2]);
+  expect(vector(items[2].embedding).toArray()).toStrictEqual([2, 2, 2]);
 
   items = await db.selectFrom('kysely_items')
     .selectAll()