feat: support npm:postgres (run-llama#1248)

.changeset/weak-camels-perform.md

@@ -0,0 +1,8 @@
+---
+"@llamaindex/core": patch
+"llamaindex": patch
+"@llamaindex/core-e2e": patch
+"pg-vector-store": patch
+---
+
+feat: support `npm:postgres`

examples/package.json

@@ -16,7 +16,8 @@
     "js-tiktoken": "^1.0.14",
     "llamaindex": "^0.6.0",
     "mongodb": "^6.7.0",
-    "pathe": "^1.1.2"
+    "pathe": "^1.1.2",
+    "postgres": "^3.4.4"
   },
   "devDependencies": {
     "@types/node": "^22.5.1",

examples/vector-store/pg/.env.template

@@ -0,0 +1,6 @@
+# neon template
+PGHOST=
+PGDATABASE=
+PGUSER=
+PGPASSWORD=
+ENDPOINT_ID=

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

@@ -1,11 +1,11 @@
 // load-docs.ts
-import fs from "fs/promises";
 import {
   PGVectorStore,
   SimpleDirectoryReader,
   storageContextFromDefaults,
   VectorStoreIndex,
 } from "llamaindex";
+import fs from "node:fs/promises";
 
 async function getSourceFilenames(sourceDir: string) {
   return await fs

examples/vector-store/pg/neon.ts

@@ -0,0 +1,45 @@
+/* eslint-disable turbo/no-undeclared-env-vars */
+import dotenv from "dotenv";
+import { Document, PGVectorStore, VectorStoreQueryMode } from "llamaindex";
+import postgres from "postgres";
+
+dotenv.config();
+
+const { PGHOST, PGDATABASE, PGUSER, ENDPOINT_ID } = process.env;
+const PGPASSWORD = decodeURIComponent(process.env.PGPASSWORD!);
+
+const sql = postgres({
+  host: PGHOST,
+  database: PGDATABASE,
+  username: PGUSER,
+  password: PGPASSWORD,
+  port: 5432,
+  ssl: "require",
+  connection: {
+    options: `project=${ENDPOINT_ID}`,
+  },
+});
+
+await sql`CREATE EXTENSION IF NOT EXISTS vector`;
+
+const vectorStore = new PGVectorStore({
+  dimensions: 3,
+  client: sql,
+});
+
+await vectorStore.add([
+  new Document({
+    text: "hello, world",
+    embedding: [1, 2, 3],
+  }),
+]);
+
+const results = await vectorStore.query({
+  mode: VectorStoreQueryMode.DEFAULT,
+  similarityTopK: 1,
+  queryEmbedding: [1, 2, 3],
+});
+
+console.log("result", results);
+
+await sql.end();

examples/vector-store/pg/package.json

@@ -0,0 +1,5 @@
+{
+  "name": "pg-vector-store",
+  "type": "module",
+  "private": true
+}

examples/vector-store/pg/tsconfig.json

@@ -0,0 +1,9 @@
+{
+  "extends": "../../tsconfig.json",
+  "compilerOptions": {
+    "outDir": "./dist",
+    "types": ["node"],
+    "skipLibCheck": true
+  },
+  "include": ["./**/*.ts"]
+}

packages/core/package.json

@@ -255,6 +255,20 @@
         "types": "./retriever/dist/index.d.ts",
         "default": "./retriever/dist/index.js"
       }
+    },
+    "./vector-store": {
+      "require": {
+        "types": "./dist/vector-store/index.d.cts",
+        "default": "./dist/vector-store/index.cjs"
+      },
+      "import": {
+        "types": "./dist/vector-store/index.d.ts",
+        "default": "./dist/vector-store/index.js"
+      },
+      "default": {
+        "types": "./dist/vector-store/index.d.ts",
+        "default": "./dist/vector-store/index.js"
+      }
     }
   },
   "files": [

packages/core/src/vector-store/index.ts

@@ -0,0 +1,13 @@
+/**
+ * should compatible with npm:pg and npm:postgres
+ */
+export interface IsomorphicDB {
+  query: (sql: string, params?: any[]) => Promise<any[]>;
+  // begin will wrap the multiple queries in a transaction
+  begin: <T>(fn: (query: IsomorphicDB["query"]) => Promise<T>) => Promise<T>;
+
+  // event handler
+  connect: () => Promise<void>;
+  close: () => Promise<void>;
+  onCloseEvent: (listener: () => void) => void;
+}

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

@@ -27,7 +27,7 @@ await test("init with client", async (t) => {
     client: pgClient,
     shouldConnect: false,
   });
-  assert.deepStrictEqual(await vectorStore.client(), pgClient);
+  assert.notDeepStrictEqual(await vectorStore.client(), undefined);
 });
 
 await test("init with pool", async (t) => {
@@ -44,16 +44,16 @@ await test("init with pool", async (t) => {
     shouldConnect: false,
     client,
   });
-  assert.deepStrictEqual(await vectorStore.client(), client);
+  assert.notDeepStrictEqual(await vectorStore.client(), undefined);
 });
 
 await test("init without client", async (t) => {
   const vectorStore = new PGVectorStore({ clientConfig: pgConfig });
-  const pgClient = (await vectorStore.client()) as pg.Client;
+  const db = await vectorStore.client();
   t.after(async () => {
-    await pgClient.end();
+    await db.close();
   });
-  assert.notDeepStrictEqual(pgClient, undefined);
+  assert.notDeepStrictEqual(db, undefined);
 });
 
 await test("simple node", async (t) => {
@@ -71,9 +71,9 @@ await test("simple node", async (t) => {
     dimensions,
     schemaName,
   });
-  const pgClient = (await vectorStore.client()) as pg.Client;
+  const db = await vectorStore.client();
   t.after(async () => {
-    await pgClient.end();
+    await db.close();
   });
 
   await vectorStore.add([node]);

packages/llamaindex/package.json

@@ -97,6 +97,7 @@
     "glob": "^11.0.0",
     "pg": "^8.12.0",
     "pgvector": "0.2.0",
+    "postgres": "^3.4.4",
     "typescript": "^5.6.2"
   },
   "engines": {