fix: update function name

docs/api/openapi.yml

@@ -11020,6 +11020,10 @@ paths:
         name: source
         schema:
           $ref: '#/components/schemas/TelemetrytypesSource'
+      - in: query
+        name: type
+        schema:
+          type: string
       - in: query
         name: limit
         schema:
@@ -11109,6 +11113,10 @@ paths:
         name: source
         schema:
           $ref: '#/components/schemas/TelemetrytypesSource'
+      - in: query
+        name: type
+        schema:
+          type: string
       - in: query
         name: limit
         schema:
@@ -21311,6 +21319,10 @@ paths:
         name: source
         schema:
           $ref: '#/components/schemas/TelemetrytypesSource'
+      - in: query
+        name: type
+        schema:
+          type: string
       - in: query
         name: limit
         schema:
@@ -21412,6 +21424,10 @@ paths:
         name: source
         schema:
           $ref: '#/components/schemas/TelemetrytypesSource'
+      - in: query
+        name: type
+        schema:
+          type: string
       - in: query
         name: limit
         schema:

frontend/src/api/generated/services/sigNoz.schemas.ts

@@ -10426,6 +10426,11 @@ export type GetFieldsKeysParams = {
 	 * @description undefined
 	 */
 	source?: TelemetrytypesSourceDTO;
+	/**
+	 * @type string
+	 * @description undefined
+	 */
+	type?: string;
 	/**
 	 * @type integer
 	 * @description undefined
@@ -10485,6 +10490,11 @@ export type GetFieldsValuesParams = {
 	 * @description undefined
 	 */
 	source?: TelemetrytypesSourceDTO;
+	/**
+	 * @type string
+	 * @description undefined
+	 */
+	type?: string;
 	/**
 	 * @type integer
 	 * @description undefined
@@ -11792,6 +11802,11 @@ export type GetRuleHistoryFilterKeysParams = {
 	 * @description undefined
 	 */
 	source?: TelemetrytypesSourceDTO;
+	/**
+	 * @type string
+	 * @description undefined
+	 */
+	type?: string;
 	/**
 	 * @type integer
 	 * @description undefined
@@ -11854,6 +11869,11 @@ export type GetRuleHistoryFilterValuesParams = {
 	 * @description undefined
 	 */
 	source?: TelemetrytypesSourceDTO;
+	/**
+	 * @type string
+	 * @description undefined
+	 */
+	type?: string;
 	/**
 	 * @type integer
 	 * @description undefined

pkg/statementbuilder/aistatementbuilder/statement_builder.go

@@ -49,20 +49,21 @@ func Scope() scopedtraces.TraceScope {
 	str := telemetrytypes.FieldDataTypeString
 	columns := append(scopedtraces.CommonTraceColumns(),
 		// LLM calls only (request model present), not the full gate.
-		scopedtraces.TraceColumn{Alias: "llm_call_count", Orderable: true, Expr: scopedtraces.CountExists(&reqModel)},
-		scopedtraces.TraceColumn{Alias: "tool_call_count", Orderable: true, Expr: scopedtraces.CountExists(&toolName)},
-		scopedtraces.TraceColumn{Alias: "distinct_tool_count", Orderable: true, Expr: scopedtraces.UniqCount(&toolName, str)},
+		scopedtraces.TraceColumn{Alias: "llm_call_count", Orderable: true, Filterable: true, Expr: scopedtraces.CountExists(&reqModel)},
+		scopedtraces.TraceColumn{Alias: "tool_call_count", Orderable: true, Filterable: true, Expr: scopedtraces.CountExists(&toolName)},
+		scopedtraces.TraceColumn{Alias: "distinct_tool_count", Orderable: true, Filterable: true, Expr: scopedtraces.UniqCount(&toolName, str)},
 		// tokens live only on LLM spans, so a plain sum needs no gate scoping.
-		scopedtraces.TraceColumn{Alias: "input_tokens", Orderable: true, Expr: scopedtraces.Reduce(scopedtraces.AggSum, &inTok)},
-		scopedtraces.TraceColumn{Alias: "output_tokens", Orderable: true, Expr: scopedtraces.Reduce(scopedtraces.AggSum, &outTok)},
-		scopedtraces.TraceColumn{Alias: "total_tokens", Orderable: true, Expr: scopedtraces.SumOfKeys(telemetrytypes.FieldDataTypeFloat64, &inTok, &outTok)},
+		scopedtraces.TraceColumn{Alias: "input_tokens", Orderable: true, Filterable: true, Expr: scopedtraces.Reduce(scopedtraces.AggSum, &inTok)},
+		scopedtraces.TraceColumn{Alias: "output_tokens", Orderable: true, Filterable: true, Expr: scopedtraces.Reduce(scopedtraces.AggSum, &outTok)},
+		scopedtraces.TraceColumn{Alias: "total_tokens", Orderable: true, Filterable: true, Expr: scopedtraces.SumOfKeys(telemetrytypes.FieldDataTypeFloat64, &inTok, &outTok)},
 		// per-span cost attached by the SigNoz LLM pricing processor.
-		scopedtraces.TraceColumn{Alias: "estimated_total_cost", Orderable: true, Expr: scopedtraces.Reduce(scopedtraces.AggSum, &cost)},
+		scopedtraces.TraceColumn{Alias: "estimated_total_cost", Orderable: true, Filterable: true, Expr: scopedtraces.Reduce(scopedtraces.AggSum, &cost)},
 		// slowest single LLM call in the trace.
-		scopedtraces.TraceColumn{Alias: "max_llm_duration_nano", Orderable: true, Expr: scopedtraces.ScopedToKeyColumn(scopedtraces.AggMax, scopedtraces.IntrinsicSpanKey("duration_nano"), &reqModel)},
+		scopedtraces.TraceColumn{Alias: "max_llm_duration_nano", Orderable: true, Filterable: true, Expr: scopedtraces.ScopedToKeyColumn(scopedtraces.AggMax, scopedtraces.IntrinsicSpanKey("duration_nano"), &reqModel)},
 		// errors across the whole trace (any span), so display-only.
 		scopedtraces.TraceColumn{Alias: "error_count", Expr: scopedtraces.CondCount(scopedtraces.IntrinsicSpanKey("has_error"), qbtypes.FilterOperatorEqual, true)},
-		// timestamp of the last gen_ai span (LLM/tool/agent), hence gate-scoped.
+		// timestamp of the last gen_ai span (LLM/tool/agent), hence gate-scoped;
+		// order-only: a raw-nanos threshold makes no sense in the filter bar.
 		scopedtraces.TraceColumn{Alias: "last_activity_time", Orderable: true, Expr: scopedtraces.ScopedReduce(scopedtraces.AggMax, scopedtraces.IntrinsicSpanKey("timestamp"))},
 		// previews: first call's input (the prompt), last call's output (the answer).
 		scopedtraces.TraceColumn{Alias: "input", SpanLevel: true, Expr: scopedtraces.PickBy(&inMsg, str, scopedtraces.IntrinsicSpanKey("timestamp"), scopedtraces.PickEarliest)},
@@ -76,3 +77,22 @@ func Scope() scopedtraces.TraceScope {
 		DefaultOrderAlias: "last_activity_time",
 	}
 }
+
+// MetadataFieldKeys returns the aggregates the metadata store surfaces as trace-context
+// keys for builder_ai_query suggestions; only filterable columns qualify.
+func MetadataFieldKeys() []*telemetrytypes.TelemetryFieldKey {
+	cols := Scope().Columns
+	keys := make([]*telemetrytypes.TelemetryFieldKey, 0, len(cols))
+	for _, c := range cols {
+		if !c.Orderable || !c.Filterable {
+			continue
+		}
+		keys = append(keys, &telemetrytypes.TelemetryFieldKey{
+			Name:          c.Alias,
+			Signal:        telemetrytypes.SignalTraces,
+			FieldContext:  telemetrytypes.FieldContextTrace,
+			FieldDataType: telemetrytypes.FieldDataTypeFloat64,
+		})
+	}
+	return keys
+}

pkg/statementbuilder/scopedtracesstatementbuilder/scope.go

@@ -23,10 +23,9 @@ type TraceColumn struct {
 	// Alias must not reuse a physical span-index column name (e.g. duration_nano):
 	// ClickHouse resolves bare identifiers to same-SELECT aliases first, so any
 	// expression referencing that column would silently bind to the alias.
-	Alias string
-	// Orderable columns can be used in ORDER BY and the aggregate filter; all-span
-	// aggregates are display-only and set false.
-	Orderable bool
+	Alias      string
+	Orderable  bool
+	Filterable bool
 	// SpanLevel columns surface a real span/resource attribute; a filter on them is
 	// applied span-level, so they are excluded from the trace-level aliases.
 	SpanLevel bool

pkg/statementbuilder/scopedtracesstatementbuilder/statement_builder.go

@@ -185,18 +185,19 @@ func (b *scopedTraceStatementBuilder) buildTraceListQuery(
 		return nil, err
 	}
 	orderableSet := orderableAliasSet(resolved)
+	filterableSet := filterableAliasSet(resolved)
 
 	resourceFrag, resourceArgs, resourcePred, err := b.maybeAttachResourceFilter(ctx, orgID, query, start, end, variables)
 	if err != nil {
 		return nil, err
 	}
 
-	fp, err := b.splitFilter(ctx, orgID, query, b.aggregateAliasSet(), orderableSet, start, end, variables, matchedSB)
+	fp, err := b.splitFilter(ctx, orgID, query, b.aggregateAliasSet(), filterableSet, start, end, variables, matchedSB)
 	if err != nil {
 		return nil, err
 	}
 
-	matchedFrag, matchedArgs, err := b.buildMatchedCTE(matchedSB, start, end, startBucket, endBucket, resolved, orders, orderableSet, maskExpr, fp, resourcePred, limit, query.Offset)
+	matchedFrag, matchedArgs, err := b.buildMatchedCTE(matchedSB, start, end, startBucket, endBucket, resolved, orders, orderableSet, filterableSet, maskExpr, fp, resourcePred, limit, query.Offset)
 	if err != nil {
 		return nil, err
 	}
@@ -326,9 +327,10 @@ func (b *scopedTraceStatementBuilder) resolveMask(ctx context.Context, orgID val
 }
 
 type resolvedColumn struct {
-	alias     string
-	expr      string
-	orderable bool
+	alias      string
+	expr       string
+	orderable  bool
+	filterable bool
 }
 
 func (b *scopedTraceStatementBuilder) resolveColumns(ctx context.Context, orgID valuer.UUID, start, end uint64, cols *columnResolver, preds *predicateResolver) ([]resolvedColumn, error) {
@@ -338,7 +340,7 @@ func (b *scopedTraceStatementBuilder) resolveColumns(ctx context.Context, orgID
 		if err != nil {
 			return nil, err
 		}
-		out = append(out, resolvedColumn{alias: c.Alias, expr: expr, orderable: c.Orderable})
+		out = append(out, resolvedColumn{alias: c.Alias, expr: expr, orderable: c.Orderable, filterable: c.Filterable})
 	}
 	return out, nil
 }
@@ -393,7 +395,7 @@ type filterParts struct {
 // splitFilter splits query.Filter into a span-level predicate (args bound into sb)
 // and a trace-level HAVING (explicit query.Having ANDed on), then validates the
 // trace-level part against the matched-pass aggregates.
-func (b *scopedTraceStatementBuilder) splitFilter(ctx context.Context, orgID valuer.UUID, query qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation], classifySet, orderableSet map[string]struct{}, start, end uint64, variables map[string]qbtypes.VariableItem, sb *sqlbuilder.SelectBuilder) (filterParts, error) {
+func (b *scopedTraceStatementBuilder) splitFilter(ctx context.Context, orgID valuer.UUID, query qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation], classifySet, filterableSet map[string]struct{}, start, end uint64, variables map[string]qbtypes.VariableItem, sb *sqlbuilder.SelectBuilder) (filterParts, error) {
 	var fp filterParts
 	if query.Filter != nil && strings.TrimSpace(query.Filter.Expression) != "" {
 		spanExpr, traceExpr, err := querybuilder.SplitFilterForAggregates(query.Filter.Expression, classifySet)
@@ -429,7 +431,7 @@ func (b *scopedTraceStatementBuilder) splitFilter(ctx context.Context, orgID val
 		}
 		fp.havingExpr = replaced
 	}
-	if err := validateAggregateFilter(fp.havingExpr, orderableSet); err != nil {
+	if err := validateAggregateFilter(fp.havingExpr, filterableSet); err != nil {
 		return fp, err
 	}
 	return fp, nil
@@ -473,7 +475,7 @@ func (b *scopedTraceStatementBuilder) resolveSpanPredicate(ctx context.Context,
 // span filter + HAVING + ORDER BY + LIMIT/OFFSET, selecting only the aliases ORDER BY
 // / HAVING reference. Expressions carry $n markers bound to sb, so each can appear
 // several times and every occurrence resolves to the same arg.
-func (b *scopedTraceStatementBuilder) buildMatchedCTE(sb *sqlbuilder.SelectBuilder, start, end, startBucket, endBucket uint64, resolved []resolvedColumn, orders []listOrder, orderableSet map[string]struct{}, maskExpr string, fp filterParts, resourcePred string, limit, offset int) (string, []any, error) {
+func (b *scopedTraceStatementBuilder) buildMatchedCTE(sb *sqlbuilder.SelectBuilder, start, end, startBucket, endBucket uint64, resolved []resolvedColumn, orders []listOrder, orderableSet, filterableSet map[string]struct{}, maskExpr string, fp filterParts, resourcePred string, limit, offset int) (string, []any, error) {
 	needed := neededMatchedAliases(orders, fp.havingExpr, orderableSet)
 	selects := []string{"trace_id"}
 	for _, rc := range resolved {
@@ -513,8 +515,8 @@ func (b *scopedTraceStatementBuilder) buildMatchedCTE(sb *sqlbuilder.SelectBuild
 	}
 	if strings.TrimSpace(fp.havingExpr) != "" {
 		// the rewriter matches raw key text, so map the trace. form alongside the bare name
-		columnMap := make(map[string]string, len(orderableSet)*2)
-		for a := range orderableSet {
+		columnMap := make(map[string]string, len(filterableSet)*2)
+		for a := range filterableSet {
 			columnMap[a] = quoteAlias(a)
 			columnMap[telemetrytypes.FieldContextTrace.StringValue()+"."+a] = quoteAlias(a)
 		}
@@ -603,6 +605,17 @@ func orderableAliasSet(resolved []resolvedColumn) map[string]struct{} {
 	return set
 }
 
+// filterableAliasSet is the subset of aliases usable in the trace-level filter.
+func filterableAliasSet(resolved []resolvedColumn) map[string]struct{} {
+	set := make(map[string]struct{})
+	for _, rc := range resolved {
+		if rc.filterable {
+			set[rc.alias] = struct{}{}
+		}
+	}
+	return set
+}
+
 // neededMatchedAliases is the minimal alias set the matched pass must select: those
 // in ORDER BY plus those in the aggregate HAVING.
 func neededMatchedAliases(orders []listOrder, havingExpr string, orderableSet map[string]struct{}) map[string]struct{} {
@@ -630,19 +643,19 @@ func traceAggregateNames(havingExpr string) []string {
 	return names
 }
 
-// validateAggregateFilter rejects a trace-level filter referencing an aggregate not
-// computable in the matched pass.
-func validateAggregateFilter(havingExpr string, orderableSet map[string]struct{}) error {
+// validateAggregateFilter rejects a trace-level filter referencing an aggregate that
+// is not filterable.
+func validateAggregateFilter(havingExpr string, filterableSet map[string]struct{}) error {
 	if strings.TrimSpace(havingExpr) == "" {
 		return nil
 	}
-	allowed := make([]string, 0, len(orderableSet))
-	for a := range orderableSet {
+	allowed := make([]string, 0, len(filterableSet))
+	for a := range filterableSet {
 		allowed = append(allowed, a)
 	}
 	sort.Strings(allowed)
 	for _, name := range traceAggregateNames(havingExpr) {
-		if _, ok := orderableSet[name]; !ok {
+		if _, ok := filterableSet[name]; !ok {
 			return errors.NewInvalidInputf(errors.CodeInvalidInput,
 				"aggregate %q cannot be used in the trace-list filter; filterable aggregates: %s", name, strings.Join(allowed, ", "))
 		}

pkg/telemetrymetadata/metadata.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"log/slog"
+	"slices"
 	"strings"
 	"time"
 
@@ -14,6 +15,7 @@ import (
 	"github.com/SigNoz/signoz/pkg/factory"
 	"github.com/SigNoz/signoz/pkg/flagger"
 	"github.com/SigNoz/signoz/pkg/querybuilder"
+	"github.com/SigNoz/signoz/pkg/statementbuilder/aistatementbuilder"
 	"github.com/SigNoz/signoz/pkg/telemetryschema/audittelemetryschema"
 	"github.com/SigNoz/signoz/pkg/telemetryschema/logstelemetryschema"
 	"github.com/SigNoz/signoz/pkg/telemetryschema/metertelemetryschema"
@@ -159,6 +161,13 @@ func (t *telemetryMetaStore) getTracesKeys(ctx context.Context, fieldKeySelector
 		instrumentationtypes.CodeFunctionName: "getTracesKeys",
 	})
 
+	// The trace field context never matches ingested keys — it names the computed
+	// per-trace aggregates, which enrichWithAITraceAggregateKeys serves. Without this
+	// the tagType condition below has no branch for it and the scan returns every key.
+	fieldKeySelectors = slices.DeleteFunc(slices.Clone(fieldKeySelectors), func(s *telemetrytypes.FieldKeySelector) bool {
+		return s.FieldContext == telemetrytypes.FieldContextTrace
+	})
+
 	if len(fieldKeySelectors) == 0 {
 		return nil, true, nil
 	}
@@ -1168,6 +1177,31 @@ func enrichWithIntrinsicMetricKeys(keys map[string][]*telemetrytypes.TelemetryFi
 	return keys
 }
 
+// enrichWithAITraceAggregateKeys adds the computed per-trace aggregate keys for
+// builder_ai_query selectors; they are never ingested, so the scan cannot serve them.
+func enrichWithAITraceAggregateKeys(keys map[string][]*telemetrytypes.TelemetryFieldKey, selectors []*telemetrytypes.FieldKeySelector) map[string][]*telemetrytypes.TelemetryFieldKey {
+	defs := aistatementbuilder.MetadataFieldKeys()
+	matched := make(map[string]*telemetrytypes.TelemetryFieldKey)
+	for _, selector := range selectors {
+		if selector.QueryType != qbtypes.QueryTypeBuilderAI.StringValue() {
+			continue
+		}
+		if selector.Signal != telemetrytypes.SignalTraces && selector.Signal != telemetrytypes.SignalUnspecified {
+			continue
+		}
+		for _, def := range defs {
+			if selectorMatchesIntrinsicField(selector, *def) {
+				matched[def.Name] = def
+			}
+		}
+	}
+
+	for name, def := range matched {
+		keys[name] = append(keys[name], def)
+	}
+	return keys
+}
+
 // enrichWithGenAIKeys adds keys that can be queried for GenAI signals, even though they have not been ingested yet.
 func enrichWithGenAIKeys(keys map[string][]*telemetrytypes.TelemetryFieldKey, selectors []*telemetrytypes.FieldKeySelector) map[string][]*telemetrytypes.TelemetryFieldKey {
 	for _, selector := range selectors {
@@ -1276,6 +1310,7 @@ func (t *telemetryMetaStore) GetKeys(ctx context.Context, orgID valuer.UUID, fie
 	mapOfKeys = enrichWithIntrinsicMetricKeys(mapOfKeys, selectors)
 	if t.fl.BooleanOrEmpty(ctx, flagger.FeatureEnableAIObservability, featuretypes.NewFlaggerEvaluationContext(orgID)) {
 		mapOfKeys = enrichWithGenAIKeys(mapOfKeys, selectors)
+		mapOfKeys = enrichWithAITraceAggregateKeys(mapOfKeys, selectors)
 	}
 
 	return mapOfKeys, complete, nil
@@ -1357,6 +1392,7 @@ func (t *telemetryMetaStore) GetKeysMulti(ctx context.Context, orgID valuer.UUID
 	mapOfKeys = enrichWithIntrinsicMetricKeys(mapOfKeys, fieldKeySelectors)
 	if t.fl.BooleanOrEmpty(ctx, flagger.FeatureEnableAIObservability, featuretypes.NewFlaggerEvaluationContext(orgID)) {
 		mapOfKeys = enrichWithGenAIKeys(mapOfKeys, fieldKeySelectors)
+		mapOfKeys = enrichWithAITraceAggregateKeys(mapOfKeys, fieldKeySelectors)
 	}
 
 	return mapOfKeys, complete, nil

pkg/types/telemetrytypes/field.go

@@ -238,6 +238,7 @@ type FieldKeySelector struct {
 	EndUnixMilli      int64                  `json:"endUnixMilli"`
 	Signal            Signal                 `json:"signal"`
 	Source            Source                 `json:"source"`
+	QueryType         string                 `json:"queryType"`
 	FieldContext      FieldContext           `json:"fieldContext"`
 	FieldDataType     FieldDataType          `json:"fieldDataType"`
 	Name              string                 `json:"name"`
@@ -261,6 +262,7 @@ type GettableFieldKeys struct {
 type PostableFieldKeysParams struct {
 	Signal          Signal        `query:"signal"`
 	Source          Source        `query:"source"`
+	Type            string        `query:"type"`
 	Limit           int           `query:"limit"`
 	StartUnixMilli  int64         `query:"startUnixMilli"`
 	EndUnixMilli    int64         `query:"endUnixMilli"`
@@ -297,6 +299,7 @@ func NewFieldKeySelectorFromPostableFieldKeysParams(params PostableFieldKeysPara
 
 	req.Signal = params.Signal
 	req.Source = params.Source
+	req.QueryType = params.Type
 	req.FieldContext = params.FieldContext
 	req.FieldDataType = params.FieldDataType
 	req.SelectorMatchType = FieldSelectorMatchTypeFuzzy

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

@@ -0,0 +1,79 @@
+"""Fields metadata API with type="builder_ai_query": gen_ai attributes and per-trace
+aggregate keys are served pre-ingestion, behind the conftest's AI observability flag."""
+
+from collections.abc import Callable
+from http import HTTPStatus
+
+import pytest
+import requests
+
+from fixtures import types
+from fixtures.auth import USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD
+
+AI_TRACE_AGGREGATES = {
+    "llm_call_count",
+    "tool_call_count",
+    "distinct_tool_count",
+    "input_tokens",
+    "output_tokens",
+    "total_tokens",
+    "estimated_total_cost",
+    "max_llm_duration_nano",
+}
+
+
+@pytest.fixture(name="get_keys")
+def get_keys_fixture(
+    signoz: types.SigNoz,
+    create_user_admin: None,  # pylint: disable=unused-argument
+    get_token: Callable[[str, str], str],
+) -> Callable[[dict], dict]:
+    token = get_token(USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD)
+
+    def get_keys(params: dict) -> dict:
+        response = requests.get(
+            signoz.self.host_configs["8080"].get("/api/v1/fields/keys"),
+            timeout=5,
+            headers={"authorization": f"Bearer {token}"},
+            params={"signal": "traces", **params},
+        )
+        assert response.status_code == HTTPStatus.OK, response.text
+        assert response.json()["status"] == "success"
+        return response.json()["data"]["keys"]
+
+    return get_keys
+
+
+def trace_context_names(keys: dict) -> set:
+    return {name for name, variants in keys.items() if any(k["fieldContext"] == "trace" for k in variants)}
+
+
+def test_ai_fields_trace_context_lists_only_aggregates(get_keys: Callable[[dict], dict]) -> None:
+    """fieldContext=trace (the order-by picker request) returns exactly the
+    filterable aggregates — the ingested-key scan must not leak into it."""
+    keys = get_keys({"type": "builder_ai_query", "fieldContext": "trace"})
+    assert set(keys.keys()) == AI_TRACE_AGGREGATES, keys
+    assert trace_context_names(keys) == AI_TRACE_AGGREGATES
+
+
+def test_ai_fields_trace_prefix_search(get_keys: Callable[[dict], dict]) -> None:
+    """`trace.output` in the filter bar parses into the trace context and suggests
+    the matching aggregate."""
+    keys = get_keys({"type": "builder_ai_query", "searchText": "trace.output"})
+    assert trace_context_names(keys) == {"output_tokens"}, keys
+
+
+def test_ai_fields_bare_prefix_suggests_both_classes(get_keys: Callable[[dict], dict]) -> None:
+    """A bare prefix suggests the aggregate and the gen_ai span attribute side by side."""
+    keys = get_keys({"type": "builder_ai_query", "searchText": "output_tok"})
+    assert "output_tokens" in trace_context_names(keys), keys
+    assert "gen_ai.usage.output_tokens" in keys, keys
+    assert any(k["fieldContext"] == "attribute" for k in keys["gen_ai.usage.output_tokens"])
+
+
+def test_ai_fields_aggregates_require_ai_query_type(get_keys: Callable[[dict], dict]) -> None:
+    """Without type=builder_ai_query the aggregates are not suggested; the gen_ai
+    attributes still are (flag-gated, not query-type-gated)."""
+    keys = get_keys({"searchText": "output_tok"})
+    assert not trace_context_names(keys), keys
+    assert "gen_ai.usage.output_tokens" in keys, keys

tests/integration/tests/querierai/conftest.py

@@ -0,0 +1,35 @@
+import pytest
+from testcontainers.core.container import Network
+
+from fixtures import types
+from fixtures.signoz import create_signoz
+
+
+@pytest.fixture(name="signoz", scope="package")
+def signoz_ai(
+    network: Network,
+    zeus: types.TestContainerDocker,
+    gateway: types.TestContainerDocker,
+    sqlstore: types.TestContainerSQL,
+    clickhouse: types.TestContainerClickhouse,
+    request: pytest.FixtureRequest,
+    pytestconfig: pytest.Config,
+) -> types.SigNoz:
+    """
+    Package-scoped SigNoz with AI observability enabled: the metadata store surfaces
+    the gen_ai semconv keys and the trace-aggregate suggestion keys only behind this
+    flag, so the querierai suite runs against a flag-enabled instance.
+    """
+    return create_signoz(
+        network=network,
+        zeus=zeus,
+        gateway=gateway,
+        sqlstore=sqlstore,
+        clickhouse=clickhouse,
+        request=request,
+        pytestconfig=pytestconfig,
+        cache_key="signoz-ai",
+        env_overrides={
+            "SIGNOZ_FLAGGER_CONFIG_BOOLEAN_ENABLE__AI__OBSERVABILITY": True,
+        },
+    )