[OPIK-492] [FR]: Export traces as JSON instead of CSV

apps/opik-frontend/src/components/pages/TracesPage/ExportToButton.tsx

@@ -0,0 +1,111 @@
+import React, { useCallback, useState } from "react";
+import { Download } from "lucide-react";
+import FileSaver from "file-saver";
+import { json2csv } from "json-2-csv";
+import last from "lodash/last";
+import first from "lodash/first";
+import get from "lodash/get";
+import slugify from "slugify";
+
+import { Button } from "@/components/ui/button";
+import {
+  DropdownMenu,
+  DropdownMenuContent,
+  DropdownMenuItem,
+  DropdownMenuTrigger,
+} from "@/components/ui/dropdown-menu";
+import { Span, Trace } from "@/types/traces";
+import { TRACE_DATA_TYPE } from "@/hooks/useTracesOrSpansList";
+
+const mapRowData = (rows: Array<Trace | Span>, columns: string[]) => {
+  return rows.map((row) => {
+    return columns.reduce<Record<string, unknown>>((acc, column) => {
+      // we need split by dot to parse usage into correct structure
+      const keys = column.split(".");
+      const key = last(keys) as string;
+      const keyPrefix = first(keys) as string;
+
+      if (keyPrefix === "feedback_scores") {
+        acc[key] = get(
+          row.feedback_scores?.find((f) => f.name === key),
+          "value",
+          "-",
+        );
+      } else {
+        acc[key] = get(row, keys, "");
+      }
+      return acc;
+    }, {});
+  });
+};
+
+const generateFileName = (
+  projectName: string,
+  type: TRACE_DATA_TYPE,
+  extension = "csv",
+) => {
+  return `${slugify(projectName, { lower: true })}-${
+    type === TRACE_DATA_TYPE.traces ? "traces" : "llm-calls"
+  }.${extension}`;
+};
+
+type ExportToButtonProps = {
+  projectName: string;
+  type: TRACE_DATA_TYPE;
+  disabled: boolean;
+  rows: Array<Trace | Span>;
+  columnsToExport: string[];
+};
+
+const ExportToButton: React.FC<ExportToButtonProps> = ({
+  projectName,
+  type,
+  disabled,
+  rows,
+  columnsToExport,
+}) => {
+  const [open, setOpen] = useState<boolean>(false);
+
+  const exportCSVHandler = useCallback(() => {
+    const fileName = generateFileName(projectName, type);
+    const mappedRows = mapRowData(rows, columnsToExport);
+
+    FileSaver.saveAs(
+      new Blob([json2csv(mappedRows, { arrayIndexesAsKeys: true })], {
+        type: "text/csv;charset=utf-8",
+      }),
+      fileName,
+    );
+  }, [projectName, rows, columnsToExport, type]);
+
+  const exportJSONHandler = useCallback(() => {
+    const fileName = generateFileName(projectName, type, "json");
+    const mappedRows = mapRowData(rows, columnsToExport);
+
+    FileSaver.saveAs(
+      new Blob([JSON.stringify(mappedRows, null, 2)], {
+        type: "application/json;charset=utf-8",
+      }),
+      fileName,
+    );
+  }, [projectName, rows, columnsToExport, type]);
+  return (
+    <DropdownMenu open={open} onOpenChange={setOpen}>
+      <DropdownMenuTrigger asChild>
+        <Button variant="outline" size="icon" disabled={disabled}>
+          <Download className="size-4" />
+        </Button>
+      </DropdownMenuTrigger>
+      <DropdownMenuContent align="end" className="w-52">
+        <DropdownMenuItem onClick={exportCSVHandler}>
+          Export as CSV
+        </DropdownMenuItem>
+        <DropdownMenuItem onClick={exportJSONHandler}>
+          Export as JSON
+        </DropdownMenuItem>
+      </DropdownMenuContent>
+    </DropdownMenu>
+  );
+};
+
+export default ExportToButton;

apps/opik-frontend/src/components/pages/TracesPage/TracesActionsPanel.tsx

@@ -1,12 +1,5 @@
 import React, { useCallback, useRef, useState } from "react";
-import last from "lodash/last";
-import get from "lodash/get";
-import { first } from "lodash";
-import { json2csv } from "json-2-csv";
-import FileSaver from "file-saver";
-import slugify from "slugify";
-
-import { Database, Download, Trash } from "lucide-react";
+import { Database, Trash } from "lucide-react";
 
 import { Button } from "@/components/ui/button";
 import { Span, Trace } from "@/types/traces";
