Skip to content

Commit 0f2c730

Browse files
aaronvgaaronvg
authored
Add a way to create baml clients with a default set of options (#1595)
1 parent 3500413 commit 0f2c730

23 files changed

Lines changed: 11374 additions & 5946 deletions

File tree

‎engine/language_client_codegen/src/python/templates/async_client.py.j2‎

Lines changed: 38 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -23,11 +23,35 @@ class BamlAsyncClient:
2323
__runtime: baml_py.BamlRuntime
2424
__ctx_manager: baml_py.BamlCtxManager
2525
__stream_client: "BamlStreamClient"
26+
__baml_options: BamlCallOptions
2627

27-
def __init__(self, runtime: baml_py.BamlRuntime, ctx_manager: baml_py.BamlCtxManager):
28+
def __init__(self, runtime: baml_py.BamlRuntime, ctx_manager: baml_py.BamlCtxManager, baml_options: Optional[BamlCallOptions] = None):
2829
self.__runtime = runtime
2930
self.__ctx_manager = ctx_manager
30-
self.__stream_client = BamlStreamClient(self.__runtime, self.__ctx_manager)
31+
self.__stream_client = BamlStreamClient(self.__runtime, self.__ctx_manager, baml_options)
32+
self.__baml_options = baml_options or {}
33+
34+
def with_options(
35+
self,
36+
tb: Optional[TypeBuilder] = None,
37+
client_registry: Optional[baml_py.baml_py.ClientRegistry] = None,
38+
collector: Optional[Union[baml_py.baml_py.Collector, List[baml_py.baml_py.Collector]]] = None,
39+
) -> "BamlAsyncClient":
40+
"""
41+
Returns a new instance of BamlAsyncClient with explicitly typed baml options
42+
for Python 3.8 compatibility.
43+
"""
44+
new_options = {}
45+
46+
# Override if any keyword arguments were provided.
47+
if tb is not None:
48+
new_options["tb"] = tb
49+
if client_registry is not None:
50+
new_options["client_registry"] = client_registry
51+
if collector is not None:
52+
new_options["collector"] = collector
53+
54+
return BamlAsyncClient(self.__runtime, self.__ctx_manager, new_options)
3155

3256
@property
3357
def stream(self):
@@ -42,13 +66,15 @@ class BamlAsyncClient:
4266
{%- endfor %}
4367
baml_options: BamlCallOptions = {},
4468
) -> {{fn.return_type}}:
45-
__tb__ = baml_options.get("tb", None)
69+
options = {**self.__baml_options, **(baml_options or {})}
70+
71+
__tb__ = options.get("tb", None)
4672
if __tb__ is not None:
4773
tb = __tb__._tb # type: ignore (we know how to use this private attribute)
4874
else:
4975
tb = None
50-
__cr__ = baml_options.get("client_registry", None)
51-
collector = baml_options.get("collector", None)
76+
__cr__ = options.get("client_registry", None)
77+
collector = options.get("collector", None)
5278
collectors = collector if isinstance(collector, list) else [collector] if collector is not None else []
5379
raw = await self.__runtime.call_function(
5480
"{{fn.name}}",
@@ -69,10 +95,11 @@ class BamlAsyncClient:
6995
class BamlStreamClient:
7096
__runtime: baml_py.BamlRuntime
7197
__ctx_manager: baml_py.BamlCtxManager
72-
73-
def __init__(self, runtime: baml_py.BamlRuntime, ctx_manager: baml_py.BamlCtxManager):
98+
__baml_options: BamlCallOptions
99+
def __init__(self, runtime: baml_py.BamlRuntime, ctx_manager: baml_py.BamlCtxManager, baml_options: Optional[BamlCallOptions] = None):
74100
self.__runtime = runtime
75101
self.__ctx_manager = ctx_manager
102+
self.__baml_options = baml_options or {}
76103

77104
{% for fn in funcs %}
78105
def {{ fn.name }}(
@@ -82,13 +109,14 @@ class BamlStreamClient:
82109
{%- endfor %}
83110
baml_options: BamlCallOptions = {},
84111
) -> baml_py.BamlStream[{{ fn.partial_return_type }}, {{ fn.return_type }}]:
85-
__tb__ = baml_options.get("tb", None)
112+
options = {**self.__baml_options, **(baml_options or {})}
113+
__tb__ = options.get("tb", None)
86114
if __tb__ is not None:
87115
tb = __tb__._tb # type: ignore (we know how to use this private attribute)
88116
else:
89117
tb = None
90-
__cr__ = baml_options.get("client_registry", None)
91-
collector = baml_options.get("collector", None)
118+
__cr__ = options.get("client_registry", None)
119+
collector = options.get("collector", None)
92120
collectors = collector if isinstance(collector, list) else [collector] if collector is not None else []
93121
raw = self.__runtime.stream_function(
94122
"{{fn.name}}",

‎engine/language_client_codegen/src/python/templates/sync_client.py.j2‎

Lines changed: 36 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -22,16 +22,38 @@ class BamlSyncClient:
2222
__runtime: baml_py.BamlRuntime
2323
__ctx_manager: baml_py.BamlCtxManager
2424
__stream_client: "BamlStreamClient"
25+
__baml_options: BamlCallOptions
2526

26-
def __init__(self, runtime: baml_py.BamlRuntime, ctx_manager: baml_py.BamlCtxManager):
27+
def __init__(self, runtime: baml_py.BamlRuntime, ctx_manager: baml_py.BamlCtxManager, baml_options: Optional[BamlCallOptions] = None):
2728
self.__runtime = runtime
2829
self.__ctx_manager = ctx_manager
29-
self.__stream_client = BamlStreamClient(self.__runtime, self.__ctx_manager)
30-
30+
self.__stream_client = BamlStreamClient(self.__runtime, self.__ctx_manager, baml_options)
31+
self.__baml_options = baml_options or {}
3132
@property
3233
def stream(self):
3334
return self.__stream_client
3435

36+
def with_options(
37+
self,
38+
tb: Optional[TypeBuilder] = None,
39+
client_registry: Optional[baml_py.baml_py.ClientRegistry] = None,
40+
collector: Optional[Union[baml_py.baml_py.Collector, List[baml_py.baml_py.Collector]]] = None,
41+
) -> "BamlSyncClient":
42+
"""
43+
Returns a new instance of BamlSyncClient with explicitly typed baml options
44+
for Python 3.8 compatibility.
45+
"""
46+
new_options = {}
47+
48+
# Override if any keyword arguments were provided.
49+
if tb is not None:
50+
new_options["tb"] = tb
51+
if client_registry is not None:
52+
new_options["client_registry"] = client_registry
53+
if collector is not None:
54+
new_options["collector"] = collector
55+
return BamlSyncClient(self.__runtime, self.__ctx_manager, new_options)
56+
3557
{% for fn in funcs %}
3658
def {{ fn.name }}(
3759
self,
@@ -40,13 +62,14 @@ class BamlSyncClient:
4062
{%- endfor %}
4163
baml_options: BamlCallOptions = {},
4264
) -> {{fn.return_type}}:
43-
__tb__ = baml_options.get("tb", None)
65+
options = {**self.__baml_options, **(baml_options or {})}
66+
__tb__ = options.get("tb", None)
4467
if __tb__ is not None:
4568
tb = __tb__._tb # type: ignore (we know how to use this private attribute)
4669
else:
4770
tb = None
48-
__cr__ = baml_options.get("client_registry", None)
49-
collector = baml_options.get("collector", None)
71+
__cr__ = options.get("client_registry", None)
72+
collector = options.get("collector", None)
5073
collectors = collector if isinstance(collector, list) else [collector] if collector is not None else []
5174

5275
raw = self.__runtime.call_function_sync(
@@ -69,10 +92,11 @@ class BamlSyncClient:
6992
class BamlStreamClient:
7093
__runtime: baml_py.BamlRuntime
7194
__ctx_manager: baml_py.BamlCtxManager
72-
73-
def __init__(self, runtime: baml_py.BamlRuntime, ctx_manager: baml_py.BamlCtxManager):
95+
__baml_options: BamlCallOptions
96+
def __init__(self, runtime: baml_py.BamlRuntime, ctx_manager: baml_py.BamlCtxManager, baml_options: Optional[BamlCallOptions] = None):
7497
self.__runtime = runtime
7598
self.__ctx_manager = ctx_manager
99+
self.__baml_options = baml_options or {}
76100

77101
{% for fn in funcs %}
78102
def {{ fn.name }}(
@@ -82,13 +106,14 @@ class BamlStreamClient:
82106
{%- endfor %}
83107
baml_options: BamlCallOptions = {},
84108
) -> baml_py.BamlSyncStream[{{ fn.partial_return_type }}, {{ fn.return_type }}]:
85-
__tb__ = baml_options.get("tb", None)
109+
options = {**self.__baml_options, **(baml_options or {})}
110+
__tb__ = options.get("tb", None)
86111
if __tb__ is not None:
87112
tb = __tb__._tb # type: ignore (we know how to use this private attribute)
88113
else:
89114
tb = None
90-
__cr__ = baml_options.get("client_registry", None)
91-
collector = baml_options.get("collector", None)
115+
__cr__ = options.get("client_registry", None)
116+
collector = options.get("collector", None)
92117
collectors = collector if isinstance(collector, list) else [collector] if collector is not None else []
93118

94119
raw = self.__runtime.stream_function_sync(

‎engine/language_client_codegen/src/ruby/templates/client.rb.j2‎

Lines changed: 31 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -32,11 +32,12 @@ module Baml
3232
sig { returns(BamlStreamClient) }
3333
attr_reader :stream
3434

35-
sig {params(runtime: Baml::Ffi::BamlRuntime).void}
36-
def initialize(runtime:)
35+
sig {params(runtime: Baml::Ffi::BamlRuntime, ctx_manager: Baml::Ffi::RuntimeContextManager, baml_options: T.nilable(T::Hash[Symbol, T.any(Baml::TypeBuilder, Baml::ClientRegistry, T.any(Baml::Collector, T::Array[Baml::Collector]))])).void}
36+
def initialize(runtime:, ctx_manager: nil, baml_options: nil)
3737
@runtime = runtime
38-
@ctx_manager = runtime.create_context_manager()
39-
@stream = BamlStreamClient.new(runtime: @runtime, ctx_manager: @ctx_manager)
38+
@ctx_manager = ctx_manager || runtime.create_context_manager()
39+
@stream = BamlStreamClient.new(runtime: @runtime, ctx_manager: @ctx_manager, baml_options: baml_options)
40+
@baml_options = baml_options
4041
end
4142

4243
sig {params(path: String).returns(BamlClient)}
@@ -69,6 +70,22 @@ module Baml
6970
raise ArgumentError.new("Received unknown keys in baml_options (valid keys: :client_registry, :tb, :collector): #{baml_options.keys - [:client_registry, :tb, :collector]}")
7071
end
7172

73+
# Merge options from initialization with those passed to the method
74+
# Passed options take precedence over initialization options
75+
effective_options = {}
76+
77+
if @baml_options
78+
effective_options = @baml_options.dup
79+
end
80+
81+
# Override with any options passed to this specific call
82+
baml_options.each do |key, value|
83+
effective_options[key] = value
84+
end
85+
86+
# Use the merged options for the rest of the method
87+
baml_options = effective_options
88+
7289
collector = if baml_options[:collector]
7390
baml_options[:collector].is_a?(Array) ? baml_options[:collector] : [baml_options[:collector]]
7491
else
@@ -92,15 +109,20 @@ module Baml
92109

93110
{% endfor %}
94111

112+
sig {params(collector: T.nilable(T.any(Baml::Collector, T::Array[Baml::Collector])), tb: T.nilable(Baml::TypeBuilder), client_registry: T.nilable(Baml::ClientRegistry)).returns(BamlClient)}
113+
def with_options(collector: nil, tb: nil, client_registry: nil)
114+
BamlClient.new(runtime: @runtime, ctx_manager: @ctx_manager, baml_options: {collector: collector, tb: tb, client_registry: client_registry})
115+
end
95116
end
96117

97118
class BamlStreamClient
98119
extend T::Sig
99120

100-
sig {params(runtime: Baml::Ffi::BamlRuntime, ctx_manager: Baml::Ffi::RuntimeContextManager).void}
101-
def initialize(runtime:, ctx_manager:)
121+
sig {params(runtime: Baml::Ffi::BamlRuntime, ctx_manager: Baml::Ffi::RuntimeContextManager, baml_options: T.nilable(T::Hash[Symbol, T.any(Baml::TypeBuilder, Baml::ClientRegistry, T.any(Baml::Collector, T::Array[Baml::Collector]))])).void}
122+
def initialize(runtime:, ctx_manager:, baml_options: nil)
102123
@runtime = runtime
103124
@ctx_manager = ctx_manager
125+
@baml_options = baml_options || {}
104126
end
105127

106128
{% for fn in funcs -%}
@@ -128,6 +150,9 @@ module Baml
128150
raise ArgumentError.new("Received unknown keys in baml_options (valid keys: :client_registry, :tb, :collector): #{baml_options.keys - [:client_registry, :tb, :collector]}")
129151
end
130152

153+
# Merge options from initialization with those passed to the method
154+
baml_options = (@baml_options || {}).merge(baml_options)
155+
131156
collector = if baml_options[:collector]
132157
baml_options[:collector].is_a?(Array) ? baml_options[:collector] : [baml_options[:collector]]
133158
else

‎engine/language_client_codegen/src/typescript/templates/async_client.ts.j2‎

Lines changed: 33 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -14,15 +14,27 @@ import { DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX, DO_NOT_USE_DI
1414
*/
1515
export type RecursivePartialNull<T> = MovedRecursivePartialNull<T>
1616

17+
type BamlCallOptions = {
18+
tb?: TypeBuilder
19+
clientRegistry?: ClientRegistry
20+
collector?: Collector | Collector[]
21+
}
22+
1723
export class BamlAsyncClient {
1824
private runtime: BamlRuntime
1925
private ctx_manager: BamlCtxManager
2026
private stream_client: BamlStreamClient
27+
private baml_options: BamlCallOptions
2128

22-
constructor(runtime: BamlRuntime, ctx_manager: BamlCtxManager) {
29+
constructor(runtime: BamlRuntime, ctx_manager: BamlCtxManager, baml_options?: BamlCallOptions) {
2330
this.runtime = runtime
2431
this.ctx_manager = ctx_manager
25-
this.stream_client = new BamlStreamClient(runtime, ctx_manager)
32+
this.stream_client = new BamlStreamClient(runtime, ctx_manager, baml_options)
33+
this.baml_options = baml_options || {}
34+
}
35+
36+
withOptions(baml_options: BamlCallOptions) {
37+
return new BamlAsyncClient(this.runtime, this.ctx_manager, baml_options)
2638
}
2739

2840
get stream() {
@@ -34,20 +46,21 @@ export class BamlAsyncClient {
3446
{% for (name, optional, type) in fn.args -%}
3547
{{name}}{% if optional %}?{% endif %}: {{type}},
3648
{%- endfor %}
37-
__baml_options__?: { tb?: TypeBuilder, clientRegistry?: ClientRegistry, collector?: Collector | Collector[] }
49+
__baml_options__?: BamlCallOptions
3850
): Promise<{{fn.return_type}}> {
3951
try {
40-
const collector = __baml_options__?.collector ? (Array.isArray(__baml_options__.collector) ? __baml_options__.collector : [__baml_options__.collector]) : [];
52+
const options = { ...this.baml_options, ...(__baml_options__ || {}) }
53+
const collector = options.collector ? (Array.isArray(options.collector) ? options.collector : [options.collector]) : [];
4154
const raw = await this.runtime.callFunction(
4255
"{{fn.name}}",
4356
{
4457
{% for (name, optional, type) in fn.args -%}
4558
"{{name}}": {{name}}{% if optional %}?? null{% endif %}{% if !loop.last %},{% endif %}
4659
{%- endfor %}
4760
},
48-
this.ctx_manager.cloneContext(),
49-
__baml_options__?.tb?.__tb(),
50-
__baml_options__?.clientRegistry,
61+
this.ctx_manager.cloneContext(),
62+
options.tb?.__tb(),
63+
options.clientRegistry,
5164
collector,
5265
)
5366
return raw.parsed(false) as {{fn.return_type}}
@@ -59,7 +72,15 @@ export class BamlAsyncClient {
5972
}
6073

6174
class BamlStreamClient {
62-
constructor(private runtime: BamlRuntime, private ctx_manager: BamlCtxManager) {}
75+
private runtime: BamlRuntime
76+
private ctx_manager: BamlCtxManager
77+
private baml_options: BamlCallOptions
78+
79+
constructor(runtime: BamlRuntime, ctx_manager: BamlCtxManager, baml_options?: BamlCallOptions) {
80+
this.runtime = runtime
81+
this.ctx_manager = ctx_manager
82+
this.baml_options = baml_options || {}
83+
}
6384

6485
{% for fn in funcs %}
6586
{{ fn.name }}(
@@ -69,7 +90,8 @@ class BamlStreamClient {
6990
__baml_options__?: { tb?: TypeBuilder, clientRegistry?: ClientRegistry, collector?: Collector | Collector[] }
7091
): BamlStream<{{ fn.partial_return_type }}, {{ fn.return_type }}> {
7192
try {
72-
const collector = __baml_options__?.collector ? (Array.isArray(__baml_options__.collector) ? __baml_options__.collector : [__baml_options__.collector]) : [];
93+
const options = { ...this.baml_options, ...(__baml_options__ || {}) }
94+
const collector = options.collector ? (Array.isArray(options.collector) ? options.collector : [options.collector]) : [];
7395
const raw = this.runtime.streamFunction(
7496
"{{fn.name}}",
7597
{
@@ -79,8 +101,8 @@ class BamlStreamClient {
79101
},
80102
undefined,
81103
this.ctx_manager.cloneContext(),
82-
__baml_options__?.tb?.__tb(),
83-
__baml_options__?.clientRegistry,
104+
options.tb?.__tb(),
105+
options.clientRegistry,
84106
collector,
85107
)
86108
return new BamlStream<{{ fn.partial_return_type }}, {{ fn.return_type }}>(

0 commit comments

Comments
 (0)