feat: support ai trace aggregate filtering in ai span list

pkg/telemetryai/trace_scope_test.go

@@ -0,0 +1,141 @@
+package telemetryai
+
+import (
+	"context"
+	"testing"
+
+	qbtypes "github.com/SigNoz/signoz/pkg/types/querybuildertypes/querybuildertypesv5"
+	"github.com/SigNoz/signoz/pkg/types/telemetrytypes"
+	"github.com/stretchr/testify/require"
+)
+
+// Span list with a mixed filter: gen_ai spans matching the span-level part, in
+// traces whose window-clipped aggregates satisfy the trace-level part (the
+// __trace_scope qualification on the delegated path).
+func TestBuild_FullSQL_SpanList_TraceScoped(t *testing.T) {
+	b := newTestBuilder(t)
+	stmt, err := b.Build(context.Background(), testStartMs, testEndMs, qbtypes.RequestTypeRaw,
+		qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
+			Signal: telemetrytypes.SignalTraces, Source: telemetrytypes.SourceAI,
+			Filter: &qbtypes.Filter{Expression: "gen_ai.request.model = 'gpt-4o-mini' AND trace.output_tokens > 1000"},
+			Limit:  10,
+		}, nil)
+	require.NoError(t, err)
+
+	requireSQLEqual(t, `
+WITH __trace_scope AS (
+    SELECT trace_id,
+        sum(multiIf(mapContains(attributes_number, 'gen_ai.usage.output_tokens') = true, toFloat64(attributes_number['gen_ai.usage.output_tokens']), NULL)) AS output_tokens
+    FROM signoz_traces.distributed_signoz_index_v3
+    WHERE timestamp >= '1747947419000000000'
+      AND timestamp < '1747983448000000000'
+      AND ts_bucket_start >= 1747945619
+      AND ts_bucket_start <= 1747983448
+      AND (mapContains(attributes_string, 'gen_ai.request.model') = true OR mapContains(attributes_string, 'gen_ai.tool.name') = true OR mapContains(attributes_string, 'gen_ai.agent.name') = true)
+    GROUP BY trace_id
+    HAVING output_tokens > 1000
+)
+SELECT timestamp AS timestamp, trace_id AS trace_id, span_id AS span_id,
+    trace_state AS trace_state, parent_span_id AS parent_span_id, flags AS flags,
+    name AS name, kind AS kind, kind_string AS kind_string, duration_nano AS duration_nano,
+    status_code AS status_code, status_message AS status_message,
+    status_code_string AS status_code_string, events AS events, links AS links,
+    response_status_code AS response_status_code, external_http_url AS external_http_url,
+    http_url AS http_url, external_http_method AS external_http_method,
+    http_method AS http_method, http_host AS http_host, db_name AS db_name,
+    db_operation AS db_operation, has_error AS has_error, is_remote AS is_remote,
+    attributes_string, attributes_number, attributes_bool, resources_string
+FROM signoz_traces.distributed_signoz_index_v3
+WHERE trace_id GLOBAL IN (SELECT trace_id FROM __trace_scope)
+  AND (((mapContains(attributes_string, 'gen_ai.request.model') = true
+        OR mapContains(attributes_string, 'gen_ai.tool.name') = true
+        OR mapContains(attributes_string, 'gen_ai.agent.name') = true))
+    AND ((attributes_string['gen_ai.request.model'] = 'gpt-4o-mini'
+        AND mapContains(attributes_string, 'gen_ai.request.model') = true)))
+  AND timestamp >= '1747947419000000000'
+  AND timestamp < '1747983448000000000'
+  AND ts_bucket_start >= 1747945619
+  AND ts_bucket_start <= 1747983448
+LIMIT 10
+`, stmt)
+}
+
+// Without a trace-level condition nothing changes: the span list stays a single
+// gated span scan (no __trace_scope CTE).
+func TestBuild_SpanList_NoTraceFilter_NoScope(t *testing.T) {
+	b := newTestBuilder(t)
+	stmt, err := b.Build(context.Background(), testStartMs, testEndMs, qbtypes.RequestTypeRaw,
+		qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
+			Signal: telemetrytypes.SignalTraces, Source: telemetrytypes.SourceAI,
+			Filter: &qbtypes.Filter{Expression: "gen_ai.request.model = 'gpt-4o-mini'"},
+			Limit:  10,
+		}, nil)
+	require.NoError(t, err)
+	require.NotContains(t, stmt.Query, "__trace_scope")
+}
+
+// The span-list trace-level filter shares the trace list's rules: output-only
+// aggregates are rejected, OR-mixing the two classes is rejected, and explicitly
+// trace-level order keys get a targeted error — while bare span columns that happen
+// to share a name with an aggregate alias (duration_nano) stay orderable.
+func TestBuild_SpanList_TraceFilter_Validation(t *testing.T) {
+	b := newTestBuilder(t)
+	build := func(q qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]) error {
+		q.Signal = telemetrytypes.SignalTraces
+		q.Source = telemetrytypes.SourceAI
+		_, err := b.Build(context.Background(), testStartMs, testEndMs, qbtypes.RequestTypeRaw, q, nil)
+		return err
+	}
+
+	err := build(qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
+		Filter: &qbtypes.Filter{Expression: "trace.span_count > 3"},
+	})
+	require.ErrorContains(t, err, `aggregate "span_count" cannot be used`)
+
+	err = build(qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
+		Filter: &qbtypes.Filter{Expression: "trace.output_tokens > 1000 OR kind_string = 'Client'"},
+	})
+	require.ErrorContains(t, err, "cannot be combined")
+
+	err = build(qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
+		Order: []qbtypes.OrderBy{{Key: qbtypes.OrderByKey{TelemetryFieldKey: telemetrytypes.TelemetryFieldKey{Name: "trace.output_tokens"}}}},
+	})
+	require.ErrorContains(t, err, `ordering the span list by trace-level aggregate "trace.output_tokens" is not supported`)
+
+	err = build(qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
+		Order: []qbtypes.OrderBy{{Key: qbtypes.OrderByKey{TelemetryFieldKey: telemetrytypes.TelemetryFieldKey{Name: "duration_nano"}}, Direction: qbtypes.OrderDirectionDesc}},
+		Limit: 10,
+	})
+	require.NoError(t, err, "bare duration_nano is a span column, not a trace-level key")
+}
+
+// Variables in a trace-level condition on the span list get the trace list's
+// treatment: substituted as literals, __all__ drops the condition (no scope CTE),
+// tracefield. spelling behaves like trace..
+func TestBuild_SpanList_TraceFilter_Variables(t *testing.T) {
+	b := newTestBuilder(t)
+	build := func(expr string, vars map[string]qbtypes.VariableItem) (*qbtypes.Statement, error) {
+		return b.Build(context.Background(), testStartMs, testEndMs, qbtypes.RequestTypeRaw,
+			qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]{
+				Signal: telemetrytypes.SignalTraces, Source: telemetrytypes.SourceAI,
+				Filter: &qbtypes.Filter{Expression: expr},
+				Limit:  10,
+			}, vars)
+	}
+
+	stmt, err := build("trace.output_tokens > $threshold",
+		map[string]qbtypes.VariableItem{"threshold": {Value: 700}})
+	require.NoError(t, err)
+	require.Contains(t, stmt.Query, "HAVING output_tokens > 700")
+
+	stmt, err = build("trace.output_tokens > $threshold",
+		map[string]qbtypes.VariableItem{"threshold": {Type: qbtypes.DynamicVariableType, Value: "__all__"}})
+	require.NoError(t, err)
+	require.NotContains(t, stmt.Query, "__trace_scope")
+
+	viaTrace, err := build("trace.output_tokens > 1000", nil)
+	require.NoError(t, err)
+	viaTracefield, err := build("tracefield.output_tokens > 1000", nil)
+	require.NoError(t, err)
+	require.Equal(t, viaTrace.Query, viaTracefield.Query)
+}

