fix(ai-o11y): return trace aggregate keys as number

frontend/src/container/LLMObservability/Explorer/TracesView/__tests__/TracesView.columns.test.tsx

@@ -33,7 +33,7 @@ jest.mock('../../TracesTable/TracesTable', () => ({
 			<button
 				data-testid="remove-total-tokens"
 				type="button"
-				onClick={(): void => onColumnRemove?.('trace:total_tokens:float64')}
+				onClick={(): void => onColumnRemove?.('trace:total_tokens:number')}
 			>
 				remove
 			</button>
@@ -48,7 +48,7 @@ const QUERY_RANGE_URL = `${ENVIRONMENT.baseURL}/api/v5/query_range`;
 const TOTAL_TOKENS = {
 	name: 'total_tokens',
 	fieldContext: 'trace',
-	fieldDataType: 'float64',
+	fieldDataType: 'number',
 };
 
 const storedColumnNames = (): string[] | undefined => {

frontend/src/container/LLMObservability/Explorer/TracesView/__tests__/TracesView.test.tsx

@@ -63,7 +63,7 @@ const mockFieldKeys = (names: string[]): void => {
 									{
 										name,
 										fieldContext: TelemetrytypesFieldContextDTO.trace,
-										fieldDataType: TelemetrytypesFieldDataTypeDTO.float64,
+										fieldDataType: TelemetrytypesFieldDataTypeDTO.number,
 									},
 								],
 							]),

frontend/src/container/LLMObservability/Explorer/constants.ts

@@ -73,7 +73,7 @@ export const TRACE_VIEW_DEFAULT_COLUMNS: TelemetryFieldKey[] = [
 		name: 'estimated_total_cost',
 		signal: 'traces',
 		fieldContext: 'trace',
-		fieldDataType: 'float64',
+		fieldDataType: 'number',
 	},
 	{ name: 'trace_duration_nano' },
 	{ name: 'span_count' },
@@ -81,37 +81,37 @@ export const TRACE_VIEW_DEFAULT_COLUMNS: TelemetryFieldKey[] = [
 		name: 'total_tokens',
 		signal: 'traces',
 		fieldContext: 'trace',
-		fieldDataType: 'float64',
+		fieldDataType: 'number',
 	},
 	{
 		name: 'input_tokens',
 		signal: 'traces',
 		fieldContext: 'trace',
-		fieldDataType: 'float64',
+		fieldDataType: 'number',
 	},
 	{
 		name: 'output_tokens',
 		signal: 'traces',
 		fieldContext: 'trace',
-		fieldDataType: 'float64',
+		fieldDataType: 'number',
 	},
 	{
 		name: 'distinct_tool_count',
 		signal: 'traces',
 		fieldContext: 'trace',
-		fieldDataType: 'float64',
+		fieldDataType: 'number',
 	},
 	{
 		name: 'llm_call_count',
 		signal: 'traces',
 		fieldContext: 'trace',
-		fieldDataType: 'float64',
+		fieldDataType: 'number',
 	},
 	{
 		name: 'tool_call_count',
 		signal: 'traces',
 		fieldContext: 'trace',
-		fieldDataType: 'float64',
+		fieldDataType: 'number',
 	},
 	{ name: 'start_time' },
 	{ name: 'end_time' },
@@ -122,7 +122,7 @@ export const TRACE_VIEW_DEFAULT_COLUMNS: TelemetryFieldKey[] = [
 		name: 'max_llm_duration_nano',
 		signal: 'traces',
 		fieldContext: 'trace',
-		fieldDataType: 'float64',
+		fieldDataType: 'number',
 	},
 ];
 

pkg/telemetrymetadata/static_fields.go

@@ -1,53 +1,14 @@
 package telemetrymetadata
 
 import (
-	"strings"
-
 	"github.com/SigNoz/signoz/pkg/types/telemetrytypes"
 )
 
 func staticFieldMatchesAny(field telemetrytypes.TelemetryFieldKey, selectors []*telemetrytypes.FieldKeySelector) bool {
 	for _, selector := range selectors {
-		if staticFieldMatches(field, selector) {
+		if selector.MatchesKey(&field) {
 			return true
 		}
 	}
 	return false
 }
-
-// staticFieldMatches mirrors the keys-table lookup for a static field: the
-// requested context and data type, when given, must agree with the field's,
-// and the name matches case-insensitively, as a substring for fuzzy selectors
-// and as the whole name for exact ones.
-func staticFieldMatches(field telemetrytypes.TelemetryFieldKey, selector *telemetrytypes.FieldKeySelector) bool {
-	if selector.FieldContext != telemetrytypes.FieldContextUnspecified && selector.FieldContext != field.FieldContext {
-		return false
-	}
-	if selector.FieldDataType != telemetrytypes.FieldDataTypeUnspecified && !sameDataTypeFamily(selector.FieldDataType, field.FieldDataType) {
-		return false
-	}
-	if selector.Name == "" {
-		return true
-	}
-	if selector.SelectorMatchType == telemetrytypes.FieldSelectorMatchTypeExact {
-		return strings.EqualFold(field.Name, selector.Name)
-	}
-	return strings.Contains(strings.ToLower(field.Name), strings.ToLower(selector.Name))
-}
-
-// sameDataTypeFamily treats the numeric types as one family: static fields
-// declare "number" while callers may ask for int64 or float64.
-func sameDataTypeFamily(requested, actual telemetrytypes.FieldDataType) bool {
-	if requested == actual {
-		return true
-	}
-	return isNumericDataType(requested) && isNumericDataType(actual)
-}
-
-func isNumericDataType(dataType telemetrytypes.FieldDataType) bool {
-	switch dataType {
-	case telemetrytypes.FieldDataTypeNumber, telemetrytypes.FieldDataTypeInt64, telemetrytypes.FieldDataTypeFloat64:
-		return true
-	}
-	return false
-}

pkg/telemetryschema/aitelemetryschema/const.go

@@ -54,6 +54,6 @@ func traceAggregate(name string) telemetrytypes.TelemetryFieldKey {
 		Name:          name,
 		Signal:        telemetrytypes.SignalTraces,
 		FieldContext:  telemetrytypes.FieldContextTrace,
-		FieldDataType: telemetrytypes.FieldDataTypeFloat64,
+		FieldDataType: telemetrytypes.FieldDataTypeNumber,
 	}
 }

pkg/telemetryschema/aitelemetryschema/field_keys_test.go

@@ -0,0 +1,40 @@
+package aitelemetryschema
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/SigNoz/signoz/pkg/types/telemetrytypes"
+)
+
+func TestFieldKeysDataTypeFilter(t *testing.T) {
+	testCases := []struct {
+		name          string
+		fieldDataType telemetrytypes.FieldDataType
+		expectTrace   bool
+	}{
+		{name: "Number_IncludesTraceAggregates", fieldDataType: telemetrytypes.FieldDataTypeNumber, expectTrace: true},
+		{name: "Float64_IncludesTraceAggregates", fieldDataType: telemetrytypes.FieldDataTypeFloat64, expectTrace: true},
+		{name: "Int64_IncludesTraceAggregates", fieldDataType: telemetrytypes.FieldDataTypeInt64, expectTrace: true},
+		{name: "String_ExcludesTraceAggregates", fieldDataType: telemetrytypes.FieldDataTypeString, expectTrace: false},
+	}
+
+	for _, testCase := range testCases {
+		t.Run(testCase.name, func(t *testing.T) {
+			selector := &telemetrytypes.FieldKeySelector{Signal: telemetrytypes.SignalTraces, FieldDataType: testCase.fieldDataType}
+			keys := FieldKeys(map[string][]*telemetrytypes.TelemetryFieldKey{}, selector)
+
+			for name := range TraceAggregateFields {
+				if !testCase.expectTrace {
+					assert.NotContains(t, keys, name)
+					continue
+				}
+				require.Len(t, keys[name], 1)
+				assert.Equal(t, telemetrytypes.FieldContextTrace, keys[name][0].FieldContext)
+				assert.Equal(t, telemetrytypes.FieldDataTypeNumber, keys[name][0].FieldDataType)
+			}
+		})
+	}
+}

