Add s3 file loader (#599)

* Add s3 loader

* revert changes to yarn.lock / package.json

* install aws-sdk/client-s3 as devDependency

* merge

* Revert "merge"

This reverts commit 2954cd7.

* Add s3 loader to new, refactored file structure

* Fix after import paths change

* Add entrypoint, make an args aobject

---------

Co-authored-by: Nuno Campos <nuno@boringbits.io>

docs/docs/modules/indexes/document_loaders/examples/file_loaders/s3.mdx

@@ -0,0 +1,27 @@
+---
+hide_table_of_contents: true
+sidebar_class_name: node-only
+---
+
+# S3 File
+
+:::tip Compatibility
+Only available on Node.js.
+:::
+
+This covers how to load document objects from an s3 file object.
+
+## Setup
+
+To run this index you'll need to have Unstructured already set up and ready to use at an available URL endpoint. It can also be configured to run locally.
+
+See the docs [here](https://js.langchain.com/docs/modules/indexes/document_loaders/examples/file_loaders/unstructured) for information on how to do that.
+
+## Usage
+
+Once Unstructured is configured, you can use the S3 loader to load files and then convert them into a Document.
+
+import CodeBlock from "@theme/CodeBlock";
+import Example from "@examples/document_loaders/s3.ts";
+
+<CodeBlock language="typescript">{Example}</CodeBlock>

examples/src/document_loaders/s3.ts

@@ -0,0 +1,11 @@
+import { S3Loader } from "langchain/document_loaders/web/s3";
+
+const loader = new S3Loader({
+  bucket: "my-document-bucket-123",
+  key: "AccountingOverview.pdf",
+  unstructuredAPIURL: "http://localhost:8000/general/v0/general",
+});
+
+const docs = await loader.load();
+
+console.log(docs);

langchain/.gitignore

@@ -121,6 +121,9 @@ document_loaders/web/imsdb.d.ts
 document_loaders/web/github.cjs
 document_loaders/web/github.js
 document_loaders/web/github.d.ts
+document_loaders/web/s3.cjs
+document_loaders/web/s3.js
+document_loaders/web/s3.d.ts
 document_loaders/fs/directory.cjs
 document_loaders/fs/directory.js
 document_loaders/fs/directory.d.ts

langchain/package.json

@@ -133,6 +133,9 @@
     "document_loaders/web/github.cjs",
     "document_loaders/web/github.js",
     "document_loaders/web/github.d.ts",
+    "document_loaders/web/s3.cjs",
+    "document_loaders/web/s3.js",
+    "document_loaders/web/s3.d.ts",
     "document_loaders/fs/directory.cjs",
     "document_loaders/fs/directory.js",
     "document_loaders/fs/directory.d.ts",
@@ -234,6 +237,7 @@
   "author": "LangChain",
   "license": "MIT",
   "devDependencies": {
+    "@aws-sdk/client-s3": "^3.310.0",
     "@faker-js/faker": "^7.6.0",
     "@getmetal/metal-sdk": "^1.0.12",
     "@huggingface/inference": "^1.5.1",
@@ -280,6 +284,7 @@
     "typescript": "^4.9.5"
   },
   "peerDependencies": {
+    "@aws-sdk/client-s3": "^3.310.0",
     "@getmetal/metal-sdk": "*",
     "@huggingface/inference": "^1.5.1",
     "@pinecone-database/pinecone": "*",
@@ -300,6 +305,9 @@
     "typeorm": "^0.3.12"
   },
   "peerDependenciesMeta": {
+    "@aws-sdk/client-s3": {
+      "optional": true
+    },
     "@getmetal/metal-sdk": {
       "optional": true
     },
@@ -607,6 +615,11 @@
       "import": "./document_loaders/web/github.js",
       "require": "./document_loaders/web/github.cjs"
     },
+    "./document_loaders/web/s3": {
+      "types": "./document_loaders/web/s3.d.ts",
+      "import": "./document_loaders/web/s3.js",
+      "require": "./document_loaders/web/s3.cjs"
+    },
     "./document_loaders/fs/directory": {
       "types": "./document_loaders/fs/directory.d.ts",
       "import": "./document_loaders/fs/directory.js",

langchain/scripts/create-entrypoints.js

@@ -62,6 +62,7 @@ const entrypoints = {
   "document_loaders/web/hn": "document_loaders/web/hn",
   "document_loaders/web/imsdb": "document_loaders/web/imsdb",
   "document_loaders/web/github": "document_loaders/web/github",
+  "document_loaders/web/s3": "document_loaders/web/s3",
   "document_loaders/fs/directory": "document_loaders/fs/directory",
   "document_loaders/fs/buffer": "document_loaders/fs/buffer",
   "document_loaders/fs/text": "document_loaders/fs/text",
@@ -130,6 +131,7 @@ const requiresOptionalDependency = [
   "document_loaders/web/hn",
   "document_loaders/web/imsdb",
   "document_loaders/web/github",
+  "document_loaders/web/s3",
   "document_loaders/fs/directory",
   "document_loaders/fs/buffer",
   "document_loaders/fs/text",

langchain/src/document_loaders/tests/s3.test.ts

@@ -0,0 +1,60 @@
+/* eslint-disable tree-shaking/no-side-effects-in-initialization */
+import { test, jest, expect } from "@jest/globals";
+import S3Client from "@aws-sdk/client-s3";
+import * as fs from "node:fs";
+import * as path from "node:path";
+import { Readable } from "node:stream";
+import { S3Loader } from "../web/s3.js";
+import { UnstructuredLoader } from "../fs/unstructured.js";
+
+const fsMock = {
+  ...fs,
+  mkdtempSync: jest.fn().mockReturnValue("tmp/s3fileloader-12345"),
+  mkdirSync: jest.fn().mockImplementation(() => {}),
+  writeFileSync: jest.fn().mockImplementation(() => {}),
+};
+
+const UnstructuredLoaderMock = jest.fn().mockImplementation(() => ({
+  load: jest.fn().mockImplementation(() => ["fake document"]),
+}));
+
+jest.mock("@aws-sdk/client-s3", () => ({
+  S3Client: jest.fn().mockImplementation(() => ({
+    send: jest.fn().mockImplementation(() =>
+      Promise.resolve({
+        Body: new Readable({
+          read() {
+            this.push(Buffer.from("Mock file content"));
+            this.push(null);
+          },
+        }),
+      })
+    ),
+  })),
+  GetObjectCommand: jest.fn(),
+}));
+
+test("Test S3 loader", async () => {
+  if (!S3Client) {
+    // this is to avoid a linting error. S3Client is mocked above.
+  }
+
+  const loader = new S3Loader({
+    bucket: "test-bucket-123",
+    key: "AccountingOverview.pdf",
+    unstructuredAPIURL: "http://localhost:8000/general/v0/general",
+    fs: fsMock as typeof fs,
+    UnstructuredLoader: UnstructuredLoaderMock as typeof UnstructuredLoader,
+  });
+
+  const result = await loader.load();
+
+  expect(fsMock.mkdtempSync).toHaveBeenCalled();
+  expect(fsMock.mkdirSync).toHaveBeenCalled();
+  expect(fsMock.writeFileSync).toHaveBeenCalled();
+  expect(UnstructuredLoaderMock).toHaveBeenCalledWith(
+    "http://localhost:8000/general/v0/general",
+    path.join("tmp", "s3fileloader-12345", "AccountingOverview.pdf")
+  );
+  expect(result).toEqual(["fake document"]);
+});

langchain/src/document_loaders/web/s3.ts

@@ -0,0 +1,112 @@
+import * as fsDefault from "node:fs";
+import * as path from "node:path";
+import * as os from "node:os";
+import { Readable } from "node:stream";
+import { BaseDocumentLoader } from "../base.js";
+import { UnstructuredLoader as UnstructuredLoaderDefault } from "../fs/unstructured.js";
+
+export interface S3LoaderParams {
+  bucket: string;
+  key: string;
+  unstructuredAPIURL: string;
+
+  fs?: typeof fsDefault;
+  UnstructuredLoader?: typeof UnstructuredLoaderDefault;
+}
+
+export class S3Loader extends BaseDocumentLoader {
+  private bucket: string;
+
+  private key: string;
+
+  private unstructuredAPIURL: string;
+
+  private _fs: typeof fsDefault;
+
+  private _UnstructuredLoader: typeof UnstructuredLoaderDefault;
+
+  constructor({
+    bucket,
+    key,
+    unstructuredAPIURL,
+    fs = fsDefault,
+    UnstructuredLoader = UnstructuredLoaderDefault,
+  }: S3LoaderParams) {
+    super();
+    this.bucket = bucket;
+    this.key = key;
+    this.unstructuredAPIURL = unstructuredAPIURL;
+    this._fs = fs;
+    this._UnstructuredLoader = UnstructuredLoader;
+  }
+
+  public async load() {
+    const { S3Client, GetObjectCommand } = await S3LoaderImports();
+
+    const tempDir = this._fs.mkdtempSync(
+      path.join(os.tmpdir(), "s3fileloader-")
+    );
+
+    const filePath = path.join(tempDir, this.key);
+
+    try {
+      const s3Client = new S3Client({});
+
+      const getObjectCommand = new GetObjectCommand({
+        Bucket: this.bucket,
+        Key: this.key,
+      });
+
+      const response = await s3Client.send(getObjectCommand);
+
+      const objectData = await new Promise<Buffer>((resolve, reject) => {
+        const chunks: Buffer[] = [];
+
+        // eslint-disable-next-line no-instanceof/no-instanceof
+        if (response.Body instanceof Readable) {
+          response.Body.on("data", (chunk: Buffer) => chunks.push(chunk));
+          response.Body.on("end", () => resolve(Buffer.concat(chunks)));
+          response.Body.on("error", reject);
+        } else {
+          reject(new Error("Response body is not a readable stream."));
+        }
+      });
+
+      this._fs.mkdirSync(path.dirname(filePath), { recursive: true });
+
+      this._fs.writeFileSync(filePath, objectData);
+    } catch {
+      throw new Error(
+        `Failed to download file ${this.key} from S3 bucket ${this.bucket}.`
+      );
+    }
+
+    try {
+      const unstructuredLoader = new this._UnstructuredLoader(
+        this.unstructuredAPIURL,
+        filePath
+      );
+
+      const docs = await unstructuredLoader.load();
+
+      return docs;
+    } catch {
+      throw new Error(
+        `Failed to load file ${filePath} using unstructured loader.`
+      );
+    }
+  }
+}
+
+async function S3LoaderImports() {
+  try {
+    const s3Module = await import("@aws-sdk/client-s3");
+
+    return s3Module as typeof s3Module;
+  } catch (e) {
+    console.error(e);
+    throw new Error(
+      "Failed to load @aws-sdk/client-s3'. Please install it eg. `yarn add @aws-sdk/client-s3`."
+    );
+  }
+}

langchain/tsconfig.json

@@ -71,6 +71,7 @@
       "src/document_loaders/web/hn.ts",
       "src/document_loaders/web/imsdb.ts",
       "src/document_loaders/web/github.ts",
+      "src/document_loaders/web/s3.ts",
       "src/document_loaders/fs/directory.ts",
       "src/document_loaders/fs/buffer.ts",
       "src/document_loaders/fs/text.ts",