@axiom-lattice/pg-stores

PostgreSQL stores implementation for the Axiom Lattice framework.

Overview

This package provides PostgreSQL-based store implementations that conform to the store interfaces. It can be used with the StoreLatticeManager from @axiom-lattice/core to register and manage PostgreSQL-based store services.

Features

• ThreadStore: PostgreSQL implementation for thread management
• AssistantStore: PostgreSQL implementation for assistant management
• Migration System: Automatic schema migration with version tracking
• Type Safety: Full TypeScript support with proper type definitions

Installation

pnpm add @axiom-lattice/pg-stores

Prerequisites

• PostgreSQL database (version 12 or higher)
• Node.js 18 or higher

Usage

Basic Usage - Register ThreadStore

import { PostgreSQLThreadStore } from "@axiom-lattice/pg-stores";
import { registerStoreLattice, getStoreLattice } from "@axiom-lattice/core";

// Create and initialize ThreadStore with connection string
const threadStore = new PostgreSQLThreadStore({
  poolConfig: process.env.DATABASE_URL || "postgresql://localhost:5432/mydb",
});

// Ensure initialization (migrations run automatically)
await threadStore.initialize();

// Register to StoreLatticeManager
registerStoreLattice("threads", "thread", threadStore);

// Or use PoolConfig object
const threadStore2 = new PostgreSQLThreadStore({
  poolConfig: {
    host: process.env.DB_HOST || "localhost",
    port: parseInt(process.env.DB_PORT || "5432"),
    database: process.env.DB_NAME || "mydb",
    user: process.env.DB_USER,
    password: process.env.DB_PASSWORD,
  },
});

await threadStore2.initialize();
registerStoreLattice("threads", "thread", threadStore2);

Basic Usage - Register AssistantStore

import { PostgreSQLAssistantStore } from "@axiom-lattice/pg-stores";
import { registerStoreLattice, getStoreLattice } from "@axiom-lattice/core";

// Create and initialize AssistantStore with connection string
const assistantStore = new PostgreSQLAssistantStore({
  poolConfig: process.env.DATABASE_URL || "postgresql://localhost:5432/mydb",
});

// Ensure initialization (migrations run automatically)
await assistantStore.initialize();

// Register to StoreLatticeManager
registerStoreLattice("assistants", "assistant", assistantStore);

// Get and use the store (with type safety)
const storeLattice = getStoreLattice("assistants", "assistant");
const assistantStoreInstance = storeLattice.store; // TypeScript knows this is AssistantStore

// Create an assistant
const assistant = await assistantStoreInstance.createAssistant(
  "assistant-123",
  {
    name: "My Assistant",
    description: "A helpful assistant",
    graphDefinition: {
      /* ... */
    },
  }
);

// Get all assistants
const assistants = await assistantStoreInstance.getAllAssistants();

// Get a specific assistant
const foundAssistant = await assistantStoreInstance.getAssistantById(
  "assistant-123"
);

// Update assistant
const updatedAssistant = await assistantStoreInstance.updateAssistant(
  "assistant-123",
  {
    name: "Updated Name",
  }
);

// Delete assistant
await assistantStoreInstance.deleteAssistant("assistant-123");

// Get and use the store (with type safety) const storeLattice = getStoreLattice("threads", "thread"); const threadStore = storeLattice.store; // TypeScript knows this is ThreadStore

// Create a thread const thread = await threadStore.createThread("assistant-123", "thread-456", { metadata: { title: "My Thread" }, });

// Get threads by assistant const threads = await threadStore.getThreadsByAssistantId("assistant-123");

// Get a specific thread const foundThread = await threadStore.getThreadById( "assistant-123", "thread-456" );

// Update thread const updatedThread = await threadStore.updateThread( "assistant-123", "thread-456", { metadata: { title: "Updated Title" } } );

// Delete thread await threadStore.deleteThread("assistant-123", "thread-456");

### Direct Usage (Without Registration)

