From c2931c6e2a6447f8cb3a835e16133ef81cb6b677 Mon Sep 17 00:00:00 2001
From: Han Xiao <han.xiao@jina.ai>
Date: Sun, 9 Feb 2025 15:23:26 +0800
Subject: [PATCH] fix: streaming

---
 config.json   | 2 +-
 src/agent.ts  | 5 +++--
 src/server.ts | 2 +-
 3 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/config.json b/config.json
index 04be083..45eeb09 100644
--- a/config.json
+++ b/config.json
@@ -11,7 +11,7 @@
   "defaults": {
     "search_provider": "jina",
     "llm_provider": "gemini",
-    "step_sleep": 1000
+    "step_sleep": 0
   },
   "providers": {
     "gemini": {
diff --git a/src/agent.ts b/src/agent.ts
index 86f70ed..a26cd12 100644
--- a/src/agent.ts
+++ b/src/agent.ts
@@ -312,7 +312,6 @@ export async function getResponse(question: string, tokenBudget: number = 1_000_
     await sleep(STEP_SLEEP);
     step++;
     totalStep++;
-    context.actionTracker.trackAction({totalStep, thisStep, gaps, badAttempts});
     const budgetPercentage = (context.tokenTracker.getTotalUsage() / tokenBudget * 100).toFixed(2);
     console.log(`Step ${totalStep} / Budget used ${budgetPercentage}%`);
     console.log('Gaps:', gaps);
@@ -355,13 +354,15 @@ export async function getResponse(question: string, tokenBudget: number = 1_000_
       object = result.object;
       totalTokens = result.totalTokens;
     }
-    context.tokenTracker.trackUsage('agent', totalTokens);
     thisStep = object as StepAction;
     // print allowed and chose action
     const actionsStr = [allowSearch, allowRead, allowAnswer, allowReflect].map((a, i) => a ? ['search', 'read', 'answer', 'reflect'][i] : null).filter(a => a).join(', ');
     console.log(`${thisStep.action} <- [${actionsStr}]`);
     console.log(thisStep)
 
+    context.actionTracker.trackAction({totalStep, thisStep, gaps, badAttempts});
+    context.tokenTracker.trackUsage('agent', totalTokens);
+
     // reset allowAnswer to true
     allowAnswer = true;
     allowReflect = true;
diff --git a/src/server.ts b/src/server.ts
index 2eabb61..8c6cf1d 100644
--- a/src/server.ts
+++ b/src/server.ts
@@ -114,7 +114,7 @@ app.post('/v1/chat/completions', (async (req: Request, res: Response) => {
           system_fingerprint: 'fp_' + requestId,
           choices: [{
             index: 0,
-            delta: { content: `${action.thisStep.think}` },
+            delta: { content: `${action.thisStep.think}\n` },
             logprobs: null,
             finish_reason: null
           }]