@axiom-lattice/pg-stores 1.0.1

package/src/stores/PostgreSQLThreadStore.ts

@@ -0,0 +1,276 @@
+/**
+ * PostgreSQL implementation of ThreadStore
+ */
+
+import { Pool, PoolClient, PoolConfig } from "pg";
+import {
+  ThreadStore,
+  Thread,
+  CreateThreadRequest,
+} from "@axiom-lattice/protocols";
+import { MigrationManager } from "../migrations/migration";
+import { createThreadsTable } from "../migrations/thread_migrations";
+
+/**
+ * PostgreSQL ThreadStore options
+ */
+export interface PostgreSQLThreadStoreOptions {
+  /**
+   * 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 ThreadStore
+ */
+export class PostgreSQLThreadStore implements ThreadStore {
+  private pool: Pool;
+  private migrationManager: MigrationManager;
+  private initialized: boolean = false;
+  private ownsPool: boolean = true;
+
+  constructor(options: PostgreSQLThreadStoreOptions) {
+    // 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(createThreadsTable);
+
+    // Auto-migrate by default
+    if (options.autoMigrate !== false) {
+      this.initialize().catch((error) => {
+        console.error("Failed to initialize PostgreSQLThreadStore:", error);
+        throw error;
+      });
+    }
+  }
+
+  /**
+   * 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();
+    }
+  }
+
+  /**
+   * Initialize the store and run migrations
+   */
+  async initialize(): Promise<void> {
+    if (this.initialized) {
+      return;
+    }
+
+    await this.migrationManager.migrate();
+    this.initialized = true;
+  }
+
+  /**
+   * Get all threads for a specific assistant
+   */
+  async getThreadsByAssistantId(assistantId: string): Promise<Thread[]> {
+    await this.ensureInitialized();
+
+    const result = await this.pool.query<{
+      id: string;
+      assistant_id: string;
+      metadata: any;
+      created_at: Date;
+      updated_at: Date;
+    }>(
+      `
+      SELECT id, assistant_id, metadata, created_at, updated_at
+      FROM lattice_threads
+      WHERE assistant_id = $1
+      ORDER BY created_at DESC
+    `,
+      [assistantId]
+    );
+
+    return result.rows.map(this.mapRowToThread);
+  }
+
+  /**
+   * Get a thread by ID for a specific assistant
+   */
+  async getThreadById(
+    assistantId: string,
+    threadId: string
+  ): Promise<Thread | undefined> {
+    await this.ensureInitialized();
+
+    const result = await this.pool.query<{
+      id: string;
+      assistant_id: string;
+      metadata: any;
+      created_at: Date;
+      updated_at: Date;
+    }>(
+      `
+      SELECT id, assistant_id, metadata, created_at, updated_at
+      FROM lattice_threads
+      WHERE id = $1 AND assistant_id = $2
+    `,
+      [threadId, assistantId]
+    );
+
+    if (result.rows.length === 0) {
+      return undefined;
+    }
+
+    return this.mapRowToThread(result.rows[0]);
+  }
+
+  /**
+   * Create a new thread for an assistant
+   */
+  async createThread(
+    assistantId: string,
+    threadId: string,
+    data: CreateThreadRequest
+  ): Promise<Thread> {
+    await this.ensureInitialized();
+
+    const now = new Date();
+    const metadata = data.metadata || {};
+
+    await this.pool.query(
+      `
+      INSERT INTO lattice_threads (id, assistant_id, metadata, created_at, updated_at)
+      VALUES ($1, $2, $3, $4, $5)
+      ON CONFLICT (id, assistant_id) DO UPDATE SET
+        metadata = EXCLUDED.metadata,
+        updated_at = EXCLUDED.updated_at
+    `,
+      [threadId, assistantId, JSON.stringify(metadata), now, now]
+    );
+
+    return {
+      id: threadId,
+      assistantId,
+      metadata,
+      createdAt: now,
+      updatedAt: now,
+    };
+  }
+
+  /**
+   * Update an existing thread
+   */
+  async updateThread(
+    assistantId: string,
+    threadId: string,
+    updates: Partial<CreateThreadRequest>
+  ): Promise<Thread | null> {
+    await this.ensureInitialized();
+
+    // Get existing thread
+    const existing = await this.getThreadById(assistantId, threadId);
+    if (!existing) {
+      return null;
+    }
+
+    // Merge metadata
+    const updatedMetadata = {
+      ...existing.metadata,
+      ...(updates.metadata || {}),
+    };
+
+    const now = new Date();
+
+    await this.pool.query(
+      `
+      UPDATE lattice_threads
+      SET metadata = $1, updated_at = $2
+      WHERE id = $3 AND assistant_id = $4
+    `,
+      [JSON.stringify(updatedMetadata), now, threadId, assistantId]
+    );
+
+    return {
+      ...existing,
+      metadata: updatedMetadata,
+      updatedAt: now,
+    };
+  }
+
+  /**
+   * Delete a thread by ID
+   */
+  async deleteThread(assistantId: string, threadId: string): Promise<boolean> {
+    await this.ensureInitialized();
+
+    const result = await this.pool.query(
+      `
+      DELETE FROM lattice_threads
+      WHERE id = $1 AND assistant_id = $2
+    `,
+      [threadId, assistantId]
+    );
+
+    return result.rowCount !== null && result.rowCount > 0;
+  }
+
+  /**
+   * Check if thread exists
+   */
+  async hasThread(assistantId: string, threadId: string): Promise<boolean> {
+    await this.ensureInitialized();
+
+    const result = await this.pool.query(
+      `
+      SELECT 1 FROM lattice_threads
+      WHERE id = $1 AND assistant_id = $2
+      LIMIT 1
+    `,
+      [threadId, assistantId]
+    );
+
+    return result.rows.length > 0;
+  }
+
+  /**
+   * Ensure store is initialized
+   */
+  private async ensureInitialized(): Promise<void> {
+    if (!this.initialized) {
+      await this.initialize();
+    }
+  }
+
+  /**
+   * Map database row to Thread object
+   */
+  private mapRowToThread(row: {
+    id: string;
+    assistant_id: string;
+    metadata: any;
+    created_at: Date;
+    updated_at: Date;
+  }): Thread {
+    return {
+      id: row.id,
+      assistantId: row.assistant_id,
+      metadata:
+        typeof row.metadata === "string"
+          ? JSON.parse(row.metadata)
+          : row.metadata || {},
+      createdAt: row.created_at,
+      updatedAt: row.updated_at,
+    };
+  }
+}

package/src/types.ts

@@ -0,0 +1,10 @@
+/**
+ * Re-export ThreadStore types from protocols package
+ * This ensures all packages use the same type definitions
+ */
+
+export type {
+  Thread,
+  CreateThreadRequest,
+  ThreadStore,
+} from "@axiom-lattice/protocols";

package/tsconfig.json

@@ -0,0 +1,26 @@
+{
+  "compilerOptions": {
+    "target": "ES2020",
+    "module": "preserve",
+    "lib": ["ES2020"],
+    "outDir": "./dist",
+    "rootDir": "./src",
+    "strict": true,
+    "moduleResolution": "Bundler",
+    "esModuleInterop": true,
+    "skipLibCheck": true,
+    "forceConsistentCasingInFileNames": true,
+    "resolveJsonModule": true,
+    "declaration": true,
+    "declarationMap": true,
+    "types": ["node"],
+    "sourceMap": true,
+    "incremental": true,
+    "tsBuildInfoFile": "./.tsbuildinfo"
+  },
+  "include": ["src/index.ts"],
+  "exclude": ["node_modules", "dist"]
+}
+
+
+