diff --git a/examples/eval_async.rs b/examples/eval_async.rs index 130f100..f17894d 100644 --- a/examples/eval_async.rs +++ b/examples/eval_async.rs @@ -44,7 +44,7 @@ async fn main() -> Result<()> { "my-project", EvalOpts::builder() .name("uppercase-eval".to_string()) - .dataset(cases.into()) + .source(cases.into()) .task(Box::new(task(|input: &String, _hooks| { Box::pin(async move { Ok(input.to_uppercase()) }) }))) diff --git a/examples/eval_sync.rs b/examples/eval_sync.rs index ae065a5..f6fec87 100644 --- a/examples/eval_sync.rs +++ b/examples/eval_sync.rs @@ -42,7 +42,7 @@ fn main() -> anyhow::Result<()> { let result = evaluator.run( EvalOpts::builder() .name("sync-uppercase-eval".to_string()) - .dataset(cases.into()) + .source(cases.into()) .task(Box::new(task)) .scorers(scorers) .build(), diff --git a/src/eval/bt_runner.rs b/src/eval/bt_runner.rs index 1daae18..fcb2eed 100644 --- a/src/eval/bt_runner.rs +++ b/src/eval/bt_runner.rs @@ -10,7 +10,7 @@ use std::io::Write; use std::sync::{Arc, Mutex}; -use serde::{de::DeserializeOwned, Serialize}; +use serde::Serialize; use serde_json::{json, Value}; use crate::{BraintrustClient, BraintrustError, Result}; @@ -193,7 +193,7 @@ fn create_sse_writer() -> Option> { /// "my-project", /// EvalOpts::builder() /// .name("my-eval".to_string()) -/// .dataset(vec![Case { +/// .source(vec![Case { /// input: "hello".to_string(), /// expected: Some("HELLO".to_string()), /// ..Default::default() @@ -288,8 +288,8 @@ impl BtEvalRunner { /// If the eval name does not pass the configured filters, it is skipped silently. pub async fn eval(&mut self, project: &str, mut opts: EvalOpts) -> Result<()> where - I: Serialize + DeserializeOwned + Send + Sync + Clone + 'static, - O: Serialize + DeserializeOwned + Send + Sync + PartialEq + Clone + 'static, + I: Serialize + Send + Sync + 'static, + O: Serialize + Send + Sync + 'static, { let name = opts.name.clone(); @@ -349,7 +349,7 @@ impl BtEvalRunner { eprintln!("Evaluator '{}' failed: {}", name, e); } if self.config.terminate_on_failure { - return Err(e); + return Err(BraintrustError::Background(e.to_string())); } } } diff --git a/src/eval/dataset.rs b/src/eval/dataset.rs deleted file mode 100644 index 8eb9b44..0000000 --- a/src/eval/dataset.rs +++ /dev/null @@ -1,196 +0,0 @@ -use async_trait::async_trait; -use futures::stream::Stream; - -use crate::Result; - -use super::types::Case; - -/// Iterator trait for evaluation test cases -/// -/// This trait allows different dataset sources (in-memory vectors, API-backed datasets, etc.) -/// to provide test cases to the evaluator. Follows the Go SDK pattern where Dataset is -/// an interface/trait rather than a concrete type. -#[async_trait] -pub trait Dataset: Send + Sync -where - I: Clone, - O: Clone, -{ - /// Get the next case from the dataset - async fn next(&mut self) -> Result>>; - - /// Get the dataset ID (if backed by Braintrust API) - fn id(&self) -> Option { - None - } - - /// Get the dataset version (if backed by Braintrust API) - fn version(&self) -> Option { - None - } -} - -/// In-memory dataset backed by a vector -pub struct VecDataset -where - I: Clone, - O: Clone, -{ - cases: Vec>, - index: usize, -} - -impl VecDataset -where - I: Clone, - O: Clone, -{ - pub fn new(cases: Vec>) -> Self { - Self { cases, index: 0 } - } -} - -#[async_trait] -impl Dataset for VecDataset -where - I: Send + Sync + Clone, - O: Send + Sync + Clone, -{ - async fn next(&mut self) -> Result>> { - if self.index < self.cases.len() { - let case = self.cases[self.index].clone(); - self.index += 1; - Ok(Some(case)) - } else { - Ok(None) - } - } -} - -/// Implement Dataset for any sync iterator over Cases -/// -/// This allows using standard iterators as datasets without manual wrapping. -pub struct IterDataset -where - I: Clone, - O: Clone, - T: Iterator>, -{ - iter: T, -} - -impl IterDataset -where - I: Clone, - O: Clone, - T: Iterator>, -{ - pub fn new(iter: T) -> Self { - Self { iter } - } -} - -#[async_trait] -impl Dataset for IterDataset -where - T: Iterator> + Send + Sync, - I: Send + Sync + Clone, - O: Send + Sync + Clone, -{ - async fn next(&mut self) -> Result>> { - Ok(self.iter.next()) - } -} - -/// Implement Dataset for async streams -pub struct StreamDataset -where - I: Clone, - O: Clone, - S: Stream> + Send + Unpin, -{ - stream: S, -} - -impl StreamDataset -where - I: Clone, - O: Clone, - S: Stream> + Send + Unpin, -{ - pub fn new(stream: S) -> Self { - Self { stream } - } -} - -#[async_trait] -impl Dataset for StreamDataset -where - S: Stream> + Send + Sync + Unpin, - I: Send + Sync + Clone, - O: Send + Sync + Clone, -{ - async fn next(&mut self) -> Result>> { - use futures::stream::StreamExt; - Ok(self.stream.next().await) - } -} - -/// Convenience conversion from `Vec` to boxed Dataset -impl From>> for Box> -where - I: Send + Sync + Clone + 'static, - O: Send + Sync + Clone + 'static, -{ - fn from(cases: Vec>) -> Self { - Box::new(VecDataset::new(cases)) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[tokio::test] - async fn test_vec_dataset() { - let cases = vec![ - Case { - input: "hello", - expected: Some("HELLO"), - ..Default::default() - }, - Case { - input: "world", - expected: Some("WORLD"), - ..Default::default() - }, - ]; - - let mut dataset = VecDataset::new(cases); - - let case1 = dataset.next().await.unwrap(); - assert!(case1.is_some()); - assert_eq!(case1.unwrap().input, "hello"); - - let case2 = dataset.next().await.unwrap(); - assert!(case2.is_some()); - assert_eq!(case2.unwrap().input, "world"); - - let case3 = dataset.next().await.unwrap(); - assert!(case3.is_none()); - } - - #[tokio::test] - async fn test_vec_to_boxed_dataset() { - let cases = vec![Case { - input: "test", - expected: Some("TEST"), - ..Default::default() - }]; - - let mut dataset: Box> = cases.into(); - let case = dataset.next().await.unwrap(); - assert!(case.is_some()); - assert_eq!(case.unwrap().input, "test"); - } -} diff --git a/src/eval/error.rs b/src/eval/error.rs index fa98412..9798a2f 100644 --- a/src/eval/error.rs +++ b/src/eval/error.rs @@ -48,3 +48,94 @@ pub enum EvalConfigError { #[error("duplicate scorer name: {name}")] DuplicateScorerName { name: String }, } + +/// Run-level failures retain their boundary and source location. +/// Case indices are zero-based; None denotes run-level recording work. +#[derive(Debug, thiserror::Error)] +pub enum EvalRunFailure { + #[error(transparent)] + Config(#[from] EvalConfigError), + #[error("case source failed at index {case_index}: {source}")] + Source { + case_index: usize, + #[source] + source: super::SourceError, + }, + #[error("recording failed for case {case_index:?}: {source}")] + Recording { + case_index: Option, + #[source] + source: crate::BraintrustError, + }, + #[error(transparent)] + Execution(#[from] crate::ExecutionError), + #[error(transparent)] + Admission(#[from] crate::EnqueueError), +} + +/// An incomplete run preserves finished case results and every run-level failure. +/// Case computation errors remain in those results rather than aborting the run. +pub struct EvalRunError { + pub failures: Vec, + pub partial: Option>>, +} +impl EvalRunError { + fn new(failure: EvalRunFailure) -> Self { + Self { + failures: vec![failure], + partial: None, + } + } +} +impl std::fmt::Debug for EvalRunError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("EvalRunError") + .field("failures", &self.failures) + .field( + "completed_cases", + &self.partial.as_ref().map(|summary| summary.results.len()), + ) + .finish() + } +} +impl std::fmt::Display for EvalRunError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + if let Some(first) = self.failures.first() { + write!(f, "{first}")?; + } + if self.failures.len() > 1 { + write!(f, " ({} run failures)", self.failures.len())?; + } + Ok(()) + } +} +impl std::error::Error for EvalRunError { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + self.failures + .first() + .map(|failure| failure as &dyn std::error::Error) + } +} +impl From for EvalRunError { + fn from(error: EvalConfigError) -> Self { + Self::new(EvalRunFailure::Config(error)) + } +} +impl From for EvalRunError { + fn from(source: crate::BraintrustError) -> Self { + Self::new(EvalRunFailure::Recording { + case_index: None, + source, + }) + } +} +impl From for EvalRunError { + fn from(error: crate::ExecutionError) -> Self { + Self::new(EvalRunFailure::Execution(error)) + } +} +impl From for EvalRunError { + fn from(error: crate::EnqueueError) -> Self { + Self::new(EvalRunFailure::Admission(error)) + } +} diff --git a/src/eval/mod.rs b/src/eval/mod.rs index 1754c27..c785cd6 100644 --- a/src/eval/mod.rs +++ b/src/eval/mod.rs @@ -13,7 +13,7 @@ //! use braintrust_sdk_rust::*; //! use braintrust_sdk_rust::eval::*; //! -//! # async fn example() -> Result<()> { +//! # async fn example() -> anyhow::Result<()> { //! let client = BraintrustClient::builder() //! .api_key("sk-...") //! .build() @@ -33,7 +33,7 @@ //! .run( //! EvalOpts::builder() //! .name("uppercase-eval".to_string()) -//! .dataset(cases.into()) +//! .source(cases.into()) //! .task(Box::new(task(|input: &String, _hooks| Box::pin(async move { //! Ok(input.to_uppercase()) //! })))) @@ -49,19 +49,21 @@ mod bt_runner; mod case_runner; -mod dataset; mod error; mod runner; mod scorer; +mod source; mod task; mod types; -pub use error::{CaseError, EvalConfigError}; +pub use error::{CaseError, EvalConfigError, EvalRunError, EvalRunFailure}; // Re-export public API pub use bt_runner::{summary_to_value, BtEvalRunner}; -pub use dataset::{Dataset, IterDataset, StreamDataset, VecDataset}; pub use runner::{EvalOpts, Evaluator, SyncEvaluator}; pub use scorer::{scorer, ExactMatch, LevenshteinScorer, NumericDiffScorer, Scorer, SyncFnScorer}; +pub use source::{ + CaseSource, DatasetVersionRef, IterCaseSource, SourceError, StreamCaseSource, VecCaseSource, +}; pub use task::{task, SimpleFnTask, SyncFnTask, Task}; pub use types::{Case, EvalResult, EvalSummary, Score, ScoreStats, ScorerArgs, Scores, TaskHooks}; diff --git a/src/eval/runner.rs b/src/eval/runner.rs index 38ea6e5..19453c4 100644 --- a/src/eval/runner.rs +++ b/src/eval/runner.rs @@ -1,19 +1,18 @@ use futures::stream::{FuturesUnordered, StreamExt}; -use serde::{de::DeserializeOwned, Serialize}; +use serde::Serialize; use serde_json::{Map, Value}; use std::marker::PhantomData; use std::sync::Arc; -use tokio::sync::Semaphore; use crate::{BraintrustClient, Result, SpanHandle}; use super::{ case_runner::CaseRunner, - dataset::Dataset, scorer::Scorer, + source::CaseSource, task::Task, types::{Case, EvalResult, EvalSummary}, - EvalConfigError, + EvalConfigError, EvalRunError, EvalRunFailure, }; /// Options for running an evaluation @@ -21,8 +20,8 @@ use super::{ pub struct EvalOpts { /// Name of the evaluation pub name: String, - /// Dataset providing test cases - pub dataset: Box>, + /// Owned source providing test cases + pub source: Box>, /// Task to execute on each input pub task: Box>, /// Scorers to evaluate outputs @@ -73,8 +72,8 @@ pub struct Evaluator { impl Evaluator where - I: Serialize + DeserializeOwned + Send + Sync + Clone + 'static, - O: Serialize + DeserializeOwned + Send + Sync + PartialEq + Clone + 'static, + I: Serialize + Send + Sync + 'static, + O: Serialize + Send + Sync + 'static, { /// Create a new evaluator pub fn new(client: BraintrustClient) -> Self { @@ -93,9 +92,11 @@ where } /// Run the evaluation - pub async fn run(&self, mut opts: EvalOpts) -> Result> { - opts.validate() - .map_err(|error| crate::BraintrustError::InvalidConfig(error.to_string()))?; + pub async fn run( + &self, + mut opts: EvalOpts, + ) -> std::result::Result, EvalRunError> { + opts.validate()?; if !opts.quiet { println!("Running evaluation: {}", opts.name); } @@ -118,83 +119,90 @@ where log_builder = log_builder.tags(tags.clone()); } - root_span.log(log_builder.build().unwrap()); + let root_write = root_span.log(log_builder.build().unwrap()); - // Get dataset ID if available - let dataset_id = opts.dataset.id(); - - // Collect all results + let dataset = opts.source.provenance(); let mut results = Vec::new(); + let mut failures = Vec::new(); let mut case_count = 0; - - // Create semaphore for concurrency control - let semaphore = Arc::new(Semaphore::new(opts.max_concurrency)); - - // Process cases with controlled concurrency let mut futures = FuturesUnordered::new(); loop { - // Try to get next case - let case = match opts.dataset.next().await { + // Retain the next-read future while polling active cases. A source + // may wait for prior computation, and need not be cancellation-safe. + let mut next = Box::pin(opts.source.next()); + let next_case = loop { + if futures.is_empty() { + break next.await; + } + tokio::select! { + next_case = next.as_mut() => break next_case, + Some((index, result)) = futures.next() => { + Self::collect_case_result(index, result, &mut results, &mut failures); + if !failures.is_empty() { break Ok(None); } + } + } + }; + let case = match next_case { Ok(Some(case)) => case, Ok(None) => break, - Err(e) => { - eprintln!("Error fetching case: {}", e); + Err(source) => { + failures.push(EvalRunFailure::Source { + case_index: case_count, + source, + }); break; } }; - case_count += 1; - - // Clone what we need for the async task + let index = case_count - 1; let client = Arc::clone(&self.client); let root_span = root_span.clone(); let task = &opts.task; let scorers = &opts.scorers; - let semaphore = Arc::clone(&semaphore); let quiet = opts.quiet; - - // Acquire permit before spawning - let permit = semaphore.clone().acquire_owned().await.unwrap(); - - // Process this case - let future = async move { + futures.push(async move { let result = Self::process_case( case, task.as_ref(), scorers, &root_span, &client, - case_count, + index + 1, quiet, ) .await; - drop(permit); // Release permit - result - }; - - futures.push(future); - - // Collect completed futures if we have many pending - while futures.len() >= opts.max_concurrency { - if let Some(result) = futures.next().await { - results.push(result?); + (index, result) + }); + if futures.len() >= opts.max_concurrency { + if let Some((index, result)) = futures.next().await { + Self::collect_case_result(index, result, &mut results, &mut failures); + } + if !failures.is_empty() { + break; } } } - - // Collect remaining futures - while let Some(result) = futures.next().await { - results.push(result?); + // Source/recording failure does not abandon already scheduled cases. + while let Some((index, result)) = futures.next().await { + Self::collect_case_result(index, result, &mut results, &mut failures); } - - // End root span - root_span.end().await?; + let end = root_span.end(); + for result in [root_write.await, end.await] { + if let Err(error) = result { + failures.push(EvalRunFailure::Recording { + case_index: None, + source: error.into(), + }); + } + } + results.sort_unstable_by_key(|(index, _)| *index); + let results: Vec<_> = results.into_iter().map(|(_, result)| result).collect(); // Compute statistics let score_stats = EvalSummary::::compute_stats(&results); let successful = results.iter().filter(|r| r.errors.is_empty()).count(); - let failed = results.len() - successful; + let failed = case_count - successful; if !opts.quiet { println!("\nEvaluation complete!"); @@ -210,15 +218,38 @@ where } } - Ok(EvalSummary { + let summary = EvalSummary { name: opts.name, results, score_stats, total_cases: case_count, successful_cases: successful, failed_cases: failed, - dataset_id, - }) + dataset, + }; + if failures.is_empty() { + Ok(summary) + } else { + Err(EvalRunError { + failures, + partial: Some(Box::new(summary)), + }) + } + } + + fn collect_case_result( + index: usize, + result: Result>, + results: &mut Vec<(usize, EvalResult)>, + failures: &mut Vec, + ) { + match result { + Ok(result) => results.push((index, result)), + Err(source) => failures.push(EvalRunFailure::Recording { + case_index: Some(index), + source, + }), + } } /// Process a single case @@ -322,8 +353,8 @@ pub struct SyncEvaluator { impl SyncEvaluator where - I: Serialize + DeserializeOwned + Send + Sync + Clone + 'static, - O: Serialize + DeserializeOwned + Send + Sync + PartialEq + Clone + 'static, + I: Serialize + Send + Sync + 'static, + O: Serialize + Send + Sync + 'static, { /// Create a new sync evaluator pub fn new(client: BraintrustClient) -> Result { @@ -336,7 +367,10 @@ where } /// Run the evaluation synchronously - pub fn run(&self, opts: EvalOpts) -> Result> { + pub fn run( + &self, + opts: EvalOpts, + ) -> std::result::Result, EvalRunError> { self.runtime.block_on(self.evaluator.run(opts)) } } @@ -344,7 +378,7 @@ where #[cfg(test)] mod tests { use super::*; - use crate::eval::{dataset::VecDataset, scorer::ExactMatch, task::SimpleFnTask, types::Case}; + use crate::eval::{scorer::ExactMatch, source::VecCaseSource, task::SimpleFnTask, types::Case}; // Note: Full integration tests require a BraintrustClient setup // These tests verify type correctness and compilation @@ -357,14 +391,14 @@ mod tests { ..Default::default() }]; - let dataset: Box> = Box::new(VecDataset::new(cases)); + let dataset: Box> = Box::new(VecCaseSource::new(cases)); let task: Box> = Box::new(SimpleFnTask::new(|input: &String| Ok(input.to_uppercase()))); let scorers: Vec>> = vec![Box::new(ExactMatch::new())]; let _opts = EvalOpts::builder() .name("test-eval".to_string()) - .dataset(dataset) + .source(dataset) .task(task) .scorers(scorers) .build(); diff --git a/src/eval/source.rs b/src/eval/source.rs new file mode 100644 index 0000000..025f2ee --- /dev/null +++ b/src/eval/source.rs @@ -0,0 +1,105 @@ +//! Owned, fallible case sources, independent of recording and presentation. +//! +//! ```no_run +//! use braintrust_sdk_rust::eval::{Case, CaseSource, IterCaseSource}; +//! +//! # async fn example() -> anyhow::Result<()> { +//! let cases = (0..3).map(|input| Ok(Case:: { +//! input, ..Default::default() +//! })); +//! let mut source: Box> = +//! Box::new(IterCaseSource::new(cases)); +//! assert_eq!(source.next().await?.unwrap().input, 0); +//! # Ok(()) +//! # } +//! ``` +use super::Case; +use async_trait::async_trait; +use futures::{Stream, StreamExt}; +use serde::Serialize; +use std::pin::Pin; + +/// Application errors from a source retain their original cause. +pub type SourceError = anyhow::Error; + +/// Dataset identity captured by a source and retained with its run result. +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +pub struct DatasetVersionRef { + pub dataset_id: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub version: Option, +} + +/// One run owns its source and reads it sequentially. Sources need Send, +/// without requiring Sync, cloning user values, or serializing local cases. +#[async_trait] +pub trait CaseSource: Send { + async fn next(&mut self) -> Result>, SourceError>; + fn provenance(&self) -> Option { + None + } +} + +/// Move cases out of an in-memory vector without retaining consumed values. +pub struct VecCaseSource { + cases: std::vec::IntoIter>, +} +impl VecCaseSource { + pub fn new(cases: Vec>) -> Self { + Self { + cases: cases.into_iter(), + } + } +} +#[async_trait] +impl CaseSource for VecCaseSource { + async fn next(&mut self) -> Result>, SourceError> { + Ok(self.cases.next()) + } +} + +/// Adapt a fallible iterator. For infallible iterators, use `iter.map(Ok)`. +pub struct IterCaseSource { + iter: T, +} +impl IterCaseSource { + pub fn new(iter: T) -> Self { + Self { iter } + } +} +#[async_trait] +impl CaseSource for IterCaseSource +where + T: Iterator, SourceError>> + Send, +{ + async fn next(&mut self) -> Result>, SourceError> { + self.iter.next().transpose() + } +} + +/// Adapt a fallible stream, including streams that are not Unpin or Sync. +/// For infallible streams, use `stream.map(Ok)`. +pub struct StreamCaseSource { + stream: Pin>, +} +impl StreamCaseSource { + pub fn new(stream: S) -> Self { + Self { + stream: Box::pin(stream), + } + } +} +#[async_trait] +impl CaseSource for StreamCaseSource +where + S: Stream, SourceError>> + Send, +{ + async fn next(&mut self) -> Result>, SourceError> { + self.stream.next().await.transpose() + } +} +impl From>> for Box> { + fn from(cases: Vec>) -> Self { + Box::new(VecCaseSource::new(cases)) + } +} diff --git a/src/eval/types.rs b/src/eval/types.rs index c7356ce..9078148 100644 --- a/src/eval/types.rs +++ b/src/eval/types.rs @@ -2,7 +2,7 @@ use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; use std::sync::{Arc, Mutex}; -use super::CaseError; +use super::{CaseError, DatasetVersionRef}; use crate::SpanHandle; /// A single test case for evaluation @@ -159,9 +159,9 @@ pub struct EvalSummary { pub successful_cases: usize, /// Number of failed evaluations pub failed_cases: usize, - /// Dataset ID (if using Braintrust dataset) + /// Full dataset provenance captured before source consumption. #[serde(skip_serializing_if = "Option::is_none")] - pub dataset_id: Option, + pub dataset: Option, } impl EvalSummary { diff --git a/tests/eval_callables.rs b/tests/eval_callables.rs index 726b5d3..9b6f47a 100644 --- a/tests/eval_callables.rs +++ b/tests/eval_callables.rs @@ -1,10 +1,10 @@ use async_trait::async_trait; use braintrust_sdk_rust::{ eval::{ - scorer, task, Case, Dataset, EvalConfigError, EvalOpts, Evaluator, ExactMatch, Score, + scorer, task, Case, CaseSource, EvalConfigError, EvalOpts, Evaluator, ExactMatch, Score, Scorer, ScorerArgs, Scores, SimpleFnTask, Task, TaskHooks, }, - BraintrustClient, BraintrustError, + BraintrustClient, }; use futures::future::BoxFuture; use std::sync::{ @@ -139,8 +139,8 @@ async fn scorer_helpers_accept_borrowed_and_owned_names_and_captured_state() { struct UntouchedSource(Arc); #[async_trait] -impl Dataset for UntouchedSource { - async fn next(&mut self) -> braintrust_sdk_rust::Result>> { +impl CaseSource for UntouchedSource { + async fn next(&mut self) -> anyhow::Result>> { self.0.store(true, Ordering::SeqCst); panic!("invalid plans must not read their source"); } @@ -180,17 +180,19 @@ async fn invalid_concurrency_and_names_fail_before_login_recording_or_source_acc as Box> }) .collect(); - let opts = EvalOpts::builder() - .name("invalid".into()) - .dataset(Box::new(UntouchedSource(accessed.clone())) as Box>) - .task(Box::new(SimpleFnTask::new(|input: &usize| Ok(*input)))) - .scorers(scorers) - .max_concurrency(concurrency) - .quiet(true) - .build(); + let opts = + EvalOpts::builder() + .name("invalid".into()) + .source(Box::new(UntouchedSource(accessed.clone())) + as Box>) + .task(Box::new(SimpleFnTask::new(|input: &usize| Ok(*input)))) + .scorers(scorers) + .max_concurrency(concurrency) + .quiet(true) + .build(); assert_eq!(opts.validate(), Err(expected.clone())); assert!( - matches!(evaluator.run(opts).await, Err(BraintrustError::InvalidConfig(message)) if message == expected.to_string()) + matches!(evaluator.run(opts).await, Err(error) if matches!(&error.failures[0], braintrust_sdk_rust::eval::EvalRunFailure::Config(actual) if actual == &expected) && error.partial.is_none()) ); } assert!(!accessed.load(Ordering::SeqCst)); @@ -217,7 +219,7 @@ fn duplicate_names_are_rejected_across_custom_tuple_and_builtin_scorers() { ] { let opts = EvalOpts::builder() .name("duplicate".into()) - .dataset(Vec::>::new().into()) + .source(Vec::>::new().into()) .task(Box::new(SimpleFnTask::new(|input: &usize| Ok(*input)))) .scorers(vec![Box::new(ExactMatch::new()), callback]) .build(); diff --git a/tests/eval_sources.rs b/tests/eval_sources.rs new file mode 100644 index 0000000..befff5b --- /dev/null +++ b/tests/eval_sources.rs @@ -0,0 +1,113 @@ +use braintrust_sdk_rust::eval::{ + Case, CaseSource, IterCaseSource, StreamCaseSource, VecCaseSource, +}; +use std::{ + cell::Cell, + sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, + }, +}; + +fn case(input: I) -> Case { + Case { + input, + expected: None, + metadata: None, + tags: None, + id: None, + xact_id: None, + created: None, + } +} +struct OwnedInput { + value: Cell, + drops: Arc, +} +impl Drop for OwnedInput { + fn drop(&mut self) { + self.drops.fetch_add(1, Ordering::SeqCst); + } +} + +#[test] +fn vectors_move_nonclone_nonsync_cases_and_release_consumed_values() { + futures::executor::block_on(async { + let drops = Arc::new(AtomicUsize::new(0)); + let cases: Vec>> = (0..2) + .map(|value| { + case(OwnedInput { + value: Cell::new(value), + drops: drops.clone(), + }) + }) + .collect(); + let mut source: Box>> = cases.into(); + let first = source.next().await.unwrap().unwrap(); + assert_eq!(first.input.value.get(), 0); + drop(first); + assert_eq!( + drops.load(Ordering::SeqCst), + 1, + "the source must release consumed cases" + ); + drop(source); + assert_eq!(drops.load(Ordering::SeqCst), 2); + }); +} + +#[derive(Debug, thiserror::Error)] +#[error("read failed")] +struct ReadFailure; + +#[tokio::test] +async fn fallible_iterators_keep_error_causes_and_do_not_require_sync() { + type LocalCase = Case, Cell>; + let items: Vec> = + vec![Ok(case(Cell::new(7))), Err(ReadFailure.into())]; + let mut source: Box, Cell>> = + Box::new(IterCaseSource::new(items.into_iter())); + assert_eq!(source.next().await.unwrap().unwrap().input.get(), 7); + assert!(source + .next() + .await + .unwrap_err() + .downcast_ref::() + .is_some()); + assert!(source.next().await.unwrap().is_none()); +} + +#[tokio::test] +async fn pinned_streams_allow_nonunpin_nonsync_state_and_owned_error_causes() { + let value = Cell::new(8); + let stream = futures::stream::once(async move { + value.set(value.get() + 1); + tokio::task::yield_now().await; + Ok::, Cell>, anyhow::Error>(case(value)) + }); + let mut source: Box, Cell>> = + Box::new(StreamCaseSource::new(stream)); + assert_eq!(source.next().await.unwrap().unwrap().input.get(), 9); + assert!(source.next().await.unwrap().is_none()); + let stream = + futures::stream::once(async { Err::, anyhow::Error>(ReadFailure.into()) }); + let mut source = StreamCaseSource::new(stream); + assert!(source + .next() + .await + .unwrap_err() + .downcast_ref::() + .is_some()); +} + +#[tokio::test] +async fn empty_and_infallible_sources_end_without_synthetic_cases() { + let mut empty = VecCaseSource::::new(vec![]); + assert!(empty.next().await.unwrap().is_none()); + assert!(empty.next().await.unwrap().is_none()); + let mut source = + IterCaseSource::new((0..2).map(|value| Ok::<_, anyhow::Error>(case::<_, usize>(value)))); + assert_eq!(source.next().await.unwrap().unwrap().input, 0); + assert_eq!(source.next().await.unwrap().unwrap().input, 1); + assert!(source.next().await.unwrap().is_none()); +} diff --git a/tests/eval_test.rs b/tests/eval_test.rs index 36c7148..e139f1b 100644 --- a/tests/eval_test.rs +++ b/tests/eval_test.rs @@ -1,4 +1,4 @@ -use braintrust_sdk_rust::eval::Dataset; +use braintrust_sdk_rust::eval::CaseSource; use braintrust_sdk_rust::eval::*; use braintrust_sdk_rust::*; @@ -80,7 +80,7 @@ async fn test_basic_eval_flow() { .run( EvalOpts::builder() .name("test-basic-eval".to_string()) - .dataset(cases.into()) + .source(cases.into()) .task(Box::new(callback)) .scorers(scorers) .quiet(true) @@ -110,7 +110,7 @@ fn test_eval_types_compile() { ..Default::default() }]; - let _dataset: Box> = cases.into(); + let _dataset: Box> = cases.into(); let _task: Box> = Box::new(SimpleFnTask::new(|input: &String| Ok(input.to_uppercase()))); let _scorers: Vec>> = vec![Box::new(ExactMatch::new())]; @@ -183,14 +183,14 @@ fn test_eval_opts_builder() { ..Default::default() }]; - let dataset: Box> = cases.into(); + let dataset: Box> = cases.into(); let task: Box> = Box::new(SimpleFnTask::new(|input: &String| Ok(input.to_uppercase()))); let scorers: Vec>> = vec![Box::new(ExactMatch::new())]; let _opts = EvalOpts::builder() .name("test-eval".to_string()) - .dataset(dataset) + .source(dataset) .task(task) .scorers(scorers) .max_concurrency(10) @@ -212,7 +212,7 @@ async fn recorded_scorer_failure_keeps_valid_scores_and_finishes_the_case() { .run( EvalOpts::builder() .name("failed-scorer".into()) - .dataset( + .source( vec![Case { input: "a".to_owned(), expected: Some("A".to_owned()), @@ -258,3 +258,298 @@ async fn recorded_scorer_failure_keeps_valid_scores_and_finishes_the_case() { .any(|row| row["span_id"].as_str() == case.span_id.as_deref() && row["metrics"]["end"].is_number())); } + +#[tokio::test] +async fn source_failure_drains_active_cases_preserves_provenance_and_finishes_recording() { + use async_trait::async_trait; + use std::{ + cell::Cell, + sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, + }, + }; + use tokio::sync::{Notify, Semaphore}; + #[derive(serde::Serialize)] + struct Input(usize); + #[derive(serde::Serialize)] + struct Output(usize); + #[derive(Debug, thiserror::Error)] + #[error("source read failed")] + struct ReadFailure; + struct Source { + reads: Cell, + started: Arc, + failed: Arc, + } + #[async_trait] + impl CaseSource for Source { + async fn next(&mut self) -> anyhow::Result>> { + let index = self.reads.get(); + self.reads.set(index + 1); + if index == 1 { + self.started.notified().await; + } + if index == 2 { + self.failed.notify_one(); + return Err(ReadFailure.into()); + } + Ok(Some(Case { + input: Input(index), + expected: None, + metadata: None, + tags: None, + id: Some(format!("case-{index}")), + xact_id: Some(format!("row-version-{index}")), + created: None, + })) + } + fn provenance(&self) -> Option { + Some(DatasetVersionRef { + dataset_id: "source-dataset".into(), + version: Some(format!("version-{}", self.reads.get())), + }) + } + } + let (client, server) = recording_client().await; + let started = Arc::new(Notify::new()); + let failed = Arc::new(Notify::new()); + let second_done = Arc::new(Notify::new()); + let gates = Arc::new([Semaphore::new(0), Semaphore::new(0)]); + let calls = Arc::new(AtomicUsize::new(0)); + let callback = { + let started = started.clone(); + let second_done = second_done.clone(); + let gates = gates.clone(); + let calls = calls.clone(); + task(move |input: &Input, _hooks| { + let started = started.clone(); + let second_done = second_done.clone(); + let gates = gates.clone(); + let calls = calls.clone(); + Box::pin(async move { + calls.fetch_add(1, Ordering::SeqCst); + if input.0 == 0 { + started.notify_one(); + } + gates[input.0].acquire().await.unwrap().forget(); + if input.0 == 1 { + second_done.notify_one(); + } + Ok(Output(input.0 * 2)) + }) + }) + }; + let opts = EvalOpts::builder() + .name("source-failure".into()) + .source(Box::new(Source { + reads: Cell::new(0), + started, + failed: failed.clone(), + }) as Box>) + .task(Box::new(callback)) + .scorers(vec![]) + .max_concurrency(3) + .quiet(true) + .build(); + let run = tokio::spawn(async move { Evaluator::new(client).run(opts).await }); + tokio::time::timeout(std::time::Duration::from_secs(3), failed.notified()) + .await + .expect("pending source reads must allow active task progress"); + gates[1].add_permits(1); + second_done.notified().await; + assert!( + !run.is_finished(), + "source failure must still wait for the other scheduled case" + ); + gates[0].add_permits(1); + let error = match tokio::time::timeout(std::time::Duration::from_secs(3), run) + .await + .unwrap() + .unwrap() + { + Err(error) => error, + Ok(_) => panic!("source failure must not become a successful summary"), + }; + assert_eq!(calls.load(Ordering::SeqCst), 2); + let [EvalRunFailure::Source { case_index, source }] = error.failures.as_slice() else { + panic!("expected a typed source failure") + }; + assert_eq!(*case_index, 2); + assert!(source.downcast_ref::().is_some()); + let partial = error.partial.unwrap(); + assert_eq!(partial.total_cases, 2); + assert_eq!(partial.successful_cases, 2); + assert_eq!(partial.failed_cases, 0); + let dataset = partial.dataset.unwrap(); + assert_eq!(dataset.dataset_id, "source-dataset"); + assert_eq!(dataset.version.as_deref(), Some("version-0")); + assert_eq!( + partial + .results + .iter() + .map(|case| case.input.0) + .collect::>(), + [0, 1] + ); + assert_eq!(partial.results[1].case_id.as_deref(), Some("case-1")); + assert_eq!(partial.results[1].xact_id.as_deref(), Some("row-version-1")); + assert_eq!(partial.results[1].output.as_ref().unwrap().0, 2); + let rows: Vec = server + .received_requests() + .await + .unwrap() + .iter() + .filter(|request| request.url.path() == "/logs3") + .flat_map(|request| { + serde_json::from_slice::(&request.body).unwrap()["rows"] + .as_array() + .unwrap() + .clone() + }) + .collect(); + let ended: std::collections::HashSet<_> = rows + .iter() + .filter(|row| row["metrics"]["end"].is_number()) + .map(|row| row["span_id"].as_str().unwrap()) + .collect(); + assert_eq!( + ended.len(), + 3, + "both cases and the run span must finish before returning partial results" + ); +} + +#[tokio::test] +async fn failure_before_the_first_case_returns_an_empty_partial_run_and_finalizes_root() { + let (client, server) = recording_client().await; + let items: Vec>> = + vec![Err(anyhow::anyhow!("source unavailable"))]; + let opts = EvalOpts::builder() + .name("empty-failure".into()) + .source( + Box::new(IterCaseSource::new(items.into_iter())) as Box> + ) + .task(Box::new(SimpleFnTask::new( + |_: &String| -> anyhow::Result { panic!("no case was read") }, + ))) + .scorers(vec![]) + .quiet(true) + .build(); + let error = Evaluator::new(client).run(opts).await.unwrap_err(); + assert!( + matches!(&error.failures[0],EvalRunFailure::Source {case_index:0,source} if source.to_string()=="source unavailable") + ); + let partial = error.partial.unwrap(); + assert_eq!(partial.total_cases, 0); + assert_eq!(partial.successful_cases, 0); + assert_eq!(partial.failed_cases, 0); + assert!(partial.results.is_empty()); + assert!(partial.dataset.is_none()); + let ended = server + .received_requests() + .await + .unwrap() + .iter() + .filter(|r| r.url.path() == "/logs3") + .any(|request| { + serde_json::from_slice::(&request.body).unwrap()["rows"] + .as_array() + .unwrap() + .iter() + .any(|row| row["metrics"]["end"].is_number()) + }); + assert!( + ended, + "the run span must finish even when source failure precedes all cases" + ); +} + +#[tokio::test] +async fn recording_failure_stops_new_cases_but_drains_already_scheduled_cases() { + use std::sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, + }; + use tokio::sync::{Notify, Semaphore}; + use wiremock::{ + matchers::{method, path}, + Mock, ResponseTemplate, + }; + let (client, server) = recording_client().await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(400)) + .with_priority(1) + .mount(&server) + .await; + let calls = Arc::new(AtomicUsize::new(0)); + let entered = Arc::new(Notify::new()); + let gate = Arc::new(Semaphore::new(0)); + let callback = { + let calls = calls.clone(); + let entered = entered.clone(); + let gate = gate.clone(); + task(move |input: &usize, _hooks| { + let calls = calls.clone(); + let entered = entered.clone(); + let gate = gate.clone(); + Box::pin(async move { + if calls.fetch_add(1, Ordering::SeqCst) == 1 { + entered.notify_one(); + } + gate.acquire().await.unwrap().forget(); + Ok(*input) + }) + }) + }; + let opts = EvalOpts::builder() + .name("record-failure".into()) + .source( + (0..3) + .map(|input| Case { + input, + ..Default::default() + }) + .collect::>>() + .into(), + ) + .task(Box::new(callback)) + .scorers(vec![]) + .max_concurrency(2) + .quiet(true) + .build(); + let run = tokio::spawn(async move { Evaluator::new(client).run(opts).await }); + tokio::time::timeout(std::time::Duration::from_secs(3), entered.notified()) + .await + .unwrap(); + gate.add_permits(2); + let error = tokio::time::timeout(std::time::Duration::from_secs(3), run) + .await + .unwrap() + .unwrap() + .unwrap_err(); + assert_eq!( + calls.load(Ordering::SeqCst), + 2, + "the third case must remain unread after recording failure" + ); + assert!(error.failures.iter().any(|failure| matches!( + failure, + EvalRunFailure::Recording { + case_index: Some(0), + .. + } + ))); + assert!(error.failures.iter().any(|failure| matches!( + failure, + EvalRunFailure::Recording { + case_index: Some(1), + .. + } + ))); + let partial = error.partial.unwrap(); + assert_eq!(partial.total_cases, 2); + assert_eq!(partial.failed_cases, 2); + assert_eq!(partial.successful_cases, 0); +}