pkg/telemetryscopedtraces/statement_builder.go

@@ -102,14 +102,53 @@ func (b *scopedTraceStatementBuilder) Build(
 	case qbtypes.RequestTypeTrace:
 		return b.buildTraceListQuery(ctx, querybuilder.ToNanoSecs(start), querybuilder.ToNanoSecs(end), query, variables)
 	case qbtypes.RequestTypeRaw:
+		if err := b.validateRawOrderKeys(query); err != nil {
+			return nil, err
+		}
 		return b.buildDelegated(ctx, start, end, requestType, query, variables)
 	default:
 		return nil, ErrUnsupportedRequestType
 	}
 }
 
-// buildDelegated ANDs the base gate into the user filter and delegates to the
-// standard trace builder (the span-list / raw path).
+// traceScopedStatementBuilder is the delegate's optional capability of constraining a
+// query to a set of trace ids (implemented by the telemetrytraces builder).
+type traceScopedStatementBuilder interface {
+	BuildTraceScoped(ctx context.Context, start, end uint64, requestType qbtypes.RequestType, query qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation], variables map[string]qbtypes.VariableItem, traceScope *qbtypes.Statement) (*qbtypes.Statement, error)
+}
+
+// substituteTraceLevelVariables resolves query variables in a trace-level expression.
+// The span-level parts bind variables via PrepareWhereClause; trace-level parts are
+// text rewrites, so variables become literals (a dynamic __all__ drops its condition).
+func substituteTraceLevelVariables(expr string, variables map[string]qbtypes.VariableItem) (string, error) {
+	if strings.TrimSpace(expr) == "" || len(variables) == 0 {
+		return expr, nil
+	}
+	return qbvariables.ReplaceVariablesInExpression(expr, variables)
+}
+
+// validateRawOrderKeys rejects ordering the span list by an explicitly trace-level
+// aggregate (trace./tracefield. prefix or trace field context) — the per-trace value
+// does not exist on span rows. Bare names pass through: they may legitimately be span
+// columns (duration_nano, timestamp).
+func (b *scopedTraceStatementBuilder) validateRawOrderKeys(query qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]) error {
+	aliases := b.aggregateAliasSet()
+	for _, o := range query.Order {
+		name := strings.TrimPrefix(strings.TrimPrefix(o.Key.Name, "tracefield."), "trace.")
+		if _, ok := aliases[name]; !ok {
+			continue
+		}
+		if name != o.Key.Name || o.Key.FieldContext == telemetrytypes.FieldContextTrace {
+			return errors.NewInvalidInputf(errors.CodeInvalidInput,
+				"ordering the span list by trace-level aggregate %q is not supported; order by span columns instead (e.g. timestamp, duration_nano)", o.Key.Name)
+		}
+	}
+	return nil
+}
+
+// buildDelegated ANDs the base gate into the span-level filter part and delegates to
+// the standard trace builder. A trace-level part (trace.output_tokens > 1000) becomes
+// a __trace_scope qualification the delegate constrains trace_id by.
 func (b *scopedTraceStatementBuilder) buildDelegated(
 	ctx context.Context,
 	start, end uint64,
@@ -117,17 +156,98 @@ func (b *scopedTraceStatementBuilder) buildDelegated(
 	query qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation],
 	variables map[string]qbtypes.VariableItem,
 ) (*qbtypes.Statement, error) {
+	spanExpr, traceExpr := "", ""
+	if query.Filter != nil && strings.TrimSpace(query.Filter.Expression) != "" {
+		var err error
+		spanExpr, traceExpr, err = querybuilder.SplitFilterForAggregates(query.Filter.Expression, b.aggregateAliasSet())
+		if err != nil {
+			return nil, err
+		}
+	}
+	traceExpr, err := substituteTraceLevelVariables(traceExpr, variables)
+	if err != nil {
+		return nil, err
+	}
+
 	gate := b.baseCond.FilterExpression()
 	expr := gate
-	if query.Filter != nil && strings.TrimSpace(query.Filter.Expression) != "" {
-		expr = fmt.Sprintf("(%s) AND (%s)", gate, query.Filter.Expression)
+	if strings.TrimSpace(spanExpr) != "" {
+		expr = fmt.Sprintf("(%s) AND (%s)", gate, spanExpr)
 	}
 
 	// shallow copy; only Filter is replaced, caller's query untouched
 	gated := query
 	gated.Filter = &qbtypes.Filter{Expression: expr}
 
-	return b.traceStmtBuilder.Build(ctx, start, end, requestType, gated, variables)
+	if strings.TrimSpace(traceExpr) == "" {
+		return b.traceStmtBuilder.Build(ctx, start, end, requestType, gated, variables)
+	}
+
+	scoped, ok := b.traceStmtBuilder.(traceScopedStatementBuilder)
+	if !ok {
+		return nil, errors.NewInternalf(errors.CodeInternal, "trace statement builder does not support trace-scoped queries")
+	}
+	scope, err := b.buildQualifiedStatement(ctx, querybuilder.ToNanoSecs(start), querybuilder.ToNanoSecs(end), traceExpr)
+	if err != nil {
+		return nil, err
+	}
+	return scoped.BuildTraceScoped(ctx, start, end, requestType, gated, variables, scope)
+}
+
+// buildQualifiedStatement builds the __trace_scope statement: trace ids whose
+// window-clipped per-trace aggregates satisfy traceExpr, from one windowed,
+// mask-pruned GROUP BY trace_id scan. start/end are ns.
+func (b *scopedTraceStatementBuilder) buildQualifiedStatement(ctx context.Context, start, end uint64, traceExpr string) (*qbtypes.Statement, error) {
+	keys, err := b.fetchKeys(ctx)
+	if err != nil {
+		return nil, err
+	}
+	maskExpr, maskArgs, err := b.resolveMask(ctx, start, end, keys)
+	if err != nil {
+		return nil, err
+	}
+	resolved, err := b.resolveColumns(ctx, start, end, keys, maskExpr, maskArgs)
+	if err != nil {
+		return nil, err
+	}
+	orderableSet := orderableAliasSet(resolved)
+	if err := validateAggregateFilter(traceExpr, orderableSet); err != nil {
+		return nil, err
+	}
+	having, err := b.buildHaving(traceExpr, orderableSet)
+	if err != nil {
+		return nil, err
+	}
+
+	startBucket := start/querybuilder.NsToSeconds - querybuilder.BucketAdjustment
+	endBucket := end / querybuilder.NsToSeconds
+
+	sb := sqlbuilder.NewSelectBuilder()
+	needed := neededMatchedAliases(nil, traceExpr, orderableSet)
+	selects := []string{"trace_id"}
+	for _, rc := range resolved {
+		if _, ok := needed[rc.alias]; !ok {
+			continue
+		}
+		colExpr, err := embedExpr(sb, rc.expr, rc.args)
+		if err != nil {
+			return nil, err
+		}
+		selects = append(selects, colExpr+" AS "+quoteAlias(rc.alias))
+	}
+	sb.Select(selects...)
+	sb.From(spanTable())
+	mask, err := embedExpr(sb, maskExpr, maskArgs)
+	if err != nil {
+		return nil, err
+	}
+	sb.Where(append(windowWhere(sb, start, end, startBucket, endBucket), mask)...)
+	sb.GroupBy("trace_id")
+	if strings.TrimSpace(having) != "" {
+		sb.Having(having)
+	}
+	sql, args := sb.BuildWithFlavor(sqlbuilder.ClickHouse)
+	return &qbtypes.Statement{Query: sql, Args: args}, nil
 }
 
 // buildTraceListQuery wires the CTE pipeline (benchmarked, see ai-qb-handoff.md): one
@@ -482,16 +602,11 @@ func (b *scopedTraceStatementBuilder) splitFilter(ctx context.Context, query qbt
 			fp.havingExpr = query.Having.Expression
 		}
 	}
