From a822c665cda0f5e47709e8aa3e9ba42d890fa110 Mon Sep 17 00:00:00 2001
From: ashrafchowdury <ashrafdev365@gmail.com>
Date: Tue, 3 Dec 2024 21:13:46 +0600
Subject: [PATCH] refactor(frontend): made the states name more meaningful

---
 .../observability/ObservabilityDashboard.tsx  | 30 ++++---
 .../observability/drawer/TestsetDrawer.tsx    | 80 ++++++++++---------
 .../observability/drawer/TraceContent.tsx     |  9 +--
 3 files changed, 60 insertions(+), 59 deletions(-)

diff --git a/agenta-web/src/components/pages/observability/ObservabilityDashboard.tsx b/agenta-web/src/components/pages/observability/ObservabilityDashboard.tsx
index 4e90496723..88500c3bf8 100644
--- a/agenta-web/src/components/pages/observability/ObservabilityDashboard.tsx
+++ b/agenta-web/src/components/pages/observability/ObservabilityDashboard.tsx
@@ -41,7 +41,7 @@ import {convertToCsv, downloadCsv} from "@/lib/helpers/fileManipulations"
 import {useUpdateEffect} from "usehooks-ts"
 import {getStringOrJson} from "@/lib/helpers/utils"
 import ObservabilityContextProvider, {useObservabilityData} from "@/contexts/observability.context"
