@kernl-sdk/pg 0.1.10 → 0.1.12

package/src/storage.ts

@@ -2,20 +2,73 @@ import assert from "assert";
 import type { Pool, PoolClient } from "pg";
 
 /* workspace */
-import type { Table, Column, IndexConstraint } from "@kernl-sdk/storage";
-import { SCHEMA_NAME, TABLE_MIGRATIONS } from "@kernl-sdk/storage";
 import type {
   AgentRegistry,
   ModelRegistry,
   KernlStorage,
   Transaction,
 } from "kernl";
+import type { Table, Column, IndexConstraint } from "@kernl-sdk/storage";
+import { KERNL_SCHEMA_NAME, TABLE_MIGRATIONS } from "@kernl-sdk/storage";
 import { UnimplementedError } from "@kernl-sdk/shared/lib";
 
 /* pg */
 import { PGThreadStore } from "./thread/store";
+import { PGMemoryStore } from "./memory/store";
+import { MIGRATIONS } from "./migrations";
 import { SQL_IDENTIFIER_REGEX } from "./sql";
-import { migrations } from "./migrations";
+
+/**
+ * Vector similarity metric for pgvector.
+ */
+export type VectorSimilarity = "cosine" | "euclidean" | "dot_product";
+
+/**
+ * pgvector configuration options.
+ */
+export interface PGVectorConfig {
+  /**
+   * Vector dimensions.
+   * @default 1536 (OpenAI text-embedding-3-small)
+   */
+  dimensions?: number;
+
+  /**
+   * Distance metric for similarity search.
+   * @default "cosine"
+   */
+  similarity?: VectorSimilarity;
+}
+
+/**
+ * Resolved vector configuration with defaults applied.
+ */
+export interface ResolvedVectorConfig {
+  dimensions: number;
+  similarity: VectorSimilarity;
+}
+
+/**
+ * Default vector configuration.
+ */
+export const DEFAULT_VECTOR_CONFIG: ResolvedVectorConfig = {
+  dimensions: 1536,
+  similarity: "cosine",
+};
+
+/**
+ * Resolve vector config, applying defaults.
+ */
+export function resolveVectorConfig(
+  config: boolean | PGVectorConfig | undefined,
+): ResolvedVectorConfig | undefined {
+  if (!config) return undefined;
+  if (config === true) return DEFAULT_VECTOR_CONFIG;
+  return {
+    dimensions: config.dimensions ?? DEFAULT_VECTOR_CONFIG.dimensions,
+    similarity: config.similarity ?? DEFAULT_VECTOR_CONFIG.similarity,
+  };
+}
 
 /**
  * PostgreSQL storage configuration.
@@ -25,19 +78,57 @@ export interface PGStorageConfig {
    * Pool instance for database connections.
    */
   pool: Pool;
