pass in logger to bigquery client

Author: charleslien

backend/src/admin/relabelRuns.ts

@@ -196,7 +196,7 @@ export async function relabelForUserHandler(req: Request, res: Response) {
             }
 
             // Store the relabel
-            await insertRelabel(relabel)
+            await insertRelabel({ relabel, logger })
             logger.info(`Successfully stored relabel for trace ${trace.id}`)
 
             return {
@@ -352,7 +352,7 @@ async function relabelWithRelace(
     },
   }
 
-  await insertRelabel(relabel)
+  await insertRelabel({ relabel, logger })
 
   return relaced
 }
@@ -364,7 +364,7 @@ export async function relabelWithClaudeWithFullFileContext(
   dataset?: string,
 ) {
   if (dataset) {
-    await setupBigQuery(dataset)
+    await setupBigQuery({ dataset, logger })
   }
   logger.info(`Relabeling ${trace.id} with Claude with full file context`)
   const filesWithPath = Object.entries(fileBlobs.payload.files).map(
@@ -418,7 +418,7 @@ export async function relabelWithClaudeWithFullFileContext(
     },
   } as Relabel
 
-  await insertRelabel(relabel, dataset)
+  await insertRelabel({ relabel, dataset, logger })
 
   return relabel
 }

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

@@ -373,7 +373,7 @@ async function getRelevantFiles(
     },
   }
 
-  insertTrace(trace).catch((error: Error) => {
+  insertTrace({ trace, logger }).catch((error: Error) => {
     logger.error({ error }, 'Failed to insert trace')
   })
 
@@ -443,7 +443,7 @@ async function getRelevantFilesForTraining(
     },
   }
 
-  insertTrace(trace).catch((error: Error) => {
+  insertTrace({ trace, logger }).catch((error: Error) => {
     logger.error({ error }, 'Failed to insert trace')
   })
 

backend/src/index.ts

@@ -72,7 +72,7 @@ app.use(
 )
 
 // Initialize BigQuery before starting the server
-setupBigQuery().catch((err) => {
+setupBigQuery({ logger }).catch((err) => {
   logger.error(
     {
       error: err,

backend/src/run-agent-step.ts

@@ -322,7 +322,7 @@ export const runAgentStep = async (
     },
   }
 
-  insertTrace(agentResponseTrace)
+  insertTrace({ trace: agentResponseTrace, logger })
 
   const newAgentContext = state.agentContext as AgentState['agentContext']
   // Use the updated agent state from tool execution

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

@@ -70,7 +70,7 @@ export function getSearchSystemPrompt(
     },
   }
 
-  insertTrace(trace).catch((error: Error) => {
+  insertTrace({ trace, logger }).catch((error: Error) => {
     logger.error({ error }, 'Failed to insert file trees trace')
   })
   const fileTreeTokens = countTokensJson(projectFileTreePrompt)

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

@@ -217,5 +217,8 @@ async function uploadExpandedFileContextForTraining(
   }
 
   // Upload the files to bigquery
-  await insertTrace(trace)
+  await insertTrace({
+    trace,
+    logger,
+  })
 }

packages/bigquery/package.json

@@ -27,7 +27,8 @@
   },
   "dependencies": {
     "@google-cloud/bigquery": "^7.9.4",
-    "@codebuff/common": "workspace:*"
+    "@codebuff/common": "workspace:*",
+    "@codebuff/types": "workspace:*"
   },
   "devDependencies": {
     "@types/node": "22",

packages/bigquery/src/client.ts

@@ -1,5 +1,4 @@
 import { getErrorObject } from '@codebuff/common/util/error'
-import { logger } from '@codebuff/common/util/logger'
 import { BigQuery } from '@google-cloud/bigquery'
 
 import { MESSAGE_SCHEMA, RELABELS_SCHEMA, TRACES_SCHEMA } from './schema'
@@ -11,6 +10,7 @@ import type {
   Relabel,
   Trace,
 } from './schema'
+import type { Logger } from '@codebuff/types/logger'
 
 const DATASET =
   process.env.NEXT_PUBLIC_CB_ENVIRONMENT === 'prod'
@@ -33,7 +33,13 @@ function getClient(): BigQuery {
   return client
 }
 
-export async function setupBigQuery(dataset: string = DATASET) {
+export async function setupBigQuery({
+  dataset = DATASET,
+  logger,
+}: {
+  dataset?: string
+  logger: Logger
+}) {
   if (client) {
     return
   }
@@ -93,10 +99,15 @@ export async function setupBigQuery(dataset: string = DATASET) {
   }
 }
 
-export async function insertMessage(
-  row: MessageRow,
-  dataset: string = DATASET,
-) {
+export async function insertMessage({
+  row,
+  dataset = DATASET,
+  logger,
+}: {
+  row: MessageRow
+  dataset?: string
+  logger: Logger
+}) {
   try {
     await getClient()
       .dataset(dataset)
@@ -121,7 +132,15 @@ export async function insertMessage(
   }
 }
 
-export async function insertTrace(trace: Trace, dataset: string = DATASET) {
+export async function insertTrace({
+  trace,
+  dataset = DATASET,
+  logger,
+}: {
+  trace: Trace
+  dataset?: string
+  logger: Logger
+}) {
   try {
     // Create a copy of the trace and stringify payload if needed
     const traceToInsert = {
@@ -148,10 +167,15 @@ export async function insertTrace(trace: Trace, dataset: string = DATASET) {
   }
 }
 
-export async function insertRelabel(
-  relabel: Relabel,
-  dataset: string = DATASET,
-) {
+export async function insertRelabel({
+  relabel,
+  dataset = DATASET,
+  logger,
+}: {
+  relabel: Relabel
+  dataset?: string
+  logger: Logger
+}) {
   try {
     // Stringify payload if needed
     const relabelToInsert = {

scripts/ft-file-selection/collect-tuning-data.ts

@@ -416,7 +416,7 @@ async function main() {
       process.exit(1)
     }
 
-    await setupBigQuery(DATASET)
+    await setupBigQuery({ dataset: DATASET, logger: console })
     console.log(`Using dataset: ${DATASET}`)
 
     // Get traces for the specified model from BigQuery

scripts/ft-file-selection/grade-traces.ts

@@ -1,17 +1,18 @@
 import { getTracesAndRelabelsForUser, setupBigQuery } from '@codebuff/bigquery'
-import { logger } from '@codebuff/common/util/logger'
 
 import { gradeRun } from '../../backend/src/admin/grade-runs'
 
+import type { Logger } from '@codebuff/types/logger'
+
 // Parse command line arguments to check for --prod flag
 const isProd = process.argv.includes('--prod')
 const DATASET = isProd ? 'codebuff_data' : 'codebuff_data_dev'
 const MAX_PARALLEL = 1 // Maximum number of traces to process in parallel
 
-async function gradeTraces() {
+async function gradeTraces({ logger }: { logger: Logger }) {
   try {
     // Initialize BigQuery
-    await setupBigQuery(DATASET)
+    await setupBigQuery({ dataset: DATASET, logger: console })
 
     console.log(`\nGrading traces from dataset: ${DATASET}`)
 
@@ -53,6 +54,7 @@ async function gradeTraces() {
             }
           } catch (error) {
             logger.error(
+              {},
               `Error grading trace ${traceAndRelabels.trace.id}:`,
               error,
             )
@@ -87,4 +89,4 @@ async function gradeTraces() {
 }
 
 // Run the script
-gradeTraces()
+gradeTraces({ logger: console })