Function getTracesWithRelabels

packages/bigquery/src/client.ts:357–412

(
  model: string,
  limit: number = 100,
  dataset: string = DATASET,
)

}

export async function getTracesWithRelabels(
  model: string,
  limit: number = 100,
  dataset: string = DATASET,
) {
  // Get traces that DO have matching relabels for the specified model
  const query = `
  SELECT
    ANY_VALUE(t) as trace,
    ARRAY_AGG(r ORDER BY r.created_at DESC LIMIT 1)[OFFSET(0)] as relabel
  FROM \`${dataset}.${TRACES_TABLE}\` t
  INNER JOIN (
    SELECT *
    FROM \`${dataset}.${RELABELS_TABLE}\` r
    WHERE r.model = '${model}'
  ) r
  ON t.agent_step_id = r.agent_step_id
     AND t.user_id = r.user_id
     AND JSON_EXTRACT_SCALAR(t.payload, '$.user_input_id') = JSON_EXTRACT_SCALAR(r.payload, '$.user_input_id')
  WHERE t.type = 'get-relevant-files'
    AND JSON_EXTRACT_SCALAR(t.payload, '$.output') IS NOT NULL
    AND JSON_EXTRACT_SCALAR(r.payload, '$.output') IS NOT NULL
  GROUP BY t.agent_step_id
  ORDER BY MAX(t.created_at) DESC
  LIMIT ${limit}
  `

  const [rows] = await getClient().query(query)

  // Filter out any results where either trace or relabel data is missing
  const res = rows
    .filter((row) => row.trace && row.relabel)
    .map((row) => ({
      trace: row.trace as GetRelevantFilesTrace,
      relabel: row.relabel as Relabel,
    }))

  // Parse the payload as JSON if it's a string
  return res.map((row) => ({
    ...row,
    trace: {
      ...row.trace,
      payload:
        typeof row.trace.payload === 'string'
          ? JSON.parse(row.trace.payload)
          : row.trace.payload,
    },
    relabel: {
      ...row.relabel,
      payload:
        typeof row.relabel.payload === 'string'
          ? JSON.parse(row.relabel.payload)
          : row.relabel.payload,
    },
  })) as { trace: GetRelevantFilesTrace; relabel: Relabel }[]
}

export async function getTracesAndRelabelsForUser(