Add the table_id as a label

When we send data to prometheus we provide labels, which we can use
to create dashboards that are helpful. Currently the labels we send
are top-k and month. We also need to add the table_id so that we can
see which metrics are for clickstream data, and which are for binary.

I've called the label "dataset" rather than table_id as it feels a bit
easier to understand.

Author: hannako

app/services/discovery_engine/quality/evaluations.rb

@@ -18,7 +18,7 @@ def collect_quality_metrics(month_label)
       all_sample_query_sets(month_label).each do |set|
         e = DiscoveryEngine::Quality::Evaluation.new(set).fetch_quality_metrics
         Rails.logger.info(e)
-        metric_collector.record_evaluations(e, month_label)
+        metric_collector.record_evaluations(e, month_label, set.table_id)
       end
     end
 

app/services/discovery_engine/quality/sample_query_set.rb

@@ -3,6 +3,8 @@ module Quality
     class SampleQuerySet
       BIGQUERY_DATASET_ID = "automated_evaluation_input".freeze
 
+      attr_reader :table_id
+
       def initialize(table_id:, month_label: nil, month: nil, year: nil)
         @month_label = month_label
         @month = month
@@ -21,7 +23,7 @@ def name
 
     private
 
-      attr_reader :month_label, :table_id, :month, :year
+      attr_reader :month_label, :month, :year
 
       def create
         DiscoveryEngine::Clients

app/services/metrics/evaluation.rb

@@ -6,22 +6,22 @@ def initialize(registry)
       @doc_recall = registry.gauge(
         :search_api_v2_evaluation_monitoring_recall,
         docstring: "Vertex AI search evaluation recall",
-        labels: %i[top month],
+        labels: %i[top month dataset],
       )
       @doc_precision = registry.gauge(
         :search_api_v2_evaluation_monitoring_precision,
         docstring: "Vertex AI search evaluation precision",
-        labels: %i[top month],
+        labels: %i[top month dataset],
       )
       @doc_ndcg = registry.gauge(
         :search_api_v2_evaluation_monitoring_ndcg,
         docstring: "Vertex AI search evaluation ndcg",
-        labels: %i[top month],
+        labels: %i[top month dataset],
       )
     end
 
-    def record_evaluations(evaluation_result, month)
-      metrics.each { |key, registry| record_evaluation(key, registry, month, evaluation_result) }
+    def record_evaluations(evaluation_result, month, dataset)
+      metrics.each { |key, registry| record_evaluation(key, registry, month, dataset, evaluation_result) }
     end
 
   private
@@ -32,10 +32,10 @@ def metrics
       { doc_recall:, doc_precision:, doc_ndcg: }
     end
 
-    def record_evaluation(key, registry, month, evaluation_result)
+    def record_evaluation(key, registry, month, dataset, evaluation_result)
       TOP_K_LEVELS.each do |k|
         value = evaluation_result[key][:"top_#{k}"]
-        registry.set(value, labels: { top: k, month: })
+        registry.set(value, labels: { top: k, month:, dataset: })
       end
     end
   end

spec/services/discovery_engine/quality/evaluations_spec.rb

@@ -5,7 +5,7 @@
   let(:evaluation) { double("evaluation") }
   let(:evaluation_response) { "amything" }
   let(:sample_query_sets) { double("sample_query_sets") }
-  let(:sample_query_set) { double("sample_query_set") }
+  let(:sample_query_set) { double("sample_query_set", table_id: "clickstream") }
 
   before do
     allow(DiscoveryEngine::Quality::Evaluation)
@@ -19,11 +19,11 @@
 
     allow(metric_collector)
       .to receive(:record_evaluations)
-      .with(evaluation_response, :last_month)
+      .with(evaluation_response, :last_month, "clickstream")
 
     allow(metric_collector)
       .to receive(:record_evaluations)
-      .with(evaluation_response, :month_before_last)
+      .with(evaluation_response, :month_before_last, "clickstream")
 
     allow(DiscoveryEngine::Quality::SampleQuerySets)
       .to receive(:new)

spec/services/metrics/evaluation_spec.rb

@@ -3,6 +3,7 @@
 
   let(:registry) { double("registry") }
   let(:month) { :last_month }
+  let(:table_id) { "clickstream" }
 
   let(:evaluation_response) do
     {
@@ -44,33 +45,33 @@
   describe "#record_evaluations" do
     it "records the recall, precision and ndcg score" do
       expect(recall_gauge).to receive(:set)
-        .with(0.988, { labels: { top: "1", month: } })
+        .with(0.988, { labels: { top: "1", month:, dataset: "clickstream" } })
       expect(recall_gauge).to receive(:set)
-        .with(0.995, { labels: { top: "3", month: } })
+        .with(0.995, { labels: { top: "3", month:, dataset: "clickstream" } })
       expect(recall_gauge).to receive(:set)
-        .with(0.998, { labels: { top: "5", month: } })
+        .with(0.998, { labels: { top: "5", month:, dataset: "clickstream" } })
       expect(recall_gauge).to receive(:set)
-        .with(0.999, { labels: { top: "10", month: } })
+        .with(0.999, { labels: { top: "10", month:, dataset: "clickstream" } })
 
       expect(precision_gauge).to receive(:set)
-        .with(0.988, { labels: { top: "1", month: } })
+        .with(0.988, { labels: { top: "1", month:, dataset: "clickstream" } })
       expect(precision_gauge).to receive(:set)
-        .with(0.953, { labels: { top: "3", month: } })
+        .with(0.953, { labels: { top: "3", month:, dataset: "clickstream" } })
       expect(precision_gauge).to receive(:set)
-        .with(0.896, { labels: { top: "5", month: } })
+        .with(0.896, { labels: { top: "5", month:, dataset: "clickstream" } })
       expect(precision_gauge).to receive(:set)
-        .with(0.752, { labels: { top: "10", month: } })
+        .with(0.752, { labels: { top: "10", month:, dataset: "clickstream" } })
 
       expect(ndcg_gauge).to receive(:set)
-        .with(0.988, { labels: { top: "1", month: } })
+        .with(0.988, { labels: { top: "1", month:, dataset: "clickstream" } })
       expect(ndcg_gauge).to receive(:set)
-        .with(0.961, { labels: { top: "3", month: } })
+        .with(0.961, { labels: { top: "3", month:, dataset: "clickstream" } })
       expect(ndcg_gauge).to receive(:set)
-        .with(0.929, { labels: { top: "5", month: } })
+        .with(0.929, { labels: { top: "5", month:, dataset: "clickstream" } })
       expect(ndcg_gauge).to receive(:set)
-        .with(0.887, { labels: { top: "10", month: } })
+        .with(0.887, { labels: { top: "10", month:, dataset: "clickstream" } })
 
-      evaluation.record_evaluations(evaluation_response, month)
+      evaluation.record_evaluations(evaluation_response, month, table_id)
     end
   end
 end