-import TestsetDrawer from "./drawer/TestsetDrawer"
+import TestsetDrawer, {TestsetTraceData} from "./drawer/TestsetDrawer"
 
 const useStyles = createUseStyles((theme: JSSTheme) => ({
     title: {
@@ -80,11 +80,8 @@ const ObservabilityDashboard = () => {
     const [selectedTraceId, setSelectedTraceId] = useQueryParam("trace", "")
     const [editColumns, setEditColumns] = useState<string[]>(["span_type", "key", "usage"])
     const [isFilterColsDropdownOpen, setIsFilterColsDropdownOpen] = useState(false)
-    const [isTestsetDrawerOpen, setIsTestsetDrawerOpen] = useState(false)
     const [selectedRowKeys, setSelectedRowKeys] = useState<React.Key[]>([])
-    const [testsetDrawerData, setTestsetDrawerData] = useState<
-        {key: string; data: KeyValuePair; id: number}[]
-    >([])
+    const [testsetDrawerData, setTestsetDrawerData] = useState<TestsetTraceData[]>([])
     const [columns, setColumns] = useState<ColumnsType<_AgentaRootsResponse>>([
         {
             title: "ID",
@@ -499,16 +496,18 @@ const ObservabilityDashboard = () => {
         setSort({type, sorted, customRange})
     }, [])
 
-    const getMatchingTracesByDataKeys = () => {
+    const getTestsetTraceData = () => {
         if (!traces?.length) return []
 
-        const extractData = traces
-            .filter((trace) => selectedRowKeys.includes(trace.key))
-            .flatMap((trace, idx) => ({data: trace.data, key: trace.key, id: idx + 1}))
+        const extractData = traces.reduce<TestsetTraceData[]>((acc, trace, idx) => {
+            if (selectedRowKeys.includes(trace.key)) {
+                acc.push({data: trace.data as KeyValuePair, key: trace.key, id: idx + 1})
+            }
+            return acc
+        }, [])
 
         if (extractData.length > 0) {
-            setTestsetDrawerData(extractData as any)
-            setIsTestsetDrawerOpen(true)
+            setTestsetDrawerData(extractData)
         }
     }
 
@@ -562,7 +561,7 @@ const ObservabilityDashboard = () => {
                             Export as CSV
                         </Button>
                         <Button
-                            onClick={() => getMatchingTracesByDataKeys()}
+                            onClick={() => getTestsetTraceData()}
                             icon={<Database size={14} />}
                             disabled={traces.length === 0 || selectedRowKeys.length === 0}
                         >
@@ -653,14 +652,13 @@ const ObservabilityDashboard = () => {
                 />
             </div>
 
-            {isTestsetDrawerOpen && (
+            {testsetDrawerData.length > 0 && (
                 <TestsetDrawer
-                    open={isTestsetDrawerOpen}
+                    open={testsetDrawerData.length > 0}
                     data={testsetDrawerData}
                     onClose={() => {
-                        setSelectedRowKeys([])
                         setTestsetDrawerData([])
-                        setIsTestsetDrawerOpen(false)
+                        setSelectedRowKeys([])
                     }}
                 />
             )}
diff --git a/agenta-web/src/components/pages/observability/drawer/TestsetDrawer.tsx b/agenta-web/src/components/pages/observability/drawer/TestsetDrawer.tsx
index 68d4777167..4c686d02a7 100644
--- a/agenta-web/src/components/pages/observability/drawer/TestsetDrawer.tsx
+++ b/agenta-web/src/components/pages/observability/drawer/TestsetDrawer.tsx
@@ -1,4 +1,4 @@
-import React, {useEffect, useMemo, useState} from "react"
+import {useMemo, useState} from "react"
 import GenericDrawer from "@/components/GenericDrawer"
 import {ArrowRight, PencilSimple, Plus, Trash} from "@phosphor-icons/react"
 import {
@@ -55,11 +55,11 @@ const useStyles = createUseStyles((theme: JSSTheme) => ({
 }))
 
 type Mapping = {data: string; column: string; newColumn?: string}
-type TraceData = {key: string; data: KeyValuePair; id: number}
 type Preview = {key: string; data: KeyValuePair[]}
+export type TestsetTraceData = {key: string; data: KeyValuePair; id: number}
 type Props = {
     onClose: () => void
-    data: TraceData[]
+    data: TestsetTraceData[]
 } & React.ComponentProps<typeof Drawer>
 
 const TestsetDrawer = ({onClose, data, ...props}: Props) => {
@@ -74,17 +74,18 @@ const TestsetDrawer = ({onClose, data, ...props}: Props) => {
     const [isLoading, setIsLoading] = useState(false)
     const [traceData, setTraceData] = useState(data.length > 0 ? data : [])
     const [testset, setTestset] = useState({name: "", id: ""})
-    const [testsetName, setTestsetName] = useState("")
+    const [newTestsetName, setNewTestsetName] = useState("")
     const [editorFormat, setEditorFormat] = useState("JSON")
-    const [tableColumns, setTableColumns] = useState<string[]>([])
-    const [tableRows, setTableRows] = useState<KeyValuePair[]>([])
+    const [selectedTestsetColumns, setSelectedTestsetColumns] = useState<string[]>([])
+    const [selectedTestsetRows, setSelectedTestsetRows] = useState<KeyValuePair[]>([])
     const [showLastFiveRows, setShowLastFiveRows] = useState(false)
-    const [dataPreview, setDataPreview] = useState(traceData[0]?.key || "")
+    const [rowDataPreview, setRowDataPreview] = useState(traceData[0]?.key || "")
     const [mappingData, setMappingData] = useState<Mapping[]>([])
     const [preview, setPreview] = useState<Preview>({key: traceData[0]?.key || "", data: []})
 
     const isNewTestset = testset.id === "create"
     const elementWidth = isDrawerExtended ? 200 * 2 : 200
+    const selectedTestsetTestCases = selectedTestsetRows.slice(-5)
     const isMapColumnExist = mappingData.some((mapping) =>
         mapping.column === "create" || !mapping.column ? !!mapping?.newColumn : !!mapping.column,
     )
@@ -113,8 +114,8 @@ const TestsetDrawer = ({onClose, data, ...props}: Props) => {
             if (value && value !== "create") {
                 const data = await fetchTestset(value)
                 if (data?.csvdata?.length) {
-                    setTableColumns(Object.keys(data.csvdata[0]))
-                    setTableRows(data.csvdata)
+                    setSelectedTestsetColumns(Object.keys(data.csvdata[0]))
+                    setSelectedTestsetRows(data.csvdata)
                 }
             }
 
@@ -132,22 +133,22 @@ const TestsetDrawer = ({onClose, data, ...props}: Props) => {
     }
 
     const onRemoveTraceData = () => {
-        const removeTrace = traceData.filter((trace) => trace.key !== dataPreview)
+        const removeTrace = traceData.filter((trace) => trace.key !== rowDataPreview)
         setTraceData(removeTrace)
 
         if (removeTrace.length > 0) {
-            const currentIndex = traceData.findIndex((trace) => trace.key === dataPreview)
+            const currentIndex = traceData.findIndex((trace) => trace.key === rowDataPreview)
             // [currentIndex]: Next option in list | [currentIndex - 1]: Previous option if next doesn't exist | [0]: Default to first option
             const nextPreview =
                 removeTrace[currentIndex] || removeTrace[currentIndex - 1] || removeTrace[0]
 
-            setDataPreview(nextPreview.key)
+            setRowDataPreview(nextPreview.key)
 
-            if (dataPreview === preview.key) {
+            if (rowDataPreview === preview.key) {
                 onPreviewOptionChange(nextPreview.key)
             }
         } else {
-            setDataPreview("")
+            setRowDataPreview("")
         }
     }
 
@@ -155,7 +156,9 @@ const TestsetDrawer = ({onClose, data, ...props}: Props) => {
         if (!traceData?.length) return ""
 
         const jsonObject = {
-            data: traceData.find((trace) => trace?.key === dataPreview)?.data || traceData[0]?.data,
+            data:
+                traceData.find((trace) => trace?.key === rowDataPreview)?.data ||
+                traceData[0]?.data,
         }
         if (!jsonObject) return ""
 
@@ -165,7 +168,7 @@ const TestsetDrawer = ({onClose, data, ...props}: Props) => {
             message.error("Failed to convert JSON to YAML. Please ensure the data is valid.")
             return getStringOrJson(jsonObject)
         }
-    }, [editorFormat, traceData, dataPreview])
+    }, [editorFormat, traceData, rowDataPreview])
 
     const mappingOptions = useMemo(() => {
         const uniquePaths = new Set<string>()
@@ -176,14 +179,14 @@ const TestsetDrawer = ({onClose, data, ...props}: Props) => {
         })
 
         return Array.from(uniquePaths).map((item) => ({value: item}))
-    }, [data])
+    }, [traceData])
 
     const columnOptions = useMemo(() => {
         const selectedColumns = mappingData
             .map((item) => item.column)
             .filter((col) => col !== "create")
-        return tableColumns.filter((column) => !selectedColumns.includes(column))
-    }, [mappingData, tableColumns])
+        return selectedTestsetColumns.filter((column) => !selectedColumns.includes(column))
+    }, [mappingData, selectedTestsetColumns])
 
     const onMappingOptionChange = ({
         pathName,
@@ -220,15 +223,15 @@ const TestsetDrawer = ({onClose, data, ...props}: Props) => {
     }, [mappingData])
 
     const resetStates = () => {
-        setTableColumns([])
-        setTableRows([])
+        setSelectedTestsetColumns([])
+        setSelectedTestsetRows([])
         setShowLastFiveRows(false)
         setMappingData((prev) => prev.map((item) => ({...item, column: "", newColumn: ""})))
         setPreview({key: traceData[0]?.key || "", data: []})
-        setTestsetName("")
+        setNewTestsetName("")
     }
 
-    const mapAndConvertDataInCsvFormat = (traceData: TraceData[]) => {
+    const mapAndConvertDataInCsvFormat = (traceData: TestsetTraceData[]) => {
         return traceData.map((item) => {
             const formattedItem: Record<string, any> = {}
 
@@ -262,17 +265,17 @@ const TestsetDrawer = ({onClose, data, ...props}: Props) => {
             const newTestsetData = mapAndConvertDataInCsvFormat(traceData)
 
             if (isNewTestset) {
-                if (!testsetName) {
+                if (!newTestsetName) {
                     message.error("Please add a Test set name before saving it")
                     return
                 }
 
-                await createNewTestset(testsetName, newTestsetData)
+                await createNewTestset(newTestsetName, newTestsetData)
                 message.success("Test set created successfully")
             } else {
                 await updateTestset(testset.id as string, testset.name, [
                     ...newTestsetData,
-                    ...tableRows,
+                    ...selectedTestsetRows,
                 ])
                 message.success("Test set updated successfully")
             }
@@ -348,8 +351,8 @@ const TestsetDrawer = ({onClose, data, ...props}: Props) => {
                                     <div className="relative">
                                         <Input
                                             style={{width: elementWidth}}
-                                            value={testsetName}
-                                            onChange={(e) => setTestsetName(e.target.value)}
+                                            value={newTestsetName}
+                                            onChange={(e) => setNewTestsetName(e.target.value)}
                                             placeholder="Test set name"
                                         />
                                         <PencilSimple
@@ -369,8 +372,8 @@ const TestsetDrawer = ({onClose, data, ...props}: Props) => {
                             <div className="flex justify-between items-center mb-2">
                                 <Select
                                     style={{width: elementWidth}}
-                                    value={dataPreview}
-                                    onChange={(value) => setDataPreview(value)}
+                                    value={rowDataPreview}
+                                    onChange={(value) => setRowDataPreview(value)}
                                     options={traceData.map((trace) => ({
                                         value: trace?.key,
                                         label: `Span ${trace.id}`,
@@ -386,8 +389,9 @@ const TestsetDrawer = ({onClose, data, ...props}: Props) => {
                                         >
                                             Remove span{" "}
                                             {
-                                                traceData.find((trace) => trace.key === dataPreview)
-                                                    ?.id
+                                                traceData.find(
+                                                    (trace) => trace.key === rowDataPreview,
+                                                )?.id
                                             }
                                         </Button>
                                     )}
@@ -564,8 +568,8 @@ const TestsetDrawer = ({onClose, data, ...props}: Props) => {
                                                     setShowLastFiveRows(!showLastFiveRows)
                                                 }
                                             >
-                                                Show last {tableRows.slice(-5).length} test cases in
-                                                test set
+                                                Show last {selectedTestsetTestCases.length} test
+                                                cases in test set
                                             </Checkbox>
                                         )}
                                     </div>
@@ -590,17 +594,19 @@ const TestsetDrawer = ({onClose, data, ...props}: Props) => {
                                             })}
                                             dataSource={[
                                                 ...preview.data,
-                                                ...(showLastFiveRows ? tableRows.slice(-5) : []),
+                                                ...(showLastFiveRows
+                                                    ? selectedTestsetTestCases
+                                                    : []),
                                             ]}
                                             rowClassName={(_, index) => {
                                                 if (showLastFiveRows) {
                                                     const totalRows =
                                                         preview.data.length +
-                                                        tableRows.slice(-5).length
+                                                        selectedTestsetTestCases.length
 
                                                     if (
                                                         index >=
-                                                        totalRows - tableRows.slice(-5).length
+                                                        totalRows - selectedTestsetTestCases.length
                                                     ) {
                                                         return "!bg-[#fafafa]"
                                                     }
diff --git a/agenta-web/src/components/pages/observability/drawer/TraceContent.tsx b/agenta-web/src/components/pages/observability/drawer/TraceContent.tsx
index eca6c96a05..766689fc74 100644
--- a/agenta-web/src/components/pages/observability/drawer/TraceContent.tsx
+++ b/agenta-web/src/components/pages/observability/drawer/TraceContent.tsx
@@ -1,5 +1,5 @@
 import ResultTag from "@/components/ResultTag/ResultTag"
-import {JSSTheme} from "@/lib/Types"
+import {JSSTheme, KeyValuePair} from "@/lib/Types"
 import {ArrowRight, Database, PlusCircle, Rocket, Timer} from "@phosphor-icons/react"
 import {Button, Divider, Space, Tabs, TabsProps, Typography} from "antd"
 import React, {useState} from "react"
@@ -88,10 +88,6 @@ const TraceContent = ({activeTrace}: TraceContentProps) => {
     const {icon, bgColor, color} = statusMapper(activeTrace.node.type)
     const [isTestsetDrawerOpen, setIsTestsetDrawerOpen] = useState(false)
 
-    const testsetDrawerData = () => {
-        return [{data: activeTrace.data, key: activeTrace.key, id: 1}]
-    }
-
     const transformDataInputs = (data: any) => {
         return Object.keys(data).reduce((acc, curr) => {
             if (curr === "prompt") {
@@ -279,6 +275,7 @@ const TraceContent = ({activeTrace}: TraceContentProps) => {
                             <Button
                                 className="flex items-center"
                                 onClick={() => setIsTestsetDrawerOpen(true)}
+                                disabled={!activeTrace.key}
                             >
                                 <Database size={14} />
                                 Add to testset
@@ -368,7 +365,7 @@ const TraceContent = ({activeTrace}: TraceContentProps) => {
             {isTestsetDrawerOpen && (
                 <TestsetDrawer
                     open={isTestsetDrawerOpen}
-                    data={testsetDrawerData() as any}
+                    data={[{data: activeTrace.data as KeyValuePair, key: activeTrace.key, id: 1}]}
                     onClose={() => setIsTestsetDrawerOpen(false)}
                 />
             )}