From f742478d159c6fb3db2b33b56ddd583fa54e2acd Mon Sep 17 00:00:00 2001
From: Han Xiao <han.xiao@jina.ai>
Date: Fri, 14 Mar 2025 13:41:29 +0800
Subject: [PATCH] feat: add num urls

---
 src/agent.ts | 5 +++--
 src/app.ts   | 2 +-
 src/types.ts | 1 +
 3 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/src/agent.ts b/src/agent.ts
index fbb1979..9922385 100644
--- a/src/agent.ts
+++ b/src/agent.ts
@@ -238,7 +238,8 @@ export async function getResponse(question?: string,
                                   tokenBudget: number = 1_000_000,
                                   maxBadAttempts: number = 3,
                                   existingContext?: Partial<TrackerContext>,
-                                  messages?: Array<CoreMessage>
+                                  messages?: Array<CoreMessage>,
+                                  numReturnedURLs: number = 100
 ): Promise<{ result: StepAction; context: TrackerContext; visitedURLs: string[], readURLs: string[], allURLs: string[] }> {
 
   let step = 0;
@@ -845,7 +846,7 @@ But unfortunately, you failed to solve the issue. You need to think out of the b
   }, totalStep);
 
   // max return 300 urls
-  const returnedURLs = weightedURLs.slice(0, 50).map(r => r.url);
+  const returnedURLs = weightedURLs.slice(0, numReturnedURLs).map(r => r.url);
   return {
     result: thisStep,
     context,
diff --git a/src/app.ts b/src/app.ts
index 9a4c08b..ffcceee 100644
--- a/src/app.ts
+++ b/src/app.ts
@@ -548,7 +548,7 @@ app.post('/v1/chat/completions', (async (req: Request, res: Response) => {
       visitedURLs,
       readURLs,
       allURLs
-    } = await getResponse(undefined, tokenBudget, maxBadAttempts, context, body.messages)
+    } = await getResponse(undefined, tokenBudget, maxBadAttempts, context, body.messages, body.numReturnedURLs)
     let finalAnswer = (finalStep as AnswerAction).mdAnswer;
 
     const annotations = (finalStep as AnswerAction).references?.map(ref => ({
diff --git a/src/types.ts b/src/types.ts
index 222ed52..d66cde8 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -215,6 +215,7 @@ export interface ChatCompletionRequest {
   max_attempts?: number;
 
   response_format?: ResponseFormat;
+  numReturnedURLs?: number;
 }
 
 export interface URLAnnotation {