fix: support for quick filters

pkg/signoz/provider.go

@@ -220,6 +220,7 @@ func NewSQLMigrationProviderFactories(
 		sqlmigration.NewRemoveOrganizationTuplesFactory(sqlstore),
 		sqlmigration.NewAddRoleTransactionGroupsFactory(sqlstore, sqlschema),
 		sqlmigration.NewAddTagUniqueIndexFactory(sqlstore, sqlschema),
+		sqlmigration.NewAddAiObservabilityQuickFiltersFactory(sqlstore),
 	)
 }
 

pkg/sqlmigration/101_add_ai_observability_quickfilters.go

@@ -0,0 +1,134 @@
+package sqlmigration
+
+import (
+	"context"
+	"database/sql"
+	"encoding/json"
+	"time"
+
+	"github.com/SigNoz/signoz/pkg/errors"
+	"github.com/SigNoz/signoz/pkg/factory"
+	"github.com/SigNoz/signoz/pkg/sqlstore"
+	"github.com/SigNoz/signoz/pkg/types/telemetrytypes"
+	"github.com/SigNoz/signoz/pkg/valuer"
+	"github.com/uptrace/bun"
+	"github.com/uptrace/bun/migrate"
+)
+
+type addAiObservabilityQuickFilters struct {
+	sqlstore sqlstore.SQLStore
+}
+
+func NewAddAiObservabilityQuickFiltersFactory(sqlstore sqlstore.SQLStore) factory.ProviderFactory[SQLMigration, Config] {
+	return factory.NewProviderFactory(factory.MustNewName("add_ai_observability_filters"), func(ctx context.Context, providerSettings factory.ProviderSettings, config Config) (SQLMigration, error) {
+		return &addAiObservabilityQuickFilters{sqlstore: sqlstore}, nil
+	})
+}
+
+func (migration *addAiObservabilityQuickFilters) Register(migrations *migrate.Migrations) error {
+	if err := migrations.Register(migration.Up, migration.Down); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (migration *addAiObservabilityQuickFilters) Up(ctx context.Context, db *bun.DB) error {
+	// keep in sync with the ai_observability defaults in quickfiltertypes.NewDefaultQuickFilter
+	aiObservabilityFilters := []map[string]interface{}{
+		{"key": telemetrytypes.GenAIRequestModel, "dataType": "string", "type": "tag"},
+		{"key": telemetrytypes.GenAIProviderName, "dataType": "string", "type": "tag"},
+		{"key": telemetrytypes.GenAIToolName, "dataType": "string", "type": "tag"},
+		{"key": telemetrytypes.GenAIAgentName, "dataType": "string", "type": "tag"},
+		{"key": "deployment.environment", "dataType": "string", "type": "resource"},
+		{"key": "service.name", "dataType": "string", "type": "resource"},
+		{"key": "hasError", "dataType": "bool", "type": "tag"},
+	}
+
+	aiObservabilityJSON, err := json.Marshal(aiObservabilityFilters)
+	if err != nil {
+		return errors.Wrapf(err, errors.TypeInternal, errors.CodeInternal, "failed to marshal ai monitoring filters")
+	}
+
+	type signal struct {
+		valuer.String
+	}
+
+	type identifiable struct {
+		ID valuer.UUID `json:"id" bun:"id,pk,type:text"`
+	}
+
+	type timeAuditable struct {
+		CreatedAt time.Time `bun:"created_at" json:"createdAt"`
+		UpdatedAt time.Time `bun:"updated_at" json:"updatedAt"`
+	}
+
+	type quickFilterType struct {
+		bun.BaseModel `bun:"table:quick_filter"`
+		identifiable
+		OrgID  valuer.UUID `bun:"org_id,type:text,notnull"`
+		Filter string      `bun:"filter,type:text,notnull"`
+		Signal signal      `bun:"signal,type:text,notnull"`
+		timeAuditable
+	}
+
+	tx, err := db.BeginTx(ctx, nil)
+	if err != nil {
+		return err
+	}
+
+	defer func() {
+		_ = tx.Rollback()
+	}()
+
+	var orgIDs []string
+	err = tx.NewSelect().
+		Table("organizations").
+		Column("id").
+		Scan(ctx, &orgIDs)
+	if err != nil && err != sql.ErrNoRows {
+		return err
+	}
+
+	var filtersToInsert []quickFilterType
+	for _, orgIDStr := range orgIDs {
+		orgID, err := valuer.NewUUID(orgIDStr)
+		if err != nil {
+			return err
+		}
+
+		filtersToInsert = append(filtersToInsert, quickFilterType{
+			identifiable: identifiable{
+				ID: valuer.GenerateUUID(),
+			},
+			OrgID:  orgID,
+			Filter: string(aiObservabilityJSON),
+			Signal: signal{valuer.NewString("ai_observability")},
+			timeAuditable: timeAuditable{
+				CreatedAt: time.Now(),
+				UpdatedAt: time.Now(),
+			},
+		})
+	}
+
+	if len(filtersToInsert) > 0 {
+		_, err = tx.NewInsert().
+			Model(&filtersToInsert).
+			On("CONFLICT (org_id, signal) DO UPDATE").
+			Set("filter = EXCLUDED.filter, updated_at = EXCLUDED.updated_at").
+			Exec(ctx)
+		if err != nil {
+			return err
+		}
+	}
+
+	if err := tx.Commit(); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (migration *addAiObservabilityQuickFilters) Down(ctx context.Context, db *bun.DB) error {
+	return nil
+}

pkg/telemetryai/field_mapper.go

@@ -97,3 +97,23 @@ func (p genAIColumnProvider) AggregateAliases() []string {
 	}
 	return aliases
 }
+
+// TraceAggregateFieldKeys returns the filterable/orderable per-trace aggregate columns
+// as trace-context field keys; the metadata store surfaces them for source=ai key
+// suggestions (`trace.` autocomplete in the filter bar and the order-by picker).
+func TraceAggregateFieldKeys() []*telemetrytypes.TelemetryFieldKey {
+	cols := genAIColumnProvider{}.Columns()
+	keys := make([]*telemetrytypes.TelemetryFieldKey, 0, len(cols))
+	for _, c := range cols {
+		if !c.Orderable {
+			continue
+		}
+		keys = append(keys, &telemetrytypes.TelemetryFieldKey{
+			Name:          c.Alias,
+			Signal:        telemetrytypes.SignalTraces,
+			FieldContext:  telemetrytypes.FieldContextTrace,
+			FieldDataType: telemetrytypes.FieldDataTypeFloat64,
+		})
+	}
+	return keys
+}

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/telemetryai"
 	"github.com/SigNoz/signoz/pkg/telemetryaudit"
 	"github.com/SigNoz/signoz/pkg/telemetrylogs"
 	"github.com/SigNoz/signoz/pkg/telemetrymetrics"
@@ -179,6 +181,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
 	}
