Skip to content

Commit

Permalink
Revert changes to call stats
Browse files Browse the repository at this point in the history
  • Loading branch information
m-rgba committed Dec 13, 2024
1 parent 3f74736 commit 13110dc
Showing 1 changed file with 27 additions and 104 deletions.
Original file line number Diff line number Diff line change
@@ -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();

Expand All @@ -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"
Expand Down

0 comments on commit 13110dc

Please sign in to comment.