@useverk/drizzle-pgvector
Version:
52 lines (44 loc) • 1.17 kB
text/typescript
import { AnyColumn, SQLWrapper, sql } from "drizzle-orm";
import { customType } from "drizzle-orm/pg-core";
export const customVector = customType<{
data: number[];
driverData: string;
config: { dimensions: number };
}>({
dataType(config) {
if (!config) {
return `vector(1536)`;
}
return `vector(${config.dimensions})`;
},
toDriver(value: number[]): string {
return toSql(value);
},
fromDriver(value) {
return fromSql(value);
},
});
function fromSql(value: string) {
return value
.substring(1, value.length - 1)
.split(",")
.map((v) => parseFloat(v));
}
function toSql(value: number[]): string {
return JSON.stringify(value);
}
export function l2Distance(column: SQLWrapper | AnyColumn, value: number[]) {
return sql`${column} <-> ${toSql(value)}`;
}
export function maxInnerProduct(
column: SQLWrapper | AnyColumn,
value: number[]
) {
return sql`${column} <#> ${toSql(value)}`;
}
export function cosineDistance(
column: SQLWrapper | AnyColumn,
value: number[]
) {
return sql`${column} <=> ${toSql(value)}`;
}