Serialize pydantic and other tricky objects correctly from Rust.

From the Rust bindings, import the `_serialize_json` function from `langsmith._internal._serde`, then use it as the default fallback if `orjson` serialization can't handle some object. This makes the Rust serialization code equivalent to the `_orjson.dumps()` call inside `langsmith._internal._serde.dumps_json`.

Open question: do we need to handle UTF surrogate characters similarly to what `langsmith._internal._serde.dumps_json` does? I wasn't able to find enough information on the issue encountered there. If that's still needed, we should do that in a new PR.

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)):

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_ffi::PyObject> = {
+        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<Vec<u8>, 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())
+            }
         }
-    }
+    })
 }