[segmentation] Allow clients to record ukm source
This lets the CollectTrainingData API callers pass in an UKM source ID
so that the recorded training data can be attached to the right UKM
source.
Bug: 362515254
Change-Id: I1a0046c92ef2093635c93aa0a26eda53d123180c
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/5809858
Reviewed-by: Salvador Guerrero Ramos <[email protected]>
Auto-Submit: Siddhartha S <[email protected]>
Reviewed-by: Trevor Perrier <[email protected]>
Commit-Queue: Siddhartha S <[email protected]>
Cr-Commit-Position: refs/heads/main@{#1347686}
diff --git a/components/segmentation_platform/internal/segmentation_ukm_helper.cc b/components/segmentation_platform/internal/segmentation_ukm_helper.cc
index 572a759..00909eca 100644
--- a/components/segmentation_platform/internal/segmentation_ukm_helper.cc
+++ b/components/segmentation_platform/internal/segmentation_ukm_helper.cc
@@ -24,6 +24,7 @@
#include "components/segmentation_platform/public/local_state_helper.h"
#include "services/metrics/public/cpp/ukm_builders.h"
#include "services/metrics/public/cpp/ukm_recorder.h"
+#include "services/metrics/public/cpp/ukm_source_id.h"
#define CALL_MEMBER_FN(obj, func) ((obj).*(func))
#define ARRAY_SIZE(x) (sizeof(x) / sizeof(x)[0])
@@ -208,12 +209,15 @@
ukm::SourceId SegmentationUkmHelper::RecordTrainingData(
SegmentId segment_id,
int64_t model_version,
+ ukm::SourceId ukm_source_id,
const ModelProvider::Request& input_tensor,
const ModelProvider::Response& outputs,
const std::vector<int>& output_indexes,
std::optional<proto::PredictionResult> prediction_result,
std::optional<SelectedSegment> selected_segment) {
- ukm::SourceId source_id = ukm::NoURLSourceId();
+ ukm::SourceId source_id = ukm_source_id != ukm::kInvalidSourceId
+ ? ukm_source_id
+ : ukm::NoURLSourceId();
ukm::builders::Segmentation_ModelExecution execution_result(source_id);
if (!AddInputsToUkm(&execution_result, segment_id, model_version,
input_tensor)) {