pkg/types/telemetrytypes/field.go

@@ -260,7 +260,7 @@ func (s *FieldKeySelector) MatchesKey(key *TelemetryFieldKey) bool {
 		return false
 	}
 
-	if s.FieldDataType != FieldDataTypeUnspecified && s.FieldDataType != key.FieldDataType {
+	if s.FieldDataType != FieldDataTypeUnspecified && !sameDataTypeFamily(s.FieldDataType, key.FieldDataType) {
 		return false
 	}
 

pkg/types/telemetrytypes/field_datatype.go

@@ -198,3 +198,20 @@ func (FieldDataType) Enum() []any {
 		// FieldDataTypeArrayNumber,
 	}
 }
+
+// sameDataTypeFamily treats the numeric types as one family: static fields
+// declare "number" while callers may ask for int64 or float64.
+func sameDataTypeFamily(requested, actual FieldDataType) bool {
+	if requested == actual {
+		return true
+	}
+	return isNumericDataType(requested) && isNumericDataType(actual)
+}
+
+func isNumericDataType(dataType FieldDataType) bool {
+	switch dataType {
+	case FieldDataTypeNumber, FieldDataTypeInt64, FieldDataTypeFloat64:
+		return true
+	}
+	return false
+}

tests/integration/tests/querierai/03_ai_metadata.py

@@ -59,6 +59,24 @@ def test_ai_fields_trace_prefix_search(
     assert trace_context_names == {"output_tokens"}, keys
 
 
+@pytest.mark.parametrize("field_data_type", ["number", "float64", "int64"])
+def test_ai_fields_numeric_data_type_lists_aggregates(
+    signoz: types.SigNoz,
+    create_user_admin: None,  # pylint: disable=unused-argument
+    get_token: Callable[[str, str], str],
+    field_data_type: str,
+) -> None:
+    token = get_token(USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD)
+
+    response = get_field_keys(signoz, token, {"fieldDataType": field_data_type}, AI_KEYS_PATH)
+    assert response.status_code == HTTPStatus.OK, response.text
+
+    keys = response.json()["data"]["keys"]
+    trace_keys = [key for variants in keys.values() for key in variants if key["fieldContext"] == "trace"]
+    assert {key["name"] for key in trace_keys} == AI_TRACE_AGGREGATES, keys
+    assert {key["fieldDataType"] for key in trace_keys} == {"number"}, keys
+
+
 def test_ai_fields_bare_prefix_suggests_aggregate_and_attribute(
     signoz: types.SigNoz,
     create_user_admin: None,  # pylint: disable=unused-argument