@@ -1188,6 +1197,31 @@ func enrichWithIntrinsicMetricKeys(keys map[string][]*telemetrytypes.TelemetryFi
 	return keys
 }
 
+// enrichWithAITraceAggregateKeys adds keys that can be queried for AI trace aggregate signals.
+func enrichWithAITraceAggregateKeys(keys map[string][]*telemetrytypes.TelemetryFieldKey, selectors []*telemetrytypes.FieldKeySelector) map[string][]*telemetrytypes.TelemetryFieldKey {
+	for _, selector := range selectors {
+		if selector.Source != telemetrytypes.SourceAI {
+			continue
+		}
+		if selector.Signal != telemetrytypes.SignalTraces && selector.Signal != telemetrytypes.SignalUnspecified {
+			continue
+		}
+		for _, def := range telemetryai.TraceAggregateFieldKeys() {
+			if !selectorMatchesIntrinsicField(selector, *def) {
+				continue
+			}
+			if slices.ContainsFunc(keys[def.Name], func(k *telemetrytypes.TelemetryFieldKey) bool {
+				return k.FieldContext == telemetrytypes.FieldContextTrace
+			}) {
+				continue // already added by an earlier selector
+			}
+			keys[def.Name] = append(keys[def.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 {
@@ -1296,6 +1330,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
@@ -1377,6 +1412,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/quickfiltertypes/filter.go

@@ -7,6 +7,7 @@ import (
 	"github.com/SigNoz/signoz/pkg/errors"
 	v3 "github.com/SigNoz/signoz/pkg/query-service/model/v3"
 	"github.com/SigNoz/signoz/pkg/types"
+	"github.com/SigNoz/signoz/pkg/types/telemetrytypes"
 	"github.com/SigNoz/signoz/pkg/valuer"
 	"github.com/uptrace/bun"
 )
@@ -31,11 +32,12 @@ func (enum *Signal) UnmarshalJSON(data []byte) error {
 }
 
 var (
-	SignalTraces        = Signal{valuer.NewString("traces")}
-	SignalLogs          = Signal{valuer.NewString("logs")}
-	SignalApiMonitoring = Signal{valuer.NewString("api_monitoring")}
-	SignalExceptions    = Signal{valuer.NewString("exceptions")}
-	SignalMeter         = Signal{valuer.NewString("meter")}
+	SignalTraces          = Signal{valuer.NewString("traces")}
+	SignalLogs            = Signal{valuer.NewString("logs")}
+	SignalApiMonitoring   = Signal{valuer.NewString("api_monitoring")}
+	SignalExceptions      = Signal{valuer.NewString("exceptions")}
+	SignalMeter           = Signal{valuer.NewString("meter")}
+	SignalAiObservability = Signal{valuer.NewString("ai_observability")}
 )
 
 // NewSignal creates a Signal from a string.
@@ -51,6 +53,8 @@ func NewSignal(s string) (Signal, error) {
 		return SignalExceptions, nil
 	case "meter":
 		return SignalMeter, nil
+	case "ai_observability":
+		return SignalAiObservability, nil
 	default:
 		return Signal{}, errors.Newf(errors.TypeInvalidInput, errors.CodeInvalidInput, "invalid signal: %s", s)
 	}
@@ -187,6 +191,20 @@ func NewDefaultQuickFilter(orgID valuer.UUID) ([]*StorableQuickFilter, error) {
 		{"key": "host.name", "dataType": "float64", "type": "Sum"},
 	}
 
+	// AI observability (source=ai trace explorer): categorical gen_ai span attributes
+	// plus the usual service/environment/error narrowing. The per-trace aggregates
+	// (trace.output_tokens, …) are threshold filters with no value list, so they are
+	// not quick filters.
+	aiObservabilityFilters := []map[string]interface{}{
+		{"key": telemetrytypes.GenAIRequestModel, "dataType": "string", "type": "tag"},
+		{"key": telemetrytypes.GenAIProviderName, "dataType": "string", "type": "tag"},
+		{"key": telemetrytypes.GenAIToolName, "dataType": "string", "type": "tag"},
+		{"key": telemetrytypes.GenAIAgentName, "dataType": "string", "type": "tag"},
+		{"key": "deployment.environment", "dataType": "string", "type": "resource"},
+		{"key": "service.name", "dataType": "string", "type": "resource"},
+		{"key": "hasError", "dataType": "bool", "type": "tag"},
+	}
+
 	tracesJSON, err := json.Marshal(tracesFilters)
 	if err != nil {
 		return nil, errors.Wrapf(err, errors.TypeInternal, errors.CodeInternal, "failed to marshal traces filters")
@@ -212,6 +230,11 @@ func NewDefaultQuickFilter(orgID valuer.UUID) ([]*StorableQuickFilter, error) {
 		return nil, errors.Wrapf(err, errors.TypeInternal, errors.CodeInternal, "failed to marshal meter filters")
 	}
 
+	aiObservabilityJSON, err := json.Marshal(aiObservabilityFilters)
+	if err != nil {
+		return nil, errors.Wrapf(err, errors.TypeInternal, errors.CodeInternal, "failed to marshal ai observability filters")
+	}
+
 	timeRightNow := time.Now()
 
 	return []*StorableQuickFilter{
@@ -275,5 +298,17 @@ func NewDefaultQuickFilter(orgID valuer.UUID) ([]*StorableQuickFilter, error) {
 				UpdatedAt: timeRightNow,
 			},
 		},
+		{
+			Identifiable: types.Identifiable{
+				ID: valuer.GenerateUUID(),
+			},
+			OrgID:  orgID,
+			Filter: string(aiObservabilityJSON),
+			Signal: SignalAiObservability,
+			TimeAuditable: types.TimeAuditable{
+				CreatedAt: timeRightNow,
+				UpdatedAt: timeRightNow,
+			},
+		},
 	}, nil
 }

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

@@ -0,0 +1,101 @@
+"""
+Integration tests for the fields metadata API with source="ai".
+
+The metadata store serves two synthetic key sets for AI-enabled orgs (the suite's
+conftest boots SigNoz with the enable_ai_observability flag): the gen_ai semconv
+span attributes and the per-trace aggregate columns as trace-context keys. Both are
+served before any gen_ai data is ingested, so no traces are inserted here — that is
+the designed cold-start behavior.
+"""
+
+from http import HTTPStatus
+from collections.abc import Callable
+
+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_cost_usd",
+    "max_llm_latency_ns",
+    "last_activity_time",
+}
+
+
+def _get_keys(signoz: types.SigNoz, token: str, 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"]
+
+
+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(
+    signoz: types.SigNoz,
+    create_user_admin: None,  # pylint: disable=unused-argument
+    get_token: Callable[[str, str], str],
+) -> None:
+    """fieldContext=trace (the order-by picker request) returns exactly the
+    filterable aggregates — the ingested-key scan must not leak into it."""
+    token = get_token(USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD)
+
+    keys = _get_keys(signoz, token, {"source": "ai", "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(
+    signoz: types.SigNoz,
+    create_user_admin: None,  # pylint: disable=unused-argument
+    get_token: Callable[[str, str], str],
+) -> None:
+    """Typing `trace.output` in the filter bar parses into the trace context and
+    suggests the matching aggregate."""
+    token = get_token(USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD)
+
+    keys = _get_keys(signoz, token, {"source": "ai", "searchText": "trace.output"})
+    assert _trace_context_names(keys) == {"output_tokens"}, keys
+
+
+def test_ai_fields_bare_prefix_suggests_both_classes(
+    signoz: types.SigNoz,
+    create_user_admin: None,  # pylint: disable=unused-argument
+    get_token: Callable[[str, str], str],
+) -> None:
+    """A bare prefix suggests the aggregate and the gen_ai span attribute side by
+    side — both served pre-ingestion for flag-enabled orgs."""
+    token = get_token(USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD)
+
+    keys = _get_keys(signoz, token, {"source": "ai", "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_source(
+    signoz: types.SigNoz,
+    create_user_admin: None,  # pylint: disable=unused-argument
+    get_token: Callable[[str, str], str],
+) -> None:
+    """Without source=ai the trace aggregates are not suggested; the gen_ai semconv
+    attributes still are (flag-gated, not source-gated)."""
+    token = get_token(USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD)
+
+    keys = _get_keys(signoz, token, {"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,
+        },
+    )