diff --git a/apps/api/src/controllers/v1/batch-scrape.ts b/apps/api/src/controllers/v1/batch-scrape.ts index b4da87c5..d5896f6d 100644 --- a/apps/api/src/controllers/v1/batch-scrape.ts +++ b/apps/api/src/controllers/v1/batch-scrape.ts @@ -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>, diff --git a/apps/api/src/controllers/v1/scrape.ts b/apps/api/src/controllers/v1/scrape.ts index de33a7b9..86ca905b 100644 --- a/apps/api/src/controllers/v1/scrape.ts +++ b/apps/api/src/controllers/v1/scrape.ts @@ -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, diff --git a/apps/api/src/controllers/v1/search.ts b/apps/api/src/controllers/v1/search.ts index 06f26700..1c902df4 100644 --- a/apps/api/src/controllers/v1/search.ts +++ b/apps/api/src/controllers/v1/search.ts @@ -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 { 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 { 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({ diff --git a/apps/api/src/controllers/v1/types.ts b/apps/api/src/controllers/v1/types.ts index 92c95e20..38881dd2 100644 --- a/apps/api/src/controllers/v1/types.ts +++ b/apps/api/src/controllers/v1/types.ts @@ -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?: { diff --git a/apps/api/src/lib/deep-research/deep-research-service.ts b/apps/api/src/lib/deep-research/deep-research-service.ts index 64d94946..296661c9 100644 --- a/apps/api/src/lib/deep-research/deep-research-service.ts +++ b/apps/api/src/lib/deep-research/deep-research-service.ts @@ -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", diff --git a/apps/api/src/lib/deep-research/research-manager.ts b/apps/api/src/lib/deep-research/research-manager.ts index 87655fac..0f9618a9 100644 --- a/apps/api/src/lib/deep-research/research-manager.ts +++ b/apps/api/src/lib/deep-research/research-manager.ts @@ -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 { 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 { @@ -312,6 +330,13 @@ export class ResearchLLMService { }, markdown: "", model: getModel("o3-mini"), + costTrackingOptions: { + costTracking, + metadata: { + module: "deep-research", + method: "generateFinalAnalysis", + }, + }, }); return extract; diff --git a/apps/api/src/lib/extract/completions/analyzeSchemaAndPrompt.ts b/apps/api/src/lib/extract/completions/analyzeSchemaAndPrompt.ts index 98e3ccc0..2626c50b 100644 --- a/apps/api/src/lib/extract/completions/analyzeSchemaAndPrompt.ts +++ b/apps/api/src/lib/extract/completions/analyzeSchemaAndPrompt.ts @@ -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, }; } diff --git a/apps/api/src/lib/extract/completions/batchExtract.ts b/apps/api/src/lib/extract/completions/batchExtract.ts index c57bcd7a..b7075c6e 100644 --- a/apps/api/src/lib/extract/completions/batchExtract.ts +++ b/apps/api/src/lib/extract/completions/batchExtract.ts @@ -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 }); } diff --git a/apps/api/src/lib/extract/completions/checkShouldExtract.ts b/apps/api/src/lib/extract/completions/checkShouldExtract.ts index 3bff4fc7..e2c10ade 100644 --- a/apps/api/src/lib/extract/completions/checkShouldExtract.ts +++ b/apps/api/src/lib/extract/completions/checkShouldExtract.ts @@ -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, }; } diff --git a/apps/api/src/lib/extract/completions/singleAnswer.ts b/apps/api/src/lib/extract/completions/singleAnswer.ts index 3f15edfb..866b6049 100644 --- a/apps/api/src/lib/extract/completions/singleAnswer.ts +++ b/apps/api/src/lib/extract/completions/singleAnswer.ts @@ -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, }; } diff --git a/apps/api/src/lib/extract/extraction-service.ts b/apps/api/src/lib/extract/extraction-service.ts index 50358328..ae3df261 100644 --- a/apps/api/src/lib/extract/extraction-service.ts +++ b/apps/api/src/lib/extract/extraction-service.ts @@ -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, + cost: number, + tokens?: { + input: number, + output: number, + }, + stack: string, + }[] = []; + + constructor() {} + + public addCall(call: Omit) { + 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 = {}; - 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); diff --git a/apps/api/src/lib/extract/fire-0/reranker-f0.ts b/apps/api/src/lib/extract/fire-0/reranker-f0.ts index 87e673df..155df0c6 100644 --- a/apps/api/src/lib/extract/fire-0/reranker-f0.ts +++ b/apps/api/src/lib/extract/fire-0/reranker-f0.ts @@ -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 { +export async function rerankLinksWithLLM_F0(options: RerankerOptions, costTracking: CostTracking): Promise { const { links, searchQuery, urlTraces } = options; const chunkSize = 100; const chunks: MapDocument[][] = []; @@ -231,7 +231,14 @@ export async function rerankLinksWithLLM_F0(options: RerankerOptions): Promise { +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)", { diff --git a/apps/api/src/lib/generate-llmstxt/generate-llmstxt-service.ts b/apps/api/src/lib/generate-llmstxt/generate-llmstxt-service.ts index fa72d582..3528f48f 100644 --- a/apps/api/src/lib/generate-llmstxt/generate-llmstxt-service.ts +++ b/apps/api/src/lib/generate-llmstxt/generate-llmstxt-service.ts @@ -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 diff --git a/apps/api/src/main/runWebScraper.ts b/apps/api/src/main/runWebScraper.ts index ba983e07..8e287180 100644 --- a/apps/api/src/main/runWebScraper.ts +++ b/apps/api/src/main/runWebScraper.ts @@ -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 & { 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 { 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; diff --git a/apps/api/src/scraper/WebScraper/sitemap.ts b/apps/api/src/scraper/WebScraper/sitemap.ts index f945cd22..67d3a2ff 100644 --- a/apps/api/src/scraper/WebScraper/sitemap.ts +++ b/apps/api/src/scraper/WebScraper/sitemap.ts @@ -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 ( diff --git a/apps/api/src/scraper/scrapeURL/index.ts b/apps/api/src/scraper/scrapeURL/index.ts index cedd275e..b20fbdfd 100644 --- a/apps/api/src/scraper/scrapeURL/index.ts +++ b/apps/api/src/scraper/scrapeURL/index.ts @@ -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 { 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 { - const meta = await buildMetaObject(id, url, options, internalOptions); + const meta = await buildMetaObject(id, url, options, internalOptions, costTracking); try { while (true) { try { diff --git a/apps/api/src/scraper/scrapeURL/lib/extractSmartScrape.ts b/apps/api/src/scraper/scrapeURL/lib/extractSmartScrape.ts index 82f16d12..e02f4c83 100644 --- a/apps/api/src/scraper/scrapeURL/lib/extractSmartScrape.ts +++ b/apps/api/src/scraper/scrapeURL/lib/extractSmartScrape.ts @@ -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, }; } diff --git a/apps/api/src/scraper/scrapeURL/lib/smartScrape.ts b/apps/api/src/scraper/scrapeURL/lib/smartScrape.ts index a913ec27..8458e506 100644 --- a/apps/api/src/scraper/scrapeURL/lib/smartScrape.ts +++ b/apps/api/src/scraper/scrapeURL/lib/smartScrape.ts @@ -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 { 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) { diff --git a/apps/api/src/scraper/scrapeURL/scrapeURL.test.ts b/apps/api/src/scraper/scrapeURL/scrapeURL.test.ts index b545266f..2f11d945 100644 --- a/apps/api/src/scraper/scrapeURL/scrapeURL.test.ts +++ b/apps/api/src/scraper/scrapeURL/scrapeURL.test.ts @@ -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) { diff --git a/apps/api/src/scraper/scrapeURL/transformers/agent.ts b/apps/api/src/scraper/scrapeURL/transformers/agent.ts index 5ad304d3..7f98bee1 100644 --- a/apps/api/src/scraper/scrapeURL/transformers/agent.ts +++ b/apps/api/src/scraper/scrapeURL/transformers/agent.ts @@ -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; diff --git a/apps/api/src/scraper/scrapeURL/transformers/diff.ts b/apps/api/src/scraper/scrapeURL/transformers/diff.ts index 8cdea891..94df276b 100644 --- a/apps/api/src/scraper/scrapeURL/transformers/diff.ts +++ b/apps/api/src/scraper/scrapeURL/transformers/diff.ts @@ -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; + }; }; 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}`); diff --git a/apps/api/src/services/queue-worker.ts b/apps/api/src/services/queue-worker.ts index 1118a0f9..15af06bd 100644 --- a/apps/api/src/services/queue-worker.ts +++ b/apps/api/src/services/queue-worker.ts @@ -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, ); diff --git a/apps/api/src/types.ts b/apps/api/src/types.ts index 03f0a015..34e2f60f 100644 --- a/apps/api/src/types.ts +++ b/apps/api/src/types.ts @@ -62,6 +62,7 @@ export interface RunWebScraperParams { is_scrape?: boolean; is_crawl?: boolean; urlInvisibleInCurrentCrawl?: boolean; + costTracking: CostTracking; } export type RunWebScraperResult =