From 13110dcc78c192b29c8bcd541d6e8ddedc7976ef Mon Sep 17 00:00:00 2001
From: Martin Mark <martin@martinmark.com>
Date: Thu, 12 Dec 2024 23:00:23 -0500
Subject: [PATCH] Revert changes to call stats

---
 .../PlaygroundChat/PlaygroundCallStats.tsx    | 131 ++++--------------
 1 file changed, 27 insertions(+), 104 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 f710ae08563..9888ea88681 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,99 +1,31 @@
 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, {useMemo} from 'react';
+import React from 'react';
 import {useHistory} from 'react-router-dom';
 
 import {useWeaveflowRouteContext} from '../../../context';
 import {Reactions} from '../../../feedback/Reactions';
-import {addCostsToCallResults} from '../../CallPage/cost/costUtils';
-import {TraceCostStats} from '../../CallPage/cost/TraceCostStats';
-import {useWFHooks} from '../../wfReactInterface/context';
 import {TraceCallSchema} from '../../wfReactInterface/traceServerClientTypes';
-import {CallSchema} from '../../wfReactInterface/wfDataModelHooksInterface';
-
-type StatusCode = 'OK' | 'ERROR';
-const traceCallStatusCode = (call: TraceCallSchema): StatusCode => {
-  return call.exception ? 'ERROR' : 'OK';
-};
 
 export const PlaygroundCallStats = ({call}: {call: TraceCallSchema}) => {
-  const {useCalls} = useWFHooks();
-  const [entityName, projectName] = call?.project_id?.split('/') || [];
-  const callId = call?.id || '';
-
-  // Convert TraceCallSchema to CallSchema
-  const callSchema: CallSchema = useMemo(
-    () => ({
-      entity: entityName,
-      project: projectName,
-      callId,
-      traceId: call.trace_id,
-      parentId: call.parent_id ?? null,
-      userId: call.wb_user_id ?? null,
-      runId: call.wb_run_id ?? null,
-      spanName: call.op_name ?? '',
-      displayName: call.display_name ?? null,
-      opVersionRef: call.op_name ?? null,
-      rawSpan: {
-        name: call.op_name,
-        inputs: call.inputs,
-        output: (call.output ?? {}) as {[key: string]: any; _keys?: string[]},
-        status_code: traceCallStatusCode(call),
-        exception: call.exception,
-        attributes: call.attributes,
-        summary: {
-          ...(call.summary ?? {}),
-          latency_s: (call.summary?.weave?.latency_ms ?? 0) / 1000,
-        },
-        span_id: call.id,
-        trace_id: call.trace_id,
-        parent_id: call.parent_id ?? undefined,
-        timestamp: new Date(call.started_at).getTime(),
-        start_time_ms: new Date(call.started_at).getTime(),
-        end_time_ms: call.ended_at
-          ? new Date(call.ended_at).getTime()
-          : undefined,
-      },
-      traceCall: call,
-    }),
-    [call, entityName, projectName, callId]
-  );
-
-  // Fetch cost data
-  const costCols = useMemo(() => ['id'], []);
-  const costs = useCalls(
-    entityName,
-    projectName,
-    {
-      callIds: [callId],
-    },
-    undefined,
-    undefined,
-    undefined,
-    undefined,
-    costCols,
-    undefined,
-    {
-      includeCosts: true,
-    }
-  );
-
-  // Merge cost data with the call
-  const callWithCosts = useMemo(() => {
-    if (!costs.result || costs.result.length === 0) {
-      return callSchema;
+  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 [updatedCall] = addCostsToCallResults([callSchema], costs.result);
-    return updatedCall;
-  }, [callSchema, costs.result]);
-
-  const latency = callWithCosts?.traceCall?.summary?.weave?.latency_ms ?? 0;
-  const usageData = callWithCosts?.traceCall?.summary?.usage;
-  const costData = callWithCosts?.traceCall?.summary?.weave?.costs;
+  }
 
+  const [entityName, projectName] = call?.project_id?.split('/') || [];
+  const callId = call?.id || '';
+  const latency = call?.summary?.weave?.latency_ms;
   const {peekingRouter} = useWeaveflowRouteContext();
   const history = useHistory();
 
@@ -113,28 +45,19 @@ export const PlaygroundCallStats = ({call}: {call: TraceCallSchema}) => {
 
   return (
     <Tailwind>
-      <div className="flex w-full items-center justify-center gap-8 py-8">
-        <TraceCostStats
-          usageData={usageData}
-          costData={costData}
-          latency_ms={latency}
-          costLoading={costs.loading}
-        />
+      <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>
         {(call.output as any)?.choices?.[0]?.finish_reason && (
-          <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>
+              Finish reason: {(call.output as any).choices[0].finish_reason}
+            </span>
+            <span>•</span>
+          </>
         )}
+        <span>{totalTokens} tokens</span>
+        <span>•</span>
         {callLink && (
           <Button
             size="small"