@kernl-sdk/pg 0.1.10 → 0.1.12

package/src/pgvector/__tests__/search.test.ts

@@ -0,0 +1,366 @@
+import { describe, it, expect, beforeAll, afterAll } from "vitest";
+import { Pool } from "pg";
+import { PGSearchIndex } from "../search";
+
+const TEST_DB_URL = process.env.KERNL_PG_TEST_URL;
+const SCHEMA = "kernl_search_idx_test";
+
+describe.sequential("PGSearchIndex", () => {
+  if (!TEST_DB_URL) {
+    it.skip("requires KERNL_PG_TEST_URL environment variable", () => {});
+    return;
+  }
+
+  let pool: Pool;
+  let pgvec: PGSearchIndex;
+
+  beforeAll(async () => {
+    pool = new Pool({ connectionString: TEST_DB_URL });
+
+    // Ensure pgvector extension exists
+    await pool.query(`CREATE EXTENSION IF NOT EXISTS vector`);
+
+    // Clean slate
+    await pool.query(`DROP SCHEMA IF EXISTS "${SCHEMA}" CASCADE`);
+    await pool.query(`CREATE SCHEMA "${SCHEMA}"`);
+
+    pgvec = new PGSearchIndex({ pool });
+  });
+
+  afterAll(async () => {
+    await pool.query(`DROP SCHEMA IF EXISTS "${SCHEMA}" CASCADE`);
+    await pool.end();
+  });
+
+  describe("createIndex", () => {
+    it("creates a table with correct column types", async () => {
+      await pgvec.createIndex({
+        id: "articles",
+        schema: {
+          id: { type: "string", pk: true },
+          title: { type: "string" },
+          views: { type: "int" },
+          rating: { type: "float" },
+          published: { type: "boolean" },
+          created_at: { type: "date" },
+        },
+        providerOptions: { schema: SCHEMA },
+      });
+
+      // Verify table exists with correct columns
+      const result = await pool.query<{
+        column_name: string;
+        data_type: string;
+      }>(
+        `SELECT column_name, data_type
+         FROM information_schema.columns
+         WHERE table_schema = $1 AND table_name = $2
+         ORDER BY ordinal_position`,
+        [SCHEMA, "articles"],
+      );
+
+      const columns = Object.fromEntries(
+        result.rows.map((r) => [r.column_name, r.data_type]),
+      );
+
+      expect(columns.id).toBe("text");
+      expect(columns.title).toBe("text");
+      expect(columns.views).toBe("integer");
+      expect(columns.rating).toBe("double precision");
+      expect(columns.published).toBe("boolean");
+      expect(columns.created_at).toBe("timestamp with time zone");
+    });
+
+    it("creates vector columns with correct dimensions", async () => {
+      await pgvec.createIndex({
+        id: "embeddings",
+        schema: {
+          id: { type: "string", pk: true },
+          embedding: { type: "vector", dimensions: 384 },
+        },
+        providerOptions: { schema: SCHEMA },
+      });
+
+      // Check the vector column type via pg_attribute
+      const result = await pool.query<{ typname: string; atttypmod: number }>(
+        `SELECT t.typname, a.atttypmod
+         FROM pg_attribute a
+         JOIN pg_class c ON a.attrelid = c.oid
+         JOIN pg_namespace n ON c.relnamespace = n.oid
+         JOIN pg_type t ON a.atttypid = t.oid
+         WHERE n.nspname = $1
+           AND c.relname = $2
+           AND a.attname = $3`,
+        [SCHEMA, "embeddings", "embedding"],
+      );
+
+      expect(result.rows[0]?.typname).toBe("vector");
+      // atttypmod encodes dimensions for vector type
+      expect(result.rows[0]?.atttypmod).toBe(384);
+    });
+
+    it("creates HNSW index for vector fields", async () => {
+      await pgvec.createIndex({
+        id: "searchable",
+        schema: {
+          id: { type: "string", pk: true },
+          embedding: { type: "vector", dimensions: 128, similarity: "cosine" },
+        },
+        providerOptions: { schema: SCHEMA },
+      });
+
+      const result = await pool.query<{ indexname: string; indexdef: string }>(
+        `SELECT indexname, indexdef
+         FROM pg_indexes
+         WHERE schemaname = $1 AND tablename = $2`,
+        [SCHEMA, "searchable"],
+      );
+
+      const hnswIndex = result.rows.find(
+        (r) => r.indexname === "searchable_embedding_idx",
+      );
+
+      expect(hnswIndex).toBeDefined();
+      expect(hnswIndex?.indexdef).toContain("hnsw");
+      expect(hnswIndex?.indexdef).toContain("vector_cosine_ops");
+    });
+
+    it("uses correct operator class for each similarity metric", async () => {
+      // Cosine
+      await pgvec.createIndex({
+        id: "cosine_idx",
+        schema: {
+          id: { type: "string", pk: true },
+          vec: { type: "vector", dimensions: 8, similarity: "cosine" },
+        },
+        providerOptions: { schema: SCHEMA },
+      });
+
+      // Euclidean
+      await pgvec.createIndex({
+        id: "euclidean_idx",
+        schema: {
+          id: { type: "string", pk: true },
+          vec: { type: "vector", dimensions: 8, similarity: "euclidean" },
+        },
+        providerOptions: { schema: SCHEMA },
+      });
+
+      // Dot product
+      await pgvec.createIndex({
+        id: "dot_idx",
+        schema: {
+          id: { type: "string", pk: true },
+          vec: { type: "vector", dimensions: 8, similarity: "dot_product" },
+        },
+        providerOptions: { schema: SCHEMA },
+      });
+
+      const result = await pool.query<{ indexname: string; indexdef: string }>(
+        `SELECT indexname, indexdef
+         FROM pg_indexes
+         WHERE schemaname = $1
+         ORDER BY indexname`,
+        [SCHEMA],
+      );
+
+      const indexes = Object.fromEntries(
+        result.rows.map((r) => [r.indexname, r.indexdef]),
+      );
+
+      expect(indexes["cosine_idx_vec_idx"]).toContain("vector_cosine_ops");
+      expect(indexes["euclidean_idx_vec_idx"]).toContain("vector_l2_ops");
+      expect(indexes["dot_idx_vec_idx"]).toContain("vector_ip_ops");
+    });
+
+    it("throws if schema has no pk field", async () => {
+      await expect(
+        pgvec.createIndex({
+          id: "no_pk",
+          schema: {
+            title: { type: "string" },
+            content: { type: "string" },
+          },
+          providerOptions: { schema: SCHEMA },
+        }),
+      ).rejects.toThrow("schema must have a field with pk: true");
+    });
+
+    it("auto-binds the created index for immediate use", async () => {
+      await pgvec.createIndex({
+        id: "auto_bound",
+        schema: {
+          id: { type: "string", pk: true },
+          name: { type: "string" },
+          embedding: { type: "vector", dimensions: 3, similarity: "cosine" },
+        },
+        providerOptions: { schema: SCHEMA },
+      });
+
+      // Should be able to use the index immediately without bindIndex
+      const handle = pgvec.index("auto_bound");
+      expect(handle.id).toBe("auto_bound");
+
+      // Insert and query should work
+      await handle.upsert({
+        id: "test-1",
+        name: "Test Doc",
+        embedding: [0.1, 0.2, 0.3],
+      });
+
+      const results = await handle.query({
+        query: [{ embedding: [0.1, 0.2, 0.3] }],
+        topK: 1,
+      });
+
+      expect(results).toHaveLength(1);
+      expect(results[0].id).toBe("test-1");
+    });
+  });
+
+  describe("deleteIndex", () => {
+    it("drops the table and removes binding", async () => {
+      await pgvec.createIndex({
+        id: "to_delete",
+        schema: {
+          id: { type: "string", pk: true },
+          name: { type: "string" },
+        },
+        providerOptions: { schema: SCHEMA },
+      });
+
+      // Verify table exists
+      const before = await pool.query(
+        `SELECT 1 FROM information_schema.tables
+         WHERE table_schema = $1 AND table_name = $2`,
+        [SCHEMA, "to_delete"],
+      );
+      expect(before.rows).toHaveLength(1);
+
+      await pgvec.deleteIndex("to_delete");
+
+      // Verify table is gone
+      const after = await pool.query(
+        `SELECT 1 FROM information_schema.tables
+         WHERE table_schema = $1 AND table_name = $2`,
+        [SCHEMA, "to_delete"],
+      );
+      expect(after.rows).toHaveLength(0);
+    });
+
+    it("throws if index is not bound", async () => {
+      await expect(pgvec.deleteIndex("nonexistent")).rejects.toThrow(
+        'Index "nonexistent" not bound',
+      );
+    });
+  });
+
+  describe("listIndexes", () => {
+    it("returns empty page when no indexes match prefix", async () => {
+      const page = await pgvec.listIndexes({ prefix: "nonexistent_prefix_" });
+
+      expect(page.data).toEqual([]);
+      expect(page.last).toBe(true);
+    });
+
+    it("lists created indexes", async () => {
+      await pgvec.createIndex({
+        id: "list_test_a",
+        schema: { id: { type: "string", pk: true } },
+        providerOptions: { schema: SCHEMA },
+      });
+      await pgvec.createIndex({
+        id: "list_test_b",
+        schema: { id: { type: "string", pk: true } },
+        providerOptions: { schema: SCHEMA },
+      });
+
+      const page = await pgvec.listIndexes();
+
+      const ids = page.data.map((s) => s.id);
+      expect(ids).toContain("list_test_a");
+      expect(ids).toContain("list_test_b");
+    });
+
+    it("filters by prefix", async () => {
+      await pgvec.createIndex({
+        id: "prefix_foo_1",
+        schema: { id: { type: "string", pk: true } },
+        providerOptions: { schema: SCHEMA },
+      });
+      await pgvec.createIndex({
+        id: "prefix_bar_1",
+        schema: { id: { type: "string", pk: true } },
+        providerOptions: { schema: SCHEMA },
+      });
+
+      const page = await pgvec.listIndexes({ prefix: "prefix_foo" });
+
+      expect(page.data).toHaveLength(1);
+      expect(page.data[0].id).toBe("prefix_foo_1");
+    });
+
+    it("respects limit and provides cursor for pagination", async () => {
+      await pgvec.createIndex({
+        id: "page_1",
+        schema: { id: { type: "string", pk: true } },
+        providerOptions: { schema: SCHEMA },
+      });
+      await pgvec.createIndex({
+        id: "page_2",
+        schema: { id: { type: "string", pk: true } },
+        providerOptions: { schema: SCHEMA },
+      });
+      await pgvec.createIndex({
+        id: "page_3",
+        schema: { id: { type: "string", pk: true } },
+        providerOptions: { schema: SCHEMA },
+      });
+
+      const page1 = await pgvec.listIndexes({ prefix: "page_", limit: 2 });
+
+      expect(page1.data).toHaveLength(2);
+      expect(page1.last).toBe(false);
+
+      const page2 = await page1.next();
+
+      expect(page2).not.toBeNull();
+      expect(page2!.data).toHaveLength(1);
+      expect(page2!.last).toBe(true);
+    });
+  });
+
+  describe("describeIndex", () => {
+    it("returns stats for an index", async () => {
+      await pgvec.createIndex({
+        id: "describe_test",
+        schema: {
+          id: { type: "string", pk: true },
+          content: { type: "string" },
+          embedding: { type: "vector", dimensions: 128, similarity: "cosine" },
+        },
+        providerOptions: { schema: SCHEMA },
+      });
+
+      const handle = pgvec.index("describe_test");
+      const vec = new Array(128).fill(0.1);
+      await handle.upsert({ id: "doc-1", content: "hello", embedding: vec });
+      await handle.upsert({ id: "doc-2", content: "world", embedding: vec });
+
+      const stats = await pgvec.describeIndex("describe_test");
+
+      expect(stats.id).toBe("describe_test");
+      expect(stats.count).toBe(2);
+      expect(stats.sizeb).toBeGreaterThan(0);
+      expect(stats.dimensions).toBe(128);
+      expect(stats.similarity).toBe("cosine");
+      expect(stats.status).toBe("ready");
+    });
+
+    it("throws if index is not bound", async () => {
+      await expect(pgvec.describeIndex("nonexistent")).rejects.toThrow(
+        'Index "nonexistent" not bound',
+      );
+    });
+  });
+});

