@axiom-lattice/pg-stores 1.0.1

package/src/index.ts

@@ -0,0 +1,25 @@
+/**
+ * PostgreSQL Stores Package
+ *
+ * Provides PostgreSQL implementations for various stores
+ */
+
+export * from "./stores/PostgreSQLThreadStore";
+export * from "./stores/PostgreSQLAssistantStore";
+export * from "./migrations/migration";
+export * from "./migrations/thread_migrations";
+export * from "./migrations/assistant_migrations";
+
+// Re-export for convenience
+export { PostgreSQLThreadStore } from "./stores/PostgreSQLThreadStore";
+export { PostgreSQLAssistantStore } from "./stores/PostgreSQLAssistantStore";
+
+// Re-export types from protocols
+export type {
+  ThreadStore,
+  Thread,
+  CreateThreadRequest,
+  AssistantStore,
+  Assistant,
+  CreateAssistantRequest,
+} from "@axiom-lattice/protocols";

package/src/migrations/assistant_migrations.ts

@@ -0,0 +1,44 @@
+/**
+ * Assistant table migrations
+ */
+
+import { PoolClient } from "pg";
+import { Migration } from "./migration";
+
+/**
+ * Initial migration: Create assistants table
+ */
+export const createAssistantsTable: Migration = {
+  version: 1,
+  name: "create_assistants_table",
+  up: async (client: PoolClient) => {
+    await client.query(`
+      CREATE TABLE IF NOT EXISTS lattice_assistants (
+        id VARCHAR(255) PRIMARY KEY,
+        name VARCHAR(255) NOT NULL,
+        description TEXT,
+        graph_definition JSONB NOT NULL,
+        created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+        updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
+      )
+    `);
+
+    // Create indexes for better query performance
+    await client.query(`
+      CREATE INDEX IF NOT EXISTS idx_lattice_assistants_name 
+      ON lattice_assistants(name)
+    `);
+
+    await client.query(`
+      CREATE INDEX IF NOT EXISTS idx_lattice_assistants_created_at 
+      ON lattice_assistants(created_at DESC)
+    `);
+  },
+  down: async (client: PoolClient) => {
+    await client.query(
+      "DROP INDEX IF EXISTS idx_lattice_assistants_created_at"
+    );
+    await client.query("DROP INDEX IF EXISTS idx_lattice_assistants_name");
+    await client.query("DROP TABLE IF EXISTS lattice_assistants");
+  },
+};

package/src/migrations/migration.ts

