mirror of
https://github.com/jina-ai/node-DeepResearch.git
synced 2025-12-26 06:28:56 +08:00
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 <han.xiao@jina.ai> * refactor: simplify sleep function and use console.log consistently Co-Authored-By: Han Xiao <han.xiao@jina.ai> * refactor: remove color modifiers from console.log statements Co-Authored-By: Han Xiao <han.xiao@jina.ai> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Han Xiao <han.xiao@jina.ai>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
966ef5d026
commit
f99608909c
20
.eslintrc.js
Normal file
20
.eslintrc.js
Normal file
@ -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'
|
||||
}
|
||||
};
|
||||
@ -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"
|
||||
}
|
||||
|
||||
95
src/agent.ts
95
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);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@ -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);
|
||||
}
|
||||
|
||||
@ -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'
|
||||
});
|
||||
|
||||
@ -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);
|
||||
}
|
||||
|
||||
@ -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);
|
||||
}
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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
|
||||
});
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user