From defc80af323dd05ac7e16d3414cdeed672b542e6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gerg=C5=91=20M=C3=B3ricz?= <mo.geryy@gmail.com>
Date: Wed, 16 Apr 2025 03:23:18 -0700
Subject: [PATCH] stream session IDs for single URLs

---
 apps/api/src/lib/extract/completions/singleAnswer.ts | 3 +++
 apps/api/src/lib/extract/extraction-service.ts       | 4 ++++
 2 files changed, 7 insertions(+)

diff --git a/apps/api/src/lib/extract/completions/singleAnswer.ts b/apps/api/src/lib/extract/completions/singleAnswer.ts
index 4b36a58f..3f15edfb 100644
--- a/apps/api/src/lib/extract/completions/singleAnswer.ts
+++ b/apps/api/src/lib/extract/completions/singleAnswer.ts
@@ -16,6 +16,7 @@ export async function singleAnswerCompletion({
   systemPrompt,
   useAgent,
   extractId,
+  sessionId,
 }: {
   singleAnswerDocs: Document[];
   rSchema: any;
@@ -24,6 +25,7 @@ export async function singleAnswerCompletion({
   systemPrompt: string;
   useAgent: boolean;
   extractId: string;
+  sessionId: string;
 }): Promise<{
   extract: any;
   tokenUsage: TokenUsage;
@@ -58,6 +60,7 @@ export async function singleAnswerCompletion({
     urls: singleAnswerDocs.map(doc => doc.metadata.url || doc.metadata.sourceURL || ""),
     useAgent,
     extractId,
+    sessionId,
   });
 
   const completion = {
diff --git a/apps/api/src/lib/extract/extraction-service.ts b/apps/api/src/lib/extract/extraction-service.ts
index 6f14306a..01035151 100644
--- a/apps/api/src/lib/extract/extraction-service.ts
+++ b/apps/api/src/lib/extract/extraction-service.ts
@@ -750,6 +750,8 @@ export async function performExtraction(
         };
       }
 
+      let thisSessionId = 'fc-' + crypto.randomUUID();
+
       await updateExtract(extractId, {
         status: "processing",
         steps: [
@@ -760,6 +762,7 @@ export async function performExtraction(
             discoveredLinks: links,
           },
         ],
+        sessionIds: [thisSessionId],
       });
 
       // Generate completions
@@ -781,6 +784,7 @@ export async function performExtraction(
         systemPrompt: request.systemPrompt ?? "",
         useAgent: isAgentExtractModelValid(request.agent?.model),
         extractId,
+        sessionId: thisSessionId,
       });
       costTracking.smartScrapeCost += singleAnswerSmartScrapeCost;
       costTracking.smartScrapeCallCount += singleAnswerSmartScrapeCallCount;