-	// The span predicate binds variables via PrepareWhereClause; the HAVING is a plain
-	// text rewrite, so substitute variables here (list/IN quoting, __all__ drops the
-	// condition) before validating.
-	if strings.TrimSpace(fp.havingExpr) != "" && len(variables) > 0 {
-		replaced, err := qbvariables.ReplaceVariablesInExpression(fp.havingExpr, variables)
-		if err != nil {
-			return fp, err
-		}
-		fp.havingExpr = replaced
+	havingExpr, err := substituteTraceLevelVariables(fp.havingExpr, variables)
+	if err != nil {
+		return fp, err
 	}
+	fp.havingExpr = havingExpr
 	if err := validateAggregateFilter(fp.havingExpr, orderableSet); err != nil {
 		return fp, err
 	}

pkg/telemetrytraces/statement_builder.go

@@ -29,6 +29,9 @@ type traceQueryStatementBuilder struct {
 	resourceFilterResolver         *telemetryresourcefilter.ResourceFingerprintResolver[qbtypes.TraceAggregation]
 	aggExprRewriter                qbtypes.AggExprRewriter
 	skipResourceFingerprintEnabled bool
+	// traceScope is set only on the per-call copy made by BuildTraceScoped; it
+	// constrains the query to trace ids selected by the __trace_scope CTE.
+	traceScope *qbtypes.Statement
 }
 
 var _ qbtypes.StatementBuilder[qbtypes.TraceAggregation] = (*traceQueryStatementBuilder)(nil)
@@ -70,6 +73,32 @@ func NewTraceQueryStatementBuilder(
 	}
 }
 
