add docs and unit test

.changeset/cold-gifts-help.md

@@ -0,0 +1,6 @@
+---
+"llamaindex": patch
+"@llamaindex/core": patch
+---
+
+Fix context not being sent using ContextChatEngine

apps/docs/CHANGELOG.md

@@ -1,12 +1,5 @@
 # docs
 
-## 0.0.74
-
-### Patch Changes
-
-- Updated dependencies [e9714db]
-  - llamaindex@0.6.5
-
 ## 0.0.73
 
 ### Patch Changes

apps/docs/package.json

@@ -1,6 +1,6 @@
 {
   "name": "docs",
-  "version": "0.0.74",
+  "version": "0.0.73",
   "private": true,
   "scripts": {
     "docusaurus": "docusaurus",

examples/pg-vector-store/load-docs.ts

@@ -40,11 +40,7 @@ async function main(args: any) {
     const rdr = new SimpleDirectoryReader(callback);
     const docs = await rdr.loadData({ directoryPath: sourceDir });
 
-    const pgvs = new PGVectorStore({
-      clientConfig: {
-        connectionString: process.env.PG_CONNECTION_STRING,
-      },
-    });
+    const pgvs = new PGVectorStore();
     pgvs.setCollection(sourceDir);
     await pgvs.clearCollection();
 

examples/pg-vector-store/query.ts

@@ -7,11 +7,7 @@ async function main() {
   });
 
   try {
-    const pgvs = new PGVectorStore({
-      clientConfig: {
-        connectionString: process.env.PG_CONNECTION_STRING,
-      },
-    });
+    const pgvs = new PGVectorStore();
     // Optional - set your collection name, default is no filter on this field.
     // pgvs.setCollection();
 

packages/autotool/CHANGELOG.md

@@ -1,12 +1,5 @@
 # @llamaindex/autotool
 
-## 3.0.5
-
-### Patch Changes
-
-- Updated dependencies [e9714db]
-  - llamaindex@0.6.5
-
 ## 3.0.4
 
 ### Patch Changes

packages/autotool/examples/01_node/CHANGELOG.md

@@ -1,13 +1,5 @@
 # @llamaindex/autotool-01-node-example
 
-## 0.0.14
-
-### Patch Changes
-
-- Updated dependencies [e9714db]
-  - llamaindex@0.6.5
-  - @llamaindex/autotool@3.0.5
-
 ## 0.0.13
 
 ### Patch Changes

packages/autotool/examples/01_node/package.json

@@ -13,5 +13,5 @@
   "scripts": {
     "start": "node --import tsx --import @llamaindex/autotool/node ./src/index.ts"
   },
-  "version": "0.0.14"
+  "version": "0.0.13"
 }

packages/autotool/examples/02_nextjs/CHANGELOG.md

@@ -1,13 +1,5 @@
 # @llamaindex/autotool-02-next-example
 
-## 0.1.58
-
-### Patch Changes
-
-- Updated dependencies [e9714db]
-  - llamaindex@0.6.5
-  - @llamaindex/autotool@3.0.5
-
 ## 0.1.57
 
 ### Patch Changes

packages/autotool/examples/02_nextjs/package.json

@@ -1,7 +1,7 @@
 {
   "name": "@llamaindex/autotool-02-next-example",
   "private": true,
-  "version": "0.1.58",
+  "version": "0.1.57",
   "scripts": {
     "dev": "next dev",
     "build": "next build",

packages/autotool/package.json

@@ -1,7 +1,7 @@
 {
   "name": "@llamaindex/autotool",
   "type": "module",
-  "version": "3.0.5",
+  "version": "3.0.4",
   "description": "auto transpile your JS function to LLM Agent compatible",
   "files": [
     "dist",

packages/core/src/memory/base.ts

@@ -1,5 +1,5 @@
 import { Settings } from "../global";
-import type { ChatMessage, MessageContent } from "../llms";
+import type { ChatMessage } from "../llms";
 import { type BaseChatStore, SimpleChatStore } from "../storage/chat-store";
 import { extractText } from "../utils";
 
@@ -12,15 +12,36 @@ export const DEFAULT_CHAT_STORE_KEY = "chat_history";
 export abstract class BaseMemory<
   AdditionalMessageOptions extends object = object,
 > {
+  /**
+   * Retrieves messages from the memory, optionally including transient messages.
+   * Compared to getAllMessages, this method a) allows for transient messages to be included in the retrieval and b) may return a subset of the total messages by applying a token limit.
+   * @param transientMessages Optional array of temporary messages to be included in the retrieval.
+   * These messages are not stored in the memory but are considered for the current interaction.
+   * @returns An array of chat messages, either synchronously or as a Promise.
+   */
   abstract getMessages(
-    input?: MessageContent | undefined,
+    transientMessages?: ChatMessage<AdditionalMessageOptions>[] | undefined,
   ):
     | ChatMessage<AdditionalMessageOptions>[]
     | Promise<ChatMessage<AdditionalMessageOptions>[]>;
+
+  /**
+   * Retrieves all messages stored in the memory.
+   * @returns An array of all chat messages, either synchronously or as a Promise.
+   */
   abstract getAllMessages():
     | ChatMessage<AdditionalMessageOptions>[]
     | Promise<ChatMessage<AdditionalMessageOptions>[]>;
+
+  /**
+   * Adds a new message to the memory.
+   * @param messages The chat message to be added to the memory.
+   */
   abstract put(messages: ChatMessage<AdditionalMessageOptions>): void;
+
+  /**
+   * Clears all messages from the memory.
+   */
   abstract reset(): void;
 
   protected _tokenCountForMessages(messages: ChatMessage[]): number {

packages/core/src/memory/chat-memory-buffer.ts

@@ -1,5 +1,5 @@
 import { Settings } from "../global";
-import type { ChatMessage, LLM, MessageContent } from "../llms";
+import type { ChatMessage, LLM } from "../llms";
 import { type BaseChatStore } from "../storage/chat-store";
 import { BaseChatStoreMemory, DEFAULT_TOKEN_LIMIT_RATIO } from "./base";
 
@@ -34,7 +34,7 @@ export class ChatMemoryBuffer<
   }
 
   getMessages(
-    input?: MessageContent | undefined,
+    transientMessages?: ChatMessage<AdditionalMessageOptions>[] | undefined,
     initialTokenCount: number = 0,
   ) {
     const messages = this.getAllMessages();
@@ -43,16 +43,22 @@ export class ChatMemoryBuffer<
       throw new Error("Initial token count exceeds token limit");
     }
 
-    let messageCount = messages.length;
-    let currentMessages = messages.slice(-messageCount);
-    let tokenCount = this._tokenCountForMessages(messages) + initialTokenCount;
+    // Add input messages as transient messages
+    const messagesWithInput = transientMessages
+      ? [...transientMessages, ...messages]
+      : messages;
+
+    let messageCount = messagesWithInput.length;
+    let currentMessages = messagesWithInput.slice(-messageCount);
+    let tokenCount =
+      this._tokenCountForMessages(messagesWithInput) + initialTokenCount;
 
     while (tokenCount > this.tokenLimit && messageCount > 1) {
       messageCount -= 1;
-      if (messages.at(-messageCount)!.role === "assistant") {
+      if (messagesWithInput.at(-messageCount)!.role === "assistant") {
         messageCount -= 1;
       }
-      currentMessages = messages.slice(-messageCount);
+      currentMessages = messagesWithInput.slice(-messageCount);
       tokenCount =
         this._tokenCountForMessages(currentMessages) + initialTokenCount;
     }
@@ -60,6 +66,6 @@ export class ChatMemoryBuffer<
     if (tokenCount > this.tokenLimit && messageCount <= 0) {
       return [];
     }
-    return messages.slice(-messageCount);
+    return messagesWithInput.slice(-messageCount);
   }
 }

packages/core/src/memory/summary-memory.ts

@@ -114,18 +114,22 @@ export class ChatSummaryMemoryBuffer extends BaseMemory {
     }
   }
 
-  private calcCurrentRequestMessages() {
-    // TODO: check order: currently, we're sending:
+  private calcCurrentRequestMessages(transientMessages?: ChatMessage[]) {
+    // currently, we're sending:
     // system messages first, then transient messages and then the messages that describe the conversation so far
-    return [...this.systemMessages, ...this.calcConversationMessages(true)];
+    return [
+      ...this.systemMessages,
+      ...(transientMessages ? transientMessages : []),
+      ...this.calcConversationMessages(true),
+    ];
   }
 
   reset() {
     this.messages = [];
   }
 
-  async getMessages(): Promise<ChatMessage[]> {
-    const requestMessages = this.calcCurrentRequestMessages();
+  async getMessages(transientMessages?: ChatMessage[]): Promise<ChatMessage[]> {
+    const requestMessages = this.calcCurrentRequestMessages(transientMessages);
 
     // get tokens of current request messages and the transient messages
     const tokens = requestMessages.reduce(
@@ -149,7 +153,7 @@ export class ChatSummaryMemoryBuffer extends BaseMemory {
       // TODO: we still might have too many tokens
       // e.g. too large system messages or transient messages
       // how should we deal with that?
-      return this.calcCurrentRequestMessages();
+      return this.calcCurrentRequestMessages(transientMessages);
     }
     return requestMessages;
   }

packages/core/tests/memory/chat-memory-buffer.test.ts

@@ -0,0 +1,74 @@
+import { Settings } from "@llamaindex/core/global";
+import type { ChatMessage } from "@llamaindex/core/llms";
+import { ChatMemoryBuffer } from "@llamaindex/core/memory";
+import { beforeEach, describe, expect, test } from "vitest";
+
+describe("ChatMemoryBuffer", () => {
+  beforeEach(() => {
+    // Mock the Settings.llm
+    (Settings.llm as any) = {
+      metadata: {
+        contextWindow: 1000,
+      },
+    };
+  });
+
+  test("constructor initializes with custom token limit", () => {
+    const buffer = new ChatMemoryBuffer({ tokenLimit: 500 });
+    expect(buffer.tokenLimit).toBe(500);
+  });
+
+  test("getMessages returns all messages when under token limit", () => {
+    const messages: ChatMessage[] = [
+      { role: "user", content: "Hello" },
+      { role: "assistant", content: "Hi there!" },
+      { role: "user", content: "How are you?" },
+    ];
+    const buffer = new ChatMemoryBuffer({
+      tokenLimit: 1000,
+      chatHistory: messages,
+    });
+
+    const result = buffer.getMessages();
+    expect(result).toEqual(messages);
+  });
+
+  test("getMessages truncates messages when over token limit", () => {
+    const messages: ChatMessage[] = [
+      { role: "user", content: "This is a long message" },
+      { role: "assistant", content: "This is also a long reply" },
+      { role: "user", content: "Short" },
+    ];
+    const buffer = new ChatMemoryBuffer({
+      tokenLimit: 5, // limit to only allow the last message
+      chatHistory: messages,
+    });
+
+    const result = buffer.getMessages();
+    expect(result).toEqual([{ role: "user", content: "Short" }]);
+  });
+
+  test("getMessages handles input messages", () => {
+    const storedMessages: ChatMessage[] = [
+      { role: "user", content: "Hello" },
+      { role: "assistant", content: "Hi there!" },
+    ];
+    const buffer = new ChatMemoryBuffer({
+      tokenLimit: 50,
+      chatHistory: storedMessages,
+    });
+
+    const inputMessages: ChatMessage[] = [
+      { role: "user", content: "New message" },
+    ];
+    const result = buffer.getMessages(inputMessages);
+    expect(result).toEqual([...inputMessages, ...storedMessages]);
+  });
+
+  test("getMessages throws error when initial token count exceeds limit", () => {
+    const buffer = new ChatMemoryBuffer({ tokenLimit: 10 });
+    expect(() => buffer.getMessages(undefined, 20)).toThrow(
+      "Initial token count exceeds token limit",
+    );
+  });
+});

packages/experimental/CHANGELOG.md

@@ -1,12 +1,5 @@
 # @llamaindex/experimental
 
-## 0.0.83
-
-### Patch Changes
-
-- Updated dependencies [e9714db]
-  - llamaindex@0.6.5
-
 ## 0.0.82
 
 ### Patch Changes

packages/experimental/package.json

@@ -1,7 +1,7 @@
 {
   "name": "@llamaindex/experimental",
   "description": "Experimental package for LlamaIndexTS",
-  "version": "0.0.83",
+  "version": "0.0.82",
   "type": "module",
   "types": "dist/type/index.d.ts",
   "main": "dist/cjs/index.js",

packages/llamaindex/CHANGELOG.md

@@ -1,15 +1,5 @@
 # llamaindex
 
-## 0.6.5
-
-### Patch Changes
-
-- e9714db: feat: update `PGVectorStore`
-
-  - move constructor parameter `config.user` | `config.database` | `config.password` | `config.connectionString` into `config.clientConfig`
-  - if you pass `pg.Client` or `pg.Pool` instance to `PGVectorStore`, move it to `config.client`, setting `config.shouldConnect` to false if it's already connected
-  - default value of `PGVectorStore.collection` is now `"data"` instead of `""` (empty string)
-
 ## 0.6.4
 
 ### Patch Changes

packages/llamaindex/e2e/examples/cloudflare-worker-agent/CHANGELOG.md

@@ -1,12 +1,5 @@
 # @llamaindex/cloudflare-worker-agent-test
 
-## 0.0.67
-
-### Patch Changes
-
-- Updated dependencies [e9714db]
-  - llamaindex@0.6.5
-
 ## 0.0.66
 
 ### Patch Changes

packages/llamaindex/e2e/examples/cloudflare-worker-agent/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@llamaindex/cloudflare-worker-agent-test",
-  "version": "0.0.67",
+  "version": "0.0.66",
   "type": "module",
   "private": true,
   "scripts": {

packages/llamaindex/e2e/examples/nextjs-agent/CHANGELOG.md

@@ -1,12 +1,5 @@
 # @llamaindex/next-agent-test
 
-## 0.1.67
-
-### Patch Changes
-
-- Updated dependencies [e9714db]
-  - llamaindex@0.6.5
-
 ## 0.1.66
 
 ### Patch Changes

packages/llamaindex/e2e/examples/nextjs-agent/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@llamaindex/next-agent-test",
-  "version": "0.1.67",
+  "version": "0.1.66",
   "private": true,
   "scripts": {
     "dev": "next dev",

packages/llamaindex/e2e/examples/nextjs-edge-runtime/CHANGELOG.md

@@ -1,12 +1,5 @@
 # test-edge-runtime
 
-## 0.1.66
-
-### Patch Changes
-
-- Updated dependencies [e9714db]
-  - llamaindex@0.6.5
-
 ## 0.1.65
 
 ### Patch Changes

packages/llamaindex/e2e/examples/nextjs-edge-runtime/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@llamaindex/nextjs-edge-runtime-test",
-  "version": "0.1.66",
+  "version": "0.1.65",
   "private": true,
   "scripts": {
     "dev": "next dev",

packages/llamaindex/e2e/examples/nextjs-node-runtime/CHANGELOG.md

@@ -1,12 +1,5 @@
 # @llamaindex/next-node-runtime
 
-## 0.0.48
-
-### Patch Changes
-
-- Updated dependencies [e9714db]
-  - llamaindex@0.6.5
-
 ## 0.0.47
 
 ### Patch Changes

packages/llamaindex/e2e/examples/nextjs-node-runtime/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@llamaindex/next-node-runtime-test",
-  "version": "0.0.48",
+  "version": "0.0.47",
   "private": true,
   "scripts": {
     "dev": "next dev",

packages/llamaindex/e2e/examples/waku-query-engine/CHANGELOG.md

@@ -1,12 +1,5 @@
 # @llamaindex/waku-query-engine-test
 
-## 0.0.67
-
-### Patch Changes
-
-- Updated dependencies [e9714db]
-  - llamaindex@0.6.5
-
 ## 0.0.66
 
 ### Patch Changes

packages/llamaindex/e2e/examples/waku-query-engine/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@llamaindex/waku-query-engine-test",
-  "version": "0.0.67",
+  "version": "0.0.66",
   "type": "module",
   "private": true,
   "scripts": {

packages/llamaindex/e2e/node/vector-store/pg-vector-store.e2e.ts

@@ -9,54 +9,43 @@ import { registerTypes } from "pgvector/pg";
 
 config({ path: [".env.local", ".env", ".env.ci"] });
 
+let pgClient: pg.Client | pg.Pool;
+test.afterEach(async () => {
+  await pgClient.end();
+});
+
 const pgConfig = {
   user: process.env.POSTGRES_USER ?? "user",
   password: process.env.POSTGRES_PASSWORD ?? "password",
   database: "llamaindex_node_test",
 };
 
-await test("init with client", async (t) => {
-  const pgClient = new pg.Client(pgConfig);
+await test("init with client", async () => {
+  pgClient = new pg.Client(pgConfig);
   await pgClient.connect();
   await pgClient.query("CREATE EXTENSION IF NOT EXISTS vector");
   await registerTypes(pgClient);
-  t.after(async () => {
-    await pgClient.end();
-  });
-  const vectorStore = new PGVectorStore({
-    client: pgClient,
-    shouldConnect: false,
-  });
+  const vectorStore = new PGVectorStore(pgClient);
   assert.deepStrictEqual(await vectorStore.client(), pgClient);
 });
 
-await test("init with pool", async (t) => {
-  const pgClient = new pg.Pool(pgConfig);
+await test("init with pool", async () => {
+  pgClient = new pg.Pool(pgConfig);
   await pgClient.query("CREATE EXTENSION IF NOT EXISTS vector");
   const client = await pgClient.connect();
-  await client.query("CREATE EXTENSION IF NOT EXISTS vector");
   await registerTypes(client);
-  t.after(async () => {
-    client.release();
-    await pgClient.end();
-  });
-  const vectorStore = new PGVectorStore({
-    shouldConnect: false,
-    client,
-  });
+  const vectorStore = new PGVectorStore(client);
   assert.deepStrictEqual(await vectorStore.client(), client);
+  client.release();
 });
 
-await test("init without client", async (t) => {
-  const vectorStore = new PGVectorStore({ clientConfig: pgConfig });
-  const pgClient = (await vectorStore.client()) as pg.Client;
-  t.after(async () => {
-    await pgClient.end();
-  });
+await test("init without client", async () => {
+  const vectorStore = new PGVectorStore(pgConfig);
+  pgClient = (await vectorStore.client()) as pg.Client;
   assert.notDeepStrictEqual(pgClient, undefined);
 });
 
-await test("simple node", async (t) => {
+await test("simple node", async () => {
   const dimensions = 3;
   const schemaName =
     "llamaindex_vector_store_test_" + Math.random().toString(36).substring(7);
@@ -67,14 +56,10 @@ await test("simple node", async (t) => {
     embedding: [0.1, 0.2, 0.3],
   });
   const vectorStore = new PGVectorStore({
-    clientConfig: pgConfig,
+    ...pgConfig,
     dimensions,
     schemaName,
   });
-  const pgClient = (await vectorStore.client()) as pg.Client;
-  t.after(async () => {
-    await pgClient.end();
-  });
 
   await vectorStore.add([node]);
 
@@ -104,4 +89,6 @@ await test("simple node", async (t) => {
     });
     assert.deepStrictEqual(result.nodes, []);
   }
+
+  pgClient = (await vectorStore.client()) as pg.Client;
 });

packages/llamaindex/package.json

@@ -1,6 +1,6 @@
 {
   "name": "llamaindex",
-  "version": "0.6.5",
+  "version": "0.6.4",
   "license": "MIT",
   "type": "module",
   "keywords": [

packages/llamaindex/src/agent/base.ts

@@ -356,9 +356,8 @@ export abstract class AgentRunner<
     let chatHistory: ChatMessage<AdditionalMessageOptions>[] = [];
 
     if (params.chatHistory instanceof BaseMemory) {
-      chatHistory = (await params.chatHistory.getMessages(
-        params.message,
-      )) as ChatMessage<AdditionalMessageOptions>[];
+      chatHistory =
+        (await params.chatHistory.getMessages()) as ChatMessage<AdditionalMessageOptions>[];
     } else {
       chatHistory =
         params.chatHistory as ChatMessage<AdditionalMessageOptions>[];

packages/llamaindex/src/engines/chat/CondenseQuestionChatEngine.ts

@@ -78,9 +78,7 @@ export class CondenseQuestionChatEngine
   }
 
   private async condenseQuestion(chatHistory: BaseMemory, question: string) {
-    const chatHistoryStr = messagesToHistory(
-      await chatHistory.getMessages(question),
-    );
+    const chatHistoryStr = messagesToHistory(await chatHistory.getMessages());
 
     return this.llm.complete({
       prompt: this.condenseMessagePrompt.format({
@@ -103,7 +101,7 @@ export class CondenseQuestionChatEngine
       ? new ChatMemoryBuffer({
           chatHistory:
             params.chatHistory instanceof BaseMemory
-              ? await params.chatHistory.getMessages(message)
+              ? await params.chatHistory.getMessages()
               : params.chatHistory,
         })
       : this.chatHistory;

packages/llamaindex/src/engines/chat/ContextChatEngine.ts

@@ -92,7 +92,7 @@ export class ContextChatEngine extends PromptMixin implements ChatEngine {
       ? new ChatMemoryBuffer({
           chatHistory:
             params.chatHistory instanceof BaseMemory
-              ? await params.chatHistory.getMessages(message)
+              ? await params.chatHistory.getMessages()
               : params.chatHistory,
         })
       : this.chatHistory;
@@ -139,7 +139,7 @@ export class ContextChatEngine extends PromptMixin implements ChatEngine {
     const textOnly = extractText(message);
     const context = await this.contextGenerator.generate(textOnly);
     const systemMessage = this.prependSystemPrompt(context.message);
-    const messages = await chatHistory.getMessages(systemMessage.content);
+    const messages = await chatHistory.getMessages([systemMessage]);
     return { nodes: context.nodes, messages };
   }
 

packages/llamaindex/src/engines/chat/SimpleChatEngine.ts

@@ -40,7 +40,7 @@ export class SimpleChatEngine implements ChatEngine {
       ? new ChatMemoryBuffer({
           chatHistory:
             params.chatHistory instanceof BaseMemory
-              ? await params.chatHistory.getMessages(message)
+              ? await params.chatHistory.getMessages()
               : params.chatHistory,
         })
       : this.chatHistory;
@@ -48,7 +48,7 @@ export class SimpleChatEngine implements ChatEngine {
 
     if (stream) {
       const stream = await this.llm.chat({
-        messages: await chatHistory.getMessages(params.message),
+        messages: await chatHistory.getMessages(),
         stream: true,
       });
       return streamConverter(
@@ -66,7 +66,7 @@ export class SimpleChatEngine implements ChatEngine {
 
     const response = await this.llm.chat({
       stream: false,
-      messages: await chatHistory.getMessages(params.message),
+      messages: await chatHistory.getMessages(),
     });
     chatHistory.put(response.message);
     return EngineResponse.fromChatResponse(response);

packages/llamaindex/src/vector-store/PGVectorStore.ts

@@ -14,44 +14,25 @@ import {
 import { escapeLikeString } from "./utils.js";
 
 import type { BaseEmbedding } from "@llamaindex/core/embeddings";
-import { DEFAULT_COLLECTION } from "@llamaindex/core/global";
 import type { BaseNode, Metadata } from "@llamaindex/core/schema";
 import { Document, MetadataMode } from "@llamaindex/core/schema";
 
 export const PGVECTOR_SCHEMA = "public";
 export const PGVECTOR_TABLE = "llamaindex_embedding";
-export const DEFAULT_DIMENSIONS = 1536;
 
-type PGVectorStoreBaseConfig = {
+export type PGVectorStoreConfig = Pick<
+  pg.ClientConfig,
+  "user" | "database" | "password" | "connectionString"
+> & {
   schemaName?: string | undefined;
   tableName?: string | undefined;
   dimensions?: number | undefined;
   embedModel?: BaseEmbedding | undefined;
 };
 
-export type PGVectorStoreConfig = PGVectorStoreBaseConfig &
-  (
-    | {
-        /**
-         * Client configuration options for the pg client.
-         *
-         * {@link https://node-postgres.com/apis/client#new-client PostgresSQL Client API}
-         */
-        clientConfig: pg.ClientConfig;
-      }
-    | {
-        /**
-         * A pg client or pool client instance.
-         * If provided, make sure it is not connected to the database yet, or it will throw an error.
-         */
-        shouldConnect?: boolean | undefined;
-        client: pg.Client | pg.PoolClient;
-      }
-  );
-
 /**
  * Provides support for writing and querying vector data in Postgres.
- * Note: Can't be used with data created using the Python version of the vector store (https://docs.llamaindex.ai/en/stable/examples/vector_stores/postgres/)
+ * Note: Can't be used with data created using the Python version of the vector store (https://docs.llamaindex.ai/en/stable/examples/vector_stores/postgres.html)
  */
 export class PGVectorStore
   extends VectorStoreBase
@@ -59,26 +40,52 @@ export class PGVectorStore
 {
   storesText: boolean = true;
 
-  private collection: string = DEFAULT_COLLECTION;
-  private readonly schemaName: string = PGVECTOR_SCHEMA;
-  private readonly tableName: string = PGVECTOR_TABLE;
-  private readonly dimensions: number = DEFAULT_DIMENSIONS;
+  private collection: string = "";
+  private schemaName: string = PGVECTOR_SCHEMA;
+  private tableName: string = PGVECTOR_TABLE;
 
-  private isDBConnected: boolean = false;
-  private db: pg.ClientBase | null = null;
-  private readonly clientConfig: pg.ClientConfig | null = null;
+  private user: pg.ClientConfig["user"] | undefined = undefined;
+  private password: pg.ClientConfig["password"] | undefined = undefined;
+  private database: pg.ClientConfig["database"] | undefined = undefined;
+  private connectionString: pg.ClientConfig["connectionString"] | undefined =
+    undefined;
 
-  constructor(config: PGVectorStoreConfig) {
-    super(config?.embedModel);
-    this.schemaName = config?.schemaName ?? PGVECTOR_SCHEMA;
-    this.tableName = config?.tableName ?? PGVECTOR_TABLE;
-    this.dimensions = config?.dimensions ?? DEFAULT_DIMENSIONS;
-    if ("clientConfig" in config) {
-      this.clientConfig = config.clientConfig;
+  private dimensions: number = 1536;
+
+  private db?: pg.ClientBase;
+
+  /**
+   * Constructs a new instance of the PGVectorStore
+   *
+   * If the `connectionString` is not provided the following env variables are
+   * used to connect to the DB:
+   * PGHOST=your database host
+   * PGUSER=your database user
+   * PGPASSWORD=your database password
+   * PGDATABASE=your database name
+   * PGPORT=your database port
+   */
+  constructor(configOrClient?: PGVectorStoreConfig | pg.ClientBase) {
+    // We cannot import pg from top level, it might have side effects
+    //  so we only check if the config.connect function exists
+    if (
+      configOrClient &&
+      "connect" in configOrClient &&
+      typeof configOrClient.connect === "function"
+    ) {
+      const db = configOrClient as pg.ClientBase;
+      super();
+      this.db = db;
     } else {
-      this.isDBConnected =
-        config.shouldConnect !== undefined ? !config.shouldConnect : false;
-      this.db = config.client;
+      const config = configOrClient as PGVectorStoreConfig;
+      super(config?.embedModel);
+      this.schemaName = config?.schemaName ?? PGVECTOR_SCHEMA;
+      this.tableName = config?.tableName ?? PGVECTOR_TABLE;
+      this.user = config?.user;
+      this.password = config?.password;
+      this.database = config?.database;
+      this.connectionString = config?.connectionString;
+      this.dimensions = config?.dimensions ?? 1536;
     }
   }
 
@@ -106,41 +113,39 @@ export class PGVectorStore
 
   private async getDb(): Promise<pg.ClientBase> {
     if (!this.db) {
-      const pg = await import("pg");
-      const { Client } = pg.default ? pg.default : pg;
+      try {
+        const pg = await import("pg");
+        const { Client } = pg.default ? pg.default : pg;
 
-      const { registerTypes } = await import("pgvector/pg");
-      // Create DB connection
-      // Read connection params from env - see comment block above
-      const db = new Client({
-        ...this.clientConfig,
-      });
+        const { registerType } = await import("pgvector/pg");
+        // Create DB connection
+        // Read connection params from env - see comment block above
+        const db = new Client({
+          user: this.user,
+          password: this.password,
+          database: this.database,
+          connectionString: this.connectionString,
+        });
+        await db.connect();
 
-      await db.connect();
-      this.isDBConnected = true;
+        // Check vector extension
+        await db.query("CREATE EXTENSION IF NOT EXISTS vector");
+        await registerType(db);
 
-      // Check vector extension
-      await db.query("CREATE EXTENSION IF NOT EXISTS vector");
-      await registerTypes(db);
-
-      // All good?  Keep the connection reference
-      this.db = db;
+        // All good?  Keep the connection reference
+        this.db = db;
+      } catch (err) {
+        console.error(err);
+        return Promise.reject(err instanceof Error ? err : new Error(`${err}`));
+      }
     }
 
-    if (this.db && !this.isDBConnected) {
-      await this.db.connect();
-      this.isDBConnected = true;
-    }
-
-    this.db.on("end", () => {
-      // Connection closed
-      this.isDBConnected = false;
-    });
+    const db = this.db;
 
     // Check schema, table(s), index(es)
-    await this.checkSchema(this.db);
+    await this.checkSchema(db);
 
-    return this.db;
+    return Promise.resolve(this.db);
   }
 
   private async checkSchema(db: pg.ClientBase) {