```typescript
import { PostgreSQLThreadStore } from "@axiom-lattice/pg-stores";

// Create store instance with connection string (migrations run automatically)
const threadStore = new PostgreSQLThreadStore({
  poolConfig: process.env.DATABASE_URL || "postgresql://localhost:5432/mydb",
  autoMigrate: true, // Default: true
});

// Ensure initialization
await threadStore.initialize();

// When done, dispose the store to close connections
// await threadStore.dispose();

// Use the store
const threads = await threadStore.getThreadsByAssistantId("assistant-123");

Manual Migration Control

import { PostgreSQLThreadStore } from "@axiom-lattice/pg-stores";

// Create store without auto-migration
const threadStore = new PostgreSQLThreadStore({
  poolConfig: process.env.DATABASE_URL || "postgresql://localhost:5432/mydb",
  autoMigrate: false,
});

// Manually run migrations when ready
await threadStore.initialize();

Migration System

The package includes a migration system that automatically:

1. Creates the lattice_schema_migrations table to track applied migrations
2. Applies pending migrations in order
3. Tracks migration versions to prevent duplicate applications

How Migrations Work

Migrations run automatically when you create a store instance (unless autoMigrate: false):

// Migrations run automatically on initialization
const threadStore = new PostgreSQLThreadStore({
  poolConfig: process.env.DATABASE_URL,
  autoMigrate: true, // Default: true
});

await threadStore.initialize(); // Migrations are applied here

Migration Structure

Migrations are defined with:

• version: Sequential version number (must be unique)
• name: Descriptive migration name
• up: Migration function to apply changes (creates/modifies tables)
• down: Optional rollback function (for undoing migrations)

Using MigrationManager Directly

For more control, you can use MigrationManager directly:

import { Pool } from "pg";
import { MigrationManager } from "@axiom-lattice/pg-stores";
import { createThreadsTable } from "@axiom-lattice/pg-stores";

const pool = new Pool({ connectionString: process.env.DATABASE_URL });
const migrationManager = new MigrationManager(pool);

// Register migrations
migrationManager.register(createThreadsTable);

// Apply pending migrations
await migrationManager.migrate();

// Check current version
const version = await migrationManager.getCurrentVersion();

// Rollback last migration (if down migration exists)
await migrationManager.rollback();

Current Migrations

• ThreadStore Version 1: create_threads_table - Creates the threads table with indexes
• AssistantStore Version 1: create_assistants_table - Creates the assistants table with indexes

For detailed migration usage, see MIGRATION_GUIDE.md.

Database Schema

lattice_threads Table

CREATE TABLE 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 INDEX idx_lattice_threads_assistant_id ON lattice_threads(assistant_id);
CREATE INDEX idx_lattice_threads_created_at ON lattice_threads(created_at DESC);

lattice_assistants Table

CREATE TABLE 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 INDEX idx_lattice_assistants_name ON lattice_assistants(name);
CREATE INDEX idx_lattice_assistants_created_at ON lattice_assistants(created_at DESC);

Configuration

The PostgreSQL store can be configured via:

1. Constructor options: Pass poolConfig (connection string or PoolConfig object) and autoMigrate directly
2. Environment variables:
   • DATABASE_URL - PostgreSQL connection string (recommended)
   • DB_HOST - Database host (if using PoolConfig object)
   • DB_PORT - Database port (if using PoolConfig object)
   • DB_USER - Database user (if using PoolConfig object)
   • DB_PASSWORD - Database password (if using PoolConfig object)
   • DB_NAME - Database name (if using PoolConfig object)

PoolConfig Options

You can pass either:

• Connection string: "postgresql://user:password@host:port/database"
• PoolConfig object: { host, port, database, user, password, ... } (see pg PoolConfig for all options)

API Reference

ThreadStore Interface

interface ThreadStore {
  getThreadsByAssistantId(assistantId: string): Promise<Thread[]>;
  getThreadById(
    assistantId: string,
    threadId: string
  ): Promise<Thread | undefined>;
  createThread(
    assistantId: string,
    threadId: string,
    data: CreateThreadRequest
  ): Promise<Thread>;
  updateThread(
    assistantId: string,
    threadId: string,
    updates: Partial<CreateThreadRequest>
  ): Promise<Thread | null>;
  deleteThread(assistantId: string, threadId: string): Promise<boolean>;
  hasThread(assistantId: string, threadId: string): Promise<boolean>;
}

AssistantStore Interface

interface AssistantStore {
  getAllAssistants(): Promise<Assistant[]>;
  getAssistantById(id: string): Promise<Assistant | undefined>;
  createAssistant(id: string, data: CreateAssistantRequest): Promise<Assistant>;
  updateAssistant(
    id: string,
    updates: Partial<CreateAssistantRequest>
  ): Promise<Assistant | null>;
  deleteAssistant(id: string): Promise<boolean>;
  hasAssistant(id: string): Promise<boolean>;
}

License

MIT