+// BuildTraceScoped is Build additionally constrained to spans whose trace_id is
+// selected by traceScope. The receiver is copied so the shared builder stays stateless.
+func (b *traceQueryStatementBuilder) BuildTraceScoped(
+	ctx context.Context,
+	start uint64,
+	end uint64,
+	requestType qbtypes.RequestType,
+	query qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation],
+	variables map[string]qbtypes.VariableItem,
+	traceScope *qbtypes.Statement,
+) (*qbtypes.Statement, error) {
+	scoped := *b
+	scoped.traceScope = traceScope
+	return scoped.Build(ctx, start, end, requestType, query, variables)
+}
+
+// attachTraceScope adds the trace-scope condition to sb and returns the CTE fragment
+// + args to prepend; both empty when no scope is set.
+func (b *traceQueryStatementBuilder) attachTraceScope(sb *sqlbuilder.SelectBuilder) (string, []any) {
+	if b.traceScope == nil {
+		return "", nil
+	}
+	sb.Where("trace_id GLOBAL IN (SELECT trace_id FROM __trace_scope)")
+	return fmt.Sprintf("__trace_scope AS (%s)", b.traceScope.Query), b.traceScope.Args
+}
+
 // Build builds a SQL query for traces based on the given parameters.
 func (b *traceQueryStatementBuilder) Build(
 	ctx context.Context,
@@ -292,6 +321,11 @@ func (b *traceQueryStatementBuilder) buildListQuery(
 		cteArgs = append(cteArgs, args)
 	}
 
+	if scopeFrag, scopeArgs := b.attachTraceScope(sb); scopeFrag != "" {
+		cteFragments = append(cteFragments, scopeFrag)
+		cteArgs = append(cteArgs, scopeArgs)
+	}
+
 	for _, field := range query.SelectFields {
 		colExpr, err := b.fm.ColumnExpressionFor(ctx, start, end, &field, keys)
 		if err != nil {

tests/integration/tests/querierai/01_ai_traces.py

@@ -369,6 +369,44 @@ def test_ai_span_list_excludes_non_gen_ai_spans(
     assert "POST /api/chat" not in names  # root span excluded
 
 
+def test_ai_span_list_trace_level_filter(
+    signoz: types.SigNoz,
+    create_user_admin: None,  # pylint: disable=unused-argument
+    get_token: Callable[[str, str], str],
+    insert_traces: Callable[[list[Traces]], None],
+) -> None:
+    """
+    Span list (raw) with a trace-level condition: only gen_ai spans of traces whose
+    window-clipped aggregates qualify come back (the __trace_scope qualification on
+    the delegated path). Two traces with out-tokens 100 / 300: `trace.output_tokens
+    > 100` keeps only the large trace's LLM span.
+    """
+    now = datetime.now(tz=UTC).replace(second=0, microsecond=0)
+    service = "ai-it-spanlist-tracefilter"
+    small = _ai_trace(now=now, service=service, user="a", in_tokens=10, out_tokens=100, cost=0.1)
+    large = _ai_trace(now=now, service=service, user="b", in_tokens=30, out_tokens=300, cost=0.2)
+    insert_traces(small + large)
+
+    token = get_token(USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD)
+    start_ms, end_ms = _window_ms(now)
+
+    query = BuilderQuery(
+        signal="traces",
+        source="ai",
+        name="A",
+        filter_expression=f"service.name = '{service}' AND trace.output_tokens > 100",
+        limit=10,
+    )
+    resp = make_query_request(signoz, token, start_ms, end_ms, [query.to_dict()], request_type=RequestType.RAW)
+    assert resp.status_code == HTTPStatus.OK, resp.text
+
+    rows = resp.json()["data"]["data"]["results"][0]["rows"]
+    assert len(rows) == 1, f"expected only the large trace's LLM span, got {len(rows)} rows"
+    body = json.dumps(rows)
+    assert large[0].trace_id in body
+    assert small[0].trace_id not in body
+
+
 def test_ai_list_having_or_aggregates(
     signoz: types.SigNoz,
     create_user_admin: None,  # pylint: disable=unused-argument

tests/integration/tests/querierai/conftest.py

@@ -0,0 +1,38 @@
+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_observability(
+    network: Network,
+    migrator: types.Operation,  # pylint: disable=unused-argument
+    zeus: types.TestContainerDocker,
+    gateway: types.TestContainerDocker,
+    sqlstore: types.TestContainerSQL,
+    clickhouse: types.TestContainerClickhouse,
+    request: pytest.FixtureRequest,
+    pytestconfig: pytest.Config,
+) -> types.SigNoz:
+    """
+    Package-scoped SigNoz instance with AI observability enabled. source=ai
+    queries rely on the metadata store surfacing the static gen_ai key
+    definitions (enrichWithGenAIKeys), which is gated on this flag — without it
+    the gate keys (gen_ai.tool.name, gen_ai.agent.name, ...) only resolve once
+    a span carrying them has been ingested.
+    """
+    return create_signoz(
+        network=network,
+        zeus=zeus,
+        gateway=gateway,
+        sqlstore=sqlstore,
+        clickhouse=clickhouse,
+        request=request,
+        pytestconfig=pytestconfig,
+        cache_key="signoz-ai-observability",
+        env_overrides={
+            "SIGNOZ_FLAGGER_CONFIG_BOOLEAN_ENABLE__AI__OBSERVABILITY": True,
+        },
+    )