episteme-local 0.1.1

// This file was generated by BAML: do not edit it.
// Instead, edit the BAML source files.
//
// Learn more at https://docs.boundaryml.com

//! Synchronous BAML client with function-object pattern.

use crate::baml_client::{
    runtime::{FunctionOptions, get_runtime},
    stream_types, types,
};
use baml::{BamlEncode, BamlError, StreamingCall};

// =============================================================================
// Convenience Builder Methods Macro
// =============================================================================

/// Implements convenience builder methods that delegate to `with_options`.
/// Requires the type to have:
/// - `options: FunctionOptions` field
/// - `fn with_options(&self, options: FunctionOptions) -> Self` method
macro_rules! impl_options_convenience_methods {
    ($name:ident) => {
        impl $name {
            pub fn with_collector(&self, collector: &baml::Collector) -> Self {
                self.with_options(self.options.clone().with_collector(collector))
            }

            pub fn with_collectors(&self, collectors: &[baml::Collector]) -> Self {
                self.with_options(self.options.clone().with_collectors(collectors))
            }

            pub fn with_cancellation_token(&self, token: Option<baml::CancellationToken>) -> Self {
                self.with_options(self.options.clone().with_cancellation_token(token))
            }

            pub fn with_type_builder(&self, tb: &super::super::type_builder::TypeBuilder) -> Self {
                self.with_options(self.options.clone().with_type_builder(tb))
            }

            pub fn with_env_var(&self, key: impl Into<String>, value: impl Into<String>) -> Self {
                self.with_options(self.options.clone().with_env_var(key, value))
            }

            pub fn with_tag(&self, key: impl Into<String>, value: impl Into<String>) -> Self {
                self.with_options(self.options.clone().with_tag(key, value))
            }

            pub fn with_client(&self, client_name: impl Into<String>) -> Self {
                self.with_options(self.options.clone().with_client(client_name))
            }

            pub fn with_client_registry(&self, registry: &baml::ClientRegistry) -> Self {
                self.with_options(self.options.clone().with_client_registry(registry))
            }

            pub fn with_on_tick<F>(&self, callback: F) -> Self
            where
                F: Fn(&baml::FunctionLog) + Send + Sync + 'static,
            {
                self.with_options(self.options.clone().with_on_tick(callback))
            }
        }
    };
}

// =============================================================================
// Function Struct Macro (generates per-function struct)
// =============================================================================

macro_rules! baml_function_sync {
    ($name:ident($($param:ident : $ptype:ty),* $(,)?) -> ($stream_ret:ty, $final_ret:ty)) => {
        #[derive(Clone)]
        pub struct $name {
            pub(crate) options: FunctionOptions,
        }

        impl $name {
            pub const fn new() -> Self {
                Self { options: FunctionOptions::new() }
            }

            fn with_options(&self, options: FunctionOptions) -> Self {
                Self { options }
            }

            // Terminal methods
            pub fn call(&self, $($param: $ptype),*) -> Result<$final_ret, BamlError> {
                let args = self.options.to_baml_args()
                    $(.arg(stringify!($param), $param.baml_encode()))*;
                get_runtime().call_function(stringify!($name), &args)
            }

            pub fn stream(&self, $($param: $ptype),*) -> Result<StreamingCall<$stream_ret, $final_ret>, BamlError> {
                let args = self.options.to_baml_args()
                    $(.arg(stringify!($param), $param.baml_encode()))*;
                get_runtime().call_function_stream(stringify!($name), &args)
            }

            pub fn parse(&self, response: &str) -> Result<$final_ret, BamlError> {
                get_runtime().parse(stringify!($name), response, false)
            }

            pub fn parse_stream(&self, response: &str) -> Result<$stream_ret, BamlError> {
                get_runtime().parse(stringify!($name), response, true)
            }

            pub fn build_request(&self, $($param: $ptype),*) -> Result<baml::HTTPRequest, BamlError> {
                let args = self.options.to_baml_args()
                    $(.arg(stringify!($param), $param.baml_encode()))*
                    .arg("stream", false);
                get_runtime().build_request(stringify!($name), &args)
            }

            pub fn build_request_stream(&self, $($param: $ptype),*) -> Result<baml::HTTPRequest, BamlError> {
                let args = self.options.to_baml_args()
                    $(.arg(stringify!($param), $param.baml_encode()))*
                    .arg("stream", true);
                get_runtime().build_request_stream(stringify!($name), &args)
            }
        }

        // Builder convenience methods (separate impl block via macro)
        impl_options_convenience_methods!($name);
    };
}

// =============================================================================
// Generate function structs
// =============================================================================

baml_function_sync!(AggregateResearchChunks(classification: &types::DocumentClassification, chunks: &[types::ResearchChunkOutput], ) -> (stream_types::AggregatedResearchOutput, types::AggregatedResearchOutput));

baml_function_sync!(ClassifyDocument(document: impl AsRef<str> + BamlEncode, ) -> (stream_types::DocumentClassification, types::DocumentClassification));

baml_function_sync!(DistillResearch(document: impl AsRef<str> + BamlEncode, classification: &types::DocumentClassification, ) -> (stream_types::ResearchDraftOutput, types::ResearchDraftOutput));

baml_function_sync!(DistillResearchChunk(spans: &[types::SourceSpanInput], classification: &types::DocumentClassification, ) -> (stream_types::ResearchChunkOutput, types::ResearchChunkOutput));

// =============================================================================
// Client Struct
// =============================================================================

#[derive(Clone)]
pub struct BamlSyncClient {
    options: FunctionOptions,

    pub AggregateResearchChunks: AggregateResearchChunks,

    pub ClassifyDocument: ClassifyDocument,

    pub DistillResearch: DistillResearch,

    pub DistillResearchChunk: DistillResearchChunk,
}

impl BamlSyncClient {
    pub const fn new() -> Self {
        Self {
            options: FunctionOptions::new(),

            AggregateResearchChunks: AggregateResearchChunks::new(),

            ClassifyDocument: ClassifyDocument::new(),

            DistillResearch: DistillResearch::new(),

            DistillResearchChunk: DistillResearchChunk::new(),
        }
    }

    /// Apply options to all functions. Returns a new client with the options set.
    pub fn with_options(&self, options: FunctionOptions) -> Self {
        Self {
            options: options.clone(),

            AggregateResearchChunks: AggregateResearchChunks {
                options: options.clone(),
            },

            ClassifyDocument: ClassifyDocument {
                options: options.clone(),
            },

            DistillResearch: DistillResearch {
                options: options.clone(),
            },

            DistillResearchChunk: DistillResearchChunk {
                options: options.clone(),
            },
        }
    }
}

// Builder convenience methods
impl_options_convenience_methods!(BamlSyncClient);

pub static B: BamlSyncClient = BamlSyncClient::new();