diff --git a/src/agent.ts b/src/agent.ts index 32dafec..4959883 100644 --- a/src/agent.ts +++ b/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); \ No newline at end of file +getResponse(question); diff --git a/src/config.ts b/src/config.ts new file mode 100644 index 0000000..2cd1754 --- /dev/null +++ b/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"); diff --git a/src/tools/dedup.ts b/src/tools/dedup.ts index 9cc65e0..b333a29 100644 --- a/src/tools/dedup.ts +++ b/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 { +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); -} \ No newline at end of file +} diff --git a/src/tools/error-analyzer.ts b/src/tools/error-analyzer.ts index 82bbc10..aca6fef 100644 --- a/src/tools/error-analyzer.ts +++ b/src/tools/error-analyzer.ts @@ -136,16 +136,17 @@ ${diaryContext.join('\n')} `; } -export async function analyzeSteps(diaryContext: string[]): Promise { +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; } -} \ No newline at end of file +} diff --git a/src/tools/evaluator.ts b/src/tools/evaluator.ts index 08355b2..6b11f16 100644 --- a/src/tools/evaluator.ts +++ b/src/tools/evaluator.ts @@ -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 { +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); -} \ No newline at end of file +} diff --git a/src/tools/query-rewriter.ts b/src/tools/query-rewriter.ts index 2d0519d..8b0028b 100644 --- a/src/tools/query-rewriter.ts +++ b/src/tools/query-rewriter.ts @@ -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 { - - +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); -} \ No newline at end of file +}