Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions libs/langchain-community/src/vectorstores/pgvector.ts
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ export interface PGVectorStoreArgs {
collectionMetadata?: Metadata | null;
schemaName?: string | null;
extensionSchemaName?: string | null;
skipInitializationCheck?: boolean;
columns?: {
idColumnName?: string;
vectorColumnName?: string;
Expand Down Expand Up @@ -213,6 +214,8 @@ export class PGVectorStore extends VectorStore {

extensionSchemaName: string | null;

skipInitializationCheck: boolean;

metadataColumnName: string;

filter?: Metadata;
Expand Down Expand Up @@ -247,6 +250,7 @@ export class PGVectorStore extends VectorStore {
this.collectionMetadata = config.collectionMetadata ?? null;
this.schemaName = config.schemaName ?? null;
this.extensionSchemaName = config.extensionSchemaName ?? null;
this.skipInitializationCheck = config.skipInitializationCheck ?? false;

this.filter = config.filter;

Expand Down Expand Up @@ -741,6 +745,9 @@ export class PGVectorStore extends VectorStore {
* @returns Promise that resolves when the table has been ensured.
*/
async ensureTableInDatabase(dimensions?: number): Promise<void> {
if (this.skipInitializationCheck) {
return;
}
const vectorQuery =
this.extensionSchemaName == null
? "CREATE EXTENSION IF NOT EXISTS vector;"
Expand Down Expand Up @@ -772,6 +779,9 @@ export class PGVectorStore extends VectorStore {
*/
async ensureCollectionTableInDatabase(): Promise<void> {
try {
if (this.skipInitializationCheck) {
return;
}
const queryString = `
CREATE TABLE IF NOT EXISTS ${this.computedCollectionTableName} (
uuid uuid NOT NULL DEFAULT gen_random_uuid() PRIMARY KEY,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -821,6 +821,131 @@ describe("PGVectorStore with schema", () => {
});
});

describe("PGVectorStore with skipInitializationCheck", () => {
let pgvectorVectorStore: PGVectorStore;
const tableName = "testlangchain_skip_init";

afterEach(async () => {
const pool = new pg.Pool(postgresConnectionOptions);
await pool.query(`DROP TABLE IF EXISTS "${tableName}"`);
await pool.end();
});

test("skipInitializationCheck=false (default) should initialize tables", async () => {
const config: PGVectorStoreArgs = {
postgresConnectionOptions,
tableName,
columns: {
idColumnName: "id",
vectorColumnName: "vector",
contentColumnName: "content",
metadataColumnName: "metadata",
},
};

pgvectorVectorStore = await PGVectorStore.initialize(
embeddingsEngine,
config
);

const result = await pgvectorVectorStore.pool.query(
`SELECT EXISTS (
SELECT FROM information_schema.tables
WHERE table_schema = 'public'
AND table_name = '${tableName}'
)`
);

expect(result.rows[0].exists).toBe(true);
await pgvectorVectorStore.end();
});

test("skipInitializationCheck=true should skip table initialization", async () => {
const config: PGVectorStoreArgs = {
postgresConnectionOptions,
tableName,
skipInitializationCheck: true,
columns: {
idColumnName: "id",
vectorColumnName: "vector",
contentColumnName: "content",
metadataColumnName: "metadata",
},
};

pgvectorVectorStore = await PGVectorStore.initialize(
embeddingsEngine,
config
);

const tableExists = await pgvectorVectorStore.pool.query(
`SELECT EXISTS (
SELECT FROM information_schema.tables
WHERE table_schema = 'public'
AND table_name = '${tableName}'
)`
);

expect(tableExists.rows[0].exists).toBe(false);

await pgvectorVectorStore.end();
});

test("skipInitializationCheck=true should work with addDocuments", async () => {
const setupPool = new pg.Pool(postgresConnectionOptions);

await setupPool.query(`CREATE EXTENSION IF NOT EXISTS vector`);
await setupPool.query(`
CREATE TABLE "${tableName}" (
"id" uuid NOT NULL DEFAULT gen_random_uuid() PRIMARY KEY,
"content" text,
"metadata" jsonb,
"vector" vector(1536)
);
`);

await setupPool.end();

const config: PGVectorStoreArgs = {
postgresConnectionOptions,
tableName,
skipInitializationCheck: true,
columns: {
idColumnName: "id",
vectorColumnName: "vector",
contentColumnName: "content",
metadataColumnName: "metadata",
},
};

pgvectorVectorStore = await PGVectorStore.initialize(
embeddingsEngine,
config
);

const documents = [
{ pageContent: "Hello world", metadata: { source: "test" } },
{
pageContent: "Testing skipInitializationCheck",
metadata: { source: "test" },
},
];

await pgvectorVectorStore.addDocuments(documents);

const query = await embeddingsEngine.embedQuery("Hello");
const results = await pgvectorVectorStore.similaritySearchVectorWithScore(
query,
1
);

expect(results).toHaveLength(1);
expect(results[0][0].pageContent).toBe("Hello world");

await pgvectorVectorStore.end();
});
});

describe("PGVectorStore with HNSW index", () => {
let pgvectorVectorStore: PGVectorStore;
const tableName = "testlangchain";
Expand Down
Loading