From b4fe6ae0db581557001fd576867027c14c29ca25 Mon Sep 17 00:00:00 2001 From: Han Xiao Date: Tue, 4 Feb 2025 21:47:27 +0800 Subject: [PATCH] chore: first commit --- src/agent.ts | 99 +++++++++++++++++++++++++++++++++++---------------- src/config.ts | 2 ++ 2 files changed, 71 insertions(+), 30 deletions(-) diff --git a/src/agent.ts b/src/agent.ts index f289737..a9cd40c 100644 --- a/src/agent.ts +++ b/src/agent.ts @@ -7,7 +7,14 @@ import {rewriteQuery} from "./tools/query-rewriter"; import {dedupQueries} from "./tools/dedup"; import {evaluateAnswer} from "./tools/evaluator"; import {analyzeSteps} from "./tools/error-analyzer"; -import {GEMINI_API_KEY, JINA_API_KEY, SEARCH_PROVIDER, STEP_SLEEP, modelConfigs} from "./config"; +import { + GEMINI_API_KEY, + JINA_API_KEY, + SEARCH_PROVIDER, + STEP_SLEEP, + modelConfigs, + DEFAULT_BUDGET_SPLIT_RATIO +} from "./config"; import {TokenTracker} from "./utils/token-tracker"; import {ActionTracker} from "./utils/action-tracker"; import {StepAction, SchemaProperty, ResponseSchema, AnswerAction} from "./types"; @@ -197,7 +204,7 @@ ${urlList} } if (beastMode) { - actions.push(`**answer**: + actions.push(`**answer**: - You have gathered enough information to answer the question; they may not be perfect, but this is your very last chance to answer the question. - Try the best of the best reasoning ability, investigate every details in the context and provide the best answer you can think of. - When uncertain, educated guess is allowed and encouraged, but make sure it is based on the context and knowledge you have gathered. @@ -245,19 +252,22 @@ function removeAllLineBreaks(text: string) { export async function getResponse(question: string, tokenBudget: number = 1_000_000, maxBadAttempts: number = 3, - existingContext?: Partial): Promise<{ result: StepAction; context: TrackerContext }> { + existingContext?: Partial, + parentKnowledge: Array<{ question: string, answer: string, type: string }> = [], + budgetSplitRatio: number = DEFAULT_BUDGET_SPLIT_RATIO, + recursionLevel: number = 0): Promise<{ result: StepAction; context: TrackerContext }> { const context: TrackerContext = { tokenTracker: existingContext?.tokenTracker || new TokenTracker(tokenBudget), actionTracker: existingContext?.actionTracker || new ActionTracker() }; - context.actionTracker.trackAction({ gaps: [question], totalStep: 0, badAttempts: 0 }); + context.actionTracker.trackAction({gaps: [question], totalStep: 0, badAttempts: 0}); let step = 0; let totalStep = 0; let badAttempts = 0; 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 allKnowledge = [...parentKnowledge]; // knowledge are intermedidate questions that are answered const badContext = []; let diaryContext = []; let allowAnswer = true; @@ -275,9 +285,9 @@ export async function getResponse(question: string, tokenBudget: number = 1_000_ await sleep(STEP_SLEEP); step++; totalStep++; - context.actionTracker.trackAction({ totalStep, thisStep, gaps, badAttempts }); + context.actionTracker.trackAction({totalStep, thisStep, gaps, badAttempts}); const budgetPercentage = (context.tokenTracker.getTotalUsage() / tokenBudget * 100).toFixed(2); - console.log(`Step ${totalStep} / Budget used ${budgetPercentage}%`); + console.log(`| Recursion ${recursionLevel} | Step ${totalStep} | Budget used ${budgetPercentage}% |`); console.log('Gaps:', gaps); allowReflect = allowReflect && (gaps.length <= 1); const currentQuestion = gaps.length > 0 ? gaps.shift()! : question; @@ -298,7 +308,7 @@ export async function getResponse(question: string, tokenBudget: number = 1_000_ allKnowledge, allURLs, false - ); + ); const model = genAI.getGenerativeModel({ model: modelConfigs.agent.model, @@ -456,17 +466,43 @@ Although you solved a sub-question, you still need to find the answer to the ori newGapQuestions = (await dedupQueries(newGapQuestions, allQuestions)).unique_queries; } if (newGapQuestions.length > 0) { - // found new gap questions - diaryContext.push(` -At step ${step}, you took **reflect** and think about the knowledge gaps. You found some sub-questions are important to the question: "${currentQuestion}" -You realize you need to know the answers to the following sub-questions: -${newGapQuestions.map((q: string) => `- ${q}`).join('\n')} + for (const gapQuestion of newGapQuestions) { + // Create sub-context with remaining token budget + // Calculate sub-question budget as a fraction of current remaining budget + const currentRemaining = tokenBudget - context.tokenTracker.getTotalUsage(); + const subQuestionBudget = Math.floor(currentRemaining * budgetSplitRatio); + // Log the budget allocation for monitoring + console.log(`Allocating ${subQuestionBudget} tokens for sub-question: "${gapQuestion}"`) -You will now figure out the answers to these sub-questions and see if they can help me find the answer to the original question. + const subContext = { + tokenTracker: context.tokenTracker, // Share token tracker + actionTracker: context.actionTracker // Share action tracker + }; + + // Recursively resolve gap question + const {result: gapResult} = await getResponse( + gapQuestion, + subQuestionBudget, + maxBadAttempts, + subContext, + allKnowledge, // Pass current knowledge to sub-question + recursionLevel + 1 + ); + + // If gap was successfully answered, add to knowledge + if (gapResult.action === 'answer') { + allKnowledge.push({ + question: gapQuestion, + answer: gapResult.answer, + type: 'qa' + }); + } + } + diaryContext.push(` +At step ${step}, you took **reflect** and resolved ${newGapQuestions.length} sub-questions: +${newGapQuestions.map(q => `- ${q}`).join('\n')} +The answers to these questions have been added to your knowledge base. `); - gaps.push(...newGapQuestions); - allQuestions.push(...newGapQuestions); - gaps.push(question); // always keep the original question in the gaps } else { 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(', ')} @@ -603,13 +639,13 @@ You decided to think out of the box or cut from a completely different angle.`); } } - await storeContext(prompt, [allContext, allKeywords, allQuestions, allKnowledge], totalStep); + await storeContext(prompt, [allContext, allKeywords, allQuestions, allKnowledge], totalStep, recursionLevel); } step++; totalStep++; - await storeContext(prompt, [allContext, allKeywords, allQuestions, allKnowledge], totalStep); + await storeContext(prompt, [allContext, allKeywords, allQuestions, allKnowledge], totalStep, recursionLevel); if (isAnswered) { - return { result: thisStep, context }; + return {result: thisStep, context}; } else { console.log('Enter Beast mode!!!') const prompt = getPrompt( @@ -624,7 +660,7 @@ You decided to think out of the box or cut from a completely different angle.`); allKnowledge, allURLs, true - ); + ); const model = genAI.getGenerativeModel({ model: modelConfigs.agentBeastMode.model, @@ -647,21 +683,21 @@ You decided to think out of the box or cut from a completely different angle.`); const usage = response.usageMetadata; context.tokenTracker.trackUsage('agent', usage?.totalTokenCount || 0); - await storeContext(prompt, [allContext, allKeywords, allQuestions, allKnowledge], totalStep); + await storeContext(prompt, [allContext, allKeywords, allQuestions, allKnowledge], totalStep, recursionLevel); thisStep = JSON.parse(response.text()); console.log(thisStep) - return { result: thisStep, context }; + return {result: thisStep, context}; } } -async function storeContext(prompt: string, memory: any[][], step: number) { +async function storeContext(prompt: string, memory: any[][], step: number, recursionLevel: number) { try { - await fs.writeFile(`prompt-${step}.txt`, prompt); + await fs.writeFile(`prompt-${recursionLevel}-${step}.txt`, prompt); const [context, keywords, questions, knowledge] = memory; - await fs.writeFile('context.json', JSON.stringify(context, null, 2)); - await fs.writeFile('queries.json', JSON.stringify(keywords, null, 2)); - await fs.writeFile('questions.json', JSON.stringify(questions, null, 2)); - await fs.writeFile('knowledge.json', JSON.stringify(knowledge, null, 2)); + await fs.writeFile(`context-${recursionLevel}.json`, JSON.stringify(context, null, 2)); + await fs.writeFile(`queries-${recursionLevel}.json`, JSON.stringify(keywords, null, 2)); + await fs.writeFile(`questions-${recursionLevel}.json`, JSON.stringify(questions, null, 2)); + await fs.writeFile(`knowledge-${recursionLevel}.json`, JSON.stringify(knowledge, null, 2)); } catch (error) { console.error('Context storage failed:', error); } @@ -672,7 +708,10 @@ const genAI = new GoogleGenerativeAI(GEMINI_API_KEY); export async function main() { const question = process.argv[2] || ""; - const { result: finalStep, context: tracker } = await getResponse(question) as { result: AnswerAction; context: TrackerContext }; + const { + result: finalStep, + context: tracker + } = await getResponse(question) as { result: AnswerAction; context: TrackerContext }; console.log('Final Answer:', finalStep.answer); tracker.tokenTracker.printSummary(); diff --git a/src/config.ts b/src/config.ts index 5f9ba4f..56e558d 100644 --- a/src/config.ts +++ b/src/config.ts @@ -67,6 +67,8 @@ export const modelConfigs: ToolConfigs = { }; export const STEP_SLEEP = 1000; +export const DEFAULT_BUDGET_SPLIT_RATIO = 0.3; // Each sub-question gets 30% of parent's budget + if (!GEMINI_API_KEY) throw new Error("GEMINI_API_KEY not found"); if (!JINA_API_KEY) throw new Error("JINA_API_KEY not found");