@axiom-lattice/pg-stores 1.0.85 → 1.0.86

package/dist/index.mjs

@@ -1,8 +1,8 @@
 // src/index.ts
-import { Pool as Pool24 } from "pg";
+import { Pool as Pool26 } from "pg";
 
 // src/createPgStoreConfig.ts
-import { Pool as Pool21 } from "pg";
+import { Pool as Pool23 } from "pg";
 
 // src/migrations/migration.ts
 var MigrationManager = class {
@@ -1775,7 +1775,7 @@ var PostgreSQLMcpServerConfigStore = class {
         data.name || null,
         data.description || null,
         JSON.stringify(configWithEncryptedEnv),
-        data.selectedTools || [],
+        JSON.stringify(data.selectedTools || []),
         isEnvEncrypted,
         "disconnected",
         nowString,
@@ -1821,7 +1821,7 @@ var PostgreSQLMcpServerConfigStore = class {
       updateData.is_env_encrypted = isEnvEncrypted;
     }
     if (updates.selectedTools !== void 0) {
-      updateData.selected_tools = updates.selectedTools;
+      updateData.selected_tools = JSON.stringify(updates.selectedTools);
     }
     if (updates.status !== void 0) {
       updateData.status = updates.status;
@@ -6889,9 +6889,218 @@ var PostgresSharedResourceStore = class {
   }
 };
 
+// src/stores/PostgreSQLCollectionStore.ts
+import { Pool as Pool21 } from "pg";
+var PostgreSQLCollectionStore = class {
+  constructor(options) {
+    this.initialized = false;
+    this.ownsPool = true;
+    this.initPromise = null;
+    if (options.pool) {
+      this.pool = options.pool;
+      this.ownsPool = false;
+      this.initialized = true;
+      return;
+    }
+    if (typeof options.poolConfig === "string") {
+      this.pool = new Pool21({ connectionString: options.poolConfig });
+    } else if (options.poolConfig) {
+      this.pool = new Pool21(options.poolConfig);
+    } else {
+      throw new Error("Either pool or poolConfig must be provided");
+    }
+    if (options.autoMigrate !== false) {
+      this.initialize().catch((error) => {
+        console.error("Failed to initialize PostgreSQLCollectionStore:", error);
+      });
+    }
+  }
+  async initialize() {
+    if (this.initialized) return;
+    if (this.initPromise) return this.initPromise;
+    this.initPromise = (async () => {
+      this.initialized = true;
+    })();
+    return this.initPromise;
+  }
+  ensureInit() {
+    if (!this.initialized) {
+      throw new Error("PostgreSQLCollectionStore not initialized");
+    }
+  }
+  async getAllCollections(tenantId) {
+    this.ensureInit();
+    const result = await this.pool.query(
+      `SELECT id, name, tenant_id, label, embedding_key, schema, created_at, updated_at
+       FROM lattice_collections WHERE tenant_id = $1 ORDER BY name`,
+      [tenantId]
+    );
+    return result.rows.map(this.rowToCollection);
+  }
+  async getCollectionByName(tenantId, name) {
+    this.ensureInit();
+    const result = await this.pool.query(
+      `SELECT id, name, tenant_id, label, embedding_key, schema, created_at, updated_at
+       FROM lattice_collections WHERE tenant_id = $1 AND name = $2`,
+      [tenantId, name]
+    );
+    return result.rows.length > 0 ? this.rowToCollection(result.rows[0]) : null;
+  }
+  async createCollection(tenantId, data) {
+    this.ensureInit();
+    const schema = JSON.stringify(data.schema || { fields: [] });
+    const result = await this.pool.query(
+      `INSERT INTO lattice_collections (name, tenant_id, label, embedding_key, schema)
+       VALUES ($1, $2, $3, $4, $5)
+       ON CONFLICT (tenant_id, name) DO NOTHING
+       RETURNING id, name, tenant_id, label, embedding_key, schema, created_at, updated_at`,
+      [data.name, tenantId, data.label, data.embeddingKey, schema]
+    );
+    if (result.rows.length === 0) {
+      throw new Error(`Collection "${data.name}" already exists`);
+    }
+    return this.rowToCollection(result.rows[0]);
+  }
+  async updateCollection(tenantId, name, updates) {
+    this.ensureInit();
+    const setClauses = [];
+    const values = [];
+    let paramIndex = 1;
+    if (updates.label !== void 0) {
+      setClauses.push(`label = $${paramIndex++}`);
+      values.push(updates.label);
+    }
+    if (updates.embeddingKey !== void 0) {
+      setClauses.push(`embedding_key = $${paramIndex++}`);
+      values.push(updates.embeddingKey);
+    }
+    if (updates.schema !== void 0) {
+      setClauses.push(`schema = $${paramIndex++}`);
+      values.push(JSON.stringify(updates.schema));
+    }
+    if (setClauses.length === 0) return this.getCollectionByName(tenantId, name);
+    setClauses.push(`updated_at = NOW()`);
+    values.push(tenantId, name);
+    const result = await this.pool.query(
+      `UPDATE lattice_collections SET ${setClauses.join(", ")}
+       WHERE tenant_id = $${paramIndex++} AND name = $${paramIndex}
+       RETURNING id, name, tenant_id, label, embedding_key, schema, created_at, updated_at`,
+      values
+    );
+    return result.rows.length > 0 ? this.rowToCollection(result.rows[0]) : null;
+  }
+  async deleteCollection(tenantId, name) {
+    this.ensureInit();
+    const result = await this.pool.query(
+      `DELETE FROM lattice_collections WHERE tenant_id = $1 AND name = $2`,
+      [tenantId, name]
+    );
+    return (result.rowCount ?? 0) > 0;
+  }
+  async dispose() {
+    if (this.ownsPool && this.pool) {
+      await this.pool.end();
+    }
+  }
+  rowToCollection(row) {
+    return {
+      id: row.id,
+      name: row.name,
+      tenantId: row.tenant_id,
+      label: row.label,
+      embeddingKey: row.embedding_key,
+      schema: typeof row.schema === "string" ? JSON.parse(row.schema) : row.schema,
+      createdAt: row.created_at,
+      updatedAt: row.updated_at
+    };
+  }
+};
+
 // src/createPgStoreConfig.ts
 import { PostgresSaver } from "@langchain/langgraph-checkpoint-postgres";
 
+// src/PGVectorStoreProvider.ts
+import { Pool as Pool22 } from "pg";
+import { PGVectorStore } from "@langchain/community/vectorstores/pgvector";
+import { Document } from "@langchain/core/documents";
+import { embeddingsLatticeManager } from "@axiom-lattice/core";
+function safeTableName(name) {
+  if (!/^[a-z][a-z0-9_-]*$/.test(name)) {
+    throw new Error(`Invalid table name: ${name}`);
+  }
+  return name;
+}
+var PGVectorStoreProvider = class {
+  constructor(pool, connectionString) {
+    this.pool = pool;
+    this.connectionString = connectionString;
+  }
+  async create(params) {
+    const { tableName, embeddingKey } = params;
+    let embeddings;
+    try {
+      embeddings = embeddingsLatticeManager.getEmbeddingsClient(embeddingKey);
+    } catch {
+      try {
+        embeddings = embeddingsLatticeManager.getEmbeddingsClient("default");
+      } catch {
+        throw new Error(
+          `No embedding model registered. Register one first:
+  import { registerEmbeddingsLattice } from "@axiom-lattice/core";
+  registerEmbeddingsLattice("default", new SomeEmbeddings(...));`
+        );
+      }
+    }
+    return PGVectorStore.initialize(embeddings, {
+      postgresConnectionOptions: { connectionString: this.connectionString },
+      tableName: safeTableName(tableName),
+      columns: {
+        idColumnName: "id",
+        vectorColumnName: "vector",
+        contentColumnName: "content",
+        metadataColumnName: "metadata"
+      },
+      distanceStrategy: "cosine"
+    });
+  }
+  async getEntryCount(tableName) {
+    const name = safeTableName(tableName);
+    try {
+      const result = await this.pool.query(`SELECT COUNT(*) as count FROM "${name}"`);
+      return parseInt(result.rows[0]?.count || "0", 10);
+    } catch {
+      return 0;
+    }
+  }
+  async listEntries(tableName) {
+    const name = safeTableName(tableName);
+    try {
+      const result = await this.pool.query(
+        `SELECT content, metadata FROM "${name}" ORDER BY metadata->>'_created_at' DESC`
+      );
+      return result.rows.map((row) => new Document({
+        pageContent: row.content || "",
+        metadata: typeof row.metadata === "string" ? JSON.parse(row.metadata) : row.metadata || {}
+      }));
+    } catch {
+      return [];
+    }
+  }
+  async dropTable(tableName) {
+    const name = safeTableName(tableName);
+    await this.pool.query(`DROP TABLE IF EXISTS "${name}"`);
+  }
+  async dispose(_name) {
+  }
+};
+function createPGVectorStoreProvider(connectionString) {
+  const pool = new Pool22({ connectionString });
+  return {
+    provider: new PGVectorStoreProvider(pool, connectionString),
+    pool
+  };
+}
+
 // src/migrations/skill_migrations.ts
 var createSkillsTable = {
   version: 4,
@@ -7100,9 +7309,43 @@ var createChannelIdentityMappingTables = {
   }
 };
 
+// src/migrations/collection_migrations.ts
+var createCollectionsTable = {
+  version: 137,
+  name: "create_collections_table",
+  up: async (client) => {
+    await client.query(`
+      CREATE TABLE IF NOT EXISTS lattice_collections (
+        id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
+        name VARCHAR(255) NOT NULL,
+        tenant_id VARCHAR(255) NOT NULL DEFAULT 'default',
+        label VARCHAR(255) NOT NULL,
+        embedding_key VARCHAR(255) NOT NULL DEFAULT 'default',
+        schema JSONB DEFAULT '{"fields":[]}',
+        created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+        updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+        UNIQUE(tenant_id, name)
+      )
+    `);
+    await client.query(`
+      CREATE INDEX IF NOT EXISTS idx_lattice_collections_tenant
+      ON lattice_collections(tenant_id)
+    `);
+    await client.query(`
+      CREATE INDEX IF NOT EXISTS idx_lattice_collections_name
+      ON lattice_collections(name)
+    `);
+  },
+  down: async (client) => {
+    await client.query("DROP INDEX IF EXISTS idx_lattice_collections_name");
+    await client.query("DROP INDEX IF EXISTS idx_lattice_collections_tenant");
+    await client.query("DROP TABLE IF EXISTS lattice_collections");
+  }
+};
+
 // src/createPgStoreConfig.ts
 async function createPgStoreConfig(connectionString) {
-  const pool = new Pool21({ connectionString });
+  const pool = new Pool23({ connectionString });
   const mm = new MigrationManager(pool);
   mm.register(createThreadsTable);
   mm.register(createScheduledTasksTable);
@@ -7146,6 +7389,7 @@ async function createPgStoreConfig(connectionString) {
   mm.register(createMenuItemsTable);
   mm.register(createSharedResourcesTable);
   mm.register(addFileContentType);
+  mm.register(createCollectionsTable);
   await mm.migrate();
   const checkpoint = PostgresSaver.fromConnString(connectionString);
   checkpoint.setup().catch((err) => {
@@ -7173,12 +7417,14 @@ async function createPgStoreConfig(connectionString) {
     schedule: new PostgreSQLScheduleStorage(opts),
     menu: new MenuStore(opts),
     sharedResource: new PostgresSharedResourceStore(opts),
+    collection: new PostgreSQLCollectionStore(opts),
+    vectorStoreProvider: new PGVectorStoreProvider(pool, connectionString),
     checkpoint
   };
 }
 
 // src/stores/PostgreSQLSkillStore.ts
-import { Pool as Pool22 } from "pg";
+import { Pool as Pool24 } from "pg";
 var PostgreSQLSkillStore = class {
   constructor(options) {
     this.initialized = false;
@@ -7191,9 +7437,9 @@ var PostgreSQLSkillStore = class {
       return;
     }
     if (typeof options.poolConfig === "string") {
-      this.pool = new Pool22({ connectionString: options.poolConfig });
+      this.pool = new Pool24({ connectionString: options.poolConfig });
     } else if (options.poolConfig) {
-      this.pool = new Pool22(options.poolConfig);
+      this.pool = new Pool24(options.poolConfig);
     } else {
       throw new Error("Either pool or poolConfig must be provided");
     }
@@ -7494,7 +7740,7 @@ var PostgreSQLSkillStore = class {
 };
 
 // src/stores/ChannelIdentityMappingStore.ts
-import { Pool as Pool23 } from "pg";
+import { Pool as Pool25 } from "pg";
 var ChannelIdentityMappingStore = class {
   constructor(options) {
     this.initialized = false;
@@ -7506,7 +7752,7 @@ var ChannelIdentityMappingStore = class {
       this.initialized = true;
       return;
     }
-    this.pool = typeof options.poolConfig === "string" ? new Pool23({ connectionString: options.poolConfig }) : options.poolConfig ? new Pool23(options.poolConfig) : (() => {
+    this.pool = typeof options.poolConfig === "string" ? new Pool25({ connectionString: options.poolConfig }) : options.poolConfig ? new Pool25(options.poolConfig) : (() => {
       throw new Error("Either pool or poolConfig must be provided");
     })();
     this.migrationManager = new MigrationManager(this.pool);
@@ -7729,10 +7975,12 @@ export {
   ChannelIdentityMappingStore,
   MenuStore,
   MigrationManager,
-  Pool24 as Pool,
+  PGVectorStoreProvider,
+  Pool26 as Pool,
   PostgreSQLA2AApiKeyStore,
   PostgreSQLAssistantStore,
   PostgreSQLChannelInstallationStore,
+  PostgreSQLCollectionStore,
   PostgreSQLDatabaseConfigStore,
   PostgreSQLEvalStore,
   PostgreSQLMcpServerConfigStore,
@@ -7764,6 +8012,7 @@ export {
   createChannelBindingsTable,
   createChannelIdentityMappingTables,
   createChannelInstallationsTable,
+  createCollectionsTable,
   createDatabaseConfigsTable,
   createEvalCasesTable,
   createEvalProjectsTable,
@@ -7773,6 +8022,7 @@ export {
   createMcpServerConfigsTable,
   createMenuItemsTable,
   createMetricsConfigsTable,
+  createPGVectorStoreProvider,
   createPgStoreConfig,
   createProjectsTable,
   createScheduledTasksTable,