@axiom-lattice/pg-stores 1.0.85 → 1.0.86

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@axiom-lattice/pg-stores",
-  "version": "1.0.85",
+  "version": "1.0.86",
   "description": "PG stores implementation for Axiom Lattice framework",
   "main": "dist/index.js",
   "module": "dist/index.mjs",
@@ -20,11 +20,13 @@
   "author": "",
   "license": "MIT",
   "dependencies": {
+    "@langchain/community": "^0.3.40",
     "@langchain/langgraph-checkpoint-postgres": "1.0.0",
+    "@langchain/core": "1.1.30",
     "pg": "^8.16.3",
     "uuid": "^9.0.1",
-    "@axiom-lattice/core": "2.1.94",
-    "@axiom-lattice/protocols": "2.1.48"
+    "@axiom-lattice/core": "2.1.95",
+    "@axiom-lattice/protocols": "2.1.49"
   },
   "devDependencies": {
     "@types/node": "^20.11.24",

package/src/PGVectorStoreProvider.ts

@@ -0,0 +1,113 @@
+/**
+ * PGVectorStoreProvider
+ *
+ * VectorStoreProvider implementation using PostgreSQL with pgvector extension.
+ * Each created VectorStore maps to a separate table with the vector column.
+ */
+
+import { Pool } from "pg";
+import { PGVectorStore } from "@langchain/community/vectorstores/pgvector";
+import { Document } from "@langchain/core/documents";
+import { embeddingsLatticeManager } from "@axiom-lattice/core";
+import type { VectorStoreProvider, VectorStoreCreateParams } from "@axiom-lattice/protocols";
+import type { DocumentInterface } from "@langchain/core/documents";
+
+function safeTableName(name: string): string {
+  if (!/^[a-z][a-z0-9_-]*$/.test(name)) {
+    throw new Error(`Invalid table name: ${name}`);
+  }
+  return name;
+}
+
+export class PGVectorStoreProvider implements VectorStoreProvider {
+  private pool: Pool;
+  private connectionString: string;
+
+  constructor(pool: Pool, connectionString: string) {
+    this.pool = pool;
+    this.connectionString = connectionString;
+  }
+
+  async create(params: VectorStoreCreateParams): Promise<any> {
+    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:\n` +
+          `  import { registerEmbeddingsLattice } from "@axiom-lattice/core";\n` +
+          `  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" as any,
+    });
+  }
+
+  async getEntryCount(tableName: string): Promise<number> {
+    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: string): Promise<DocumentInterface[]> {
+    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: any) => new Document({
+        pageContent: row.content || "",
+        metadata: typeof row.metadata === "string" ? JSON.parse(row.metadata) : (row.metadata || {}),
+      }));
+    } catch {
+      return [];
+    }
+  }
+
+  async dropTable(tableName: string): Promise<void> {
+    const name = safeTableName(tableName);
+    await this.pool.query(`DROP TABLE IF EXISTS "${name}"`);
+  }
+
+  async dispose(_name: string): Promise<void> {
+    // Pool is shared, don't close it
+  }
+}
+
+/**
+ * Create a PGVectorStoreProvider from a connection string.
+ * Creates its own Pool — call pool.end() to clean up.
+ *
+ * @param connectionString PostgreSQL connection string
+ * @returns PGVectorStoreProvider instance
+ */
+export function createPGVectorStoreProvider(connectionString: string): {
+  provider: PGVectorStoreProvider;
+  pool: Pool;
+} {
+  const pool = new Pool({ connectionString });
+  return {
+    provider: new PGVectorStoreProvider(pool, connectionString),
+    pool,
+  };
+}

package/src/createPgStoreConfig.ts

@@ -20,7 +20,9 @@ import { PostgreSQLScheduleStorage } from "./stores/PostgreSQLScheduleStorage";
 import { PostgreSQLTaskStore } from "./stores/PostgreSQLTaskStore";
 import { MenuStore } from "./stores/MenuStore";
 import { PostgresSharedResourceStore } from "./stores/PostgresSharedResourceStore";
+import { PostgreSQLCollectionStore } from "./stores/PostgreSQLCollectionStore";
 import { PostgresSaver } from "@langchain/langgraph-checkpoint-postgres";
+import { PGVectorStoreProvider } from "./PGVectorStoreProvider";
 
 // Migration imports — all migrations registered centrally to guarantee
 // ordering and prevent name collisions across features.
@@ -62,6 +64,7 @@ import { createTasksTable } from "./migrations/task_migration";
 import { createMenuItemsTable } from "./migrations/menu_items_migration";
 import { addFileContentType } from "./migrations/menu_items_add_file_type";
 import { createSharedResourcesTable } from "./migrations/shared_resources_migration";
