diff --git a/src/completion.rs b/src/completion.rs index 89e7ccb..9a35ef8 100644 --- a/src/completion.rs +++ b/src/completion.rs @@ -18,7 +18,7 @@ use tokio::sync::oneshot; /// use braintrust_sdk_rust::{BraintrustClient, SpanHandle, SpanLog}; /// use std::time::Duration; /// -/// # async fn example(client: &BraintrustClient, span: &SpanHandle) +/// # async fn example(client: &BraintrustClient, span: &SpanHandle) /// # -> Result<(), Box> { /// let completion = span.log(SpanLog::default()); // admission happens now /// if let Err(error) = &*completion { diff --git a/src/dataset.rs b/src/dataset.rs index 9d80ef5..5bcacf7 100644 --- a/src/dataset.rs +++ b/src/dataset.rs @@ -27,7 +27,6 @@ use std::collections::VecDeque; use std::fmt; use std::sync::Arc; -use async_trait::async_trait; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; use tokio::sync::OnceCell; @@ -35,8 +34,8 @@ use uuid::Uuid; use crate::api::datasets::{DatasetFields, DatasetWrite}; use crate::error::Result; +use crate::internal::ports::{DatasetApi, DeliveryControl}; use crate::internal::session::CredentialScope; -use crate::span::SpanSubmitter; // ============================================================================ // Dataset Insert Types @@ -237,53 +236,8 @@ pub(crate) struct DatasetMetadata { pub project_id: String, } -pub(crate) use crate::api::btql::{BTQLQuery, BTQLQueryInner, BTQLResponse}; -pub(crate) use crate::api::datasets::{ - DatasetRegisterRequest, DatasetRegisterResponse, DatasetSummaryResponse, -}; - -// ============================================================================ -// Traits for API Operations -// ============================================================================ - -/// Trait for registering datasets. -#[async_trait] -pub(crate) trait DatasetRegistrar: Send + Sync { - /// Register a dataset with the API. - async fn register_dataset( - &self, - scope: &CredentialScope, - request: DatasetRegisterRequest, - ) -> Result; -} - -/// Trait for fetching dataset records via BTQL. -#[async_trait] -pub(crate) trait DatasetFetcher: Send + Sync { - /// Fetch dataset records via BTQL query. - async fn fetch_dataset_records( - &self, - scope: &CredentialScope, - query: BTQLQuery, - ) -> Result; -} - -/// Trait for fetching dataset summary. -#[async_trait] -pub(crate) trait DatasetSummarizer: Send + Sync { - /// Fetch dataset summary. - async fn fetch_dataset_summary( - &self, - scope: &CredentialScope, - dataset_id: &str, - ) -> Result; -} - -/// Direct dataset mutation port; writes are acknowledged by the endpoint. -#[async_trait] -pub(crate) trait DatasetWriter: Send + Sync { - async fn write_dataset_record(&self, scope: &CredentialScope, row: DatasetWrite) -> Result<()>; -} +use crate::api::btql::{BTQLQuery, BTQLQueryInner}; +use crate::api::datasets::DatasetRegisterRequest; // ============================================================================ // DatasetBuilder Error Types @@ -319,11 +273,9 @@ impl std::error::Error for DatasetBuilderError {} // ============================================================================ /// Builder for creating datasets with lazy registration. -#[allow(private_bounds)] -pub struct DatasetBuilder< - S: SpanSubmitter + DatasetRegistrar + DatasetFetcher + DatasetSummarizer + DatasetWriter, -> { - submitter: Arc, +pub struct DatasetBuilder { + api: Arc, + lifecycle: Arc, scope: CredentialScope, project_name: Option, dataset_name: Option, @@ -331,12 +283,11 @@ pub struct DatasetBuilder< metadata: Option>, } -impl Clone - for DatasetBuilder -{ +impl Clone for DatasetBuilder { fn clone(&self) -> Self { Self { - submitter: Arc::clone(&self.submitter), + api: self.api.clone(), + lifecycle: self.lifecycle.clone(), scope: self.scope.clone(), project_name: self.project_name.clone(), dataset_name: self.dataset_name.clone(), @@ -346,20 +297,16 @@ impl DatasetBuilder -{ +impl DatasetBuilder { /// Create a new DatasetBuilder (internal use). - pub(crate) fn new(submitter: Arc, scope: CredentialScope) -> Self { + pub(crate) fn new( + api: Arc, + lifecycle: Arc, + scope: CredentialScope, + ) -> Self { Self { - submitter, + api, + lifecycle, scope, project_name: None, dataset_name: None, @@ -401,7 +348,7 @@ impl< } /// Build the dataset (lazy registration). - pub fn build(self) -> std::result::Result, DatasetBuilderError> { + pub fn build(self) -> std::result::Result { let project_name = self .project_name .ok_or(DatasetBuilderError::MissingProjectName)?; @@ -410,7 +357,8 @@ impl< .ok_or(DatasetBuilderError::MissingDatasetName)?; Ok(Dataset { - submitter: self.submitter, + api: self.api, + lifecycle: self.lifecycle, scope: self.scope, project_name, dataset_name, @@ -428,11 +376,9 @@ impl< /// A dataset for storing test/evaluation data. /// /// Datasets are lazily registered with the API on first use. -#[allow(private_bounds)] -pub struct Dataset< - S: SpanSubmitter + DatasetRegistrar + DatasetFetcher + DatasetSummarizer + DatasetWriter, -> { - submitter: Arc, +pub struct Dataset { + api: Arc, + lifecycle: Arc, scope: CredentialScope, project_name: String, dataset_name: String, @@ -441,12 +387,11 @@ pub struct Dataset< lazy_metadata: Arc>, } -impl Clone - for Dataset -{ +impl Clone for Dataset { fn clone(&self) -> Self { Self { - submitter: self.submitter.clone(), + api: self.api.clone(), + lifecycle: self.lifecycle.clone(), scope: self.scope.clone(), project_name: self.project_name.clone(), dataset_name: self.dataset_name.clone(), @@ -457,9 +402,7 @@ impl - fmt::Debug for Dataset -{ +impl fmt::Debug for Dataset { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("Dataset") .field("project_name", &self.project_name) @@ -469,16 +412,7 @@ impl Dataset -{ +impl Dataset { /// Ensure the dataset is registered, returning cached metadata. async fn ensure_registered(&self) -> Result<&DatasetMetadata> { self.lazy_metadata @@ -490,10 +424,7 @@ impl< description: self.description.clone(), metadata: self.metadata.clone(), }; - let response = self - .submitter - .register_dataset(&self.scope, request) - .await?; + let response = self.api.register_dataset(&self.scope, request).await?; Ok(DatasetMetadata { dataset_id: response.dataset.id, project_id: response.project.id, @@ -521,7 +452,7 @@ impl< pub async fn insert(&self, record: DatasetInsert) -> Result { let dataset_id = self.ensure_registered().await?.dataset_id.clone(); let id = Uuid::new_v4().to_string(); - self.submitter + self.api .write_dataset_record( &self.scope, DatasetWrite::Insert { @@ -537,7 +468,7 @@ impl< /// Merge supplied fields into a record, returning after API acknowledgement. pub async fn update(&self, id: &str, record: DatasetInsert) -> Result<()> { let dataset_id = self.ensure_registered().await?.dataset_id.clone(); - self.submitter + self.api .write_dataset_record( &self.scope, DatasetWrite::Update { @@ -552,7 +483,7 @@ impl< /// Delete a record, returning after API acknowledgement. pub async fn delete(&self, id: &str) -> Result<()> { let dataset_id = self.ensure_registered().await?.dataset_id.clone(); - self.submitter + self.api .write_dataset_record( &self.scope, DatasetWrite::Delete { @@ -566,10 +497,10 @@ impl< /// Fetch records with pagination. /// /// Returns an iterator that fetches records in batches. - pub async fn fetch(&self, batch_size: Option) -> Result> { + pub async fn fetch(&self, batch_size: Option) -> Result { let metadata = self.ensure_registered().await?; Ok(DatasetIterator::new( - Arc::clone(&self.submitter), + self.api.clone(), self.scope.clone(), metadata.dataset_id.clone(), batch_size.unwrap_or(1000).max(1), @@ -582,7 +513,7 @@ impl< // Fetch summary from API let response = self - .submitter + .api .fetch_dataset_summary(&self.scope, &metadata.dataset_id) .await?; @@ -600,7 +531,7 @@ impl< /// Flush the client's shared trace queue. Dataset mutations are already /// acknowledged when their individual operation completes. pub fn flush(&self) -> crate::Completion { - self.submitter.flush() + self.lifecycle.flush() } } @@ -611,8 +542,8 @@ impl< /// Iterator for fetching dataset records with pagination. /// /// ```no_run -/// use braintrust_sdk_rust::{BraintrustClient, BraintrustError, Dataset}; -/// # async fn read(dataset: &Dataset) -> Result<(), BraintrustError> { +/// use braintrust_sdk_rust::{BraintrustError, Dataset}; +/// # async fn read(dataset: &Dataset) -> Result<(), BraintrustError> { /// let mut records = dataset.fetch(Some(100)).await?; /// while let Some(record) = records.next().await? { /// println!("{}", record.id); @@ -620,9 +551,8 @@ impl< /// # Ok(()) /// # } /// ``` -#[allow(private_bounds)] -pub struct DatasetIterator { - fetcher: Arc, +pub struct DatasetIterator { + fetcher: Arc, scope: CredentialScope, dataset_id: String, batch_size: usize, @@ -631,9 +561,13 @@ pub struct DatasetIterator { done: bool, } -#[allow(private_bounds)] -impl DatasetIterator { - fn new(fetcher: Arc, scope: CredentialScope, dataset_id: String, batch_size: usize) -> Self { +impl DatasetIterator { + fn new( + fetcher: Arc, + scope: CredentialScope, + dataset_id: String, + batch_size: usize, + ) -> Self { Self { fetcher, scope, diff --git a/src/eval/runner.rs b/src/eval/runner.rs index 96db594..052ed33 100644 --- a/src/eval/runner.rs +++ b/src/eval/runner.rs @@ -200,7 +200,7 @@ where case: Case, task: &dyn Task, scorers: &[Box>], - _root_span: &SpanHandle, + _root_span: &SpanHandle, client: &BraintrustClient, case_num: usize, quiet: bool, diff --git a/src/eval/types.rs b/src/eval/types.rs index ff0ed05..acbaf0e 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 crate::{BraintrustClient, SpanHandle}; +use crate::SpanHandle; /// A single test case for evaluation #[derive(Debug, Clone, Serialize, Deserialize)] @@ -203,7 +203,7 @@ impl EvalSummary { /// Hooks provided to tasks for instrumentation and metadata pub struct TaskHooks<'a, I, O> { /// Span handle for logging - pub span: &'a SpanHandle, + pub span: &'a SpanHandle, /// Expected output (if available) pub expected: Option<&'a O>, /// Mutable metadata map for the task to add context @@ -215,7 +215,7 @@ pub struct TaskHooks<'a, I, O> { } impl<'a, I, O> TaskHooks<'a, I, O> { - pub fn new(span: &'a SpanHandle, expected: Option<&'a O>) -> Self { + pub fn new(span: &'a SpanHandle, expected: Option<&'a O>) -> Self { Self { span, expected, diff --git a/src/experiments/api/comparison.rs b/src/experiments/api/comparison.rs deleted file mode 100644 index 173bffc..0000000 --- a/src/experiments/api/comparison.rs +++ /dev/null @@ -1 +0,0 @@ -pub(crate) use crate::api::experiments::ExperimentComparisonResponse; diff --git a/src/experiments/api/mod.rs b/src/experiments/api/mod.rs deleted file mode 100644 index 14d106c..0000000 --- a/src/experiments/api/mod.rs +++ /dev/null @@ -1,11 +0,0 @@ -//! Internal API contracts and traits. -//! -//! This module contains internal API types and traits used for communicating -//! with the Braintrust backend. These are not part of the public API. - -pub(crate) mod comparison; -pub(crate) mod registration; -pub(crate) mod traits; - -pub(crate) use registration::ExperimentRegisterRequest; -pub(crate) use traits::{BaseExperimentFetcher, ExperimentComparisonFetcher, ExperimentRegistrar}; diff --git a/src/experiments/api/registration.rs b/src/experiments/api/registration.rs deleted file mode 100644 index 4dac4cb..0000000 --- a/src/experiments/api/registration.rs +++ /dev/null @@ -1,2 +0,0 @@ -//! Compatibility imports for existing experiment callers. -pub(crate) use crate::api::experiments::{ExperimentRegisterRequest, ExperimentRegisterResponse}; diff --git a/src/experiments/api/traits.rs b/src/experiments/api/traits.rs deleted file mode 100644 index e14bd8c..0000000 --- a/src/experiments/api/traits.rs +++ /dev/null @@ -1,40 +0,0 @@ -//! Internal API trait definitions for experiments. - -use crate::error::Result; -use crate::experiments::api::comparison::ExperimentComparisonResponse; -use crate::experiments::api::registration::{ - ExperimentRegisterRequest, ExperimentRegisterResponse, -}; -use crate::experiments::metadata::BaseExperimentInfo; -use crate::internal::session::CredentialScope; - -/// Trait for registering experiments with the Braintrust API. -#[async_trait::async_trait] -pub(crate) trait ExperimentRegistrar: Send + Sync { - async fn register_experiment( - &self, - scope: &CredentialScope, - request: ExperimentRegisterRequest, - ) -> Result; -} - -/// Trait for fetching base experiment information. -#[async_trait::async_trait] -pub(crate) trait BaseExperimentFetcher: Send + Sync { - async fn fetch_base_experiment( - &self, - scope: &CredentialScope, - experiment_id: &str, - ) -> Result>; -} - -/// Trait for fetching experiment comparison data. -#[async_trait::async_trait] -pub(crate) trait ExperimentComparisonFetcher: Send + Sync { - async fn fetch_experiment_comparison( - &self, - scope: &CredentialScope, - experiment_id: &str, - base_experiment_id: Option<&str>, - ) -> Result; -} diff --git a/src/experiments/builder.rs b/src/experiments/builder.rs index 9ba3376..fe97b25 100644 --- a/src/experiments/builder.rs +++ b/src/experiments/builder.rs @@ -7,8 +7,8 @@ use serde_json::{Map, Value}; use crate::experiments::experiment::{Experiment, ExperimentState}; use crate::experiments::metadata::RepoInfo; +use crate::internal::ports::{ExperimentApi, SpanSubmitter}; use crate::internal::session::CredentialScope; -use crate::span::SpanSubmitter; /// Error type for ExperimentBuilder validation. #[derive(Debug, Clone, PartialEq)] @@ -34,9 +34,9 @@ impl std::error::Error for ExperimentBuilderError {} /// /// Registration is shared by clones of the resulting handle. Explicit /// credentials belong to the resource and do not mutate the client's session. -#[allow(private_bounds)] -pub struct ExperimentBuilder { - submitter: Arc, +pub struct ExperimentBuilder { + api: Arc, + submitter: Arc, scope: CredentialScope, project_name: Option, experiment_name: Option, @@ -48,10 +48,11 @@ pub struct ExperimentBuilder { repo_info: Option, } -impl Clone for ExperimentBuilder { +impl Clone for ExperimentBuilder { fn clone(&self) -> Self { Self { - submitter: Arc::clone(&self.submitter), + api: self.api.clone(), + submitter: self.submitter.clone(), scope: self.scope.clone(), project_name: self.project_name.clone(), experiment_name: self.experiment_name.clone(), @@ -65,11 +66,15 @@ impl Clone for ExperimentBuilder { } } -#[allow(private_bounds)] -impl ExperimentBuilder { +impl ExperimentBuilder { /// Create a new ExperimentBuilder (internal use). - pub(crate) fn new(submitter: Arc, scope: CredentialScope) -> Self { + pub(crate) fn new( + api: Arc, + submitter: Arc, + scope: CredentialScope, + ) -> Self { Self { + api, submitter, scope, project_name: None, @@ -152,12 +157,13 @@ impl ExperimentBuilder { } /// Build the experiment. Returns an error if required fields are missing. - pub fn build(self) -> std::result::Result, ExperimentBuilderError> { + pub fn build(self) -> std::result::Result { let project_name = self .project_name .ok_or(ExperimentBuilderError::MissingProjectName)?; Ok(Experiment { + api: self.api, submitter: self.submitter, scope: self.scope, project_name, diff --git a/src/experiments/experiment.rs b/src/experiments/experiment.rs index e8eb515..02c33cb 100644 --- a/src/experiments/experiment.rs +++ b/src/experiments/experiment.rs @@ -8,11 +8,8 @@ use std::time::{SystemTime, UNIX_EPOCH}; use serde_json::{Map, Value}; use tokio::sync::OnceCell; +use crate::api::experiments::ExperimentRegisterRequest; use crate::error::Result; -use crate::experiments::api::{ - BaseExperimentFetcher, ExperimentComparisonFetcher, ExperimentRegisterRequest, - ExperimentRegistrar, -}; use crate::experiments::event_log::ExperimentLog; use crate::experiments::feedback::Feedback; use crate::experiments::metadata::{ @@ -20,8 +17,9 @@ use crate::experiments::metadata::{ }; use crate::experiments::span_builder::ExperimentSpanBuilder; use crate::experiments::summary::{ExperimentSummary, MetricSummary, ScoreSummary}; +use crate::internal::ports::{ExperimentApi, SpanSubmitter}; use crate::internal::session::CredentialScope; -use crate::span::{SpanBuilder, SpanHandle, SpanLog, SpanSubmitter}; +use crate::span::{SpanBuilder, SpanHandle, SpanLog}; use crate::types::{ParentSpanInfo, SpanAttributes, SpanPayload, SpanType}; /// Metadata from experiment registration, cached after lazy initialization. @@ -35,9 +33,9 @@ pub(crate) struct ExperimentMetadata { /// Handle to a Braintrust experiment. /// /// Use this to log evaluation events, create spans, log feedback, and get summaries. -#[allow(private_bounds)] -pub struct Experiment { - pub(crate) submitter: Arc, +pub struct Experiment { + pub(crate) api: Arc, + pub(crate) submitter: Arc, pub(crate) scope: CredentialScope, pub(crate) project_name: String, pub(crate) experiment_name: Option, @@ -59,10 +57,11 @@ pub(crate) struct ExperimentState { called_start_span: AtomicBool, } -impl Clone for Experiment { +impl Clone for Experiment { fn clone(&self) -> Self { Self { - submitter: Arc::clone(&self.submitter), + api: self.api.clone(), + submitter: self.submitter.clone(), scope: self.scope.clone(), project_name: self.project_name.clone(), experiment_name: self.experiment_name.clone(), @@ -77,15 +76,7 @@ impl Clone for Experiment { } } -#[allow(private_bounds)] -impl< - S: SpanSubmitter - + ExperimentRegistrar - + BaseExperimentFetcher - + ExperimentComparisonFetcher - + 'static, - > Experiment -{ +impl Experiment { /// Log an event to the experiment, creating and ending a span internally. /// /// Returns the span ID which can be used with `log_feedback()`. @@ -142,7 +133,7 @@ impl< /// /// Use this for more control over span lifecycle. Once called, /// you cannot use `log()` on this experiment. - pub async fn start_span(&self) -> Result> { + pub async fn start_span(&self) -> Result { let metadata = self.ensure_registered().await?; self.state.called_start_span.store(true, Ordering::Release); Ok(ExperimentSpanBuilder::new(self, metadata.clone())) @@ -221,7 +212,7 @@ impl< // Fetch comparison data from API let comparison = self - .submitter + .api .fetch_experiment_comparison( &self.scope, &metadata.experiment_id, @@ -336,7 +327,7 @@ impl< /// ``` pub async fn traced(&self, callback: F) -> Result where - F: FnOnce(SpanHandle) -> Fut, + F: FnOnce(SpanHandle) -> Fut, Fut: std::future::Future, { let span = self.start_span().await?.name("root").build(); @@ -394,7 +385,7 @@ impl< /// Returns `None` if no base experiment has been configured or detected. pub async fn fetch_base_experiment(&self) -> Result> { let metadata = self.ensure_registered().await?; - self.submitter + self.api .fetch_base_experiment(&self.scope, &metadata.experiment_id) .await } @@ -427,10 +418,7 @@ impl< update: self.update, repo_info: self.repo_info.clone(), }; - let response = self - .submitter - .register_experiment(&self.scope, request) - .await?; + let response = self.api.register_experiment(&self.scope, request).await?; Ok(ExperimentMetadata { project_id: response.project.id, experiment_id: response.experiment.id, @@ -441,11 +429,7 @@ impl< } /// Create an internal span for log() calls. - fn create_internal_span( - &self, - metadata: &ExperimentMetadata, - start_time: f64, - ) -> SpanHandle { + fn create_internal_span(&self, metadata: &ExperimentMetadata, start_time: f64) -> SpanHandle { let parent_info = ParentSpanInfo::Experiment { object_id: metadata.experiment_id.clone(), }; diff --git a/src/experiments/mod.rs b/src/experiments/mod.rs index 870e762..d38f8c2 100644 --- a/src/experiments/mod.rs +++ b/src/experiments/mod.rs @@ -26,9 +26,6 @@ //! experiment.flush().await?; //! ``` -// Internal API module (not public) -pub(crate) mod api; - // Public modules mod builder; mod event_log; diff --git a/src/experiments/span_builder.rs b/src/experiments/span_builder.rs index 0efb335..1f04a6d 100644 --- a/src/experiments/span_builder.rs +++ b/src/experiments/span_builder.rs @@ -3,23 +3,24 @@ use std::sync::Arc; use crate::experiments::experiment::{Experiment, ExperimentMetadata}; -use crate::span::{SpanBuilder, SpanHandle, SpanSubmitter}; +use crate::internal::{ports::SpanSubmitter, session::CredentialScope}; +use crate::span::{SpanBuilder, SpanHandle}; use crate::types::{ParentSpanInfo, SpanType}; /// Builder for creating spans within an experiment. -#[allow(private_bounds)] -pub struct ExperimentSpanBuilder<'a, S: SpanSubmitter> { - experiment: &'a Experiment, +pub struct ExperimentSpanBuilder { + submitter: Arc, + scope: CredentialScope, metadata: ExperimentMetadata, name: Option, span_type: SpanType, } -#[allow(private_bounds)] -impl<'a, S: SpanSubmitter + 'static> ExperimentSpanBuilder<'a, S> { - pub(crate) fn new(experiment: &'a Experiment, metadata: ExperimentMetadata) -> Self { +impl ExperimentSpanBuilder { + pub(crate) fn new(experiment: &Experiment, metadata: ExperimentMetadata) -> Self { Self { - experiment, + submitter: experiment.submitter.clone(), + scope: experiment.scope.clone(), metadata, name: None, span_type: SpanType::Eval, @@ -39,19 +40,16 @@ impl<'a, S: SpanSubmitter + 'static> ExperimentSpanBuilder<'a, S> { } /// Build the span. - pub fn build(self) -> SpanHandle { + pub fn build(self) -> SpanHandle { let parent_info = ParentSpanInfo::Experiment { object_id: self.metadata.experiment_id.clone(), }; // Note: When using ParentSpanInfo::Experiment, we don't need project_name // because the destination is determined by the experiment_id. - let builder = SpanBuilder::new( - Arc::clone(&self.experiment.submitter), - self.experiment.scope.clone(), - ) - .parent_info(parent_info) - .span_type(self.span_type); + let builder = SpanBuilder::new(self.submitter, self.scope) + .parent_info(parent_info) + .span_type(self.span_type); builder.build() } diff --git a/src/internal/mod.rs b/src/internal/mod.rs index 7eb3eba..965aab2 100644 --- a/src/internal/mod.rs +++ b/src/internal/mod.rs @@ -2,6 +2,7 @@ pub(crate) mod batching; pub(crate) mod delivery; pub(crate) mod ingest; pub(crate) mod object_store; +pub(crate) mod ports; pub(crate) mod resources; pub(crate) mod retry; pub(crate) mod session; diff --git a/src/internal/ports.rs b/src/internal/ports.rs new file mode 100644 index 0000000..a5ba29b --- /dev/null +++ b/src/internal/ports.rs @@ -0,0 +1,73 @@ +//! Private dependency-injection boundaries for resource I/O and delivery. +use async_trait::async_trait; + +use crate::api::{ + btql::{BTQLQuery, BTQLResponse}, + datasets::{ + DatasetRegisterRequest, DatasetRegisterResponse, DatasetSummaryResponse, DatasetWrite, + }, + experiments::{ + ExperimentComparisonResponse, ExperimentRegisterRequest, ExperimentRegisterResponse, + }, +}; +use crate::error::Result; +use crate::experiments::BaseExperimentInfo; +use crate::internal::session::CredentialScope; +use crate::types::{ParentSpanInfo, SpanPayload}; + +/// Shared delivery lifecycle, independent of resource HTTP operations. +#[async_trait] +pub(crate) trait DeliveryControl: Send + Sync { + fn flush(&self) -> crate::Completion; + async fn trigger_flush(&self) -> Result<()>; +} + +#[async_trait] +pub(crate) trait SpanSubmitter: DeliveryControl { + fn submit( + &self, + scope: CredentialScope, + payload: SpanPayload, + parent_info: Option, + ) -> crate::Completion; +} + +#[async_trait] +pub(crate) trait DatasetApi: Send + Sync { + async fn register_dataset( + &self, + scope: &CredentialScope, + request: DatasetRegisterRequest, + ) -> Result; + async fn fetch_dataset_records( + &self, + scope: &CredentialScope, + query: BTQLQuery, + ) -> Result; + async fn fetch_dataset_summary( + &self, + scope: &CredentialScope, + dataset_id: &str, + ) -> Result; + async fn write_dataset_record(&self, scope: &CredentialScope, row: DatasetWrite) -> Result<()>; +} + +#[async_trait] +pub(crate) trait ExperimentApi: Send + Sync { + async fn register_experiment( + &self, + scope: &CredentialScope, + request: ExperimentRegisterRequest, + ) -> Result; + async fn fetch_base_experiment( + &self, + scope: &CredentialScope, + experiment_id: &str, + ) -> Result>; + async fn fetch_experiment_comparison( + &self, + scope: &CredentialScope, + experiment_id: &str, + base_experiment_id: Option<&str>, + ) -> Result; +} diff --git a/src/internal/resources.rs b/src/internal/resources.rs index 4d74ee2..557aec6 100644 --- a/src/internal/resources.rs +++ b/src/internal/resources.rs @@ -12,16 +12,26 @@ use crate::api::{ ExperimentRegisterResponse, }, }; -use crate::dataset::{DatasetFetcher, DatasetRegistrar, DatasetSummarizer, DatasetWriter}; use crate::error::{BraintrustError, Result}; -use crate::experiments::api::{ - BaseExperimentFetcher, ExperimentComparisonFetcher, ExperimentRegistrar, -}; use crate::experiments::BaseExperimentInfo; -use crate::BraintrustClient; +use crate::internal::ports::{DatasetApi, ExperimentApi}; + +/// Shared HTTP pool without retaining client/session/delivery state. +#[derive(Clone)] +pub(crate) struct ResourceTransport { + http: reqwest::Client, +} +impl ResourceTransport { + pub fn new(http: reqwest::Client) -> Self { + Self { http } + } + fn resource_api(&self, scope: &CredentialScope) -> Result { + scope.api(&self.http) + } +} #[async_trait] -impl ExperimentRegistrar for BraintrustClient { +impl ExperimentApi for ResourceTransport { async fn register_experiment( &self, scope: &CredentialScope, @@ -33,10 +43,7 @@ impl ExperimentRegistrar for BraintrustClient { .await .map_err(Into::into) } -} -#[async_trait] -impl BaseExperimentFetcher for BraintrustClient { async fn fetch_base_experiment( &self, scope: &CredentialScope, @@ -59,10 +66,7 @@ impl BaseExperimentFetcher for BraintrustClient { _ => None, }) } -} -#[async_trait] -impl ExperimentComparisonFetcher for BraintrustClient { async fn fetch_experiment_comparison( &self, scope: &CredentialScope, @@ -78,7 +82,7 @@ impl ExperimentComparisonFetcher for BraintrustClient { } #[async_trait] -impl DatasetRegistrar for BraintrustClient { +impl DatasetApi for ResourceTransport { async fn register_dataset( &self, scope: &CredentialScope, @@ -90,10 +94,7 @@ impl DatasetRegistrar for BraintrustClient { .await .map_err(Into::into) } -} -#[async_trait] -impl DatasetFetcher for BraintrustClient { async fn fetch_dataset_records( &self, scope: &CredentialScope, @@ -105,10 +106,7 @@ impl DatasetFetcher for BraintrustClient { .await .map_err(Into::into) } -} -#[async_trait] -impl DatasetSummarizer for BraintrustClient { async fn fetch_dataset_summary( &self, scope: &CredentialScope, @@ -120,10 +118,7 @@ impl DatasetSummarizer for BraintrustClient { .await .map_err(Into::into) } -} -#[async_trait] -impl DatasetWriter for BraintrustClient { async fn write_dataset_record(&self, scope: &CredentialScope, row: DatasetWrite) -> Result<()> { self.resource_api(scope)? .datasets() diff --git a/src/lib.rs b/src/lib.rs index a3b9857..15c7583 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,3 +1,5 @@ +#![deny(private_bounds, private_interfaces)] + pub mod api; mod completion; mod dataset; diff --git a/src/log_queue/queue.rs b/src/log_queue/queue.rs index 994cce0..3c1b6e2 100644 --- a/src/log_queue/queue.rs +++ b/src/log_queue/queue.rs @@ -7,6 +7,7 @@ use super::merge::merge_row_into; use super::row_key::RowKey; use crate::error::{BraintrustError, Result}; use crate::internal::delivery::WriteService; +use crate::internal::ports::{DeliveryControl, SpanSubmitter}; use crate::internal::session::CredentialScope; use crate::span_components::{project_logs_identifier, ProjectLogsIdentifier}; use crate::types::{LogDestination, Logs3Row, ParentSpanInfo, SpanObjectType, SpanPayload}; @@ -681,6 +682,27 @@ impl LogQueue { } } +#[async_trait::async_trait] +impl DeliveryControl for LogQueue { + fn flush(&self) -> Completion { + LogQueue::flush(self) + } + async fn trigger_flush(&self) -> Result<()> { + self.trigger_flush_command().await + } +} +#[async_trait::async_trait] +impl SpanSubmitter for LogQueue { + fn submit( + &self, + scope: CredentialScope, + payload: SpanPayload, + parent_info: Option, + ) -> Completion { + LogQueue::submit(self, scope, payload, parent_info) + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/src/logger.rs b/src/logger.rs index 7deaa7d..b16d847 100644 --- a/src/logger.rs +++ b/src/logger.rs @@ -11,10 +11,12 @@ use crate::dataset::DatasetBuilder; use crate::error::{BraintrustError, Result}; use crate::experiments::ExperimentBuilder; use crate::http::{build_http_client, CaBundleConfig}; +use crate::internal::ports::SpanSubmitter; +use crate::internal::resources::ResourceTransport; use crate::internal::retry::{exponential, RetryExecutor}; use crate::internal::session::CredentialScope; use crate::log_queue::{LogQueue, LogQueueConfig}; -use crate::span::{merge_span_origin_context, SpanOrigin, SpanOriginEnvironment, SpanSubmitter}; +use crate::span::{merge_span_origin_context, SpanOrigin, SpanOriginEnvironment}; use crate::span_components::SpanComponents; use crate::types::{ParentSpanInfo, SpanAttributes, SpanEventData, SpanObjectType, SpanPayload}; @@ -460,8 +462,12 @@ impl std::fmt::Debug for ClientInner { } impl BraintrustClient { - pub(crate) fn resource_api(&self, scope: &CredentialScope) -> Result { - scope.api(&self.inner.http_client) + fn resource_transport(&self) -> Arc { + Arc::new(ResourceTransport::new(self.inner.http_client.clone())) + } + + fn span_submitter(&self) -> Arc { + Arc::new(self.inner.queue.clone()) } fn explicit_scope( @@ -540,14 +546,14 @@ impl BraintrustClient { /// # Ok(()) /// # } /// ``` - pub async fn span_builder(&self) -> Result> { + pub async fn span_builder(&self) -> Result { if self.inner.login_skipped { return Err(BraintrustError::InvalidConfig( "span_builder() requires login; use span_builder_with_credentials() when skip_login is enabled".into(), )); } let state = self.wait_for_login_state().await?; - let mut builder = crate::span::SpanBuilder::new(Arc::new(self.clone()), state.scope()?); + let mut builder = crate::span::SpanBuilder::new(self.span_submitter(), state.scope()?); if let Some(environment) = &self.inner.environment { builder = builder.span_origin_environment(environment.clone()); } @@ -563,12 +569,16 @@ impl BraintrustClient { /// Create a dataset builder using the logged-in state. /// /// This waits for login to complete if it hasn't already. - pub async fn dataset_builder(&self) -> Result> { + pub async fn dataset_builder(&self) -> Result { if self.inner.login_skipped { return Err(BraintrustError::InvalidConfig("dataset_builder() requires login; use dataset_builder_with_credentials() when skip_login is enabled".into())); } let state = self.wait_for_login_state().await?; - Ok(DatasetBuilder::new(Arc::new(self.clone()), state.scope()?)) + Ok(DatasetBuilder::new( + self.resource_transport(), + Arc::new(self.inner.queue.clone()), + state.scope()?, + )) } /// Create a dataset builder with immutable explicit credentials, without @@ -577,9 +587,10 @@ impl BraintrustClient { &self, token: impl Into, org_name: impl Into, - ) -> DatasetBuilder { + ) -> DatasetBuilder { DatasetBuilder::new( - Arc::new(self.clone()), + self.resource_transport(), + Arc::new(self.inner.queue.clone()), self.explicit_scope(token.into(), String::new(), Some(org_name.into())), ) } @@ -594,11 +605,11 @@ impl BraintrustClient { &self, token: impl Into, org_id: impl Into, - ) -> crate::span::SpanBuilder { + ) -> crate::span::SpanBuilder { let token = token.into(); let org_id = org_id.into(); - let submitter = Arc::new(self.clone()); + let submitter = self.span_submitter(); let mut builder = crate::span::SpanBuilder::new(submitter, self.explicit_scope(token, org_id, None)); if let Some(environment) = &self.inner.environment { @@ -905,12 +916,16 @@ impl BraintrustClient { &self, token: impl Into, org_id: impl Into, - ) -> ExperimentBuilder { + ) -> ExperimentBuilder { let token = token.into(); let org_id = org_id.into(); - let submitter = Arc::new(self.clone()); - ExperimentBuilder::new(submitter, self.explicit_scope(token, org_id, None)) + let submitter = self.span_submitter(); + ExperimentBuilder::new( + self.resource_transport(), + submitter, + self.explicit_scope(token, org_id, None), + ) } /// Submit immediately; the token observes admission and full batch delivery. @@ -945,26 +960,6 @@ impl BraintrustClient { } } -#[async_trait::async_trait] -impl SpanSubmitter for BraintrustClient { - fn submit( - &self, - scope: CredentialScope, - payload: SpanPayload, - parent_info: Option, - ) -> crate::Completion { - self.submit_payload(scope, payload, parent_info) - } - - fn flush(&self) -> crate::Completion { - BraintrustClient::flush(self) - } - - async fn trigger_flush(&self) -> Result<()> { - BraintrustClient::trigger_flush(self).await - } -} - #[allow(dead_code)] #[derive(Debug, Default)] struct MutableSpanEvent { @@ -1151,6 +1146,32 @@ mod tests { ResponseTemplate::new(200).set_body_json(serde_json::json!({ "org_info": org_info })) } + #[tokio::test] + async fn resource_ports_do_not_retain_the_whole_client_or_login_state() { + let client = BraintrustClient::builder() + .skip_login(true) + .build() + .await + .unwrap(); + let weak = Arc::downgrade(&client.inner); + let span = client.span_builder_with_credentials("token", "org"); + let dataset = client + .dataset_builder_with_credentials("token", "tenant") + .project_name("p") + .dataset_name("d"); + let experiment = client + .experiment_builder_with_credentials("token", "org") + .project_name("p"); + drop(client); + assert!( + weak.upgrade().is_none(), + "resource dependencies must not retain ClientInner" + ); + let _span = span.build(); + let _dataset = dataset.build().unwrap(); + let _experiment = experiment.build().unwrap(); + } + #[tokio::test] async fn builder_rejects_missing_api_key() { // Clear any env vars that might be set diff --git a/src/span.rs b/src/span.rs index f0aed07..05b71a2 100644 --- a/src/span.rs +++ b/src/span.rs @@ -5,11 +5,13 @@ use std::path::PathBuf; use std::sync::{Arc, Mutex}; use std::time::{SystemTime, UNIX_EPOCH}; +#[cfg(test)] use async_trait::async_trait; use serde_json::{json, Map, Value}; use uuid::Uuid; use crate::error::Result; +use crate::internal::ports::SpanSubmitter; use crate::internal::session::CredentialScope; use crate::span_components::SpanComponents; use crate::types::{ @@ -248,27 +250,8 @@ impl SpanLogBuilder { } } -#[async_trait] -pub(crate) trait SpanSubmitter: Send + Sync { - /// Attempt admission immediately and observe the logical batch outcome. - fn submit( - &self, - scope: CredentialScope, - payload: SpanPayload, - parent_info: Option, - ) -> crate::Completion; - - /// Flush all queued span data through to the server. - fn flush(&self) -> crate::Completion; - - /// Trigger a non-blocking background flush. - async fn trigger_flush(&self) -> Result<()>; -} - -#[derive(Debug)] -#[allow(private_bounds)] -pub struct SpanBuilder { - submitter: Arc, +pub struct SpanBuilder { + submitter: Arc, scope: CredentialScope, project_name: Option, parent_info: Option, @@ -282,7 +265,17 @@ pub struct SpanBuilder { span_attributes_extra: HashMap, } -impl Clone for SpanBuilder { +impl fmt::Debug for SpanBuilder { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("SpanBuilder") + .field("scope", &self.scope) + .field("project_name", &self.project_name) + .field("parent_info", &self.parent_info) + .finish_non_exhaustive() + } +} + +impl Clone for SpanBuilder { fn clone(&self) -> Self { Self { submitter: Arc::clone(&self.submitter), @@ -369,9 +362,8 @@ impl SpanOrigin { } } -#[allow(private_bounds)] -impl SpanBuilder { - pub(crate) fn new(submitter: Arc, scope: CredentialScope) -> Self { +impl SpanBuilder { + pub(crate) fn new(submitter: Arc, scope: CredentialScope) -> Self { Self { submitter, scope, @@ -479,7 +471,7 @@ impl SpanBuilder { self } - pub fn build(self) -> SpanHandle { + pub fn build(self) -> SpanHandle { use std::time::{SystemTime, UNIX_EPOCH}; // Generate both IDs ONCE at span creation - reused for all flushes. @@ -535,9 +527,8 @@ impl SpanBuilder { } } -#[allow(private_bounds)] -pub struct SpanHandle { - submitter: Arc, +pub struct SpanHandle { + submitter: Arc, scope: CredentialScope, parent_info: Option, /// Stable identifier set at creation, accessible without locking. @@ -545,7 +536,7 @@ pub struct SpanHandle { inner: Arc>, } -impl Clone for SpanHandle { +impl Clone for SpanHandle { fn clone(&self) -> Self { Self { submitter: Arc::clone(&self.submitter), @@ -557,8 +548,7 @@ impl Clone for SpanHandle { } } -#[allow(private_bounds)] -impl SpanHandle { +impl SpanHandle { /// Return the stable row ID for this span, usable with `log_feedback()` and `update_span()`. pub fn row_id(&self) -> &str { &self.row_id @@ -1085,6 +1075,9 @@ mod tests { } completion } + } + #[async_trait] + impl crate::internal::ports::DeliveryControl for Submitter { fn flush(&self) -> crate::Completion { let (sender, completion) = crate::Completion::pair(); let _ = sender.send(Ok(crate::FlushReport::default())); diff --git a/src/stream.rs b/src/stream.rs index 248050e..e4f061c 100644 --- a/src/stream.rs +++ b/src/stream.rs @@ -66,7 +66,7 @@ use serde::{Deserialize, Serialize}; use serde_json::Value; use tokio::sync::Mutex; -use crate::span::{SpanHandle, SpanLog, SpanSubmitter}; +use crate::span::{SpanHandle, SpanLog}; use crate::types::{usage_metrics_to_map, UsageMetrics}; // ============================================================================= @@ -1021,16 +1021,13 @@ impl BraintrustStream { /// # Type Parameters /// - `S`: The stream type yielding `Result` /// - `E`: The error type (allows use with any error type) -/// - `Sub`: The span submitter type -#[allow(private_bounds)] -pub fn wrap_stream_with_span( +pub fn wrap_stream_with_span( stream: S, - span: SpanHandle, + span: SpanHandle, ) -> Pin> + Send>> where S: Stream> + Send + Unpin + 'static, E: Send + 'static, - Sub: SpanSubmitter + 'static, { use futures::StreamExt; @@ -1150,18 +1147,15 @@ enum FinalizeState { } /// A wrapper stream that logs aggregated output when the stream is exhausted. -struct SpanCompleteWrapper { +struct SpanCompleteWrapper { inner: S, - span: Option>, + span: Option, aggregator: Option>>, finalize_state: FinalizeState, } /// Finalize the stream by logging and flushing the span. -async fn finalize_span( - span: SpanHandle, - aggregator: Arc>, -) { +async fn finalize_span(span: SpanHandle, aggregator: Arc>) { let mut agg = aggregator.lock().await; if !agg.is_empty() { match agg.final_value() { @@ -1204,11 +1198,10 @@ async fn finalize_span( } } -impl Stream for SpanCompleteWrapper +impl Stream for SpanCompleteWrapper where S: Stream> + Unpin, E: Send + 'static, - Sub: SpanSubmitter + 'static, { type Item = std::result::Result; diff --git a/src/test_utils.rs b/src/test_utils.rs index c2ab142..fb12081 100644 --- a/src/test_utils.rs +++ b/src/test_utils.rs @@ -7,7 +7,7 @@ use crate::internal::session::CredentialScope; use crate::span::{SpanBuilder, SpanHandle}; use crate::types::{ParentSpanInfo, SpanPayload}; -use crate::span::SpanSubmitter; +use crate::internal::ports::{DeliveryControl, SpanSubmitter}; #[derive(Clone, Default)] pub(crate) struct TestSpanCollector { @@ -74,7 +74,10 @@ impl SpanSubmitter for MockBraintrustClient { let _ = sender.send(Ok(crate::DeliveryReceipt)); completion } +} +#[async_trait] +impl DeliveryControl for MockBraintrustClient { fn flush(&self) -> crate::Completion { let (sender, completion) = crate::Completion::pair(); let _ = sender.send(Ok(crate::FlushReport::default())); @@ -87,13 +90,13 @@ impl SpanSubmitter for MockBraintrustClient { } } -pub(crate) fn mock_span_builder() -> (SpanBuilder, TestSpanCollector) { +pub(crate) fn mock_span_builder() -> (SpanBuilder, TestSpanCollector) { let (submitter, collector) = MockBraintrustClient::new(); let builder = SpanBuilder::new(submitter, CredentialScope::test("test-token", "org-id")); (builder, collector) } -pub(crate) fn build_test_span() -> (SpanHandle, TestSpanCollector) { +pub(crate) fn build_test_span() -> (SpanHandle, TestSpanCollector) { let (builder, collector) = mock_span_builder(); (builder.build(), collector) } diff --git a/tests/dataset_test.rs b/tests/dataset_test.rs index f982236..6de0cd4 100644 --- a/tests/dataset_test.rs +++ b/tests/dataset_test.rs @@ -298,7 +298,7 @@ async fn isolated_client(server: &MockServer) -> BraintrustClient { fn registration() -> ResponseTemplate { ResponseTemplate::new(200).set_body_json(json!({"project":{"id":"project","name":"project"},"dataset":{"id":"dataset","name":"dataset"}})) } -fn dataset(client: &BraintrustClient) -> braintrust_sdk_rust::Dataset { +fn dataset(client: &BraintrustClient) -> braintrust_sdk_rust::Dataset { client .dataset_builder_with_credentials("token", "tenant") .project_name("project") diff --git a/tests/resource_handles.rs b/tests/resource_handles.rs new file mode 100644 index 0000000..7e6812c --- /dev/null +++ b/tests/resource_handles.rs @@ -0,0 +1,130 @@ +//! Consumer coverage for the concrete public resource model. +use braintrust_sdk_rust::{ + BraintrustClient, Dataset, DatasetBuilder, DatasetInsert, DatasetIterator, Experiment, + ExperimentBuilder, ExperimentSpanBuilder, ParentSpanInfo, SpanBuilder, SpanHandle, SpanLog, +}; +use serde_json::json; +use wiremock::{ + matchers::{method, path}, + Mock, MockServer, ResponseTemplate, +}; + +// Applications can store these types without naming an SDK implementation or +// carrying its private bounds through their own structs and functions. +struct Resources { + span: SpanHandle, + experiment: Experiment, + dataset: Dataset, +} + +#[test] +fn resource_handles_and_builders_are_send_sync_and_owned() { + fn assert_owned() {} + assert_owned::(); + assert_owned::(); + assert_owned::(); + assert_owned::(); + assert_owned::(); + assert_owned::(); + assert_owned::(); + assert_owned::(); +} + +#[tokio::test] +async fn resources_and_detached_span_builders_work_after_the_client_is_dropped() { + let server = MockServer::start().await; + Mock::given(method("POST")).and(path("/api/experiment/register")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"project":{"id":"project","name":"p"},"experiment":{"id":"experiment","name":"e"}}))) + .expect(1).mount(&server).await; + Mock::given(method("POST")) + .and(path("/api/dataset/register")) + .respond_with(ResponseTemplate::new(200).set_body_json( + json!({"project":{"id":"project","name":"p"},"dataset":{"id":"dataset","name":"d"}}), + )) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(204)) + .expect(3) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/btql")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"data":[{"id":"record"}]}))) + .expect(1) + .mount(&server) + .await; + let client = BraintrustClient::builder() + .skip_login(true) + .app_url(server.uri()) + .api_url(server.uri()) + .build() + .await + .unwrap(); + let span: SpanBuilder = client + .span_builder_with_credentials("token", "org") + .parent_info(ParentSpanInfo::ProjectLogs { + object_id: "project".into(), + }); + let experiment: ExperimentBuilder = client + .experiment_builder_with_credentials("token", "org") + .project_name("p"); + let dataset: DatasetBuilder = client + .dataset_builder_with_credentials("token", "tenant") + .project_name("p") + .dataset_name("d"); + drop(client); + let resources = Resources { + span: span.build(), + experiment: experiment.build().unwrap(), + dataset: dataset.build().unwrap(), + }; + resources.span.log(SpanLog::default()).await.unwrap(); + let child: ExperimentSpanBuilder = resources.experiment.start_span().await.unwrap(); + drop(resources.experiment); + let child: SpanHandle = child.build(); + child.log(SpanLog::default()).await.unwrap(); + resources + .dataset + .insert(DatasetInsert::builder().input(json!(1)).build().unwrap()) + .await + .unwrap(); + let records: DatasetIterator = resources.dataset.fetch(None).await.unwrap(); + drop(resources.dataset); + assert_eq!(records.collect().await.unwrap()[0].id, "record"); + resources.span.flush().await.unwrap(); +} + +#[tokio::test] +async fn dropping_the_last_resource_and_client_still_drains_an_accepted_write() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(204)) + .expect(1) + .mount(&server) + .await; + let client = BraintrustClient::builder() + .skip_login(true) + .app_url(server.uri()) + .api_url(server.uri()) + .build() + .await + .unwrap(); + let span: SpanHandle = client + .span_builder_with_credentials("token", "org") + .parent_info(ParentSpanInfo::ProjectLogs { + object_id: "project".into(), + }) + .build(); + let write = span.log(SpanLog::default()); + assert!(write.is_ok()); + drop(span); + drop(client); + tokio::time::timeout(std::time::Duration::from_secs(5), async { write.await }) + .await + .unwrap() + .unwrap(); +}