package/src/pgvector/handle.ts

@@ -0,0 +1,285 @@
+import type { Pool } from "pg";
+
+import type {
+  IndexHandle,
+  DocumentPatch,
+  QueryInput,
+  SearchHit,
+  FieldSchema,
+  UnknownDocument,
+  UpsertResult,
+  PatchResult,
+  DeleteResult,
+} from "@kernl-sdk/retrieval";
+import { normalizeQuery } from "@kernl-sdk/retrieval";
+
+import { SEARCH_HIT } from "./hit";
+import { sqlize, SQL_SELECT, SQL_WHERE, SQL_ORDER, SQL_LIMIT } from "./sql";
+import type { PGIndexConfig } from "./types";
+import { parseIndexId, isVector } from "./utils";
+
+/**
+ * pgvector-backed IndexHandle.
+ */
+export class PGIndexHandle<TDocument = UnknownDocument>
+  implements IndexHandle<TDocument>
+{
+  readonly id: string;
+
+  private pool: Pool;
+  private ensureInit: () => Promise<void>;
+  private config?: PGIndexConfig;
+
+  constructor(
+    pool: Pool,
+    ensureInit: () => Promise<void>,
+    id: string,
+    config?: PGIndexConfig,
+  ) {
+    this.id = id;
+    this.pool = pool;
+    this.ensureInit = ensureInit;
+    this.config = config;
+  }
+
+  /**
+   * Query the index using vector search, full-text search, or filters.
+   */
+  async query(input: QueryInput): Promise<SearchHit<TDocument>[]> {
+    await this.ensureInit();
+
+    const q = normalizeQuery(input);
+    const { schema, table, pkey } = this.table;
+
+    const sqlized = sqlize(q, { pkey, schema, table, binding: this.config });
+
+    const select = SQL_SELECT.encode(sqlized.select);
+    const where = SQL_WHERE.encode({
+      ...sqlized.where,
+      startIdx: select.params.length + 1,
+    });
+    const order = SQL_ORDER.encode(sqlized.order);
+    const limit = SQL_LIMIT.encode({
+      ...sqlized.limit,
+      startIdx: select.params.length + where.params.length + 1,
+    });
+
+    const sql = `
+      SELECT ${select.sql}
+      FROM "${schema}"."${table}"
+      ${where.sql ? `WHERE ${where.sql}` : ""}
+      ORDER BY ${order.sql}
+      ${limit.sql}
+    `;
+
+    const params = [...select.params, ...where.params, ...limit.params];
+    const result = await this.pool.query(sql, params);
+
+    return result.rows.map((row) =>
+      SEARCH_HIT.decode<TDocument>(row, this.id, this.config),
+    );
+  }
+
+  /* ---- document ops ---- */
+
+  /**
+   * Upsert one or more documents.
+   *
+   * Documents are flat objects. The pkey field (default "id") is used for
+   * conflict resolution. All other fields are written to matching columns.
+   */
+  async upsert(docs: TDocument | TDocument[]): Promise<UpsertResult> {
+    const arr = Array.isArray(docs) ? docs : [docs];
+    if (arr.length === 0) return { count: 0, inserted: 0, updated: 0 };
+
+    await this.ensureInit();
+    const { schema, table, pkey, fields } = this.table;
+
+    // Find which logical field maps to the pkey column
+    let pkeyField = pkey; // default: same name
+    for (const [fieldName, fieldCfg] of Object.entries(fields)) {
+      if (fieldCfg.column === pkey) {
+        pkeyField = fieldName;
+        break;
+      }
+    }
+
+    // collect field names from first doc (excluding the pkey field)
+    const first = arr[0] as UnknownDocument;
+    const fieldNames = Object.keys(first).filter((k) => k !== pkeyField);
+
+    // map field name → column name (from binding or same name)
+    const colFor = (name: string) => fields[name]?.column ?? name;
+    const cols = [pkey, ...fieldNames.map(colFor)];
+
+    const params: unknown[] = [];
+    const rows: string[] = [];
+
+    for (const doc of arr) {
+      const d = doc as UnknownDocument;
+      const placeholders: string[] = [];
+
+      // pkey value - use the logical field name, not column name
+      const pkval = d[pkeyField];
+      if (typeof pkval !== "string") {
+        throw new Error(`Document missing string field "${pkeyField}"`);
+      }
+      params.push(pkval);
+      placeholders.push(`$${params.length}`);
+
+      // ...other fields
+      for (const field of fieldNames) {
+        const val = d[field] ?? null;
+        const binding = fields[field];
+
+        // detect vector: explicit binding or runtime check
+        const isVec = binding?.type === "vector" || isVector(val);
+
+        params.push(isVec ? JSON.stringify(val) : val);
+        placeholders.push(
+          isVec ? `$${params.length}::vector` : `$${params.length}`,
+        );
+      }
+
+      rows.push(`(${placeholders.join(", ")})`);
+    }
+
+    // build SET clause for conflict (exclude pkey)
+    const sets = cols.slice(1).map((c) => `"${c}" = EXCLUDED."${c}"`);
+
+    // xmax = 0 means inserted, xmax != 0 means updated
+    const sql = `
+      INSERT INTO "${schema}"."${table}" (${cols.map((c) => `"${c}"`).join(", ")})
+      VALUES ${rows.join(", ")}
+      ON CONFLICT ("${pkey}") DO UPDATE SET ${sets.join(", ")}
+      RETURNING (xmax = 0) as inserted
+    `;
+
+    const result = await this.pool.query(sql, params);
+    const inserted = result.rows.filter((r) => r.inserted).length;
+
+    return {
+      count: result.rowCount ?? 0,
+      inserted,
+      updated: (result.rowCount ?? 0) - inserted,
+    };
+  }
+
+  /**
+   * Patch one or more documents.
+   *
+   * Only specified fields are updated. Set a field to `null` to unset it.
+   */
+  async patch(
+    patches: DocumentPatch<TDocument> | DocumentPatch<TDocument>[],
+  ): Promise<PatchResult> {
+    const arr = Array.isArray(patches) ? patches : [patches];
+    if (arr.length === 0) return { count: 0 };
+
+    await this.ensureInit();
+    const { schema, table, pkey, fields } = this.table;
+
+    let totalCount = 0;
+
+    // process each patch individually (different fields may be updated)
+    for (const patch of arr) {
+      const pkval = patch.id;
+      if (typeof pkval !== "string") {
+        throw new Error(`Patch missing string field "id"`);
+      }
+
+      // collect fields to update (excluding pkey)
+      const updates: string[] = [];
+      const params: unknown[] = [];
+
+      for (const [key, val] of Object.entries(patch)) {
+        if (key === "id" || key === pkey) continue;
+        if (val === undefined) continue;
+
+        const col = fields[key]?.column ?? key;
+        const binding = fields[key];
+        const isVec = binding?.type === "vector" || isVector(val);
+
+        params.push(isVec && val !== null ? JSON.stringify(val) : val);
+        updates.push(
+          `"${col}" = $${params.length}${isVec && val !== null ? "::vector" : ""}`,
+        );
+      }
+
+      if (updates.length === 0) continue;
+
+      params.push(pkval);
+      const sql = `
+        UPDATE "${schema}"."${table}"
+        SET ${updates.join(", ")}
+        WHERE "${pkey}" = $${params.length}
+      `;
+
+      const result = await this.pool.query(sql, params);
+      totalCount += result.rowCount ?? 0;
+    }
+
+    return { count: totalCount };
+  }
+
+  /**
+   * Delete one or more documents by ID.
+   */
+  async delete(ids: string | string[]): Promise<DeleteResult> {
+    const arr = Array.isArray(ids) ? ids : [ids];
+    if (arr.length === 0) return { count: 0 };
+
+    await this.ensureInit();
+    const { schema, table, pkey } = this.table;
+
+    const placeholders = arr.map((_, i) => `$${i + 1}`);
+    const sql = `
+      DELETE FROM "${schema}"."${table}"
+      WHERE "${pkey}" IN (${placeholders.join(", ")})
+    `;
+
+    const result = await this.pool.query(sql, arr);
+    return { count: result.rowCount ?? 0 };
+  }
+
+  /* ---- schema ops ---- */
+
+  /**
+   * Add a field to the index schema.
+   *
+   * Not yet implemented - left as a stub until we have a concrete use case
+   * to inform the design (e.g. vector-only vs general columns, index creation).
+   */
+  async addField(_field: string, _schema: FieldSchema): Promise<void> {
+    throw new Error("addField not yet implemented");
+  }
+
+  /* ---- internal utils ---- */
+
+  /**
+   * Resolve table config: use explicit binding or derive from conventions.
+   *
+   * Resolution order:
+   * - If a PGIndexConfig was provided or bound via PGSearchIndex.createIndex()
+   *   / bindIndex(), that config is used (schema, table, pkey, fields).
+   * - Otherwise we fall back to convention-based parsing of the index id via
+   *   parseIndexId(id):
+   *     - "docs"             → schema: "public",    table: "docs"
+   *     - "analytics.events" → schema: "analytics", table: "events"
+   *
+   * This allows callers to:
+   *   - Point at existing tables without an explicit bindIndex call by using
+   *     either "table" or "schema.table" ids, and
+   *   - Rely on explicit configs (from createIndex/bindIndex) when using the
+   *     higher-level SearchIndex API.
+   */
+  private get table() {
+    if (this.config) {
+      return this.config;
+    }
+
+    // convention-based defaults
+    const { schema, table } = parseIndexId(this.id);
+    return { schema, table, pkey: "id", fields: {} };
+  }
+}

