(frontend)[Age-1286]: Update web UI to use BaseResponse version 3.0 in the Playground

agenta-web/src/components/Playground/Views/TestView.tsx

@@ -1,4 +1,4 @@
-import React, {useContext, useEffect, useRef, useState} from "react"
+import React, {useContext, useEffect, useMemo, useRef, useState} from "react"
 import {Button, Input, Card, Row, Col, Space, Form, Modal} from "antd"
 import {CaretRightOutlined, CloseCircleOutlined, PlusOutlined} from "@ant-design/icons"
 import {callVariant} from "@/services/api"
@@ -10,9 +10,9 @@ import {
     Parameter,
     Variant,
     StyleProps,
-    BaseResponseSpans,
     BaseResponse,
-    TraceDetails,
+    TraceDetailsV2,
+    TraceDetailsV3,
 } from "@/lib/Types"
 import {batchExecute, getStringOrJson, isDemo, randString, removeKeys} from "@/lib/helpers/utils"
 import LoadTestsModal from "../LoadTestsModal"
@@ -34,12 +34,20 @@ import relativeTime from "dayjs/plugin/relativeTime"
 import duration from "dayjs/plugin/duration"
 import {useQueryParam} from "@/hooks/useQuery"
 import {formatCurrency, formatLatency, formatTokenUsage} from "@/lib/helpers/formatters"
-import {dynamicComponent, dynamicService} from "@/lib/helpers/dynamic"
+import {dynamicService} from "@/lib/helpers/dynamic"
 import {isBaseResponse, isFuncResponse} from "@/lib/helpers/playgroundResp"
-import {fromBaseResponseToTraceSpanType} from "@/lib/transformers"
 import {AgentaNodeDTO} from "@/services/observability/types"
+import {isTraceDetailsV2, isTraceDetailsV3} from "@/lib/helpers/observability_helpers"
+import GenericDrawer from "@/components/GenericDrawer"
+import TraceHeader from "@/components/pages/observability/drawer/TraceHeader"
+import TraceTree from "@/components/pages/observability/drawer/TraceTree"
+import TraceContent from "@/components/pages/observability/drawer/TraceContent"
+import {
+    buildNodeTree,
+    getNodeById,
+    observabilityTransformer,
+} from "@/lib/helpers/observability_helpers"
 
-const PlaygroundDrawer: any = dynamicComponent(`Playground/PlaygroundDrawer/PlaygroundDrawer`)
 const promptRevision: any = dynamicService("promptVersioning/api")
 
 dayjs.extend(relativeTime)