+
+  /**
+   * Enable pgvector support for semantic search.
+   *
+   * - `true`: Use default config (1536 dimensions, cosine similarity)
+   * - `PGVectorConfig`: Custom dimensions and similarity metric
+   *
+   * Requires pgvector extension to be installed by superuser:
+   * ```sql
+   * CREATE EXTENSION IF NOT EXISTS vector;
+   * ```
+   */
+  vector?: boolean | PGVectorConfig;
 }
 
 /**
  * PostgreSQL storage adapter.
+ *
+ * Storage is lazily initialized on first use via `ensureInit()`. This means
+ * callers don't need to explicitly call `init()` - it happens automatically.
+ *
+ * NOTE: If the number of store methods grows significantly, consider replacing
+ * the manual `ensureInit()` calls with a Proxy-based wrapper for foolproof
+ * auto-initialization.
  */
 export class PGStorage implements KernlStorage {
   private pool: Pool;
+  private initPromise: Promise<void> | null = null;
 
   threads: PGThreadStore;
+  memories: PGMemoryStore;
 
   constructor(config: PGStorageConfig) {
     this.pool = config.pool;
-    this.threads = new PGThreadStore(this.pool);
+    this.threads = new PGThreadStore(this.pool, () => this.ensureInit());
+    this.memories = new PGMemoryStore(this.pool, () => this.ensureInit());
+  }
+
+  /**
+   * Ensure storage is initialized before any operation.
+   *
+   * Safe to call multiple times - initialization only runs once.
+   */
+  private async ensureInit(): Promise<void> {
+    if (!this.initPromise) {
+      this.initPromise = this.init().catch((err) => {
+        this.initPromise = null;
+        throw err;
+      });
+    }
+    return this.initPromise;
   }
 
   /**
@@ -58,7 +149,7 @@ export class PGStorage implements KernlStorage {
    * Initialize the storage backend.
    */
   async init(): Promise<void> {
-    await this.pool.query(`CREATE SCHEMA IF NOT EXISTS "${SCHEMA_NAME}"`);
+    await this.pool.query(`CREATE SCHEMA IF NOT EXISTS "${KERNL_SCHEMA_NAME}"`);
     await this.createTable(TABLE_MIGRATIONS);
     await this.migrate();
   }
@@ -80,12 +171,12 @@ export class PGStorage implements KernlStorage {
 
       // read applied migration IDs
       const result = await client.query<{ id: string }>(
-        `SELECT id FROM "${SCHEMA_NAME}".migrations ORDER BY applied_at ASC`,
+        `SELECT id FROM "${KERNL_SCHEMA_NAME}".migrations ORDER BY applied_at ASC`,
       );
       const applied = new Set(result.rows.map((row) => row.id));
 
       // filter pending migrations
-      const pending = migrations.filter((m) => !applied.has(m.id));
+      const pending = MIGRATIONS.filter((m) => !applied.has(m.id));
       if (pending.length === 0) {
         await client.query("COMMIT");
         return;
@@ -100,7 +191,7 @@ export class PGStorage implements KernlStorage {
           },
         });
         await client.query(
-          `INSERT INTO "${SCHEMA_NAME}".migrations (id, applied_at) VALUES ($1, $2)`,
+          `INSERT INTO "${KERNL_SCHEMA_NAME}".migrations (id, applied_at) VALUES ($1, $2)`,
           [migration.id, Date.now()],
         );
       }
