codeflash-internal/js/cf-webapp/src/app/observability/llm-export/route.ts

import { type NextRequest, NextResponse } from "next/server"
import { auth0 } from "@/lib/auth0"
import { getTraceData, type TraceData } from "@/app/observability/lib/get-trace-data"
import { transformToTimelineSections } from "@/app/observability/components/timeline-types"
import { formatTimelineForLLM } from "@/app/observability/components/format-llm-export"

function getFunctionName(
  metadata: Record<string, unknown> | undefined,
  optimizationEvent: TraceData["optimizationEvent"],
): string | null {
  const metadataFunctionName = metadata?.function_to_optimize as string | undefined
  return metadataFunctionName ?? optimizationEvent?.function_name ?? null
}

function getFilePath(
  optimizationEvent: TraceData["optimizationEvent"],
  metadata: Record<string, unknown> | undefined,
): string | null {
  const eventFilePath = optimizationEvent?.file_path
  const metadataFilePath = metadata?.file_path as string | undefined
  return eventFilePath ?? metadataFilePath ?? null
}

export async function GET(request: NextRequest) {
  const session = await auth0.getSession()
  if (!session?.user?.sub) {
    return new NextResponse("Unauthorized", { status: 401 })
  }

  const traceId = request.nextUrl.searchParams.get("trace_id")?.trim()

  if (!traceId) {
    return new NextResponse("No trace ID provided.", {
      status: 400,
      headers: { "Content-Type": "text/plain; charset=utf-8" },
    })
  }

  const tracePrefix = traceId.substring(0, 33)
  const traceData = await getTraceData(tracePrefix)

  if (traceData.rawLlmCalls.length === 0 && traceData.errors.length === 0) {
    return new NextResponse(`No data found for trace: ${traceId}`, {
      status: 404,
      headers: { "Content-Type": "text/plain; charset=utf-8" },
    })
  }

  const { rawLlmCalls, errors, optimizationFeatures, optimizationEvent } = traceData

  const optimizationsOrigin =
    (optimizationFeatures?.optimizations_origin as Record<
      string,
      { source: string; model?: string; call_sequence?: number; parent?: string }
    >) || {}

  const candidateExplanations =
    (optimizationFeatures?.explanations_post as Record<string, string>) || {}

  const allCandidates = optimizationFeatures?.optimizations_post
    ? Object.entries(optimizationFeatures.optimizations_post as Record<string, string>).map(
        ([id, code]) => ({
          id,
          code: typeof code === "string" ? code : "",
          source: optimizationsOrigin[id]?.source || "OPTIMIZE",
          model: optimizationsOrigin[id]?.model,
          callSequence: optimizationsOrigin[id]?.call_sequence,
          explanation: candidateExplanations[id],
        }),
      )
    : []

  const optimizationCandidates = allCandidates
    .filter(c => c.source === "OPTIMIZE")
    .sort((a, b) => (a.callSequence ?? Infinity) - (b.callSequence ?? Infinity))
    .map((c, index) => ({ ...c, index: index + 1 }))

  const lineProfilerCandidates = allCandidates
    .filter(c => c.source === "OPTIMIZE_LP")
    .sort((a, b) => (a.callSequence ?? Infinity) - (b.callSequence ?? Infinity))
    .map((c, index) => ({ ...c, index: index + 1 }))

  const refinementCandidates = allCandidates
    .filter(c => c.source === "REFINE")
    .sort((a, b) => (a.callSequence ?? Infinity) - (b.callSequence ?? Infinity))
    .map((c, index) => ({
      ...c,
      index: index + 1,
      parentId: optimizationsOrigin[c.id]?.parent || null,
    }))

  const adaptiveCandidates = allCandidates
    .filter(c => c.source === "ADAPTIVE")
    .sort((a, b) => (a.callSequence ?? Infinity) - (b.callSequence ?? Infinity))
    .map((c, index) => ({
      ...c,
      index: index + 1,
      parentId: optimizationsOrigin[c.id]?.parent || null,
    }))

  const rankingData = optimizationFeatures?.ranking as {
    ranking?: string[]
    explanation?: string
  } | null
  const bestCandidateId = rankingData?.ranking?.[0] ?? null

  const pullRequestRaw = optimizationFeatures?.pull_request
  const usedForPr = Boolean(
    pullRequestRaw != null &&
    typeof pullRequestRaw === "object" &&
    !Array.isArray(pullRequestRaw) &&
    Object.keys(pullRequestRaw as Record<string, unknown>).length > 0,
  )

  const rankMap: Record<string, number> = {}
  if (rankingData?.ranking) {
    rankingData.ranking.forEach((id, index) => {
      rankMap[id] = index + 1
    })
  }

  const generatedTests = ((optimizationFeatures?.generated_test ?? []) as string[]).map(
    (code: string, index: number) => ({
      code,
      index: index + 1,
    }),
  )

  const instrumentedTests = (
    (optimizationFeatures?.instrumented_generated_test ?? []) as string[]
  ).map((code: string, index: number) => ({
    code,
    index: index + 1,
  }))

  const instrumentedPerfTests = (
    (optimizationFeatures?.instrumented_perf_test ?? []) as string[]
  ).map((code: string, index: number) => ({
    code,
    index: index + 1,
  }))

  type RawLlmCall = {
    id: string
    trace_id: string | null
    call_type: string | null
    model_name: string | null
    status: string
    latency_ms: number | null
    llm_cost: number | null
    total_tokens: number | null
    created_at: Date
    context: unknown
  }

  const llmCalls = (rawLlmCalls as RawLlmCall[]).sort((a: RawLlmCall, b: RawLlmCall) => {
    const seqA = (a.context as { call_sequence?: number } | null)?.call_sequence ?? Infinity
    const seqB = (b.context as { call_sequence?: number } | null)?.call_sequence ?? Infinity
    if (seqA !== seqB) return seqA - seqB
    return new Date(a.created_at).getTime() - new Date(b.created_at).getTime()
  })

  const transformedCalls = llmCalls.map((call: RawLlmCall) => ({
    id: call.id,
    call_type: call.call_type,
    model_name: call.model_name,
    status: call.status,
    latency_ms: call.latency_ms,
    llm_cost: call.llm_cost,
    total_tokens: call.total_tokens,
    created_at: call.created_at,
    context: call.context as { call_sequence?: number } | null,
  }))

  const { sections, totalDuration } = transformToTimelineSections({
    calls: transformedCalls,
    optimizationCandidates,
    lineProfilerCandidates,
    refinementCandidates,
    adaptiveCandidates,
    generatedTests,
    instrumentedTests,
    instrumentedPerfTests,
    originalCode: optimizationFeatures?.original_code ?? null,
    testFramework: optimizationFeatures?.test_framework ?? null,
    candidateRankMap: rankMap,
    bestCandidateId,
    rankingExplanation: rankingData?.explanation ?? null,
    usedForPr,
  })

  const metadata = optimizationFeatures?.metadata as Record<string, unknown> | undefined
  const functionName = getFunctionName(metadata, optimizationEvent)
  const filePath = getFilePath(optimizationEvent, metadata)
  const originalCode = optimizationFeatures?.original_code ?? null

  const transformedErrors = errors.map(
    (error: {
      error_type: string | null
      severity: string | null
      error_message: string | null
      context: unknown
      created_at: Date
    }) => ({
      error_type: error.error_type,
      severity: error.severity,
      error_message: error.error_message,
      context: error.context as {
        test_name?: string
        failure_reason?: string
        test_output?: string
        expected?: string
        actual?: string
      } | null,
      created_at: error.created_at,
    }),
  )

  const formattedText = formatTimelineForLLM({
    traceId,
    functionName,
    filePath,
    originalCode,
    sections,
    errors: transformedErrors,
    totalDuration,
  })

  return new NextResponse(formattedText, {
    headers: { "Content-Type": "text/plain; charset=utf-8" },
  })
}