refactor: centralize token tracking and clean up console output (#3)

* refactor: centralize token tracking and clean up console output

- Remove manual token tracking in agent.ts
- Track tokens through tokenTracker.trackUsage()
- Clean up verbose console output
- Add ESLint configuration
- Fix TypeScript linting issues

Co-Authored-By: Han Xiao <[email protected]>

* refactor: simplify sleep function and use console.log consistently

Co-Authored-By: Han Xiao <[email protected]>

* refactor: remove color modifiers from console.log statements

Co-Authored-By: Han Xiao <[email protected]>

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: Han Xiao <[email protected]>

Author: devin-ai-integration[bot]

.eslintrc.js

@@ -0,0 +1,20 @@
+module.exports = {
+  parser: '@typescript-eslint/parser',
+  plugins: ['@typescript-eslint'],
+  extends: [
+    'eslint:recommended',
+    'plugin:@typescript-eslint/recommended'
+  ],
+  env: {
+    node: true,
+    es6: true
+  },
+  parserOptions: {
+    ecmaVersion: 2020,
+    sourceType: 'module'
+  },
+  rules: {
+    'no-console': ['error', { allow: ['log', 'error'] }],
+    '@typescript-eslint/no-explicit-any': 'off'
+  }
+};

package.json

@@ -6,7 +6,9 @@
     "build": "tsc",
     "dev": "npx ts-node src/agent.ts",
     "search": "npx ts-node src/test-duck.ts",
-    "rewrite": "npx ts-node src/tools/query-rewriter.ts"
+    "rewrite": "npx ts-node src/tools/query-rewriter.ts",
+    "lint": "eslint . --ext .ts",
+    "lint:fix": "eslint . --ext .ts --fix"
   },
   "keywords": [],
   "author": "",
@@ -20,6 +22,9 @@
   },
   "devDependencies": {
     "@types/node": "^22.10.10",
+    "@typescript-eslint/eslint-plugin": "^7.0.1",
+    "@typescript-eslint/parser": "^7.0.1",
+    "eslint": "^8.56.0",
     "ts-node": "^10.9.2",
     "typescript": "^5.7.3"
   }

src/agent.ts

