Add Rust implementation for Python client's update_run() method. (#…

…1314)

Implement the `update_run()` client method in the Rust PyO3 bindings,
and call it from the Python client if the Rust bindings are enabled.

python/langsmith/client.py

@@ -1752,7 +1752,10 @@ def update_run(
             data["events"] = events
         if data["extra"]:
             self._insert_runtime_env([data])
-        if use_multipart and self.tracing_queue is not None:
+
+        if self._pyo3_client is not None:
+            self._pyo3_client.update_run(data)
+        elif use_multipart and self.tracing_queue is not None:
             # not collecting attachments currently, use empty dict
             serialized_op = serialize_run_dict(operation="patch", payload=data)
             self.tracing_queue.put(

rust/crates/langsmith-pyo3/src/blocking_tracing_client.rs

@@ -47,8 +47,7 @@ impl BlockingTracingClient {
         Ok(Self { client: Arc::from(client) })
     }
 
-    // N.B.: We use `Py<Self>` so that we don't hold the GIL while running this method.
-    //       `slf.get()` below is only valid if the `Self` type is `Sync` and `pyclass(frozen)`,
+    // N.B.: `slf.get()` below is only valid if the `Self` type is `Sync` and `pyclass(frozen)`,
     //       which is enforced at compile-time.
     pub fn create_run(
         slf: &Bound<'_, Self>,
@@ -59,6 +58,17 @@ impl BlockingTracingClient {
             .map_err(|e| into_py_err(slf.py(), e))
     }
 
+    // N.B.: `slf.get()` below is only valid if the `Self` type is `Sync` and `pyclass(frozen)`,
+    //       which is enforced at compile-time.
+    pub fn update_run(
+        slf: &Bound<'_, Self>,
+        run: super::py_run::RunUpdateExtended,
+    ) -> PyResult<()> {
+        let unpacked = slf.get();
+        Python::allow_threads(slf.py(), || unpacked.client.submit_run_update(run.into_inner()))
+            .map_err(|e| into_py_err(slf.py(), e))
+    }
+
     pub fn drain(slf: &Bound<'_, Self>) -> PyResult<()> {
         let unpacked = slf.get();
         Python::allow_threads(slf.py(), || unpacked.client.drain())

rust/crates/langsmith-pyo3/src/py_run.rs

@@ -46,43 +46,43 @@ impl FromPyObject<'_> for RunCreateExtended {
     }
 }
 
-fn extract_attachments(value: &Bound<'_, PyAny>) -> PyResult<Option<Vec<Attachment>>> {
-    if value.is_none() {
-        return Ok(None);
-    }
-
-    let mapping = value.downcast::<PyMapping>()?;
+#[derive(Debug)]
+pub struct RunUpdateExtended(langsmith_tracing_client::client::RunUpdateExtended);
 
-    let size = mapping.len()?;
-    if size == 0 {
-        return Ok(None);
+impl RunUpdateExtended {
+    #[inline]
+    pub(crate) fn into_inner(self) -> langsmith_tracing_client::client::RunUpdateExtended {
+        self.0
     }
+}
 
-    let mut attachments = Vec::with_capacity(size);
-
-    for result in mapping.items()?.iter()? {
-        let key_value_pair = result?;
-
-        let key_item = key_value_pair.get_item(0)?;
-        let key = key_item.extract::<&str>()?;
+impl FromPyObject<'_> for RunUpdateExtended {
+    fn extract_bound(value: &Bound<'_, PyAny>) -> PyResult<Self> {
+        let run_update = value.extract::<RunUpdate>()?.into_inner();
+
+        // TODO: attachments are WIP at the moment, ignore them here for now.
+        //
+        // let attachments = {
+        //     if let Ok(attachments_value) = value.get_item(pyo3::intern!(value.py(), "attachments"))
+        //     {
+        //         extract_attachments(&attachments_value)?
+        //     } else {
+        //         None
+        //     }
+        // };
+        let attachments = None;
 
-        // Each value in the attachments dict is a (mime_type, bytes) tuple.
-        let value = key_value_pair.get_item(1)?;
-        let value_tuple = value.downcast_exact::<PyTuple>()?;
-        let mime_type_value = value_tuple.get_item(0)?;
-        let bytes_value = value_tuple.get_item(1)?;
+        let io = RunIO {
+            inputs: serialize_optional_dict_value(value, pyo3::intern!(value.py(), "inputs"))?,
+            outputs: serialize_optional_dict_value(value, pyo3::intern!(value.py(), "outputs"))?,
+        };
 
-        attachments.push(Attachment {
-            // TODO: It's unclear whether the key in the attachments dict is
-            //       the `filename`` or the `ref_name`, and where the other one is coming from.
-            ref_name: key.to_string(),
-            filename: key.to_string(),
-            data: bytes_value.extract()?,
-            content_type: mime_type_value.extract()?,
-        });
+        Ok(Self(langsmith_tracing_client::client::RunUpdateExtended {
+            run_update,
+            io,
+            attachments,
+        }))
     }
-
-    Ok(Some(attachments))
 }
 
 #[derive(Debug)]
@@ -137,6 +137,26 @@ impl FromPyObject<'_> for RunCreate {
     }
 }
 
+#[derive(Debug)]
+pub(crate) struct RunUpdate(langsmith_tracing_client::client::RunUpdate);
+
+impl FromPyObject<'_> for RunUpdate {
+    fn extract_bound(value: &Bound<'_, PyAny>) -> PyResult<Self> {
+        let common = RunCommon::extract_bound(value)?.into_inner();
+
+        let end_time = extract_time_value(&value.get_item(pyo3::intern!(value.py(), "end_time"))?)?;
+
+        Ok(Self(langsmith_tracing_client::client::RunUpdate { common, end_time }))
+    }
+}
+
+impl RunUpdate {
+    #[inline]
+    pub(crate) fn into_inner(self) -> langsmith_tracing_client::client::RunUpdate {
+        self.0
+    }
+}
+
 #[derive(Debug)]
 pub(crate) struct RunCommon(langsmith_tracing_client::client::RunCommon);
 
@@ -196,6 +216,45 @@ impl FromPyObject<'_> for RunCommon {
     }
 }
 