@@ -0,0 +1,184 @@
+/**
+ * Migration system for database schema management
+ */
+
+import { Pool, PoolClient } from "pg";
+
+/**
+ * Migration record stored in database
+ */
+interface MigrationRecord {
+  version: number;
+  name: string;
+  applied_at: Date;
+}
+
+/**
+ * Migration definition
+ */
+export interface Migration {
+  version: number;
+  name: string;
+  up: (client: PoolClient) => Promise<void>;
+  down?: (client: PoolClient) => Promise<void>;
+}
+
+/**
+ * Migration manager
+ */
+export class MigrationManager {
+  private pool: Pool;
+  private migrations: Migration[] = [];
+
+  constructor(pool: Pool) {
+    this.pool = pool;
+  }
+
+  /**
+   * Register a migration
+   */
+  register(migration: Migration): void {
+    this.migrations.push(migration);
+    // Sort migrations by version
+    this.migrations.sort((a, b) => a.version - b.version);
+  }
+
+  /**
+   * Initialize migrations table if it doesn't exist
+   */
+  private async ensureMigrationsTable(client: PoolClient): Promise<void> {
+    await client.query(`
+      CREATE TABLE IF NOT EXISTS lattice_schema_migrations (
+        version INTEGER PRIMARY KEY,
+        name VARCHAR(255) NOT NULL,
+        applied_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
+      )
+    `);
+  }
+
+  /**
+   * Get applied migrations from database
+   */
+  private async getAppliedMigrations(
+    client: PoolClient
+  ): Promise<MigrationRecord[]> {
+    await client.query(`
+      CREATE TABLE IF NOT EXISTS lattice_schema_migrations (
+        version INTEGER PRIMARY KEY,
+        name VARCHAR(255) NOT NULL,
+        applied_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
+      )
+    `);
+
+    const result = await client.query<MigrationRecord>(
+      "SELECT version, name, applied_at FROM lattice_schema_migrations ORDER BY version"
+    );
+    return result.rows;
+  }
+
+  /**
+   * Apply pending migrations
+   */
+  async migrate(): Promise<void> {
+    const client = await this.pool.connect();
+    try {
+      await client.query("BEGIN");
+
+      await this.ensureMigrationsTable(client);
+      const appliedMigrations = await this.getAppliedMigrations(client);
+      const appliedVersions = new Set(appliedMigrations.map((m) => m.version));
+
+      // Find pending migrations
+      const pendingMigrations = this.migrations.filter(
+        (m) => !appliedVersions.has(m.version)
+      );
+
+      if (pendingMigrations.length === 0) {
+        console.log("No pending migrations");
+        await client.query("COMMIT");
+        return;
+      }
+
+      // Apply pending migrations
+      for (const migration of pendingMigrations) {
+        console.log(
+          `Applying migration ${migration.version}: ${migration.name}`
+        );
+        await migration.up(client);
+        await client.query(
+          "INSERT INTO lattice_schema_migrations (version, name) VALUES ($1, $2)",
+          [migration.version, migration.name]
+        );
+      }
+
+      await client.query("COMMIT");
+      console.log(`Applied ${pendingMigrations.length} migration(s)`);
+    } catch (error) {
+      await client.query("ROLLBACK");
+      throw error;
+    } finally {
+      client.release();
+    }
+  }
+
+  /**
+   * Rollback last migration
+   */
+  async rollback(): Promise<void> {
+    const client = await this.pool.connect();
+    try {
+      await client.query("BEGIN");
+
+      const appliedMigrations = await this.getAppliedMigrations(client);
+      if (appliedMigrations.length === 0) {
+        console.log("No migrations to rollback");
+        await client.query("COMMIT");
+        return;
+      }
+
+      const lastMigration = appliedMigrations[appliedMigrations.length - 1];
+      const migration = this.migrations.find(
+        (m) => m.version === lastMigration.version
+      );
+
+      if (!migration || !migration.down) {
+        throw new Error(
+          `Migration ${lastMigration.version} does not have a down migration`
+        );
+      }
+
+      console.log(
+        `Rolling back migration ${lastMigration.version}: ${lastMigration.name}`
+      );
+      await migration.down(client);
+      await client.query(
+        "DELETE FROM lattice_schema_migrations WHERE version = $1",
+        [lastMigration.version]
+      );
+
+      await client.query("COMMIT");
+      console.log("Rollback completed");
+    } catch (error) {
+      await client.query("ROLLBACK");
+      throw error;
+    } finally {
+      client.release();
+    }
+  }
+
+  /**
+   * Get current migration version
+   */
+  async getCurrentVersion(): Promise<number> {
+    const client = await this.pool.connect();
+    try {
+      const appliedMigrations = await this.getAppliedMigrations(client);
+      if (appliedMigrations.length === 0) {
+        return 0;
+      }
+      return Math.max(...appliedMigrations.map((m) => m.version));
+    } finally {
+      client.release();
+    }
+  }
+}

package/src/migrations/thread_migrations.ts

@@ -0,0 +1,42 @@
+/**
+ * Thread table migrations
+ */
+
+import { PoolClient } from "pg";
+import { Migration } from "./migration";
+
+/**
+ * Initial migration: Create threads table
+ */
+export const createThreadsTable: Migration = {
+  version: 1,
+  name: "create_threads_table",
+  up: async (client: PoolClient) => {
+    await client.query(`
+      CREATE TABLE IF NOT EXISTS lattice_threads (
+        id VARCHAR(255) NOT NULL,
+        assistant_id VARCHAR(255) NOT NULL,
+        metadata JSONB DEFAULT '{}',
+        created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+        updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+        PRIMARY KEY (id, assistant_id)
+      )
+    `);
+
+    // Create indexes for better query performance
+    await client.query(`
+      CREATE INDEX IF NOT EXISTS idx_lattice_threads_assistant_id 
+      ON lattice_threads(assistant_id)
+    `);
+
+    await client.query(`
+      CREATE INDEX IF NOT EXISTS idx_lattice_threads_created_at 
+      ON lattice_threads(created_at DESC)
+    `);
+  },
+  down: async (client: PoolClient) => {
+    await client.query("DROP INDEX IF EXISTS idx_lattice_threads_created_at");
+    await client.query("DROP INDEX IF EXISTS idx_lattice_threads_assistant_id");
+    await client.query("DROP TABLE IF EXISTS lattice_threads");
+  },
+};

