AI/firecrawl
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/mendableai/firecrawl synced 2025-07-30 08:11:59 +08:00
Code Issues Packages Projects Releases Wiki Activity

new cost tracking

Browse Source
This commit is contained in:
Gergő Móricz 2025-04-17 09:23:53 -07:00
parent ba4df67de7
commit 8546bcacc0
27 changed files with 395 additions and 224 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
apps/api/src/controllers/v1/batch-scrape.ts
View File

@ -22,6 +22,7 @@ import { getJobPriority } from "../../lib/job-priority";
import { addScrapeJobs } from "../../services/queue-jobs";
import { callWebhook } from "../../services/webhook";
import { logger as _logger } from "../../lib/logger";
import { CostTracking } from "../../lib/extract/extraction-service";
export async function batchScrapeController(
req: RequestWithAuth<{}, BatchScrapeResponse, BatchScrapeRequest>,

7
apps/api/src/controllers/v1/scrape.ts
View File

@ -15,6 +15,7 @@ import { getJobPriority } from "../../lib/job-priority";
import { getScrapeQueue } from "../../services/queue-service";
import { getJob } from "./crawl-status";
import { getJobFromGCS } from "../../lib/gcs-jobs";
import { CostTracking } from "src/lib/extract/extraction-service";
export async function scrapeController(
req: RequestWithAuth<{}, ScrapeResponse, ScrapeRequest>,
@ -128,12 +129,6 @@ export async function scrapeController(
}
}
const cost_tracking = doc?.metadata?.costTracking;
if (doc && doc.metadata) {
delete doc.metadata.costTracking;
}
return res.status(200).json({
success: true,
data: doc,

11
apps/api/src/controllers/v1/search.ts
View File

@ -21,6 +21,7 @@ import { BLOCKLISTED_URL_MESSAGE } from "../../lib/strings";
import { logger as _logger } from "../../lib/logger";
import type { Logger } from "winston";
import { getJobFromGCS } from "../../lib/gcs-jobs";
import { CostTracking } from "../../lib/extract/extraction-service";
// Used for deep research
export async function searchAndScrapeSearchResult(
@ -32,6 +33,7 @@ export async function searchAndScrapeSearchResult(
scrapeOptions: ScrapeOptions;
},
logger: Logger,
costTracking: CostTracking,
): Promise<Document[]> {
try {
const searchResults = await search({
@ -48,7 +50,8 @@ export async function searchAndScrapeSearchResult(
description: result.description
},
options,
logger
logger,
costTracking
)
)
);
@ -68,6 +71,7 @@ async function scrapeSearchResult(
scrapeOptions: ScrapeOptions;
},
logger: Logger,
costTracking: CostTracking,
): Promise<Document> {
const jobId = uuidv4();
const jobPriority = await getJobPriority({
@ -220,6 +224,8 @@ export async function searchController(
});
}
const costTracking = new CostTracking();
// Scrape each non-blocked result, handling timeouts individually
logger.info("Scraping search results");
const scrapePromises = searchResults.map((result) =>
@ -228,7 +234,7 @@ export async function searchController(
origin: req.body.origin,
timeout: req.body.timeout,
scrapeOptions: req.body.scrapeOptions,
}, logger),
}, logger, costTracking),
);
const docs = await Promise.all(scrapePromises);
@ -279,6 +285,7 @@ export async function searchController(
mode: "search",
url: req.body.query,
origin: req.body.origin,
cost_tracking: costTracking,
});
return res.status(200).json({

1
apps/api/src/controllers/v1/types.ts
View File

@ -739,7 +739,6 @@ export type Document = {
statusCode: number;
scrapeId?: string;
error?: string;
costTracking?: CostTracking;
// [key: string]: string | string[] | number | { smartScrape: number; other: number; total: number } | undefined;
};
serpResults?: {

9
apps/api/src/lib/deep-research/deep-research-service.ts
View File

@ -5,6 +5,7 @@ import { ResearchLLMService, ResearchStateManager } from "./research-manager";
import { logJob } from "../../services/logging/log_job";
import { billTeam } from "../../services/billing/credit_billing";
import { ExtractOptions } from "../../controllers/v1/types";
import { CostTracking } from "../extract/extraction-service";
interface DeepResearchServiceOptions {
researchId: string;
@ -21,6 +22,7 @@ interface DeepResearchServiceOptions {
}
export async function performDeepResearch(options: DeepResearchServiceOptions) {
const costTracking = new CostTracking();
const { researchId, teamId, timeLimit, subId, maxUrls } = options;
const startTime = Date.now();
let currentTopic = options.query;
@ -70,6 +72,7 @@ export async function performDeepResearch(options: DeepResearchServiceOptions) {
await llmService.generateSearchQueries(
nextSearchTopic,
state.getFindings(),
costTracking,
)
).slice(0, 3);
@ -109,7 +112,7 @@ export async function performDeepResearch(options: DeepResearchServiceOptions) {
fastMode: false,
blockAds: false,
},
}, logger);
}, logger, costTracking);
return response.length > 0 ? response : [];
});
@ -205,6 +208,7 @@ export async function performDeepResearch(options: DeepResearchServiceOptions) {
currentTopic,
timeRemaining,
options.systemPrompt ?? "",
costTracking,
);
if (!analysis) {
@ -268,6 +272,7 @@ export async function performDeepResearch(options: DeepResearchServiceOptions) {
state.getFindings(),
state.getSummaries(),
options.analysisPrompt,
costTracking,
options.formats,
options.jsonOptions,
);
@ -278,6 +283,7 @@ export async function performDeepResearch(options: DeepResearchServiceOptions) {
state.getFindings(),
state.getSummaries(),
options.analysisPrompt,
costTracking,
);
}
@ -307,6 +313,7 @@ export async function performDeepResearch(options: DeepResearchServiceOptions) {
origin: "api",
num_tokens: 0,
tokens_billed: 0,
cost_tracking: costTracking,
});
await updateDeepResearch(researchId, {
status: "completed",

25
apps/api/src/lib/deep-research/research-manager.ts
View File

@ -12,6 +12,7 @@ import {
import { ExtractOptions } from "../../controllers/v1/types";
import { getModel } from "../generic-ai";
import { CostTracking } from "../extract/extraction-service";
interface AnalysisResult {
gaps: string[];
nextSteps: string[];
@ -152,6 +153,7 @@ export class ResearchLLMService {
async generateSearchQueries(
topic: string,
findings: DeepResearchFinding[] = [],
costTracking: CostTracking,
): Promise<{ query: string; researchGoal: string }[]> {
const { extract } = await generateCompletions({
logger: this.logger.child({
@ -194,6 +196,13 @@ export class ResearchLLMService {
The first SERP query you generate should be a very concise, simple version of the topic. `,
},
markdown: "",
costTrackingOptions: {
costTracking,
metadata: {
module: "deep-research",
method: "generateSearchQueries",
},
},
});
return extract.queries;
@ -204,6 +213,7 @@ export class ResearchLLMService {
currentTopic: string,
timeRemaining: number,
systemPrompt: string,
costTracking: CostTracking,
): Promise<AnalysisResult | null> {
try {
const timeRemainingMinutes =
@ -246,6 +256,13 @@ export class ResearchLLMService {
).text,
},
markdown: "",
costTrackingOptions: {
costTracking,
metadata: {
module: "deep-research",
method: "analyzeAndPlan",
},
},
});
return extract.analysis;
@ -260,6 +277,7 @@ export class ResearchLLMService {
findings: DeepResearchFinding[],
summaries: string[],
analysisPrompt: string,
costTracking: CostTracking,
formats?: string[],
jsonOptions?: ExtractOptions,
): Promise<any> {
@ -312,6 +330,13 @@ export class ResearchLLMService {
},
markdown: "",
model: getModel("o3-mini"),
costTrackingOptions: {
costTracking,
metadata: {
module: "deep-research",
method: "generateFinalAnalysis",
},
},
});
return extract;

20
apps/api/src/lib/extract/completions/analyzeSchemaAndPrompt.ts
View File

@ -11,25 +11,23 @@ import {
import { jsonSchema } from "ai";
import { getModel } from "../../../lib/generic-ai";
import { Logger } from "winston";
import { CostTracking } from "../extraction-service";
export async function analyzeSchemaAndPrompt(
urls: string[],
schema: any,
prompt: string,
logger: Logger,
costTracking: CostTracking,
): Promise<{
isMultiEntity: boolean;
multiEntityKeys: string[];
reasoning: string;
keyIndicators: string[];
tokenUsage: TokenUsage;
cost: number;
}> {
let cost = 0;
if (!schema) {
const genRes = await generateSchemaFromPrompt(prompt, logger);
const genRes = await generateSchemaFromPrompt(prompt, logger, costTracking);
schema = genRes.extract;
cost = genRes.cost;
}
const schemaString = JSON.stringify(schema);
@ -49,7 +47,7 @@ export async function analyzeSchemaAndPrompt(
);
try {
const { extract: result, totalUsage, cost: cost2 } = await generateCompletions({
const { extract: result, totalUsage } = await generateCompletions({
logger,
options: {
mode: "llm",
@ -59,8 +57,14 @@ export async function analyzeSchemaAndPrompt(
},
markdown: "",
model,
costTrackingOptions: {
costTracking,
metadata: {
module: "extract",
method: "analyzeSchemaAndPrompt",
},
},
});
cost += cost2;
const { isMultiEntity, multiEntityKeys, reasoning, keyIndicators } =
checkSchema.parse(result);
@ -71,7 +75,6 @@ export async function analyzeSchemaAndPrompt(
reasoning,
keyIndicators,
tokenUsage: totalUsage,
cost,
};
} catch (e) {
logger.warn("(analyzeSchemaAndPrompt) Error parsing schema analysis", {
@ -90,6 +93,5 @@ export async function analyzeSchemaAndPrompt(
totalTokens: 0,
model: model.modelId,
},
cost: 0,
};
}

20
apps/api/src/lib/extract/completions/batchExtract.ts
View File

@ -10,7 +10,7 @@ import {
buildBatchExtractSystemPrompt,
} from "../build-prompts";
import { getModel } from "../../generic-ai";
import { CostTracking } from "../extraction-service";
import fs from "fs/promises";
import { extractData } from "../../../scraper/scrapeURL/lib/extractSmartScrape";
import type { Logger } from "winston";
@ -24,6 +24,7 @@ type BatchExtractOptions = {
useAgent: boolean;
extractId?: string;
sessionId?: string;
costTracking: CostTracking;
};
/**
@ -75,6 +76,13 @@ export async function batchExtractPromise(options: BatchExtractOptions, logger:
isExtractEndpoint: true,
model: getModel("gemini-2.5-pro-preview-03-25", "vertex"),
retryModel: getModel("gemini-2.5-pro-preview-03-25", "google"),
costTrackingOptions: {
costTracking: options.costTracking,
metadata: {
module: "extract",
method: "batchExtractPromise",
},
},
};
let extractedDataArray: any[] = [];
@ -84,23 +92,15 @@ export async function batchExtractPromise(options: BatchExtractOptions, logger:
const {
extractedDataArray: e,
warning: w,
smartScrapeCost,
otherCost,
smartScrapeCallCount,
otherCallCount
} = await extractData({
extractOptions: generationOptions,
urls: [doc.metadata.sourceURL || doc.metadata.url || ""],
useAgent,
extractId,
sessionId
sessionId,
});
extractedDataArray = e;
warning = w;
smCost = smartScrapeCost;
oCost = otherCost;
smCallCount = smartScrapeCallCount;
oCallCount = otherCallCount;
} catch (error) {
logger.error("extractData failed", { error });
}

12
apps/api/src/lib/extract/completions/checkShouldExtract.ts
View File

@ -7,12 +7,14 @@ import {
buildShouldExtractUserPrompt,
} from "../build-prompts";
import { getModel } from "../../../lib/generic-ai";
import { CostTracking } from "../extraction-service";
export async function checkShouldExtract(
prompt: string,
multiEntitySchema: any,
doc: Document,
): Promise<{ tokenUsage: TokenUsage; extract: boolean; cost: number }> {
costTracking: CostTracking,
): Promise<{ tokenUsage: TokenUsage; extract: boolean; }> {
const shouldExtractCheck = await generateCompletions({
logger: logger.child({ method: "extractService/checkShouldExtract" }),
options: {
@ -32,11 +34,17 @@ export async function checkShouldExtract(
markdown: buildDocument(doc),
isExtractEndpoint: true,
model: getModel("gpt-4o-mini"),
costTrackingOptions: {
costTracking,
metadata: {
module: "extract",
method: "checkShouldExtract",
},
},
});
return {
tokenUsage: shouldExtractCheck.totalUsage,
extract: shouldExtractCheck.extract["extract"],
cost: shouldExtractCheck.cost,
};
}

30
apps/api/src/lib/extract/completions/singleAnswer.ts
View File

@ -7,6 +7,7 @@ import { buildDocument } from "../build-document";
import { Document, TokenUsage } from "../../../controllers/v1/types";
import { getModel } from "../../../lib/generic-ai";
import { extractData } from "../../../scraper/scrapeURL/lib/extractSmartScrape";
import { CostTracking } from "../extraction-service";
export async function singleAnswerCompletion({
singleAnswerDocs,
@ -17,6 +18,7 @@ export async function singleAnswerCompletion({
useAgent,
extractId,
sessionId,
costTracking,
}: {
singleAnswerDocs: Document[];
rSchema: any;
@ -26,14 +28,11 @@ export async function singleAnswerCompletion({
useAgent: boolean;
extractId: string;
sessionId: string;
costTracking: CostTracking;
}): Promise<{
extract: any;
tokenUsage: TokenUsage;
sources: string[];
smartScrapeCallCount: number;
smartScrapeCost: number;
otherCallCount: number;
otherCost: number;
}> {
const docsPrompt = `Today is: ` + new Date().toISOString() + `.\n` + prompt;
const generationOptions: GenerateCompletionsOptions = {
@ -49,13 +48,20 @@ export async function singleAnswerCompletion({
"Always prioritize using the provided content to answer the question. Do not make up an answer. Do not hallucinate. In case you can't find the information and the string is required, instead of 'N/A' or 'Not speficied', return an empty string: '', if it's not a string and you can't find the information, return null. Be concise and follow the schema always if provided.",
prompt: docsPrompt,
schema: rSchema,
},
markdown: `${singleAnswerDocs.map((x, i) => `[START_PAGE (ID: ${i})]` + buildDocument(x)).join("\n")} [END_PAGE]\n`,
isExtractEndpoint: true,
model: getModel("gemini-2.0-flash", "google"),
costTrackingOptions: {
costTracking,
metadata: {
module: "extract",
method: "singleAnswerCompletion",
},
markdown: `${singleAnswerDocs.map((x, i) => `[START_PAGE (ID: ${i})]` + buildDocument(x)).join("\n")} [END_PAGE]\n`,
isExtractEndpoint: true,
model: getModel("gemini-2.0-flash", "google"),
};
const { extractedDataArray, warning, smartScrapeCost, otherCost, smartScrapeCallCount, otherCallCount } = await extractData({
},
};
const { extractedDataArray, warning } = await extractData({
extractOptions: generationOptions,
urls: singleAnswerDocs.map(doc => doc.metadata.url || doc.metadata.sourceURL || ""),
useAgent,
@ -100,9 +106,5 @@ export async function singleAnswerCompletion({
sources: singleAnswerDocs.map(
(doc) => doc.metadata.url || doc.metadata.sourceURL || "",
),
smartScrapeCost,
otherCost,
smartScrapeCallCount,
otherCallCount,
};
}

88
apps/api/src/lib/extract/extraction-service.ts
View File

@ -67,14 +67,39 @@ type completions = {
sources?: string[];
};
export type CostTracking = {
smartScrapeCallCount: number;
smartScrapeCost: number;
otherCallCount: number;
otherCost: number;
totalCost: number;
costLimitExceededTokenUsage?: number;
};
export class CostTracking {
calls: {
type: "smartScrape" | "other",
metadata: Record<string, any>,
cost: number,
tokens?: {
input: number,
output: number,
},
stack: string,
}[] = [];
constructor() {}
public addCall(call: Omit<typeof this.calls[number], "stack">) {
this.calls.push({
...call,
stack: new Error().stack!.split("\n").slice(2).join("\n"),
});
}
public toJSON() {
return {
calls: this.calls,
smartScrapeCallCount: this.calls.filter(c => c.type === "smartScrape").length,
smartScrapeCost: this.calls.filter(c => c.type === "smartScrape").reduce((acc, c) => acc + c.cost, 0),
otherCallCount: this.calls.filter(c => c.type === "other").length,
otherCost: this.calls.filter(c => c.type === "other").reduce((acc, c) => acc + c.cost, 0),
totalCost: this.calls.reduce((acc, c) => acc + c.cost, 0),
}
}
}
export async function performExtraction(
extractId: string,
@ -89,13 +114,7 @@ export async function performExtraction(
let singleAnswerResult: any = {};
let totalUrlsScraped = 0;
let sources: Record<string, string[]> = {};
let costTracking: CostTracking = {
smartScrapeCallCount: 0,
smartScrapeCost: 0,
otherCallCount: 0,
otherCost: 0,
totalCost: 0,
};
let costTracking: CostTracking = new CostTracking();
let log = {
extractId,
@ -118,13 +137,9 @@ export async function performExtraction(
});
const rephrasedPrompt = await generateBasicCompletion(
buildRephraseToSerpPrompt(request.prompt),
costTracking,
);
let rptxt = rephrasedPrompt?.text.replace('"', "").replace("'", "") || "";
if (rephrasedPrompt) {
costTracking.otherCallCount++;
costTracking.otherCost += rephrasedPrompt.cost;
costTracking.totalCost += rephrasedPrompt.cost;
}
const searchResults = await search({
query: rptxt,
num_results: 10,
@ -197,11 +212,9 @@ export async function performExtraction(
let reqSchema = request.schema;
if (!reqSchema && request.prompt) {
const schemaGenRes = await generateSchemaFromPrompt(request.prompt, logger);
const schemaGenRes = await generateSchemaFromPrompt(request.prompt, logger, costTracking);
reqSchema = schemaGenRes.extract;
costTracking.otherCallCount++;
costTracking.otherCost += schemaGenRes.cost;
costTracking.totalCost += schemaGenRes.cost;
logger.debug("Generated request schema.", {
originalSchema: request.schema,
@ -232,8 +245,7 @@ export async function performExtraction(
reasoning,
keyIndicators,
tokenUsage: schemaAnalysisTokenUsage,
cost: schemaAnalysisCost,
} = await analyzeSchemaAndPrompt(urls, reqSchema, request.prompt ?? "", logger);
} = await analyzeSchemaAndPrompt(urls, reqSchema, request.prompt ?? "", logger, costTracking);
logger.debug("Analyzed schema.", {
isMultiEntity,
@ -242,11 +254,6 @@ export async function performExtraction(
keyIndicators,
});
costTracking.otherCallCount++;
costTracking.otherCost += schemaAnalysisCost;
costTracking.totalCost += schemaAnalysisCost;
// Track schema analysis tokens
tokenUsage.push(schemaAnalysisTokenUsage);
let startMap = Date.now();
@ -467,7 +474,8 @@ export async function performExtraction(
doc,
useAgent: isAgentExtractModelValid(request.agent?.model),
extractId,
sessionId
sessionId,
costTracking,
}, logger);
// Race between timeout and completion
@ -481,12 +489,6 @@ export async function performExtraction(
if (multiEntityCompletion) {
tokenUsage.push(multiEntityCompletion.totalUsage);
costTracking.smartScrapeCallCount += multiEntityCompletion.smartScrapeCallCount;
costTracking.smartScrapeCost += multiEntityCompletion.smartScrapeCost;
costTracking.otherCallCount += multiEntityCompletion.otherCallCount;
costTracking.otherCost += multiEntityCompletion.otherCost;
costTracking.totalCost += multiEntityCompletion.smartScrapeCost + multiEntityCompletion.otherCost;
if (multiEntityCompletion.extract) {
return {
extract: multiEntityCompletion.extract,
@ -776,10 +778,6 @@ export async function performExtraction(
extract: completionResult,
tokenUsage: singleAnswerTokenUsage,
sources: singleAnswerSources,
smartScrapeCost: singleAnswerSmartScrapeCost,
otherCost: singleAnswerOtherCost,
smartScrapeCallCount: singleAnswerSmartScrapeCallCount,
otherCallCount: singleAnswerOtherCallCount,
} = await singleAnswerCompletion({
singleAnswerDocs,
rSchema,
@ -789,12 +787,8 @@ export async function performExtraction(
useAgent: isAgentExtractModelValid(request.agent?.model),
extractId,
sessionId: thisSessionId,
costTracking,
});
costTracking.smartScrapeCost += singleAnswerSmartScrapeCost;
costTracking.smartScrapeCallCount += singleAnswerSmartScrapeCallCount;
costTracking.otherCost += singleAnswerOtherCost;
costTracking.otherCallCount += singleAnswerOtherCallCount;
costTracking.totalCost += singleAnswerSmartScrapeCost + singleAnswerOtherCost;
logger.debug("Done generating singleAnswer completions.");
singleAnswerResult = transformArrayToObject(rSchema, completionResult);

13
apps/api/src/lib/extract/fire-0/reranker-f0.ts
View File

@ -6,7 +6,7 @@ import { extractConfig } from "../config";
import { generateCompletions } from "../../../scraper/scrapeURL/transformers/llmExtract";
import { performRanking_F0 } from "./ranker-f0";
import { buildRerankerSystemPrompt_F0, buildRerankerUserPrompt_F0 } from "./build-prompts-f0";
import { CostTracking } from "../extraction-service";
const cohere = new CohereClient({
token: process.env.COHERE_API_KEY,
});
@ -166,7 +166,7 @@ export type RerankerOptions = {
urlTraces: URLTrace[];
};
export async function rerankLinksWithLLM_F0(options: RerankerOptions): Promise<RerankerResult> {
export async function rerankLinksWithLLM_F0(options: RerankerOptions, costTracking: CostTracking): Promise<RerankerResult> {
const { links, searchQuery, urlTraces } = options;
const chunkSize = 100;
const chunks: MapDocument[][] = [];
@ -231,7 +231,14 @@ export async function rerankLinksWithLLM_F0(options: RerankerOptions): Promise<R
schema: schema,
},
markdown: linksContent,
isExtractEndpoint: true
isExtractEndpoint: true,
costTrackingOptions: {
costTracking: new CostTracking(),
metadata: {
module: "extract",
method: "rerankLinksWithLLM",
},
},
});
const completion = await Promise.race([

5
apps/api/src/lib/extract/fire-0/url-processor-f0.ts
View File

@ -8,6 +8,7 @@ import { extractConfig } from "../config";
import type { Logger } from "winston";
import { generateText } from "ai";
import { getModel } from "../../generic-ai";
import { CostTracking } from "../extraction-service";
export async function generateBasicCompletion_FO(prompt: string) {
const { text } = await generateText({
@ -211,7 +212,7 @@ export async function processUrl_F0(
links: mappedLinks,
searchQuery: rephrasedPrompt,
urlTraces,
});
}, new CostTracking());
mappedLinks = rerankerResult.mapDocument;
let tokensUsed = rerankerResult.tokensUsed;
logger.info("Reranked! (pass 1)", {
@ -225,7 +226,7 @@ export async function processUrl_F0(
links: mappedLinks,
searchQuery: rephrasedPrompt,
urlTraces,
});
}, new CostTracking());
mappedLinks = rerankerResult.mapDocument;
tokensUsed += rerankerResult.tokensUsed;
logger.info("Reranked! (pass 2)", {

9
apps/api/src/lib/extract/reranker.ts
View File

@ -11,6 +11,7 @@ import { buildRerankerSystemPrompt } from "./build-prompts";
import { dumpToFile } from "./helpers/dump-to-file";
import { getModel } from "../generic-ai";
import fs from "fs/promises";
import { CostTracking } from "./extraction-service";
const THRESHOLD_FOR_SINGLEPAGE = 0.6;
const THRESHOLD_FOR_MULTIENTITY = 0.45;
@ -177,6 +178,7 @@ export type RerankerOptions = {
reasoning: string;
multiEntityKeys: string[];
keyIndicators: string[];
costTracking: CostTracking;
};
export async function rerankLinksWithLLM(
@ -315,6 +317,13 @@ export async function rerankLinksWithLLM(
// },
markdown: linksContent,
isExtractEndpoint: true,
costTrackingOptions: {
costTracking: options.costTracking,
metadata: {
module: "extract",
method: "rerankLinksWithLLM",
},
},
});
completion = await completionPromise;

46
apps/api/src/lib/extract/url-processor.ts
View File

@ -11,7 +11,7 @@ import { getModel } from "../generic-ai";
import { calculateCost } from "../../scraper/scrapeURL/transformers/llmExtract";
import type { CostTracking } from "./extraction-service";
export async function generateBasicCompletion(prompt: string): Promise<{ text: string, cost: number } | null> {
export async function generateBasicCompletion(prompt: string, costTracking: CostTracking): Promise<{ text: string } | null> {
try {
const result = await generateText({
model: getModel("gpt-4o", "openai"),
@ -22,7 +22,19 @@ export async function generateBasicCompletion(prompt: string): Promise<{ text: s
},
}
});
return { text: result.text, cost: calculateCost("openai/gpt-4o", result.usage?.promptTokens ?? 0, result.usage?.completionTokens ?? 0) };
costTracking.addCall({
type: "other",
metadata: {
module: "extract",
method: "generateBasicCompletion",
},
cost: calculateCost("openai/gpt-4o", result.usage?.promptTokens ?? 0, result.usage?.completionTokens ?? 0),
tokens: {
input: result.usage?.promptTokens ?? 0,
output: result.usage?.completionTokens ?? 0,
},
});
return { text: result.text };
} catch (error) {
console.error("Error generating basic completion:", error);
if (error?.type == "rate_limit_error") {
@ -36,7 +48,19 @@ export async function generateBasicCompletion(prompt: string): Promise<{ text: s
},
}
});
return { text: result.text, cost: calculateCost("openai/gpt-4o-mini", result.usage?.promptTokens ?? 0, result.usage?.completionTokens ?? 0) };
costTracking.addCall({
type: "other",
metadata: {
module: "extract",
method: "generateBasicCompletion",
},
cost: calculateCost("openai/gpt-4o-mini", result.usage?.promptTokens ?? 0, result.usage?.completionTokens ?? 0),
tokens: {
input: result.usage?.promptTokens ?? 0,
output: result.usage?.completionTokens ?? 0,
},
});
return { text: result.text };
} catch (fallbackError) {
console.error("Error generating basic completion with fallback model:", fallbackError);
return null;
@ -96,13 +120,11 @@ export async function processUrl(
if (options.prompt) {
const res = await generateBasicCompletion(
buildRefrasedPrompt(options.prompt, baseUrl),
costTracking,
);
if (res) {
searchQuery = res.text.replace('"', "").replace("/", "") ?? options.prompt;
costTracking.otherCallCount++;
costTracking.otherCost += res.cost;
costTracking.totalCost += res.cost;
}
}
@ -223,13 +245,11 @@ export async function processUrl(
try {
const res = await generateBasicCompletion(
buildPreRerankPrompt(rephrasedPrompt, options.schema, baseUrl),
costTracking,
);
if (res) {
rephrasedPrompt = res.text;
costTracking.otherCallCount++;
costTracking.otherCost += res.cost;
costTracking.totalCost += res.cost;
} else {
rephrasedPrompt =
"Extract the data according to the schema: " +
@ -262,10 +282,8 @@ export async function processUrl(
reasoning: options.reasoning,
multiEntityKeys: options.multiEntityKeys,
keyIndicators: options.keyIndicators,
costTracking,
});
costTracking.otherCallCount++;
costTracking.otherCost += rerankerResult.cost;
costTracking.totalCost += rerankerResult.cost;
mappedLinks = rerankerResult.mapDocument;
let tokensUsed = rerankerResult.tokensUsed;
logger.info("Reranked! (pass 1)", {
@ -283,10 +301,8 @@ export async function processUrl(
reasoning: options.reasoning,
multiEntityKeys: options.multiEntityKeys,
keyIndicators: options.keyIndicators,
costTracking,
});
costTracking.otherCallCount++;
costTracking.otherCost += rerankerResult.cost;
costTracking.totalCost += rerankerResult.cost;
mappedLinks = rerankerResult.mapDocument;
tokensUsed += rerankerResult.tokensUsed;
logger.info("Reranked! (pass 2)", {

11
apps/api/src/lib/generate-llmstxt/generate-llmstxt-service.ts
View File

@ -11,7 +11,7 @@ import { billTeam } from "../../services/billing/credit_billing";
import { logJob } from "../../services/logging/log_job";
import { getModel } from "../generic-ai";
import { generateCompletions } from "../../scraper/scrapeURL/transformers/llmExtract";
import { CostTracking } from "../extract/extraction-service";
interface GenerateLLMsTextServiceOptions {
generationId: string;
teamId: string;
@ -71,6 +71,7 @@ export async function performGenerateLlmsTxt(
generationId,
teamId,
});
const costTracking = new CostTracking();
try {
// Enforce max URL limit
@ -167,6 +168,13 @@ export async function performGenerateLlmsTxt(
prompt: `Generate a 9-10 word description and a 3-4 word title of the entire page based on ALL the content one will find on the page for this url: ${document.metadata?.url}. This will help in a user finding the page for its intended purpose.`,
},
markdown: document.markdown,
costTrackingOptions: {
costTracking,
metadata: {
module: "generate-llmstxt",
method: "generateDescription",
},
},
});
return {
@ -229,6 +237,7 @@ export async function performGenerateLlmsTxt(
num_tokens: 0,
tokens_billed: 0,
sources: {},
cost_tracking: costTracking,
});
// Bill team for usage

7
apps/api/src/main/runWebScraper.ts
View File

@ -17,14 +17,17 @@ import {
} from "../scraper/scrapeURL";
import { Engine } from "../scraper/scrapeURL/engines";
import { indexPage } from "../lib/extract/index/pinecone";
import { CostTracking } from "../lib/extract/extraction-service";
configDotenv();
export async function startWebScraperPipeline({
job,
token,
costTracking,
}: {
job: Job<WebScraperOptions> & { id: string };
token: string;
costTracking: CostTracking;
}) {
return await runWebScraper({
url: job.data.url,
@ -52,6 +55,7 @@ export async function startWebScraperPipeline({
is_scrape: job.data.is_scrape ?? false,
is_crawl: !!(job.data.crawl_id && job.data.crawlerOptions !== null),
urlInvisibleInCurrentCrawl: job.data.crawlerOptions?.urlInvisibleInCurrentCrawl ?? false,
costTracking,
});
}
@ -68,6 +72,7 @@ export async function runWebScraper({
is_scrape = false,
is_crawl = false,
urlInvisibleInCurrentCrawl = false,
costTracking,
}: RunWebScraperParams): Promise<ScrapeUrlResponse> {
const logger = _logger.child({
method: "runWebScraper",
@ -101,7 +106,7 @@ export async function runWebScraper({
...internalOptions,
urlInvisibleInCurrentCrawl,
teamId: internalOptions?.teamId ?? team_id,
});
}, costTracking);
if (!response.success) {
if (response.error instanceof Error) {
throw response.error;

2
apps/api/src/scraper/WebScraper/sitemap.ts
View File

@ -3,6 +3,7 @@ import { WebCrawler } from "./crawler";
import { scrapeURL } from "../scrapeURL";
import { scrapeOptions, TimeoutSignal } from "../../controllers/v1/types";
import type { Logger } from "winston";
import { CostTracking } from "../../lib/extract/extraction-service";
const useFireEngine =
process.env.FIRE_ENGINE_BETA_URL !== "" &&
process.env.FIRE_ENGINE_BETA_URL !== undefined;
@ -49,6 +50,7 @@ export async function getLinksFromSitemap(
abort,
teamId: "sitemap",
},
new CostTracking(),
);
if (

9
apps/api/src/scraper/scrapeURL/index.ts
View File

@ -26,6 +26,7 @@ import { executeTransformers } from "./transformers";
import { LLMRefusalError } from "./transformers/llmExtract";
import { urlSpecificParams } from "./lib/urlSpecificParams";
import { loadMock, MockState } from "./lib/mock";
import { CostTracking } from "../../lib/extract/extraction-service";
export type ScrapeUrlResponse = (
| {
@ -55,6 +56,7 @@ export type Meta = {
url?: string;
status: number;
} | null | undefined; // undefined: no prefetch yet, null: prefetch came back empty
costTracking: CostTracking;
};
function buildFeatureFlags(
@ -127,6 +129,7 @@ async function buildMetaObject(
url: string,
options: ScrapeOptions,
internalOptions: InternalOptions,
costTracking: CostTracking,
): Promise<Meta> {
const specParams =
urlSpecificParams[new URL(url).hostname.replace(/^www\./, "")];
@ -158,12 +161,13 @@ async function buildMetaObject(
? await loadMock(options.useMock, _logger)
: null,
pdfPrefetch: undefined,
costTracking,
};
}
export type InternalOptions = {
teamId: string;
priority?: number; // Passed along to fire-engine
forceEngine?: Engine | Engine[];
atsv?: boolean; // anti-bot solver, beta
@ -389,8 +393,9 @@ export async function scrapeURL(
url: string,
options: ScrapeOptions,
internalOptions: InternalOptions,
costTracking: CostTracking,
): Promise<ScrapeUrlResponse> {
const meta = await buildMetaObject(id, url, options, internalOptions);
const meta = await buildMetaObject(id, url, options, internalOptions, costTracking);
try {
while (true) {
try {

51
apps/api/src/scraper/scrapeURL/lib/extractSmartScrape.ts
View File

@ -10,8 +10,7 @@ import { parseMarkdown } from "../../../lib/html-to-markdown";
import { getModel } from "../../../lib/generic-ai";
import { TokenUsage } from "../../../controllers/v1/types";
import type { SmartScrapeResult } from "./smartScrape";
import { ExtractStep } from "src/lib/extract/extract-redis";
import { CostTracking } from "../../../lib/extract/extraction-service";
const commonSmartScrapeProperties = {
shouldUseSmartscrape: {
type: "boolean",
@ -225,26 +224,16 @@ export async function extractData({
}): Promise<{
extractedDataArray: any[];
warning: any;
smartScrapeCallCount: number;
otherCallCount: number;
smartScrapeCost: number;
otherCost: number;
costLimitExceededTokenUsage: number | null;
}> {
let schema = extractOptions.options.schema;
const logger = extractOptions.logger;
const isSingleUrl = urls.length === 1;
let smartScrapeCost = 0;
let otherCost = 0;
let smartScrapeCallCount = 0;
let otherCallCount = 0;
let costLimitExceededTokenUsage: number | null = null;
// TODO: remove the "required" fields here!! it breaks o3-mini
if (!schema && extractOptions.options.prompt) {
const genRes = await generateSchemaFromPrompt(extractOptions.options.prompt, logger);
otherCallCount++;
otherCost += genRes.cost;
const genRes = await generateSchemaFromPrompt(extractOptions.options.prompt, logger, extractOptions.costTrackingOptions.costTracking);
schema = genRes.extract;
}
@ -278,17 +267,22 @@ export async function extractData({
extract: e,
warning: w,
totalUsage: t,
cost: c,
} = await generateCompletions({
...extractOptionsNewSchema,
model: getModel("gemini-2.5-pro-preview-03-25", "vertex"),
retryModel: getModel("gemini-2.5-pro-preview-03-25", "google"),
costTrackingOptions: {
costTracking: extractOptions.costTrackingOptions.costTracking,
metadata: {
module: "scrapeURL",
method: "extractData",
description: "Check if using smartScrape is needed for this case"
},
},
});
extract = e;
warning = w;
totalUsage = t;
otherCost += c;
otherCallCount++;
} catch (error) {
logger.error(
"failed during extractSmartScrape.ts:generateCompletions",
@ -321,10 +315,9 @@ export async function extractData({
sessionId,
extractId,
scrapeId,
costTracking: extractOptions.costTrackingOptions.costTracking,
}),
];
smartScrapeCost += smartscrapeResults[0].tokenUsage;
smartScrapeCallCount++;
} else {
const pages = extract?.smartscrapePages ?? [];
//do it async promiseall instead
@ -344,14 +337,10 @@ export async function extractData({
sessionId,
extractId,
scrapeId,
costTracking: extractOptions.costTrackingOptions.costTracking,
});
}),
);
smartScrapeCost += smartscrapeResults.reduce(
(acc, result) => acc + result.tokenUsage,
0,
);
smartScrapeCallCount += smartscrapeResults.length;
}
// console.log("smartscrapeResults", smartscrapeResults);
@ -372,11 +361,17 @@ export async function extractData({
markdown: markdown,
model: getModel("gemini-2.5-pro-preview-03-25", "vertex"),
retryModel: getModel("gemini-2.5-pro-preview-03-25", "google"),
costTrackingOptions: {
costTracking: extractOptions.costTrackingOptions.costTracking,
metadata: {
module: "scrapeURL",
method: "extractData",
description: "Extract data from markdown (smart-scape results)",
},
},
};
const { extract, warning, totalUsage, model, cost } =
const { extract } =
await generateCompletions(newExtractOptions);
otherCost += cost;
otherCallCount++;
return extract;
}),
);
@ -399,10 +394,6 @@ export async function extractData({
return {
extractedDataArray: extractedData,
warning: warning,
smartScrapeCallCount: smartScrapeCallCount,
otherCallCount: otherCallCount,
smartScrapeCost: smartScrapeCost,
otherCost: otherCost,
costLimitExceededTokenUsage: costLimitExceededTokenUsage,
};
}

14
apps/api/src/scraper/scrapeURL/lib/smartScrape.ts
View File

@ -3,7 +3,7 @@ import { logger as _logger } from "../../../lib/logger";
import { robustFetch } from "./fetch";
import fs from "fs/promises";
import { configDotenv } from "dotenv";
import { CostTracking } from "../../../lib/extract/extraction-service";
configDotenv();
// Define schemas outside the function scope
@ -52,6 +52,7 @@ export async function smartScrape({
extractId,
scrapeId,
beforeSubmission,
costTracking,
}: {
url: string,
prompt: string,
@ -59,6 +60,7 @@ export async function smartScrape({
extractId?: string,
scrapeId?: string,
beforeSubmission?: () => unknown,
costTracking: CostTracking,
}): Promise<SmartScrapeResult> {
let logger = _logger.child({
method: "smartScrape",
@ -139,6 +141,16 @@ export async function smartScrape({
});
logger.info("Smart scrape cost $" + response.tokenUsage);
costTracking.addCall({
type: "smartScrape",
cost: response.tokenUsage,
metadata: {
module: "smartScrape",
method: "smartScrape",
url,
sessionId,
},
});
return response; // The response type now matches SmartScrapeResult
} catch (error) {

20
apps/api/src/scraper/scrapeURL/scrapeURL.test.ts
View File

@ -5,6 +5,7 @@ process.env.ENV = "test";
import { scrapeURL } from ".";
import { scrapeOptions } from "../../controllers/v1/types";
import { Engine } from "./engines";
import { CostTracking } from "../../lib/extract/extraction-service";
const testEngines: (Engine | undefined)[] = [
undefined,
@ -32,6 +33,7 @@ describe("Standalone scrapeURL tests", () => {
"https://www.roastmywebsite.ai/",
scrapeOptions.parse({}),
{ forceEngine, teamId: "test" },
new CostTracking(),
);
// expect(out.logs.length).toBeGreaterThan(0);
@ -79,6 +81,7 @@ describe("Standalone scrapeURL tests", () => {
formats: ["markdown", "html"],
}),
{ forceEngine, teamId: "test" },
new CostTracking(),
);
// expect(out.logs.length).toBeGreaterThan(0);
@ -103,6 +106,7 @@ describe("Standalone scrapeURL tests", () => {
onlyMainContent: false,
}),
{ forceEngine, teamId: "test" },
new CostTracking(),
);
// expect(out.logs.length).toBeGreaterThan(0);
@ -126,6 +130,7 @@ describe("Standalone scrapeURL tests", () => {
excludeTags: [".nav", "#footer", "strong"],
}),
{ forceEngine, teamId: "test" },
new CostTracking(),
);
// expect(out.logs.length).toBeGreaterThan(0);
@ -146,6 +151,7 @@ describe("Standalone scrapeURL tests", () => {
"https://httpstat.us/400",
scrapeOptions.parse({}),
{ forceEngine, teamId: "test" },
new CostTracking(),
);
// expect(out.logs.length).toBeGreaterThan(0);
@ -164,6 +170,7 @@ describe("Standalone scrapeURL tests", () => {
"https://httpstat.us/401",
scrapeOptions.parse({}),
{ forceEngine, teamId: "test" },
new CostTracking(),
);
// expect(out.logs.length).toBeGreaterThan(0);
@ -182,6 +189,7 @@ describe("Standalone scrapeURL tests", () => {
"https://httpstat.us/403",
scrapeOptions.parse({}),
{ forceEngine, teamId: "test" },
new CostTracking(),
);
// expect(out.logs.length).toBeGreaterThan(0);
@ -200,6 +208,7 @@ describe("Standalone scrapeURL tests", () => {
"https://httpstat.us/404",
scrapeOptions.parse({}),
{ forceEngine, teamId: "test" },
new CostTracking(),
);
// expect(out.logs.length).toBeGreaterThan(0);
@ -218,6 +227,7 @@ describe("Standalone scrapeURL tests", () => {
"https://httpstat.us/405",
scrapeOptions.parse({}),
{ forceEngine, teamId: "test" },
new CostTracking(),
);
// expect(out.logs.length).toBeGreaterThan(0);
@ -236,6 +246,7 @@ describe("Standalone scrapeURL tests", () => {
"https://httpstat.us/500",
scrapeOptions.parse({}),
{ forceEngine, teamId: "test" },
new CostTracking(),
);
// expect(out.logs.length).toBeGreaterThan(0);
@ -254,6 +265,7 @@ describe("Standalone scrapeURL tests", () => {
"https://scrapethissite.com/",
scrapeOptions.parse({}),
{ forceEngine, teamId: "test" },
new CostTracking(),
);
// expect(out.logs.length).toBeGreaterThan(0);
@ -286,6 +298,7 @@ describe("Standalone scrapeURL tests", () => {
formats: ["screenshot"],
}),
{ forceEngine, teamId: "test" },
new CostTracking(),
);
// expect(out.logs.length).toBeGreaterThan(0);
@ -314,6 +327,7 @@ describe("Standalone scrapeURL tests", () => {
formats: ["screenshot@fullPage"],
}),
{ forceEngine, teamId: "test" },
new CostTracking(),
);
// expect(out.logs.length).toBeGreaterThan(0);
@ -342,6 +356,7 @@ describe("Standalone scrapeURL tests", () => {
"https://arxiv.org/pdf/astro-ph/9301001.pdf",
scrapeOptions.parse({}),
{ teamId: "test" },
new CostTracking(),
);
// expect(out.logs.length).toBeGreaterThan(0);
@ -361,6 +376,7 @@ describe("Standalone scrapeURL tests", () => {
"https://nvca.org/wp-content/uploads/2019/06/NVCA-Model-Document-Stock-Purchase-Agreement.docx",
scrapeOptions.parse({}),
{ teamId: "test" },
new CostTracking(),
);
// expect(out.logs.length).toBeGreaterThan(0);
@ -398,6 +414,7 @@ describe("Standalone scrapeURL tests", () => {
},
}),
{ teamId: "test" },
new CostTracking(),
);
// expect(out.logs.length).toBeGreaterThan(0);
@ -434,6 +451,7 @@ describe("Standalone scrapeURL tests", () => {
},
}),
{ teamId: "test" },
new CostTracking(),
);
// expect(out.logs.length).toBeGreaterThan(0);
@ -455,7 +473,7 @@ describe("Standalone scrapeURL tests", () => {
async (i) => {
const url = "https://www.scrapethissite.com/?i=" + i;
const id = "test:concurrent:" + url;
const out = await scrapeURL(id, url, scrapeOptions.parse({}), { teamId: "test" });
const out = await scrapeURL(id, url, scrapeOptions.parse({}), { teamId: "test" }, new CostTracking());
const replacer = (key: string, value: any) => {
if (value instanceof Error) {

15
apps/api/src/scraper/scrapeURL/transformers/agent.ts
View File

@ -30,6 +30,7 @@ export async function performAgent(
prompt,
sessionId,
scrapeId: meta.id,
costTracking: meta.costTracking,
})
} catch (error) {
if (error instanceof Error && error.message === "Cost limit exceeded") {
@ -50,20 +51,6 @@ export async function performAgent(
if (meta.options.formats.includes("html")) {
document.html = html
}
if (document.metadata.costTracking) {
document.metadata.costTracking.smartScrapeCallCount++;
document.metadata.costTracking.smartScrapeCost = document.metadata.costTracking.smartScrapeCost + smartscrapeResults.tokenUsage;
document.metadata.costTracking.totalCost = document.metadata.costTracking.totalCost + smartscrapeResults.tokenUsage;
} else {
document.metadata.costTracking = {
smartScrapeCallCount: 1,
smartScrapeCost: smartscrapeResults.tokenUsage,
otherCallCount: 0,
otherCost: 0,
totalCost: smartscrapeResults.tokenUsage,
}
}
}
return document;

39
apps/api/src/scraper/scrapeURL/transformers/diff.ts
View File

@ -6,9 +6,9 @@ import gitDiff from 'git-diff';
import parseDiff from 'parse-diff';
import { generateCompletions } from "./llmExtract";
async function extractDataWithSchema(content: string, meta: Meta): Promise<{ extract: any, cost: number } | null> {
async function extractDataWithSchema(content: string, meta: Meta): Promise<{ extract: any } | null> {
try {
const { extract, cost } = await generateCompletions({
const { extract } = await generateCompletions({
logger: meta.logger.child({
method: "extractDataWithSchema/generateCompletions",
}),
@ -18,9 +18,16 @@ async function extractDataWithSchema(content: string, meta: Meta): Promise<{ ext
systemPrompt: "Extract the requested information from the content based on the provided schema.",
temperature: 0
},
markdown: content
markdown: content,
costTrackingOptions: {
costTracking: meta.costTracking,
metadata: {
module: "extract",
method: "extractDataWithSchema",
},
},
});
return { extract, cost };
return { extract };
} catch (error) {
meta.logger.error("Error extracting data with schema", { error });
return null;
@ -145,19 +152,6 @@ export async function deriveDiff(meta: Meta, document: Document): Promise<Docume
if (previousData && currentData) {
document.changeTracking.json = compareExtractedData(previousData.extract, currentData.extract);
if (document.metadata.costTracking) {
document.metadata.costTracking.otherCallCount += 2;
document.metadata.costTracking.otherCost = document.metadata.costTracking.otherCost + previousData.cost + currentData.cost;
} else {
document.metadata.costTracking = {
smartScrapeCallCount: 0,
smartScrapeCost: 0,
otherCallCount: 2,
otherCost: previousData.cost + currentData.cost,
totalCost: previousData.cost + currentData.cost
}
}
} else {
const { extract } = await generateCompletions({
logger: meta.logger.child({
@ -171,9 +165,16 @@ export async function deriveDiff(meta: Meta, document: Document): Promise<Docume
temperature: 0
},
markdown: `Previous Content:\n${previousMarkdown}\n\nCurrent Content:\n${currentMarkdown}`,
previousWarning: document.warning
previousWarning: document.warning,
costTrackingOptions: {
costTracking: meta.costTracking,
metadata: {
module: "diff",
method: "deriveDiff",
},
},
});
document.changeTracking.json = extract;
}
} catch (error) {

142
apps/api/src/scraper/scrapeURL/transformers/llmExtract.ts
View File

@ -11,6 +11,7 @@ import { EngineResultsTracker, Meta } from "..";
import { logger } from "../../../lib/logger";
import { modelPrices } from "../../../lib/extract/usage/model-prices";
import {
AISDKError,
generateObject,
generateText,
LanguageModel,
@ -22,7 +23,7 @@ import { z } from "zod";
import fs from "fs/promises";
import Ajv from "ajv";
import { extractData } from "../lib/extractSmartScrape";
import { CostTracking } from "../../../lib/extract/extraction-service";
// TODO: fix this, it's horrible
type LanguageModelV1ProviderMetadata = {
anthropic?: {
@ -231,6 +232,10 @@ export type GenerateCompletionsOptions = {
mode?: "object" | "no-object";
providerOptions?: LanguageModelV1ProviderMetadata;
retryModel?: LanguageModel;
costTrackingOptions: {
costTracking: CostTracking;
metadata: Record<string, any>;
};
};
export async function generateCompletions({
logger,
@ -242,13 +247,13 @@ export async function generateCompletions({
mode = "object",
providerOptions,
retryModel = getModel("claude-3-5-sonnet-20240620", "anthropic"),
costTrackingOptions,
}: GenerateCompletionsOptions): Promise<{
extract: any;
numTokens: number;
warning: string | undefined;
totalUsage: TokenUsage;
model: string;
cost: number;
}> {
let extract: any;
let warning: string | undefined;
@ -278,6 +283,19 @@ export async function generateCompletions({
},
});
costTrackingOptions.costTracking.addCall({
type: "other",
metadata: {
...costTrackingOptions.metadata,
gcDetails: "no-object",
},
cost: calculateCost(
currentModel.modelId,
result.usage?.promptTokens ?? 0,
result.usage?.completionTokens ?? 0,
),
});
extract = result.text;
return {
@ -290,11 +308,6 @@ export async function generateCompletions({
totalTokens: result.usage?.promptTokens ?? 0 + (result.usage?.completionTokens ?? 0),
},
model: currentModel.modelId,
cost: calculateCost(
currentModel.modelId,
result.usage?.promptTokens ?? 0,
result.usage?.completionTokens ?? 0,
),
};
} catch (error) {
lastError = error as Error;
@ -321,6 +334,19 @@ export async function generateCompletions({
extract = result.text;
costTrackingOptions.costTracking.addCall({
type: "other",
metadata: {
...costTrackingOptions.metadata,
gcDetails: "no-object fallback",
},
cost: calculateCost(
currentModel.modelId,
result.usage?.promptTokens ?? 0,
result.usage?.completionTokens ?? 0,
),
});
return {
extract,
warning,
@ -331,11 +357,6 @@ export async function generateCompletions({
totalTokens: result.usage?.promptTokens ?? 0 + (result.usage?.completionTokens ?? 0),
},
model: currentModel.modelId,
cost: calculateCost(
currentModel.modelId,
result.usage?.promptTokens ?? 0,
result.usage?.completionTokens ?? 0,
),
};
} catch (retryError) {
lastError = retryError as Error;
@ -410,7 +431,7 @@ export async function generateCompletions({
}
try {
const { text: fixedText } = await generateText({
const { text: fixedText, usage: repairUsage } = await generateText({
model: currentModel,
prompt: `Fix this JSON that had the following error: ${error}\n\nOriginal text:\n${text}\n\nReturn only the fixed JSON, no explanation.`,
system:
@ -421,6 +442,23 @@ export async function generateCompletions({
},
},
});
costTrackingOptions.costTracking.addCall({
type: "other",
metadata: {
...costTrackingOptions.metadata,
gcDetails: "repairConfig",
},
cost: calculateCost(
currentModel.modelId,
repairUsage?.promptTokens ?? 0,
repairUsage?.completionTokens ?? 0,
),
tokens: {
input: repairUsage?.promptTokens ?? 0,
output: repairUsage?.completionTokens ?? 0,
},
});
logger.debug("Repaired text with LLM");
return fixedText;
} catch (repairError) {
@ -464,6 +502,23 @@ export async function generateCompletions({
let result: { object: any; usage: TokenUsage } | undefined;
try {
result = await generateObject(generateObjectConfig);
costTrackingOptions.costTracking.addCall({
type: "other",
metadata: {
...costTrackingOptions.metadata,
gcDetails: "generateObject",
gcModel: generateObjectConfig.model.modelId,
},
tokens: {
input: result.usage?.promptTokens ?? 0,
output: result.usage?.completionTokens ?? 0,
},
cost: calculateCost(
currentModel.modelId,
result.usage?.promptTokens ?? 0,
result.usage?.completionTokens ?? 0,
),
});
} catch (error) {
lastError = error as Error;
if (
@ -481,6 +536,23 @@ export async function generateCompletions({
model: currentModel,
};
result = await generateObject(retryConfig);
costTrackingOptions.costTracking.addCall({
type: "other",
metadata: {
...costTrackingOptions.metadata,
gcDetails: "generateObject fallback",
gcModel: retryConfig.model.modelId,
},
tokens: {
input: result.usage?.promptTokens ?? 0,
output: result.usage?.completionTokens ?? 0,
},
cost: calculateCost(
currentModel.modelId,
result.usage?.promptTokens ?? 0,
result.usage?.completionTokens ?? 0,
),
});
} catch (retryError) {
lastError = retryError as Error;
logger.error("Failed with fallback model", {
@ -549,7 +621,6 @@ export async function generateCompletions({
totalTokens: promptTokens + completionTokens,
},
model: currentModel.modelId,
cost: calculateCost(currentModel.modelId, promptTokens, completionTokens),
};
} catch (error) {
lastError = error as Error;
@ -589,9 +660,16 @@ export async function performLLMExtract(
// model: getModel("gemini-2.5-pro-preview-03-25", "vertex"),
model: getModel("gemini-2.5-pro-preview-03-25", "vertex"),
retryModel: getModel("gemini-2.5-pro-preview-03-25", "google"),
costTrackingOptions: {
costTracking: meta.costTracking,
metadata: {
module: "scrapeURL",
method: "performLLMExtract",
},
},
};
const { extractedDataArray, warning, smartScrapeCost, otherCost, costLimitExceededTokenUsage } =
const { extractedDataArray, warning, costLimitExceededTokenUsage } =
await extractData({
extractOptions: generationOptions,
urls: [meta.url],
@ -603,25 +681,6 @@ export async function performLLMExtract(
document.warning = warning + (document.warning ? " " + document.warning : "");
}
if (document.metadata.costTracking) {
document.metadata.costTracking.smartScrapeCallCount++;
document.metadata.costTracking.smartScrapeCost += smartScrapeCost;
document.metadata.costTracking.otherCallCount++;
document.metadata.costTracking.otherCost += otherCost;
document.metadata.costTracking.totalCost += smartScrapeCost + otherCost;
if (costLimitExceededTokenUsage) {
document.metadata.costTracking.costLimitExceededTokenUsage = costLimitExceededTokenUsage;
}
} else {
document.metadata.costTracking = {
smartScrapeCallCount: 1,
smartScrapeCost: smartScrapeCost,
otherCallCount: 1,
otherCost: otherCost,
totalCost: smartScrapeCost + otherCost,
};
}
// IMPORTANT: here it only get's the last page!!!
const extractedData =
extractedDataArray[extractedDataArray.length - 1] ?? undefined;
@ -758,7 +817,8 @@ export function removeDefaultProperty(schema: any): any {
export async function generateSchemaFromPrompt(
prompt: string,
logger: Logger,
): Promise<{ extract: any; cost: number }> {
costTracking: CostTracking,
): Promise<{ extract: any }> {
const model = getModel("gpt-4o", "openai");
const retryModel = getModel("gpt-4o-mini", "openai");
const temperatures = [0, 0.1, 0.3]; // Different temperatures to try
@ -766,7 +826,7 @@ export async function generateSchemaFromPrompt(
for (const temp of temperatures) {
try {
const { extract, cost } = await generateCompletions({
const { extract } = await generateCompletions({
logger: logger.child({
method: "generateSchemaFromPrompt/generateCompletions",
}),
@ -802,10 +862,16 @@ Return a valid JSON schema object with properties that would capture the informa
prompt: `Generate a JSON schema for extracting the following information: ${prompt}`,
// temperature: temp,
},
markdown: prompt,
costTrackingOptions: {
costTracking,
metadata: {
module: "scrapeURL",
method: "generateSchemaFromPrompt",
},
},
});
return { extract, cost };
return { extract };
} catch (error) {
lastError = error as Error;
logger.warn(`Failed attempt with temperature ${temp}: ${error.message}`);

11
apps/api/src/services/queue-worker.ts
View File

@ -75,6 +75,7 @@ import { performDeepResearch } from "../lib/deep-research/deep-research-service"
import { performGenerateLlmsTxt } from "../lib/generate-llmstxt/generate-llmstxt-service";
import { updateGeneratedLlmsTxt } from "../lib/generate-llmstxt/generate-llmstxt-redis";
import { performExtraction_F0 } from "../lib/extract/fire-0/extraction-service-f0";
import { CostTracking } from "../lib/extract/extraction-service";
configDotenv();
@ -1010,6 +1011,7 @@ async function processJob(job: Job & { id: string }, token: string) {
// };
// return data;
// }
const costTracking = new CostTracking();
try {
job.updateProgress({
@ -1030,6 +1032,7 @@ async function processJob(job: Job & { id: string }, token: string) {
startWebScraperPipeline({
job,
token,
costTracking,
}),
...(job.data.scrapeOptions.timeout !== undefined
? [
@ -1171,6 +1174,7 @@ async function processJob(job: Job & { id: string }, token: string) {
scrapeOptions: job.data.scrapeOptions,
origin: job.data.origin,
crawl_id: job.data.crawl_id,
cost_tracking: costTracking,
},
true,
);
@ -1276,10 +1280,6 @@ async function processJob(job: Job & { id: string }, token: string) {
await finishCrawlIfNeeded(job, sc);
} else {
const cost_tracking = doc?.metadata?.costTracking;
delete doc.metadata.costTracking;
await logJob({
job_id: job.id,
success: true,
@ -1293,7 +1293,7 @@ async function processJob(job: Job & { id: string }, token: string) {
scrapeOptions: job.data.scrapeOptions,
origin: job.data.origin,
num_tokens: 0, // TODO: fix
cost_tracking,
cost_tracking: costTracking,
});
indexJob(job, doc);
@ -1442,6 +1442,7 @@ async function processJob(job: Job & { id: string }, token: string) {
scrapeOptions: job.data.scrapeOptions,
origin: job.data.origin,
crawl_id: job.data.crawl_id,
cost_tracking: costTracking,
},
true,
);

1
apps/api/src/types.ts
View File

@ -62,6 +62,7 @@ export interface RunWebScraperParams {
is_scrape?: boolean;
is_crawl?: boolean;
urlInvisibleInCurrentCrawl?: boolean;
costTracking: CostTracking;
}
export type RunWebScraperResult =