diff --git a/python/langsmith/_internal/_serde.py b/python/langsmith/_internal/_serde.py index 1bf8865c1..739dd848a 100644 --- a/python/langsmith/_internal/_serde.py +++ b/python/langsmith/_internal/_serde.py @@ -79,6 +79,11 @@ def _simple_default(obj): ] +# IMPORTANT: This function is used from Rust code in `langsmith-pyo3` serialization, +# in order to handle serializing these tricky Python types *from Rust*. +# Do not cause this function to become inaccessible (e.g. by deleting +# or renaming it) without also fixing the corresponding Rust code found in: +# rust/crates/langsmith-pyo3/src/serialization/mod.rs def _serialize_json(obj: Any) -> Any: try: if isinstance(obj, (set, tuple)): diff --git a/rust/crates/langsmith-pyo3/src/serialization/mod.rs b/rust/crates/langsmith-pyo3/src/serialization/mod.rs index 74972cad9..819118d63 100644 --- a/rust/crates/langsmith-pyo3/src/serialization/mod.rs +++ b/rust/crates/langsmith-pyo3/src/serialization/mod.rs @@ -1,21 +1,44 @@ +use std::ptr::NonNull; + +use pyo3::types::PyAnyMethods as _; + mod writer; +thread_local! { + static ORJSON_DEFAULT: NonNull = { + pyo3::Python::with_gil(|py| { + let module = match py.import("langsmith._internal._serde") { + Ok(m) => m, + Err(e) => { + let _ = py.import("langsmith").expect("failed to import `langsmith` package; please make sure `langsmith-pyo3` is only used via the `langsmith` package"); + panic!("Failed to import `langsmith._internal._serde` even though `langsmith` can be imported. Did internal `langsmith` package structure change? Underlying error: {e}"); + } + }; + + let function = module.getattr("_serialize_json").expect("`_serialize_json` function not found").as_ptr(); + NonNull::new(function).expect("function was null, which shouldn't ever happen") + }) + } +} + pub(crate) fn dumps(ptr: *mut pyo3_ffi::PyObject) -> Result, String> { let mut writer = writer::BufWriter::new(); - let obj = orjson::PyObjectSerializer::new( - ptr, - orjson::SerializerState::new(Default::default()), - None, - ); + ORJSON_DEFAULT.with(|default| { + let obj = orjson::PyObjectSerializer::new( + ptr, + orjson::SerializerState::new(Default::default()), + Some(*default), + ); - let res = orjson::to_writer(&mut writer, &obj); - match res { - Ok(_) => Ok(writer.finish()), - Err(err) => { - // Make sure we drop the allocated buffer. - let _ = writer.into_inner(); - Err(err.to_string()) + let res = orjson::to_writer(&mut writer, &obj); + match res { + Ok(_) => Ok(writer.finish()), + Err(err) => { + // Make sure we drop the allocated buffer. + let _ = writer.into_inner(); + Err(err.to_string()) + } } - } + }) }