Function formatTraceResults

web/src/app/api/admin/relabel-for-user/route.ts:471–511

(traceBundles: TraceBundle[])

}

function formatTraceResults(traceBundles: TraceBundle[]) {
  return traceBundles.map(({ trace, relatedTraces, relabels }) => {
    const payload =
      typeof trace.payload === 'string'
        ? (JSON.parse(trace.payload) as GetRelevantFilesPayload)
        : (trace.payload as GetRelevantFilesPayload)

    const timestamp =
      trace.created_at instanceof Date
        ? trace.created_at.toISOString()
        : new Date(
            (trace.created_at as BigQueryTimestamp)?.value ?? trace.created_at,
          ).toISOString()

    const query = extractQueryFromMessages(payload.messages)
    const outputs: Record<string, string> = {
      base: payload.output || '',
    }

    relabels.forEach((relabel) => {
      if (relabel.model && relabel.payload?.output) {
        outputs[relabel.model] = relabel.payload.output
      }
    })

    const expandedFilesTrace = relatedTraces.find(
      (t) => t.type === 'get-expanded-file-context-for-training',
    )
    if (expandedFilesTrace?.payload) {
      outputs['files-uploaded'] = (
        expandedFilesTrace.payload as GetRelevantFilesPayload
      ).output
    }

    return {
      timestamp,
      query,
      outputs,
    }
  })
}

function extractQueryFromMessages(messages: unknown): string {
  const items = Array.isArray(messages) ? messages : []