github/node-DeepResearch
1
0
Fork 0
mirror of https://github.com/jina-ai/node-DeepResearch.git synced 2025-12-26 06:28:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

Refactor Tools to Use config.ts (#2)

Browse Source 
* refactor: remove redundant code and use config imports

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

* refactor: standardize token tracking in search and read tools

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

* refactor: standardize console logging with colors and levels

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

* feat: add centralized token tracking utility

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

* feat: integrate token tracking across all tools

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

* feat: add token tracking and colored console output to agent

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] 2025-01-31 14:52:41 +08:00 committed by GitHub
parent 0ae6c790f7
commit 966ef5d026
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
8 changed files with 112 additions and 61 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
src/agent.ts
View File

@ -8,6 +8,7 @@ import { evaluateAnswer } from "./tools/evaluator";
import { StepData } from "./tools/getURLIndex";
import { analyzeSteps } from "./tools/error-analyzer";
import { GEMINI_API_KEY, JINA_API_KEY, MODEL_NAME } from "./config";
import { tokenTracker } from "./utils/token-tracker";
async function sleep(ms: number) {
const frames = ['⠋', '⠙', '⠹', '⠸', '⠼', '⠴', '⠦', '⠧', '⠇', '⠏'];
@ -358,6 +359,8 @@ ${evaluation.reasoning}
Your journey ends here.
`);
console.info('\x1b[32m%s\x1b[0m', 'Final Answer:', action.answer);
tokenTracker.printSummary();
await storeContext(prompt, [allContext, allKeywords, allQuestions, allKnowledge], totalStep);
return action;
}
@ -378,6 +381,8 @@ ${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);
tokenTracker.printSummary();
await storeContext(prompt, [allContext, allKeywords, allQuestions, allKnowledge], totalStep);
return action;
} else {
@ -525,13 +530,13 @@ 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, JINA_API_KEY);
const { response, tokens } = await readUrl(url, JINA_API_KEY);
allKnowledge.push({
question: `What is in ${response.data.url}?`,
answer: removeAllLineBreaks(response.data.content)});
// remove that url from allURLs
delete allURLs[url];
return {url, result: response};
return {url, result: response, tokens};
})
);
diaryContext.push(`
@ -546,7 +551,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.result.data.usage.tokens, 0);
totalTokens += urlResults.reduce((sum, r) => sum + r.tokens, 0);
}
await storeContext(prompt, [allContext, allKeywords, allQuestions, allKnowledge], totalStep);

34
src/tools/dedup.ts
View File

