diff --git a/apps/api/src/lib/extract/completions/batchExtract.ts b/apps/api/src/lib/extract/completions/batchExtract.ts index aecad61f..1bffeaa2 100644 --- a/apps/api/src/lib/extract/completions/batchExtract.ts +++ b/apps/api/src/lib/extract/completions/batchExtract.ts @@ -22,6 +22,7 @@ type BatchExtractOptions = { systemPrompt: string; doc: Document; useAgent: boolean; + extractId?: string; }; /** @@ -44,7 +45,7 @@ export async function batchExtractPromise(options: BatchExtractOptions, logger: smartScrapeCallCount: number; otherCallCount: number; }> { - const { multiEntitySchema, links, prompt, systemPrompt, doc, useAgent } = options; + const { multiEntitySchema, links, prompt, systemPrompt, doc, useAgent, extractId } = options; const generationOptions: GenerateCompletionsOptions = { @@ -74,6 +75,7 @@ export async function batchExtractPromise(options: BatchExtractOptions, logger: extractOptions: generationOptions, urls: [doc.metadata.sourceURL || doc.metadata.url || ""], useAgent, + extractId, }); extractedDataArray = e; warning = w; diff --git a/apps/api/src/lib/extract/completions/singleAnswer.ts b/apps/api/src/lib/extract/completions/singleAnswer.ts index 27e3cad0..5e76954e 100644 --- a/apps/api/src/lib/extract/completions/singleAnswer.ts +++ b/apps/api/src/lib/extract/completions/singleAnswer.ts @@ -14,7 +14,8 @@ export async function singleAnswerCompletion({ links, prompt, systemPrompt, - useAgent + useAgent, + extractId, }: { singleAnswerDocs: Document[]; rSchema: any; @@ -22,6 +23,7 @@ export async function singleAnswerCompletion({ prompt: string; systemPrompt: string; useAgent: boolean; + extractId?: string; }): Promise<{ extract: any; tokenUsage: TokenUsage; @@ -51,6 +53,7 @@ export async function singleAnswerCompletion({ extractOptions: generationOptions, urls: singleAnswerDocs.map(doc => doc.metadata.url || doc.metadata.sourceURL || ""), useAgent, + extractId, }); const completion = { diff --git a/apps/api/src/lib/extract/extraction-service.ts b/apps/api/src/lib/extract/extraction-service.ts index 1ea5651f..1a6a4262 100644 --- a/apps/api/src/lib/extract/extraction-service.ts +++ b/apps/api/src/lib/extract/extraction-service.ts @@ -430,7 +430,8 @@ export async function performExtraction( prompt: request.prompt ?? "", systemPrompt: request.systemPrompt ?? "", doc, - useAgent: isAgentExtractModelValid(request.agent?.model) + useAgent: isAgentExtractModelValid(request.agent?.model), + extractId, }, logger); // Race between timeout and completion @@ -741,6 +742,7 @@ export async function performExtraction( prompt: request.prompt ?? "", systemPrompt: request.systemPrompt ?? "", useAgent: isAgentExtractModelValid(request.agent?.model), + extractId, }); costTracking.smartScrapeCost += singleAnswerSmartScrapeCost; costTracking.smartScrapeCallCount += singleAnswerSmartScrapeCallCount; diff --git a/apps/api/src/scraper/scrapeURL/lib/extractSmartScrape.ts b/apps/api/src/scraper/scrapeURL/lib/extractSmartScrape.ts index 840c56fc..7380f380 100644 --- a/apps/api/src/scraper/scrapeURL/lib/extractSmartScrape.ts +++ b/apps/api/src/scraper/scrapeURL/lib/extractSmartScrape.ts @@ -184,10 +184,12 @@ export async function extractData({ extractOptions, urls, useAgent, + extractId, }: { extractOptions: GenerateCompletionsOptions; urls: string[]; useAgent: boolean; + extractId?: string; }): Promise<{ extractedDataArray: any[]; warning: any; @@ -273,7 +275,7 @@ export async function extractData({ let smartscrapeResults: SmartScrapeResult[]; if (isSingleUrl) { smartscrapeResults = [ - await smartScrape(urls[0], extract?.smartscrape_prompt), + await smartScrape(urls[0], extract?.smartscrape_prompt, extractId), ]; smartScrapeCost += smartscrapeResults[0].tokenUsage; smartScrapeCallCount++; @@ -285,6 +287,7 @@ export async function extractData({ return await smartScrape( urls[page.page_index], page.smartscrape_prompt, + extractId, ); }), ); diff --git a/apps/api/src/scraper/scrapeURL/lib/smartScrape.ts b/apps/api/src/scraper/scrapeURL/lib/smartScrape.ts index 7ccebe75..8e0b45f6 100644 --- a/apps/api/src/scraper/scrapeURL/lib/smartScrape.ts +++ b/apps/api/src/scraper/scrapeURL/lib/smartScrape.ts @@ -49,6 +49,7 @@ export async function smartScrape( url: string, prompt: string, sessionId?: string, + extractId?: string, ): Promise { try { logger.info("Initiating smart scrape request", { url, prompt, sessionId });