@@ -176,10 +184,7 @@ interface BoxComponentProps {
     additionalData: {
         cost: number | null
         latency: number | null
-        usage:
-            | {completion_tokens: number; prompt_tokens: number; total_tokens: number}
-            | number
-            | null
+        usage: number | null
     }
     onInputParamChange: (paramName: string, newValue: any) => void
     onRun: () => void
@@ -188,7 +193,7 @@ interface BoxComponentProps {
     isChatVariant?: boolean
     variant: Variant
     onCancel: () => void
-    traceSpans: AgentaNodeDTO
+    traceSpans: TraceDetailsV2 | TraceDetailsV3 | undefined
 }
 
 const BoxComponent: React.FC<BoxComponentProps> = ({
@@ -205,8 +210,31 @@ const BoxComponent: React.FC<BoxComponentProps> = ({
     onCancel,
     traceSpans,
 }) => {
+    const [selectedTraceId, setSelectedTraceId] = useQueryParam("trace", "")
+
+    const traces = useMemo(() => {
+        if (traceSpans && isTraceDetailsV3(traceSpans)) {
+            return traceSpans.nodes
+                .flatMap((node: AgentaNodeDTO) => buildNodeTree(node))
+                .flatMap((item: any) => observabilityTransformer(item))
+        }
+    }, [traceSpans])
+
+    const activeTrace = useMemo(() => (traces ? (traces[0] ?? null) : null), [traces])
+    const [selected, setSelected] = useState("")
+
+    useEffect(() => {
+        if (!selected) {
+            setSelected(activeTrace?.node.id ?? "")
+        }
+    }, [activeTrace, selected])
+
+    const selectedItem = useMemo(
+        () => (traces?.length ? getNodeById(traces, selected) : null),
+        [selected, traces],
+    )
+
     const {appTheme} = useAppTheme()
-    const [activeSpan, setActiveSpan] = useQueryParam("activeSpan")
     const classes = useStylesBox()
     const loading = result === LOADING_TEXT
     const [form] = Form.useForm()
@@ -332,11 +360,11 @@ const BoxComponent: React.FC<BoxComponentProps> = ({
                     <span>Tokens: {formatTokenUsage(additionalData?.usage)}</span>
                     <span>Cost: {formatCurrency(additionalData?.cost)}</span>
                     <span>Latency: {formatLatency(additionalData?.latency)}</span>
-                    {traceSpans?.nodes?.length && isDemo() && (
+                    {traceSpans && isTraceDetailsV3(traceSpans) && (
                         <Button
                             type="link"
                             className={classes.viewTracesBtn}
-                            onClick={() => setActiveSpan(traceSpans.root.id)}
+                            onClick={() => setSelectedTraceId(traceSpans.nodes[0].node.id)}
                         >
                             View Trace
                         </Button>
@@ -346,21 +374,29 @@ const BoxComponent: React.FC<BoxComponentProps> = ({
                 ""
             )}
 
-            {/* TODO:
-                Replace PlaygroundDrawer with Drawer component used
-                in observability feature
-            */}
-            {/* {traceSpans?.nodes?.length && !!activeSpan && (
-                <PlaygroundDrawer
-                    placement="bottom"
-                    open={!!activeSpan}
-                    onClose={() => setActiveSpan("")}
-                    traceSpans={fromBaseResponseToTraceSpanType(
-                        traceSpans.nodes,
-                        traceSpans.root.id,
-                    )}
+            {activeTrace && !!traces?.length && (
+                <GenericDrawer
+                    open={!!selectedTraceId}
+                    onClose={() => setSelectedTraceId("")}
+                    expandable
+                    headerExtra={
+                        <TraceHeader
+                            activeTrace={activeTrace}
+                            traces={traces}
+                            setSelectedTraceId={setSelectedTraceId}
+                            activeTraceIndex={0}
+                        />
+                    }
+                    mainContent={selectedItem ? <TraceContent activeTrace={selectedItem} /> : null}
+                    sideContent={
+                        <TraceTree
+                            activeTrace={activeTrace}
+                            selected={selected}
+                            setSelected={setSelected}
+                        />
+                    }
                 />
-            )} */}
+            )}
         </Card>
     )
 }
@@ -397,7 +433,7 @@ const App: React.FC<TestViewProps> = ({
             usage: number | null
         }>
     >(testList.map(() => ({cost: null, latency: null, usage: null})))
-    const [traceSpans, setTraceSpans] = useState<AgentaNodeDTO[] | TraceDetails | null>(null)
+    const [traceSpans, setTraceSpans] = useState<TraceDetailsV2 | TraceDetailsV3>()
     const [revisionNum] = useQueryParam("revision")
 
     useEffect(() => {
@@ -554,7 +590,7 @@ const App: React.FC<TestViewProps> = ({
                 res = {version: "3.0", data: result} as BaseResponse
                 setResultForIndex(getStringOrJson(res.data), index)
             } else if (isFuncResponse(result)) {
-                const res: BaseResponse = {version: "3.0", data: result.message}
+                const res = {version: "3.0", data: result.message}
                 setResultForIndex(getStringOrJson(res.data), index)
 
                 const {message, cost, latency, usage} = result
@@ -573,28 +609,30 @@ const App: React.FC<TestViewProps> = ({
                 res = result as BaseResponse
                 setResultForIndex(getStringOrJson(res.data), index)
 
-                const {trace} = result
+                const {trace, version} = result
+
+                // Main update logic
                 setAdditionalDataList((prev) => {
                     const newDataList = [...prev]
-                    if (result.version === "2.0") {
-                        // Version 2.0 logic
+                    if (version === "2.0" && isTraceDetailsV2(result.trace)) {
                         newDataList[index] = {
-                            cost: trace?.cost || null,
-                            latency: trace?.latency || null,
-                            usage: trace?.usage?.total_tokens || null,
+                            cost: result.trace?.cost ?? null,
+                            latency: result.trace?.latency ?? null,
+                            usage: result.trace?.usage?.total_tokens ?? null,
                         }
-                    } else if (result.version === "3.0") {
-                        // Version 3.0 logic
+                    } else if (version === "3.0" && isTraceDetailsV3(result.trace)) {
+                        const firstTraceNode = result.trace.nodes[0]
                         newDataList[index] = {
-                            cost: trace?.nodes?.[0]?.metrics?.acc?.costs?.total ?? null,
-                            latency: trace?.nodes?.[0]?.time?.span
-                                ? trace?.nodes?.[0]?.time?.span / 1_000_000
+                            cost: firstTraceNode?.metrics?.acc?.costs?.total ?? null,
+                            latency: firstTraceNode?.time?.span
+                                ? firstTraceNode.time.span / 1_000_000
                                 : null,
-                            usage: trace?.nodes?.[0]?.metrics?.acc?.tokens?.total ?? null,
+                            usage: firstTraceNode?.metrics?.acc?.tokens?.total ?? null,
                         }
                     }
                     return newDataList
                 })
+
                 if (trace && isDemo()) {
                     setTraceSpans(trace)
                 }
@@ -744,7 +782,7 @@ const App: React.FC<TestViewProps> = ({
                     isChatVariant={isChatVariant}
                     variant={variant}
                     onCancel={() => handleCancel(index)}
-                    traceSpans={Array.isArray(traceSpans) ? traceSpans : traceSpans}
+                    traceSpans={traceSpans}
                 />
             ))}
             <Button

agenta-web/src/components/pages/observability/drawer/TraceContent.tsx

@@ -297,7 +297,11 @@ const TraceContent = ({activeTrace}: TraceContentProps) => {
                         value1={
                             <div className={classes.resultTag}>
                                 <Timer size={14} />{" "}
-                                {formatLatency(activeTrace?.metrics?.acc?.duration.total / 1000)}
+                                {formatLatency(
+                                    activeTrace?.metrics?.acc?.duration.total
+                                        ? activeTrace?.metrics?.acc?.duration.total
+                                        : activeTrace?.metrics?.acc?.tokens.total / 1000,
+                                )}
                             </div>
                         }
                     />

agenta-web/src/components/pages/observability/drawer/TraceHeader.tsx

@@ -11,9 +11,9 @@ interface TraceHeaderProps {
     activeTrace: _AgentaRootsResponse
     traces: _AgentaRootsResponse[]
     setSelectedTraceId: (val: string) => void
-    activeTraceIndex: number
-    handleNextTrace: () => void
-    handlePrevTrace: () => void
+    activeTraceIndex?: number
+    handleNextTrace?: () => void
+    handlePrevTrace?: () => void
 }
 
 const useStyles = createUseStyles((theme: JSSTheme) => ({

agenta-web/src/components/pages/observability/drawer/TraceTree.tsx

@@ -96,7 +96,11 @@ const TreeContent = ({value}: {value: _AgentaRootsResponse}) => {
                 <Space className={classes.treeContentContainer}>
                     <div className={classes.treeContent}>
                         <Timer />
-                        {formatLatency(metrics?.acc?.duration.total / 1000)}
+                        {formatLatency(
+                            metrics?.acc?.duration.total
+                                ? metrics?.acc?.duration.total
+                                : metrics?.acc?.tokens.total / 1000,
+                        )}
                     </div>
 
                     {metrics?.acc?.costs?.total && (

agenta-web/src/lib/Types.ts

@@ -547,18 +547,24 @@ export type FuncResponse = {
     usage: {completion_tokens: number; prompt_tokens: number; total_tokens: number}
 }
 
-export interface TraceDetails {
+export interface TraceDetailsV2 {
     trace_id: string
     cost?: number
     latency?: number
-    usage?: number
+    usage: {completion_tokens: number; prompt_tokens: number; total_tokens: number}
     spans?: BaseResponseSpans[]
 }
 
-export type BaseResponse = {
+export interface TraceDetailsV3 {
     version: string
+    nodes: AgentaNodeDTO[]
+    count?: number | null
+}
+
+export type BaseResponse = {
+    version?: string | null
     data: string | Record<string, any>
-    trace?: TraceDetails | AgentaNodeDTO[]
+    trace: TraceDetailsV2 | TraceDetailsV3
 }
 
 export type BaseResponseSpans = {

agenta-web/src/lib/helpers/observability_helpers.ts

@@ -1,4 +1,5 @@
 import {_AgentaRootsResponse, AgentaNodeDTO, AgentaTreeDTO} from "@/services/observability/types"
+import {BaseResponse, TraceDetailsV2, TraceDetailsV3} from "../Types"
 
 export const observabilityTransformer = (
     item: AgentaTreeDTO | AgentaNodeDTO,
@@ -74,3 +75,11 @@ export const getNodeById = (
     }
     return null
 }
+
+export const isTraceDetailsV2 = (trace: BaseResponse["trace"]): trace is TraceDetailsV2 => {
+    return (trace as TraceDetailsV2)?.trace_id !== undefined
+}
+
+export const isTraceDetailsV3 = (trace: BaseResponse["trace"]): trace is TraceDetailsV3 => {
+    return (trace as TraceDetailsV3)?.nodes !== undefined
+}

agenta-web/src/lib/transformers.ts

@@ -129,87 +129,6 @@ export const singleModelTestEvaluationTransformer = ({
     variant_revision_ids: item.variant_revision_ids,
 })
 
-export const fromBaseResponseToTraceSpanType = (
-    spans: BaseResponseSpans[],
-    traceId: string,
-): [TraceSpan[], Record<string, TraceSpan>] => {
-    const TRACE_DEFAULT_KEY = "__default__"
-
-    const all_spans = spans.map((span) => ({
-        id: span.id,
-        name: span.name,
-        created_at: span.start_time,
-        variant: {
-            variant_id: span.variant_id || null,
-            variant_name: span.variant_name || null,
-
-            revision:
-                span.environment == "playground"
-                    ? null
-                    : span.config
-                      ? span.config?.revision
-                      : null,
-        },
-        environment: span.environment || null,
-        spankind: span.spankind,
-        status: span.status,
-        metadata: {
-            cost: span.cost,
-            latency:
-                (new Date(span.end_time).getTime() - new Date(span.start_time).getTime()) / 1000,
-            usage: span.tokens,
-        },
-        content: {
-            inputs: span.inputs,
-            internals: span.internals,
-            outputs:
-                Array.isArray(span.outputs) ||
-                span.outputs == undefined ||
-                !span.outputs.hasOwnProperty(TRACE_DEFAULT_KEY)
-                    ? span.outputs
-                    : span.outputs[TRACE_DEFAULT_KEY],
-        } as {
-            inputs: Record<string, any> | null
-            internals: Record<string, any> | null
-            outputs: string[] | Record<string, any> | null
-        },
-        user_id: span.user,
-        trace_id: traceId,
-        parent_span_id: span.parent_span_id,
-
-        children: null,
-    }))
-
-    let spans_dict: Record<string, TraceSpan> = {}
-    for (const span of all_spans) {
-        spans_dict[span.id] = span
-    }
-
-    let child_spans: Array<string> = []
-
-    for (const span of all_spans) {
-        if (span.parent_span_id) {
-            const parent_span: TraceSpan = spans_dict[span.parent_span_id]
-            const child_span: TraceSpan = spans_dict[span.id]
-
-            if (parent_span) {
-                if (parent_span?.children === null) {
-                    parent_span.children = []
-                }
-
-                parent_span.children?.push(child_span)
-                child_spans.push(child_span.id)
-            }
-        }
-    }
-
-    const top_level_spans: Array<TraceSpan> = all_spans.filter(
-        (span) => !child_spans.includes(span.id),
-    )
-
-    return [top_level_spans, spans_dict]
-}
-
 export const transformTraceTreeToJson = (tree: TraceSpan[]) => {
     const nodeMap: Record<string, any> = {}