From 4cedd57fcc3928e923712820c3673cef4a5647ee Mon Sep 17 00:00:00 2001
From: Han Xiao <han.xiao@jina.ai>
Date: Wed, 26 Feb 2025 15:34:54 +0800
Subject: [PATCH] feat: add api params

---
 src/agent.ts | 8 +++++---
 src/app.ts   | 9 ++++++++-
 src/types.ts | 3 +++
 3 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/src/agent.ts b/src/agent.ts
index f09c6ed..0e9a862 100644
--- a/src/agent.ts
+++ b/src/agent.ts
@@ -261,6 +261,8 @@ export async function getResponse(question?: string,
     actionTracker: existingContext?.actionTracker || new ActionTracker()
   };
 
+  const generator = new ObjectGeneratorSafe(context.tokenTracker);
+
   let schema: ZodObject<any> = SchemaGen.getAgentSchema(true, true, true, true, true)
   const gaps: string[] = [question];  // All questions to be answered including the orginal question
   const allQuestions = [question];
@@ -280,7 +282,9 @@ export async function getResponse(question?: string,
   const allURLs: Record<string, SearchResult> = {};
   const visitedURLs: string[] = [];
   const evaluationMetrics: Record<string, EvaluationType[]> = {};
-  while (context.tokenTracker.getTotalUsage().totalTokens < tokenBudget && badAttempts <= maxBadAttempts) {
+  // reserve the 10% final budget for the beast mode
+  const regularBudget = tokenBudget * 0.9;
+  while (context.tokenTracker.getTotalUsage().totalTokens < regularBudget && badAttempts <= maxBadAttempts) {
     // add 1s delay to avoid rate limiting
     step++;
     totalStep++;
@@ -314,7 +318,6 @@ export async function getResponse(question?: string,
       false,
     );
     schema = SchemaGen.getAgentSchema(allowReflect, allowRead, allowAnswer, allowSearch, allowCoding)
-    const generator = new ObjectGeneratorSafe(context.tokenTracker);
     const result = await generator.generateObject({
       model: 'agent',
       schema,
@@ -721,7 +724,6 @@ But unfortunately, you failed to solve the issue. You need to think out of the b
     );
 
     schema = SchemaGen.getAgentSchema(false, false, true, false, false);
-    const generator = new ObjectGeneratorSafe(context.tokenTracker);
     const result = await generator.generateObject({
       model: 'agentBeastMode',
       schema,
diff --git a/src/app.ts b/src/app.ts
index 87d90ac..3b337c3 100644
--- a/src/app.ts
+++ b/src/app.ts
@@ -392,11 +392,18 @@ app.post('/v1/chat/completions', (async (req: Request, res: Response) => {
   });
   console.log('messages', body.messages);
 
-  const {tokenBudget, maxBadAttempts} = getTokenBudgetAndMaxAttempts(
+  let {tokenBudget, maxBadAttempts} = getTokenBudgetAndMaxAttempts(
     body.reasoning_effort,
     body.max_completion_tokens
   );
 
+  if (body.budget_tokens) {
+    tokenBudget = body.budget_tokens;
+  }
+  if (body.max_attempts) {
+    maxBadAttempts = body.max_attempts;
+  }
+
   const requestId = Date.now().toString();
   const created = Math.floor(Date.now() / 1000);
   const context: TrackerContext = {
diff --git a/src/types.ts b/src/types.ts
index 0ce8761..1ea97eb 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -187,6 +187,9 @@ export interface ChatCompletionRequest {
   stream?: boolean;
   reasoning_effort?: 'low' | 'medium' | 'high' | null;
   max_completion_tokens?: number | null;
+
+  budget_tokens?: number | null;
+  max_attempts?: number | null;
 }
 
 export interface ChatCompletionResponse {