feat: improve token tracking and centralize config

- Add centralized config.ts for API keys and model settings
- Track token usage from all Gemini model calls
- Update tools to return token usage alongside results
- Fix const reassignment in agent.ts

Co-Authored-By: Han Xiao <han.xiao@jina.ai>

src/agent.ts

@@ -1,26 +1,13 @@
-import {GoogleGenerativeAI, SchemaType} from "@google/generative-ai";
+import { GoogleGenerativeAI, SchemaType } from "@google/generative-ai";
-import dotenv from 'dotenv';
+import { readUrl } from "./tools/read";
-import {ProxyAgent, setGlobalDispatcher} from "undici";
-import {readUrl} from "./tools/read";
 import fs from 'fs/promises';
-import {SafeSearchType, search} from "duck-duck-scrape";
+import { SafeSearchType, search } from "duck-duck-scrape";
-import {rewriteQuery} from "./tools/query-rewriter";
+import { rewriteQuery } from "./tools/query-rewriter";
-import {dedupQueries} from "./tools/dedup";
+import { dedupQueries } from "./tools/dedup";
-import {evaluateAnswer} from "./tools/evaluator";
+import { evaluateAnswer } from "./tools/evaluator";
-import {StepData} from "./tools/getURLIndex";
+import { StepData } from "./tools/getURLIndex";
-import {analyzeSteps} from "./tools/error-analyzer";
+import { analyzeSteps } from "./tools/error-analyzer";
-
+import { GEMINI_API_KEY, JINA_API_KEY, MODEL_NAME } from "./config";
-// Proxy setup remains the same
-if (process.env.https_proxy) {
-  try {
-    const proxyUrl = new URL(process.env.https_proxy).toString();
-    const dispatcher = new ProxyAgent({uri: proxyUrl});
-    setGlobalDispatcher(dispatcher);
-  } catch (error) {
-    console.error('Failed to set proxy:', error);
-  }
-}
-dotenv.config();
 
 async function sleep(ms: number) {
   const frames = ['⠋', '⠙', '⠹', '⠸', '⠼', '⠴', '⠦', '⠧', '⠇', '⠏'];
@@ -326,7 +313,7 @@ async function getResponse(question: string, tokenBudget: number = 1000000) {
     console.log('Prompt len:', prompt.length)
 
     const model = genAI.getGenerativeModel({
-      model: modelName,
+      model: MODEL_NAME,
       generationConfig: {
         temperature: 0.7,
         responseMimeType: "application/json",
@@ -351,7 +338,8 @@ async function getResponse(question: string, tokenBudget: number = 1000000) {
         ...action,
       });
 
-      const evaluation = await evaluateAnswer(currentQuestion, action.answer);
+      const { response: evaluation, tokens: evalTokens } = await evaluateAnswer(currentQuestion, action.answer);
+      totalTokens += evalTokens;
 
       if (currentQuestion === question) {
         if (badAttempts >= 3) {
@@ -420,7 +408,8 @@ The evaluator thinks your answer is bad because:
 ${evaluation.reasoning}
 `);
           // store the bad context and reset the diary context
-          const errorAnalysis = await analyzeSteps(diaryContext);
+          const { response: errorAnalysis, tokens: analyzeTokens } = await analyzeSteps(diaryContext);
+          totalTokens += analyzeTokens;
           badContext.push(errorAnalysis);
           badAttempts++;
           diaryContext = [];
@@ -477,11 +466,14 @@ 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 keywordsQueries = await rewriteQuery(action.searchQuery);
+        let { keywords: keywordsQueries, tokens: rewriteTokens } = await rewriteQuery(action.searchQuery);
+        totalTokens += rewriteTokens;
         const oldKeywords = keywordsQueries;
         // avoid exisitng searched queries
         if (allKeywords.length) {
-          keywordsQueries = await dedupQueries(keywordsQueries, allKeywords)
+          const { unique_queries: dedupedQueries, tokens: dedupTokens } = await dedupQueries(keywordsQueries, allKeywords);
+          totalTokens += dedupTokens;
+          keywordsQueries = dedupedQueries;
         }
         if (keywordsQueries.length > 0) {
           const searchResults = [];
@@ -533,7 +525,7 @@ You decided to think out of the box or cut from a completely different angle.
     else if (action.action === 'visit' && action.URLTargets?.length) {
         const urlResults = await Promise.all(
           action.URLTargets.map(async (url: string) => {
-            const response = await readUrl(url, jinaToken);
+            const response = await readUrl(url, JINA_API_KEY);
             allKnowledge.push({
               question: `What is in ${response.data.url}?`,
               answer: removeAllLineBreaks(response.data.content)});
@@ -574,13 +566,7 @@ async function storeContext(prompt: string, memory: any[][], step: number) {
   }
 }
 
-const apiKey = process.env.GEMINI_API_KEY as string;
+const genAI = new GoogleGenerativeAI(GEMINI_API_KEY);
-const jinaToken = process.env.JINA_API_KEY as string;
-if (!apiKey) throw new Error("GEMINI_API_KEY not found");
-if (!jinaToken) throw new Error("JINA_API_KEY not found");
-
-const modelName = 'gemini-1.5-flash';
-const genAI = new GoogleGenerativeAI(apiKey);
 
 const question = process.argv[2] || "";
-getResponse(question);
+getResponse(question);

src/config.ts

@@ -0,0 +1,22 @@
+import dotenv from 'dotenv';
+import { ProxyAgent, setGlobalDispatcher } from 'undici';
+
+dotenv.config();
+
+// Setup the proxy globally if present
+if (process.env.https_proxy) {
+  try {
+    const proxyUrl = new URL(process.env.https_proxy).toString();
+    const dispatcher = new ProxyAgent({ uri: proxyUrl });
+    setGlobalDispatcher(dispatcher);
+  } catch (error) {
+    console.error('Failed to set proxy:', error);
+  }
+}
+
+export const GEMINI_API_KEY = process.env.GEMINI_API_KEY as string;
+export const JINA_API_KEY = process.env.JINA_API_KEY as string;
+export const MODEL_NAME = 'gemini-1.5-flash';
+
+if (!GEMINI_API_KEY) throw new Error("GEMINI_API_KEY not found");
+if (!JINA_API_KEY) throw new Error("JINA_API_KEY not found");

src/tools/dedup.ts

@@ -95,14 +95,15 @@ Set A: ${JSON.stringify(newQueries)}
 Set B: ${JSON.stringify(existingQueries)}`;
 }
 
-export async function dedupQueries(newQueries: string[], existingQueries: string[]): Promise<string[]> {
+export async function dedupQueries(newQueries: string[], existingQueries: string[]): Promise<{ unique_queries: string[], tokens: number }> {
   try {
     const prompt = getPrompt(newQueries, existingQueries);
     const result = await model.generateContent(prompt);
     const response = await result.response;
+    const usage = response.usageMetadata;
     const json = JSON.parse(response.text()) as DedupResponse;
     console.log('Dedup:', json);
-    return json.unique_queries;
+    return { unique_queries: json.unique_queries, tokens: usage?.totalTokenCount || 0 };
   } catch (error) {
     console.error('Error in deduplication analysis:', error);
     throw error;
@@ -127,4 +128,4 @@ async function main() {
 
 if (require.main === module) {
   main().catch(console.error);
-}
+}

src/tools/error-analyzer.ts

@@ -136,16 +136,17 @@ ${diaryContext.join('\n')}
 `;
 }
 
-export async function analyzeSteps(diaryContext: string[]): Promise<EvaluationResponse> {
+export async function analyzeSteps(diaryContext: string[]): Promise<{ response: EvaluationResponse, tokens: number }> {
   try {
     const prompt = getPrompt(diaryContext);
     const result = await model.generateContent(prompt);
     const response = await result.response;
+    const usage = response.usageMetadata;
     const json = JSON.parse(response.text()) as EvaluationResponse;
     console.log('Rejection analysis:', json);
-    return json;
+    return { response: json, tokens: usage?.totalTokenCount || 0 };
   } catch (error) {
     console.error('Error in answer evaluation:', error);
     throw error;
   }
-}
+}

src/tools/evaluator.ts

@@ -1,23 +1,5 @@
 import { GoogleGenerativeAI, SchemaType } from "@google/generative-ai";
-import dotenv from 'dotenv';
+import { GEMINI_API_KEY, MODEL_NAME } from "../config";
-import { ProxyAgent, setGlobalDispatcher } from "undici";
-
-// Proxy setup
-if (process.env.https_proxy) {
-  try {
-    const proxyUrl = new URL(process.env.https_proxy).toString();
-    const dispatcher = new ProxyAgent({ uri: proxyUrl });
-    setGlobalDispatcher(dispatcher);
-  } catch (error) {
-    console.error('Failed to set proxy:', error);
-  }
-}
-dotenv.config();
-
-const apiKey = process.env.GEMINI_API_KEY;
-if (!apiKey) {
-  throw new Error("GEMINI_API_KEY not found in environment variables");
-}
 
 type EvaluationResponse = {
   is_valid_answer: boolean;
@@ -39,11 +21,9 @@ const responseSchema = {
   required: ["is_valid_answer", "reasoning"]
 };
 
-const modelName = 'gemini-1.5-flash';
+const genAI = new GoogleGenerativeAI(GEMINI_API_KEY);
-
-const genAI = new GoogleGenerativeAI(apiKey);
 const model = genAI.getGenerativeModel({
-  model: modelName,
+  model: MODEL_NAME,
   generationConfig: {
     temperature: 0,
     responseMimeType: "application/json",
@@ -93,14 +73,15 @@ Question: ${JSON.stringify(question)}
 Answer: ${JSON.stringify(answer)}`;
 }
 
-export async function evaluateAnswer(question: string, answer: string): Promise<EvaluationResponse> {
+export async function evaluateAnswer(question: string, answer: string): Promise<{ response: EvaluationResponse, tokens: number }> {
   try {
     const prompt = getPrompt(question, answer);
     const result = await model.generateContent(prompt);
     const response = await result.response;
+    const usage = response.usageMetadata;
     const json = JSON.parse(response.text()) as EvaluationResponse;
     console.log('Evaluation:', json);
-    return json;
+    return { response: json, tokens: usage?.totalTokenCount || 0 };
   } catch (error) {
     console.error('Error in answer evaluation:', error);
     throw error;
@@ -130,4 +111,4 @@ async function main() {
 
 if (require.main === module) {
   main().catch(console.error);
-}
+}

src/tools/query-rewriter.ts

@@ -1,23 +1,5 @@
 import {GoogleGenerativeAI, SchemaType} from "@google/generative-ai";
-import dotenv from 'dotenv';
+import { GEMINI_API_KEY, MODEL_NAME } from "../config";
-import {ProxyAgent, setGlobalDispatcher} from "undici";
-
-// Proxy setup remains the same
-if (process.env.https_proxy) {
-  try {
-    const proxyUrl = new URL(process.env.https_proxy).toString();
-    const dispatcher = new ProxyAgent({uri: proxyUrl});
-    setGlobalDispatcher(dispatcher);
-  } catch (error) {
-    console.error('Failed to set proxy:', error);
-  }
-}
-dotenv.config();
-
-const apiKey = process.env.GEMINI_API_KEY;
-if (!apiKey) {
-  throw new Error("GEMINI_API_KEY not found in environment variables");
-}
 
 type KeywordsResponse = {
   keywords: string[];
@@ -44,11 +26,9 @@ const responseSchema = {
   required: ["thought", "keywords"]
 };
 
-const modelName = 'gemini-1.5-flash';
+const genAI = new GoogleGenerativeAI(GEMINI_API_KEY);
-
-const genAI = new GoogleGenerativeAI(apiKey);
 const model = genAI.getGenerativeModel({
-  model: modelName,
+  model: MODEL_NAME,
   generationConfig: {
     temperature: 0.1,
     responseMimeType: "application/json",
@@ -103,16 +83,15 @@ Now, process this query:
 Input Query: ${query}`;
 }
 
-export async function rewriteQuery(query: string): Promise<string[]> {
+export async function rewriteQuery(query: string): Promise<{ keywords: string[], tokens: number }> {
-
-
   try {
     const prompt = getPrompt(query);
     const result = await model.generateContent(prompt);
     const response = await result.response;
+    const usage = response.usageMetadata;
     const json = JSON.parse(response.text()) as KeywordsResponse;
     console.log('Rewriter:', json)
-    return json.keywords;
+    return { keywords: json.keywords, tokens: usage?.totalTokenCount || 0 };
   } catch (error) {
     console.error('Error in query rewriting:', error);
     throw error;
@@ -134,4 +113,4 @@ async function main() {
 
 if (require.main === module) {
   main().catch(console.error);
-}
+}