langchain-ai/langgraphjs
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraphjs.git synced 2026-07-25 21:26:05 -04:00
Code Issues 108 Packages Projects Releases 20 Wiki Activity

feat(checkpoint-postgres): Add PostgresStore for LangGraph.js (#1457)

Browse Source 
Co-authored-by: Tat Dat Duong <david@duong.cz>
This commit is contained in:
Andrzej
2025-07-30 16:08:37 +02:00
committed by GitHub
parent 77b21d534a
commit 42ced3a11d
20 changed files with 3701 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.changeset/nine-plums-teach.md
+5
View File
@@ -0,0 +1,5 @@
---
"@langchain/langgraph-checkpoint-postgres": patch
---
Add Store implemention for Postgres
int-test-deps-docker-compose.yml
+1 -1
View File
@@ -2,7 +2,7 @@ version: "3.8"
services:
postgres:
image: postgres:latest
image: pgvector/pgvector:pg16
container_name: langgraphjs-postgres-test
volumes:
- ./tmp/integration_tests/pgvector:/var/lib/postgresql/data
libs/checkpoint-postgres/.env.example
+7 -1
View File
@@ -3,4 +3,10 @@ LANGCHAIN_TRACING_V2=true
LANGCHAIN_ENDPOINT="https://api.smith.langchain.com"
LANGCHAIN_API_KEY=
LANGCHAIN_PROJECT=
# -----------------------------------------------------
# -----------------------------------------------------
# PostgreSQL connection string for the store
POSTGRES_URL=postgresql://user:password@localhost:5432/database
# Test database connection string
TEST_POSTGRES_URL=postgresql://postgres:postgres@localhost:5434/testdb
libs/checkpoint-postgres/.gitignore
+4
View File
@@ -2,6 +2,10 @@ index.cjs
index.js
index.d.ts
index.d.cts
store.cjs
store.js
store.d.ts
store.d.cts
node_modules
dist
.yarn
libs/checkpoint-postgres/docker-compose.yml
+1 -1
View File
@@ -2,7 +2,7 @@ version: "3.8"
services:
postgres:
image: postgres:latest
image: pgvector/pgvector:pg17
container_name: langgraphjs-postgres-test
environment:
POSTGRES_USER: user
libs/checkpoint-postgres/langchain.config.js
+2 -1
View File
@@ -12,7 +12,8 @@ function abs(relativePath) {
export const config = {
internals: [/node\:/, /@langchain\/core\//, /async_hooks/],
entrypoints: {
index: "index"
index: "index",
store: "store/index",
},
tsConfigPath: resolve("./tsconfig.json"),
cjsSource: "./dist-cjs",
libs/checkpoint-postgres/package.json
+14 -1
View File
@@ -73,6 +73,15 @@
"import": "./index.js",
"require": "./index.cjs"
},
"./store": {
"types": {
"import": "./store.d.ts",
"require": "./store.d.cts",
"default": "./store.d.ts"
},
"import": "./store.js",
"require": "./store.cjs"
},
"./package.json": "./package.json"
},
"files": [
@@ -80,6 +89,10 @@
"index.cjs",
"index.js",
"index.d.ts",
"index.d.cts"
"index.d.cts",
"store.cjs",
"store.js",
"store.d.ts",
"store.d.cts"
]
}
libs/checkpoint-postgres/src/store/index.ts
+652
View File
@@ -0,0 +1,652 @@
import pg from "pg";
import {
BaseStore,
type Operation,
type OperationResults,
type ListNamespacesOperation,
type PutOperation,
type Item,
type MatchCondition,
} from "@langchain/langgraph-checkpoint";
// Import types
import type {
PostgresStoreConfig,
SearchOptions,
SearchItem,
FilterOperators,
} from "./modules/types.js";
// Import modules
import { DatabaseCore } from "./modules/database-core.js";
import { VectorOperations } from "./modules/vector-operations.js";
import { CrudOperations } from "./modules/crud-operations.js";
import { SearchOperations } from "./modules/search-operations.js";
import { TTLManager } from "./modules/ttl-manager.js";
import {
getStoreMigrations,
StoreMigrationConfig,
} from "./store-migrations.js";
import { getStoreTablesWithSchema } from "./sql.js";
export type * from "./modules/types.js";
const { Pool } = pg;
/**
* PostgreSQL implementation of the BaseStore interface.
* This is now a lightweight orchestrator that delegates to specialized modules.
*/
export class PostgresStore extends BaseStore {
private core: DatabaseCore;
private vectorOps: VectorOperations;
private crudOps: CrudOperations;
private searchOps: SearchOperations;
private ttlManager: TTLManager;
private isSetup: boolean = false;
private isClosed: boolean = false;
private ensureTables: boolean;
constructor(config: PostgresStoreConfig) {
super();
// Create connection pool
const pool =
typeof config.connectionOptions === "string"
? new Pool({ connectionString: config.connectionOptions })
: new Pool(config.connectionOptions);
// Initialize core and modules
this.core = new DatabaseCore(
pool,
config.schema || "public",
config.ttl,
config.index,
config.textSearchLanguage
);
this.vectorOps = new VectorOperations(this.core);
this.crudOps = new CrudOperations(this.core, this.vectorOps);
this.searchOps = new SearchOperations(this.core, this.vectorOps);
this.ttlManager = new TTLManager(this.core);
this.ensureTables = config.ensureTables ?? true;
}
/**
* Put an item with optional indexing configuration and TTL.
*/
async put(
namespace: string[],
key: string,
value: Record<string, unknown>,
index?: false | string[],
options?: { ttl?: number }
): Promise<void> {
if (!this.isSetup && this.ensureTables) {
await this.setup();
}
return this.core.withClient(async (client) => {
const operation: PutOperation & { options?: { ttl?: number } } = {
namespace,
key,
value,
index,
options,
};
await this.crudOps.executePut(client, operation);
});
}
/**
* Get an item by namespace and key.
*/
async get(namespace: string[], key: string): Promise<Item | null> {
if (!this.isSetup && this.ensureTables) {
await this.setup();
}
return this.core.withClient(async (client) => {
return this.crudOps.executeGet(client, { namespace, key });
});
}
/**
* Delete an item by namespace and key.
*/
async delete(namespace: string[], key: string): Promise<void> {
if (!this.isSetup && this.ensureTables) {
await this.setup();
}
return this.core.withClient(async (client) => {
const operation: PutOperation = { namespace, key, value: null };
await this.crudOps.executePut(client, operation);
});
}
/**
* List namespaces with optional filtering.
*/
async listNamespaces(
options: {
prefix?: string[];
suffix?: string[];
maxDepth?: number;
limit?: number;
offset?: number;
} = {}
): Promise<string[][]> {
if (!this.isSetup && this.ensureTables) {
await this.setup();
}
const { prefix, suffix, maxDepth, limit = 100, offset = 0 } = options;
// Convert options to match conditions format
const matchConditions: MatchCondition[] = [];
if (prefix) {
matchConditions.push({
matchType: "prefix",
path: prefix,
});
}
if (suffix) {
matchConditions.push({
matchType: "suffix",
path: suffix,
});
}
const operation: ListNamespacesOperation = {
matchConditions,
maxDepth,
limit,
offset,
};
return this.core.withClient(async (client) => {
return this.executeListNamespaces(client, operation);
});
}
/**
* Creates a PostgresStore instance from a connection string.
*/
static fromConnString(
connectionString: string,
options?: Omit<PostgresStoreConfig, "connectionOptions">
): PostgresStore {
return new PostgresStore({
connectionOptions: connectionString,
...options,
});
}
/**
* Initialize the store by running migrations to create necessary tables and indexes.
*/
async setup(): Promise<void> {
if (this.isSetup) return;
await this.runStoreMigrations();
this.isSetup = true;
// Start TTL sweeper if configured
if (this.core.ttlConfig?.sweepIntervalMinutes) {
this.ttlManager.start();
}
}
/**
* Run store migrations to set up the database schema
*/
private async runStoreMigrations(): Promise<void> {
const client = await this.core.pool.connect();
const STORE_TABLES = getStoreTablesWithSchema(this.core.schema);
try {
await client.query(`CREATE SCHEMA IF NOT EXISTS ${this.core.schema}`);
let version = -1;
const migrationConfig: StoreMigrationConfig = {
schema: this.core.schema,
indexConfig: this.core.indexConfig
? {
dims: this.core.indexConfig.dims,
indexType: this.core.indexConfig.indexType,
distanceMetric: this.core.indexConfig.distanceMetric,
createAllMetricIndexes:
this.core.indexConfig.createAllMetricIndexes,
hnsw: this.core.indexConfig.hnsw,
ivfflat: this.core.indexConfig.ivfflat,
}
: undefined,
};
const migrations = getStoreMigrations(migrationConfig);
// Check current migration version using the same pattern as checkpoints
try {
const result = await client.query(
`SELECT v FROM ${STORE_TABLES.store_migrations} ORDER BY v DESC LIMIT 1`
);
if (result.rows.length > 0) {
version = result.rows[0].v;
}
} catch (error: unknown) {
// Assume table doesn't exist if there's an error
if (
typeof error === "object" &&
error !== null &&
"code" in error &&
typeof error.code === "string" &&
error.code === "42P01" // Postgres error code for undefined_table
) {
version = -1;
} else {
throw error;
}
}
// Run migrations starting from the next version
for (let v = version + 1; v < migrations.length; v += 1) {
await client.query(migrations[v]);
await client.query(
`INSERT INTO ${STORE_TABLES.store_migrations} (v) VALUES ($1)`,
[v]
);
}
} finally {
client.release();
}
}
/**
* Execute multiple operations in a single batch.
*/
async batch<Op extends Operation[]>(
operations: Op
): Promise<OperationResults<Op>> {
if (!this.isSetup && this.ensureTables) {
await this.setup();
}
return this.core.withClient(async (client) => {
const results: unknown[] = [];
for (const operation of operations) {
if ("namespacePrefix" in operation) {
// SearchOperation
results.push(await this.searchOps.executeSearch(client, operation));
} else if ("key" in operation && !("value" in operation)) {
// GetOperation
results.push(await this.crudOps.executeGet(client, operation));
} else if ("value" in operation) {
// PutOperation
results.push(await this.crudOps.executePut(client, operation));
} else if ("matchConditions" in operation) {
// ListNamespacesOperation
results.push(await this.executeListNamespaces(client, operation));
} else {
throw new Error(
`Unsupported operation type: ${JSON.stringify(operation)}`
);
}
}
return results as OperationResults<Op>;
});
}
private async executeListNamespaces(
client: pg.PoolClient,
operation: ListNamespacesOperation
): Promise<string[][]> {
const { matchConditions, maxDepth, limit = 100, offset = 0 } = operation;
let sqlQuery = `
SELECT DISTINCT namespace_path
FROM ${this.core.schema}.store
`;
const params: unknown[] = [];
const conditions: string[] = [];
let paramIndex = 1;
// Add match conditions
if (matchConditions && matchConditions.length > 0) {
for (const condition of matchConditions) {
if (condition.matchType === "prefix") {
const prefix = condition.path.join(":");
conditions.push(`namespace_path LIKE $${paramIndex}`);
params.push(`${prefix}%`);
paramIndex += 1;
} else if (condition.matchType === "suffix") {
const suffix = condition.path.join(":");
conditions.push(`namespace_path LIKE $${paramIndex}`);
params.push(`%${suffix}`);
paramIndex += 1;
}
}
}
if (conditions.length > 0) {
sqlQuery += ` WHERE ${conditions.join(" AND ")}`;
}
sqlQuery += ` ORDER BY namespace_path LIMIT $${paramIndex} OFFSET $${
paramIndex + 1
}`;
params.push(limit, offset);
const result = await client.query(sqlQuery, params);
let namespaces = result.rows.map((row) => row.namespace_path.split(":"));
// Apply maxDepth filter if specified
if (maxDepth !== undefined) {
namespaces = namespaces.filter((ns) => ns.length <= maxDepth);
}
return namespaces;
}
/**
* Start the store. Calls setup() if ensureTables is true.
*/
async start(): Promise<void> {
if (this.ensureTables && !this.isSetup) {
await this.setup();
}
}
/**
* Stop the store and close all database connections.
*/
async stop(): Promise<void> {
if (this.isClosed) return;
this.ttlManager.stop();
await this.core.pool.end();
this.isClosed = true;
}
/**
* Manually sweep expired items from the store.
*/
async sweepExpiredItems(): Promise<number> {
if (!this.isSetup && this.ensureTables) {
await this.setup();
}
return this.ttlManager.sweepExpiredItems();
}
/**
* Enhanced search with advanced filtering and similarity scoring.
* @private Internal method used by search.
*/
private async textSearch(
namespacePrefix: string[],
options: SearchOptions = {}
): Promise<SearchItem[]> {
return this.searchOps.textSearch(namespacePrefix, options);
}
/**
* Get statistics about the store.
*/
async getStats(): Promise<{
totalItems: number;
expiredItems: number;
namespaceCount: number;
oldestItem: Date | null;
newestItem: Date | null;
}> {
if (!this.isSetup && this.ensureTables) {
await this.setup();
}
return this.core.withClient(async (client) => {
const result = await client.query(`
SELECT
COUNT(*) as total_items,
COUNT(CASE WHEN expires_at IS NOT NULL AND expires_at <= CURRENT_TIMESTAMP THEN 1 END) as expired_items,
COUNT(DISTINCT namespace_path) as namespace_count,
MIN(created_at) as oldest_item,
MAX(created_at) as newest_item
FROM ${this.core.schema}.store
`);
const row = result.rows[0];
return {
totalItems: parseInt(row.total_items, 10),
expiredItems: parseInt(row.expired_items, 10),
namespaceCount: parseInt(row.namespace_count, 10),
oldestItem: row.oldest_item,
newestItem: row.newest_item,
};
});
}
/**
* Performs vector similarity search using embeddings.
*
* @param namespacePrefix - The namespace prefix to search within
* @param query - The text query to embed and search for similar items
* @param options - Search options including filter, similarity threshold, and distance metric
* @returns Promise resolving to an array of search results with similarity scores
*/
protected async vectorSearch(
namespacePrefix: string[],
query: string,
options: {
filter?: Record<string, unknown>;
limit?: number;
offset?: number;
similarityThreshold?: number;
distanceMetric?: "cosine" | "l2" | "inner_product";
} = {}
): Promise<SearchItem[]> {
if (!this.isSetup && this.ensureTables) {
await this.setup();
}
if (!this.core.indexConfig) {
throw new Error(
"Vector search not configured. Please provide an IndexConfig when creating the store."
);
}
return this.searchOps.vectorSearch(namespacePrefix, query, options);
}
/**
* Performs hybrid search combining vector similarity and text search.
*
* @param namespacePrefix - The namespace prefix to search within
* @param query - The text query to search for
* @param options - Search options including filter, vector weight, and similarity threshold
* @returns Promise resolving to an array of search results with combined similarity scores
*/
protected async hybridSearch(
namespacePrefix: string[],
query: string,
options: {
filter?: Record<string, unknown>;
limit?: number;
offset?: number;
vectorWeight?: number;
similarityThreshold?: number;
} = {}
): Promise<SearchItem[]> {
if (!this.isSetup && this.ensureTables) {
await this.setup();
}
if (!this.core.indexConfig) {
throw new Error(
"Vector search not configured. Please provide an IndexConfig when creating the store."
);
}
return this.searchOps.hybridSearch(namespacePrefix, query, options);
}
/**
* Search for items in the store with support for text search, vector search, and filtering.
*
* @param namespacePrefix - The namespace prefix to search within
* @param options - Search options including search mode, filters, query text, and pagination
* @returns Promise resolving to an array of search results with optional similarity scores
*
* @example
* ```typescript
* // Basic text search
* const results = await store.search(["documents"], {
* query: "machine learning",
* mode: "text"
* });
*
* // Vector search
* const results = await store.search(["documents"], {
* query: "machine learning",
* mode: "vector",
* similarityThreshold: 0.7
* });
*
* // Hybrid search (combining vector and text)
* const results = await store.search(["documents"], {
* query: "machine learning",
* mode: "hybrid",
* vectorWeight: 0.7
* });
*
* // Filtered search
* const results = await store.search(["products"], {
* filter: { category: "electronics", price: { $lt: 100 } }
* });
* ```
*/
async search(
namespacePrefix: string[],
options: {
/**
* Filter conditions with support for advanced operators.
*/
filter?: Record<
string,
string | number | boolean | null | FilterOperators
>;
/**
* Natural language search query.
*/
query?: string;
/**
* Maximum number of results to return.
* @default 10
*/
limit?: number;
/**
* Number of results to skip for pagination.
* @default 0
*/
offset?: number;
/**
* Whether to refresh TTL for returned items.
*/
refreshTtl?: boolean;
/**
* Search mode.
* @default "auto"
*/
mode?: "text" | "vector" | "hybrid" | "auto";
/**
* Similarity threshold for vector search.
*/
similarityThreshold?: number;
/**
* Distance metric for vector search.
* @default "cosine"
*/
distanceMetric?: "cosine" | "l2" | "inner_product";
/**
* Weight for vector search in hybrid mode.
* @default 0.7
*/
vectorWeight?: number;
} = {}
): Promise<SearchItem[]> {
if (!this.isSetup && this.ensureTables) {
await this.setup();
}
const { mode = "auto", query, ...restOptions } = options;
// No query provided - just do metadata filtering
if (!query) {
return this.textSearch(namespacePrefix, restOptions);
}
const hasVectorSearch = Boolean(this.core.indexConfig);
// Determine search mode based on configuration and options
let effectiveMode = mode;
if (mode === "auto") {
effectiveMode = hasVectorSearch ? "vector" : "text";
}
// Execute appropriate search based on mode
switch (effectiveMode) {
case "vector":
if (!hasVectorSearch) {
throw new Error(
"Vector search requested but not configured. Please provide an IndexConfig when creating the store."
);
}
return this.vectorSearch(namespacePrefix, query, {
...restOptions,
similarityThreshold: options.similarityThreshold,
distanceMetric: options.distanceMetric,
});
case "hybrid":
if (!hasVectorSearch) {
throw new Error(
"Hybrid search requested but vector search not configured. Please provide an IndexConfig when creating the store."
);
}
return this.hybridSearch(namespacePrefix, query, {
...restOptions,
vectorWeight: options.vectorWeight,
similarityThreshold: options.similarityThreshold,
});
case "text":
return this.textSearch(namespacePrefix, { query, ...restOptions });
default:
throw new Error(`Unknown search mode: ${mode}`);
}
}
}
libs/checkpoint-postgres/src/store/modules/crud-operations.ts
+103
View File
@@ -0,0 +1,103 @@
import pg from "pg";
import {
type Item,
type GetOperation,
type PutOperation,
} from "@langchain/langgraph-checkpoint";
import { DatabaseCore } from "./database-core.js";
import { VectorOperations } from "./vector-operations.js";
import { PutOptions } from "./types.js";
import { validateNamespace } from "./utils.js";
/**
* Handles basic CRUD operations: get, put, delete.
*/
export class CrudOperations {
constructor(
private core: DatabaseCore,
private vectorOps: VectorOperations
) {}
async executeGet(
client: pg.PoolClient,
operation: GetOperation
): Promise<Item | null> {
validateNamespace(operation.namespace);
const namespacePath = operation.namespace.join(":");
const result = await client.query(
`
SELECT namespace_path, key, value, created_at, updated_at
FROM ${this.core.schema}.store
WHERE namespace_path = $1 AND key = $2
AND (expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP)
`,
[namespacePath, operation.key]
);
if (result.rows.length === 0) return null;
const row = result.rows[0];
// Refresh TTL if configured
if (this.core.ttlConfig?.refreshOnRead) {
await this.core.refreshTtl(client, namespacePath, operation.key);
}
return {
namespace: row.namespace_path.split(":"),
key: row.key,
value: row.value,
createdAt: row.created_at,
updatedAt: row.updated_at,
};
}
async executePut(
client: pg.PoolClient,
operation: PutOperation & { options?: PutOptions }
): Promise<void> {
validateNamespace(operation.namespace);
const { namespace, key, value, options, index } = operation;
const namespacePath = namespace.join(":");
if (value === null) {
// Delete operation
await client.query(
`
DELETE FROM ${this.core.schema}.store
WHERE namespace_path = $1 AND key = $2
`,
[namespacePath, key]
);
} else {
const expiresAt = this.core.calculateExpiresAt(options?.ttl);
await client.query(
`
INSERT INTO ${this.core.schema}.store (namespace_path, key, value, expires_at)
VALUES ($1, $2, $3, $4)
ON CONFLICT (namespace_path, key)
DO UPDATE SET
value = $3,
expires_at = $4,
updated_at = CURRENT_TIMESTAMP
`,
[namespacePath, key, JSON.stringify(value), expiresAt]
);
// Handle vector indexing if configured
if (this.core.indexConfig && index !== false) {
await this.vectorOps.indexItemVectors(
client,
namespacePath,
key,
value,
index // Pass the index parameter to control which fields get indexed
);
}
}
}
}
libs/checkpoint-postgres/src/store/modules/database-core.ts
+69
View File
@@ -0,0 +1,69 @@
import pg from "pg";
import { TTLConfig, IndexConfig } from "./types.js";
/**
* Core database operations and connection management.
* Shared by all modules to avoid duplication.
*/
export class DatabaseCore {
public readonly pool: pg.Pool;
public readonly schema: string;
public readonly ttlConfig?: TTLConfig;
public readonly indexConfig?: IndexConfig;
public readonly textSearchLanguage: string;
constructor(
pool: pg.Pool,
schema: string,
ttlConfig?: TTLConfig,
indexConfig?: IndexConfig,
textSearchLanguage?: string
) {
this.pool = pool;
this.schema = schema;
this.ttlConfig = ttlConfig;
this.indexConfig = indexConfig;
this.textSearchLanguage = textSearchLanguage || "english";
}
async withClient<T>(
operation: (client: pg.PoolClient) => Promise<T>
): Promise<T> {
const client = await this.pool.connect();
try {
return await operation(client);
} finally {
client.release();
}
}
calculateExpiresAt(ttl?: number): Date | null {
const effectiveTtl = ttl ?? this.ttlConfig?.defaultTtl;
if (!effectiveTtl) return null;
return new Date(Date.now() + effectiveTtl * 60 * 1000);
}
async refreshTtl(
client: pg.PoolClient,
namespacePath: string,
key: string
): Promise<void> {
if (!this.ttlConfig?.refreshOnRead) return;
const expiresAt = this.calculateExpiresAt();
if (expiresAt) {
await client.query(
`
UPDATE ${this.schema}.store
SET expires_at = $3, updated_at = CURRENT_TIMESTAMP
WHERE namespace_path = $1 AND key = $2
`,
[namespacePath, key, expiresAt]
);
}
}
}
libs/checkpoint-postgres/src/store/modules/query-builder.ts
+146
View File
@@ -0,0 +1,146 @@
/**
* Shared query building logic for CRUD and search operations.
*/
export class QueryBuilder {
static buildFilterConditions(
filter: Record<string, unknown>,
params: unknown[],
paramIndex: number
): { conditions: string[]; newParamIndex: number } {
const conditions: string[] = [];
let currentParamIndex = paramIndex;
for (const [key, value] of Object.entries(filter)) {
if (
typeof value === "object" &&
value !== null &&
!Array.isArray(value)
) {
const operators = Object.keys(value);
const isOperatorObject = operators.some((op) => op.startsWith("$"));
if (isOperatorObject) {
// Handle advanced operators
for (const [operator, operatorValue] of Object.entries(value)) {
const result = this.buildOperatorCondition(
key,
operator,
operatorValue,
params,
currentParamIndex
);
if (result.condition) {
conditions.push(result.condition);
currentParamIndex = result.newParamIndex;
}
}
} else {
// Handle nested object queries
conditions.push(`value @> $${currentParamIndex}::jsonb`);
params.push(JSON.stringify({ [key]: value }));
currentParamIndex += 1;
}
} else {
// Handle simple value queries
conditions.push(
`value ->> $${currentParamIndex} = $${currentParamIndex + 1}`
);
params.push(key, String(value));
currentParamIndex += 2;
}
}
return { conditions, newParamIndex: currentParamIndex };
}
private static buildOperatorCondition(
key: string,
operator: string,
operatorValue: unknown,
params: unknown[],
paramIndex: number
): { condition?: string; newParamIndex: number } {
switch (operator) {
case "$eq":
params.push(key, String(operatorValue));
return {
condition: `value ->> $${paramIndex} = $${paramIndex + 1}`,
newParamIndex: paramIndex + 2,
};
case "$ne":
params.push(key, String(operatorValue));
return {
condition: `value ->> $${paramIndex} != $${paramIndex + 1}`,
newParamIndex: paramIndex + 2,
};
case "$gt":
params.push(key, operatorValue);
return {
condition: `(value ->> $${paramIndex})::numeric > $${paramIndex + 1}`,
newParamIndex: paramIndex + 2,
};
case "$gte":
params.push(key, operatorValue);
return {
condition: `(value ->> $${paramIndex})::numeric >= $${
paramIndex + 1
}`,
newParamIndex: paramIndex + 2,
};
case "$lt":
params.push(key, operatorValue);
return {
condition: `(value ->> $${paramIndex})::numeric < $${paramIndex + 1}`,
newParamIndex: paramIndex + 2,
};
case "$lte":
params.push(key, operatorValue);
return {
condition: `(value ->> $${paramIndex})::numeric <= $${
paramIndex + 1
}`,
newParamIndex: paramIndex + 2,
};
case "$in":
if (Array.isArray(operatorValue) && operatorValue.length > 0) {
const placeholders = operatorValue.map(
(_, i) => `$${paramIndex + 1 + i}`
);
params.push(key, ...operatorValue.map(String));
return {
condition: `value ->> $${paramIndex} = ANY(ARRAY[${placeholders.join(
","
)}])`,
newParamIndex: paramIndex + 1 + operatorValue.length,
};
}
break;
case "$nin":
if (Array.isArray(operatorValue) && operatorValue.length > 0) {
const placeholders = operatorValue.map(
(_, i) => `$${paramIndex + 1 + i}`
);
params.push(key, ...operatorValue.map(String));
return {
condition: `value ->> $${paramIndex} != ALL(ARRAY[${placeholders.join(
","
)}])`,
newParamIndex: paramIndex + 1 + operatorValue.length,
};
}
break;
case "$exists":
params.push(key);
return {
condition: operatorValue
? `value ? $${paramIndex}`
: `NOT (value ? $${paramIndex})`,
newParamIndex: paramIndex + 1,
};
default:
// Unknown operator, ignore
break;
}
return { newParamIndex: paramIndex };
}
}
libs/checkpoint-postgres/src/store/modules/search-operations.ts
+491
View File
@@ -0,0 +1,491 @@
import pg from "pg";
import {
type Item,
type SearchOperation,
} from "@langchain/langgraph-checkpoint";
import { DatabaseCore } from "./database-core.js";
import { VectorOperations } from "./vector-operations.js";
import { QueryBuilder } from "./query-builder.js";
import { SearchOptions, SearchItem } from "./types.js";
import { validateNamespace } from "./utils.js";
/**
* Handles all search operations: basic search, vector search, hybrid search.
*/
export class SearchOperations {
constructor(
private core: DatabaseCore,
private vectorOps: VectorOperations
) {}
async executeSearch(
client: pg.PoolClient,
operation: SearchOperation
): Promise<Item[]> {
validateNamespace(operation.namespacePrefix);
const namespacePath = operation.namespacePrefix.join(":");
const { filter, limit = 10, offset = 0, query } = operation;
// If vector search is configured and query is provided, use vector similarity search
if (query && this.core.indexConfig) {
return this.executeVectorSearch(client, operation);
}
// If query is provided but no vector search, use text search for better results
if (query && !this.core.indexConfig) {
const results = await this.textSearch(operation.namespacePrefix, {
query,
filter,
limit,
offset,
});
// Convert SearchItem[] to Item[] for consistent return type
return results.map((item) => ({
namespace: item.namespace,
key: item.key,
value: item.value,
createdAt: item.createdAt,
updatedAt: item.updatedAt,
}));
}
// Basic metadata search without query
let sqlQuery = `
SELECT namespace_path, key, value, created_at, updated_at
FROM ${this.core.schema}.store
WHERE namespace_path LIKE $1
AND (expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP)
`;
const params: unknown[] = [`${namespacePath}%`];
let paramIndex = 2;
// Add filter conditions using advanced filtering
if (filter && Object.keys(filter).length > 0) {
const { conditions, newParamIndex } = QueryBuilder.buildFilterConditions(
filter,
params,
paramIndex
);
if (conditions.length > 0) {
sqlQuery += ` AND (${conditions.join(" AND ")})`;
paramIndex = newParamIndex;
}
}
sqlQuery += ` ORDER BY created_at DESC LIMIT $${paramIndex} OFFSET $${
paramIndex + 1
}`;
params.push(limit, offset);
const result = await client.query(sqlQuery, params);
return result.rows.map((row) => ({
namespace: row.namespace_path.split(":"),
key: row.key,
value: row.value,
createdAt: row.created_at,
updatedAt: row.updated_at,
}));
}
async executeVectorSearch(
client: pg.PoolClient,
operation: SearchOperation
): Promise<Item[]> {
if (!this.core.indexConfig || !operation.query) return [];
validateNamespace(operation.namespacePrefix);
const namespacePath = operation.namespacePrefix.join(":");
const { filter, limit = 10, offset = 0, query } = operation;
// Generate query embedding
const queryEmbedding = await this.vectorOps.generateQueryEmbedding(query);
if (queryEmbedding.length !== this.core.indexConfig.dims) {
throw new Error(
`Query embedding dimension mismatch: expected ${this.core.indexConfig.dims}, got ${queryEmbedding.length}`
);
}
let sqlQuery = `
SELECT DISTINCT
s.namespace_path,
s.key,
s.value,
s.created_at,
s.updated_at,
MIN(v.embedding <=> $2) as similarity_score
FROM ${this.core.schema}.store s
JOIN ${this.core.schema}.store_vectors v ON s.namespace_path = v.namespace_path AND s.key = v.key
WHERE s.namespace_path LIKE $1
AND (s.expires_at IS NULL OR s.expires_at > CURRENT_TIMESTAMP)
`;
const params: unknown[] = [
`${namespacePath}%`,
`[${queryEmbedding.join(",")}]`,
];
let paramIndex = 3;
// Add filter conditions
if (filter && Object.keys(filter).length > 0) {
const { conditions, newParamIndex } = QueryBuilder.buildFilterConditions(
filter,
params,
paramIndex
);
if (conditions.length > 0) {
// Adjust conditions to use 's.' prefix for store table columns
const adjustedConditions = conditions.map((condition: string) =>
condition.replace(/value ->/g, "s.value ->")
);
sqlQuery += ` AND (${adjustedConditions.join(" AND ")})`;
paramIndex = newParamIndex;
}
}
sqlQuery += `
GROUP BY s.namespace_path, s.key, s.value, s.created_at, s.updated_at
ORDER BY similarity_score ASC
LIMIT $${paramIndex} OFFSET $${paramIndex + 1}
`;
params.push(limit, offset);
const result = await client.query(sqlQuery, params);
return result.rows.map((row) => ({
namespace: row.namespace_path.split(":"),
key: row.key,
value: row.value,
createdAt: row.created_at,
updatedAt: row.updated_at,
score: 1 - parseFloat(row.similarity_score), // Convert distance to similarity
}));
}
async textSearch(
namespacePrefix: string[],
options: SearchOptions = {}
): Promise<SearchItem[]> {
return this.core.withClient(async (client) => {
validateNamespace(namespacePrefix);
const namespacePath = namespacePrefix.join(":");
const { filter, limit = 10, offset = 0, query, refreshTtl } = options;
let sqlQuery = `
SELECT
namespace_path,
key,
value,
created_at,
updated_at,
CASE
WHEN $2::text IS NOT NULL THEN
ts_rank(to_tsvector($3::regconfig, value::text), plainto_tsquery($3::regconfig, $2::text))
ELSE 0
END as score
FROM ${this.core.schema}.store
WHERE namespace_path LIKE $1
AND (expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP)
`;
const params: unknown[] = [
`${namespacePath}%`,
query || null,
this.core.textSearchLanguage,
];
let paramIndex = 4;
// Add filter conditions using advanced filtering
if (filter && Object.keys(filter).length > 0) {
const { conditions, newParamIndex } =
QueryBuilder.buildFilterConditions(filter, params, paramIndex);
if (conditions.length > 0) {
sqlQuery += ` AND (${conditions.join(" AND ")})`;
paramIndex = newParamIndex;
}
}
// Add full-text search if query is provided
if (query) {
sqlQuery += ` AND (
to_tsvector($3::regconfig, value::text) @@ plainto_tsquery($3::regconfig, $2::text)
OR value::text ILIKE $${paramIndex}
)`;
params.push(`%${query}%`);
paramIndex += 1;
}
// Add ordering by score if query provided, otherwise by updated_at
if (query) {
sqlQuery += ` ORDER BY score DESC, updated_at DESC`;
} else {
sqlQuery += ` ORDER BY updated_at DESC`;
}
sqlQuery += ` LIMIT $${paramIndex} OFFSET $${paramIndex + 1}`;
params.push(limit, offset);
const result = await client.query(sqlQuery, params);
const items: SearchItem[] = result.rows.map((row) => ({
namespace: row.namespace_path.split(":"),
key: row.key,
value: row.value,
createdAt: row.created_at || new Date(),
updatedAt: row.updated_at || new Date(),
score: row.score || undefined,
}));
// Refresh TTL for returned items if requested
if (refreshTtl || this.core.ttlConfig?.refreshOnRead) {
for (const item of items) {
await this.core.refreshTtl(
client,
item.namespace.join(":"),
item.key
);
}
}
return items;
});
}
async vectorSearch(
namespacePrefix: string[],
query: string,
options: {
filter?: Record<string, unknown>;
limit?: number;
offset?: number;
similarityThreshold?: number;
distanceMetric?: "cosine" | "l2" | "inner_product";
} = {}
): Promise<SearchItem[]> {
if (!this.core.indexConfig) {
throw new Error(
"Vector search not configured. Please provide an IndexConfig when creating the store."
);
}
return this.core.withClient(async (client) => {
validateNamespace(namespacePrefix);
const namespacePath = namespacePrefix.join(":");
const {
filter,
limit = 10,
offset = 0,
similarityThreshold = 0.0,
distanceMetric = "cosine",
} = options;
// Generate query embedding
const queryEmbedding = await this.vectorOps.generateQueryEmbedding(query);
if (queryEmbedding.length !== this.core.indexConfig!.dims) {
throw new Error(
`Query embedding dimension mismatch: expected ${
this.core.indexConfig!.dims
}, got ${queryEmbedding.length}`
);
}
// Choose distance operator based on metric
let distanceOp: string;
let scoreTransform: string;
switch (distanceMetric) {
case "l2":
distanceOp = "<->";
scoreTransform = "1 / (1 + MIN(v.embedding <-> $2))"; // Convert L2 distance to similarity
break;
case "inner_product":
distanceOp = "<#>";
scoreTransform = "MIN(v.embedding <#> $2)"; // Inner product (higher is better)
break;
case "cosine":
default:
distanceOp = "<=>";
scoreTransform = "1 - MIN(v.embedding <=> $2)"; // Convert cosine distance to similarity
break;
}
let sqlQuery = `
SELECT DISTINCT
s.namespace_path,
s.key,
s.value,
s.created_at,
s.updated_at,
${scoreTransform} as similarity_score
FROM ${this.core.schema}.store s
JOIN ${this.core.schema}.store_vectors v ON s.namespace_path = v.namespace_path AND s.key = v.key
WHERE s.namespace_path LIKE $1
AND (s.expires_at IS NULL OR s.expires_at > CURRENT_TIMESTAMP)
`;
const params: unknown[] = [
`${namespacePath}%`,
`[${queryEmbedding.join(",")}]`,
];
let paramIndex = 3;
// Add similarity threshold
if (similarityThreshold > 0) {
if (distanceMetric === "inner_product") {
sqlQuery += ` AND v.embedding <#> $2 >= $${paramIndex}`;
} else {
sqlQuery += ` AND v.embedding ${distanceOp} $2 <= $${paramIndex}`;
}
params.push(
distanceMetric === "cosine"
? 1 - similarityThreshold
: similarityThreshold
);
paramIndex += 1;
}
// Add filter conditions
if (filter && Object.keys(filter).length > 0) {
const { conditions, newParamIndex } =
QueryBuilder.buildFilterConditions(filter, params, paramIndex);
if (conditions.length > 0) {
const adjustedConditions = conditions.map((condition: string) =>
condition.replace(/value ->/g, "s.value ->")
);
sqlQuery += ` AND (${adjustedConditions.join(" AND ")})`;
paramIndex = newParamIndex;
}
}
const orderDirection = "DESC"; // From most similar to least similar
sqlQuery += `
GROUP BY s.namespace_path, s.key, s.value, s.created_at, s.updated_at
ORDER BY similarity_score ${orderDirection}
LIMIT $${paramIndex} OFFSET $${paramIndex + 1}
`;
params.push(limit, offset);
const result = await client.query(sqlQuery, params);
return result.rows.map((row) => ({
namespace: row.namespace_path.split(":"),
key: row.key,
value: row.value,
createdAt: row.created_at || new Date(),
updatedAt: row.updated_at || new Date(),
score: parseFloat(row.similarity_score),
}));
});
}
async hybridSearch(
namespacePrefix: string[],
query: string,
options: {
filter?: Record<string, unknown>;
limit?: number;
offset?: number;
vectorWeight?: number; // 0.0 to 1.0, weight for vector search vs text search
similarityThreshold?: number;
} = {}
): Promise<SearchItem[]> {
if (!this.core.indexConfig) {
throw new Error(
"Vector search not configured. Please provide an IndexConfig when creating the store."
);
}
return this.core.withClient(async (client) => {
validateNamespace(namespacePrefix);
const namespacePath = namespacePrefix.join(":");
const {
filter,
limit = 10,
offset = 0,
vectorWeight = 0.7,
similarityThreshold = 0.0,
} = options;
// Generate query embedding
const queryEmbedding = await this.vectorOps.generateQueryEmbedding(query);
if (queryEmbedding.length !== this.core.indexConfig!.dims) {
throw new Error(
`Query embedding dimension mismatch: expected ${
this.core.indexConfig!.dims
}, got ${queryEmbedding.length}`
);
}
let sqlQuery = `
SELECT DISTINCT
s.namespace_path,
s.key,
s.value,
s.created_at,
s.updated_at,
(
$3 * (1 - MIN(v.embedding <=> $2)) +
(1 - $3) * ts_rank(to_tsvector($6::regconfig, s.value::text), plainto_tsquery($6::regconfig, $4))
) as hybrid_score
FROM ${this.core.schema}.store s
JOIN ${this.core.schema}.store_vectors v ON s.namespace_path = v.namespace_path AND s.key = v.key
WHERE s.namespace_path LIKE $1
AND (s.expires_at IS NULL OR s.expires_at > CURRENT_TIMESTAMP)
AND (
to_tsvector($6::regconfig, s.value::text) @@ plainto_tsquery($6::regconfig, $4)
OR v.embedding <=> $2 <= $5
)
`;
const params: unknown[] = [
`${namespacePath}%`,
`[${queryEmbedding.join(",")}]`,
vectorWeight,
query,
1 - similarityThreshold,
this.core.textSearchLanguage,
];
let paramIndex = 7;
// Add filter conditions
if (filter && Object.keys(filter).length > 0) {
const { conditions, newParamIndex } =
QueryBuilder.buildFilterConditions(filter, params, paramIndex);
if (conditions.length > 0) {
const adjustedConditions = conditions.map((condition: string) =>
condition.replace(/value ->/g, "s.value ->")
);
sqlQuery += ` AND (${adjustedConditions.join(" AND ")})`;
paramIndex = newParamIndex;
}
}
sqlQuery += `
GROUP BY s.namespace_path, s.key, s.value, s.created_at, s.updated_at
ORDER BY hybrid_score DESC
LIMIT $${paramIndex} OFFSET $${paramIndex + 1}
`;
params.push(limit, offset);
const result = await client.query(sqlQuery, params);
return result.rows.map((row) => ({
namespace: row.namespace_path.split(":"),
key: row.key,
value: row.value,
createdAt: row.created_at || new Date(),
updatedAt: row.updated_at || new Date(),
score: parseFloat(row.hybrid_score),
}));
});
}
}
libs/checkpoint-postgres/src/store/modules/ttl-manager.ts
+39
View File
@@ -0,0 +1,39 @@
import { DatabaseCore } from "./database-core.js";
/**
* Handles TTL (Time-To-Live) operations: expiration, cleanup, sweeping.
*/
export class TTLManager {
private sweepInterval?: NodeJS.Timeout;
constructor(private core: DatabaseCore) {}
start(): void {
if (this.sweepInterval) return; // Already running
const intervalMs =
(this.core.ttlConfig?.sweepIntervalMinutes || 60) * 60 * 1000;
this.sweepInterval = setInterval(() => {
this.sweepExpiredItems().catch((error) => {
console.error("Error during TTL sweep:", error);
});
}, intervalMs);
}
stop(): void {
if (this.sweepInterval) {
clearInterval(this.sweepInterval);
this.sweepInterval = undefined;
}
}
async sweepExpiredItems(): Promise<number> {
return this.core.withClient(async (client) => {
const result = await client.query(`
DELETE FROM ${this.core.schema}.store
WHERE expires_at IS NOT NULL AND expires_at <= CURRENT_TIMESTAMP
`);
return result.rowCount || 0;
});
}
}
libs/checkpoint-postgres/src/store/modules/types.ts
+231
View File
@@ -0,0 +1,231 @@
import pg from "pg";
import type { Item } from "@langchain/langgraph-checkpoint";
/**
* Interface for embedding models.
* Compatible with LangChain's Embeddings interface.
*/
export interface Embeddings {
embedDocuments(texts: string[]): Promise<number[][]>;
embedQuery(text: string): Promise<number[]>;
}
export type EmbeddingsFunc = (texts: string[]) => Promise<number[][]>;
export type VectorIndexType = "ivfflat" | "hnsw";
export type DistanceMetric = "cosine" | "l2" | "inner_product";
export interface HNSWConfig {
/**
* Maximum number of connections for each node in the graph.
* Higher values improve recall but increase memory usage and build time.
* @default 16
*/
m?: number;
/**
* Size of the dynamic candidate list during construction.
* Higher values improve index quality but increase build time.
* @default 200
*/
efConstruction?: number;
/**
* Size of the dynamic candidate list during search.
* Higher values improve recall but increase search time.
* Can be adjusted per query for performance tuning.
* @default 40
*/
ef?: number;
}
export interface IVFFlatConfig {
/**
* Number of inverted lists (clusters).
* Rule of thumb: sqrt(number_of_rows) for good performance.
* @default 100
*/
lists?: number;
/**
* Number of probes to use during search.
* Higher values improve recall but increase search time.
* @default 1
*/
probes?: number;
}
export interface IndexConfig {
/**
* Number of dimensions in the embedding vectors.
*/
dims: number;
/**
* Embedding function to generate embeddings from text.
* Can be a LangChain Embeddings instance or a function.
*/
embed: Embeddings | EmbeddingsFunc;
/**
* Fields to extract text from for embedding generation.
* Uses JSON path syntax. Defaults to ["$"] (entire document).
*/
fields?: string[];
/**
* Vector index type to use.
* - 'hnsw': Hierarchical Navigable Small World (best for most use cases)
* - 'ivfflat': Inverted File with Flat compression (good for large datasets)
* @default 'hnsw'
*/
indexType?: VectorIndexType;
/**
* Distance metric for vector similarity.
* @default 'cosine'
*/
distanceMetric?: DistanceMetric;
/**
* HNSW-specific configuration parameters.
* Only used when indexType is 'hnsw'.
*/
hnsw?: HNSWConfig;
/**
* IVFFlat-specific configuration parameters.
* Only used when indexType is 'ivfflat'.
*/
ivfflat?: IVFFlatConfig;
/**
* Whether to create indexes for all distance metrics.
* If false, only creates index for the specified distanceMetric.
* @default false
*/
createAllMetricIndexes?: boolean;
}
export interface TTLConfig {
/**
* Default TTL in minutes for new items.
*/
defaultTtl?: number;
/**
* Whether to refresh TTL on read operations by default.
* @default true
*/
refreshOnRead?: boolean;
/**
* Interval in minutes between TTL sweep operations.
* @default 60
*/
sweepIntervalMinutes?: number;
}
export interface FilterOperators {
$eq?: string | number | boolean | null;
$ne?: string | number | boolean | null;
$gt?: string | number;
$gte?: string | number;
$lt?: string | number;
$lte?: string | number;
$in?: (string | number | boolean | null)[];
$nin?: (string | number | boolean | null)[];
$exists?: boolean;
}
export interface SearchItem extends Item {
/**
* Relevance/similarity score for the search result.
*/
score?: number;
}
export interface SearchOptions {
/**
* Filter conditions with support for advanced operators.
*/
filter?: Record<string, string | number | boolean | null | FilterOperators>;
/**
* Natural language search query.
*/
query?: string;
/**
* Maximum number of results to return.
* @default 10
*/
limit?: number;
/**
* Number of results to skip for pagination.
* @default 0
*/
offset?: number;
/**
* Whether to refresh TTL for returned items.
*/
refreshTtl?: boolean;
}
export interface PutOptions {
/**
* TTL in minutes for this item.
*/
ttl?: number;
}
export interface PostgresStoreConfig {
/**
* PostgreSQL connection string or connection configuration object.
*
* @example
* // Connection string
* "postgresql://user:password@localhost:5432/database"
*
* // Configuration object
* {
* host: "localhost",
* port: 5432,
* database: "mydb",
* user: "postgres",
* password: "password"
* }
*/
connectionOptions: string | pg.PoolConfig;
/**
* Database schema name to use for store tables.
* @default "public"
*/
schema?: string;
/**
* Whether to automatically create tables if they don't exist.
* @default true
*/
ensureTables?: boolean;
/**
* TTL configuration for automatic expiration of items.
*/
ttl?: TTLConfig;
/**
* Vector search configuration for semantic search capabilities.
* If provided, enables vector similarity search using pgvector extension.
*/
index?: IndexConfig;
/**
* Language for PostgreSQL full-text search operations.
* Supports any language configuration available in PostgreSQL.
* @default "english"
*/
textSearchLanguage?: string;
}
libs/checkpoint-postgres/src/store/modules/utils.ts
+33
View File
@@ -0,0 +1,33 @@
/**
* Validates the provided namespace.
* @param namespace The namespace to validate.
* @throws {Error} If the namespace is invalid.
*/
export function validateNamespace(namespace: string[]): void {
if (namespace.length === 0) {
throw new Error("Namespace cannot be empty.");
}
for (const label of namespace) {
if (typeof label !== "string") {
throw new Error(
`Invalid namespace label '${label}' found in ${namespace}. Namespace labels ` +
`must be strings, but got ${typeof label}.`
);
}
if (label.includes(".")) {
throw new Error(
`Invalid namespace label '${label}' found in ${namespace}. Namespace labels cannot contain periods ('.').`
);
}
if (label === "") {
throw new Error(
`Namespace labels cannot be empty strings. Got ${label} in ${namespace}`
);
}
}
if (namespace[0] === "langgraph") {
throw new Error(
`Root label for namespace cannot be "langgraph". Got: ${namespace}`
);
}
}
libs/checkpoint-postgres/src/store/modules/vector-operations.ts
+195
View File
@@ -0,0 +1,195 @@
import pg from "pg";
import { DatabaseCore } from "./database-core.js";
import { Embeddings } from "./types.js";
/**
* Handles vector processing, embeddings, and text extraction.
*/
export class VectorOperations {
constructor(private core: DatabaseCore) {}
async indexItemVectors(
client: pg.PoolClient,
namespacePath: string,
key: string,
value: Record<string, unknown>,
index?: string[] | false
): Promise<void> {
// Early exit if vector indexing is not configured or explicitly disabled
if (!this.core.indexConfig || index === false) return;
// Delete existing vectors for this item
await client.query(
`
DELETE FROM ${this.core.schema}.store_vectors
WHERE namespace_path = $1 AND key = $2
`,
[namespacePath, key]
);
// Use provided index fields if specified, otherwise use the configured default
const fields = index || this.core.indexConfig.fields || ["$"];
const textsToEmbed: { fieldPath: string; text: string }[] = [];
// Extract text from configured fields
for (const fieldPath of fields) {
const extractedTexts = this.extractTextAtPath(value, fieldPath);
extractedTexts.forEach((text, i) => {
const trimmedText = text?.trim();
if (trimmedText) {
const actualFieldPath =
extractedTexts.length > 1 ? `${fieldPath}[${i}]` : fieldPath;
textsToEmbed.push({ fieldPath: actualFieldPath, text: trimmedText });
}
});
}
if (textsToEmbed.length === 0) return;
// Generate embeddings
const texts = textsToEmbed.map((item) => item.text);
const embeddings = await this.generateEmbeddings(texts);
// Insert vectors
for (let i = 0; i < textsToEmbed.length; i += 1) {
const { fieldPath, text } = textsToEmbed[i];
const embedding = embeddings[i];
if (embedding?.length === this.core.indexConfig.dims) {
await client.query(
`
INSERT INTO ${this.core.schema}.store_vectors
(namespace_path, key, field_path, text_content, embedding)
VALUES ($1, $2, $3, $4, $5)
`,
[namespacePath, key, fieldPath, text, `[${embedding.join(",")}]`]
);
}
}
}
async generateEmbeddings(texts: string[]): Promise<number[][]> {
if (!this.core.indexConfig) {
throw new Error("Vector search not configured");
}
const { embed } = this.core.indexConfig;
if (typeof embed === "function") {
return await embed(texts);
}
// LangChain Embeddings interface
if (embed && typeof embed === "object" && "embedDocuments" in embed) {
return await (embed as Embeddings).embedDocuments(texts);
}
throw new Error("Invalid embedding configuration");
}
async generateQueryEmbedding(text: string): Promise<number[]> {
if (!this.core.indexConfig) {
throw new Error("Vector search not configured");
}
const { embed } = this.core.indexConfig;
if (typeof embed === "function") {
const embeddings = await embed([text]);
return embeddings[0] || [];
}
// LangChain Embeddings interface
if (embed && typeof embed === "object" && "embedQuery" in embed) {
return await (embed as Embeddings).embedQuery(text);
}
if (embed && typeof embed === "object" && "embedDocuments" in embed) {
const embeddings = await (embed as Embeddings).embedDocuments([text]);
return embeddings[0] || [];
}
throw new Error("Invalid embedding configuration");
}
private extractTextAtPath(obj: unknown, path: string): string[] {
if (path === "$") {
return [JSON.stringify(obj)];
}
const parts = path.split(".");
let current = obj;
const results: string[] = [];
try {
for (let i = 0; i < parts.length; i += 1) {
const part = parts[i];
if (part.includes("[")) {
const [field, arrayPart] = part.split("[");
const arrayIndex = arrayPart.replace("]", "");
if (field && typeof current === "object" && current !== null) {
current = (current as Record<string, unknown>)[field];
}
if (arrayIndex === "*") {
if (Array.isArray(current)) {
const remainingPath = parts.slice(i + 1).join(".");
if (remainingPath) {
for (const item of current) {
if (item != null) {
results.push(
...this.extractTextAtPath(item, remainingPath)
);
}
}
} else {
for (const item of current) {
if (typeof item === "string") {
results.push(item);
} else if (typeof item === "object" && item !== null) {
results.push(JSON.stringify(item));
} else if (item != null) {
results.push(String(item));
}
}
}
}
return results;
} else if (arrayIndex === "-1") {
if (Array.isArray(current) && current.length > 0) {
current = current[current.length - 1];
}
} else {
const index = parseInt(arrayIndex, 10);
if (
Array.isArray(current) &&
index >= 0 &&
index < current.length
) {
current = current[index];
}
}
} else if (typeof current === "object" && current !== null) {
current = (current as Record<string, unknown>)[part];
}
if (current == null) return [];
}
if (typeof current === "string") {
results.push(current);
} else if (typeof current === "object" && current !== null) {
results.push(JSON.stringify(current));
} else {
results.push(String(current));
}
} catch (error) {
return [];
}
return results;
}
}
libs/checkpoint-postgres/src/store/sql.ts
+13
View File
@@ -0,0 +1,13 @@
interface STORE_TABLES {
store: string;
store_vectors: string;
store_migrations: string;
}
export const getStoreTablesWithSchema = (schema: string): STORE_TABLES => {
const tables = ["store", "store_vectors", "store_migrations"];
return tables.reduce((acc, table) => {
acc[table as keyof STORE_TABLES] = `${schema}.${table}`;
return acc;
}, {} as STORE_TABLES);
};
libs/checkpoint-postgres/src/store/store-migrations.ts
+141
View File
@@ -0,0 +1,141 @@
import { getStoreTablesWithSchema } from "./sql.js";
import { VectorIndexType, DistanceMetric } from "./modules/types.js";
/**
* Store migration configuration
*/
export interface StoreMigrationConfig {
schema: string;
indexConfig?: {
dims: number;
indexType?: VectorIndexType;
distanceMetric?: DistanceMetric;
createAllMetricIndexes?: boolean;
hnsw?: { m?: number; efConstruction?: number };
ivfflat?: { lists?: number };
};
}
/**
* To add a new store migration, add a new string to the list returned by the getStoreMigrations function.
* The position of the migration in the list is the version number.
*/
export const getStoreMigrations = (config: StoreMigrationConfig): string[] => {
const { schema, indexConfig } = config;
const STORE_TABLES = getStoreTablesWithSchema(schema);
const migrations = [];
// Migration 1: Create store migrations table
migrations.push(`CREATE TABLE IF NOT EXISTS ${STORE_TABLES.store_migrations} (
v INTEGER PRIMARY KEY
);`);
// Migration 2: Create main store table
migrations.push(`CREATE TABLE IF NOT EXISTS ${STORE_TABLES.store} (
namespace_path TEXT NOT NULL,
key TEXT NOT NULL,
value JSONB NOT NULL,
created_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP,
expires_at TIMESTAMPTZ,
PRIMARY KEY (namespace_path, key)
);`);
// Migration 3: Create basic indexes
migrations.push(`
CREATE INDEX IF NOT EXISTS idx_store_namespace_path
ON ${STORE_TABLES.store} USING btree (namespace_path);
CREATE INDEX IF NOT EXISTS idx_store_value_gin
ON ${STORE_TABLES.store} USING gin (value);
CREATE INDEX IF NOT EXISTS idx_store_expires_at
ON ${STORE_TABLES.store} USING btree (expires_at)
WHERE expires_at IS NOT NULL;
`);
// Migration 4: Create update trigger
migrations.push(`
CREATE OR REPLACE FUNCTION ${schema}.update_updated_at_column()
RETURNS TRIGGER AS $$
BEGIN
NEW.updated_at = CURRENT_TIMESTAMP;
RETURN NEW;
END;
$$ language 'plpgsql';
DROP TRIGGER IF EXISTS update_store_updated_at ON ${STORE_TABLES.store};
CREATE TRIGGER update_store_updated_at
BEFORE UPDATE ON ${STORE_TABLES.store}
FOR EACH ROW EXECUTE FUNCTION ${schema}.update_updated_at_column();
`);
// Vector-related migrations (only if indexConfig is provided)
if (indexConfig) {
// Migration 5: Enable vector extension
migrations.push(`CREATE EXTENSION IF NOT EXISTS vector;`);
// Migration 6: Create vector table
migrations.push(`CREATE TABLE IF NOT EXISTS ${STORE_TABLES.store_vectors} (
namespace_path TEXT NOT NULL,
key TEXT NOT NULL,
field_path TEXT NOT NULL,
text_content TEXT NOT NULL,
embedding vector(${indexConfig.dims}) NOT NULL,
created_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP,
PRIMARY KEY (namespace_path, key, field_path),
FOREIGN KEY (namespace_path, key) REFERENCES ${STORE_TABLES.store}(namespace_path, key) ON DELETE CASCADE
);`);
// Migration 7: Create vector indexes
const {
indexType = "hnsw",
distanceMetric = "cosine",
createAllMetricIndexes = false,
} = indexConfig;
const metricsToIndex: DistanceMetric[] = createAllMetricIndexes
? ["cosine", "l2", "inner_product"]
: [distanceMetric];
for (const metric of metricsToIndex) {
let metricSuffix: string;
if (metric === "cosine") {
metricSuffix = "cosine";
} else if (metric === "l2") {
metricSuffix = "l2";
} else {
metricSuffix = "ip";
}
const indexName = `idx_store_vectors_embedding_${metricSuffix}_${indexType}`;
const operatorClassMap = {
cosine: "vector_cosine_ops",
l2: "vector_l2_ops",
inner_product: "vector_ip_ops",
} as const;
const operatorClass = operatorClassMap[metric];
let vectorIndexSql: string;
if (indexType === "hnsw") {
const m = indexConfig.hnsw?.m || 16;
const efConstruction = indexConfig.hnsw?.efConstruction || 200;
vectorIndexSql = `CREATE INDEX IF NOT EXISTS ${indexName}
ON ${STORE_TABLES.store_vectors} USING hnsw (embedding ${operatorClass})
WITH (m = ${m}, ef_construction = ${efConstruction});`;
} else if (indexType === "ivfflat") {
const lists = indexConfig.ivfflat?.lists || 100;
vectorIndexSql = `CREATE INDEX IF NOT EXISTS ${indexName}
ON ${STORE_TABLES.store_vectors} USING ivfflat (embedding ${operatorClass})
WITH (lists = ${lists});`;
} else {
throw new Error(`Unsupported vector index type: ${indexType}`);
}
migrations.push(vectorIndexSql);
}
}
return migrations;
};
libs/checkpoint-postgres/src/tests/checkpoints.int.test.ts
+25 -28
View File
@@ -60,12 +60,35 @@ if (!TEST_POSTGRES_URL) {
let postgresSavers: PostgresSaver[] = [];
afterAll(async () => {
await Promise.all(postgresSavers.map((saver) => saver.end()));
// clear the ended savers to clean up for the next test
postgresSavers = [];
// Drop all test databases
const pool = new Pool({ connectionString: TEST_POSTGRES_URL });
try {
const result = await pool.query(`
SELECT datname FROM pg_database
WHERE datname LIKE 'lg_test_db_%'
`);
for (const row of result.rows) {
const dbName = row.datname;
await pool.query(`DROP DATABASE ${dbName} WITH (FORCE)`);
console.log(`🗑️ Dropped database: ${dbName}`);
}
} finally {
await pool.end();
}
}, 30_000);
describe.each([
{ schema: undefined, description: "the default schema" },
// { schema: "custom_schema", description: "a custom schema" },
{ schema: "custom_schema", description: "a custom schema" },
])("PostgresSaver with $description", ({ schema }) => {
let postgresSaver: PostgresSaver;
let currentDbConnectionString: string;
beforeEach(async () => {
@@ -97,31 +120,6 @@ describe.each([
}
});
afterAll(async () => {
await Promise.all(postgresSavers.map((saver) => saver.end()));
// clear the ended savers to clean up for the next test
postgresSavers = [];
// Drop all test databases
const pool = new Pool({
connectionString: TEST_POSTGRES_URL,
});
try {
const result = await pool.query(`
SELECT datname FROM pg_database
WHERE datname LIKE 'lg_test_db_%'
`);
for (const row of result.rows) {
const dbName = row.datname;
await pool.query(`DROP DATABASE ${dbName} WITH (FORCE)`);
console.log(`🗑️ Dropped database: ${dbName}`);
}
} finally {
await pool.end();
}
});
it("should properly initialize and setup the database", async () => {
// Verify that the database is properly initialized
const pool = new Pool({
@@ -382,7 +380,6 @@ describe.each([
expect(firstCheckpointTuple?.metadata).toEqual({
source: "update",
step: -1,
writes: null,
parents: {},
});
expect(firstCheckpointTuple?.parentConfig).toBeUndefined();
libs/checkpoint-postgres/src/tests/store.int.test.ts
+1529
View File
File diff suppressed because it is too large Load Diff