fix: limit values to just 1

pkg/modules/llmpricingrule/impllmpricingrule/module.go

@@ -4,9 +4,12 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"log/slog"
+	"maps"
 	"time"
 
 	"github.com/SigNoz/signoz/pkg/errors"
+	"github.com/SigNoz/signoz/pkg/factory"
 	"github.com/SigNoz/signoz/pkg/modules/llmpricingrule"
 	"github.com/SigNoz/signoz/pkg/querier"
 	"github.com/SigNoz/signoz/pkg/query-service/agentConf"
@@ -22,12 +25,14 @@ import (
 const unmappedModelsLookback = time.Hour
 
 type module struct {
-	store   llmpricingruletypes.Store
-	querier querier.Querier
+	store                  llmpricingruletypes.Store
+	querier                querier.Querier
+	telemetryMetadataStore telemetrytypes.MetadataStore
+	logger                 *slog.Logger
 }
 
-func NewModule(store llmpricingruletypes.Store, querier querier.Querier) llmpricingrule.Module {
-	return &module{store: store, querier: querier}
+func NewModule(store llmpricingruletypes.Store, querier querier.Querier, telemetryMetadataStore telemetrytypes.MetadataStore, providerSettings factory.ProviderSettings) llmpricingrule.Module {
+	return &module{store: store, querier: querier, telemetryMetadataStore: telemetryMetadataStore, logger: providerSettings.Logger}
 }
 
 func (module *module) List(ctx context.Context, orgID valuer.UUID, offset, limit int, search string, isOverride *bool) ([]*llmpricingruletypes.LLMPricingRule, int, error) {
@@ -51,13 +56,36 @@ func (module *module) ListUnmappedModels(ctx context.Context, orgID valuer.UUID)
 		return nil, err
 	}
 
-	unmapped := make([]*llmpricingruletypes.UnmappedModel, 0, len(models))
-	for _, m := range models {
-		if !llmpricingruletypes.ModelMatchesAnyRule(m.ModelName, rules) {
-			unmapped = append(unmapped, m)
-		}
+	return llmpricingruletypes.UnmappedModels(models, rules), nil
+}
+
+// Collect keeps the rule counts even when ClickHouse is unreachable: the
+// counts that depend on it are skipped so the reporter retains their last value.
+func (module *module) Collect(ctx context.Context, orgID valuer.UUID) (map[string]any, error) {
+	rules, err := module.listAllRules(ctx, orgID)
+	if err != nil {
+		return nil, err
 	}
-	return unmapped, nil
+	stats := llmpricingruletypes.NewStatsFromRules(rules)
+
+	if models, err := module.discoverModels(ctx, orgID); err != nil {
+		module.logger.DebugContext(ctx, "failed to collect unmapped model stats", errors.Attr(err))
+	} else {
+		maps.Copy(stats, llmpricingruletypes.NewStatsFromUnmappedModels(llmpricingruletypes.UnmappedModels(models, rules)))
+	}
+
+	values := make(map[string]*telemetrytypes.TelemetryFieldValues)
+	for _, selector := range aiobservabilitytypes.StatsFieldValueSelectors(time.Now()) {
+		fieldValues, _, err := module.telemetryMetadataStore.GetAllValues(ctx, orgID, selector)
+		if err != nil {
+			module.logger.DebugContext(ctx, "failed to collect gen_ai attribute stats", slog.String("key", selector.Name), errors.Attr(err))
+			continue
+		}
+		values[selector.Name] = fieldValues
+	}
+	maps.Copy(stats, aiobservabilitytypes.NewStatsFromFieldValues(values))
+
+	return stats, nil
 }
 
 // CreateOrUpdate saves a batch of pricing rules. isOverride decides how a rule

pkg/modules/llmpricingrule/llmpricingrule.go

@@ -5,6 +5,7 @@ import (
 	"net/http"
 
 	"github.com/SigNoz/signoz/pkg/query-service/agentConf"
+	"github.com/SigNoz/signoz/pkg/statsreporter"
 	"github.com/SigNoz/signoz/pkg/types/llmpricingruletypes"
 	"github.com/SigNoz/signoz/pkg/valuer"
 )
@@ -13,6 +14,8 @@ type Module interface {
 	// Since this module interacts with OpAMP, it must implement the AgentFeature interface.
 	agentConf.AgentFeature
 
+	statsreporter.StatsCollector
+
 	List(ctx context.Context, orgID valuer.UUID, offset, limit int, search string, isOverride *bool) ([]*llmpricingruletypes.LLMPricingRule, int, error)
 	Get(ctx context.Context, orgID valuer.UUID, id valuer.UUID) (*llmpricingruletypes.LLMPricingRule, error)
 	CreateOrUpdate(ctx context.Context, orgID valuer.UUID, userEmail string, rules []*llmpricingruletypes.UpdatableLLMPricingRule) (err error)

pkg/signoz/module.go

@@ -163,7 +163,7 @@ func NewModules(
 		CloudIntegration:     cloudIntegrationModule,
 		TraceDetail:          impltracedetail.NewModule(impltracedetail.NewTraceStore(telemetryStore, telemetryMetadataStore, fl), providerSettings, config.TraceDetail),
 		SpanMapper:           spanMapper,
-		LLMPricingRule:       impllmpricingrule.NewModule(impllmpricingrule.NewStore(sqlstore), querier),
+		LLMPricingRule:       impllmpricingrule.NewModule(impllmpricingrule.NewStore(sqlstore), querier, telemetryMetadataStore, providerSettings),
 		Tag:                  tagModule,
 	}
 }

pkg/signoz/signoz.go

@@ -570,6 +570,7 @@ func New(
 		cloudIntegrationModule,
 		modules.LogsPipeline,
 		modules.InfraMonitoring,
+		modules.LLMPricingRule,
 		querier,
 		authz,
 	}

pkg/types/aiobservabilitytypes/stats.go

@@ -0,0 +1,46 @@
+package aiobservabilitytypes
+
+import (
+	"time"
+
+	"github.com/SigNoz/signoz/pkg/types/telemetrytypes"
+)
+
+// statsLookback reaches into the previous day; the metadata store rounds the
+// start down to a day boundary.
+const statsLookback = 24 * time.Hour
+
+// statsAttributeKeys are the gen_ai span attributes whose recent presence the
+// stats reporter tracks under telemetry.traces.<key>.exists.
+var statsAttributeKeys = []string{GenAIRequestModel, GenAIToolName, GenAIAgentName}
+
+// StatsFieldValueSelectors look up whether each tracked attribute had a string
+// value since the start of the previous day. One value is enough.
+func StatsFieldValueSelectors(now time.Time) []*telemetrytypes.FieldValueSelector {
+	selectors := make([]*telemetrytypes.FieldValueSelector, 0, len(statsAttributeKeys))
+	for _, key := range statsAttributeKeys {
+		selectors = append(selectors, &telemetrytypes.FieldValueSelector{
+			FieldKeySelector: &telemetrytypes.FieldKeySelector{
+				StartUnixMilli: now.Add(-statsLookback).UnixMilli(),
+				Signal:         telemetrytypes.SignalTraces,
+				FieldContext:   telemetrytypes.FieldContextAttribute,
+				FieldDataType:  telemetrytypes.FieldDataTypeString,
+				Name:           key,
+			},
+			Limit: 1,
+		})
+	}
+	return selectors
+}
+
+// NewStatsFromFieldValues reports whether each looked-up key had a value. Keys
+// absent from values are left out, so a failed lookup never reports false.
+func NewStatsFromFieldValues(values map[string]*telemetrytypes.TelemetryFieldValues) map[string]any {
+	stats := make(map[string]any, len(values))
+	for _, key := range statsAttributeKeys {
+		if fieldValues := values[key]; fieldValues != nil {
+			stats["telemetry.traces."+key+".exists"] = len(fieldValues.StringValues) > 0
+		}
+	}
+	return stats
+}

pkg/types/aiobservabilitytypes/stats_test.go

@@ -0,0 +1,53 @@
+package aiobservabilitytypes
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/SigNoz/signoz/pkg/types/telemetrytypes"
+)
+
+func TestNewStatsFromFieldValues(t *testing.T) {
+	testCases := []struct {
+		name   string
+		values map[string]*telemetrytypes.TelemetryFieldValues
+		want   map[string]any
+	}{
+		{
+			name: "NoLookups_NothingReported",
+			want: map[string]any{},
+		},
+		{
+			name: "EveryKeySeen_AllTrue",
+			values: map[string]*telemetrytypes.TelemetryFieldValues{
+				GenAIRequestModel: {StringValues: []string{"gpt-5-nano"}},
+				GenAIToolName:     {StringValues: []string{"search"}},
+				GenAIAgentName:    {StringValues: []string{"planner"}},
+			},
+			want: map[string]any{
+				"telemetry.traces.gen_ai.request.model.exists": true,
+				"telemetry.traces.gen_ai.tool.name.exists":     true,
+				"telemetry.traces.gen_ai.agent.name.exists":    true,
+			},
+		},
+		{
+			name: "EmptyAndMissingLookups_OnlyLookedUpKeysReported",
+			values: map[string]*telemetrytypes.TelemetryFieldValues{
+				GenAIRequestModel:  {StringValues: []string{"gpt-5-nano"}},
+				GenAIToolName:      {},
+				GenAIOperationName: {StringValues: []string{"chat"}},
+			},
+			want: map[string]any{
+				"telemetry.traces.gen_ai.request.model.exists": true,
+				"telemetry.traces.gen_ai.tool.name.exists":     false,
+			},
+		},
+	}
+
+	for _, testCase := range testCases {
+		t.Run(testCase.name, func(t *testing.T) {
+			assert.Equal(t, testCase.want, NewStatsFromFieldValues(testCase.values))
+		})
+	}
+}

pkg/types/dashboardtypes/perses_dashboard_stats.go

@@ -1,6 +1,7 @@
 package dashboardtypes
 
 import (
+	qb "github.com/SigNoz/signoz/pkg/types/querybuildertypes/querybuildertypesv5"
 	"github.com/SigNoz/signoz/pkg/types/telemetrytypes"
 )
 
@@ -10,6 +11,7 @@ const (
 	statKeyPanelTracesCount  = "dashboard.panels.traces.count"
 	statKeyPanelMetricsCount = "dashboard.panels.metrics.count"
 	statKeyPanelLogsCount    = "dashboard.panels.logs.count"
+	statKeyPanelAICount      = "dashboard.panels.ai.count"
 )
 
 // panelSignalStatKeys maps a builder query's signal to the stat it contributes
@@ -22,14 +24,17 @@ var panelSignalStatKeys = map[telemetrytypes.Signal]string{
 }
 
 // NewStatsFromStorableDashboards reports the stats of stored dashboards. Rows that
-// do not decode as v2 contribute to dashboard.count only.
+// do not decode as v2 contribute to dashboard.count only. System dashboards are
+// provisioned into every org, so they are left out of every count.
 func NewStatsFromStorableDashboards(dashboards []*StorableDashboard) map[string]any {
-	stats := newPanelStats()
+	stats := newStats()
 
 	for _, dashboard := range dashboards {
-		if dashboard == nil {
+		if dashboard == nil || dashboard.Source == SourceSystem {
 			continue
 		}
+		incrementStat(stats, statKeyDashboardCount)
+
 		dashboardV2, err := dashboard.ToDashboardV2(nil)
 		if err != nil {
 			continue
@@ -37,7 +42,6 @@ func NewStatsFromStorableDashboards(dashboards []*StorableDashboard) map[string]
 		addPanelStats(&dashboardV2.Spec, stats)
 	}
 
-	stats[statKeyDashboardCount] = int64(len(dashboards))
 	return stats
 }
 
@@ -45,24 +49,26 @@ func NewStatsFromStorableDashboards(dashboards []*StorableDashboard) map[string]
 // created, straight off the postable spec — the create path has no reason to make
 // a storable round-trip just to be counted.
 func NewStatsFromPostableDashboardV2(postable PostableDashboardV2) map[string]any {
-	stats := newPanelStats()
+	stats := newStats()
+	stats[statKeyDashboardCount] = int64(1)
 	addPanelStats(&postable.Spec, stats)
 
-	stats[statKeyDashboardCount] = int64(1)
 	return stats
 }
 
-func newPanelStats() map[string]any {
+func newStats() map[string]any {
 	return map[string]any{
+		statKeyDashboardCount:    int64(0),
 		statKeyPanelCount:        int64(0),
 		statKeyPanelTracesCount:  int64(0),
 		statKeyPanelMetricsCount: int64(0),
 		statKeyPanelLogsCount:    int64(0),
+		statKeyPanelAICount:      int64(0),
 	}
 }
 
 // addPanelStats counts the panels of a v2 spec, and each panel's queries against
-// the signal they read.
+// the signal they read. AI queries count only towards the AI stat.
 func addPanelStats(spec *DashboardSpec, stats map[string]any) {
 	for _, panel := range spec.Panels {
 		if panel == nil {
@@ -76,6 +82,10 @@ func addPanelStats(spec *DashboardSpec, stats map[string]any) {
 				continue
 			}
 			for _, envelope := range composite.Queries {
+				if envelope.Type == qb.QueryTypeBuilderAI {
+					incrementStat(stats, statKeyPanelAICount)
+					continue
+				}
 				if key, ok := panelSignalStatKeys[envelope.GetSignal()]; ok {
 					incrementStat(stats, key)
 				}

pkg/types/dashboardtypes/perses_dashboard_stats_test.go

@@ -117,20 +117,60 @@ func TestNewStatsFromStorableDashboardsCountsCompositeSubQueries(t *testing.T) {
 	assert.Equal(t, int64(1), stats[statKeyPanelLogsCount])
 }
 
-// An AI builder query is always a traces query, so it counts towards traces.
-func TestNewStatsFromStorableDashboardsCountsAIBuilderQueries(t *testing.T) {
-	aiBuilder := `{
+func statsAIBuilderQuery() string {
+	return `{
 		"kind": "time_series",
 		"spec": {"plugin": {"kind": "signoz/AIBuilderQuery", "spec": {"name": "A", "aggregations": [{"expression": "count()"}]}}}
 	}`
-	dashboard := newStatsStorableV2(t, `"p1": `+statsPanel(aiBuilder))
+}
+
+// An AI builder query reads traces but counts only towards the AI stat.
+func TestNewStatsFromStorableDashboardsCountsAIBuilderQueries(t *testing.T) {
+	dashboard := newStatsStorableV2(t, `"p1": `+statsPanel(statsAIBuilderQuery()))
 
 	stats := NewStatsFromStorableDashboards([]*StorableDashboard{dashboard})
 
 	assert.Equal(t, int64(1), stats[statKeyPanelCount])
-	assert.Equal(t, int64(1), stats[statKeyPanelTracesCount])
+	assert.Equal(t, int64(0), stats[statKeyPanelTracesCount])
 	assert.Equal(t, int64(0), stats[statKeyPanelMetricsCount])
 	assert.Equal(t, int64(0), stats[statKeyPanelLogsCount])
+	assert.Equal(t, int64(1), stats[statKeyPanelAICount])
+}
+
+func TestNewStatsFromStorableDashboardsCountsAICompositeSubQueries(t *testing.T) {
+	composite := `{
+		"kind": "time_series",
+		"spec": {"plugin": {"kind": "signoz/CompositeQuery", "spec": {"queries": [
+			{"type": "builder_ai_query", "spec": {"name": "A", "aggregations": [{"expression": "count()"}]}},
+			{"type": "builder_query", "spec": ` + statsBuilderQuerySpec("B", "logs") + `}
+		]}}}
+	}`
+	dashboard := newStatsStorableV2(t, `
+		"p1": `+statsPanel(composite)+`,
+		"p2": `+statsPanel(statsAIBuilderQuery())+`
+	`)
+
+	stats := NewStatsFromStorableDashboards([]*StorableDashboard{dashboard})
+
+	assert.Equal(t, int64(2), stats[statKeyPanelCount])
+	assert.Equal(t, int64(2), stats[statKeyPanelAICount])
+	assert.Equal(t, int64(0), stats[statKeyPanelTracesCount])
+	assert.Equal(t, int64(1), stats[statKeyPanelLogsCount])
+}
+
+func TestNewStatsFromStorableDashboardsSkipsSystemDashboards(t *testing.T) {
+	system := newStatsStorableV2(t, `"p1": `+statsPanel(statsAIBuilderQuery()))
+	system.Source = SourceSystem
+	system.Name = SystemDashboardNamePrefix + "overview"
+	user := newStatsStorableV2(t, `"p1": `+statsPanel(statsBuilderQuery("logs")))
+
+	stats := NewStatsFromStorableDashboards([]*StorableDashboard{system, user})
+
+	assert.Equal(t, int64(1), stats[statKeyDashboardCount])
+	assert.Equal(t, int64(1), stats[statKeyPanelCount])
+	assert.Equal(t, int64(0), stats[statKeyPanelTracesCount])
+	assert.Equal(t, int64(0), stats[statKeyPanelAICount])
+	assert.Equal(t, int64(1), stats[statKeyPanelLogsCount])
 }
 
 // promql and clickhouse queries carry no signal, so they land in the panel total
@@ -215,6 +255,17 @@ func TestNewStatsFromPostableDashboardV2(t *testing.T) {
 	assert.Equal(t, int64(1), stats[statKeyPanelLogsCount])
 	assert.Equal(t, int64(1), stats[statKeyPanelTracesCount])
 	assert.Equal(t, int64(0), stats[statKeyPanelMetricsCount])
+	assert.Equal(t, int64(0), stats[statKeyPanelAICount])
+}
+
+func TestNewStatsFromPostableDashboardV2CountsAIBuilderQueries(t *testing.T) {
+	postable := newStatsPostableV2(t, `"p1": `+statsPanel(statsAIBuilderQuery()))
+
+	stats := NewStatsFromPostableDashboardV2(postable)
+
+	assert.Equal(t, int64(1), stats[statKeyDashboardCount])
+	assert.Equal(t, int64(1), stats[statKeyPanelCount])
+	assert.Equal(t, int64(1), stats[statKeyPanelAICount])
 }
 
 func TestNewStatsFromPostableDashboardV2WithNoPanels(t *testing.T) {
@@ -233,4 +284,5 @@ func TestNewStatsFromStorableDashboardsWithNoDashboards(t *testing.T) {
 	assert.Equal(t, int64(0), stats[statKeyPanelTracesCount])
 	assert.Equal(t, int64(0), stats[statKeyPanelMetricsCount])
 	assert.Equal(t, int64(0), stats[statKeyPanelLogsCount])
+	assert.Equal(t, int64(0), stats[statKeyPanelAICount])
 }

pkg/types/llmpricingruletypes/pricing.go

@@ -277,3 +277,14 @@ func ModelMatchesAnyRule(model string, rules []*LLMPricingRule) bool {
 	}
 	return false
 }
+
+// UnmappedModels keeps the models that no rule pattern matches.
+func UnmappedModels(models []*UnmappedModel, rules []*LLMPricingRule) []*UnmappedModel {
+	unmapped := make([]*UnmappedModel, 0, len(models))
+	for _, model := range models {
+		if !ModelMatchesAnyRule(model.ModelName, rules) {
+			unmapped = append(unmapped, model)
+		}
+	}
+	return unmapped
+}

pkg/types/llmpricingruletypes/stats.go

@@ -0,0 +1,24 @@
+package llmpricingruletypes
+
+// NewStatsFromRules reports the rules users added themselves and the synced
+// defaults users pinned.
+func NewStatsFromRules(rules []*LLMPricingRule) map[string]any {
+	custom, overrides := int64(0), int64(0)
+	for _, rule := range rules {
+		switch {
+		case rule.SourceID == nil:
+			custom++
+		case rule.IsOverride:
+			overrides++
+		}
+	}
+
+	return map[string]any{
+		"llmpricingrule.custom.count":   custom,
+		"llmpricingrule.override.count": overrides,
+	}
+}
+
+func NewStatsFromUnmappedModels(unmapped []*UnmappedModel) map[string]any {
+	return map[string]any{"llmpricingrule.unmapped_model.count": int64(len(unmapped))}
+}

pkg/types/llmpricingruletypes/stats_test.go

@@ -0,0 +1,53 @@
+package llmpricingruletypes
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/SigNoz/signoz/pkg/valuer"
+)
+
+func TestNewStatsFromRules(t *testing.T) {
+	sourceID := valuer.GenerateUUID()
+
+	testCases := []struct {
+		name  string
+		rules []*LLMPricingRule
+		want  map[string]any
+	}{
+		{
+			name: "NoRules",
+			want: map[string]any{
+				"llmpricingrule.custom.count":   int64(0),
+				"llmpricingrule.override.count": int64(0),
+			},
+		},
+		{
+			name: "CustomPinnedAndSyncedRules_CountedSeparately",
+			rules: []*LLMPricingRule{
+				{IsOverride: true},
+				{IsOverride: false},
+				{SourceID: &sourceID, IsOverride: true},
+				{SourceID: &sourceID, IsOverride: false},
+			},
+			want: map[string]any{
+				"llmpricingrule.custom.count":   int64(2),
+				"llmpricingrule.override.count": int64(1),
+			},
+		},
+	}
+
+	for _, testCase := range testCases {
+		t.Run(testCase.name, func(t *testing.T) {
+			assert.Equal(t, testCase.want, NewStatsFromRules(testCase.rules))
+		})
+	}
+}
+
+func TestNewStatsFromUnmappedModels(t *testing.T) {
+	unmapped := []*UnmappedModel{{ModelName: "gpt-5-nano"}, {ModelName: "claude-haiku-5-5"}}
+
+	assert.Equal(t, map[string]any{"llmpricingrule.unmapped_model.count": int64(2)}, NewStatsFromUnmappedModels(unmapped))
+	assert.Equal(t, map[string]any{"llmpricingrule.unmapped_model.count": int64(0)}, NewStatsFromUnmappedModels(nil))
+}