From a562949954bf639fb3d640e984dfdee9be5cee4c Mon Sep 17 00:00:00 2001
From: jwlee64 <josiah.lee@wandb.com>
Date: Thu, 12 Dec 2024 15:11:14 -0800
Subject: [PATCH 1/2] add icons to trace tree

---
 .../PlaygroundChat/PlaygroundCallStats.tsx    | 53 ++++++++++---------
 .../pages/PlaygroundPage/PlaygroundPage.tsx   | 24 +++++++--
 2 files changed, 47 insertions(+), 30 deletions(-)

diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/PlaygroundPage/PlaygroundChat/PlaygroundCallStats.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/PlaygroundPage/PlaygroundChat/PlaygroundCallStats.tsx
index 9888ea886819..f060d6859a77 100644
--- a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/PlaygroundPage/PlaygroundChat/PlaygroundCallStats.tsx
+++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/PlaygroundPage/PlaygroundChat/PlaygroundCallStats.tsx
@@ -1,31 +1,19 @@
 import {Button} from '@wandb/weave/components/Button';
+import {Pill} from '@wandb/weave/components/Tag';
 import {Tailwind} from '@wandb/weave/components/Tailwind';
+import {Tooltip} from '@wandb/weave/components/Tooltip';
 import {makeRefCall} from '@wandb/weave/util/refs';
 import React from 'react';
 import {useHistory} from 'react-router-dom';
 
 import {useWeaveflowRouteContext} from '../../../context';
 import {Reactions} from '../../../feedback/Reactions';
+import {TraceCostStats} from '../../CallPage/cost';
 import {TraceCallSchema} from '../../wfReactInterface/traceServerClientTypes';
 
 export const PlaygroundCallStats = ({call}: {call: TraceCallSchema}) => {
-  let totalTokens = 0;
-  if (call?.summary?.usage) {
-    for (const key of Object.keys(call.summary.usage)) {
-      totalTokens +=
-        call.summary.usage[key].prompt_tokens ||
-        call.summary.usage[key].input_tokens ||
-        0;
-      totalTokens +=
-        call.summary.usage[key].completion_tokens ||
-        call.summary.usage[key].output_tokens ||
-        0;
-    }
-  }
-
   const [entityName, projectName] = call?.project_id?.split('/') || [];
   const callId = call?.id || '';
-  const latency = call?.summary?.weave?.latency_ms;
   const {peekingRouter} = useWeaveflowRouteContext();
   const history = useHistory();
 
@@ -43,21 +31,34 @@ export const PlaygroundCallStats = ({call}: {call: TraceCallSchema}) => {
     false
   );
 
+  const latency = call?.summary?.weave?.latency_ms ?? 0;
+  const usageData = call?.summary?.usage;
+  const costData = call?.summary?.weave?.costs;
+
   return (
     <Tailwind>
-      <div className="flex w-full flex-wrap items-center justify-center gap-8 py-8 text-sm text-moon-500">
-        <span>Latency: {latency}ms</span>
-        <span>•</span>
+      <div className="flex w-full items-center justify-center gap-8 py-8">
+        <TraceCostStats
+          usageData={usageData}
+          costData={costData}
+          latency_ms={latency}
+          costLoading={false}
+        />
         {(call.output as any)?.choices?.[0]?.finish_reason && (
-          <>
-            <span>
-              Finish reason: {(call.output as any).choices[0].finish_reason}
-            </span>
-            <span>•</span>
-          </>
+          <Tooltip
+            content="Finish Reason"
+            trigger={
+              <span>
+                <Pill
+                  icon="checkmark-circle"
+                  label={(call.output as any).choices[0].finish_reason}
+                  color="moon"
+                  className="-ml-[8px] bg-transparent text-moon-500 dark:bg-transparent dark:text-moon-500"
+                />
+              </span>
+            }
+          />
         )}
-        <span>{totalTokens} tokens</span>
-        <span>•</span>
         {callLink && (
           <Button
             size="small"
diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/PlaygroundPage/PlaygroundPage.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/PlaygroundPage/PlaygroundPage.tsx
index c6232631e4eb..7a628e1f8efa 100644
--- a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/PlaygroundPage/PlaygroundPage.tsx
+++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/PlaygroundPage/PlaygroundPage.tsx
@@ -56,12 +56,28 @@ export const PlaygroundPageInner = (props: PlaygroundPageProps) => {
             callId: props.callId,
           }
         : null;
-    }, [props.entity, props.project, props.callId])
+    }, [props.entity, props.project, props.callId]),
+    {
+      includeCosts: true,
+    }
   );
 
-  const {result: calls} = useCalls(props.entity, props.project, {
-    callIds: playgroundStates.map(state => state.traceCall.id || ''),
-  });
+  const {result: calls} = useCalls(
+    props.entity,
+    props.project,
+    {
+      callIds: playgroundStates.map(state => state.traceCall.id || ''),
+    },
+    undefined,
+    undefined,
+    undefined,
+    undefined,
+    undefined,
+    undefined,
+    {
+      includeCosts: true,
+    }
+  );
 
   useEffect(() => {
     if (!call.loading && call.result) {

From 6d4d436d88b47b579d5cb53e7716b84b32209aba Mon Sep 17 00:00:00 2001
From: jwlee64 <josiah.lee@wandb.com>
Date: Fri, 13 Dec 2024 13:55:48 -0800
Subject: [PATCH 2/2] sentence case

---
 .../PlaygroundPage/PlaygroundChat/PlaygroundCallStats.tsx      | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/PlaygroundPage/PlaygroundChat/PlaygroundCallStats.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/PlaygroundPage/PlaygroundChat/PlaygroundCallStats.tsx
index f060d6859a77..d969a91136ed 100644
--- a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/PlaygroundPage/PlaygroundChat/PlaygroundCallStats.tsx
+++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/PlaygroundPage/PlaygroundChat/PlaygroundCallStats.tsx
@@ -46,8 +46,9 @@ export const PlaygroundCallStats = ({call}: {call: TraceCallSchema}) => {
         />
         {(call.output as any)?.choices?.[0]?.finish_reason && (
           <Tooltip
-            content="Finish Reason"
+            content="Finish reason"
             trigger={
+              // Placing in span so tooltip shows up
               <span>
                 <Pill
                   icon="checkmark-circle"