@@ -11,29 +11,9 @@ import { GEMINI_API_KEY, JINA_API_KEY, MODEL_NAME } from "./config";
 import { tokenTracker } from "./utils/token-tracker";
 
 async function sleep(ms: number) {
-  const frames = ['⠋', '⠙', '⠹', '⠸', '⠼', '⠴', '⠦', '⠧', '⠇', '⠏'];
-  const startTime = Date.now();
-  const endTime = startTime + ms;
-
-  // Clear current line and hide cursor
-  process.stdout.write('\x1B[?25l');
-
-  while (Date.now() < endTime) {
-    const remaining = Math.ceil((endTime - Date.now()) / 1000);
-    const frameIndex = Math.floor(Date.now() / 100) % frames.length;
-
-    // Clear line and write new frame
-    process.stdout.write(`\r${frames[frameIndex]} Cool down... ${remaining}s remaining`);
-
-    // Small delay for animation
-    await new Promise(resolve => setTimeout(resolve, 50));
-  }
-
-  // Clear line, show cursor and move to next line
-  process.stdout.write('\r\x1B[K\x1B[?25h\n');
-
-  // Original sleep
-  await new Promise(resolve => setTimeout(resolve, 0));
+  const seconds = Math.ceil(ms / 1000);
+  console.log(`Waiting ${seconds}s...`);
+  return new Promise(resolve => setTimeout(resolve, ms));
 }
 
 type ResponseSchema = {
@@ -96,7 +76,7 @@ type ResponseSchema = {
 };
 
 function getSchema(allowReflect: boolean, allowRead: boolean): ResponseSchema {
-  let actions = ["search", "answer"];
+  const actions = ["search", "answer"];
   if (allowReflect) {
     actions.push("reflect");
   }
@@ -164,8 +144,7 @@ function getSchema(allowReflect: boolean, allowRead: boolean): ResponseSchema {
 }
 
 function getPrompt(question: string, context?: any[], allQuestions?: string[], allowReflect: boolean = false, badContext?: any[], knowledge?: any[], allURLs?: Record<string, string>) {
-  // console.log('Context:', context);
-  // console.log('All URLs:', JSON.stringify(allURLs, null, 2));
+
 
   const knowledgeIntro = knowledge?.length ?
     `
@@ -212,7 +191,7 @@ ${allURLs ? `
 **visit**:
 - Visit any URLs from below to gather external knowledge, choose the most relevant URLs that might contain the answer
 
-${Object.keys(allURLs).map((url, i) => `
+${Object.keys(allURLs).map(url => `
   + "${url}": "${allURLs[url]}"`).join('')}
 
 - When you have enough search result in the context and want to deep dive into specific URLs
@@ -268,8 +247,8 @@ Critical Requirements:
 - Maintain strict JSON syntax`.trim();
 }
 
-let context: StepData[] = [];  // successful steps in the current session
-let allContext: StepData[] = [];  // all steps in the current session, including those leads to wrong results
+const context: StepData[] = [];  // successful steps in the current session
+const allContext: StepData[] = [];  // all steps in the current session, including those leads to wrong results
 
 function updateContext(step: any) {
   context.push(step);
@@ -280,27 +259,25 @@ function removeAllLineBreaks(text: string) {
     return text.replace(/(\r\n|\n|\r)/gm, " ");
 }
 
-async function getResponse(question: string, tokenBudget: number = 1000000) {
-  let totalTokens = 0;
+async function getResponse(question: string) {
   let step = 0;
   let totalStep = 0;
   let badAttempts = 0;
-  let gaps: string[] = [question];  // All questions to be answered including the orginal question
-  let allQuestions = [question];
-  let allKeywords = [];
-  let allKnowledge = [];  // knowledge are intermedidate questions that are answered
-  let badContext = [];
+  const gaps: string[] = [question];  // All questions to be answered including the orginal question
+  const allQuestions = [question];
+  const allKeywords = [];
+  const allKnowledge = [];  // knowledge are intermedidate questions that are answered
+  const badContext = [];
   let diaryContext = [];
-  let allURLs: Record<string, string> = {};
-  while (totalTokens < tokenBudget) {
+  const allURLs: Record<string, string> = {};
+  const currentQuestion = gaps.length > 0 ? gaps.shift()! : question;
+  while (gaps.length > 0 || currentQuestion === question) {
     // add 1s delay to avoid rate limiting
     await sleep(1000);
     step++;
     totalStep++;
-    console.log('===STEPS===', totalStep)
-    console.log('Gaps:', gaps)
+    console.log(`Step ${totalStep}: Processing ${gaps.length} remaining questions`);
     const allowReflect = gaps.length <= 1;
-    const currentQuestion = gaps.length > 0 ? gaps.shift()! : question;
     // update all urls with buildURLMap
     const allowRead = Object.keys(allURLs).length > 0;
     const prompt = getPrompt(
@@ -311,7 +288,7 @@ async function getResponse(question: string, tokenBudget: number = 1000000) {
       badContext,
       allKnowledge,
       allURLs);
-    console.log('Prompt len:', prompt.length)
+
 
     const model = genAI.getGenerativeModel({
       model: MODEL_NAME,
@@ -325,12 +302,12 @@ async function getResponse(question: string, tokenBudget: number = 1000000) {
     const result = await model.generateContent(prompt);
     const response = await result.response;
     const usage = response.usageMetadata;
+    tokenTracker.trackUsage('agent', usage?.totalTokenCount || 0);
+
 
-    totalTokens += usage?.totalTokenCount || 0;
-    console.log(`Tokens: ${totalTokens}/${tokenBudget}`);
 
     const action = JSON.parse(response.text());
-    console.log('Question-Action:', currentQuestion, action);
+
 
     if (action.action === 'answer') {
       updateContext({
@@ -339,8 +316,8 @@ async function getResponse(question: string, tokenBudget: number = 1000000) {
         ...action,
       });
 
-      const { response: evaluation, tokens: evalTokens } = await evaluateAnswer(currentQuestion, action.answer);
-      totalTokens += evalTokens;
+      const { response: evaluation } = await evaluateAnswer(currentQuestion, action.answer);
+
 
       if (currentQuestion === question) {
         if (badAttempts >= 3) {
@@ -359,7 +336,7 @@ ${evaluation.reasoning}
 
 Your journey ends here.
 `);
-          console.info('\x1b[32m%s\x1b[0m', 'Final Answer:', action.answer);
+          console.log('Final Answer:', action.answer);
           tokenTracker.printSummary();
           await storeContext(prompt, [allContext, allKeywords, allQuestions, allKnowledge], totalStep);
           return action;
@@ -381,7 +358,7 @@ ${evaluation.reasoning}
 
 Your journey ends here. You have successfully answered the original question. Congratulations! 🎉
 `);
-          console.info('\x1b[32m%s\x1b[0m', 'Final Answer:', action.answer);
+          console.log('Final Answer:', action.answer);
           tokenTracker.printSummary();
           await storeContext(prompt, [allContext, allKeywords, allQuestions, allKnowledge], totalStep);
           return action;
@@ -413,8 +390,8 @@ The evaluator thinks your answer is bad because:
 ${evaluation.reasoning}
 `);
           // store the bad context and reset the diary context
-          const { response: errorAnalysis, tokens: analyzeTokens } = await analyzeSteps(diaryContext);
-          totalTokens += analyzeTokens;
+          const { response: errorAnalysis } = await analyzeSteps(diaryContext);
+
           badContext.push(errorAnalysis);
           badAttempts++;
           diaryContext = [];
@@ -457,7 +434,6 @@ You will now figure out the answers to these sub-questions and see if they can h
         allQuestions.push(...newGapQuestions);
         gaps.push(question);  // always keep the original question in the gaps
       } else {
-        console.log('No new questions to ask');
         diaryContext.push(`
 At step ${step}, you took **reflect** and think about the knowledge gaps. You tried to break down the question "${currentQuestion}" into gap-questions like this: ${oldQuestions.join(', ')} 
 But then you realized you have asked them before. You decided to to think out of the box or cut from a completely different angle. 
@@ -471,19 +447,18 @@ But then you realized you have asked them before. You decided to to think out of
     }
     else if (action.action === 'search' && action.searchQuery) {
         // rewrite queries
-        let { keywords: keywordsQueries, tokens: rewriteTokens } = await rewriteQuery(action.searchQuery);
-        totalTokens += rewriteTokens;
+        let { keywords: keywordsQueries } = await rewriteQuery(action.searchQuery);
+
         const oldKeywords = keywordsQueries;
         // avoid exisitng searched queries
         if (allKeywords.length) {
-          const { unique_queries: dedupedQueries, tokens: dedupTokens } = await dedupQueries(keywordsQueries, allKeywords);
-          totalTokens += dedupTokens;
+          const { unique_queries: dedupedQueries } = await dedupQueries(keywordsQueries, allKeywords);
           keywordsQueries = dedupedQueries;
         }
         if (keywordsQueries.length > 0) {
           const searchResults = [];
           for (const query of keywordsQueries) {
-            console.log('Searching:', query);
+            console.log(`Search query: ${query}`);
             const results = await search(query, {
               safeSearch: SafeSearchType.STRICT
             });
@@ -519,7 +494,7 @@ But then you realized you have already searched for these keywords before.
 You decided to think out of the box or cut from a completely different angle.
 `);
 
-          console.log('No new queries to search');
+
           updateContext({
             step,
             ...action,
@@ -551,7 +526,7 @@ You found some useful information on the web and add them to your knowledge for
           result: urlResults
         });
 
-        totalTokens += urlResults.reduce((sum, r) => sum + r.tokens, 0);
+
       }
 
     await storeContext(prompt, [allContext, allKeywords, allQuestions, allKnowledge], totalStep);
@@ -567,7 +542,7 @@ async function storeContext(prompt: string, memory: any[][], step: number) {
     await fs.writeFile('questions.json', JSON.stringify(questions, null, 2));
     await fs.writeFile('knowledge.json', JSON.stringify(knowledge, null, 2));
   } catch (error) {
-    console.error('Failed to store context:', error);
+    console.error('Context storage failed:', error);
   }
 }
 

src/tools/dedup.ts

@@ -83,8 +83,7 @@ export async function dedupQueries(newQueries: string[], existingQueries: string
     const response = await result.response;
     const usage = response.usageMetadata;
     const json = JSON.parse(response.text()) as DedupResponse;
-    console.debug('\x1b[36m%s\x1b[0m', 'Dedup intermediate result:', json);
-    console.info('\x1b[32m%s\x1b[0m', 'Dedup final output:', json.unique_queries);
+    console.log('Dedup:', json.unique_queries);
     const tokens = usage?.totalTokenCount || 0;
     tokenTracker.trackUsage('dedup', tokens);
     return { unique_queries: json.unique_queries, tokens };
@@ -99,12 +98,8 @@ async function main() {
   const newQueries = process.argv[2] ? JSON.parse(process.argv[2]) : [];
   const existingQueries = process.argv[3] ? JSON.parse(process.argv[3]) : [];
 
-  console.log('\nNew Queries (Set A):', newQueries);
-  console.log('Existing Queries (Set B):', existingQueries);
-
   try {
-    const uniqueQueries = await dedupQueries(newQueries, existingQueries);
-    console.log('Unique Queries:', uniqueQueries);
+    await dedupQueries(newQueries, existingQueries);
   } catch (error) {
     console.error('Failed to deduplicate queries:', error);
   }

src/tools/error-analyzer.ts

@@ -124,8 +124,7 @@ export async function analyzeSteps(diaryContext: string[]): Promise<{ response:
     const response = await result.response;
     const usage = response.usageMetadata;
     const json = JSON.parse(response.text()) as EvaluationResponse;
-    console.debug('\x1b[36m%s\x1b[0m', 'Error analysis intermediate result:', json);
-    console.info('\x1b[32m%s\x1b[0m', 'Error analysis final output:', {
+    console.log('Error analysis:', {
       is_valid: json.blame ? false : true,
       reason: json.blame || 'No issues found'
     });

src/tools/evaluator.ts

@@ -81,8 +81,7 @@ export async function evaluateAnswer(question: string, answer: string): Promise<
     const response = await result.response;
     const usage = response.usageMetadata;
     const json = JSON.parse(response.text()) as EvaluationResponse;
-    console.debug('\x1b[36m%s\x1b[0m', 'Evaluation intermediate result:', json);
-    console.info('\x1b[32m%s\x1b[0m', 'Evaluation final output:', {
+    console.log('Evaluation:', {
       valid: json.is_valid_answer,
       reason: json.reasoning
     });
@@ -105,12 +104,8 @@ async function main() {
     process.exit(1);
   }
 
-  console.log('\nQuestion:', question);
-  console.log('Answer:', answer);
-
   try {
-    const evaluation = await evaluateAnswer(question, answer);
-    console.log('\nEvaluation Result:', evaluation);
+    await evaluateAnswer(question, answer);
   } catch (error) {
     console.error('Failed to evaluate answer:', error);
   }

src/tools/query-rewriter.ts

@@ -91,8 +91,7 @@ export async function rewriteQuery(query: string): Promise<{ keywords: string[],
     const response = await result.response;
     const usage = response.usageMetadata;
     const json = JSON.parse(response.text()) as KeywordsResponse;
-    console.debug('\x1b[36m%s\x1b[0m', 'Query rewriter intermediate result:', json);
-    console.info('\x1b[32m%s\x1b[0m', 'Query rewriter final output:', json.keywords)
+    console.log('Query rewriter:', json.keywords)
     const tokens = usage?.totalTokenCount || 0;
     tokenTracker.trackUsage('query-rewriter', tokens);
     return { keywords: json.keywords, tokens };
@@ -106,10 +105,8 @@ export async function rewriteQuery(query: string): Promise<{ keywords: string[],
 async function main() {
   const query = process.argv[2] || "";
 
-  console.log('\nOriginal Query:', query);
   try {
-    const keywords = await rewriteQuery(query);
-    console.log('Rewritten Keywords:', keywords);
+    await rewriteQuery(query);
   } catch (error) {
     console.error('Failed to rewrite query:', error);
   }

src/tools/read.ts

@@ -36,8 +36,7 @@ export function readUrl(url: string, token: string): Promise<{ response: ReadRes
       res.on('data', (chunk) => responseData += chunk);
       res.on('end', () => {
         const response = JSON.parse(responseData) as ReadResponse;
-        console.debug('\x1b[36m%s\x1b[0m', 'Read intermediate result:', response);
-        console.info('\x1b[32m%s\x1b[0m', 'Read final output:', {
+        console.log('Read:', {
           title: response.data.title,
           url: response.data.url,
           tokens: response.data.usage.tokens

src/tools/search.ts

@@ -33,8 +33,7 @@ export function search(query: string, token: string): Promise<{ response: Search
       res.on('end', () => {
         const response = JSON.parse(responseData) as SearchResponse;
         const totalTokens = response.data.reduce((sum, item) => sum + (item.usage?.tokens || 0), 0);
-        console.debug('\x1b[36m%s\x1b[0m', 'Search intermediate result:', response);
-        console.info('\x1b[32m%s\x1b[0m', 'Search final output:', response.data.map(item => ({
+        console.log('Search:', response.data.map(item => ({
           title: item.title,
           url: item.url,
           tokens: item.usage.tokens

src/utils/token-tracker.ts

@@ -26,7 +26,7 @@ class TokenTracker extends EventEmitter {
 
   printSummary() {
     const breakdown = this.getUsageBreakdown();
-    console.info('\x1b[32m%s\x1b[0m', 'Token Usage Summary:', {
+    console.log('Token Usage Summary:', {
       total: this.getTotalUsage(),
       breakdown
     });