+fn extract_attachments(value: &Bound<'_, PyAny>) -> PyResult<Option<Vec<Attachment>>> {
+    if value.is_none() {
+        return Ok(None);
+    }
+
+    let mapping = value.downcast::<PyMapping>()?;
+
+    let size = mapping.len()?;
+    if size == 0 {
+        return Ok(None);
+    }
+
+    let mut attachments = Vec::with_capacity(size);
+
+    for result in mapping.items()?.iter()? {
+        let key_value_pair = result?;
+
+        let key_item = key_value_pair.get_item(0)?;
+        let key = key_item.extract::<&str>()?;
+
+        // Each value in the attachments dict is a (mime_type, bytes) tuple.
+        let value = key_value_pair.get_item(1)?;
+        let value_tuple = value.downcast_exact::<PyTuple>()?;
+        let mime_type_value = value_tuple.get_item(0)?;
+        let bytes_value = value_tuple.get_item(1)?;
+
+        attachments.push(Attachment {
+            // TODO: It's unclear whether the key in the attachments dict is
+            //       the `filename`` or the `ref_name`, and where the other one is coming from.
+            ref_name: key.to_string(),
+            filename: key.to_string(),
+            data: bytes_value.extract()?,
+            content_type: mime_type_value.extract()?,
+        });
+    }
+
+    Ok(Some(attachments))
+}
+
 /// Get an optional string from a Python `None`, string, or string-like object such as a UUID value.
 fn extract_string_like_or_none(value: Option<&Bound<'_, PyAny>>) -> PyResult<Option<String>> {
     match value {