Merge pull request #2328 from Agenta-AI/AGE-1337/implement-testset-dr…

…awer

[Feat]: implement Test set drawer UI

agenta-web/src/components/GenericDrawer/index.tsx

@@ -7,10 +7,12 @@ type GenericDrawerProps = {
     headerExtra?: ReactNode
     mainContent: ReactNode
     sideContent?: ReactNode
+    initialWidth?: number
 } & React.ComponentProps<typeof Drawer>
 
 const GenericDrawer = ({...props}: GenericDrawerProps) => {
-    const [drawerWidth, setDrawerWidth] = useState(1200)
+    const initialWidth = props.initialWidth || 1200
+    const [drawerWidth, setDrawerWidth] = useState(initialWidth)
 
     return (
         <Drawer
@@ -28,15 +30,15 @@ const GenericDrawer = ({...props}: GenericDrawerProps) => {
                     {props.expandable && (
                         <Button
                             onClick={() => {
-                                if (drawerWidth === 1200) {
+                                if (drawerWidth === initialWidth) {
                                     setDrawerWidth(1920)
                                 } else {
-                                    setDrawerWidth(1200)
+                                    setDrawerWidth(initialWidth)
                                 }
                             }}
                             type="text"
                             icon={
-                                drawerWidth === 1200 ? (
+                                drawerWidth === initialWidth ? (
                                     <FullscreenOutlined />
                                 ) : (
                                     <FullscreenExitOutlined />

agenta-web/src/components/pages/observability/ObservabilityDashboard.tsx

@@ -14,7 +14,7 @@ import {useAppId} from "@/hooks/useAppId"
 import {useQueryParam} from "@/hooks/useQuery"
 import {formatCurrency, formatLatency, formatTokenUsage} from "@/lib/helpers/formatters"
 import {getNodeById} from "@/lib/helpers/observability_helpers"
-import {Filter, FilterConditions, JSSTheme} from "@/lib/Types"
+import {Filter, FilterConditions, JSSTheme, KeyValuePair} from "@/lib/Types"
 import {_AgentaRootsResponse} from "@/services/observability/types"
 import {ReloadOutlined, SwapOutlined} from "@ant-design/icons"
 import {
@@ -35,12 +35,13 @@ import dayjs from "dayjs"
 import {useRouter} from "next/router"
 import React, {useCallback, useEffect, useMemo, useState} from "react"
 import {createUseStyles} from "react-jss"
-import {Export} from "@phosphor-icons/react"
+import {Database, Export} from "@phosphor-icons/react"
 import {getAppValues} from "@/contexts/app.context"
 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, {TestsetTraceData} from "./drawer/TestsetDrawer"
 
 const useStyles = createUseStyles((theme: JSSTheme) => ({
     title: {
@@ -79,6 +80,8 @@ const ObservabilityDashboard = () => {
     const [selectedTraceId, setSelectedTraceId] = useQueryParam("trace", "")
     const [editColumns, setEditColumns] = useState<string[]>(["span_type", "key", "usage"])
     const [isFilterColsDropdownOpen, setIsFilterColsDropdownOpen] = useState(false)
+    const [selectedRowKeys, setSelectedRowKeys] = useState<React.Key[]>([])
+    const [testsetDrawerData, setTestsetDrawerData] = useState<TestsetTraceData[]>([])
     const [columns, setColumns] = useState<ColumnsType<_AgentaRootsResponse>>([
         {
             title: "ID",
@@ -263,6 +266,12 @@ const ObservabilityDashboard = () => {
         return () => clearInterval(interval)
     }, [])
 
+    const rowSelection = {
+        onChange: (selectedRowKeys: React.Key[]) => {
+            setSelectedRowKeys(selectedRowKeys)
+        },
+    }
+
     const selectedItem = useMemo(
         () => (traces?.length ? getNodeById(traces, selected) : null),
         [selected, traces],
@@ -487,6 +496,19 @@ const ObservabilityDashboard = () => {
         setSort({type, sorted, customRange})
     }, [])
 
+    const getTestsetTraceData = () => {
+        if (!traces?.length) return []
+
+        const extractData = selectedRowKeys.map((key, idx) => {
+            const node = getNodeById(traces, key as string)
+            return {data: node?.data as KeyValuePair, key: node?.key, id: idx + 1}
+        })
+
+        if (extractData.length > 0) {
+            setTestsetDrawerData(extractData as TestsetTraceData[])
+        }
+    }
+
     return (
         <div className="flex flex-col gap-6">
             <Typography.Text className={classes.title}>Observability</Typography.Text>
@@ -536,6 +558,13 @@ const ObservabilityDashboard = () => {
                         >
                             Export as CSV
                         </Button>
+                        <Button
+                            onClick={() => getTestsetTraceData()}
+                            icon={<Database size={14} />}
+                            disabled={traces.length === 0 || selectedRowKeys.length === 0}
+                        >
+                            Add test set
+                        </Button>
                         <EditColumns
                             isOpen={isFilterColsDropdownOpen}
                             handleOpenChange={setIsFilterColsDropdownOpen}
@@ -549,6 +578,12 @@ const ObservabilityDashboard = () => {
 
             <div className="flex flex-col gap-2">
                 <Table
+                    rowSelection={{
+                        type: "checkbox",
+                        columnWidth: 48,
+                        selectedRowKeys,
+                        ...rowSelection,
+                    }}
                     loading={isLoading}
                     columns={mergedColumns as TableColumnType<_AgentaRootsResponse>[]}
                     dataSource={traces}
@@ -615,6 +650,17 @@ const ObservabilityDashboard = () => {
                 />
             </div>
 
+            {testsetDrawerData.length > 0 && (
+                <TestsetDrawer
+                    open={testsetDrawerData.length > 0}
+                    data={testsetDrawerData}
+                    onClose={() => {
+                        setTestsetDrawerData([])
+                        setSelectedRowKeys([])
+                    }}
+                />
+            )}
+
             {activeTrace && !!traces?.length && (
                 <GenericDrawer
                     open={!!selectedTraceId}