import styled from '@emotion/styled';
import type {Location} from 'history';
import * as qs from 'query-string';
import GridEditable, {
COL_WIDTH_UNDEFINED,
type GridColumnHeader,
} from 'sentry/components/gridEditable';
import Link from 'sentry/components/links/link';
import type {CursorHandler} from 'sentry/components/pagination';
import Pagination from 'sentry/components/pagination';
import SearchBar from 'sentry/components/searchBar';
import {Tooltip} from 'sentry/components/tooltip';
import {IconInfo} from 'sentry/icons';
import {t} from 'sentry/locale';
import {space} from 'sentry/styles/space';
import type {Organization} from 'sentry/types/organization';
import {browserHistory} from 'sentry/utils/browserHistory';
import type {EventsMetaType} from 'sentry/utils/discover/eventView';
import {getFieldRenderer} from 'sentry/utils/discover/fieldRenderers';
import type {Sort} from 'sentry/utils/discover/fields';
import {RATE_UNIT_TITLE, RateUnit} from 'sentry/utils/discover/fields';
import {VisuallyCompleteWithData} from 'sentry/utils/performanceForSentry';
import {decodeScalar, decodeSorts} from 'sentry/utils/queryString';
import {MutableSearch} from 'sentry/utils/tokenizeSearch';
import {useLocation} from 'sentry/utils/useLocation';
import useOrganization from 'sentry/utils/useOrganization';
import {useModuleURL} from 'sentry/views/performance/utils/useModuleURL';
import {renderHeadCell} from 'sentry/views/starfish/components/tableCells/renderHeadCell';
import {useSpanMetrics} from 'sentry/views/starfish/queries/useDiscover';
import type {SpanMetricsResponse} from 'sentry/views/starfish/types';
import {QueryParameterNames} from 'sentry/views/starfish/views/queryParameters';
type Row = Pick<
SpanMetricsResponse,
| 'project.id'
| 'span.description'
| 'span.group'
| 'spm()'
| 'avg(span.duration)'
| 'sum(span.duration)'
| 'sum(ai.total_tokens.used)'
| 'sum(ai.total_cost)'
>;
type Column = GridColumnHeader<
| 'span.description'
| 'spm()'
| 'avg(span.duration)'
| 'sum(ai.total_tokens.used)'
| 'sum(ai.total_cost)'
>;
const COLUMN_ORDER: Column[] = [
{
key: 'span.description',
name: t('AI Pipeline Name'),
width: COL_WIDTH_UNDEFINED,
},
{
key: 'sum(ai.total_tokens.used)',
name: t('Total tokens used'),
width: 180,
},
{
key: 'sum(ai.total_cost)',
name: t('Total cost'),
width: 180,
},
{
key: `avg(span.duration)`,
name: t('Pipeline Duration'),
width: COL_WIDTH_UNDEFINED,
},
{
key: 'spm()',
name: `${t('Pipeline runs')} ${RATE_UNIT_TITLE[RateUnit.PER_MINUTE]}`,
width: COL_WIDTH_UNDEFINED,
},
];
const SORTABLE_FIELDS = ['sum(ai.total_tokens.used)', 'avg(span.duration)', 'spm()'];
type ValidSort = Sort & {
field: 'spm()' | 'avg(span.duration)';
};
export function isAValidSort(sort: Sort): sort is ValidSort {
return (SORTABLE_FIELDS as unknown as string[]).includes(sort.field);
}
export function PipelinesTable() {
const location = useLocation();
const moduleURL = useModuleURL('ai');
const organization = useOrganization();
const cursor = decodeScalar(location.query?.[QueryParameterNames.SPANS_CURSOR]);
const sortField = decodeScalar(location.query?.[QueryParameterNames.SPANS_SORT]);
const spanDescription = decodeScalar(location.query?.['span.description'], '');
let sort = decodeSorts(sortField).filter(isAValidSort)[0];
if (!sort) {
sort = {field: 'spm()', kind: 'desc'};
}
const {data, isLoading, meta, pageLinks, error} = useSpanMetrics(
{
search: MutableSearch.fromQueryObject({
'span.category': 'ai.pipeline',
'span.description': spanDescription ? `*${spanDescription}*` : undefined,
}),
fields: [
'project.id',
'span.group',
'span.description',
'spm()',
'avg(span.duration)',
'sum(span.duration)',
],
sorts: [sort],
limit: 25,
cursor,
},
'api.ai-pipelines.view'
);
const {data: tokensUsedData, isLoading: tokensUsedLoading} = useSpanMetrics(
{
search: new MutableSearch(
`span.category:ai span.ai.pipeline.group:[${(data as Row[])?.map(x => x['span.group']).join(',')}]`
),
fields: [
'span.ai.pipeline.group',
'sum(ai.total_tokens.used)',
'sum(ai.total_cost)',
],
},
'api.performance.ai-analytics.token-usage-chart'
);
const rows: Row[] = (data as Row[]).map(baseRow => {
const row: Row = {
...baseRow,
'sum(ai.total_tokens.used)': 0,
'sum(ai.total_cost)': 0,
};
if (!tokensUsedLoading) {
const tokenUsedDataPoint = tokensUsedData.find(
tokenRow => tokenRow['span.ai.pipeline.group'] === row['span.group']
);
if (tokenUsedDataPoint) {
row['sum(ai.total_tokens.used)'] =
tokenUsedDataPoint['sum(ai.total_tokens.used)'];
row['sum(ai.total_cost)'] = tokenUsedDataPoint['sum(ai.total_cost)'];
}
}
return row;
});
const handleCursor: CursorHandler = (newCursor, pathname, query) => {
browserHistory.push({
pathname,
query: {...query, [QueryParameterNames.SPANS_CURSOR]: newCursor},
});
};
const handleSearch = (newQuery: string) => {
browserHistory.push({
...location,
query: {
...location.query,
'span.description': newQuery === '' ? undefined : newQuery,
[QueryParameterNames.SPANS_CURSOR]: undefined,
},
});
};
return (
0}
isLoading={isLoading}
>
renderHeadCell({
column,
sort,
location,
sortParameterName: QueryParameterNames.SPANS_SORT,
}),
renderBodyCell: (column, row) =>
renderBodyCell(moduleURL, column, row, meta, location, organization),
}}
/>
);
}
function renderBodyCell(
moduleURL: string,
column: Column,
row: Row,
meta: EventsMetaType | undefined,
location: Location,
organization: Organization
) {
if (column.key === 'span.description') {
if (!row['span.description']) {
return (unknown);
}
if (!row['span.group']) {
return {row['span.description']};
}
const queryString = {
...location.query,
'span.description': row['span.description'],
};
return (
{row['span.description']}
);
}
if (column.key === 'sum(ai.total_cost)') {
const cost = row[column.key];
if (cost) {
return US ${cost.toFixed(3)};
}
return (
Unknown{' '}
);
}
if (!meta || !meta?.fields) {
return row[column.key];
}
const renderer = getFieldRenderer(column.key, meta.fields, false);
const rendered = renderer(row, {
location,
organization,
unit: meta.units?.[column.key],
});
return rendered;
}
const Container = styled('div')`
display: flex;
flex-direction: column;
gap: ${space(1)};
`;