1- use baml_runtime:: client_registry:: ClientRegistry ;
1+ use baml_runtime:: client_registry:: { ClientProperty , ClientProvider , ClientRegistry } ;
22use baml_runtime:: tracingv2:: storage:: storage:: Collector ;
33use baml_runtime:: type_builder:: TypeBuilder ;
44use baml_runtime:: { BamlRuntime , FunctionResult , RuntimeContextManager , TripWire } ;
@@ -12,6 +12,7 @@ use rustler::{
1212} ;
1313use std:: collections:: HashMap ;
1414use std:: path:: Path ;
15+ use std:: str:: FromStr ;
1516use std:: sync:: Arc ;
1617mod atoms {
1718 rustler:: atoms! {
@@ -83,6 +84,73 @@ fn term_to_baml_value<'a>(term: Term<'a>) -> Result<BamlValue, Error> {
8384 ) ) ) )
8485}
8586
87+ fn term_to_optional_string ( term : Term ) -> Result < Option < String > , Error > {
88+ if term. is_atom ( ) && term. decode :: < rustler:: Atom > ( ) ? == atom:: nil ( ) {
89+ Ok ( None )
90+ } else {
91+ Ok ( Some ( term_to_string ( term) ?) )
92+ }
93+ }
94+
95+ fn term_to_baml_map ( term : Term ) -> Result < BamlMap < String , BamlValue > , Error > {
96+ if term. is_atom ( ) && term. decode :: < rustler:: Atom > ( ) ? == atom:: nil ( ) {
97+ return Ok ( BamlMap :: new ( ) ) ;
98+ }
99+ if !term. is_map ( ) {
100+ return Err ( Error :: Term ( Box :: new ( "Expected a map" ) ) ) ;
101+ }
102+ let mut map = BamlMap :: new ( ) ;
103+ for ( key_term, value_term) in
104+ MapIterator :: new ( term) . ok_or ( Error :: Term ( Box :: new ( "Invalid map" ) ) ) ?
105+ {
106+ let key = term_to_string ( key_term) ?;
107+ let value = term_to_baml_value ( value_term) ?;
108+ map. insert ( key, value) ;
109+ }
110+ Ok ( map)
111+ }
112+
113+ fn term_to_client_property ( term : Term , name_override : Option < String > ) -> Result < ClientProperty , Error > {
114+ if !term. is_map ( ) {
115+ return Err ( Error :: Term ( Box :: new ( "Client must be a map" ) ) ) ;
116+ }
117+
118+ let mut name: Option < String > = name_override;
119+ let mut provider: Option < ClientProvider > = None ;
120+ let mut retry_policy: Option < String > = None ;
121+ let mut options: BamlMap < String , BamlValue > = BamlMap :: new ( ) ;
122+
123+ let iter = MapIterator :: new ( term) . ok_or ( Error :: Term ( Box :: new ( "Invalid client map" ) ) ) ?;
124+ for ( key_term, value_term) in iter {
125+ let key = term_to_string ( key_term) ?;
126+ match key. as_str ( ) {
127+ "name" => {
128+ name = Some ( term_to_string ( value_term) ?) ;
129+ }
130+ "provider" => {
131+ let provider_str = term_to_string ( value_term) ?;
132+ provider = Some (
133+ ClientProvider :: from_str ( & provider_str) . map_err ( |e| {
134+ Error :: Term ( Box :: new ( format ! ( "Invalid client provider: {e}" ) ) )
135+ } ) ?,
136+ ) ;
137+ }
138+ "retry_policy" => {
139+ retry_policy = term_to_optional_string ( value_term) ?;
140+ }
141+ "options" => {
142+ options = term_to_baml_map ( value_term) ?;
143+ }
144+ _ => { }
145+ }
146+ }
147+
148+ let name = name. ok_or ( Error :: Term ( Box :: new ( "Client missing required key: name" ) ) ) ?;
149+ let provider = provider. ok_or ( Error :: Term ( Box :: new ( "Client missing required key: provider" ) ) ) ?;
150+
151+ Ok ( ClientProperty :: new ( name, provider, retry_policy, options) )
152+ }
153+
86154fn baml_value_to_term < ' a > ( env : Env < ' a > , value : & BamlValue ) -> NifResult < Term < ' a > > {
87155 match value {
88156 BamlValue :: String ( s) => Ok ( s. encode ( env) ) ,
@@ -218,6 +286,32 @@ fn prepare_request<'a>(
218286 if key == "primary" {
219287 let primary = term_to_string ( value_term) ?;
220288 registry. set_primary ( primary) ;
289+ } else if key == "clients" {
290+ // Accept either:
291+ // - a list of client maps: [%{name: ..., provider: ..., ...}, ...]
292+ // - a map of name => client map: %{ "name" => %{provider: ..., ...}, ... }
293+ if let Ok ( list) = value_term. decode :: < Vec < Term > > ( ) {
294+ for client_term in list {
295+ let client = term_to_client_property ( client_term, None ) ?;
296+ registry. add_client ( client) ;
297+ }
298+ } else if value_term. is_map ( ) {
299+ let client_iter = MapIterator :: new ( value_term)
300+ . ok_or ( Error :: Term ( Box :: new ( "Invalid clients map" ) ) ) ?;
301+ for ( name_term, client_term) in client_iter {
302+ let name = term_to_string ( name_term) ?;
303+ let client = term_to_client_property ( client_term, Some ( name) ) ?;
304+ registry. add_client ( client) ;
305+ }
306+ } else if value_term. is_atom ( )
307+ && value_term. decode :: < rustler:: Atom > ( ) ? == atom:: nil ( )
308+ {
309+ // allow nil clients
310+ } else {
311+ return Err ( Error :: Term ( Box :: new (
312+ "Client registry clients must be a list, a map, or nil" ,
313+ ) ) ) ;
314+ }
221315 }
222316 }
223317 Some ( registry)
0 commit comments