package/src/pgvector/hit.ts

@@ -0,0 +1,56 @@
+import { FieldValue, SearchHit, UnknownDocument } from "@kernl-sdk/retrieval";
+
+import type { PGIndexConfig } from "./types";
+
+/**
+ * Codec for converting DB row ←→ SearchHit.
+ */
+export const SEARCH_HIT = {
+  encode: (_hit: SearchHit): Record<string, unknown> => {
+    throw new Error("SEARCH_HIT.encode: not implemented");
+  },
+
+  decode: <TDocument = UnknownDocument>(
+    row: Record<string, unknown>,
+    index: string,
+    config?: PGIndexConfig,
+  ): SearchHit<TDocument> => {
+    const { id, score, ...rest } = row;
+    const doc: Record<string, FieldValue> = {};
+
+    // Helper to parse pgvector strings like "[0.1,0.2,0.3]" back to arrays
+    const parseValue = (val: unknown, isVector: boolean): FieldValue => {
+      if (isVector && typeof val === "string" && val.startsWith("[")) {
+        return JSON.parse(val);
+      }
+      return val as FieldValue;
+    };
+
+    if (config) {
+      // map columns back to logical field names
+      for (const [field, cfg] of Object.entries(config.fields)) {
+        const col = cfg.column;
+        if (col in rest) {
+          const isVector = cfg.type === "vector";
+          doc[field] = parseValue(rest[col], isVector);
+        }
+      }
+    } else {
+      // no config - parse all values, detecting vectors by format
+      for (const [k, v] of Object.entries(rest)) {
+        const isVector = typeof v === "string" && v.startsWith("[") && v.endsWith("]");
+        doc[k] = parseValue(v, isVector);
+      }
+    }
+
+    // Always include id in document for consistency
+    doc.id = String(id);
+
+    return {
+      id: String(id),
+      index,
+      score: typeof score === "number" ? score : 0,
+      document: doc as unknown as Partial<TDocument>,
+    };
+  },
+};