Skip to content

Commit b83557c

Browse files
tbermantberman
authored
Fix arguments for functions from go (#1884)
This is tested to actually work, no longer relying on maps for arguments. <!-- ELLIPSIS_HIDDEN --> ---- > [!IMPORTANT] > This PR replaces map-based function argument handling with `BamlFunctionArguments`, introducing `CFFIFunctionArguments` and updating related files and tests. > > - **Behavior**: > - Introduces `CFFIFunctionArguments` in `cffi_generated.rs` to handle function arguments. > - Replaces map-based argument handling with `BamlFunctionArguments` in `client.go.j2` and `ctypes.rs`. > - Adds `buffer_to_cffi_function_arguments()` in `ctypes.rs` for argument conversion. > - **Testing**: > - Adds `TestEncodeFunctionArguments` in `encode_test.go` to verify new argument handling. > - **Enums and Constants**: > - Updates `ENUM_MAX_CFFIVALUE_UNION` and `ENUM_VALUES_CFFIVALUE_UNION` in `cffi_generated.rs` to include `CFFIFunctionArguments`. > > <sup>This description was created by </sup>[<img alt="Ellipsis" src="https://img.shields.io/badge/Ellipsis-blue?color=175173">](https://www.ellipsis.dev?ref=BoundaryML%2Fbaml&utm_source=github&utm_medium=referral)<sup> for 9a86bef. You can [customize](https://app.ellipsis.dev/BoundaryML/settings/summaries) this summary. It will automatically update as commits are pushed.</sup> <!-- ELLIPSIS_HIDDEN -->
1 parent c9ae9ce commit b83557c

14 files changed

Lines changed: 1381 additions & 387 deletions

File tree

‎engine/language_client_cffi/src/cffi/cffi_generated.rs‎

Lines changed: 127 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,10 +21,10 @@ pub mod cffi {
2121
#[deprecated(since = "2.0.0", note = "Use associated constants instead. This will no longer be generated in 2021.")]
2222
pub const ENUM_MIN_CFFIVALUE_UNION: u8 = 0;
2323
#[deprecated(since = "2.0.0", note = "Use associated constants instead. This will no longer be generated in 2021.")]
24-
pub const ENUM_MAX_CFFIVALUE_UNION: u8 = 13;
24+
pub const ENUM_MAX_CFFIVALUE_UNION: u8 = 14;
2525
#[deprecated(since = "2.0.0", note = "Use associated constants instead. This will no longer be generated in 2021.")]
2626
#[allow(non_camel_case_types)]
27-
pub const ENUM_VALUES_CFFIVALUE_UNION: [CFFIValueUnion; 14] = [
27+
pub const ENUM_VALUES_CFFIVALUE_UNION: [CFFIValueUnion; 15] = [
2828
CFFIValueUnion::NONE,
2929
CFFIValueUnion::CFFIValueString,
3030
CFFIValueUnion::CFFIValueInt,
@@ -39,6 +39,7 @@ pub const ENUM_VALUES_CFFIVALUE_UNION: [CFFIValueUnion; 14] = [
3939
CFFIValueUnion::CFFIValueUnionVariant,
4040
CFFIValueUnion::CFFIValueChecked,
4141
CFFIValueUnion::CFFIValueStreamingState,
42+
CFFIValueUnion::CFFIFunctionArguments,
4243
];
4344

4445
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Default)]
@@ -60,9 +61,10 @@ impl CFFIValueUnion {
6061
pub const CFFIValueUnionVariant: Self = Self(11);
6162
pub const CFFIValueChecked: Self = Self(12);
6263
pub const CFFIValueStreamingState: Self = Self(13);
64+
pub const CFFIFunctionArguments: Self = Self(14);
6365

6466
pub const ENUM_MIN: u8 = 0;
65-
pub const ENUM_MAX: u8 = 13;
67+
pub const ENUM_MAX: u8 = 14;
6668
pub const ENUM_VALUES: &'static [Self] = &[
6769
Self::NONE,
6870
Self::CFFIValueString,
@@ -78,6 +80,7 @@ impl CFFIValueUnion {
7880
Self::CFFIValueUnionVariant,
7981
Self::CFFIValueChecked,
8082
Self::CFFIValueStreamingState,
83+
Self::CFFIFunctionArguments,
8184
];
8285
/// Returns the variant's name or "" if unknown.
8386
pub fn variant_name(self) -> Option<&'static str> {
@@ -96,6 +99,7 @@ impl CFFIValueUnion {
9699
Self::CFFIValueUnionVariant => Some("CFFIValueUnionVariant"),
97100
Self::CFFIValueChecked => Some("CFFIValueChecked"),
98101
Self::CFFIValueStreamingState => Some("CFFIValueStreamingState"),
102+
Self::CFFIFunctionArguments => Some("CFFIFunctionArguments"),
99103
_ => None,
100104
}
101105
}
@@ -912,6 +916,21 @@ impl<'a> CFFIValueHolder<'a> {
912916
}
913917
}
914918

919+
#[inline]
920+
#[allow(non_snake_case)]
921+
pub fn value_as_cffifunction_arguments(&self) -> Option<CFFIFunctionArguments<'a>> {
922+
if self.value_type() == CFFIValueUnion::CFFIFunctionArguments {
923+
self.value().map(|t| {
924+
// Safety:
925+
// Created from a valid Table for this object
926+
// Which contains a valid union in this slot
927+
unsafe { CFFIFunctionArguments::init_from_table(t) }
928+
})
929+
} else {
930+
None
931+
}
932+
}
933+
915934
}
916935

