mirror of
https://github.com/jina-ai/node-DeepResearch.git
synced 2025-12-25 22:16:49 +08:00
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>
This commit is contained in:
Devin AI
parent
d12bc09b8b
commit
738af73010
58
src/agent.ts
58
src/agent.ts
@ -1,26 +1,13 @@
|
||||
import {GoogleGenerativeAI, SchemaType} from "@google/generative-ai";
|
||||
import dotenv from 'dotenv';
|
||||
import {ProxyAgent, setGlobalDispatcher} from "undici";
|
||||
import {readUrl} from "./tools/read";
|
||||
import { GoogleGenerativeAI, SchemaType } from "@google/generative-ai";
|
||||
import { readUrl } from "./tools/read";
|
||||
import fs from 'fs/promises';
|
||||
import {SafeSearchType, search} from "duck-duck-scrape";
|
||||
import {rewriteQuery} from "./tools/query-rewriter";
|
||||
import {dedupQueries} from "./tools/dedup";
|
||||
import {evaluateAnswer} from "./tools/evaluator";
|
||||
import {StepData} from "./tools/getURLIndex";
|
||||
import {analyzeSteps} from "./tools/error-analyzer";
|
||||
|
||||
// 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();
|
||||
import { SafeSearchType, search } from "duck-duck-scrape";
|
||||
import { rewriteQuery } from "./tools/query-rewriter";
|
||||
import { dedupQueries } from "./tools/dedup";
|
||||
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";
|
||||
|
||||
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 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 genAI = new GoogleGenerativeAI(GEMINI_API_KEY);
|
||||
|
||||
const question = process.argv[2] || "";
|
||||
getResponse(question);
|
||||
getResponse(question);
|
||||
|
||||
22
src/config.ts
Normal file
22
src/config.ts
Normal file
@ -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");
|
||||
@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
@ -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;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@ -1,23 +1,5 @@
|
||||
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";
|
||||
|
||||
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(apiKey);
|
||||
const genAI = new GoogleGenerativeAI(GEMINI_API_KEY);
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
@ -1,23 +1,5 @@
|
||||
import {GoogleGenerativeAI, SchemaType} from "@google/generative-ai";
|
||||
import dotenv from 'dotenv';
|
||||
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");
|
||||
}
|
||||
import { GEMINI_API_KEY, MODEL_NAME } from "../config";
|
||||
|
||||
type KeywordsResponse = {
|
||||
keywords: string[];
|
||||
@ -44,11 +26,9 @@ const responseSchema = {
|
||||
required: ["thought", "keywords"]
|
||||
};
|
||||
|
||||
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",
|
||||
@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user