Associate recorded evaluations with dataset selections and row origins

src/api/experiments.rs

@@ -44,6 +44,12 @@ pub struct ExperimentRegisterRequest {
     pub description: Option<String>,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub base_experiment: Option<String>,
+    /// Dataset associated with this experiment.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub dataset_id: Option<String>,
+    /// Explicit dataset snapshot; omitted for a live selection.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub dataset_version: Option<String>,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub metadata: Option<Map<String, Value>>,
     #[serde(skip_serializing_if = "Option::is_none")]

src/eval/executor.rs

@@ -54,6 +54,10 @@ impl<I, O> EvalPlan<I, O> {
 
 #[async_trait]
 pub(crate) trait EvalRecorder<I, O>: Send + Sync {
+    /// Initialize recording only after source identity is frozen, before reading cases.
+    async fn prepare(&mut self, _dataset: Option<&super::DatasetVersionRef>) -> Result<()> {
+        Ok(())
+    }
     fn target(&self) -> Option<super::EvalRecording> {
         None
     }
@@ -82,7 +86,7 @@ pub(crate) struct EvalExecutor;
 impl EvalExecutor {
     pub async fn run<I: Send + Sync + 'static, O: Send + Sync + 'static>(
         mut plan: EvalPlan<I, O>,
-        recorder: Option<&dyn EvalRecorder<I, O>>,
+        mut recorder: Option<&mut dyn EvalRecorder<I, O>>,
         reporter: &mut impl EvalReporter<I, O>,
         binding: Option<&crate::internal::execution::ExecutionBinding>,
     ) -> std::result::Result<EvalSummary<I, O>, EvalRunError<I, O>> {
@@ -97,6 +101,17 @@ impl EvalExecutor {
                 source,
             });
         }
+        if failures.is_empty() {
+            if let Some(recorder) = recorder.as_deref_mut() {
+                if let Err(source) = recorder.prepare(dataset.as_ref()).await {
+                    failures.push(EvalRunFailure::Recording {
+                        case_index: None,
+                        source,
+                    });
+                }
+            }
+        }
+        let recorder = recorder.as_deref();
         let mut case_count = 0;
         let mut pending = FuturesUnordered::new();
         while failures.is_empty() {
@@ -446,7 +461,7 @@ mod tests {
     #[test]
     fn recorder_start_and_finish_errors_preserve_computation_and_finalize_run() {
         let finished = Arc::new(AtomicUsize::new(0));
-        let recorder = FailingRecorder {
+        let mut recorder = FailingRecorder {
             finished: finished.clone(),
         };
         let cases = (0..3)
@@ -475,7 +490,7 @@ mod tests {
         let mut reporter = Observed::default();
         let error = futures::executor::block_on(EvalExecutor::run(
             plan,
-            Some(&recorder),
+            Some(&mut recorder),
             &mut reporter,
             None,
         ))

src/eval/recorder.rs

@@ -1,7 +1,7 @@
 //! Braintrust serialization and write completions, outside the run engine.
 use super::{
     executor::{CaseRecording, EvalRecorder},
-    Case, EvalRecording, EvalResult,
+    Case, DatasetVersionRef, EvalRecording, EvalResult,
 };
 use crate::{
     BraintrustClient, BraintrustError, Completion, DeliveryReceipt, Result, SpanHandle, SpanLog,
@@ -15,23 +15,47 @@ use std::sync::Arc;
 type Write = Completion<DeliveryReceipt, WriteError>;
 
 pub(crate) struct BraintrustRecorder {
+    client: BraintrustClient,
+    submitter: Arc<dyn crate::internal::ports::SpanSubmitter>,
+    project_name: Option<String>,
+    experiment_name: String,
+    metadata: Option<Map<String, Value>>,
+    ready: Option<RecordingState>,
+}
+struct RecordingState {
     template: crate::span::SpanBuilder,
     target: EvalRecording,
+    dataset: Option<DatasetVersionRef>,
 }
 impl BraintrustRecorder {
-    pub async fn new(
+    pub fn new(
         client: BraintrustClient,
         submitter: Arc<dyn crate::internal::ports::SpanSubmitter>,
         project_name: Option<String>,
         experiment_name: String,
         metadata: Option<Map<String, Value>>,
-    ) -> Result<Self> {
-        let mut builder = client
-            .experiment_builder_for_operation(submitter.clone(), project_name)
+    ) -> Self {
+        Self {
+            client,
+            submitter,
+            project_name,
+            experiment_name,
+            metadata,
+            ready: None,
+        }
+    }
+    async fn initialize(&mut self, dataset: Option<&DatasetVersionRef>) -> Result<()> {
+        let mut builder = self
+            .client
+            .experiment_builder_for_operation(self.submitter.clone(), self.project_name.clone())
             .await?
-            .experiment_name(experiment_name);
-        if let Some(metadata) = metadata {
-            builder = builder.metadata(metadata);
+            .experiment_name(self.experiment_name.clone());
+        if let Some(metadata) = &self.metadata {
+            builder = builder.metadata(metadata.clone());
+        }
+        if let Some(dataset) = dataset {
+            builder =
+                builder.dataset_reference(dataset.dataset_id.clone(), dataset.version.clone());
         }
         let experiment = builder
             .build()
@@ -42,16 +66,22 @@ impl BraintrustRecorder {
             experiment_id: resolved.experiment_id.clone(),
             experiment_name: resolved.experiment_name.clone(),
         };
-        let template = client
+        let template = self
+            .client
             .span_builder_for_operation(
-                submitter,
+                self.submitter.clone(),
                 Some(crate::types::ParentSpanInfo::Experiment {
                     object_id: target.experiment_id.clone(),
                 }),
             )
             .await?
             .span_type(crate::SpanType::Eval);
-        Ok(Self { template, target })
+        self.ready = Some(RecordingState {
+            template,
+            target,
+            dataset: dataset.cloned(),
+        });
+        Ok(())
     }
 }
 
@@ -79,12 +109,30 @@ fn serialize<T: Serialize>(
 impl<I: Serialize + Send + Sync + 'static, O: Serialize + Send + Sync + 'static> EvalRecorder<I, O>
     for BraintrustRecorder
 {
+    async fn prepare(&mut self, dataset: Option<&DatasetVersionRef>) -> Result<()> {
+        self.initialize(dataset).await
+    }
     async fn start(&self, index: usize, case: &Case<I, O>) -> Result<Box<dyn CaseRecording<I, O>>> {
-        let span = self
-            .template
-            .clone()
-            .name(format!("case-{}", index + 1))
-            .build();
+        let ready = self.ready.as_ref().ok_or_else(|| {
+            BraintrustError::Background("evaluation recorder is not initialized".into())
+        })?;
+        let mut builder = ready.template.clone().name(format!("case-{}", index + 1));
+        // Match the dataset object-reference contract: both row identity and row
+        // transaction are required. The selected dataset version is run metadata,
+        // not a substitute for a particular record's transaction.
+        if let (Some(dataset), Some(id), Some(xact_id)) = (
+            &ready.dataset,
+            case.id.as_ref().filter(|id| !id.is_empty()),
+            case.xact_id.as_ref().filter(|id| !id.is_empty()),
+        ) {
+            builder = builder.dataset_record_origin(crate::types::DatasetRecordOrigin {
+                object_id: dataset.dataset_id.clone(),
+                id: id.clone(),
+                xact_id: xact_id.clone(),
+                created: case.created.clone(),
+            });
+        }
+        let span = builder.build();
         let mut failures = Vec::new();
         let mut log = SpanLog::builder().name(format!("case-{}", index + 1));
         if let Some(input) = serialize(&case.input, "input", &mut failures) {
@@ -109,7 +157,7 @@ impl<I: Serialize + Send + Sync + 'static, O: Serialize + Send + Sync + 'static>
         }))
     }
     fn target(&self) -> Option<EvalRecording> {
-        Some(self.target.clone())
+        self.ready.as_ref().map(|ready| ready.target.clone())
     }
 
     async fn finish(&self) -> Vec<BraintrustError> {

src/eval/run.rs

@@ -117,7 +117,7 @@ fn recorded<I: Serialize + Send + Sync + 'static, O: Serialize + Send + Sync + '
             let mut reporter = TerminalReporter {
                 quiet: options.quiet,
             };
-            let recorder = BraintrustRecorder::new(
+            let mut recorder = BraintrustRecorder::new(
                 client,
                 operation.submitter.clone(),
                 options.project_name.clone(),
@@ -126,11 +126,10 @@ fn recorded<I: Serialize + Send + Sync + 'static, O: Serialize + Send + Sync + '
                     .clone()
                     .unwrap_or_else(|| options.name.clone()),
                 options.metadata.clone(),
-            )
-            .await?;
+            );
             let result = EvalExecutor::run(
                 EvalPlan::new(options)?,
-                Some(&recorder),
+                Some(&mut recorder),
                 &mut reporter,
                 Some(&callbacks),
             )

src/eval/source.rs

@@ -23,6 +23,9 @@ use std::pin::Pin;
 pub type SourceError = anyhow::Error;
 
 /// Dataset identity captured by a source and retained with its run result.
+/// Recorded runs also associate their experiment with this selection. Case row
+/// IDs and transactions identify individual origins; `version` selects the run's
+/// dataset snapshot and is omitted for live reads.
 #[derive(Debug, Clone, PartialEq, Eq, Serialize)]
 pub struct DatasetVersionRef {
     pub dataset_id: String,
@@ -35,11 +38,15 @@ pub struct DatasetVersionRef {
 #[async_trait]
 pub trait CaseSource<I, O>: Send {
     /// Resolve source identity before provenance is captured or cases are read.
+    /// The executor calls this once, before initializing recorded experiments.
     async fn prepare(&mut self) -> Result<(), SourceError> {
         Ok(())
     }
 
     async fn next(&mut self) -> Result<Option<Case<I, O>>, SourceError>;
+    /// The executor captures this once after preparation. A recorded run uses
+    /// it for experiment association and for case origins with nonempty `id`
+    /// and `xact_id`; ordinary case IDs alone never imply a dataset association.
     fn provenance(&self) -> Option<DatasetVersionRef> {
         None
     }

src/experiments/builder.rs

@@ -44,6 +44,8 @@ pub struct ExperimentBuilder {
     experiment_name: Option<String>,
     description: Option<String>,
     base_experiment: Option<String>,
+    dataset_id: Option<String>,
+    dataset_version: Option<String>,
     metadata: Option<Map<String, Value>>,
     public: Option<bool>,
     update: Option<bool>,
@@ -61,6 +63,8 @@ impl Clone for ExperimentBuilder {
             experiment_name: self.experiment_name.clone(),
             description: self.description.clone(),
             base_experiment: self.base_experiment.clone(),
+            dataset_id: self.dataset_id.clone(),
+            dataset_version: self.dataset_version.clone(),
             metadata: self.metadata.clone(),
             public: self.public,
             update: self.update,
@@ -86,6 +90,8 @@ impl ExperimentBuilder {
             experiment_name: None,
             description: None,
             base_experiment: None,
+            dataset_id: None,
+            dataset_version: None,
             metadata: None,
             public: None,
             update: None,
@@ -123,6 +129,12 @@ impl ExperimentBuilder {
         self
     }
 
+    pub(crate) fn dataset_reference(mut self, id: String, version: Option<String>) -> Self {
+        self.dataset_id = Some(id);
+        self.dataset_version = version;
+        self
+    }
+
     /// Set the experiment metadata.
     pub fn metadata(mut self, metadata: Map<String, Value>) -> Self {
         self.metadata = Some(metadata);
@@ -176,6 +188,8 @@ impl ExperimentBuilder {
             experiment_name: self.experiment_name,
             description: self.description,
             base_experiment: self.base_experiment,
+            dataset_id: self.dataset_id,
+            dataset_version: self.dataset_version,
             metadata: self.metadata,
             public: self.public,
             update: self.update,

src/experiments/experiment.rs

@@ -45,6 +45,8 @@ pub struct Experiment {
     pub(crate) experiment_name: Option<String>,
     pub(crate) description: Option<String>,
     pub(crate) base_experiment: Option<String>,
+    pub(crate) dataset_id: Option<String>,
+    pub(crate) dataset_version: Option<String>,
     pub(crate) metadata: Option<Map<String, Value>>,
     pub(crate) public: Option<bool>,
     pub(crate) update: Option<bool>,
@@ -96,6 +98,8 @@ impl Clone for Experiment {
             experiment_name: self.experiment_name.clone(),
             description: self.description.clone(),
             base_experiment: self.base_experiment.clone(),
+            dataset_id: self.dataset_id.clone(),
+            dataset_version: self.dataset_version.clone(),
             update: self.update,
             repo_info: self.repo_info.clone(),
             metadata: self.metadata.clone(),
@@ -555,6 +559,8 @@ impl Experiment {
                     experiment_name: self.experiment_name.clone(),
                     description: self.description.clone(),
                     base_experiment: self.base_experiment.clone(),
+                    dataset_id: self.dataset_id.clone(),
+                    dataset_version: self.dataset_version.clone(),
                     metadata: self.metadata.clone(),
                     public: self.public,
                     update: self.update,

src/span.rs

@@ -250,6 +250,7 @@ impl SpanLogBuilder {
 }
 
 pub struct SpanBuilder {
+    dataset_record_origin: Option<crate::types::DatasetRecordOrigin>,
     submitter: Arc<dyn SpanSubmitter>,
     scope: CredentialScope,
     project_name: Option<String>,
@@ -290,6 +291,7 @@ impl Clone for SpanBuilder {
             span_id_override: self.span_id_override.clone(),
             row_id_override: self.row_id_override.clone(),
             span_origin: self.span_origin.clone(),
+            dataset_record_origin: self.dataset_record_origin.clone(),
             span_attributes_extra: self.span_attributes_extra.clone(),
         }
     }
@@ -378,10 +380,19 @@ impl SpanBuilder {
             span_id_override: None,
             row_id_override: None,
             span_origin: None,
+            dataset_record_origin: None,
             span_attributes_extra: HashMap::new(),
         }
     }
 
+    pub(crate) fn dataset_record_origin(
+        mut self,
+        origin: crate::types::DatasetRecordOrigin,
+    ) -> Self {
+        self.dataset_record_origin = Some(origin);
+        self
+    }
+
     /// Set the span name before its first snapshot is submitted.
     pub fn name(mut self, name: impl Into<String>) -> Self {
         self.name = Some(name.into());
@@ -539,6 +550,7 @@ impl SpanBuilder {
                 name: self.name,
                 environment: self.environment,
                 span_origin: self.span_origin,
+                dataset_record_origin: self.dataset_record_origin,
                 span_attributes_extra: self.span_attributes_extra,
                 propagated_event,
                 ..Default::default()
@@ -807,6 +819,7 @@ fn current_span_payload(inner: &mut SpanData) -> SpanPayload {
 
 #[derive(Clone, Default)]
 struct SpanData {
+    dataset_record_origin: Option<crate::types::DatasetRecordOrigin>,
     row_id: String,
     span_id: String,
     has_flushed: bool,
@@ -873,6 +886,12 @@ impl From<SpanData> for SpanPayload {
             event_data.apply_propagated_event(&propagated_event);
         }
 
+        if let Some(origin) = data.dataset_record_origin {
+            event_data
+                .extra
+                .insert("origin".into(), origin.into_value());
+        }
+
         Self {
             row_id: data.row_id,
             span_id: data.span_id,