917936
impl flatbuffers::Verifiable for CFFIValueHolder<'_> {
@@ -936,6 +955,7 @@ impl flatbuffers::Verifiable for CFFIValueHolder<'_> {
936955
CFFIValueUnion::CFFIValueUnionVariant => v.verify_union_variant::<flatbuffers::ForwardsUOffset<CFFIValueUnionVariant>>("CFFIValueUnion::CFFIValueUnionVariant", pos),
937956
CFFIValueUnion::CFFIValueChecked => v.verify_union_variant::<flatbuffers::ForwardsUOffset<CFFIValueChecked>>("CFFIValueUnion::CFFIValueChecked", pos),
938957
CFFIValueUnion::CFFIValueStreamingState => v.verify_union_variant::<flatbuffers::ForwardsUOffset<CFFIValueStreamingState>>("CFFIValueUnion::CFFIValueStreamingState", pos),
958+
CFFIValueUnion::CFFIFunctionArguments => v.verify_union_variant::<flatbuffers::ForwardsUOffset<CFFIFunctionArguments>>("CFFIValueUnion::CFFIFunctionArguments", pos),
939959
_ => Ok(()),
940960
}
941961
})?
@@ -1081,6 +1101,13 @@ impl core::fmt::Debug for CFFIValueHolder<'_> {
10811101
ds.field("value", &"InvalidFlatbuffer: Union discriminant does not match value.")
10821102
}
10831103
},
1104+
CFFIValueUnion::CFFIFunctionArguments => {
1105+
if let Some(x) = self.value_as_cffifunction_arguments() {
1106+
ds.field("value", &x)
1107+
} else {
1108+
ds.field("value", &"InvalidFlatbuffer: Union discriminant does not match value.")
1109+
}
1110+
},
10841111
_ => {
10851112
let x: Option<()> = None;
10861113
ds.field("value", &x)
@@ -5538,6 +5565,103 @@ impl core::fmt::Debug for CFFIFieldTypeStreamState<'_> {
55385565
ds.finish()
55395566
}
55405567
}
5568+
pub enum CFFIFunctionArgumentsOffset {}
5569+
#[derive(Copy, Clone, PartialEq)]
5570+
5571+
pub struct CFFIFunctionArguments<'a> {
5572+
pub _tab: flatbuffers::Table<'a>,
5573+
}
5574+
5575+
impl<'a> flatbuffers::Follow<'a> for CFFIFunctionArguments<'a> {
5576+
type Inner = CFFIFunctionArguments<'a>;
5577+
#[inline]
5578+
unsafe fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
5579+
Self { _tab: flatbuffers::Table::new(buf, loc) }
5580+
}
5581+
}
5582+
5583+
impl<'a> CFFIFunctionArguments<'a> {
5584+
pub const VT_KWARGS: flatbuffers::VOffsetT = 4;
5585+
5586+
#[inline]
5587+
pub unsafe fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
5588+
CFFIFunctionArguments { _tab: table }
5589+
}
5590+
#[allow(unused_mut)]
5591+
pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>(
5592+
_fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>,
5593+
args: &'args CFFIFunctionArgumentsArgs<'args>
5594+
) -> flatbuffers::WIPOffset<CFFIFunctionArguments<'bldr>> {
5595+
let mut builder = CFFIFunctionArgumentsBuilder::new(_fbb);
5596+
if let Some(x) = args.kwargs { builder.add_kwargs(x); }
5597+
builder.finish()
5598+
}
5599+
5600+
5601+
#[inline]
5602+
pub fn kwargs(&self) -> Option<flatbuffers::Vector<'a, flatbuffers::ForwardsUOffset<CFFIMapEntry<'a>>>> {
5603+
// Safety:
5604+
// Created from valid Table for this object
5605+
// which contains a valid value in this slot
5606+
unsafe { self._tab.get::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'a, flatbuffers::ForwardsUOffset<CFFIMapEntry>>>>(CFFIFunctionArguments::VT_KWARGS, None)}
5607+
}
5608+
}
5609+
5610+
impl flatbuffers::Verifiable for CFFIFunctionArguments<'_> {
5611+
#[inline]
5612+
fn run_verifier(
5613+
v: &mut flatbuffers::Verifier, pos: usize
5614+
) -> Result<(), flatbuffers::InvalidFlatbuffer> {
5615+
use self::flatbuffers::Verifiable;
5616+
v.visit_table(pos)?
5617+
.visit_field::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'_, flatbuffers::ForwardsUOffset<CFFIMapEntry>>>>("kwargs", Self::VT_KWARGS, false)?
5618+
.finish();
5619+
Ok(())
5620+
}
5621+
}
5622+
pub struct CFFIFunctionArgumentsArgs<'a> {
5623+
pub kwargs: Option<flatbuffers::WIPOffset<flatbuffers::Vector<'a, flatbuffers::ForwardsUOffset<CFFIMapEntry<'a>>>>>,
5624+
}
5625+
impl<'a> Default for CFFIFunctionArgumentsArgs<'a> {
5626+
#[inline]
5627+
fn default() -> Self {
5628+
CFFIFunctionArgumentsArgs {
5629+
kwargs: None,
5630+
}
5631+
}
5632+
}
5633+
5634+
pub struct CFFIFunctionArgumentsBuilder<'a: 'b, 'b> {
5635+
fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>,
5636+
start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
5637+
}
5638+
impl<'a: 'b, 'b> CFFIFunctionArgumentsBuilder<'a, 'b> {
5639+
#[inline]
5640+
pub fn add_kwargs(&mut self, kwargs: flatbuffers::WIPOffset<flatbuffers::Vector<'b , flatbuffers::ForwardsUOffset<CFFIMapEntry<'b >>>>) {
5641+
self.fbb_.push_slot_always::<flatbuffers::WIPOffset<_>>(CFFIFunctionArguments::VT_KWARGS, kwargs);
5642+
}
5643+
#[inline]
5644+
pub fn new(_fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>) -> CFFIFunctionArgumentsBuilder<'a, 'b> {
5645+
let start = _fbb.start_table();
5646+
CFFIFunctionArgumentsBuilder {
5647+
fbb_: _fbb,
5648+
start_: start,
5649+
}
5650+
}
5651+
#[inline]
5652+
pub fn finish(self) -> flatbuffers::WIPOffset<CFFIFunctionArguments<'a>> {
5653+
let o = self.fbb_.end_table(self.start_);
5654+
flatbuffers::WIPOffset::new(o.value())
5655+
}
5656+
}
5657+
5658+
impl core::fmt::Debug for CFFIFunctionArguments<'_> {
5659+
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
5660+
let mut ds = f.debug_struct("CFFIFunctionArguments");
5661+
ds.field("kwargs", &self.kwargs());
5662+
ds.finish()
5663+
}
5664+
}
55415665
pub enum CFFICheckTypeOffset {}
55425666
#[derive(Copy, Clone, PartialEq)]
55435667