package/src/stores/PostgreSQLAssistantStore.ts

@@ -0,0 +1,303 @@
+/**
+ * PostgreSQL implementation of AssistantStore
+ */
+
+import { Pool, PoolClient, PoolConfig } from "pg";
+import {
+  AssistantStore,
+  Assistant,
+  CreateAssistantRequest,
+} from "@axiom-lattice/protocols";
+import { MigrationManager } from "../migrations/migration";
+import { createAssistantsTable } from "../migrations/assistant_migrations";
+
+/**
+ * PostgreSQL AssistantStore options
+ */
+export interface PostgreSQLAssistantStoreOptions {
+  /**
+   * PostgreSQL connection pool configuration
+   * Can be a connection string or PoolConfig object
+   */
+  poolConfig: string | PoolConfig;
+
+  /**
+   * Whether to run migrations automatically on initialization
+   * @default true
+   */
+  autoMigrate?: boolean;
+}
+
+/**
+ * PostgreSQL implementation of AssistantStore
+ */
+export class PostgreSQLAssistantStore implements AssistantStore {
+  private pool: Pool;
+  private migrationManager: MigrationManager;
+  private initialized: boolean = false;
+  private ownsPool: boolean = true;
+
+  constructor(options: PostgreSQLAssistantStoreOptions) {
+    // Create Pool from config
+    if (typeof options.poolConfig === "string") {
+      this.pool = new Pool({ connectionString: options.poolConfig });
+    } else {
+      this.pool = new Pool(options.poolConfig);
+    }
+
+    this.migrationManager = new MigrationManager(this.pool);
+    this.migrationManager.register(createAssistantsTable);
+
+    // Auto-migrate by default
+    if (options.autoMigrate !== false) {
+      this.initialize().catch((error) => {
+        console.error("Failed to initialize PostgreSQLAssistantStore:", error);
+        throw error;
+      });
+    }
+  }
+
+  /**
+   * Initialize the store and run migrations
+   */
+  async initialize(): Promise<void> {
+    if (this.initialized) {
+      return;
+    }
+
+    await this.migrationManager.migrate();
+    this.initialized = true;
+  }
+
+  /**
+   * Get all assistants
+   */
+  async getAllAssistants(): Promise<Assistant[]> {
+    await this.ensureInitialized();
+
+    const result = await this.pool.query<{
+      id: string;
+      name: string;
+      description: string | null;
+      graph_definition: any;
+      created_at: Date;
+      updated_at: Date;
+    }>(
+      `
+      SELECT id, name, description, graph_definition, created_at, updated_at
+      FROM lattice_assistants
+      ORDER BY created_at DESC
+    `
+    );
+
+    return result.rows.map(this.mapRowToAssistant);
+  }
+
+  /**
+   * Get assistant by ID
+   */
+  async getAssistantById(id: string): Promise<Assistant | null> {
+    await this.ensureInitialized();
+
+    const result = await this.pool.query<{
+      id: string;
+      name: string;
+      description: string | null;
+      graph_definition: any;
+      created_at: Date;
+      updated_at: Date;
+    }>(
+      `
+      SELECT id, name, description, graph_definition, created_at, updated_at
+      FROM lattice_assistants
+      WHERE id = $1
+    `,
+      [id]
+    );
+
+    if (result.rows.length === 0) {
+      return null;
+    }
+
+    return this.mapRowToAssistant(result.rows[0]);
+  }
+
+  /**
+   * Create a new assistant
+   */
+  async createAssistant(
+    id: string,
+    data: CreateAssistantRequest
+  ): Promise<Assistant> {
+    await this.ensureInitialized();
+
+    const now = new Date();
+
+    await this.pool.query(
+      `
+      INSERT INTO lattice_assistants (id, name, description, graph_definition, created_at, updated_at)
+      VALUES ($1, $2, $3, $4, $5, $6)
+      ON CONFLICT (id) DO UPDATE SET
+        name = EXCLUDED.name,
+        description = EXCLUDED.description,
+        graph_definition = EXCLUDED.graph_definition,
+        updated_at = EXCLUDED.updated_at
+    `,
+      [
+        id,
+        data.name,
+        data.description || null,
+        JSON.stringify(data.graphDefinition),
+        now,
+        now,
+      ]
+    );
+
+    return {
+      id,
+      name: data.name,
+      description: data.description,
+      graphDefinition: data.graphDefinition,
+      createdAt: now,
+      updatedAt: now,
+    };
+  }
+
+  /**
+   * Update an existing assistant
+   */
+  async updateAssistant(
+    id: string,
+    updates: Partial<CreateAssistantRequest>
+  ): Promise<Assistant | null> {
+    await this.ensureInitialized();
+
+    // Get existing assistant
+    const existing = await this.getAssistantById(id);
+    if (!existing) {
+      return null;
+    }
+
+    // Build update query dynamically based on provided fields
+    const updateFields: string[] = [];
+    const updateValues: any[] = [];
+    let paramIndex = 1;
+
+    if (updates.name !== undefined) {
+      updateFields.push(`name = $${paramIndex++}`);
+      updateValues.push(updates.name);
+    }
+
+    if (updates.description !== undefined) {
+      updateFields.push(`description = $${paramIndex++}`);
+      updateValues.push(updates.description || null);
+    }
+
+    if (updates.graphDefinition !== undefined) {
+      updateFields.push(`graph_definition = $${paramIndex++}`);
+      updateValues.push(JSON.stringify(updates.graphDefinition));
+    }
+
+    if (updateFields.length === 0) {
+      // No fields to update
+      return existing;
+    }
+
+    // Always update updated_at
+    updateFields.push(`updated_at = $${paramIndex++}`);
+    updateValues.push(new Date());
+
+    // Add id for WHERE clause
+    updateValues.push(id);
+
+    await this.pool.query(
+      `
+      UPDATE lattice_assistants
+      SET ${updateFields.join(", ")}
+      WHERE id = $${paramIndex}
+    `,
+      updateValues
+    );
+
+    // Return updated assistant
+    return await this.getAssistantById(id);
+  }
+
+  /**
+   * Delete an assistant by ID
+   */
+  async deleteAssistant(id: string): Promise<boolean> {
+    await this.ensureInitialized();
+
+    const result = await this.pool.query(
+      `
+      DELETE FROM lattice_assistants
+      WHERE id = $1
+    `,
+      [id]
+    );
+
+    return result.rowCount !== null && result.rowCount > 0;
+  }
+
+  /**
+   * Check if assistant exists
+   */
+  async hasAssistant(id: string): Promise<boolean> {
+    await this.ensureInitialized();
+
+    const result = await this.pool.query(
+      `
+      SELECT 1 FROM lattice_assistants
+      WHERE id = $1
+      LIMIT 1
+    `,
+      [id]
+    );
+
+    return result.rows.length > 0;
+  }
+
+  /**
+   * Dispose resources and close the connection pool
+   * Should be called when the store is no longer needed
+   */
+  async dispose(): Promise<void> {
+    if (this.ownsPool && this.pool) {
+      await this.pool.end();
+    }
+  }
+
+  /**
+   * Ensure store is initialized
+   */
+  private async ensureInitialized(): Promise<void> {
+    if (!this.initialized) {
+      await this.initialize();
+    }
+  }
+
+  /**
+   * Map database row to Assistant object
+   */
+  private mapRowToAssistant(row: {
+    id: string;
+    name: string;
+    description: string | null;
+    graph_definition: any;
+    created_at: Date;
+    updated_at: Date;
+  }): Assistant {
+    return {
+      id: row.id,
+      name: row.name,
+      description: row.description || undefined,
+      graphDefinition:
+        typeof row.graph_definition === "string"
+          ? JSON.parse(row.graph_definition)
+          : row.graph_definition,
+      createdAt: row.created_at,
+      updatedAt: row.updated_at,
+    };
+  }
+}