@@ -15,19 +8,20 @@ import AddToDatasetDialog from "@/components/pages/TracesPage/AddToDataset/AddTo
 import ConfirmDialog from "@/components/shared/ConfirmDialog/ConfirmDialog";
 import useTracesBatchDeleteMutation from "@/api/traces/useTraceBatchDeleteMutation";
 import TooltipWrapper from "@/components/shared/TooltipWrapper/TooltipWrapper";
+import ExportToButton from "@/components/pages/TracesPage/ExportToButton";
 
 type TracesActionsPanelProps = {
   type: TRACE_DATA_TYPE;
   rows: Array<Trace | Span>;
-  selectedColumns: string[];
+  columnsToExport: string[];
   projectName: string;
   projectId: string;
 };
 
 const TracesActionsPanel: React.FunctionComponent<TracesActionsPanelProps> = ({
   rows,
   type,
-  selectedColumns,
+  columnsToExport,
   projectName,
   projectId,
 }) => {
@@ -45,38 +39,6 @@ const TracesActionsPanel: React.FunctionComponent<TracesActionsPanelProps> = ({
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [projectId, rows]);
 
-  const exportCSVHandler = useCallback(() => {
-    const fileName = `${slugify(projectName, { lower: true })}-${
-      type === TRACE_DATA_TYPE.traces ? "traces" : "llm-calls"
-    }.csv`;
-    const mappedRows = rows.map((row) => {
-      return selectedColumns.reduce<Record<string, unknown>>((acc, column) => {
-        // we need split by dot to parse usage into correct structure
-        const keys = column.split(".");
-        const key = last(keys) as string;
-        const keyPrefix = first(keys) as string;
-
-        if (keyPrefix === "feedback_scores") {
-          acc[key] = get(
-            row.feedback_scores?.find((f) => f.name === key),
-            "value",
-            "-",
-          );
-        } else {
-          acc[key] = get(row, keys, "");
-        }
-        return acc;
-      }, {});
-    });
-
-    FileSaver.saveAs(
-      new Blob([json2csv(mappedRows, { arrayIndexesAsKeys: true })], {
-        type: "text/csv;charset=utf-8",
-      }),
-      fileName,
-    );
-  }, [projectName, rows, selectedColumns, type]);
-
   return (
     <div className="flex items-center gap-2">
       <AddToDatasetDialog
@@ -107,16 +69,14 @@ const TracesActionsPanel: React.FunctionComponent<TracesActionsPanelProps> = ({
           Add to dataset
         </Button>
       </TooltipWrapper>
-      <TooltipWrapper content="Export CSV">
-        <Button
-          variant="outline"
-          size="icon"
-          onClick={exportCSVHandler}
-          disabled={disabled || selectedColumns.length === 0}
-        >
-          <Download className="size-4" />
-        </Button>
-      </TooltipWrapper>
+      <ExportToButton
+        disabled={disabled || columnsToExport.length === 0}
+        type={type}
+        columnsToExport={columnsToExport}
+        rows={rows}
+        projectName={projectName}
+      />
+
       {type === TRACE_DATA_TYPE.traces && (
         <TooltipWrapper content="Delete">
           <Button

apps/opik-frontend/src/components/pages/TracesPage/TracesSpansTab.tsx

@@ -12,6 +12,7 @@ import findIndex from "lodash/findIndex";
 import isObject from "lodash/isObject";
 import difference from "lodash/difference";
 import union from "lodash/union";
+import get from "lodash/get";
 
 import useTracesOrSpansList, {
   TRACE_DATA_TYPE,
@@ -388,6 +389,12 @@ export const TracesSpansTab: React.FC<TracesSpansTabProps> = ({
     scoresColumnsOrder,
   ]);
 
+  const columnsToExport = useMemo(() => {
+    return columns
+      .map((c) => get(c, "accessorKey", ""))
+      .filter((c) => selectedColumns.includes(c));
+  }, [columns, selectedColumns]);
+
   const activeRowId = type === TRACE_DATA_TYPE.traces ? traceId : spanId;
   const rowIndex = findIndex(rows, (row) => activeRowId === row.id);
 
@@ -441,7 +448,7 @@ export const TracesSpansTab: React.FC<TracesSpansTabProps> = ({
             projectId={projectId}
             projectName={projectName}
             rows={selectedRows}
-            selectedColumns={selectedColumns}
+            columnsToExport={columnsToExport}
             type={type as TRACE_DATA_TYPE}
           />
           <Separator orientation="vertical" className="ml-2 mr-2.5 h-6" />