singulio-postgres 1.1.0

package/src/vector.ts

@@ -0,0 +1,312 @@
+/**
+ * @singulio/postgres - pgvector Extension Support
+ * Vector operations for AI/ML embeddings
+ */
+
+import { sql } from 'bun';
+import type {
+  VectorDistanceOperator,
+  VectorIndexType,
+  VectorSearchOptions,
+  QueryResult,
+  Logger,
+} from './types.js';
+
+/** Vector type (array of numbers) */
+export type Vector = number[] | Float32Array | Float64Array;
+
+/**
+ * Format a vector for PostgreSQL insertion
+ * Converts array to '[1,2,3]' format expected by pgvector
+ */
+export function formatVector(vector: Vector): string {
+  const arr = Array.isArray(vector) ? vector : Array.from(vector);
+  return `[${arr.join(',')}]`;
+}
+
+/**
+ * Parse a vector from PostgreSQL result
+ * Converts '[1,2,3]' string back to number array
+ */
+export function parseVector(value: string | number[] | null): number[] | null {
+  if (value === null) return null;
+
+  // Already an array (some drivers return it parsed)
+  if (Array.isArray(value)) return value;
+
+  // Parse string format '[1,2,3]'
+  if (typeof value === 'string') {
+    const trimmed = value.replace(/^\[|\]$/g, '');
+    if (!trimmed) return [];
+    return trimmed.split(',').map(Number);
+  }
+
+  return null;
+}
+
+/**
+ * Calculate vector dimension
+ */
+export function vectorDimension(vector: Vector): number {
+  return Array.isArray(vector) ? vector.length : vector.length;
+}
+
+/**
+ * Normalize a vector to unit length (for cosine similarity)
+ */
+export function normalizeVector(vector: Vector): number[] {
+  const arr = Array.isArray(vector) ? vector : Array.from(vector);
+  const magnitude = Math.sqrt(arr.reduce((sum, val) => sum + val * val, 0));
+  if (magnitude === 0) return arr;
+  return arr.map(val => val / magnitude);
+}
+
+/**
+ * Get the distance operator SQL string
+ */
+export function getDistanceOperator(op: VectorDistanceOperator): string {
+  return op;
+}
+
+/**
+ * Get human-readable name for distance operator
+ */
+export function getDistanceOperatorName(op: VectorDistanceOperator): string {
+  switch (op) {
+    case '<->':
+      return 'L2 (Euclidean)';
+    case '<=>':
+      return 'Cosine';
+    case '<#>':
+      return 'Inner Product';
+    case '<+>':
+      return 'L1 (Manhattan)';
+    default:
+      return 'Unknown';
+  }
+}
+
+/**
+ * Ensure pgvector extension is installed
+ */
+export async function ensureVectorExtension(): Promise<void> {
+  await sql`CREATE EXTENSION IF NOT EXISTS vector`;
+}
+
+/**
+ * Create a vector column on a table
+ */
+export async function createVectorColumn(
+  table: string,
+  column: string,
+  dimensions: number
+): Promise<void> {
+  await sql.unsafe(`ALTER TABLE ${table} ADD COLUMN IF NOT EXISTS ${column} vector(${dimensions})`);
+}
+
+/**
+ * Create a vector index for similarity search
+ */
+export async function createVectorIndex(
+  table: string,
+  column: string,
+  indexType: VectorIndexType = 'hnsw',
+  operator: VectorDistanceOperator = '<->',
+  options?: {
+    /** HNSW: max connections per layer (default: 16) */
+    m?: number;
+    /** HNSW: size of dynamic candidate list (default: 64) */
+    efConstruction?: number;
+    /** IVFFlat: number of lists (default: 100) */
+    lists?: number;
+  }
+): Promise<void> {
+  const indexName = `idx_${table}_${column}_${indexType}`;
+
+  // Determine operator class
+  let opClass: string;
+  switch (operator) {
+    case '<=>':
+      opClass = 'vector_cosine_ops';
+      break;
+    case '<#>':
+      opClass = 'vector_ip_ops';
+      break;
+    case '<+>':
+      opClass = 'vector_l1_ops';
+      break;
+    default:
+      opClass = 'vector_l2_ops';
+  }
+
+  if (indexType === 'hnsw') {
+    const m = options?.m ?? 16;
+    const efConstruction = options?.efConstruction ?? 64;
+    await sql.unsafe(`
+      CREATE INDEX IF NOT EXISTS ${indexName}
+      ON ${table} USING hnsw (${column} ${opClass})
+      WITH (m = ${m}, ef_construction = ${efConstruction})
+    `);
+  } else {
+    const lists = options?.lists ?? 100;
+    await sql.unsafe(`
+      CREATE INDEX IF NOT EXISTS ${indexName}
+      ON ${table} USING ivfflat (${column} ${opClass})
+      WITH (lists = ${lists})
+    `);
+  }
+}
+
+/**
+ * Search for similar vectors
+ */
+export async function vectorSearch<T = Record<string, unknown>>(
+  table: string,
+  column: string,
+  queryVector: Vector,
+  options: VectorSearchOptions = {},
+  logger?: Logger
+): Promise<QueryResult<T & { distance: number }>> {
+  const { operator = '<->', limit = 10, threshold, filter, filterParams = [] } = options;
+
+  const vectorStr = formatVector(queryVector);
+  const params: unknown[] = [vectorStr, ...filterParams, limit];
+
+  let whereClause = '';
+  if (filter) {
+    whereClause = `WHERE ${filter}`;
+  }
+  if (threshold !== undefined) {
+    const thresholdClause = `${column} ${operator} $1 < ${threshold}`;
+    whereClause = whereClause
+      ? `${whereClause} AND ${thresholdClause}`
+      : `WHERE ${thresholdClause}`;
+  }
+
+  const start = performance.now();
+  const result = await sql.unsafe(
+    `
+    SELECT *, ${column} ${operator} $1 AS distance
+    FROM ${table}
+    ${whereClause}
+    ORDER BY ${column} ${operator} $1
+    LIMIT $${params.length}
+  `,
+    params
+  );
+
+  logger?.debug('Vector search completed', {
+    table,
+    column,
+    operator: getDistanceOperatorName(operator),
+    results: result.length,
+    latencyMs: (performance.now() - start).toFixed(2),
+  });
+
+  return {
+    rows: result as (T & { distance: number })[],
+    rowCount: result.length,
+  };
+}
+
+/**
+ * Insert a row with vector data
+ */
+export async function insertWithVector(
+  table: string,
+  data: Record<string, unknown>,
+  vectorColumn: string,
+  vector: Vector
+): Promise<QueryResult> {
+  const columns = Object.keys(data);
+  const values = Object.values(data);
+
+  columns.push(vectorColumn);
+  values.push(formatVector(vector));
+
+  const placeholders = values.map((_, i) => `$${i + 1}`).join(', ');
+
+  const result = await sql.unsafe(
+    `
+    INSERT INTO ${table} (${columns.join(', ')})
+    VALUES (${placeholders})
+    RETURNING *
+  `,
+    values
+  );
+
+  return {
+    rows: result,
+    rowCount: result.length,
+  };
+}
+
+/**
+ * Update vector data for a row
+ */
+export async function updateVector(
+  table: string,
+  idColumn: string,
+  idValue: unknown,
+  vectorColumn: string,
+  vector: Vector
+): Promise<QueryResult> {
+  const result = await sql.unsafe(
+    `
+    UPDATE ${table}
+    SET ${vectorColumn} = $1
+    WHERE ${idColumn} = $2
+    RETURNING *
+  `,
+    [formatVector(vector), idValue]
+  );
+
+  return {
+    rows: result,
+    rowCount: result.length,
+  };
+}
+
+/**
+ * Batch insert rows with vectors
+ */
+export async function batchInsertWithVectors(
+  table: string,
+  columns: string[],
+  vectorColumn: string,
+  rows: Array<{ data: unknown[]; vector: Vector }>
+): Promise<QueryResult> {
+  if (rows.length === 0) {
+    return { rows: [], rowCount: 0 };
+  }
+
+  const allColumns = [...columns, vectorColumn];
+  const valueSets: string[] = [];
+  const params: unknown[] = [];
+  let paramIndex = 1;
+
+  for (const row of rows) {
+    const placeholders: string[] = [];
+    for (const value of row.data) {
+      placeholders.push(`$${paramIndex++}`);
+      params.push(value);
+    }
+    placeholders.push(`$${paramIndex++}`);
+    params.push(formatVector(row.vector));
+    valueSets.push(`(${placeholders.join(', ')})`);
+  }
+
+  const result = await sql.unsafe(
+    `
+    INSERT INTO ${table} (${allColumns.join(', ')})
+    VALUES ${valueSets.join(', ')}
+    RETURNING *
+  `,
+    params
+  );
+
+  return {
+    rows: result,
+    rowCount: result.length,
+  };
+}