From 7c7cd34908aa89b486be312db7e7e8142e78f5e6 Mon Sep 17 00:00:00 2001
From: Alex Yang <himself65@outlook.com>
Date: Sun, 20 Oct 2024 14:01:24 -0700
Subject: [PATCH] fix(pg): allow passing perform setup (#1350)
---
.../node/vector-store/pg-vector-store.e2e.ts | 19 +++++++++++++++++++
.../src/vector-store/PGVectorStore.ts | 9 +++++++--
2 files changed, 26 insertions(+), 2 deletions(-)
diff --git a/packages/llamaindex/e2e/node/vector-store/pg-vector-store.e2e.ts b/packages/llamaindex/e2e/node/vector-store/pg-vector-store.e2e.ts
index 44e5a10fb..4baecefe6 100644
--- a/packages/llamaindex/e2e/node/vector-store/pg-vector-store.e2e.ts
+++ b/packages/llamaindex/e2e/node/vector-store/pg-vector-store.e2e.ts
@@ -105,3 +105,22 @@ await test("simple node", async (t) => {
assert.deepStrictEqual(result.nodes, []);
}
});
+
+await test("no setup", async (t) => {
+ // @ts-expect-error private method
+ assert.ok(PGVectorStore.prototype.checkSchema);
+ // @ts-expect-error private method
+ const Mock = class extends PGVectorStore {
+ private override async checkSchema(): Promise<any> {
+ throw new Error("should not be called");
+ }
+ };
+ const vectorStore = new Mock({
+ clientConfig: pgConfig,
+ performSetup: false,
+ });
+ const db = await vectorStore.client();
+ t.after(async () => {
+ await db.close();
+ });
+});
diff --git a/packages/llamaindex/src/vector-store/PGVectorStore.ts b/packages/llamaindex/src/vector-store/PGVectorStore.ts
index 11cee6b5d..6970bf1c2 100644
--- a/packages/llamaindex/src/vector-store/PGVectorStore.ts
+++ b/packages/llamaindex/src/vector-store/PGVectorStore.ts
@@ -114,6 +114,7 @@ type PGVectorStoreBaseConfig = {
tableName?: string | undefined;
dimensions?: number | undefined;
embedModel?: BaseEmbedding | undefined;
+ performSetup?: boolean | undefined;
};
export type PGVectorStoreConfig = VectorStoreBaseParams &
@@ -159,12 +160,14 @@ export class PGVectorStore extends BaseVectorStore {
private isDBConnected: boolean = false;
private db: IsomorphicDB | null = null;
private readonly clientConfig: pg.ClientConfig | null = null;
+ private readonly performSetup: boolean = true;
constructor(config: PGVectorStoreConfig) {
super(config);
this.schemaName = config?.schemaName ?? PGVECTOR_SCHEMA;
this.tableName = config?.tableName ?? PGVECTOR_TABLE;
this.dimensions = config?.dimensions ?? DEFAULT_DIMENSIONS;
+ this.performSetup = config?.performSetup ?? true;
if ("clientConfig" in config) {
this.clientConfig = config.clientConfig;
} else {
@@ -236,8 +239,10 @@ export class PGVectorStore extends BaseVectorStore {
this.isDBConnected = false;
});
- // Check schema, table(s), index(es)
- await this.checkSchema(this.db);
+ if (this.performSetup) {
+ // Check schema, table(s), index(es)
+ await this.checkSchema(this.db);
+ }
return this.db;
}
--
GitLab