Release 0.5.26 (#1193)

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

.github/workflows/test.yml

@@ -12,6 +12,10 @@ concurrency:
   group: ${{ github.workflow }}-${{ github.ref }}
   cancel-in-progress: true
 
+env:
+  POSTGRES_USER: runneradmin
+  POSTGRES_HOST_AUTH_METHOD: trust
+
 jobs:
   e2e:
     strategy:
@@ -19,17 +23,20 @@ jobs:
       matrix:
         node-version: [18.x, 20.x, 22.x]
     name: E2E on Node.js ${{ matrix.node-version }}
-
-    env: POSTGRES_DB=vectordb
-      POSTGRES_USER=testuser
-      POSTGRES_PASSWORD=testpwd
-      POSTGRES_HOST_AUTH_METHOD=trust
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
-
+      - uses: ankane/setup-postgres@v1
+        with:
+          database: llamaindex_node_test
+          dev-files: true
+      - run: |
+          cd /tmp
+          git clone --branch v0.7.0 https://github.com/pgvector/pgvector.git
+          cd pgvector
+          make
+          sudo make install
       - uses: pnpm/action-setup@v4
-
       - name: Setup Node.js
         uses: actions/setup-node@v4
         with:
@@ -47,19 +54,8 @@ jobs:
         node-version: [18.x, 20.x, 22.x]
     name: Test on Node.js ${{ matrix.node-version }}
     runs-on: ubuntu-latest
-
     steps:
       - uses: actions/checkout@v4
-      - uses: ankane/setup-postgres@v1
-        with:
-          database: llamaindex_node_test
-          dev-files: true
-      - run: |
-          cd /tmp
-          git clone --branch v0.7.0 https://github.com/pgvector/pgvector.git
-          cd pgvector
-          make
-          sudo make install
       - uses: pnpm/action-setup@v4
       - name: Setup Node.js
         uses: actions/setup-node@v4
@@ -107,7 +103,7 @@ jobs:
           - nextjs-agent
           - nextjs-edge-runtime
           - nextjs-node-runtime
-    #          - waku-query-engine
+          - waku-query-engine
     runs-on: ubuntu-latest
     name: Build LlamaIndex Example (${{ matrix.packages }})
     steps:

README.md

@@ -189,6 +189,21 @@ export async function chatWithAgent(
 }
 ```
 
+### Vite
+
+We have some wasm dependencies for better performance. You can use `vite-plugin-wasm` to load them.
+
+```ts
+import wasm from "vite-plugin-wasm";
+
+export default {
+  plugins: [wasm()],
+  ssr: {
+    external: ["tiktoken"],
+  },
+};
+```
+
 ## Playground
 
 Check out our NextJS playground at https://llama-playground.vercel.app/. The source is available at https://github.com/run-llama/ts-playground

apps/docs/CHANGELOG.md

@@ -1,5 +1,13 @@
 # docs
 
+## 0.0.67
+
+### Patch Changes
+
+- Updated dependencies [ffe0cd1]
+- Updated dependencies [ffe0cd1]
+  - llamaindex@0.5.26
+
 ## 0.0.66
 
 ### Patch Changes

apps/docs/docs/modules/storage.md

@@ -4,12 +4,19 @@ sidebar_position: 7
 
 # Storage
 
-Storage in LlamaIndex.TS works automatically once you've configured a `StorageContext` object. Just configure the `persistDir` and attach it to an index.
+Storage in LlamaIndex.TS works automatically once you've configured a
+`StorageContext` object.
 
-Right now, only saving and loading from disk is supported, with future integrations planned!
+## Local Storage
+
+You can configure the `persistDir` and attach it to an index.
 
 ```typescript
-import { Document, VectorStoreIndex, storageContextFromDefaults } from "./src";
+import {
+  Document,
+  VectorStoreIndex,
+  storageContextFromDefaults,
+} from "llamaindex";
 
 const storageContext = await storageContextFromDefaults({
   persistDir: "./storage",
@@ -21,6 +28,33 @@ const index = await VectorStoreIndex.fromDocuments([document], {
 });
 ```
 
+## PostgreSQL Storage
+
+You can configure the `schemaName`, `tableName`, `namespace`, and
+`connectionString`. If a `connectionString` is not
+provided, it will use the environment variables `PGHOST`, `PGUSER`,
+`PGPASSWORD`, `PGDATABASE` and `PGPORT`.
+
+```typescript
+import {
+  Document,
+  VectorStoreIndex,
+  PostgresDocumentStore,
+  PostgresIndexStore,
+  storageContextFromDefaults,
+} from "llamaindex";
+
+const storageContext = await storageContextFromDefaults({
+  docStore: new PostgresDocumentStore(),
+  indexStore: new PostgresIndexStore(),
+});
+
+const document = new Document({ text: "Test Text" });
+const index = await VectorStoreIndex.fromDocuments([document], {
+  storageContext,
+});
+```
+
 ## API Reference
 
 - [StorageContext](../api/interfaces/StorageContext.md)

apps/docs/package.json

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

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

@@ -1,5 +1,14 @@
 # @llamaindex/autotool-01-node-example
 
+## 0.0.7
+
+### Patch Changes
+
+- Updated dependencies [ffe0cd1]
+- Updated dependencies [ffe0cd1]
+  - llamaindex@0.5.26
+  - @llamaindex/autotool@2.0.1
+
 ## 0.0.6
 
 ### 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.6"
+  "version": "0.0.7"
 }

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

@@ -1,5 +1,14 @@
 # @llamaindex/autotool-02-next-example
 
+## 0.1.51
+
+### Patch Changes
+
+- Updated dependencies [ffe0cd1]
+- Updated dependencies [ffe0cd1]
+  - llamaindex@0.5.26
+  - @llamaindex/autotool@2.0.1
+
 ## 0.1.50
 
 ### Patch Changes

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

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

packages/autotool/package.json

@@ -51,7 +51,7 @@
     "unplugin": "^1.12.2"
   },
   "peerDependencies": {
-    "llamaindex": "^0.5.25",
+    "llamaindex": "^0.5.26",
     "openai": "^4",
     "typescript": "^4"
   },

packages/experimental/CHANGELOG.md

@@ -1,5 +1,13 @@
 # @llamaindex/experimental
 
+## 0.0.76
+
+### Patch Changes
+
+- Updated dependencies [ffe0cd1]
+- Updated dependencies [ffe0cd1]
+  - llamaindex@0.5.26
+
 ## 0.0.75
 
 ### Patch Changes

packages/experimental/package.json

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

packages/llamaindex/CHANGELOG.md

@@ -1,5 +1,12 @@
 # llamaindex
 
+## 0.5.26
+
+### Patch Changes
+
+- ffe0cd1: faet: add openai o1 support
+- ffe0cd1: feat: add PostgreSQL storage
+
 ## 0.5.25
 
 ### Patch Changes

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

@@ -1,5 +1,13 @@
 # @llamaindex/cloudflare-worker-agent-test
 
+## 0.0.60
+
+### Patch Changes
+
+- Updated dependencies [ffe0cd1]
+- Updated dependencies [ffe0cd1]
+  - llamaindex@0.5.26
+
 ## 0.0.59
 
 ### Patch Changes

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

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

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

@@ -1,5 +1,13 @@
 # @llamaindex/next-agent-test
 
+## 0.1.60
+
+### Patch Changes
+
+- Updated dependencies [ffe0cd1]
+- Updated dependencies [ffe0cd1]
+  - llamaindex@0.5.26
+
 ## 0.1.59
 
 ### Patch Changes

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

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

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

@@ -1,5 +1,13 @@
 # test-edge-runtime
 
+## 0.1.59
+
+### Patch Changes
+
+- Updated dependencies [ffe0cd1]
+- Updated dependencies [ffe0cd1]
+  - llamaindex@0.5.26
+
 ## 0.1.58
 
 ### Patch Changes

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

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

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

@@ -1,5 +1,13 @@
 # @llamaindex/next-node-runtime
 
+## 0.0.41
+
+### Patch Changes
+
+- Updated dependencies [ffe0cd1]
+- Updated dependencies [ffe0cd1]
+  - llamaindex@0.5.26
+
 ## 0.0.40
 
 ### Patch Changes

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

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

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

@@ -1,5 +1,13 @@
 # @llamaindex/waku-query-engine-test
 
+## 0.0.60
+
+### Patch Changes
+
+- Updated dependencies [ffe0cd1]
+- Updated dependencies [ffe0cd1]
+  - llamaindex@0.5.26
+
 ## 0.0.59
 
 ### Patch Changes

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

@@ -1,6 +1,6 @@
 {
   "name": "@llamaindex/waku-query-engine-test",
-  "version": "0.0.59",
+  "version": "0.0.60",
   "type": "module",
   "private": true,
   "scripts": {
@@ -20,6 +20,7 @@
     "@types/react-dom": "18.3.0",
     "autoprefixer": "10.4.20",
     "tailwindcss": "3.4.10",
-    "typescript": "5.5.4"
+    "typescript": "5.5.4",
+    "vite-plugin-wasm": "^3.3.0"
   }
 }

packages/llamaindex/e2e/examples/waku-query-engine/src/pages/_layout.tsx

@@ -11,15 +11,22 @@ export default async function RootLayout({ children }: RootLayoutProps) {
   const data = await getData();
 
   return (
-    <div className="font-['Nunito']">
-      <meta property="description" content={data.description} />
-      <link rel="icon" type="image/png" href={data.icon} />
-      <Header />
-      <main className="m-6 flex items-center *:min-h-64 *:min-w-64 lg:m-0 lg:min-h-svh lg:justify-center">
-        {children}
-      </main>
-      <Footer />
-    </div>
+    <html>
+      <head>
+        <meta property="description" content={data.description} />
+        <link rel="icon" type="image/png" href={data.icon} />
+        <title>LlamaIndex Waku Example</title>
+      </head>
+      <body>
+        <div className="font-['Nunito']">
+          <Header />
+          <main className="m-6 flex items-center *:min-h-64 *:min-w-64 lg:m-0 lg:min-h-svh lg:justify-center">
+            {children}
+          </main>
+          <Footer />
+        </div>
+      </body>
+    </html>
   );
 }
 
@@ -31,7 +38,6 @@ const getData = async () => {
 
   return data;
 };
-
 export const getConfig = async () => {
   return {
     render: "static",

packages/llamaindex/e2e/examples/waku-query-engine/vite.config.ts

@@ -0,0 +1,8 @@
+import wasm from "vite-plugin-wasm";
+
+export default {
+  plugins: [wasm()],
+  ssr: {
+    external: ["tiktoken"],
+  },
+};

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

@@ -0,0 +1,90 @@
+import { Document, VectorStoreQueryMode } from "llamaindex";
+import { PGVectorStore } from "llamaindex/vector-store/PGVectorStore";
+import assert from "node:assert";
+import { test } from "node:test";
+import pg from "pg";
+import { registerTypes } from "pgvector/pg";
+
+let pgClient: pg.Client | pg.Pool;
+test.afterEach(async () => {
+  await pgClient.end();
+});
+
+await test("init with client", async () => {
+  pgClient = new pg.Client({
+    database: "llamaindex_node_test",
+  });
+  await pgClient.connect();
+  await pgClient.query("CREATE EXTENSION IF NOT EXISTS vector");
+  await registerTypes(pgClient);
+  const vectorStore = new PGVectorStore(pgClient);
+  assert.deepStrictEqual(await vectorStore.client(), pgClient);
+});
+
+await test("init with pool", async () => {
+  pgClient = new pg.Pool({
+    database: "llamaindex_node_test",
+  });
+  await pgClient.query("CREATE EXTENSION IF NOT EXISTS vector");
+  const client = await pgClient.connect();
+  await registerTypes(client);
+  const vectorStore = new PGVectorStore(client);
+  assert.deepStrictEqual(await vectorStore.client(), client);
+  client.release();
+});
+
+await test("init without client", async () => {
+  const vectorStore = new PGVectorStore({
+    database: "llamaindex_node_test",
+  });
+  pgClient = (await vectorStore.client()) as pg.Client;
+  assert.notDeepStrictEqual(pgClient, undefined);
+});
+
+await test("simple node", async () => {
+  const dimensions = 3;
+  const schemaName =
+    "llamaindex_vector_store_test_" + Math.random().toString(36).substring(7);
+  const nodeId = "5bb16627-f6c0-459c-bb18-71642813ef21";
+  const node = new Document({
+    text: "hello world",
+    id_: nodeId,
+    embedding: [0.1, 0.2, 0.3],
+  });
+  const vectorStore = new PGVectorStore({
+    database: "llamaindex_node_test",
+    dimensions,
+    schemaName,
+  });
+
+  await vectorStore.add([node]);
+
+  {
+    const result = await vectorStore.query({
+      mode: VectorStoreQueryMode.DEFAULT,
+      similarityTopK: 1,
+      queryEmbedding: [1, 2, 3],
+    });
+    const actualJSON = result.nodes![0]!.toJSON();
+    assert.deepStrictEqual(actualJSON, {
+      ...node.toJSON(),
+      hash: actualJSON.hash,
+      metadata: actualJSON.metadata,
+    });
+    assert.deepStrictEqual(result.ids, [nodeId]);
+    assert.deepStrictEqual(result.similarities, [1]);
+  }
+
+  await vectorStore.delete(nodeId);
+
+  {
+    const result = await vectorStore.query({
+      mode: VectorStoreQueryMode.DEFAULT,
+      similarityTopK: 1,
+      queryEmbedding: [1, 2, 3],
+    });
+    assert.deepStrictEqual(result.nodes, []);
+  }
+
+  pgClient = (await vectorStore.client()) as pg.Client;
+});

packages/llamaindex/package.json

@@ -1,6 +1,6 @@
 {
   "name": "llamaindex",
-  "version": "0.5.25",
+  "version": "0.5.26",
   "license": "MIT",
   "type": "module",
   "keywords": [
@@ -56,7 +56,7 @@
     "md-utils-ts": "^2.0.0",
     "mongodb": "^6.7.0",
     "notion-md-crawler": "^1.0.0",
-    "openai": "^4.57.0",
+    "openai": "^4.60.0",
     "papaparse": "^5.4.1",
     "pathe": "^1.1.2",
     "portkey-ai": "0.1.16",

packages/llamaindex/src/llm/openai.ts

@@ -97,6 +97,9 @@ export function getOpenAISession(
 }
 
 export const GPT4_MODELS = {
+  "chatgpt-4o-latest": {
+    contextWindow: 128000,
+  },
   "gpt-4": { contextWindow: 8192 },
   "gpt-4-32k": { contextWindow: 32768 },
   "gpt-4-32k-0613": { contextWindow: 32768 },
@@ -129,12 +132,28 @@ export const GPT35_MODELS = {
   "gpt-3.5-turbo-0301": { contextWindow: 16385 },
 };
 
+export const O1_MODELS = {
+  "o1-preview": {
+    contextWindow: 128000,
+  },
+  "o1-preview-2024-09-12": {
+    contextWindow: 128000,
+  },
+  "o1-mini": {
+    contextWindow: 128000,
+  },
+  "o1-mini-2024-09-12": {
+    contextWindow: 128000,
+  },
+};
+
 /**
  * We currently support GPT-3.5 and GPT-4 models
  */
 export const ALL_AVAILABLE_OPENAI_MODELS = {
   ...GPT4_MODELS,
   ...GPT35_MODELS,
+  ...O1_MODELS,
 } satisfies Record<ChatModel, { contextWindow: number }>;
 
 export function isFunctionCallingModel(llm: LLM): llm is OpenAI {
@@ -148,7 +167,8 @@ export function isFunctionCallingModel(llm: LLM): llm is OpenAI {
   }
   const isChatModel = Object.keys(ALL_AVAILABLE_OPENAI_MODELS).includes(model);
   const isOld = model.includes("0314") || model.includes("0301");
-  return isChatModel && !isOld;
+  const isO1 = model.startsWith("o1");
+  return isChatModel && !isOld && !isO1;
 }
 
 export type OpenAIAdditionalMetadata = {};

packages/llamaindex/src/storage/docStore/PostgresDocumentStore.ts

@@ -0,0 +1,21 @@
+import { DEFAULT_NAMESPACE } from "@llamaindex/core/global";
+import { PostgresKVStore } from "../kvStore/PostgresKVStore.js";
+import { KVDocumentStore } from "./KVDocumentStore.js";
+
+const DEFAULT_TABLE_NAME = "llamaindex_doc_store";
+
+export class PostgresDocumentStore extends KVDocumentStore {
+  constructor(config?: {
+    schemaName?: string;
+    tableName?: string;
+    connectionString?: string;
+    namespace?: string;
+  }) {
+    const kvStore = new PostgresKVStore({
+      schemaName: config?.schemaName,
+      tableName: config?.tableName || DEFAULT_TABLE_NAME,
+    });
+    const namespace = config?.namespace || DEFAULT_NAMESPACE;
+    super(kvStore, namespace);
+  }
+}

packages/llamaindex/src/storage/index.ts

@@ -1,10 +1,13 @@
 export { SimpleChatStore } from "./chatStore/SimpleChatStore.js";
 export * from "./chatStore/types.js";
+export { PostgresDocumentStore } from "./docStore/PostgresDocumentStore.js";
 export { SimpleDocumentStore } from "./docStore/SimpleDocumentStore.js";
 export * from "./docStore/types.js";
 export * from "./FileSystem.js";
+export { PostgresIndexStore } from "./indexStore/PostgresIndexStore.js";
 export { SimpleIndexStore } from "./indexStore/SimpleIndexStore.js";
 export * from "./indexStore/types.js";
+export { PostgresKVStore } from "./kvStore/PostgresKVStore.js";
 export { SimpleKVStore } from "./kvStore/SimpleKVStore.js";
 export * from "./kvStore/types.js";
 export * from "./StorageContext.js";

packages/llamaindex/src/storage/indexStore/PostgresIndexStore.ts

@@ -0,0 +1,21 @@
+import { DEFAULT_NAMESPACE } from "@llamaindex/core/global";
+import { PostgresKVStore } from "../kvStore/PostgresKVStore.js";
+import { KVIndexStore } from "./KVIndexStore.js";
+
+const DEFAULT_TABLE_NAME = "llamaindex_index_store";
+
+export class PostgresIndexStore extends KVIndexStore {
+  constructor(config?: {
+    schemaName?: string;
+    tableName?: string;
+    connectionString?: string;
+    namespace?: string;
+  }) {
+    const kvStore = new PostgresKVStore({
+      schemaName: config?.schemaName,
+      tableName: config?.tableName || DEFAULT_TABLE_NAME,
+    });
+    const namespace = config?.namespace || DEFAULT_NAMESPACE;
+    super(kvStore, namespace);
+  }
+}

packages/llamaindex/src/storage/kvStore/PostgresKVStore.ts

@@ -0,0 +1,140 @@
+import { DEFAULT_COLLECTION } from "@llamaindex/core/global";
+import type pg from "pg";
+import { BaseKVStore } from "./types.js";
+
+export type DataType = Record<string, Record<string, any>>;
+
+const DEFAULT_SCHEMA_NAME = "public";
+const DEFAULT_TABLE_NAME = "llamaindex_kv_store";
+
+export class PostgresKVStore extends BaseKVStore {
+  private schemaName: string;
+  private tableName: string;
+  private connectionString: string | undefined = undefined;
+  private db?: pg.Client;
+
+  constructor(config?: {
+    schemaName?: string | undefined;
+    tableName?: string | undefined;
+    connectionString?: string | undefined;
+  }) {
+    super();
+    this.schemaName = config?.schemaName || DEFAULT_SCHEMA_NAME;
+    this.tableName = config?.tableName || DEFAULT_TABLE_NAME;
+    this.connectionString = config?.connectionString;
+  }
+
+  private async getDb(): Promise<pg.Client> {
+    if (!this.db) {
+      try {
+        const pg = await import("pg");
+        const { Client } = pg.default ? pg.default : pg;
+        const db = new Client({ connectionString: this.connectionString });
+        await db.connect();
+        await this.checkSchema(db);
+        this.db = db;
+      } catch (err) {
+        console.error(err);
+        return Promise.reject(err instanceof Error ? err : new Error(`${err}`));
+      }
+    }
+    return Promise.resolve(this.db);
+  }
+
+  private async checkSchema(db: pg.Client) {
+    await db.query(`CREATE SCHEMA IF NOT EXISTS ${this.schemaName}`);
+    const tbl = `CREATE TABLE IF NOT EXISTS ${this.schemaName}.${this.tableName} (
+      id uuid DEFAULT gen_random_uuid() PRIMARY KEY,
+      collection VARCHAR,
+      key VARCHAR,
+      value JSONB DEFAULT '{}'
+    )`;
+    await db.query(tbl);
+    const idxs = `CREATE INDEX IF NOT EXISTS idx_${this.tableName}_collection ON ${this.schemaName}.${this.tableName} (collection);
+      CREATE INDEX IF NOT EXISTS idx_${this.tableName}_key ON ${this.schemaName}.${this.tableName} (key);`;
+    await db.query(idxs);
+    return db;
+  }
+
+  client() {
+    return this.getDb();
+  }
+
+  async put(
+    key: string,
+    val: any,
+    collection: string = DEFAULT_COLLECTION,
+  ): Promise<void> {
+    const db = await this.getDb();
+    try {
+      await db.query("BEGIN");
+      const sql = `
+        INSERT INTO ${this.schemaName}.${this.tableName}
+          (collection, key, value)
+        VALUES ($1, $2, $3)
+        ON CONFLICT (id) DO UPDATE SET
+          collection = EXCLUDED.collection,
+          key = EXCLUDED.key,
+          value = EXCLUDED.value
+        RETURNING id
+      `;
+      const values = [collection, key, val];
+      await db.query(sql, values);
+      await db.query("COMMIT");
+    } catch (error) {
+      await db.query("ROLLBACK");
+      throw error;
+    }
+  }
+
+  async get(
+    key: string,
+    collection: string = DEFAULT_COLLECTION,
+  ): Promise<any> {
+    const db = await this.getDb();
+    try {
+      await db.query("BEGIN");
+      const sql = `SELECT * FROM ${this.schemaName}.${this.tableName} WHERE key = $1 AND collection = $2`;
+      const result = await db.query(sql, [key, collection]);
+      await db.query("COMMIT");
+      return result.rows[0].value;
+    } catch (error) {
+      await db.query("ROLLBACK");
+      throw error;
+    }
+  }
+
+  async getAll(collection: string = DEFAULT_COLLECTION): Promise<DataType> {
+    const db = await this.getDb();
+    try {
+      await db.query("BEGIN");
+      const sql = `SELECT * FROM ${this.schemaName}.${this.tableName} WHERE collection = $1`;
+      const result = await db.query(sql, [collection]);
+      await db.query("COMMIT");
+      return result.rows.reduce((acc, row) => {
+        acc[row.key] = row.value;
+        return acc;
+      }, {});
+    } catch (error) {
+      await db.query("ROLLBACK");
+      throw error;
+    }
+  }
+
+  async delete(
+    key: string,
+    collection: string = DEFAULT_COLLECTION,
+  ): Promise<boolean> {
+    const db = await this.getDb();
+    try {
+      await db.query("BEGIN");
+      const sql = `DELETE FROM ${this.schemaName}.${this.tableName} WHERE key = $1 AND collection = $2`;
+      const result = await db.query(sql, [key, collection]);
+      await db.query("COMMIT");
+      return !!result.rowCount && result.rowCount > 0;
+    } catch (error) {
+      await db.query("ROLLBACK");
+      throw error;
+    }
+  }
+}

packages/llamaindex/tests/vector-stores/pg-vector-store.test.ts

@@ -1,85 +0,0 @@
-import { Document } from "llamaindex";
-import { PGVectorStore } from "llamaindex/vector-store/PGVectorStore";
-import pg from "pg";
-import { registerTypes } from "pgvector/pg";
-import { describe, expect, test } from "vitest";
-import { VectorStoreQueryMode } from "../../src/index.js";
-
-describe("pg - init", () => {
-  test("init with client", async () => {
-    const client = new pg.Client({
-      database: "llamaindex_node_test",
-    });
-    await client.connect();
-    await client.query("CREATE EXTENSION IF NOT EXISTS vector");
-    await registerTypes(client);
-    const vectorStore = new PGVectorStore(client);
-    expect(await vectorStore.client()).toBe(client);
-  });
-
-  test("init with pool", async () => {
-    const pool = new pg.Pool({
-      database: "llamaindex_node_test",
-    });
-    await pool.query("CREATE EXTENSION IF NOT EXISTS vector");
-    const client = await pool.connect();
-    await registerTypes(client);
-    const vectorStore = new PGVectorStore(client);
-    expect(await vectorStore.client()).toBe(client);
-  });
-
-  test("init without client", async () => {
-    const vectorStore = new PGVectorStore({
-      database: "llamaindex_node_test",
-    });
-    expect(await vectorStore.client()).toBeDefined();
-  });
-});
-
-describe("pg - save data", () => {
-  test("simple node", async () => {
-    const dimensions = 3;
-    const schemaName =
-      "llamaindex_vector_store_test_" + Math.random().toString(36).substring(7);
-    const nodeId = "5bb16627-f6c0-459c-bb18-71642813ef21";
-    const node = new Document({
-      text: "hello world",
-      id_: nodeId,
-      embedding: [0.1, 0.2, 0.3],
-    });
-    const vectorStore = new PGVectorStore({
-      database: "llamaindex_node_test",
-      dimensions,
-      schemaName,
-    });
-
-    await vectorStore.add([node]);
-
-    {
-      const result = await vectorStore.query({
-        mode: VectorStoreQueryMode.DEFAULT,
-        similarityTopK: 1,
-        queryEmbedding: [1, 2, 3],
-      });
-      const actualJSON = result.nodes![0]!.toJSON();
-      expect(actualJSON).toEqual({
-        ...node.toJSON(),
-        hash: actualJSON.hash,
-        metadata: actualJSON.metadata,
-      });
-      expect(result.ids).toEqual([nodeId]);
-      expect(result.similarities).toEqual([1]);
-    }
-
-    await vectorStore.delete(nodeId);
-
-    {
-      const result = await vectorStore.query({
-        mode: VectorStoreQueryMode.DEFAULT,
-        similarityTopK: 1,
-        queryEmbedding: [1, 2, 3],
-      });
-      expect(result.nodes).toEqual([]);
-    }
-  });
-});

pnpm-lock.yaml

@@ -152,7 +152,7 @@ importers:
         version: 2.4.6
       chromadb:
         specifier: ^1.8.1
-        version: 1.8.1(@google/generative-ai@0.12.0)(cohere-ai@7.13.0(@aws-sdk/client-sso-oidc@3.637.0(@aws-sdk/client-sts@3.637.0))(encoding@0.1.13))(encoding@0.1.13)(openai@4.57.0(encoding@0.1.13)(zod@3.23.8))
+        version: 1.8.1(@google/generative-ai@0.12.0)(cohere-ai@7.13.0(@aws-sdk/client-sso-oidc@3.637.0(@aws-sdk/client-sts@3.637.0))(encoding@0.1.13))(encoding@0.1.13)(openai@4.60.0(encoding@0.1.13)(zod@3.23.8))
       commander:
         specifier: ^12.1.0
         version: 12.1.0
@@ -167,7 +167,7 @@ importers:
         version: link:../packages/llamaindex
       mongodb:
         specifier: ^6.7.0
-        version: 6.8.0(@aws-sdk/credential-providers@3.637.0)
+        version: 6.8.0(@aws-sdk/credential-providers@3.637.0(@aws-sdk/client-sso-oidc@3.637.0(@aws-sdk/client-sts@3.637.0)))
       pathe:
         specifier: ^1.1.2
         version: 1.1.2
@@ -276,7 +276,7 @@ importers:
         version: 1.1.0(@types/react@18.3.5)(react@18.3.1)
       ai:
         specifier: ^3.3.21
-        version: 3.3.21(openai@4.57.0(zod@3.23.8))(react@18.3.1)(sswr@2.1.0(svelte@4.2.19))(svelte@4.2.19)(vue@3.4.38(typescript@5.5.4))(zod@3.23.8)
+        version: 3.3.21(openai@4.60.0(zod@3.23.8))(react@18.3.1)(sswr@2.1.0(svelte@4.2.19))(svelte@4.2.19)(vue@3.4.38(typescript@5.5.4))(zod@3.23.8)
       class-variance-authority:
         specifier: ^0.7.0
         version: 0.7.0
@@ -554,7 +554,7 @@ importers:
         version: 4.7.0
       chromadb:
         specifier: 1.8.1
-        version: 1.8.1(@google/generative-ai@0.12.0)(cohere-ai@7.13.0(@aws-sdk/client-sso-oidc@3.637.0(@aws-sdk/client-sts@3.637.0))(encoding@0.1.13))(encoding@0.1.13)(openai@4.57.0(encoding@0.1.13)(zod@3.23.8))
+        version: 1.8.1(@google/generative-ai@0.12.0)(cohere-ai@7.13.0(@aws-sdk/client-sso-oidc@3.637.0(@aws-sdk/client-sts@3.637.0))(encoding@0.1.13))(encoding@0.1.13)(openai@4.60.0(encoding@0.1.13)(zod@3.23.8))
       cohere-ai:
         specifier: 7.13.0
         version: 7.13.0(@aws-sdk/client-sso-oidc@3.637.0(@aws-sdk/client-sts@3.637.0))(encoding@0.1.13)
@@ -581,13 +581,13 @@ importers:
         version: 2.0.0
       mongodb:
         specifier: ^6.7.0
-        version: 6.8.0(@aws-sdk/credential-providers@3.637.0)
+        version: 6.8.0(@aws-sdk/credential-providers@3.637.0(@aws-sdk/client-sso-oidc@3.637.0(@aws-sdk/client-sts@3.637.0)))
       notion-md-crawler:
         specifier: ^1.0.0
         version: 1.0.0(encoding@0.1.13)
       openai:
-        specifier: ^4.57.0
-        version: 4.57.0(encoding@0.1.13)(zod@3.23.8)
+        specifier: ^4.60.0
+        version: 4.60.0(encoding@0.1.13)(zod@3.23.8)
       papaparse:
         specifier: ^5.4.1
         version: 5.4.1
@@ -700,7 +700,7 @@ importers:
     dependencies:
       ai:
         specifier: ^3.3.21
-        version: 3.3.21(openai@4.57.0(zod@3.23.8))(react@18.3.1)(sswr@2.1.0(svelte@4.2.19))(svelte@4.2.19)(vue@3.4.38(typescript@5.5.4))(zod@3.23.8)
+        version: 3.3.21(openai@4.60.0(zod@3.23.8))(react@18.3.1)(sswr@2.1.0(svelte@4.2.19))(svelte@4.2.19)(vue@3.4.38(typescript@5.5.4))(zod@3.23.8)
       llamaindex:
         specifier: workspace:*
         version: link:../../..
@@ -840,6 +840,9 @@ importers:
       typescript:
         specifier: 5.5.4
         version: 5.5.4
+      vite-plugin-wasm:
+        specifier: ^3.3.0
+        version: 3.3.0(vite@5.4.2(@types/node@22.5.1)(terser@5.31.6))
 
   packages/llamaindex/tests:
     devDependencies:
@@ -8385,6 +8388,15 @@ packages:
       zod:
         optional: true
 
+  openai@4.60.0:
+    resolution: {integrity: sha512-U/wNmrUPdfsvU1GrKRP5mY5YHR3ev6vtdfNID6Sauz+oquWD8r+cXPL1xiUlYniosPKajy33muVHhGS/9/t6KA==}
+    hasBin: true
+    peerDependencies:
+      zod: ^3.23.8
+    peerDependenciesMeta:
+      zod:
+        optional: true
+
   opener@1.5.2:
     resolution: {integrity: sha512-ur5UIdyw5Y7yEj9wLzhqXiy6GZ3Mwx0yGI+5sMn2r0N0v3cKJvUmFH5yPP+WXh9e0xfyzyJX95D8l088DNFj7A==}
     hasBin: true
@@ -10689,6 +10701,11 @@ packages:
     engines: {node: ^18.0.0 || >=20.0.0}
     hasBin: true
 
+  vite-plugin-wasm@3.3.0:
+    resolution: {integrity: sha512-tVhz6w+W9MVsOCHzxo6SSMSswCeIw4HTrXEi6qL3IRzATl83jl09JVO1djBqPSwfjgnpVHNLYcaMbaDX5WB/pg==}
+    peerDependencies:
+      vite: ^2 || ^3 || ^4 || ^5
+
   vite@5.4.2:
     resolution: {integrity: sha512-dDrQTRHp5C1fTFzcSaMxjk6vdpKvT+2/mIdE07Gw2ykehT49O0z/VHS3zZ8iV/Gh8BJJKHWOe5RjaNrW5xf/GA==}
     engines: {node: ^18.0.0 || >=20.0.0}
@@ -15958,7 +15975,7 @@ snapshots:
       clean-stack: 2.2.0
       indent-string: 4.0.0
 
-  ai@3.3.21(openai@4.57.0(zod@3.23.8))(react@18.3.1)(sswr@2.1.0(svelte@4.2.19))(svelte@4.2.19)(vue@3.4.38(typescript@5.5.4))(zod@3.23.8):
+  ai@3.3.21(openai@4.60.0(zod@3.23.8))(react@18.3.1)(sswr@2.1.0(svelte@4.2.19))(svelte@4.2.19)(vue@3.4.38(typescript@5.5.4))(zod@3.23.8):
     dependencies:
       '@ai-sdk/provider': 0.0.22
       '@ai-sdk/provider-utils': 1.0.17(zod@3.23.8)
@@ -15975,7 +15992,7 @@ snapshots:
       secure-json-parse: 2.7.0
       zod-to-json-schema: 3.23.2(zod@3.23.8)
     optionalDependencies:
-      openai: 4.57.0(zod@3.23.8)
+      openai: 4.60.0(zod@3.23.8)
       react: 18.3.1
       sswr: 2.1.0(svelte@4.2.19)
       svelte: 4.2.19
@@ -16652,14 +16669,14 @@ snapshots:
 
   chownr@2.0.0: {}
 
-  chromadb@1.8.1(@google/generative-ai@0.12.0)(cohere-ai@7.13.0(@aws-sdk/client-sso-oidc@3.637.0(@aws-sdk/client-sts@3.637.0))(encoding@0.1.13))(encoding@0.1.13)(openai@4.57.0(encoding@0.1.13)(zod@3.23.8)):
+  chromadb@1.8.1(@google/generative-ai@0.12.0)(cohere-ai@7.13.0(@aws-sdk/client-sso-oidc@3.637.0(@aws-sdk/client-sts@3.637.0))(encoding@0.1.13))(encoding@0.1.13)(openai@4.60.0(encoding@0.1.13)(zod@3.23.8)):
     dependencies:
       cliui: 8.0.1
       isomorphic-fetch: 3.0.0(encoding@0.1.13)
     optionalDependencies:
       '@google/generative-ai': 0.12.0
       cohere-ai: 7.13.0(@aws-sdk/client-sso-oidc@3.637.0(@aws-sdk/client-sts@3.637.0))(encoding@0.1.13)
-      openai: 4.57.0(encoding@0.1.13)(zod@3.23.8)
+      openai: 4.60.0(encoding@0.1.13)(zod@3.23.8)
     transitivePeerDependencies:
       - encoding
 
@@ -17747,6 +17764,16 @@ snapshots:
     transitivePeerDependencies:
       - supports-color
 
+  eslint-module-utils@2.8.2(@typescript-eslint/parser@8.3.0(eslint@8.57.0)(typescript@5.5.4))(eslint-import-resolver-node@0.3.9)(eslint@8.57.0):
+    dependencies:
+      debug: 3.2.7
+    optionalDependencies:
+      '@typescript-eslint/parser': 8.3.0(eslint@8.57.0)(typescript@5.5.4)
+      eslint: 8.57.0
+      eslint-import-resolver-node: 0.3.9
+    transitivePeerDependencies:
+      - supports-color
+
   eslint-plugin-import@2.29.1(@typescript-eslint/parser@8.3.0(eslint@8.57.0)(typescript@5.5.4))(eslint@8.57.0):
     dependencies:
       array-includes: 3.1.8
@@ -17757,7 +17784,7 @@ snapshots:
       doctrine: 2.1.0
       eslint: 8.57.0
       eslint-import-resolver-node: 0.3.9
-      eslint-module-utils: 2.8.2(@typescript-eslint/parser@7.2.0(eslint@8.57.0)(typescript@5.5.4))(eslint-import-resolver-node@0.3.9)(eslint-import-resolver-typescript@3.6.3(@typescript-eslint/parser@7.2.0(eslint@8.57.0)(typescript@5.5.4))(eslint-import-resolver-node@0.3.9)(eslint-plugin-import@2.29.1)(eslint@8.57.0))(eslint@8.57.0)
+      eslint-module-utils: 2.8.2(@typescript-eslint/parser@8.3.0(eslint@8.57.0)(typescript@5.5.4))(eslint-import-resolver-node@0.3.9)(eslint@8.57.0)
       hasown: 2.0.2
       is-core-module: 2.15.1
       is-glob: 4.0.3
@@ -20485,7 +20512,7 @@ snapshots:
     optionalDependencies:
       '@aws-sdk/credential-providers': 3.637.0(@aws-sdk/client-sso-oidc@3.637.0(@aws-sdk/client-sts@3.637.0))
 
-  mongodb@6.8.0(@aws-sdk/credential-providers@3.637.0):
+  mongodb@6.8.0(@aws-sdk/credential-providers@3.637.0(@aws-sdk/client-sso-oidc@3.637.0(@aws-sdk/client-sts@3.637.0))):
     dependencies:
       '@mongodb-js/saslprep': 1.1.7
       bson: 6.8.0
@@ -20915,7 +20942,23 @@ snapshots:
     transitivePeerDependencies:
       - encoding
 
-  openai@4.57.0(zod@3.23.8):
+  openai@4.60.0(encoding@0.1.13)(zod@3.23.8):
+    dependencies:
+      '@types/node': 18.19.47
+      '@types/node-fetch': 2.6.11
+      '@types/qs': 6.9.15
+      abort-controller: 3.0.0
+      agentkeepalive: 4.5.0
+      form-data-encoder: 1.7.2
+      formdata-node: 4.4.1
+      node-fetch: 2.7.0(encoding@0.1.13)
+      qs: 6.13.0
+    optionalDependencies:
+      zod: 3.23.8
+    transitivePeerDependencies:
+      - encoding
+
+  openai@4.60.0(zod@3.23.8):
     dependencies:
       '@types/node': 18.19.47
       '@types/node-fetch': 2.6.11
@@ -23340,7 +23383,7 @@ snapshots:
 
   union@0.5.0:
     dependencies:
-      qs: 6.11.2
+      qs: 6.13.0
 
   unique-string@3.0.0:
     dependencies:
@@ -23522,6 +23565,10 @@ snapshots:
       - supports-color
       - terser
 
+  vite-plugin-wasm@3.3.0(vite@5.4.2(@types/node@22.5.1)(terser@5.31.6)):
+    dependencies:
+      vite: 5.4.2(@types/node@22.5.1)(terser@5.31.6)
+
   vite@5.4.2(@types/node@22.5.1)(terser@5.31.6):
     dependencies:
       esbuild: 0.21.5