@ -1,23 +1,6 @@
import { GoogleGenerativeAI, SchemaType } from "@google/generative-ai";
import dotenv from 'dotenv';
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");
}
import { GEMINI_API_KEY, MODEL_NAME } from "../config";
import { tokenTracker } from "../utils/token-tracker";
type DedupResponse = {
thought: string;
@ -42,11 +25,9 @@ const responseSchema = {
required: ["thought", "unique_queries"]
};
const modelName = 'gemini-1.5-flash';
const genAI = new GoogleGenerativeAI(apiKey);
const genAI = new GoogleGenerativeAI(GEMINI_API_KEY);
const model = genAI.getGenerativeModel({
model: modelName,
model: MODEL_NAME,
generationConfig: {
temperature: 0.1,
responseMimeType: "application/json",
@ -102,8 +83,11 @@ 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.log('Dedup:', json);
return { unique_queries: json.unique_queries, tokens: usage?.totalTokenCount || 0 };
console.debug('\x1b[36m%s\x1b[0m', 'Dedup intermediate result:', json);
console.info('\x1b[32m%s\x1b[0m', 'Dedup final output:', json.unique_queries);
const tokens = usage?.totalTokenCount || 0;
tokenTracker.trackUsage('dedup', tokens);
return { unique_queries: json.unique_queries, tokens };
} catch (error) {
console.error('Error in deduplication analysis:', error);
throw error;

37
src/tools/error-analyzer.ts
View File

@ -1,23 +1,6 @@
import {GoogleGenerativeAI, SchemaType} from "@google/generative-ai";
import dotenv from 'dotenv';
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");
}
import { GEMINI_API_KEY, MODEL_NAME } from "../config";
import { tokenTracker } from "../utils/token-tracker";
type EvaluationResponse = {
recap: string;
@ -44,11 +27,9 @@ const responseSchema = {
required: ["recap", "blame", "improvement"]
};
const modelName = 'gemini-1.5-flash';
const genAI = new GoogleGenerativeAI(apiKey);
const genAI = new GoogleGenerativeAI(GEMINI_API_KEY);
const model = genAI.getGenerativeModel({
model: modelName,
model: MODEL_NAME,
generationConfig: {
temperature: 0,
responseMimeType: "application/json",
@ -143,8 +124,14 @@ 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.log('Rejection analysis:', json);
return { response: json, tokens: usage?.totalTokenCount || 0 };
console.debug('\x1b[36m%s\x1b[0m', 'Error analysis intermediate result:', json);
console.info('\x1b[32m%s\x1b[0m', 'Error analysis final output:', {
is_valid: json.blame ? false : true,
reason: json.blame || 'No issues found'
});
const tokens = usage?.totalTokenCount || 0;
tokenTracker.trackUsage('error-analyzer', tokens);
return { response: json, tokens };
} catch (error) {
console.error('Error in answer evaluation:', error);
throw error;

11
src/tools/evaluator.ts
View File

@ -1,5 +1,6 @@
import { GoogleGenerativeAI, SchemaType } from "@google/generative-ai";
import { GEMINI_API_KEY, MODEL_NAME } from "../config";
import { tokenTracker } from "../utils/token-tracker";
type EvaluationResponse = {
is_valid_answer: boolean;
@ -80,8 +81,14 @@ 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.log('Evaluation:', json);
return { response: json, tokens: usage?.totalTokenCount || 0 };
console.debug('\x1b[36m%s\x1b[0m', 'Evaluation intermediate result:', json);
console.info('\x1b[32m%s\x1b[0m', 'Evaluation final output:', {
valid: json.is_valid_answer,
reason: json.reasoning
});
const tokens = usage?.totalTokenCount || 0;
tokenTracker.trackUsage('evaluator', tokens);
return { response: json, tokens };
} catch (error) {
console.error('Error in answer evaluation:', error);
throw error;

8
src/tools/query-rewriter.ts
View File

@ -1,5 +1,6 @@
import {GoogleGenerativeAI, SchemaType} from "@google/generative-ai";
import { GEMINI_API_KEY, MODEL_NAME } from "../config";
import { tokenTracker } from "../utils/token-tracker";
type KeywordsResponse = {
keywords: string[];
@ -90,8 +91,11 @@ 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.log('Rewriter:', json)
return { keywords: json.keywords, tokens: usage?.totalTokenCount || 0 };
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)
const tokens = usage?.totalTokenCount || 0;
tokenTracker.trackUsage('query-rewriter', tokens);
return { keywords: json.keywords, tokens };
} catch (error) {
console.error('Error in query rewriting:', error);
throw error;

16
src/tools/read.ts
View File

@ -1,4 +1,5 @@
import https from 'https';
import { tokenTracker } from "../utils/token-tracker";
interface ReadResponse {
code: number;
@ -12,7 +13,7 @@ interface ReadResponse {
};
}
export function readUrl(url: string, token: string): Promise<ReadResponse> {
export function readUrl(url: string, token: string): Promise<{ response: ReadResponse, tokens: number }> {
return new Promise((resolve, reject) => {
const data = JSON.stringify({url});
@ -33,7 +34,18 @@ export function readUrl(url: string, token: string): Promise<ReadResponse> {
const req = https.request(options, (res) => {
let responseData = '';
res.on('data', (chunk) => responseData += chunk);
res.on('end', () => resolve(JSON.parse(responseData)));
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:', {
title: response.data.title,
url: response.data.url,
tokens: response.data.usage.tokens
});
const tokens = response.data?.usage?.tokens || 0;
tokenTracker.trackUsage('read', tokens);
resolve({ response, tokens });
});
});
req.on('error', reject);

16
src/tools/search.ts
View File

@ -1,4 +1,5 @@
import https from 'https';
import { tokenTracker } from "../utils/token-tracker";
interface SearchResponse {
code: number;
@ -12,7 +13,7 @@ interface SearchResponse {
}>;
}
export function search(query: string, token: string): Promise<SearchResponse> {
export function search(query: string, token: string): Promise<{ response: SearchResponse, tokens: number }> {
return new Promise((resolve, reject) => {
const options = {
hostname: 's.jina.ai',
@ -29,7 +30,18 @@ export function search(query: string, token: string): Promise<SearchResponse> {
const req = https.request(options, (res) => {
let responseData = '';
res.on('data', (chunk) => responseData += chunk);
res.on('end', () => resolve(JSON.parse(responseData)));
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 => ({
title: item.title,
url: item.url,
tokens: item.usage.tokens
})));
tokenTracker.trackUsage('search', totalTokens);
resolve({ response, tokens: totalTokens });
});
});
req.on('error', reject);

40
src/utils/token-tracker.ts Normal file
View File

@ -0,0 +1,40 @@
import { EventEmitter } from 'events';
interface TokenUsage {
tool: string;
tokens: number;
}
class TokenTracker extends EventEmitter {
private usages: TokenUsage[] = [];
trackUsage(tool: string, tokens: number) {
this.usages.push({ tool, tokens });
this.emit('usage', { tool, tokens });
}
getTotalUsage(): number {
return this.usages.reduce((sum, usage) => sum + usage.tokens, 0);
}
getUsageBreakdown(): Record<string, number> {
return this.usages.reduce((acc, { tool, tokens }) => {
acc[tool] = (acc[tool] || 0) + tokens;
return acc;
}, {} as Record<string, number>);
}
printSummary() {
const breakdown = this.getUsageBreakdown();
console.info('\x1b[32m%s\x1b[0m', 'Token Usage Summary:', {
total: this.getTotalUsage(),
breakdown
});
}
reset() {
this.usages = [];
}
}
export const tokenTracker = new TokenTracker();