@@ -151,7 +242,7 @@ export class PGStorage implements KernlStorage {
 
       // foreign key reference
       if (col._fk) {
-        let ref = `REFERENCES "${SCHEMA_NAME}"."${col._fk.table}" ("${col._fk.column}")`;
+        let ref = `REFERENCES "${KERNL_SCHEMA_NAME}"."${col._fk.table}" ("${col._fk.column}")`;
         if (col._onDelete) {
           ref += ` ON DELETE ${col._onDelete}`;
         }
@@ -210,7 +301,7 @@ export class PGStorage implements KernlStorage {
     const constraints = [...columns, ...tableConstraints];
 
     const sql = `
-      CREATE TABLE IF NOT EXISTS "${SCHEMA_NAME}"."${table.name}" (
+      CREATE TABLE IF NOT EXISTS "${KERNL_SCHEMA_NAME}"."${table.name}" (
         ${constraints.join(",\n  ")}
       )
     `.trim();
@@ -262,7 +353,7 @@ export class PGStorage implements KernlStorage {
 
     const sql = `
       CREATE ${uniqueKeyword} INDEX IF NOT EXISTS "${indexName}"
-      ON "${SCHEMA_NAME}"."${tableName}" (${columns})
+      ON "${KERNL_SCHEMA_NAME}"."${tableName}" (${columns})
     `.trim();
 
     await client.query(sql);

package/src/thread/sql.ts

@@ -0,0 +1,159 @@
+/**
+ * Thread SQL conversion codecs.
+ *
+ * TODO: generalize object -> SQL conversion into a shared utility
+ */
+
+import type { Codec } from "@kernl-sdk/shared/lib";
+import type { ThreadFilter, ThreadUpdate, SortOrder } from "kernl";
+
+export interface SQLClause {
+  sql: string;
+  params: unknown[];
+}
+
+export interface WhereInput {
+  filter?: ThreadFilter;
+  startIdx: number;
+}
+
+/**
+ * Encode ThreadFilter to SQL WHERE clause.
+ */
+export const SQL_WHERE: Codec<WhereInput, SQLClause> = {
+  encode({ filter, startIdx }) {
+    if (!filter) {
+      return { sql: "", params: [] };
+    }
+
+    const conditions: string[] = [];
+    const params: unknown[] = [];
+    let idx = startIdx;
+
+    if (filter.namespace !== undefined) {
+      conditions.push(`namespace = $${idx++}`);
+      params.push(filter.namespace);
+    }
+
+    if (filter.state !== undefined) {
+      if (Array.isArray(filter.state)) {
+        conditions.push(`state = ANY($${idx++})`);
+        params.push(filter.state);
+      } else {
+        conditions.push(`state = $${idx++}`);
+        params.push(filter.state);
+      }
+    }
+
+    if (filter.agentId !== undefined) {
+      conditions.push(`agent_id = $${idx++}`);
+      params.push(filter.agentId);
+    }
+
+    if (filter.parentTaskId !== undefined) {
+      conditions.push(`parent_task_id = $${idx++}`);
+      params.push(filter.parentTaskId);
+    }
+
+    if (filter.createdAfter !== undefined) {
+      conditions.push(`created_at > $${idx++}`);
+      params.push(filter.createdAfter.getTime());
+    }
+
+    if (filter.createdBefore !== undefined) {
+      conditions.push(`created_at < $${idx++}`);
+      params.push(filter.createdBefore.getTime());
+    }
+
+    return {
+      sql: conditions.length > 0 ? conditions.join(" AND ") : "",
+      params,
+    };
+  },
+
+  decode() {
+    throw new Error("SQL_WHERE.decode not implemented");
+  },
+};
+
+export interface OrderInput {
+  order?: {
+    createdAt?: SortOrder;
+    updatedAt?: SortOrder;
+  };
+}
+
+/**
+ * Encode order options to SQL ORDER BY clause.
+ */
+export const SQL_ORDER: Codec<OrderInput, string> = {
+  encode({ order }) {
+    const clauses: string[] = [];
+
+    if (order?.createdAt) {
+      clauses.push(`created_at ${order.createdAt.toUpperCase()}`);
+    }
+    if (order?.updatedAt) {
+      clauses.push(`updated_at ${order.updatedAt.toUpperCase()}`);
+    }
+
+    if (clauses.length === 0) {
+      return "created_at DESC";
+    }
+
+    return clauses.join(", ");
+  },
+
+  decode() {
+    throw new Error("SQL_ORDER.decode not implemented");
+  },
+};
+
+export interface UpdateInput {
+  patch: ThreadUpdate;
+  startIdx: number;
+}
+
+/**
+ * Encode ThreadUpdate to SQL SET clause.
+ */
+export const SQL_UPDATE: Codec<UpdateInput, SQLClause> = {
+  encode({ patch, startIdx }) {
+    const sets: string[] = [];
+    const params: unknown[] = [];
+    let idx = startIdx;
+
+    if (patch.tick !== undefined) {
+      sets.push(`tick = $${idx++}`);
+      params.push(patch.tick);
+    }
+
+    if (patch.state !== undefined) {
+      sets.push(`state = $${idx++}`);
+      params.push(patch.state);
+    }
+
+    if (patch.context !== undefined) {
+      sets.push(`context = $${idx++}`);
+      params.push(JSON.stringify(patch.context.context));
+    }
+
+    if (patch.metadata !== undefined) {
+      sets.push(`metadata = $${idx++}`);
+      params.push(patch.metadata ? JSON.stringify(patch.metadata) : null);
+    }
+
+    // always update updated_at
+    sets.push(`updated_at = $${idx++}`);
+    params.push(Date.now());
+
+    return {
+      sql: sets.join(", "),
+      params,
+    };
+  },
+
+  decode() {
+    throw new Error("SQL_UPDATE.decode not implemented");
+  },
+};

package/src/thread/store.ts

@@ -1,14 +1,6 @@
 import assert from "assert";
 import type { Pool, PoolClient } from "pg";
 
-import {
-  SCHEMA_NAME,
-  NewThreadCodec,
-  ThreadEventRecordCodec,
-  type ThreadRecord,
-  type ThreadEventRecord,
-} from "@kernl-sdk/storage";
-import { Thread, type ThreadEvent } from "kernl/internal";
 import {
   Context,
   type AgentRegistry,
@@ -20,16 +12,31 @@ import {
   type ThreadListOptions,
   type ThreadHistoryOptions,
 } from "kernl";
+import { Thread, type ThreadEvent } from "kernl/internal";
+import {
+  KERNL_SCHEMA_NAME,
+  NewThreadCodec,
+  ThreadEventRecordCodec,
+  type ThreadRecord,
+  type ThreadEventRecord,
+} from "@kernl-sdk/storage";
+
+import { SQL_WHERE, SQL_ORDER, SQL_UPDATE } from "./sql";
 
 /**
  * PostgreSQL Thread store implementation.
+ *
+ * IMPORTANT: All async methods must call `await this.ensureInit()` before
+ * any database operations. This ensures schema/tables exist.
  */
 export class PGThreadStore implements ThreadStore {
   private db: Pool | PoolClient;
   private registries: { agents: AgentRegistry; models: ModelRegistry } | null;
+  private ensureInit: () => Promise<void>;
 
-  constructor(db: Pool | PoolClient) {
+  constructor(db: Pool | PoolClient, ensureInit: () => Promise<void>) {
     this.db = db;
+    this.ensureInit = ensureInit;
     this.registries = null;
   }
 
@@ -46,6 +53,8 @@ export class PGThreadStore implements ThreadStore {
    * Get a thread by id.
    */
   async get(tid: string, include?: ThreadInclude): Promise<Thread | null> {
+    await this.ensureInit();
+
     // JOIN with thread_events if include.history
     if (include?.history) {
       const opts =
@@ -78,8 +87,8 @@ export class PGThreadStore implements ThreadStore {
           e.timestamp,
           e.data,
           e.metadata as event_metadata
-        FROM ${SCHEMA_NAME}.threads t
-        LEFT JOIN ${SCHEMA_NAME}.thread_events e ON t.id = e.tid${eventFilter}
+        FROM ${KERNL_SCHEMA_NAME}.threads t
+        LEFT JOIN ${KERNL_SCHEMA_NAME}.thread_events e ON t.id = e.tid${eventFilter}
         WHERE t.id = $1
         ORDER BY e.seq ${order.toUpperCase()}
         ${limit}
@@ -131,7 +140,7 @@ export class PGThreadStore implements ThreadStore {
 
     // simple query without events
     const result = await this.db.query<ThreadRecord>(
-      `SELECT * FROM ${SCHEMA_NAME}.threads WHERE id = $1`,
+      `SELECT * FROM ${KERNL_SCHEMA_NAME}.threads WHERE id = $1`,
       [tid],
     );
 
@@ -150,88 +159,30 @@ export class PGThreadStore implements ThreadStore {
    * List threads matching the filter.
    */
   async list(options?: ThreadListOptions): Promise<Thread[]> {
-    let query = `SELECT * FROM ${SCHEMA_NAME}.threads`;
-    const values: any[] = [];
-    let paramIndex = 1;
-
-    // build WHERE clause
-    const conditions: string[] = [];
-    if (options?.filter) {
-      const {
-        state,
-        agentId,
-        parentTaskId,
-        createdAfter,
-        createdBefore,
-        namespace,
-      } = options.filter;
-
-      if (namespace) {
-        conditions.push(`namespace = $${paramIndex++}`);
-        values.push(namespace);
-      }
+    await this.ensureInit();
 
-      if (state) {
-        if (Array.isArray(state)) {
-          conditions.push(`state = ANY($${paramIndex++})`);
-          values.push(state);
-        } else {
-          conditions.push(`state = $${paramIndex++}`);
-          values.push(state);
-        }
-      }
-
-      if (agentId) {
-        conditions.push(`agent_id = $${paramIndex++}`);
-        values.push(agentId);
-      }
-
-      if (parentTaskId) {
-        conditions.push(`parent_task_id = $${paramIndex++}`);
-        values.push(parentTaskId);
-      }
-
-      if (createdAfter) {
-        conditions.push(`created_at > $${paramIndex++}`);
-        values.push(createdAfter.getTime());
-      }
+    const { sql: where, params } = SQL_WHERE.encode({
+      filter: options?.filter,
+      startIdx: 1,
+    });
 
-      if (createdBefore) {
-        conditions.push(`created_at < $${paramIndex++}`);
-        values.push(createdBefore.getTime());
-      }
-    }
+    let idx = params.length + 1;
+    let query = `SELECT * FROM ${KERNL_SCHEMA_NAME}.threads`;
 
-    if (conditions.length > 0) {
-      query += ` WHERE ${conditions.join(" AND ")}`;
-    }
-
-    // build ORDER BY clause
-    const orderClauses: string[] = [];
-    if (options?.order?.createdAt) {
-      orderClauses.push(`created_at ${options.order.createdAt.toUpperCase()}`);
-    }
-    if (options?.order?.updatedAt) {
-      orderClauses.push(`updated_at ${options.order.updatedAt.toUpperCase()}`);
-    }
-    if (orderClauses.length > 0) {
-      query += ` ORDER BY ${orderClauses.join(", ")}`;
-    } else {
-      // default: most recent first
-      query += ` ORDER BY created_at DESC`;
-    }
+    if (where) query += ` WHERE ${where}`;
+    query += ` ORDER BY ${SQL_ORDER.encode({ order: options?.order })}`;
 
     if (options?.limit) {
-      query += ` LIMIT $${paramIndex++}`;
-      values.push(options.limit);
+      query += ` LIMIT $${idx++}`;
+      params.push(options.limit);
     }
 
     if (options?.offset) {
-      query += ` OFFSET $${paramIndex++}`;
-      values.push(options.offset);
+      query += ` OFFSET $${idx++}`;
+      params.push(options.offset);
     }
 
-    const result = await this.db.query<ThreadRecord>(query, values);
+    const result = await this.db.query<ThreadRecord>(query, params);
     return result.rows
       .map((record) => {
         try {
@@ -250,10 +201,12 @@ export class PGThreadStore implements ThreadStore {
    * Insert a new thread into the store.
    */
   async insert(thread: NewThread): Promise<Thread> {
+    await this.ensureInit();
+
     const record = NewThreadCodec.encode(thread);
 
     const result = await this.db.query<ThreadRecord>(
-      `INSERT INTO ${SCHEMA_NAME}.threads
+      `INSERT INTO ${KERNL_SCHEMA_NAME}.threads
        (id, namespace, agent_id, model, context, tick, state, parent_task_id, metadata, created_at, updated_at)
        VALUES ($1, $2, $3, $4, $5::jsonb, $6, $7, $8, $9::jsonb, $10, $11)
        RETURNING *`,
@@ -279,45 +232,18 @@ export class PGThreadStore implements ThreadStore {
    * Update thread runtime state.
    */
   async update(tid: string, patch: ThreadUpdate): Promise<Thread> {
-    const updates: string[] = [];
-    const values: any[] = [];
-    let paramIndex = 1;
+    await this.ensureInit();
 
-    if (patch.tick !== undefined) {
-      updates.push(`tick = $${paramIndex++}`);
-      values.push(patch.tick);
-    }
-
-    if (patch.state !== undefined) {
-      updates.push(`state = $${paramIndex++}`);
-      values.push(patch.state);
-    }
-
-    if (patch.context !== undefined) {
-      updates.push(`context = $${paramIndex++}`);
-      // NOTE: Store the raw context value, not the Context wrapper. 
-      //
-      // THis may change in the future depending on Context implementation.
-      values.push(JSON.stringify(patch.context.context));
-    }
-
-    if (patch.metadata !== undefined) {
-      updates.push(`metadata = $${paramIndex++}`);
-      values.push(patch.metadata ? JSON.stringify(patch.metadata) : null);
-    }
-
-    // always update `updated_at`
-    updates.push(`updated_at = $${paramIndex++}`);
-    values.push(Date.now());
-
-    values.push(tid); // WHERE id = $N
+    const { sql: updates, params } = SQL_UPDATE.encode({ patch, startIdx: 1 });
+    const idx = params.length + 1;
+    params.push(tid);
 
     const result = await this.db.query<ThreadRecord>(
-      `UPDATE ${SCHEMA_NAME}.threads
-       SET ${updates.join(", ")}
-       WHERE id = $${paramIndex}
+      `UPDATE ${KERNL_SCHEMA_NAME}.threads
+       SET ${updates}
+       WHERE id = $${idx}
        RETURNING *`,
-      values,
+      params,
     );
 
     return this.hydrate({ record: result.rows[0] });
@@ -327,9 +253,11 @@ export class PGThreadStore implements ThreadStore {
    * Delete a thread and cascade to thread_events.
    */
   async delete(tid: string): Promise<void> {
-    await this.db.query(`DELETE FROM ${SCHEMA_NAME}.threads WHERE id = $1`, [
-      tid,
-    ]);
+    await this.ensureInit();
+    await this.db.query(
+      `DELETE FROM ${KERNL_SCHEMA_NAME}.threads WHERE id = $1`,
+      [tid],
+    );
   }
 
   /**
@@ -339,7 +267,9 @@ export class PGThreadStore implements ThreadStore {
     tid: string,
     opts?: ThreadHistoryOptions,
   ): Promise<ThreadEvent[]> {
-    let query = `SELECT * FROM ${SCHEMA_NAME}.thread_events WHERE tid = $1`;
+    await this.ensureInit();
+
+    let query = `SELECT * FROM ${KERNL_SCHEMA_NAME}.thread_events WHERE tid = $1`;
     const values: any[] = [tid];
     let paramIndex = 2;
 
@@ -370,8 +300,7 @@ export class PGThreadStore implements ThreadStore {
     return result.rows.map((record) =>
       ThreadEventRecordCodec.decode({
         ...record,
-        // Normalize BIGINT (string) to number for zod schema
-        timestamp: Number(record.timestamp),
+        timestamp: Number(record.timestamp), // normalize BIGINT (string) to number for zod schema
       } as ThreadEventRecord),
     );
   }
@@ -388,6 +317,7 @@ export class PGThreadStore implements ThreadStore {
    */
   async append(events: ThreadEvent[]): Promise<void> {
     if (events.length === 0) return;
+    await this.ensureInit();
 
     const records = events.map((e) => ThreadEventRecordCodec.encode(e));
 
@@ -412,7 +342,7 @@ export class PGThreadStore implements ThreadStore {
 
     // insert with ON CONFLICT DO NOTHING for idempotency
     await this.db.query(
-      `INSERT INTO ${SCHEMA_NAME}.thread_events
+      `INSERT INTO ${KERNL_SCHEMA_NAME}.thread_events
        (id, tid, seq, kind, timestamp, data, metadata)
        VALUES ${placeholders.join(", ")}
        ON CONFLICT (tid, id) DO NOTHING`,
@@ -434,6 +364,7 @@ export class PGThreadStore implements ThreadStore {
     const agent = this.registries.agents.get(record.agent_id);
     const model = this.registries.models.get(record.model);
 
+    // (TODO): we might want to allow this in the future, unclear how it would look though..
     if (!agent || !model) {
       throw new Error(
         `Thread ${record.id} references non-existent agent/model (agent: ${record.agent_id}, model: ${record.model})`,

package/tsconfig.tsbuildinfo

@@ -0,0 +1 @@
+{"root":["./src/index.ts","./src/migrations.ts","./src/postgres.ts","./src/sql.ts","./src/storage.ts","./src/__tests__/integration.test.ts","./src/__tests__/memory.test.ts","./src/__tests__/thread.test.ts","./src/memory/sql.ts","./src/memory/store.ts","./src/thread/sql.ts","./src/thread/store.ts"],"version":"5.9.2"}