pass in logger to gemini with fallbacks

Author: charleslien

backend/src/__tests__/read-docs-tool.test.ts

@@ -314,12 +314,11 @@ describe('read_docs tool with researcher agent', () => {
       logger,
     })
 
-    expect(context7Api.fetchContext7LibraryDocumentation).toHaveBeenCalledWith(
-      'React',
-      {
-        topic: 'hooks',
-      },
-    )
+    expect(context7Api.fetchContext7LibraryDocumentation).toHaveBeenCalledWith({
+      query: 'React',
+      topic: 'hooks',
+      logger: expect.anything(),
+    })
 
     // Check that the documentation was added to the message history
     const toolResultMessages = newAgentState.messageHistory.filter(
@@ -388,13 +387,12 @@ describe('read_docs tool with researcher agent', () => {
       logger,
     })
 
-    expect(context7Api.fetchContext7LibraryDocumentation).toHaveBeenCalledWith(
-      'React',
-      {
-        topic: 'hooks',
-        tokens: 5000,
-      },
-    )
+    expect(context7Api.fetchContext7LibraryDocumentation).toHaveBeenCalledWith({
+      query: 'React',
+      topic: 'hooks',
+      tokens: 5000,
+      logger: expect.anything(),
+    })
   }, 10000)
 
   test('should handle case when no documentation is found', async () => {

backend/src/__tests__/request-files-prompt.test.ts

@@ -228,7 +228,6 @@ describe('requestRelevantFiles', () => {
     expect(
       geminiWithFallbacksModule.promptFlashWithFallbacks,
     ).toHaveBeenCalledWith(
-      expect.anything(),
       expect.objectContaining({
         useFinetunedModel: finetunedVertexModels.ft_filepicker_010,
       }),
@@ -259,7 +258,6 @@ describe('requestRelevantFiles', () => {
     expect(
       geminiWithFallbacksModule.promptFlashWithFallbacks,
     ).toHaveBeenCalledWith(
-      expect.anything(),
       expect.objectContaining({
         useFinetunedModel: expectedModel,
       }),

backend/src/fast-rewrite.ts

@@ -7,12 +7,12 @@ import { promptFlashWithFallbacks } from './llm-apis/gemini-with-fallbacks'
 import { promptRelaceAI } from './llm-apis/relace-api'
 import { promptAiSdk } from './llm-apis/vercel-ai-sdk/ai-sdk'
 
-import type { Logger } from '@codebuff/types/logger'
 import type { CodebuffToolMessage } from '@codebuff/common/tools/list'
 import type {
   Message,
   ToolMessage,
 } from '@codebuff/common/types/messages/codebuff-message'
+import type { Logger } from '@codebuff/types/logger'
 
 export async function fastRewrite(params: {
   initialContent: string
@@ -243,12 +243,14 @@ Do not write anything else.
       content: prompt,
     },
   )
-  const response = await promptFlashWithFallbacks(messages, {
+  const response = await promptFlashWithFallbacks({
+    messages,
     clientSessionId,
     fingerprintId,
     userInputId,
     model: models.openrouter_gemini2_5_flash,
     userId,
+    logger,
   })
   const shouldAddPlaceholderComments = response.includes('LOCAL_CHANGE_ONLY')
   logger.debug(

backend/src/find-files/request-files-prompt.ts

@@ -15,7 +15,6 @@ import { range, shuffle, uniq } from 'lodash'
 import { CustomFilePickerConfigSchema } from './custom-file-picker-config'
 import { promptFlashWithFallbacks } from '../llm-apis/gemini-with-fallbacks'
 import { promptAiSdk } from '../llm-apis/vercel-ai-sdk/ai-sdk'
-import type { Logger } from '@codebuff/types/logger'
 import {
   castAssistantMessage,
   messagesWithSystem,
@@ -31,6 +30,7 @@ import type {
 } from '@codebuff/bigquery'
 import type { Message } from '@codebuff/common/types/messages/codebuff-message'
 import type { ProjectFileContext } from '@codebuff/common/util/file'
+import type { Logger } from '@codebuff/types/logger'
 
 const NUMBER_OF_EXAMPLE_FILES = 100
 const MAX_FILES_PER_REQUEST = 30
@@ -356,7 +356,10 @@ async function getRelevantFiles(params: {
     logger,
   })
   const start = performance.now()
-  let codebuffMessages = messagesWithSystem({ messages: messagesWithPrompt, system })
+  let codebuffMessages = messagesWithSystem({
+    messages: messagesWithPrompt,
+    system,
+  })
 
   // Converts assistant messages to user messages for finetuned model
   codebuffMessages = codebuffMessages
@@ -370,13 +373,15 @@ async function getRelevantFiles(params: {
     .filter((msg) => msg !== null)
   const finetunedModel = modelId ?? finetunedVertexModels.ft_filepicker_010
 
-  let response = await promptFlashWithFallbacks(codebuffMessages, {
+  let response = await promptFlashWithFallbacks({
+    messages: codebuffMessages,
     clientSessionId,
     userInputId,
     model: models.openrouter_gemini2_5_flash,
     userId,
     useFinetunedModel: finetunedModel,
     fingerprintId,
+    logger,
   })
   const end = performance.now()
   const duration = end - start
@@ -528,7 +533,10 @@ function generateNonObviousRequestFilesPrompt(
   fileContext: ProjectFileContext,
   count: number,
 ): string {
-  const exampleFiles = getExampleFileList({ fileContext, count: NUMBER_OF_EXAMPLE_FILES })
+  const exampleFiles = getExampleFileList({
+    fileContext,
+    count: NUMBER_OF_EXAMPLE_FILES,
+  })
   return `
 Your task is to find the second-order relevant files for the following user request (in quotes).
 
@@ -583,7 +591,10 @@ function generateKeyRequestFilesPrompt(
   fileContext: ProjectFileContext,
   count: number,
 ): string {
-  const exampleFiles = getExampleFileList({ fileContext, count: NUMBER_OF_EXAMPLE_FILES })
+  const exampleFiles = getExampleFileList({
+    fileContext,
+    count: NUMBER_OF_EXAMPLE_FILES,
+  })
 
   return `
 Your task is to find the most relevant files for the following user request (in quotes).

backend/src/llm-apis/gemini-with-fallbacks.ts

@@ -1,6 +1,5 @@
 import { openaiModels, openrouterModels } from '@codebuff/common/old-constants'
 
-import { logger } from '../util/logger'
 import { promptAiSdk } from './vercel-ai-sdk/ai-sdk'
 
 import type {
@@ -9,6 +8,7 @@ import type {
   Model,
 } from '@codebuff/common/old-constants'
 import type { Message } from '@codebuff/common/types/messages/codebuff-message'
+import type { Logger } from '@codebuff/types/logger'
 
 /**
  * Prompts a Gemini model with fallback logic.
@@ -35,28 +35,29 @@ import type { Message } from '@codebuff/common/types/messages/codebuff-message'
  * @returns A promise that resolves to the complete response string from the successful API call.
  * @throws If all API calls (primary and fallbacks) fail.
  */
-export async function promptFlashWithFallbacks(
-  messages: Message[],
-  options: {
-    clientSessionId: string
-    fingerprintId: string
-    userInputId: string
-    model: Model
-    userId: string | undefined
-    maxTokens?: number
-    temperature?: number
-    costMode?: CostMode
-    useGPT4oInsteadOfClaude?: boolean
-    thinkingBudget?: number
-    useFinetunedModel?: FinetunedVertexModel | undefined
-  },
-): Promise<string> {
+export async function promptFlashWithFallbacks(params: {
+  messages: Message[]
+  clientSessionId: string
+  fingerprintId: string
+  userInputId: string
+  model: Model
+  userId: string | undefined
+  maxTokens?: number
+  temperature?: number
+  costMode?: CostMode
+  useGPT4oInsteadOfClaude?: boolean
+  thinkingBudget?: number
+  useFinetunedModel?: FinetunedVertexModel | undefined
+  logger: Logger
+}): Promise<string> {
   const {
+    messages,
     costMode,
     useGPT4oInsteadOfClaude,
     useFinetunedModel,
+    logger,
     ...geminiOptions
-  } = options
+  } = params
 
   // Try finetuned model first if enabled
   if (useFinetunedModel) {