fix (ai-o11y): use ai endpoint in query aggregate suggestion (#13521)

#### Description

- AI queries now suggest the aggregates in the aggregation input with
/ai endpoint
- Keys come from `useFieldKeysSuggestion`, which routes AI queries to
`/api/v1/ai_observability/fields/keys`, replacing the deprecated
`getKeySuggestions`. AI queries send the same `fieldDataType=number`
filter as other query types.

#### Issues closed by this PR


https://github.com/orgs/SigNoz/projects/59/views/18?pane=issue&itemId=257298696&issue=SigNoz%7Cnerve-pod%7C377


#### Screenshots / Screen Recordings



https://github.com/user-attachments/assets/a0493d43-9c1b-4768-8ce8-4d4a67a67f25


#### Additional Information

- The `trace.*` suggestions need the backend to return the trace
aggregates for `fieldDataType=number`. Until it does, AI queries suggest
numeric span keys only.

---------

Co-authored-by: Gaurav Tewari <tewarig@users.noreply.github.com>

frontend/src/components/QueryBuilderV2/QueryV2/QueryAggregation/QueryAggregationSelect.tsx

@@ -1,7 +1,6 @@
 /* eslint-disable import/no-extraneous-dependencies */
 /* eslint-disable sonarjs/cognitive-complexity */
 import { useCallback, useEffect, useMemo, useRef, useState } from 'react';
-import { useQuery } from 'react-query';
 import {
 	autocompletion,
 	closeCompletion,
@@ -24,14 +23,16 @@ import CodeMirror, {
 	ViewUpdate,
 } from '@uiw/react-codemirror';
 import { Button, Popover, Tooltip } from 'antd';
-import { getKeySuggestions } from 'api/querySuggestions/getKeySuggestions';
-import { QueryBuilderKeys } from 'constants/queryBuilder';
+import { TelemetrytypesFieldDataTypeDTO } from 'api/generated/services/sigNoz.schemas';
 import { tracesAggregateOperatorOptions } from 'constants/queryBuilderOperators';
+import { useFieldKeysSuggestion } from 'hooks/querySuggestions/useFieldKeysSuggestion';
 import { useIsDarkMode } from 'hooks/useDarkMode';
 import { Info, TriangleAlert } from '@signozhq/icons';
 import { IBuilderQuery } from 'types/api/queryBuilder/queryBuilderData';
-import { FieldDataType } from 'types/api/v5/queryRange';
-import { TracesAggregatorOperator } from 'types/common/queryBuilder';
+import {
+	DATA_SOURCE_TO_SIGNAL,
+	TracesAggregatorOperator,
+} from 'types/common/queryBuilder';
 
 import { useQueryBuilderV2Context } from '../../QueryBuilderV2Context';
 
@@ -41,6 +42,12 @@ const chipDecoration = Decoration.mark({
 	class: 'chip-decorator',
 });
 
+const OPERATORS_WITHOUT_DATA_TYPE: string[] = [
+	TracesAggregatorOperator.COUNT,
+	TracesAggregatorOperator.COUNT_DISTINCT,
+	TracesAggregatorOperator.RATE,
+];
+
 const operatorArgMeta: Record<
 	string,
 	{ acceptsArgs: boolean; multiple: boolean }
@@ -310,31 +317,19 @@ function QueryAggregationSelect({
 	// Find function context for fetching suggestions
 	const functionContextForFetch = getFunctionContextAtCursor(input, cursorPos);
 
-	const { data: aggregateAttributeData, isLoading: isLoadingFields } = useQuery(
-		[
-			QueryBuilderKeys.GET_AGGREGATE_ATTRIBUTE,
-			functionContextForFetch,
-			queryData.dataSource,
-		],
-		() => {
-			const operatorsWithoutDataType: (string | undefined)[] = [
-				TracesAggregatorOperator.COUNT,
-				TracesAggregatorOperator.COUNT_DISTINCT,
-				TracesAggregatorOperator.RATE,
-			];
-
-			const fieldDataType: FieldDataType | undefined =
-				functionContextForFetch &&
-				operatorsWithoutDataType.includes(functionContextForFetch)
-					? undefined
-					: 'number';
-
-			return getKeySuggestions({
-				signal: queryData.dataSource,
-				searchText: '',
-				fieldDataType,
-			});
+	const numericKeysOnly = !(
+		functionContextForFetch &&
+		OPERATORS_WITHOUT_DATA_TYPE.includes(functionContextForFetch)
+	);
+	const { data: fieldKeys, isLoading: isLoadingFields } = useFieldKeysSuggestion(
+		{
+			signal: DATA_SOURCE_TO_SIGNAL[queryData.dataSource],
+			searchText: '',
+			fieldDataType: numericKeysOnly
+				? TelemetrytypesFieldDataTypeDTO.number
+				: undefined,
 		},
+		queryData.builderQueryType,
 		{
 			enabled:
 				!!functionContextForFetch &&
@@ -427,48 +422,42 @@ function QueryAggregationSelect({
 		}),
 	);
 
-	// Memoize field suggestions from API (no filtering here)
+	// Memoize field suggestions from API (no search filtering here)
 	const fieldSuggestions = useMemo(
 		() =>
-			Object.keys(aggregateAttributeData?.data.data.keys || {}).flatMap((key) => {
-				const attributeKeys = aggregateAttributeData?.data.data.keys[key];
-				if (!attributeKeys) {
-					return [];
-				}
+			(fieldKeys ?? []).map((key) => ({
+				// The backend reads trace.-prefixed columns as per-trace aggregates.
+				label: key.fieldContext === 'trace' ? `trace.${key.name}` : key.name,
+				type: 'variable',
+				info: key.fieldDataType,
+				apply: (
+					view: EditorView,
+					completion: Completion,
+					from: number,
+					to: number,
+				): void => {
+					const text = view.state.sliceDoc(0, from);
+					const funcName = getFunctionContextAtCursor(text, from);
+					const multiple = funcName ? operatorArgMeta[funcName]?.multiple : false;
 
-				return attributeKeys.map((attributeKey) => ({
-					label: attributeKey.name,
-					type: 'variable',
-					info: attributeKey.fieldDataType,
-					apply: (
-						view: EditorView,
-						completion: Completion,
-						from: number,
-						to: number,
-					): void => {
-						const text = view.state.sliceDoc(0, from);
-						const funcName = getFunctionContextAtCursor(text, from);
-						const multiple = funcName ? operatorArgMeta[funcName]?.multiple : false;
+					// Insert the selected key followed by either a comma or closing parenthesis
+					const insertText = multiple
+						? `${completion.label},`
+						: `${completion.label}) `;
+					const cursorPos = from + insertText.length; // Use insertText.length instead of hardcoded values
 
-						// Insert the selected key followed by either a comma or closing parenthesis
-						const insertText = multiple
-							? `${completion.label},`
-							: `${completion.label}) `;
-						const cursorPos = from + insertText.length; // Use insertText.length instead of hardcoded values
+					view.dispatch({
+						changes: { from, to, insert: insertText },
+						selection: { anchor: cursorPos },
+					});
 
-						view.dispatch({
-							changes: { from, to, insert: insertText },
-							selection: { anchor: cursorPos },
-						});
-
-						// Trigger next suggestions after a small delay
-						setTimeout(() => {
-							safeStartCompletion();
-						}, 50);
-					},
-				}));
-			}) || [],
-		[aggregateAttributeData, safeStartCompletion],
+					// Trigger next suggestions after a small delay
+					setTimeout(() => {
+						safeStartCompletion();
+					}, 50);
+				},
+			})),
+		[fieldKeys, safeStartCompletion],
 	);
 
 	const aggregatorAutocomplete = useMemo(

frontend/src/components/QueryBuilderV2/QueryV2/__tests__/QueryAggregationSelect.test.tsx

@@ -0,0 +1,168 @@
+import {
+	GetAIObservabilityFieldsKeys200,
+	GetFieldsKeys200,
+	TelemetrytypesFieldContextDTO,
+	TelemetrytypesFieldDataTypeDTO,
+	TelemetrytypesSignalDTO,
+	TelemetrytypesTelemetryFieldKeyDTO,
+} from 'api/generated/services/sigNoz.schemas';
+import { ENVIRONMENT } from 'constants/env';
+import { initialQueriesMap } from 'constants/queryBuilder';
+import { rest, server } from 'mocks-server/server';
+import { render, screen, userEvent, waitFor } from 'tests/test-utils';
+import type { IBuilderQuery } from 'types/api/queryBuilder/queryBuilderData';
+
+import { QueryBuilderV2Provider } from '../../QueryBuilderV2Context';
+import QueryAggregationSelect from '../QueryAggregation/QueryAggregationSelect';
+import { mockCodeMirrorDomApis } from './codemirrorDomMocks';
+
+const EDITOR_TEST_ID = 'query-aggregation-editor';
+// CodeMirror's suggestion popup is slow to open under a parallel jest run.
+const SUGGESTIONS_TIMEOUT_MS = 5000;
+const TEST_TIMEOUT_MS = 15000;
+
+const fieldKey = (
+	name: string,
+	fieldContext: TelemetrytypesFieldContextDTO,
+	fieldDataType: TelemetrytypesFieldDataTypeDTO,
+): TelemetrytypesTelemetryFieldKeyDTO => ({
+	name,
+	fieldContext,
+	fieldDataType,
+	signal: TelemetrytypesSignalDTO.traces,
+});
+
+const durationNano = fieldKey(
+	'duration_nano',
+	TelemetrytypesFieldContextDTO.span,
+	TelemetrytypesFieldDataTypeDTO.float64,
+);
+
+const serviceName = fieldKey(
+	'service.name',
+	TelemetrytypesFieldContextDTO.resource,
+	TelemetrytypesFieldDataTypeDTO.string,
+);
+
+const totalTokens = fieldKey(
+	'total_tokens',
+	TelemetrytypesFieldContextDTO.trace,
+	TelemetrytypesFieldDataTypeDTO.float64,
+);
+
+const aiKeysResponse = (
+	keys: TelemetrytypesTelemetryFieldKeyDTO[],
+): GetAIObservabilityFieldsKeys200 => ({
+	status: 'success',
+	data: {
+		complete: true,
+		keys: Object.fromEntries(keys.map((key) => [key.name, [key]])),
+	},
+});
+
+// Assumes the AI endpoint's number filter also returns the float64 trace aggregates.
+const aiKeysFor = (
+	params: URLSearchParams,
+): TelemetrytypesTelemetryFieldKeyDTO[] =>
+	params.get('fieldDataType') === TelemetrytypesFieldDataTypeDTO.number
+		? [durationNano, totalTokens]
+		: [durationNano, serviceName, totalTokens];
+
+const keysResponse: GetFieldsKeys200 = {
+	status: 'success',
+	data: { complete: true, keys: { duration_nano: [durationNano] } },
+};
+
+const aiQuery: IBuilderQuery = {
+	...(initialQueriesMap.traces.builder.queryData[0] as IBuilderQuery),
+	builderQueryType: 'builder_ai_query',
+	aggregations: [],
+};
+
+const aiKeysRequests: URLSearchParams[] = [];
+const keysRequests: URLSearchParams[] = [];
+
+beforeAll(() => {
+	mockCodeMirrorDomApis();
+});
+
+beforeEach(() => {
+	aiKeysRequests.length = 0;
+	keysRequests.length = 0;
+	server.use(
+		rest.get(
+			`${ENVIRONMENT.baseURL}/api/v1/ai_observability/fields/keys`,
+			(req, res, ctx) => {
+				aiKeysRequests.push(req.url.searchParams);
+				return res(
+					ctx.status(200),
+					ctx.json(aiKeysResponse(aiKeysFor(req.url.searchParams))),
+				);
+			},
+		),
+		rest.get(`${ENVIRONMENT.baseURL}/api/v1/fields/keys`, (req, res, ctx) => {
+			keysRequests.push(req.url.searchParams);
+			return res(ctx.status(200), ctx.json(keysResponse));
+		}),
+	);
+});
+
+async function enterAggregation(text: string): Promise<void> {
+	const user = userEvent.setup({ pointerEventsCheck: 0 });
+	render(
+		<QueryBuilderV2Provider>
+			<QueryAggregationSelect queryData={aiQuery} />
+		</QueryBuilderV2Provider>,
+	);
+
+	const content = (
+		await screen.findByTestId(EDITOR_TEST_ID)
+	).querySelector<HTMLElement>('.cm-content');
+	if (!content) {
+		throw new Error('CodeMirror content element not rendered');
+	}
+	await user.click(content);
+	// Typed keys drop under load: a re-render's dispatch discards DOM input CodeMirror hasn't read yet.
+	await user.paste(text);
+	await user.keyboard('{Control>} {/Control}');
+}
+
+function suggestionLabels(): string[] {
+	return Array.from(
+		screen.getByTestId(EDITOR_TEST_ID).querySelectorAll('.cm-completionLabel'),
+		(label) => label.textContent ?? '',
+	).sort();
+}
+
+describe('QueryAggregationSelect field suggestions for an AI query', () => {
+	it.each([
+		{
+			aggregation: 'avg(',
+			fieldDataType: TelemetrytypesFieldDataTypeDTO.number,
+			labels: ['duration_nano', 'trace.total_tokens'],
+		},
+		{
+			aggregation: 'count_distinct(',
+			fieldDataType: null,
+			labels: ['duration_nano', 'service.name', 'trace.total_tokens'],
+		},
+	])(
+		'suggests keys from the AI keys endpoint, with trace aggregates as trace.<name>, for $aggregation',
+		async ({ aggregation, fieldDataType, labels }) => {
+			await enterAggregation(aggregation);
+
+			await waitFor(() => expect(suggestionLabels()).toStrictEqual(labels), {
+				timeout: SUGGESTIONS_TIMEOUT_MS,
+			});
+			expect(keysRequests).toHaveLength(0);
+			expect(
+				aiKeysRequests.map((params) => ({
+					signal: params.get('signal'),
+					fieldDataType: params.get('fieldDataType'),
+					fieldContext: params.get('fieldContext'),
+				})),
+			).toStrictEqual([{ signal: 'traces', fieldDataType, fieldContext: null }]);
+		},
+		TEST_TIMEOUT_MS,
+	);
+});