‎engine/language_client_cffi/src/ctypes.rs‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,11 +7,21 @@ mod cffi_generated;
77

88
use cffi_generated::cffi::*;
99

10+
use crate::BamlFunctionArguments;
11+
1012
pub fn buffer_to_cffi_value_holder(buffer: &[u8]) -> Result<BamlValue> {
1113
let root = flatbuffers::root::<CFFIValueHolder>(buffer)?;
1214
Ok(root.into())
1315
}
1416

17+
pub fn buffer_to_cffi_function_arguments(buffer: &[u8]) -> Result<BamlFunctionArguments> {
18+
let root = flatbuffers::root::<CFFIValueHolder>(buffer)?;
19+
Ok(root
20+
.value_as_cffifunction_arguments()
21+
.expect("Failed to convert CFFIValueHolder to CFFIFunctionArguments")
22+
.into())
23+
}
24+
1525
impl From<cffi_generated::cffi::CFFIValueHolder<'_>> for BamlValue {
1626
fn from(value: cffi_generated::cffi::CFFIValueHolder) -> Self {
1727
let value_type = value.value_type();
@@ -73,6 +83,9 @@ impl From<cffi_generated::cffi::CFFIValueHolder<'_>> for BamlValue {
7383
.value_as_cffivalue_streaming_state()
7484
.expect("Failed to convert CFFIValueStreamingState to BamlValue")
7585
.into(),
86+
CFFIValueUnion::CFFIFunctionArguments => {
87+
panic!("CFFIFunctionArguments is not supported in BamlValue");
88+
}
7689
other => {
7790
panic!("Unsupported value type: {:?}", other);
7891
}
@@ -203,6 +216,18 @@ impl From<CFFIValueUnionVariant<'_>> for BamlValue {
203216
}
204217
}
205218

