pass in logger to relace api

Author: charleslien

backend/src/admin/relabelRuns.ts

@@ -16,8 +16,6 @@ import { rerank } from '../llm-apis/relace-api'
 import { promptAiSdk } from '../llm-apis/vercel-ai-sdk/ai-sdk'
 import { messagesWithSystem } from '../util/messages'
 
-import type { Logger } from '@codebuff/types/logger'
-
 import type { System } from '../llm-apis/claude'
 import type {
   GetExpandedFileContextForTrainingBlobTrace,
@@ -27,6 +25,7 @@ import type {
   Relabel,
 } from '@codebuff/bigquery'
 import type { Message } from '@codebuff/common/types/messages/codebuff-message'
+import type { Logger } from '@codebuff/types/logger'
 import type { Request, Response } from 'express'
 
 // --- GET Handler Logic ---
@@ -151,7 +150,11 @@ export async function relabelForUserHandler(params: {
 
     const allResults = []
 
-    const relaceResults = relabelUsingFullFilesForUser({ userId, limit, logger })
+    const relaceResults = relabelUsingFullFilesForUser({
+      userId,
+      limit,
+      logger,
+    })
 
     // Process each model
     for (const model of modelsToRelabel) {
@@ -303,7 +306,12 @@ async function relabelUsingFullFilesForUser(params: {
         )
       ) {
         relabelPromises.push(
-          relabelWithClaudeWithFullFileContext({ trace, fileBlobs, model, logger }),
+          relabelWithClaudeWithFullFileContext({
+            trace,
+            fileBlobs,
+            model,
+            logger,
+          }),
         )
         didRelabel = true
       }
@@ -348,12 +356,15 @@ async function relabelWithRelace(params: {
     }),
   )
 
-  const relaced = await rerank(filesWithPath, query, {
+  const relaced = await rerank({
+    files: filesWithPath,
+    prompt: query,
     clientSessionId: trace.payload.client_session_id,
     fingerprintId: trace.payload.fingerprint_id,
     userInputId: trace.payload.user_input_id,
     userId: 'test-user-id', // Make sure we don't bill em for it!!
     messageId: trace.id,
+    logger,
   })
 
   const relabel = {
@@ -415,7 +426,10 @@ export async function relabelWithClaudeWithFullFileContext(params: {
   }
 
   const output = await promptAiSdk({
-    messages: messagesWithSystem({ messages: trace.payload.messages as Message[], system }),
+    messages: messagesWithSystem({
+      messages: trace.payload.messages as Message[],
+      system,
+    }),
     model: model as any, // Model type is string here for flexibility
     clientSessionId: 'relabel-trace-api',
     fingerprintId: 'relabel-trace-api',

backend/src/fast-rewrite.ts

@@ -40,19 +40,18 @@ export async function fastRewrite(params: {
   } = params
   const relaceStartTime = Date.now()
   const messageId = generateCompactId('cb-')
-  let response = await promptRelaceAI(
-    initialContent,
+  let response = await promptRelaceAI({
+    initialCode: initialContent,
     editSnippet,
     instructions,
-    {
-      clientSessionId,
-      fingerprintId,
-      userInputId,
-      userId,
-      userMessage,
-      messageId,
-    },
-  )
+    clientSessionId,
+    fingerprintId,
+    userInputId,
+    userId,
+    userMessage,
+    messageId,
+    logger,
+  })
   const relaceDuration = Date.now() - relaceStartTime
 
   // Check if response still contains lazy edits

backend/src/llm-apis/relace-api.ts

@@ -6,36 +6,40 @@ import {
 import { env } from '@codebuff/internal'
 
 import { saveMessage } from '../llm-apis/message-cost-tracker'
-import { logger } from '../util/logger'
 import { countTokens } from '../util/token-counter'
 import { promptAiSdk } from './vercel-ai-sdk/ai-sdk'
 
+import type { Logger } from '@codebuff/types/logger'
+
 const timeoutPromise = (ms: number) =>
   new Promise((_, reject) =>
     setTimeout(() => reject(new Error('Relace API request timed out')), ms),
   )
 
-export async function promptRelaceAI(
-  initialCode: string,
-  editSnippet: string,
-  instructions: string | undefined,
-  options: {
-    clientSessionId: string
-    fingerprintId: string
-    userInputId: string
-    userId: string | undefined
-    messageId: string
-    userMessage?: string
-  },
-) {
+export async function promptRelaceAI(params: {
+  initialCode: string
+  editSnippet: string
+  instructions: string | undefined
+  clientSessionId: string
+  fingerprintId: string
+  userInputId: string
+  userId: string | undefined
+  messageId: string
+  userMessage?: string
+  logger: Logger
+}) {
   const {
+    initialCode,
+    editSnippet,
+    instructions,
     clientSessionId,
     fingerprintId,
     userInputId,
     userId,
     userMessage,
     messageId,
-  } = options
+    logger,
+  } = params
   const startTime = Date.now()
 
   try {
@@ -151,19 +155,26 @@ export type FileWithPath = {
   content: string
 }
 
-export async function rerank(
-  files: FileWithPath[],
-  prompt: string,
-  options: {
-    clientSessionId: string
-    fingerprintId: string
-    userInputId: string
-    userId: string | undefined
-    messageId: string
-  },
-) {
-  const { clientSessionId, fingerprintId, userInputId, userId, messageId } =
-    options
+export async function rerank(params: {
+  files: FileWithPath[]
+  prompt: string
+  clientSessionId: string
+  fingerprintId: string
+  userInputId: string
+  userId: string | undefined
+  messageId: string
+  logger: Logger
+}) {
+  const {
+    files,
+    prompt,
+    clientSessionId,
+    fingerprintId,
+    userInputId,
+    userId,
+    messageId,
+    logger,
+  } = params
   const startTime = Date.now()
 
   if (!prompt || !files.length) {