use crate::baml_client::{
runtime::{FunctionOptions, get_runtime},
stream_types, types,
};
use baml::{BamlEncode, BamlError, StreamingCall};
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))
}
}
};
}
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 }
}
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)
}
}
impl_options_convenience_methods!($name);
};
}
baml_function_sync!(Chat(system: impl AsRef<str> + BamlEncode, messages: &[types::ChatMessage], ) -> (String, String));
#[derive(Clone)]
pub struct BamlSyncClient {
options: FunctionOptions,
pub Chat: Chat,
}
impl BamlSyncClient {
pub const fn new() -> Self {
Self {
options: FunctionOptions::new(),
Chat: Chat::new(),
}
}
pub fn with_options(&self, options: FunctionOptions) -> Self {
Self {
options: options.clone(),
Chat: Chat {
options: options.clone(),
},
}
}
}
impl_options_convenience_methods!(BamlSyncClient);
pub static B: BamlSyncClient = BamlSyncClient::new();