Merge pull request #728 from hughestech/main

feat(document_loader): Add optional selector

Author: nfcampos

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

@@ -1,4 +1,4 @@
-import { test } from "@jest/globals";
+import { expect, test } from "@jest/globals";
 import { CheerioWebBaseLoader } from "../web/cheerio.js";
 
 test("Test cheerio web scraper loader", async () => {
@@ -7,3 +7,13 @@ test("Test cheerio web scraper loader", async () => {
   );
   await loader.load();
 });
+
+test("Test cheerio web scraper loader with selector", async () => {
+  const selectH1 = "h1";
+  const loader = new CheerioWebBaseLoader("https://about.google/commitments/", {}, selectH1);
+  
+  const doc = await loader.load();
+  expect(doc[0].pageContent.trim()).toBe(
+    "Committed to significantly improving the lives of as many people as possible."
+  );
+});

langchain/src/document_loaders/web/cheerio.ts

@@ -1,8 +1,8 @@
-import type { CheerioAPI, load as LoadT } from "cheerio";
-import { Document } from "../../document.js";
-import { BaseDocumentLoader } from "../base.js";
-import type { DocumentLoader } from "../base.js";
-import { AsyncCaller, AsyncCallerParams } from "../../util/async_caller.js";
+import type { CheerioAPI, load as LoadT, SelectorType } from "cheerio";
+import { Document } from "../document.js";
+import { BaseDocumentLoader } from "./base.js";
+import type { DocumentLoader } from "./base.js";
+import { AsyncCaller, AsyncCallerParams } from "../util/async_caller.js";
 
 export interface WebBaseLoaderParams extends AsyncCallerParams {
   /**
@@ -19,11 +19,14 @@ export class CheerioWebBaseLoader
 
   caller: AsyncCaller;
 
-  constructor(public webPath: string, fields?: WebBaseLoaderParams) {
+  selector?: SelectorType;
+
+  constructor(public webPath: string, fields?: WebBaseLoaderParams, selector?: SelectorType) {
     super();
     const { timeout, ...rest } = fields ?? {};
     this.timeout = timeout ?? 10000;
     this.caller = new AsyncCaller(rest);
+    this.selector = selector;
   }
 
   static async _scrape(
@@ -49,7 +52,7 @@ export class CheerioWebBaseLoader
 
   async load(): Promise<Document[]> {
     const $ = await this.scrape();
-    const text = $("body").text();
+    const text = $(this.selector ?? "body").text();
     const metadata = { source: this.webPath };
     return [new Document({ pageContent: text, metadata })];
   }