+import { createCollectionsTable } from "./migrations/collection_migrations";
 
 export async function createPgStoreConfig(connectionString: string) {
   const pool = new Pool({ connectionString });
@@ -112,6 +115,7 @@ export async function createPgStoreConfig(connectionString: string) {
   mm.register(createMenuItemsTable);             // v134
   mm.register(createSharedResourcesTable);       // v135
   mm.register(addFileContentType);               // v136
+  mm.register(createCollectionsTable);           // v137
 
   await mm.migrate();
 
@@ -143,6 +147,8 @@ export async function createPgStoreConfig(connectionString: string) {
     schedule: new PostgreSQLScheduleStorage(opts),
     menu: new MenuStore(opts),
     sharedResource: new PostgresSharedResourceStore(opts),
+    collection: new PostgreSQLCollectionStore(opts),
+    vectorStoreProvider: new PGVectorStoreProvider(pool, connectionString),
     checkpoint,
   };
 }

package/src/index.ts

@@ -60,9 +60,13 @@ export * from "./stores/PostgreSQLEvalStore";
 export * from "./migrations/eval_migrations";
 export * from "./migrations/menu_items_migration";
 export * from "./migrations/menu_items_add_file_type";
+export * from "./migrations/collection_migrations";
 export * from "./stores/MenuStore";
 export * from "./stores/SharedResourceStore";
 export * from "./stores/PostgresSharedResourceStore";
+export * from "./stores/PostgreSQLCollectionStore";
+
+export * from "./PGVectorStoreProvider";
 
 // Re-export for convenience
 export { PostgreSQLThreadStore } from "./stores/PostgreSQLThreadStore";

package/src/migrations/collection_migrations.ts

@@ -0,0 +1,41 @@
+/**
+ * Collection table migrations
+ */
+
+import { PoolClient } from "pg";
+import { Migration } from "./migration";
+
+export const createCollectionsTable: Migration = {
+  version: 137,
+  name: "create_collections_table",
+  up: async (client: PoolClient) => {
+    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: PoolClient) => {
+    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");
+  },
+};

package/src/stores/PostgreSQLCollectionStore.ts

@@ -0,0 +1,165 @@
+/**
+ * PostgreSQL implementation of CollectionStore with tenant isolation
+ */
+
+import { Pool } from "pg";
+import type { PoolConfig } from "pg";
+import {
+  type CollectionStore,
+  type Collection,
+  type CreateCollectionRequest,
+  type UpdateCollectionRequest,
+} from "@axiom-lattice/protocols";
+
+export interface PostgreSQLCollectionStoreOptions {
+  pool?: Pool;
+  poolConfig?: string | PoolConfig;
+  autoMigrate?: boolean;
+}
+
+export class PostgreSQLCollectionStore implements CollectionStore {
+  private pool: Pool;
+  private initialized: boolean = false;
+  private ownsPool: boolean = true;
+  private initPromise: Promise<void> | null = null;
+
+  constructor(options: PostgreSQLCollectionStoreOptions) {
+    if (options.pool) {
+      this.pool = options.pool;
+      this.ownsPool = false;
+      this.initialized = true;
+      return;
+    }
+
+    if (typeof options.poolConfig === "string") {
+      this.pool = new Pool({ connectionString: options.poolConfig });
+    } else if (options.poolConfig) {
+      this.pool = new Pool(options.poolConfig as 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(): Promise<void> {
+    if (this.initialized) return;
+    if (this.initPromise) return this.initPromise;
+    this.initPromise = (async () => {
+      // Table is created by the centralized migration in createPgStoreConfig
+      this.initialized = true;
+    })();
+    return this.initPromise;
+  }
+
+  private ensureInit(): void {
+    if (!this.initialized) {
+      throw new Error("PostgreSQLCollectionStore not initialized");
+    }
+  }
+
+  async getAllCollections(tenantId: string): Promise<Collection[]> {
+    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: string, name: string): Promise<Collection | null> {
+    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: string, data: CreateCollectionRequest): Promise<Collection> {
+    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: string,
+    name: string,
+    updates: UpdateCollectionRequest
+  ): Promise<Collection | null> {
+    this.ensureInit();
+    const setClauses: string[] = [];
+    const values: (string | number | boolean)[] = [];
+    let paramIndex = 1;
+
+    if (updates.label !== undefined) {
+      setClauses.push(`label = $${paramIndex++}`);
+      values.push(updates.label);
+    }
+    if (updates.embeddingKey !== undefined) {
+      setClauses.push(`embedding_key = $${paramIndex++}`);
+      values.push(updates.embeddingKey);
+    }
+    if (updates.schema !== undefined) {
+      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: string, name: string): Promise<boolean> {
+    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(): Promise<void> {
+    if (this.ownsPool && this.pool) {
+      await this.pool.end();
+    }
+  }
+
+  private rowToCollection(row: any): Collection {
+    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,
+    };
+  }
+}

package/src/stores/PostgreSQLMcpServerConfigStore.ts

@@ -271,7 +271,7 @@ export class PostgreSQLMcpServerConfigStore implements McpServerConfigStore {
         data.name || null,
         data.description || null,
         JSON.stringify(configWithEncryptedEnv),
-        data.selectedTools || [],
+        JSON.stringify(data.selectedTools || []),
         isEnvEncrypted,
         'disconnected',
         nowString,
@@ -331,7 +331,7 @@ export class PostgreSQLMcpServerConfigStore implements McpServerConfigStore {
     }
 
     if (updates.selectedTools !== undefined) {
-      updateData.selected_tools = updates.selectedTools;
+      updateData.selected_tools = JSON.stringify(updates.selectedTools);
     }
 
     if (updates.status !== undefined) {