@lg-tools/crawler/src/utils/processSingleUrl.ts

Crawler for MongoDB documentation and other sites

/* eslint-disable no-console */
import chalk from 'chalk';
import { RecursiveCharacterTextSplitter } from 'langchain/text_splitter';
import { type MongoClient } from 'mongodb';

import { MDB_DB } from '../constants';
import { CrawlerDocument } from '../types';

import { checkRobots } from './checkRobots';
import { createChecksum } from './createChecksum';
import { loadPageContents } from './loadPageContents';
import { makeLangChainDocumentMapper } from './mapLangChainDocument';

interface ProcessSingleUrlOptions {
  href: string;
  collectionName: string;
  mongoClient: MongoClient;
  verbose?: boolean;
  dryRun?: boolean;
}

interface ProcessSingleUrlResult {
  docCount: number;
  links: Array<string>;
  documents?: Array<CrawlerDocument>;
}

/**
 * Process a single URL using LangChain and store results in MongoDB.
 * Additionally, extracts and returns all links found on the page
 */
export async function processSingleUrl({
  href,
  collectionName,
  mongoClient,
  verbose = false,
  dryRun = false,
}: ProcessSingleUrlOptions): Promise<ProcessSingleUrlResult> {
  try {
    verbose && console.log(chalk.gray(`Processing URL:`), chalk.blue(href));
    const { hostname, pathname } = new URL(href);

    const isCrawlingAllowed = checkRobots(hostname, pathname, verbose);

    if (!isCrawlingAllowed) {
      verbose &&
        console.log(
          chalk.red(
            `Crawling disallowed by robots.txt for ${hostname}. Skipping...`,
          ),
        );
      return {
        docCount: 0,
        links: [],
      };
    }

    const { doc, title, links } = await loadPageContents(href);

    verbose &&
      console.log(
        chalk.gray(
          `Loaded document`,
          `"${chalk.bold(title)}"`,
          `from`,
          chalk.blue(href),
        ),
      );

    // Split text into chunks for better processing/storage
    const textSplitter = new RecursiveCharacterTextSplitter({
      chunkSize: 1000,
      chunkOverlap: 200,
    });
    const chunkedDocs = await textSplitter.splitDocuments([doc]);

    verbose &&
      console.log(chalk.gray(`Split into ${chunkedDocs.length} chunks`));

    const checksum = createChecksum(href, doc.pageContent);

    const documents = chunkedDocs.map(
      makeLangChainDocumentMapper({ title, href, checksum }),
    );

    if (chunkedDocs.length <= 0) {
      verbose &&
        console.log(chalk.gray(`No content to process for URL: ${href}`));
      return {
        docCount: 0,
        links,
        documents: [],
      };
    }

    let docCount = 0;

    // Skip MongoDB insertion if in dry run mode
    if (dryRun) {
      if (verbose) {
        console.log(
          chalk.yellow(
            `[DRY RUN] Would insert ${documents.length} documents into collection: "${collectionName}"`,
          ),
        );
      }

      // For dry run, we'll count the documents that would have been processed
      docCount = documents.length;
    }
    // Normal mode - store in MongoDB
    else {
      const db = mongoClient.db(MDB_DB);
      const collection = db.collection(collectionName);

      if (verbose) {
        console.log(
          chalk.green(
            `Inserting ${documents.length} documents into collection: "${collectionName}"`,
          ),
        );
      }

      if (documents.length > 0) {
        const result = await collection.insertMany(documents);
        docCount = result.insertedCount;

        verbose &&
          console.log(
            chalk.gray(
              `Successfully inserted ${result.insertedCount} documents into collection ${collectionName}`,
            ),
          );
      } else {
        verbose &&
          console.log(chalk.gray(`No documents to insert for URL: ${href}`));
      }
    }

    return {
      docCount,
      links,
      documents,
    };
  } catch (error) {
    console.error(
      chalk.red(`Error processing URL ${href}:`),
      chalk.gray(error),
    );
    return {
      docCount: 0,
      links: [],
    };
  }
}