Skip to content

Commit 6717ba0

Browse files
obostjancicobostjancic
authored
feat(agents-insights): token usage widget (#92253)
1 parent 234f4ba commit 6717ba0

File tree

5 files changed

+194
-34
lines changed

5 files changed

+194
-34
lines changed

‎static/app/views/insights/agentMonitoring/components/llmGenerationsWidget.tsx

Lines changed: 8 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -43,9 +43,8 @@ export default function LLMGenerationsWidget() {
4343

4444
const generationsRequest = useEAPSpans(
4545
{
46-
// @ts-expect-error TODO(telex): Add model id attribute to Fields
47-
fields: [AI_MODEL_ID_ATTRIBUTE, 'count(span.duration)'],
48-
sorts: [{field: 'count(span.duration)', kind: 'desc'}],
46+
fields: [AI_MODEL_ID_ATTRIBUTE, 'count()'],
47+
sorts: [{field: 'count()', kind: 'desc'}],
4948
search: fullQuery,
5049
limit: 3,
5150
},
@@ -71,12 +70,9 @@ export default function LLMGenerationsWidget() {
7170
const error = timeSeriesRequest.error || generationsRequest.error;
7271

7372
// TODO(telex): Add model id attribute to Fields and get rid of this cast
74-
const models = generationsRequest.data as unknown as
75-
| Array<{
76-
[AI_MODEL_ID_ATTRIBUTE]: string;
77-
'count(span.duration)': number;
78-
}>
79-
| undefined;
73+
const models = generationsRequest.data as unknown as Array<
74+
Record<string, string | number>
75+
>;
8076

8177
const hasData = models && models.length > 0 && timeSeries.length > 0;
8278

@@ -132,14 +128,16 @@ export default function LLMGenerationsWidget() {
132128
hasData && (
133129
<Toolbar
134130
exploreParams={{
135-
mode: Mode.SAMPLES,
131+
mode: Mode.AGGREGATE,
136132
visualize: [
137133
{
138134
chartType: ChartType.BAR,
139135
yAxes: ['count(span.duration)'],
140136
},
141137
],
138+
groupBy: [AI_MODEL_ID_ATTRIBUTE],
142139
query: fullQuery,
140+
sort: `-count(span.duration)`,
143141
interval: pageFilterChartParams.interval,
144142
}}
145143
onOpenFullScreen={() => {

‎static/app/views/insights/agentMonitoring/components/tokenUsageWidget.tsx

Lines changed: 169 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,169 @@
1+
import {Fragment} from 'react';
2+
import {useTheme} from '@emotion/react';
3+
import styled from '@emotion/styled';
4+
5+
import {openInsightChartModal} from 'sentry/actionCreators/modal';
6+
import {t} from 'sentry/locale';
7+
import {formatAbbreviatedNumber} from 'sentry/utils/formatters';
8+
import useOrganization from 'sentry/utils/useOrganization';
9+
import {Bars} from 'sentry/views/dashboards/widgets/timeSeriesWidget/plottables/bars';
10+
import {TimeSeriesWidgetVisualization} from 'sentry/views/dashboards/widgets/timeSeriesWidget/timeSeriesWidgetVisualization';
11+
import {Widget} from 'sentry/views/dashboards/widgets/widget/widget';
12+
import {Mode} from 'sentry/views/explore/contexts/pageParamsContext/mode';
13+
import {
14+
AI_MODEL_ID_ATTRIBUTE,
15+
AI_TOKEN_USAGE_ATTRIBUTE_SUM,
16+
getLLMGenerationsFilter,
17+
} from 'sentry/views/insights/agentMonitoring/utils/query';
18+
import {ChartType} from 'sentry/views/insights/common/components/chart';
19+
import {useEAPSpans} from 'sentry/views/insights/common/queries/useDiscover';
20+
import {useTopNSpanEAPSeries} from 'sentry/views/insights/common/queries/useTopNDiscoverSeries';
21+
import {convertSeriesToTimeseries} from 'sentry/views/insights/common/utils/convertSeriesToTimeseries';
22+
import {Referrer} from 'sentry/views/insights/pages/platform/laravel/referrers';
23+
import {usePageFilterChartParams} from 'sentry/views/insights/pages/platform/laravel/utils';
24+
import {WidgetVisualizationStates} from 'sentry/views/insights/pages/platform/laravel/widgetVisualizationStates';
25+
import {
26+
ModalChartContainer,
27+
ModalTableWrapper,
28+
SeriesColorIndicator,
29+
WidgetFooterTable,
30+
} from 'sentry/views/insights/pages/platform/shared/styles';
31+
import {Toolbar} from 'sentry/views/insights/pages/platform/shared/toolbar';
32+
import {useTransactionNameQuery} from 'sentry/views/insights/pages/platform/shared/useTransactionNameQuery';
33+
import {TimeSpentInDatabaseWidgetEmptyStateWarning} from 'sentry/views/performance/landing/widgets/components/selectableList';
34+
35+
export default function TokenUsageWidget() {
36+
const theme = useTheme();
37+
const organization = useOrganization();
38+
const {query} = useTransactionNameQuery();
39+
const pageFilterChartParams = usePageFilterChartParams({
40+
granularity: 'spans-low',
41+
});
42+
43+
const fullQuery = `${getLLMGenerationsFilter()} ${query}`.trim();
44+
45+
const tokensRequest = useEAPSpans(
46+
{
47+
fields: [AI_MODEL_ID_ATTRIBUTE, AI_TOKEN_USAGE_ATTRIBUTE_SUM],
48+
sorts: [{field: AI_TOKEN_USAGE_ATTRIBUTE_SUM, kind: 'desc'}],
49+
search: fullQuery,
50+
limit: 3,
51+
},
52+
Referrer.QUERIES_CHART // TODO
53+
);
54+
55+
const timeSeriesRequest = useTopNSpanEAPSeries(
56+
{
57+
...pageFilterChartParams,
58+
search: fullQuery,
59+
fields: [AI_MODEL_ID_ATTRIBUTE, AI_TOKEN_USAGE_ATTRIBUTE_SUM],
60+
yAxis: [AI_TOKEN_USAGE_ATTRIBUTE_SUM],
61+
sort: {field: AI_TOKEN_USAGE_ATTRIBUTE_SUM, kind: 'desc'},
62+
topN: 3,
63+
enabled: !!tokensRequest.data,
64+
},
65+
Referrer.QUERIES_CHART // TODO
66+
);
67+
68+
const timeSeries = timeSeriesRequest.data.filter(ts => ts.seriesName !== 'Other');
69+
70+
const isLoading = timeSeriesRequest.isLoading || tokensRequest.isLoading;
71+
const error = timeSeriesRequest.error || tokensRequest.error;
72+
73+
const tokens = tokensRequest.data as unknown as
74+
| Array<Record<string, string | number>>
75+
| undefined;
76+
77+
const hasData = tokens && tokens.length > 0 && timeSeries.length > 0;
78+
79+
const colorPalette = theme.chart.getColorPalette(timeSeries.length - 2);
80+
81+
const visualization = (
82+
<WidgetVisualizationStates
83+
isEmpty={!hasData}
84+
isLoading={isLoading}
85+
error={error}
86+
emptyMessage={<TimeSpentInDatabaseWidgetEmptyStateWarning />}
87+
VisualizationType={TimeSeriesWidgetVisualization}
88+
visualizationProps={{
89+
showLegend: 'never',
90+
plottables: timeSeries.map(
91+
(ts, index) =>
92+
new Bars(convertSeriesToTimeseries(ts), {
93+
color: colorPalette[index],
94+
alias: ts.seriesName,
95+
stack: 'stack',
96+
})
97+
),
98+
}}
99+
/>
100+
);
101+
102+
const footer = hasData && (
103+
<WidgetFooterTable>
104+
{tokens?.map((item, index) => (
105+
<Fragment key={item[AI_MODEL_ID_ATTRIBUTE]}>
106+
<div>
107+
<SeriesColorIndicator
108+
style={{
109+
backgroundColor: colorPalette[index],
110+
}}
111+
/>
112+
</div>
113+
<div>
114+
<ModelText>{item[AI_MODEL_ID_ATTRIBUTE]}</ModelText>
115+
</div>
116+
<span>{formatAbbreviatedNumber(item['sum(ai.total_tokens.used)'] || 0)}</span>
117+
</Fragment>
118+
))}
119+
</WidgetFooterTable>
120+
);
121+
122+
return (
123+
<Widget
124+
Title={<Widget.WidgetTitle title={t('Token Usage')} />}
125+
Visualization={visualization}
126+
Actions={
127+
organization.features.includes('visibility-explore-view') &&
128+
timeSeries && (
129+
<Toolbar
130+
exploreParams={{
131+
mode: Mode.AGGREGATE,
132+
visualize: [
133+
{
134+
chartType: ChartType.BAR,
135+
yAxes: ['sum(ai.total_tokens.used)'],
136+
},
137+
],
138+
groupBy: [AI_MODEL_ID_ATTRIBUTE],
139+
query: fullQuery,
140+
sort: `-${AI_TOKEN_USAGE_ATTRIBUTE_SUM}`,
141+
interval: pageFilterChartParams.interval,
142+
}}
143+
onOpenFullScreen={() => {
144+
openInsightChartModal({
145+
title: t('Token Usage'),
146+
children: (
147+
<Fragment>
148+
<ModalChartContainer>{visualization}</ModalChartContainer>
149+
<ModalTableWrapper>{footer}</ModalTableWrapper>
150+
</Fragment>
151+
),
152+
});
153+
}}
154+
/>
155+
)
156+
}
157+
noFooterPadding
158+
Footer={footer}
159+
/>
160+
);
161+
}
162+
163+
const ModelText = styled('div')`
164+
${p => p.theme.overflowEllipsis};
165+
color: ${p => p.theme.subText};
166+
font-size: ${p => p.theme.fontSizeSmall};
167+
line-height: 1.2;
168+
min-width: 0px;
169+
`;

‎static/app/views/insights/agentMonitoring/components/toolUsageWidget.tsx

Lines changed: 7 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -43,9 +43,8 @@ export default function ToolUsageWidget() {
4343

4444
const toolsRequest = useEAPSpans(
4545
{
46-
// @ts-expect-error TODO(telex): Add tool name attribute to Fields
47-
fields: [AI_TOOL_NAME_ATTRIBUTE, 'count(span.duration)'],
48-
sorts: [{field: 'count(span.duration)', kind: 'desc'}],
46+
fields: [AI_TOOL_NAME_ATTRIBUTE, 'count()'],
47+
sorts: [{field: 'count()', kind: 'desc'}],
4948
search: fullQuery,
5049
limit: 3,
5150
},
@@ -71,12 +70,7 @@ export default function ToolUsageWidget() {
7170
const error = timeSeriesRequest.error || toolsRequest.error;
7271

7372
// TODO(telex): Add tool name attribute to Fields and get rid of this cast
74-
const tools = toolsRequest.data as unknown as
75-
| Array<{
76-
[AI_TOOL_NAME_ATTRIBUTE]: string;
77-
'count(span.duration)': number;
78-
}>
79-
| undefined;
73+
const tools = toolsRequest.data as unknown as Array<Record<string, string | number>>;
8074

8175
const hasData = tools && tools.length > 0 && timeSeries.length > 0;
8276

@@ -115,7 +109,7 @@ export default function ToolUsageWidget() {
115109
/>
116110
</div>
117111
<div>
118-
<ModelText>{item[AI_TOOL_NAME_ATTRIBUTE]}</ModelText>
112+
<ToolText>{item[AI_TOOL_NAME_ATTRIBUTE]}</ToolText>
119113
</div>
120114
<span>{item['count(span.duration)']}</span>
121115
</Fragment>
@@ -139,7 +133,9 @@ export default function ToolUsageWidget() {
139133
yAxes: ['count(span.duration)'],
140134
},
141135
],
136+
groupBy: [AI_TOOL_NAME_ATTRIBUTE],
142137
query: fullQuery,
138+
sort: `-count(span.duration)`,
143139
interval: pageFilterChartParams.interval,
144140
}}
145141
onOpenFullScreen={() => {
@@ -162,7 +158,7 @@ export default function ToolUsageWidget() {
162158
);
163159
}
164160

165-
const ModelText = styled('div')`
161+
const ToolText = styled('div')`
166162
${p => p.theme.overflowEllipsis};
167163
color: ${p => p.theme.subText};
168164
font-size: ${p => p.theme.fontSizeSmall};

‎static/app/views/insights/agentMonitoring/utils/query.tsx

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,17 @@
11
// These are the span op we are currently ingesting.
2+
3+
import type {EAPSpanProperty} from 'sentry/views/insights/types';
4+
25
// They will probably change and maybe it will be enough to inline them in the widgets.
36
const AI_PIPELINE_OPS = ['ai.pipeline.generateText', 'ai.pipeline.generateObject'];
47
const AI_GENERATION_OPS = ['ai.run.doGenerate'];
58
const AI_TOOL_CALL_OPS = ['ai.toolCall'];
69

7-
export const AI_MODEL_ID_ATTRIBUTE = 'ai.model_id';
8-
export const AI_TOOL_NAME_ATTRIBUTE = 'ai.toolCall.name';
10+
export const AI_MODEL_ID_ATTRIBUTE = 'ai.model.id' as EAPSpanProperty;
11+
export const AI_TOOL_NAME_ATTRIBUTE = 'ai.toolCall.name' as EAPSpanProperty;
12+
const AI_TOKEN_USAGE_ATTRIBUTE = 'ai.total_tokens.used';
13+
14+
export const AI_TOKEN_USAGE_ATTRIBUTE_SUM = `sum(${AI_TOKEN_USAGE_ATTRIBUTE})`;
915

1016
export const getAgentRunsFilter = () => {
1117
return `span.op:[${AI_PIPELINE_OPS.join(',')}]`;

‎static/app/views/insights/agentMonitoring/views/agentsOverviewPage.tsx

Lines changed: 2 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -12,10 +12,10 @@ import {t} from 'sentry/locale';
1212
import {space} from 'sentry/styles/space';
1313
import {useLocation} from 'sentry/utils/useLocation';
1414
import useOrganization from 'sentry/utils/useOrganization';
15-
import {Widget} from 'sentry/views/dashboards/widgets/widget/widget';
1615
import {limitMaxPickableDays} from 'sentry/views/explore/utils';
1716
import LLMGenerationsWidget from 'sentry/views/insights/agentMonitoring/components/llmGenerationsWidget';
1817
import {ModelsTable} from 'sentry/views/insights/agentMonitoring/components/modelsTable';
18+
import TokenUsageWidget from 'sentry/views/insights/agentMonitoring/components/tokenUsageWidget';
1919
import {ToolsTable} from 'sentry/views/insights/agentMonitoring/components/toolsTable';
2020
import ToolUsageWidget from 'sentry/views/insights/agentMonitoring/components/toolUsageWidget';
2121
import {TracesTable} from 'sentry/views/insights/agentMonitoring/components/tracesTable';
@@ -96,7 +96,7 @@ function AgentsMonitoringPage() {
9696
<ToolUsageWidget />
9797
</WidgetGrid.Position5>
9898
<WidgetGrid.Position6>
99-
<PlaceholderWidget title="Token Usage" />
99+
<TokenUsageWidget />
100100
</WidgetGrid.Position6>
101101
</WidgetGrid>
102102
<ControlsWrapper>
@@ -145,15 +145,6 @@ const StyledTransactionNameSearchBar = styled(TransactionNameSearchBar)`
145145
flex: 2;
146146
`;
147147

148-
function PlaceholderWidget({title}: {title?: string}) {
149-
return (
150-
<Widget
151-
Title={<Widget.WidgetTitle title={title ?? 'Placeholder'} />}
152-
Visualization={null}
153-
/>
154-
);
155-
}
156-
157148
const ControlsWrapper = styled('div')`
158149
display: flex;
159150
justify-content: space-between;

0 commit comments

Comments
 (0)