firecrawl/apps/api/src/lib/extract/completions/batchExtract.ts

import {
  generateCompletions,
  GenerateCompletionsOptions,
} from "../../../scraper/scrapeURL/transformers/llmExtract";
import { buildDocument } from "../build-document";
import { ExtractResponse, TokenUsage } from "../../../controllers/v1/types";
import { Document } from "../../../controllers/v1/types";
import {
  buildBatchExtractPrompt,
  buildBatchExtractSystemPrompt,
} from "../build-prompts";
import { getModel } from "../../generic-ai";

import fs from "fs/promises";
import { extractData } from "../../../scraper/scrapeURL/lib/extractSmartScrape";
import type { Logger } from "winston";

type BatchExtractOptions = {
  multiEntitySchema: any;
  links: string[];
  prompt: string;
  systemPrompt: string;
  doc: Document;
  useAgent: boolean;
  extractId?: string;
};

/**
 * Batch extract information from a list of URLs using a multi-entity schema.
 * @param multiEntitySchema - The schema for the multi-entity extraction
 * @param links - The URLs to extract information from
 * @param prompt - The prompt for the extraction
 * @param systemPrompt - The system prompt for the extraction
 * @param doc - The document to extract information from
 * @returns The completion promise
 */
export async function batchExtractPromise(options: BatchExtractOptions, logger: Logger): Promise<{
  extract: any; // array of extracted data
  numTokens: number;
  totalUsage: TokenUsage;
  warning?: string;
  sources: string[];
  smartScrapeCost: number;
  otherCost: number;
  smartScrapeCallCount: number;
  otherCallCount: number;
}> {
  const { multiEntitySchema, links, prompt, systemPrompt, doc, useAgent, extractId } = options;


  const generationOptions: GenerateCompletionsOptions = {
    logger: logger.child({
      method: "extractService/generateCompletions",
    }),
    options: {
      mode: "llm",
      systemPrompt: buildBatchExtractSystemPrompt(
        systemPrompt,
        multiEntitySchema,
        links,
      ),
      prompt: buildBatchExtractPrompt(prompt),
      schema: multiEntitySchema,
    },
    markdown: buildDocument(doc),
    isExtractEndpoint: true,
    model: getModel("gemini-2.0-flash", "google"),
  };

  let extractedDataArray: any[] = [];
  let warning: string | undefined;
  let smCost = 0, oCost = 0, smCallCount = 0, oCallCount = 0;
  try {
    const { extractedDataArray: e, warning: w, smartScrapeCost, otherCost, smartScrapeCallCount, otherCallCount } = await extractData({
      extractOptions: generationOptions,
      urls: [doc.metadata.sourceURL || doc.metadata.url || ""],
      useAgent,
      extractId,
    });
    extractedDataArray = e;
    warning = w;
    smCost = smartScrapeCost;
    oCost = otherCost;
    smCallCount = smartScrapeCallCount;
    oCallCount = otherCallCount;
  } catch (error) {
    logger.error("extractData failed", { error });
  }

  // await fs.writeFile(
  //   `logs/extractedDataArray-${crypto.randomUUID()}.json`,
  //   JSON.stringify(extractedDataArray, null, 2),
  // );

  // TODO: fix this
  return {
    extract: extractedDataArray,
    numTokens: 0,
    totalUsage: {
      promptTokens: 0,
      completionTokens: 0,
      totalTokens: 0,
      model: "gemini-2.0-flash",
    },
    warning: warning,
    sources: [doc.metadata.url || doc.metadata.sourceURL || ""],
    smartScrapeCost: smCost,
    otherCost: oCost,
    smartScrapeCallCount: smCallCount,
    otherCallCount: oCallCount,
  };
}