Refactor file tree truncation functions to pass logger as parameter

🤖 Generated with Codebuff
Co-Authored-By: Codebuff <noreply@codebuff.com>

Author: charleslien

backend/src/run-agent-step.ts

@@ -183,6 +183,7 @@ export const runAgentStep = async (
     fileContext,
     agentState,
     agentTemplates: localAgentTemplates,
+    logger,
     additionalToolDefinitions: () => {
       const additionalToolDefinitions = cloneDeep(
         Object.fromEntries(
@@ -460,6 +461,7 @@ export const loopAgentSteps = async (
         fileContext,
         agentState,
         agentTemplates: localAgentTemplates,
+        logger,
         additionalToolDefinitions: () => {
           const additionalToolDefinitions = cloneDeep(
             Object.fromEntries(
@@ -489,6 +491,7 @@ export const loopAgentSteps = async (
           fileContext,
           agentState,
           agentTemplates: localAgentTemplates,
+          logger,
           additionalToolDefinitions: () => {
             const additionalToolDefinitions = cloneDeep(
               Object.fromEntries(

backend/src/system-prompt/prompts.ts

@@ -13,6 +13,7 @@ import { schemaToJsonStr } from '@codebuff/common/util/zod-schema'
 
 import { truncateFileTreeBasedOnTokenBudget } from './truncate-file-tree'
 
+import type { Logger } from '@codebuff/types/logger'
 import type { ProjectFileContext } from '@codebuff/common/util/file'
 
 export const configSchemaPrompt = `
@@ -146,16 +147,19 @@ export const additionalSystemPrompts = {
   compact: compactPrompt,
 } as const
 
-export const getProjectFileTreePrompt = (
-  fileContext: ProjectFileContext,
-  fileTreeTokenBudget: number,
-  mode: 'search' | 'agent',
-) => {
+export const getProjectFileTreePrompt = (params: {
+  fileContext: ProjectFileContext
+  fileTreeTokenBudget: number
+  mode: 'search' | 'agent'
+  logger: Logger
+}) => {
+  const { fileContext, fileTreeTokenBudget, mode, logger } = params
   const { projectRoot } = fileContext
-  const { printedTree, truncationLevel } = truncateFileTreeBasedOnTokenBudget(
+  const { printedTree, truncationLevel } = truncateFileTreeBasedOnTokenBudget({
     fileContext,
-    Math.max(0, fileTreeTokenBudget),
-  )
+    tokenBudget: Math.max(0, fileTreeTokenBudget),
+    logger,
+  })
 
   const truncationNote =
     truncationLevel === 'none'

backend/src/system-prompt/search-system-prompt.ts

@@ -6,14 +6,15 @@ import {
   getProjectFileTreePrompt,
   getSystemInfoPrompt,
 } from './prompts'
-import { logger } from '../util/logger'
 import { countTokens, countTokensJson } from '../util/token-counter'
 
+import type { Logger } from '@codebuff/types/logger'
 import type { ProjectFileContext } from '@codebuff/common/util/file'
 
 export function getSearchSystemPrompt(params: {
   fileContext: ProjectFileContext
   messagesTokens: number
+  logger: Logger
   options: {
     agentStepId: string
     clientSessionId: string
@@ -22,7 +23,7 @@ export function getSearchSystemPrompt(params: {
     userId: string | undefined
   }
 }): string {
-  const { fileContext, messagesTokens, options } = params
+  const { fileContext, messagesTokens, logger, options } = params
   const startTime = Date.now()
 
   const maxTokens = 500_000 // costMode === 'lite' ? 64_000 :
@@ -41,17 +42,23 @@ export function getSearchSystemPrompt(params: {
         20_000,
     ) * 20_000
 
-  const projectFileTreePrompt = getProjectFileTreePrompt(
+  const projectFileTreePrompt = getProjectFileTreePrompt({
     fileContext,
     fileTreeTokenBudget,
-    'search',
-  )
+    mode: 'search',
+    logger,
+  })
 
   const t = Date.now()
   const truncationBudgets = [5_000, 20_000, 40_000, 100_000, 500_000]
   const truncatedTrees = truncationBudgets.reduce(
     (acc, budget) => {
-      acc[budget] = getProjectFileTreePrompt(fileContext, budget, 'search')
+      acc[budget] = getProjectFileTreePrompt({
+        fileContext,
+        fileTreeTokenBudget: budget,
+        mode: 'search',
+        logger,
+      })
       return acc
     },
     {} as Record<number, string>,

backend/src/system-prompt/truncate-file-tree.ts

@@ -4,9 +4,10 @@ import {
 } from '@codebuff/common/util/file'
 import { sampleSizeWithSeed } from '@codebuff/common/util/random'
 
-import { logger } from '../util/logger'
 import { countTokens, countTokensJson } from '../util/token-counter'
 
+import type { Logger } from '@codebuff/types/logger'
+
 import type {
   FileTreeNode,
   ProjectFileContext,
@@ -15,14 +16,16 @@ import type {
 type TruncationLevel = 'none' | 'unimportant-files' | 'tokens' | 'depth-based'
 const DEBUG = false
 
-export const truncateFileTreeBasedOnTokenBudget = (
-  fileContext: ProjectFileContext,
-  tokenBudget: number,
-): {
+export const truncateFileTreeBasedOnTokenBudget = (params: {
+  fileContext: ProjectFileContext
+  tokenBudget: number
+  logger: Logger
+}): {
   printedTree: string
   tokenCount: number
   truncationLevel: TruncationLevel
 } => {
+  const { fileContext, tokenBudget, logger } = params
   const startTime = performance.now()
   const { fileTree, fileTokenScores } = fileContext
 
@@ -69,6 +72,7 @@ export const truncateFileTreeBasedOnTokenBudget = (
       fileTree: filteredTree,
       fileTokenScores,
       tokenBudget,
+      logger,
     })
 
     if (tokenCount <= tokenBudget) {
@@ -220,8 +224,9 @@ function pruneFileTokenScores(params: {
   fileTree: FileTreeNode[]
   fileTokenScores: Record<string, Record<string, number>>
   tokenBudget: number
+  logger: Logger
 }) {
-  const { fileTree, fileTokenScores, tokenBudget } = params
+  const { fileTree, fileTokenScores, tokenBudget, logger } = params
   const startTime = performance.now()
 
   // Create sorted array of tokens by score

backend/src/templates/strings.ts

@@ -18,6 +18,7 @@ import {
 } from '../tools/prompts'
 import { parseUserMessage } from '../util/messages'
 
+import type { Logger } from '@codebuff/types/logger'
 import type { AgentTemplate, PlaceholderValue } from './types'
 import type {
   AgentState,
@@ -34,6 +35,7 @@ export async function formatPrompt({
   agentTemplates,
   intitialAgentPrompt,
   additionalToolDefinitions,
+  logger,
 }: {
   prompt: string
   fileContext: ProjectFileContext
@@ -45,6 +47,7 @@ export async function formatPrompt({
   additionalToolDefinitions: () => Promise<
     ProjectFileContext['customToolDefinitions']
   >
+  logger: Logger
 }): Promise<string> {
   const { messageHistory } = agentState
   const lastUserMessage = messageHistory.findLast(
@@ -66,9 +69,19 @@ export async function formatPrompt({
       agentTemplate ? agentTemplate.displayName || 'Unknown Agent' : 'Buffy',
     [PLACEHOLDER.CONFIG_SCHEMA]: () => schemaToJsonStr(CodebuffConfigSchema),
     [PLACEHOLDER.FILE_TREE_PROMPT_SMALL]: () =>
-      getProjectFileTreePrompt(fileContext, 2_500, 'agent'),
+      getProjectFileTreePrompt({
+        fileContext,
+        fileTreeTokenBudget: 2_500,
+        mode: 'agent',
+        logger,
+      }),
     [PLACEHOLDER.FILE_TREE_PROMPT]: () =>
-      getProjectFileTreePrompt(fileContext, 10_000, 'agent'),
+      getProjectFileTreePrompt({
+        fileContext,
+        fileTreeTokenBudget: 10_000,
+        mode: 'agent',
+        logger,
+      }),
     [PLACEHOLDER.GIT_CHANGES_PROMPT]: () => getGitChangesPrompt(fileContext),
     [PLACEHOLDER.REMAINING_STEPS]: () => `${agentState.stepsRemaining!}`,
     [PLACEHOLDER.PROJECT_ROOT]: () => fileContext.projectRoot,
@@ -142,6 +155,7 @@ export async function getAgentPrompt<T extends StringField>({
   agentState,
   agentTemplates,
   additionalToolDefinitions,
+  logger,
 }: {
   agentTemplate: AgentTemplate
   promptType: { type: T }
@@ -151,6 +165,7 @@ export async function getAgentPrompt<T extends StringField>({
   additionalToolDefinitions: () => Promise<
     ProjectFileContext['customToolDefinitions']
   >
+  logger: Logger
 }): Promise<string | undefined> {
   let promptValue = agentTemplate[promptType.type]
   for (const placeholder of additionalPlaceholders[promptType.type]) {
@@ -171,6 +186,7 @@ export async function getAgentPrompt<T extends StringField>({
     spawnableAgents: agentTemplate.spawnableAgents,
     agentTemplates,
     additionalToolDefinitions,
+    logger,
   })
 
   let addendum = ''

backend/src/tools/handlers/tool/find-files.ts

@@ -71,6 +71,7 @@ export const handleFindFiles = ((params: {
   const system = getSearchSystemPrompt({
     fileContext,
     messagesTokens: fileRequestMessagesTokens,
+    logger,
     options: {
       agentStepId,
       clientSessionId,