219+
impl From<CFFIFunctionArguments<'_>> for BamlFunctionArguments {
220+
fn from(value: CFFIFunctionArguments) -> Self {
221+
let kwargs = value
222+
.kwargs()
223+
.expect("Failed to have CFFIFunctionArguments kwargs")
224+
.into_iter()
225+
.map(|v| v.into())
226+
.collect();
227+
BamlFunctionArguments { kwargs }
228+
}
229+
}
230+
206231
impl From<CFFIValueChecked<'_>> for BamlValue {
207232
fn from(_value: CFFIValueChecked) -> Self {
208233
unimplemented!("CFFIValueChecked is not supported");

‎engine/language_client_cffi/src/lib.rs‎

Lines changed: 17 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,10 @@ use once_cell::sync::{Lazy, OnceCell};
99

1010
const VERSION: &str = env!("CARGO_PKG_VERSION");
1111

12+
struct BamlFunctionArguments {
13+
kwargs: baml_types::BamlMap<String, BamlValue>,
14+
}
15+
1216
#[no_mangle]
1317
pub extern "C" fn version() -> *const libc::c_char {
1418
let version = CString::new(VERSION).unwrap();
@@ -81,21 +85,7 @@ pub extern "C" fn invoke_runtime_cli(args: *const *const libc::c_char) -> libc::
8185
use std::ffi::CString;
8286
use std::os::raw::c_char;
8387

84-
use baml_types::{BamlMap, BamlValue};
85-
86-
/// Convert Flatbuffers encoded arguments to a BamlMap<String, BamlValue>
87-
fn ckwargs_to_map(
88-
encoded_args: *const libc::c_char,
89-
length: usize,
90-
) -> Result<BamlMap<String, BamlValue>> {
91-
let buffer = unsafe { std::slice::from_raw_parts(encoded_args as *const u8, length) };
92-
let value = ctypes::buffer_to_cffi_value_holder(buffer)?;
93-
if let Some(map) = value.as_map_owned() {
94-
Ok(map)
95-
} else {
96-
Err(anyhow::anyhow!("Invalid encoded arguments"))
97-
}
98-
}
88+
use baml_types::BamlValue;
9989

10090
pub type CallbackFn = extern "C" fn(call_id: u32, is_done: bool, content: *const i8, length: usize);
10191

@@ -192,7 +182,8 @@ fn call_function_from_c_inner(
192182
};
193183

194184
// Convert keyword arguments.
195-
let keyword_args = ckwargs_to_map(encoded_args, length)?;
185+
let buffer = unsafe { std::slice::from_raw_parts(encoded_args as *const u8, length) };
186+
let function_args = ctypes::buffer_to_cffi_function_arguments(buffer)?;
196187

197188
let ctx = runtime.create_ctx_manager(BamlValue::String("cffi".to_string()), None);
198189

@@ -201,7 +192,7 @@ fn call_function_from_c_inner(
201192
let rt = RUNTIME.clone();
202193
rt.spawn(async move {
203194
let (result, _) = runtime
204-
.call_function(func_name, &keyword_args, &ctx, None, None, None)
195+
.call_function(func_name, &function_args.kwargs, &ctx, None, None, None)
205196
.await;
206197
safe_trigger_callback(id, true, result);
207198
});
@@ -246,16 +237,17 @@ fn call_function_stream_from_c_inner(
246237
};
247238

248239
// Convert keyword arguments.
249-
let keyword_args = ckwargs_to_map(encoded_args, length)?;
240+
let buffer = unsafe { std::slice::from_raw_parts(encoded_args as *const u8, length) };
241+
let function_args = ctypes::buffer_to_cffi_function_arguments(buffer)?;
250242

251243
let ctx = runtime.create_ctx_manager(BamlValue::String("cffi".to_string()), None);
252-
let mut stream = match runtime.stream_function(func_name, &keyword_args, &ctx, None, None, None)
253-
{
254-
Ok(stream) => stream,
255-
Err(e) => {
256-
return Err(anyhow::anyhow!("Failed to stream function: {}", e));
257-
}
258-
};
244+
let mut stream =
245+
match runtime.stream_function(func_name, &function_args.kwargs, &ctx, None, None, None) {
246+
Ok(stream) => stream,
247+
Err(e) => {
248+
return Err(anyhow::anyhow!("Failed to stream function: {}", e));
249+
}
250+
};
259251

260252
let ctx = runtime.create_ctx_manager(BamlValue::String("cffi".to_string()), None);
261253

‎engine/language_client_cffi/types/cffi.fbs‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,8 @@ union CFFIValueUnion {
1818
CFFIValueTuple,
1919
CFFIValueUnionVariant, // corresponds to the Rust `Union(Vec<CFFIFieldType>, Box<CFFIValue>)`
2020
CFFIValueChecked,
21-
CFFIValueStreamingState
21+
CFFIValueStreamingState,
22+
CFFIFunctionArguments
2223
}
2324

2425
// The wrapper table for CFFIValue.
@@ -231,6 +232,10 @@ table CFFIFieldTypeStreamState {
231232
value: CFFIFieldTypeHolder;
232233
}
233234

235+
table CFFIFunctionArguments {
236+
kwargs: [CFFIMapEntry];
237+
}
238+
234239
// -----------------------------------------------------------------------------
235240
// CFFICheck definitions
236241
// -----------------------------------------------------------------------------

‎engine/language_client_codegen/src/go/templates/client.go.j2‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,9 @@ func castOptional[T any](result any, castResult func(any) T) *T {
5151
func {{ fn.name }}(ctx context.Context{% for (name, type) in fn.args -%}
5252
, {{name}} {{type}}
5353
{%- endfor %}) (*{{ fn.return_type }}, error) {
54-
args := map[string]any{ {% for (name, type) in fn.args -%}"{{name}}": {{name}},{% endfor %} }
54+
args := baml.BamlFunctionArguments{
55+
Kwargs: map[string]any{ {% for (name, type) in fn.args -%}"{{name}}": {{name}},{% endfor %} },
56+
}
5557
encoded, err := baml.EncodeRoot(args)
5658
if err != nil {
5759
panic(err)
@@ -77,7 +79,9 @@ func {{ fn.name }}(ctx context.Context{% for (name, type) in fn.args -%}
7779
func (*stream) {{ fn.name }}(ctx context.Context{% for (name, type) in fn.args -%}
7880
, {{name}} {{type}}
7981
{%- endfor %}) <-chan {{ fn.return_type }} {
80-
args := map[string]any{ {% for (name, type) in fn.args -%}"{{name}}": {{name}},{% endfor %} }
82+
args := baml.BamlFunctionArguments{
83+
Kwargs: map[string]any{ {% for (name, type) in fn.args -%}"{{name}}": {{name}},{% endfor %} },
84+
}
8185
encoded, err := baml.EncodeRoot(args)
8286
if err != nil {
8387
panic(err)

0 commit comments

Comments
 (0)