From 7ec22a23d1cce08351ad9114770332645fd60a78 Mon Sep 17 00:00:00 2001
From: Han Xiao <han.xiao@jina.ai>
Date: Sat, 22 Feb 2025 16:06:41 +0800
Subject: [PATCH] fix: streaming

---
 src/agent.ts            | 10 ++++++----
 src/app.ts              | 15 +++++++++------
 src/tools/jina-dedup.ts |  2 +-
 src/types.ts            |  2 ++
 4 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/src/agent.ts b/src/agent.ts
index 3905538..6b52df0 100644
--- a/src/agent.ts
+++ b/src/agent.ts
@@ -297,7 +297,7 @@ export async function getResponse(question?: string,
                                   maxBadAttempts: number = 3,
                                   existingContext?: Partial<TrackerContext>,
                                   messages?: Array<CoreAssistantMessage | CoreUserMessage>
-): Promise<{ result: StepAction; context: TrackerContext }> {
+): Promise<{ result: StepAction; context: TrackerContext; visitedURLs: string[] }> {
   const context: TrackerContext = {
     tokenTracker: existingContext?.tokenTracker || new TokenTracker(tokenBudget),
     actionTracker: existingContext?.actionTracker || new ActionTracker()
@@ -799,7 +799,7 @@ But unfortunately, you failed to solve the issue. You need to think out of the b
   console.log(thisStep)
 
   await storeContext(system, schema, [allContext, allKeywords, allQuestions, allKnowledge], totalStep);
-  return {result: thisStep, context};
+  return {result: thisStep, context, visitedURLs};
 
 }
 
@@ -841,9 +841,11 @@ export async function main() {
   const question = process.argv[2] || "";
   const {
     result: finalStep,
-    context: tracker
-  } = await getResponse(question) as { result: AnswerAction; context: TrackerContext };
+    context: tracker,
+    visitedURLs: visitedURLs
+  } = await getResponse(question) as { result: AnswerAction; context: TrackerContext; visitedURLs: string[] };
   console.log('Final Answer:', finalStep.answer);
+  console.log('Visited URLs:', visitedURLs);
 
   tracker.tokenTracker.printSummary();
 }
diff --git a/src/app.ts b/src/app.ts
index 04382be..7f7441f 100644
--- a/src/app.ts
+++ b/src/app.ts
@@ -559,7 +559,7 @@ app.post('/v1/chat/completions', (async (req: Request, res: Response) => {
   }
 
   try {
-    const {result: finalStep} = await getResponse(undefined, tokenBudget, maxBadAttempts, context, body.messages)
+    const {result: finalStep, visitedURLs: visitedURLs} = await getResponse(undefined, tokenBudget, maxBadAttempts, context, body.messages)
 
     const usage = context.tokenTracker.getTotalUsageSnakeCase();
     if (body.stream) {
@@ -595,7 +595,8 @@ app.post('/v1/chat/completions', (async (req: Request, res: Response) => {
           logprobs: null,
           finish_reason: 'stop'
         }],
-        usage
+        usage,
+        visitedURLs
       };
       res.write(`data: ${JSON.stringify(finalChunk)}\n\n`);
       res.end();
@@ -616,7 +617,8 @@ app.post('/v1/chat/completions', (async (req: Request, res: Response) => {
           logprobs: null,
           finish_reason: 'stop'
         }],
-        usage
+        usage,
+        visitedURLs
       };
 
       // Log final response (excluding full content for brevity)
@@ -624,7 +626,8 @@ app.post('/v1/chat/completions', (async (req: Request, res: Response) => {
         id: response.id,
         status: 200,
         contentLength: response.choices[0].message.content.length,
-        usage: response.usage
+        usage: response.usage,
+        visitedURLs: response.visitedURLs
       });
 
       res.json(response);
@@ -662,7 +665,7 @@ app.post('/v1/chat/completions', (async (req: Request, res: Response) => {
           logprobs: null,
           finish_reason: null
         }],
-        usage
+        usage,
       };
       res.write(`data: ${JSON.stringify(closeThinkChunk)}\n\n`);
 
@@ -700,7 +703,7 @@ app.post('/v1/chat/completions', (async (req: Request, res: Response) => {
           logprobs: null,
           finish_reason: 'stop'
         }],
-        usage
+        usage,
       };
       res.json(response);
     }
diff --git a/src/tools/jina-dedup.ts b/src/tools/jina-dedup.ts
index c7e304d..abcbdeb 100644
--- a/src/tools/jina-dedup.ts
+++ b/src/tools/jina-dedup.ts
@@ -3,7 +3,7 @@ import {TokenTracker} from "../utils/token-tracker";
 import {JINA_API_KEY} from "../config";
 
 const JINA_API_URL = 'https://api.jina.ai/v1/embeddings';
-const SIMILARITY_THRESHOLD = 0.88; // Adjustable threshold for cosine similarity
+const SIMILARITY_THRESHOLD = 0.85; // Adjustable threshold for cosine similarity
 
 const JINA_API_CONFIG = {
   MODEL: 'jina-embeddings-v3',
diff --git a/src/types.ts b/src/types.ts
index 21d8a05..af1c23f 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -231,6 +231,7 @@ export interface ChatCompletionResponse {
     completion_tokens: number;
     total_tokens: number;
   };
+  visitedURLs?: string[];
 }
 
 export interface ChatCompletionChunk {
@@ -249,6 +250,7 @@ export interface ChatCompletionChunk {
     finish_reason: null | 'stop';
   }>;
   usage?: any;
+  visitedURLs?: string[];
 }
 
 // Tracker Types