diff --git a/rust-core/Cargo.lock b/rust-core/Cargo.lock index 082228bd7..8f3dadaf4 100644 --- a/rust-core/Cargo.lock +++ b/rust-core/Cargo.lock @@ -2,12 +2,6 @@ # It is not intended for manual editing. version = 4 -[[package]] -name = "autocfg" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" - [[package]] name = "bitflags" version = "2.11.0" @@ -20,21 +14,6 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" -[[package]] -name = "heck" -version = "0.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" - -[[package]] -name = "indoc" -version = "2.0.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "79cf5c93f93228cf8efb3ba362535fb11199ac548a09ce117c9b1adc3030d706" -dependencies = [ - "rustversion", -] - [[package]] name = "itoa" version = "1.0.17" @@ -47,7 +26,6 @@ version = "0.1.0" dependencies = [ "libc", "parking_lot", - "pyo3", "serde", "serde_json", "tokio", @@ -74,21 +52,6 @@ version = "2.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79" -[[package]] -name = "memoffset" -version = "0.9.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a" -dependencies = [ - "autocfg", -] - -[[package]] -name = "once_cell" -version = "1.21.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" - [[package]] name = "parking_lot" version = "0.12.5" @@ -118,12 +81,6 @@ version = "0.2.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" -[[package]] -name = "portable-atomic" -version = "1.13.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" - [[package]] name = "proc-macro2" version = "1.0.106" @@ -133,69 +90,6 @@ dependencies = [ "unicode-ident", ] -[[package]] -name = "pyo3" -version = "0.23.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7778bffd85cf38175ac1f545509665d0b9b92a198ca7941f131f85f7a4f9a872" -dependencies = [ - "cfg-if", - "indoc", - "libc", - "memoffset", - "once_cell", - "portable-atomic", - "pyo3-build-config", - "pyo3-ffi", - "pyo3-macros", - "unindent", -] - -[[package]] -name = "pyo3-build-config" -version = "0.23.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "94f6cbe86ef3bf18998d9df6e0f3fc1050a8c5efa409bf712e661a4366e010fb" -dependencies = [ - "once_cell", - "target-lexicon", -] - -[[package]] -name = "pyo3-ffi" -version = "0.23.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e9f1b4c431c0bb1c8fb0a338709859eed0d030ff6daa34368d3b152a63dfdd8d" -dependencies = [ - "libc", - "pyo3-build-config", -] - -[[package]] -name = "pyo3-macros" -version = "0.23.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fbc2201328f63c4710f68abdf653c89d8dbc2858b88c5d88b0ff38a75288a9da" -dependencies = [ - "proc-macro2", - "pyo3-macros-backend", - "quote", - "syn", -] - -[[package]] -name = "pyo3-macros-backend" -version = "0.23.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fca6726ad0f3da9c9de093d6f116a93c1a38e417ed73bf138472cf4064f72028" -dependencies = [ - "heck", - "proc-macro2", - "pyo3-build-config", - "quote", - "syn", -] - [[package]] name = "quote" version = "1.0.45" @@ -214,12 +108,6 @@ dependencies = [ "bitflags", ] -[[package]] -name = "rustversion" -version = "1.0.22" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" - [[package]] name = "scopeguard" version = "1.2.0" @@ -286,12 +174,6 @@ dependencies = [ "unicode-ident", ] -[[package]] -name = "target-lexicon" -version = "0.12.16" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1" - [[package]] name = "tokio" version = "1.50.0" @@ -319,12 +201,6 @@ version = "1.0.24" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" -[[package]] -name = "unindent" -version = "0.2.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3" - [[package]] name = "windows-link" version = "0.2.1" diff --git a/rust-core/Cargo.toml b/rust-core/Cargo.toml index 70fcbd3ac..083cfcd63 100644 --- a/rust-core/Cargo.toml +++ b/rust-core/Cargo.toml @@ -7,12 +7,7 @@ edition = "2021" name = "langgraph_rust_core" crate-type = ["cdylib", "rlib"] -[features] -default = ["python-bindings"] -python-bindings = ["dep:pyo3"] - [dependencies] -pyo3 = { version = "0.23.5", features = ["extension-module"], optional = true } serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" parking_lot = "0.12" diff --git a/rust-core/src/engine.rs b/rust-core/src/engine.rs index 724ef657e..7b0b48297 100644 --- a/rust-core/src/engine.rs +++ b/rust-core/src/engine.rs @@ -279,6 +279,22 @@ pub fn merge_json_update(state: &mut Value, update: Option) { } } +enum SchedulerEventJson { + Node(Result), + Resume { + node: String, + arg: Value, + event: WaitEvent, + }, + WaitError(String), +} + +struct NodeExecutionJson { + node: String, + arg: Value, + outcome: NodeOutcome, +} + #[derive(Clone, Default)] pub struct Engine { channels: Arc>>>, @@ -491,3 +507,145 @@ impl Engine { })) } } + +fn spawn_json_node_task( + node: String, + arg: Value, + state_snapshot: Value, + tx: tokio_mpsc::UnboundedSender, + callback: Arc, +) -> Result<(), String> +where + F: Fn(String, Value, Value) -> Result, String> + + Send + + Sync + + 'static, +{ + node_pool_execute(move || { + let node_for_result = node.clone(); + let arg_for_result = arg.clone(); + let result = callback(node.clone(), arg, state_snapshot).map(|payload| NodeExecutionJson { + node: node_for_result, + arg: arg_for_result, + outcome: payload, + }); + let _ = tx.send(SchedulerEventJson::Node(result)); + }) +} + +pub async fn run_graph_json_with_callback( + entry_point: String, + finish_point: String, + initial_state: Value, + initial_input: Value, + engine: Engine, + callback: F, +) -> Result +where + F: Fn(String, Value, Value) -> Result, String> + + Send + + Sync + + 'static, +{ + let callback = Arc::new(callback); + let (tx, mut rx) = tokio_mpsc::unbounded_channel::(); + let state = Arc::new(StdMutex::new(initial_state)); + let tx_for_spawn = tx.clone(); + let state_for_spawn = Arc::clone(&state); + let state_for_merge = Arc::clone(&state); + let mut active: usize = 1; + let mut waiting: usize = 0; + spawn_json_node_task( + entry_point, + initial_input, + state_for_spawn + .lock() + .expect("state mutex poisoned") + .clone(), + tx_for_spawn.clone(), + Arc::clone(&callback), + )?; + + while active > 0 || waiting > 0 { + let evt = rx + .recv() + .await + .ok_or_else(|| "scheduler event channel closed".to_string())?; + match evt { + SchedulerEventJson::Node(result) => { + active = active.saturating_sub(1); + let exec = result?; + match exec.outcome { + NodeOutcome::Completed(node_result) => { + let mut guard = state_for_merge.lock().expect("state mutex poisoned"); + merge_json_update(&mut guard, node_result.update); + drop(guard); + if exec.node == finish_point { + break; + } + for send in node_result.sends { + active += 1; + let snapshot = state_for_spawn + .lock() + .expect("state mutex poisoned") + .clone(); + spawn_json_node_task( + send.node, + send.arg, + snapshot, + tx_for_spawn.clone(), + Arc::clone(&callback), + )?; + } + } + NodeOutcome::Suspended { wait } => { + waiting += 1; + let tx_wait = tx_for_spawn.clone(); + let node = exec.node; + let arg = exec.arg; + let engine_for_wait = engine.clone(); + tokio::spawn(async move { + match engine_for_wait.wait_request_async(&wait).await { + Ok(event) => { + let _ = tx_wait.send(SchedulerEventJson::Resume { + node, + arg, + event, + }); + } + Err(e) => { + let _ = tx_wait.send(SchedulerEventJson::WaitError(e)); + } + } + }); + } + } + } + SchedulerEventJson::Resume { node, arg, event } => { + waiting = waiting.saturating_sub(1); + active += 1; + let snapshot = state_for_spawn + .lock() + .expect("state mutex poisoned") + .clone(); + spawn_json_node_task( + node, + wrap_resume_arg(arg, event), + snapshot, + tx_for_spawn.clone(), + Arc::clone(&callback), + )?; + } + SchedulerEventJson::WaitError(e) => return Err(e), + } + } + let final_state = state.lock().expect("state mutex poisoned").clone(); + Ok(final_state) +} + +fn wrap_resume_arg(arg: Value, event: WaitEvent) -> Value { + serde_json::json!({ + "__lg_resume_arg__": arg, + "__lg_resume_event__": event, + }) +} diff --git a/rust-core/src/lib.rs b/rust-core/src/lib.rs index 7bffe5eac..29b471b11 100644 --- a/rust-core/src/lib.rs +++ b/rust-core/src/lib.rs @@ -1,4 +1,2 @@ mod engine; mod lib_c; -#[cfg(feature = "python-bindings")] -mod lib_py; diff --git a/rust-core/src/lib_c.rs b/rust-core/src/lib_c.rs index b68bd08d7..42597cfe4 100644 --- a/rust-core/src/lib_c.rs +++ b/rust-core/src/lib_c.rs @@ -1,14 +1,12 @@ use crate::engine::{ - merge_json_update, node_pool_execute, run_loop_block_on, run_loop_spawn, AnyOfCondition, - Engine, NodeExecResult, NodeOutcome, SendPayload, WaitEvent, WaitRequest, + run_graph_json_with_callback, run_loop_block_on, run_loop_spawn, AnyOfCondition, Engine, + NodeExecResult, NodeOutcome, SendPayload, WaitRequest, }; use serde::Deserialize; use serde_json::Value; use std::ffi::{CStr, CString}; use std::os::raw::c_char; use std::sync::mpsc; -use std::sync::{Arc, Mutex}; -use tokio::sync::mpsc as tokio_mpsc; #[derive(Debug, Deserialize)] struct SendPayloadJson { @@ -94,186 +92,6 @@ fn parse_c_callback_result( })) } -enum SchedulerEventJson { - Node(Result), - Resume { - node: String, - arg: Value, - event: WaitEvent, - }, - WaitError(String), -} - -struct NodeExecutionJson { - node: String, - arg: Value, - outcome: NodeOutcome, -} - -fn spawn_json_node_task( - node: String, - arg: Value, - state_snapshot: Value, - tx: tokio_mpsc::UnboundedSender, - user_data_bits: libc::c_ulong, - callback: CNodeCallback, -) -> Result<(), String> { - node_pool_execute(move || { - let node_for_result = node.clone(); - let arg_for_result = arg.clone(); - let result = (|| -> Result { - let node_c = - CString::new(node.clone()).map_err(|e| format!("invalid node name: {e}"))?; - let arg_json = serde_json::to_string(&arg) - .map_err(|e| format!("serialize arg for `{node}` failed: {e}"))?; - let state_json = serde_json::to_string(&state_snapshot) - .map_err(|e| format!("serialize state for `{node}` failed: {e}"))?; - let arg_c = - CString::new(arg_json).map_err(|e| format!("invalid arg JSON bytes: {e}"))?; - let state_c = - CString::new(state_json).map_err(|e| format!("invalid state JSON bytes: {e}"))?; - let out_ptr = unsafe { - callback( - user_data_bits, - node_c.as_ptr() as *mut c_char, - arg_c.as_ptr() as *mut c_char, - state_c.as_ptr() as *mut c_char, - ) - }; - if out_ptr.is_null() { - return Err(format!("callback returned null for `{node}`")); - } - let out_raw = unsafe { CStr::from_ptr(out_ptr) } - .to_string_lossy() - .into_owned(); - unsafe { - libc::free(out_ptr.cast()); - } - let payload = parse_c_callback_result(out_raw, &node)?; - Ok(NodeExecutionJson { - node: node_for_result, - arg: arg_for_result, - outcome: payload, - }) - })(); - let _ = tx.send(SchedulerEventJson::Node(result)); - }) -} - -async fn run_graph_scheduler_json( - entry_point: String, - finish_point: String, - initial_state: Value, - initial_input: Value, - engine: Engine, - user_data: CUserData, - callback: CNodeCallback, -) -> Result { - let (tx, mut rx) = tokio_mpsc::unbounded_channel::(); - let state = Arc::new(Mutex::new(initial_state)); - let user_data_bits = user_data.0; - let tx_for_spawn = tx.clone(); - let state_for_spawn = Arc::clone(&state); - let state_for_merge = Arc::clone(&state); - let mut active: usize = 1; - let mut waiting: usize = 0; - spawn_json_node_task( - entry_point, - initial_input, - state_for_spawn - .lock() - .expect("state mutex poisoned") - .clone(), - tx_for_spawn.clone(), - user_data_bits, - callback, - )?; - while active > 0 || waiting > 0 { - let evt = rx - .recv() - .await - .ok_or_else(|| "scheduler event channel closed".to_string())?; - match evt { - SchedulerEventJson::Node(result) => { - active = active.saturating_sub(1); - let exec = result?; - match exec.outcome { - NodeOutcome::Completed(node_result) => { - let mut guard = state_for_merge.lock().expect("state mutex poisoned"); - merge_json_update(&mut guard, node_result.update); - drop(guard); - if exec.node == finish_point { - break; - } - for send in node_result.sends { - active += 1; - let snapshot = state_for_spawn - .lock() - .expect("state mutex poisoned") - .clone(); - spawn_json_node_task( - send.node, - send.arg, - snapshot, - tx_for_spawn.clone(), - user_data_bits, - callback, - )?; - } - } - NodeOutcome::Suspended { wait } => { - waiting += 1; - let tx_wait = tx_for_spawn.clone(); - let node = exec.node; - let arg = exec.arg; - let engine_for_wait = engine.clone(); - tokio::spawn(async move { - match engine_for_wait.wait_request_async(&wait).await { - Ok(event) => { - let _ = tx_wait.send(SchedulerEventJson::Resume { - node, - arg, - event, - }); - } - Err(e) => { - let _ = tx_wait.send(SchedulerEventJson::WaitError(e)); - } - } - }); - } - } - } - SchedulerEventJson::Resume { node, arg, event } => { - waiting = waiting.saturating_sub(1); - active += 1; - let snapshot = state_for_spawn - .lock() - .expect("state mutex poisoned") - .clone(); - spawn_json_node_task( - node, - wrap_resume_arg(arg, event), - snapshot, - tx_for_spawn.clone(), - user_data_bits, - callback, - )?; - } - SchedulerEventJson::WaitError(e) => return Err(e), - } - } - let final_state = state.lock().expect("state mutex poisoned").clone(); - Ok(final_state) -} - -fn wrap_resume_arg(arg: Value, event: WaitEvent) -> Value { - serde_json::json!({ - "__lg_resume_arg__": arg, - "__lg_resume_event__": event, - }) -} - #[no_mangle] pub extern "C" fn rc_engine_new() -> *mut Engine { Box::into_raw(Box::new(Engine::new())) @@ -536,14 +354,46 @@ pub unsafe extern "C" fn rc_run_graph_json( let user_data = CUserData(user_data); let run_engine = (*ptr).clone(); let submit = run_loop_spawn(async move { - let out = run_graph_scheduler_json( + let callback_wrapper = move |node: String, + arg: Value, + state_snapshot: Value| + -> Result, String> { + let node_c = + CString::new(node.clone()).map_err(|e| format!("invalid node name: {e}"))?; + let arg_json = serde_json::to_string(&arg) + .map_err(|e| format!("serialize arg for `{node}` failed: {e}"))?; + let state_json = serde_json::to_string(&state_snapshot) + .map_err(|e| format!("serialize state for `{node}` failed: {e}"))?; + let arg_c = + CString::new(arg_json).map_err(|e| format!("invalid arg JSON bytes: {e}"))?; + let state_c = + CString::new(state_json).map_err(|e| format!("invalid state JSON bytes: {e}"))?; + let out_ptr = unsafe { + callback( + user_data.0, + node_c.as_ptr() as *mut c_char, + arg_c.as_ptr() as *mut c_char, + state_c.as_ptr() as *mut c_char, + ) + }; + if out_ptr.is_null() { + return Err(format!("callback returned null for `{node}`")); + } + let out_raw = unsafe { CStr::from_ptr(out_ptr) } + .to_string_lossy() + .into_owned(); + unsafe { + libc::free(out_ptr.cast()); + } + parse_c_callback_result(out_raw, &node) + }; + let out = run_graph_json_with_callback( entry_point, finish_point, initial_state, initial_input, run_engine.clone(), - user_data, - callback, + callback_wrapper, ) .await; run_engine.close_stream(); diff --git a/rust-core/src/lib_py.rs b/rust-core/src/lib_py.rs deleted file mode 100644 index d93d23a67..000000000 --- a/rust-core/src/lib_py.rs +++ /dev/null @@ -1,461 +0,0 @@ -#[cfg(feature = "python-bindings")] -use crate::engine::{ - node_pool_execute, run_loop_block_on, run_loop_spawn, AnyOfCondition, Engine, NodeExecResult, - NodeOutcome, SendPayload, WaitCondition, WaitEvent, WaitRequest, -}; -#[cfg(feature = "python-bindings")] -use pyo3::exceptions::PyValueError; -#[cfg(feature = "python-bindings")] -use pyo3::prelude::*; -#[cfg(feature = "python-bindings")] -use pyo3::types::PyAny; -#[cfg(feature = "python-bindings")] -use pyo3::types::{PyDict, PyList, PyTuple}; -#[cfg(feature = "python-bindings")] -use serde_json::Value; -#[cfg(feature = "python-bindings")] -use std::sync::mpsc; -#[cfg(feature = "python-bindings")] -use std::sync::Arc; -#[cfg(feature = "python-bindings")] -use tokio::sync::mpsc as tokio_mpsc; - -#[cfg(feature = "python-bindings")] -#[pyclass] -struct PyRustEngine { - inner: Engine, -} - -#[cfg(feature = "python-bindings")] -#[pymethods] -impl PyRustEngine { - #[new] - fn new() -> Self { - Self { - inner: Engine::new(), - } - } - - fn add_async_channel(&self, name: &str) { - self.inner.add_async_channel(name); - } - - fn publish_json(&self, channel: &str, value_json: &str) -> PyResult<()> { - let value: Value = serde_json::from_str(value_json) - .map_err(|e| PyValueError::new_err(format!("Invalid JSON value: {e}")))?; - self.inner - .publish_json(channel, value) - .map_err(PyValueError::new_err) - } - - fn publish_obj(&self, py: Python<'_>, channel: &str, value: Py) -> PyResult<()> { - let value_json = py_obj_to_json_string(py, &value.bind(py))?; - let parsed: Value = serde_json::from_str(&value_json) - .map_err(|e| PyValueError::new_err(format!("Invalid Python JSON value: {e}")))?; - self.inner - .publish_json(channel, parsed) - .map_err(PyValueError::new_err) - } - - fn wait_any_of_json(&self, any_of_json: &str) -> PyResult { - let any_of: AnyOfCondition = serde_json::from_str(any_of_json) - .map_err(|e| PyValueError::new_err(format!("Invalid any_of JSON: {e}")))?; - let event = run_loop_block_on(self.inner.wait_for_any_of_async(&any_of)) - .map_err(PyValueError::new_err)?; - serde_json::to_string(&event) - .map_err(|e| PyValueError::new_err(format!("Serialize event failed: {e}"))) - } - - fn wait_channel(&self, py: Python<'_>, channel: &str, n: usize) -> PyResult> { - let cond = WaitCondition::Channel { - channel: channel.to_string(), - n, - }; - let event = - run_loop_block_on(self.inner.wait_for_async(&cond)).map_err(PyValueError::new_err)?; - let event_json = serde_json::to_string(&event) - .map_err(|e| PyValueError::new_err(format!("Serialize event failed: {e}")))?; - json_string_to_py_obj(py, &event_json) - } - - fn wait_timer(&self, py: Python<'_>, seconds: f64) -> PyResult> { - let cond = WaitCondition::Timer { seconds }; - let event = - run_loop_block_on(self.inner.wait_for_async(&cond)).map_err(PyValueError::new_err)?; - let event_json = serde_json::to_string(&event) - .map_err(|e| PyValueError::new_err(format!("Serialize event failed: {e}")))?; - json_string_to_py_obj(py, &event_json) - } - - fn wait_any_of_obj(&self, py: Python<'_>, any_of_payload: Py) -> PyResult> { - let payload_json = py_obj_to_json_string(py, &any_of_payload.bind(py))?; - let any_of: AnyOfCondition = serde_json::from_str(&payload_json) - .map_err(|e| PyValueError::new_err(format!("Invalid any_of payload: {e}")))?; - let event = run_loop_block_on(self.inner.wait_for_any_of_async(&any_of)) - .map_err(PyValueError::new_err)?; - let event_json = serde_json::to_string(&event) - .map_err(|e| PyValueError::new_err(format!("Serialize event failed: {e}")))?; - json_string_to_py_obj(py, &event_json) - } - - fn wait_condition_json(&self, cond_json: &str) -> PyResult { - let cond: WaitCondition = serde_json::from_str(cond_json) - .map_err(|e| PyValueError::new_err(format!("Invalid condition JSON: {e}")))?; - let event = - run_loop_block_on(self.inner.wait_for_async(&cond)).map_err(PyValueError::new_err)?; - serde_json::to_string(&event) - .map_err(|e| PyValueError::new_err(format!("Serialize event failed: {e}"))) - } - - fn start_stream(&self, stream_mode: Option<&str>) -> PyResult<()> { - self.inner - .start_stream(stream_mode) - .map_err(PyValueError::new_err) - } - - fn receive_stream_obj(&self, py: Python<'_>) -> PyResult> { - match run_loop_block_on(self.inner.receive_stream_async()) { - Some(value) => { - let event_json = serde_json::to_string(&value) - .map_err(|e| PyValueError::new_err(format!("Serialize event failed: {e}")))?; - json_string_to_py_obj(py, &event_json) - } - None => Ok(py.None()), - } - } - - fn send_custom_stream_event_obj(&self, py: Python<'_>, value: Py) -> PyResult<()> { - let value_json = py_obj_to_json_string(py, &value.bind(py))?; - let parsed: Value = serde_json::from_str(&value_json) - .map_err(|e| PyValueError::new_err(format!("Invalid Python JSON value: {e}")))?; - self.inner.send_custom_stream_event(parsed); - Ok(()) - } - - fn close_stream(&self) { - self.inner.close_stream(); - } - - fn run_graph_py( - &self, - py: Python<'_>, - entry_point: &str, - finish_point: &str, - initial_state: Py, - callback: Py, - stream_mode: Option<&str>, - ) -> PyResult> { - let state = Arc::new(initial_state); - let callback = Arc::new(callback); - let entry_point = entry_point.to_string(); - let finish_point = finish_point.to_string(); - self.inner - .start_stream(stream_mode) - .map_err(PyValueError::new_err)?; - let (done_tx, done_rx) = mpsc::channel::>(); - let state_for_run = Arc::clone(&state); - let engine = self.inner.clone(); - run_loop_spawn(async move { - let run_result = run_graph_scheduler( - entry_point, - finish_point, - callback, - state_for_run, - engine.clone(), - ) - .await; - engine.close_stream(); - let _ = done_tx.send(run_result); - }) - .map_err(PyValueError::new_err)?; - let run_result = py - .allow_threads(move || done_rx.recv()) - .map_err(|e| PyValueError::new_err(format!("run-loop recv failed: {e}")))?; - run_result.map_err(PyValueError::new_err)?; - Ok((*state).clone_ref(py)) - } -} - -#[cfg(feature = "python-bindings")] -enum SchedulerEventPy { - Node(Result), - Resume { - node: String, - arg: Py, - event: WaitEvent, - }, - WaitError(String), -} - -struct NodeExecutionPy { - node: String, - arg: Py, - outcome: NodeOutcome, Py>, -} - -#[cfg(feature = "python-bindings")] -async fn run_graph_scheduler( - entry_point: String, - finish_point: String, - callback: Arc>, - state: Arc>, - engine: Engine, -) -> Result<(), String> { - let (tx, mut rx) = tokio_mpsc::unbounded_channel::(); - let initial_arg = Python::with_gil(|py| (*state).clone_ref(py)); - let callback_for_spawn = Arc::clone(&callback); - let state_for_spawn = Arc::clone(&state); - let tx_for_spawn = tx.clone(); - let state_for_merge = Arc::clone(&state); - let mut active: usize = 1; - let mut waiting: usize = 0; - spawn_node_task( - entry_point, - initial_arg, - tx_for_spawn.clone(), - Arc::clone(&callback_for_spawn), - Arc::clone(&state_for_spawn), - )?; - - while active > 0 || waiting > 0 { - let event = rx - .recv() - .await - .ok_or_else(|| "scheduler event channel closed".to_string())?; - match event { - SchedulerEventPy::Node(result) => { - active = active.saturating_sub(1); - let exec = result?; - match exec.outcome { - NodeOutcome::Completed(node_result) => { - Python::with_gil(|py| -> Result<(), String> { - if let Some(update) = node_result.update { - apply_update_to_state(py, state_for_merge.as_ref(), &update) - .map_err(|e| { - format!("state merge failed for `{}`: {e}", exec.node) - })?; - } - Ok(()) - })?; - if exec.node == finish_point { - break; - } - for send in node_result.sends { - active += 1; - spawn_node_task( - send.node, - send.arg, - tx_for_spawn.clone(), - Arc::clone(&callback_for_spawn), - Arc::clone(&state_for_spawn), - )?; - } - } - NodeOutcome::Suspended { wait } => { - waiting += 1; - let tx_wait = tx_for_spawn.clone(); - let node = exec.node; - let arg = exec.arg; - let engine_for_wait = engine.clone(); - tokio::spawn(async move { - let outcome = engine_for_wait.wait_request_async(&wait).await; - match outcome { - Ok(event) => { - let _ = - tx_wait.send(SchedulerEventPy::Resume { node, arg, event }); - } - Err(e) => { - let _ = tx_wait.send(SchedulerEventPy::WaitError(e)); - } - } - }); - } - } - } - SchedulerEventPy::Resume { node, arg, event } => { - waiting = waiting.saturating_sub(1); - active += 1; - let resume_arg = wrap_resume_arg(&arg, &event)?; - spawn_node_task( - node, - resume_arg, - tx_for_spawn.clone(), - Arc::clone(&callback_for_spawn), - Arc::clone(&state_for_spawn), - )?; - } - SchedulerEventPy::WaitError(e) => return Err(e), - } - } - Ok(()) -} - -#[cfg(feature = "python-bindings")] -fn spawn_node_task( - node: String, - arg: Py, - tx: tokio_mpsc::UnboundedSender, - callback: Arc>, - state_for_task: Arc>, -) -> Result<(), String> { - node_pool_execute(move || { - let node_for_result = node.clone(); - let arg_for_result = Python::with_gil(|py| arg.clone_ref(py)); - let outcome = Python::with_gil(|py| -> Result { - let callback_bound = callback.as_ref().bind(py); - let payload_obj = callback_bound - .call1((node.as_str(), arg, (*state_for_task).clone_ref(py))) - .map_err(|e| format!("callback failed for node `{node}`: {e}"))?; - let payload = parse_node_outcome(py, &payload_obj) - .map_err(|e| format!("invalid callback payload for `{node}`: {e}"))?; - Ok(NodeExecutionPy { - node: node_for_result, - arg: arg_for_result, - outcome: payload, - }) - }); - let _ = tx.send(SchedulerEventPy::Node(outcome)); - }) -} - -#[cfg(feature = "python-bindings")] -fn parse_node_outcome( - py: Python<'_>, - payload_obj: &Bound<'_, PyAny>, -) -> Result, Py>, String> { - let payload_dict = payload_obj - .downcast::() - .map_err(|_| "payload must be a dict".to_string())?; - let suspended_item = payload_dict - .get_item("suspend") - .map_err(|e| format!("failed to read suspend: {e}"))?; - if let Some(wait_obj) = suspended_item { - let wait_json = py_obj_to_json_string(py, &wait_obj) - .map_err(|e| format!("failed to encode suspend payload: {e}"))?; - let wait: WaitRequest = serde_json::from_str(&wait_json) - .map_err(|e| format!("invalid suspend payload: {e}"))?; - return Ok(NodeOutcome::Suspended { wait }); - } - - let update_item = payload_dict - .get_item("update") - .map_err(|e| format!("failed to read update: {e}"))?; - let update = match update_item { - Some(v) if !v.is_none() => Some(v.unbind()), - _ => None, - }; - - let sends_obj = payload_dict - .get_item("sends") - .map_err(|e| format!("failed to read sends: {e}"))? - .ok_or_else(|| "missing sends".to_string())?; - let sends_list = sends_obj - .downcast::() - .map_err(|_| "sends must be a list".to_string())?; - - let mut sends = Vec::with_capacity(sends_list.len()); - for item in sends_list.iter() { - let send_dict = item - .downcast::() - .map_err(|_| "send item must be a dict".to_string())?; - let node_obj = send_dict - .get_item("node") - .map_err(|e| format!("failed to read send.node: {e}"))? - .ok_or_else(|| "send.node is required".to_string())?; - let node = node_obj - .extract::() - .map_err(|e| format!("send.node must be string: {e}"))?; - let arg: Py = match send_dict.get_item("arg") { - Ok(Some(v)) => v.unbind(), - Ok(None) => Python::with_gil(|py| py.None()), - Err(e) => return Err(format!("failed to read send.arg: {e}")), - }; - sends.push(SendPayload { node, arg }); - } - - Ok(NodeOutcome::Completed(NodeExecResult { update, sends })) -} - -#[cfg(feature = "python-bindings")] -fn wrap_resume_arg(arg: &Py, event: &WaitEvent) -> Result, String> { - Python::with_gil(|py| -> Result, String> { - let wrapper = PyDict::new(py); - wrapper - .set_item("__lg_resume_arg__", arg.clone_ref(py)) - .map_err(|e| format!("failed to set resume arg: {e}"))?; - let event_json = serde_json::to_string(event) - .map_err(|e| format!("failed to encode wait event: {e}"))?; - let event_obj = json_string_to_py_obj(py, &event_json) - .map_err(|e| format!("failed to parse event: {e}"))?; - wrapper - .set_item("__lg_resume_event__", event_obj.bind(py)) - .map_err(|e| format!("failed to set resume event: {e}"))?; - Ok(wrapper.unbind().into_any()) - }) -} - -#[cfg(feature = "python-bindings")] -fn apply_update_to_state(py: Python<'_>, state: &Py, update: &Py) -> PyResult<()> { - let state_obj = state.bind(py); - let update_obj = update.bind(py); - - if update_obj.is_none() { - return Ok(()); - } - if state_obj.is_instance_of::() && update_obj.is_instance_of::() { - let state_dict = state_obj.downcast::()?; - let update_dict = update_obj.downcast::()?; - state_dict.call_method1("update", (update_dict,))?; - return Ok(()); - } - if let Ok(tuple_like) = update_obj.downcast::() { - apply_pair_updates(state_obj, tuple_like)?; - return Ok(()); - } - if let Ok(tuple_like) = update_obj.downcast::() { - let list = PyList::new(py, tuple_like)?; - apply_pair_updates(state_obj, &list)?; - } - Ok(()) -} - -#[cfg(feature = "python-bindings")] -fn apply_pair_updates(state_obj: &Bound<'_, PyAny>, entries: &Bound<'_, PyList>) -> PyResult<()> { - if !state_obj.is_instance_of::() { - return Ok(()); - } - let state_dict = state_obj.downcast::()?; - for entry in entries.iter() { - if let Ok(pair) = entry.downcast::() { - if pair.len() == 2 { - let key_obj = pair.get_item(0)?; - if let Ok(key) = key_obj.extract::() { - let value_obj = pair.get_item(1)?; - state_dict.set_item(key, value_obj)?; - } - } - } - } - Ok(()) -} - -#[cfg(feature = "python-bindings")] -fn py_obj_to_json_string(py: Python<'_>, obj: &Bound<'_, PyAny>) -> PyResult { - let json_mod = py.import("json")?; - let dumped = json_mod.call_method1("dumps", (obj,))?; - dumped.extract::() -} - -#[cfg(feature = "python-bindings")] -fn json_string_to_py_obj(py: Python<'_>, value: &str) -> PyResult> { - let json_mod = py.import("json")?; - let loaded = json_mod.call_method1("loads", (value,))?; - Ok(loaded.unbind()) -} - -#[cfg(feature = "python-bindings")] -#[pymodule] -fn langgraph_rust_core(m: &Bound<'_, PyModule>) -> PyResult<()> { - m.add_class::()?; - Ok(()) -} diff --git a/saf-python-sdk/Cargo.lock b/saf-python-sdk/Cargo.lock index 082228bd7..8f3dadaf4 100644 --- a/saf-python-sdk/Cargo.lock +++ b/saf-python-sdk/Cargo.lock @@ -2,12 +2,6 @@ # It is not intended for manual editing. version = 4 -[[package]] -name = "autocfg" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" - [[package]] name = "bitflags" version = "2.11.0" @@ -20,21 +14,6 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" -[[package]] -name = "heck" -version = "0.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" - -[[package]] -name = "indoc" -version = "2.0.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "79cf5c93f93228cf8efb3ba362535fb11199ac548a09ce117c9b1adc3030d706" -dependencies = [ - "rustversion", -] - [[package]] name = "itoa" version = "1.0.17" @@ -47,7 +26,6 @@ version = "0.1.0" dependencies = [ "libc", "parking_lot", - "pyo3", "serde", "serde_json", "tokio", @@ -74,21 +52,6 @@ version = "2.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79" -[[package]] -name = "memoffset" -version = "0.9.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a" -dependencies = [ - "autocfg", -] - -[[package]] -name = "once_cell" -version = "1.21.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" - [[package]] name = "parking_lot" version = "0.12.5" @@ -118,12 +81,6 @@ version = "0.2.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" -[[package]] -name = "portable-atomic" -version = "1.13.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" - [[package]] name = "proc-macro2" version = "1.0.106" @@ -133,69 +90,6 @@ dependencies = [ "unicode-ident", ] -[[package]] -name = "pyo3" -version = "0.23.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7778bffd85cf38175ac1f545509665d0b9b92a198ca7941f131f85f7a4f9a872" -dependencies = [ - "cfg-if", - "indoc", - "libc", - "memoffset", - "once_cell", - "portable-atomic", - "pyo3-build-config", - "pyo3-ffi", - "pyo3-macros", - "unindent", -] - -[[package]] -name = "pyo3-build-config" -version = "0.23.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "94f6cbe86ef3bf18998d9df6e0f3fc1050a8c5efa409bf712e661a4366e010fb" -dependencies = [ - "once_cell", - "target-lexicon", -] - -[[package]] -name = "pyo3-ffi" -version = "0.23.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e9f1b4c431c0bb1c8fb0a338709859eed0d030ff6daa34368d3b152a63dfdd8d" -dependencies = [ - "libc", - "pyo3-build-config", -] - -[[package]] -name = "pyo3-macros" -version = "0.23.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fbc2201328f63c4710f68abdf653c89d8dbc2858b88c5d88b0ff38a75288a9da" -dependencies = [ - "proc-macro2", - "pyo3-macros-backend", - "quote", - "syn", -] - -[[package]] -name = "pyo3-macros-backend" -version = "0.23.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fca6726ad0f3da9c9de093d6f116a93c1a38e417ed73bf138472cf4064f72028" -dependencies = [ - "heck", - "proc-macro2", - "pyo3-build-config", - "quote", - "syn", -] - [[package]] name = "quote" version = "1.0.45" @@ -214,12 +108,6 @@ dependencies = [ "bitflags", ] -[[package]] -name = "rustversion" -version = "1.0.22" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" - [[package]] name = "scopeguard" version = "1.2.0" @@ -286,12 +174,6 @@ dependencies = [ "unicode-ident", ] -[[package]] -name = "target-lexicon" -version = "0.12.16" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1" - [[package]] name = "tokio" version = "1.50.0" @@ -319,12 +201,6 @@ version = "1.0.24" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" -[[package]] -name = "unindent" -version = "0.2.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3" - [[package]] name = "windows-link" version = "0.2.1" diff --git a/saf-python-sdk/Cargo.toml b/saf-python-sdk/Cargo.toml index d506f921d..fec8c74b3 100644 --- a/saf-python-sdk/Cargo.toml +++ b/saf-python-sdk/Cargo.toml @@ -8,12 +8,7 @@ name = "langgraph_rust_core" path = "../rust-core/src/lib.rs" crate-type = ["cdylib", "rlib"] -[features] -default = ["python-bindings"] -python-bindings = ["dep:pyo3"] - [dependencies] -pyo3 = { version = "0.23.5", features = ["extension-module"], optional = true } serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" parking_lot = "0.12" diff --git a/saf-python-sdk/README.md b/saf-python-sdk/README.md index 45ada8435..e7d50f8fb 100644 --- a/saf-python-sdk/README.md +++ b/saf-python-sdk/README.md @@ -5,7 +5,7 @@ Standalone Python SDK for the `advanced_graph` runtime backed by the Rust engine This package intentionally contains only: - `saf_python_sdk.advanced_graph` (Python API) -- `langgraph_rust_core` (Rust execution engine via PyO3) +- Rust core engine via C bindings (`ctypes`) It does not package the original `langgraph` `stategraph` stack. diff --git a/saf-python-sdk/python/saf_python_sdk/advanced_graph/state.py b/saf-python-sdk/python/saf_python_sdk/advanced_graph/state.py index b89985081..a32c5c697 100644 --- a/saf-python-sdk/python/saf_python_sdk/advanced_graph/state.py +++ b/saf-python-sdk/python/saf_python_sdk/advanced_graph/state.py @@ -11,7 +11,7 @@ from dataclasses import dataclass from datetime import timedelta from typing import Any, Generic, TypeVar, cast -from saf_python_sdk.langgraph_rust_core import PyRustEngine # type: ignore[import-untyped] +from saf_python_sdk.rust_core_cffi import PyRustEngine from saf_python_sdk.types import Command, Send @@ -211,11 +211,15 @@ class GraphRunHandler(Generic[StateT]): await self._run.publish(channel, value) async def receive_stream(self) -> Any | None: - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - _advanced_graph_executor(), - self._run.receive_stream_sync, - ) + while True: + loop = asyncio.get_running_loop() + event = await loop.run_in_executor( + _advanced_graph_executor(), + self._run.receive_stream_sync, + ) + if event is not None or self._task.done(): + return event + await asyncio.sleep(0.005) def close_stream(self) -> None: self._run.close_stream_sync() diff --git a/saf-python-sdk/python/saf_python_sdk/rust_core_cffi.py b/saf-python-sdk/python/saf_python_sdk/rust_core_cffi.py new file mode 100644 index 000000000..7e53332ea --- /dev/null +++ b/saf-python-sdk/python/saf_python_sdk/rust_core_cffi.py @@ -0,0 +1,347 @@ +from __future__ import annotations + +import ctypes +import dataclasses +import json +import subprocess +from copy import deepcopy +from pathlib import Path +import threading +from typing import Any, Callable, get_args, get_origin + + +class PyRustEngine: + def __init__(self) -> None: + self._lib = _load_rust_lib() + self._engine = self._lib.rc_engine_new() + if not self._engine: + raise RuntimeError("failed to create rust engine") + + def __del__(self) -> None: + engine = getattr(self, "_engine", None) + if engine: + self._lib.rc_engine_free(engine) + self._engine = None + + def add_async_channel(self, name: str) -> None: + self._call_status(self._lib.rc_add_async_channel, name.encode()) + + def publish_obj(self, channel: str, value: Any) -> None: + payload = json.dumps(value, ensure_ascii=False).encode() + self._call_status(self._lib.rc_publish_json, channel.encode(), payload) + + def wait_any_of_obj(self, any_of_payload: Any) -> Any: + payload = json.dumps(any_of_payload, ensure_ascii=False).encode() + raw = self._consume_json_ptr(self._lib.rc_wait_any_of_json(self._engine, payload)) + if not raw.get("ok"): + raise ValueError(raw.get("error", "rust wait_any_of failed")) + return raw["event"] + + def wait_channel(self, channel: str, n: int) -> Any: + return self.wait_any_of_obj({"conditions": [{"kind": "channel", "channel": channel, "n": n}]}) + + def wait_timer(self, seconds: float) -> Any: + return self.wait_any_of_obj({"conditions": [{"kind": "timer", "seconds": seconds}]}) + + def start_stream(self, stream_mode: str | None) -> None: + encoded = stream_mode.encode() if stream_mode is not None else None + self._call_status(self._lib.rc_start_stream, encoded) + + def receive_stream_obj(self) -> Any | None: + raw = self._consume_json_ptr(self._lib.rc_receive_stream_json(self._engine)) + if not raw.get("ok"): + raise ValueError(raw.get("error", "rust receive_stream failed")) + if not raw.get("has_event", False): + return None + return raw.get("event") + + def send_custom_stream_event_obj(self, value: Any) -> None: + payload = json.dumps(value, ensure_ascii=False).encode() + self._call_status(self._lib.rc_send_custom_stream_event, payload) + + def close_stream(self) -> None: + self._call_status(self._lib.rc_close_stream) + + def run_graph_py( + self, + entry_point: str, + finish_point: str, + initial_state: Any, + callback: Callable[[str, Any, Any], dict[str, Any]], + stream_mode: str | None = None, + ) -> Any: + shared_state = initial_state + shared_state_lock = threading.Lock() + state_type = type(initial_state) + use_shared_state = dataclasses.is_dataclass(initial_state) + initial_state_json = json.dumps(_to_jsonable(shared_state), ensure_ascii=False).encode() + callback_c = _make_node_callback( + callback, + state_type, + use_shared_state, + shared_state, + shared_state_lock, + ) + stream_mode_encoded = stream_mode.encode() if stream_mode is not None else None + out = self._consume_json_ptr( + self._lib.rc_run_graph_json( + self._engine, + entry_point.encode(), + finish_point.encode(), + initial_state_json, + initial_state_json, + stream_mode_encoded, + ctypes.c_ulong(0), + callback_c, + ) + ) + if not out.get("ok"): + raise ValueError(out.get("error", "rust run_graph failed")) + if use_shared_state: + return shared_state + return _coerce_for_type(out["state"], state_type) + + def _call_status(self, func: Any, *args: Any) -> None: + raw = self._consume_json_ptr(func(self._engine, *args)) + if not raw.get("ok"): + raise ValueError(raw.get("error", "rust call failed")) + + def _consume_json_ptr(self, ptr: ctypes.c_void_p) -> dict[str, Any]: + if not ptr: + raise RuntimeError("rust returned null string pointer") + try: + text = ctypes.cast(ptr, ctypes.c_char_p).value + if text is None: + raise RuntimeError("rust returned empty string pointer") + return json.loads(text.decode()) + finally: + self._lib.rc_string_free(ptr) + + +def _make_node_callback( + callback: Callable[[str, Any, Any], dict[str, Any]], + state_type: type[Any], + use_shared_state: bool, + shared_state: Any, + shared_state_lock: threading.Lock, +) -> ctypes.CFUNCTYPE: # type: ignore[type-arg] + cb_type = ctypes.CFUNCTYPE( + ctypes.c_void_p, + ctypes.c_ulong, + ctypes.c_char_p, + ctypes.c_char_p, + ctypes.c_char_p, + ) + libc = ctypes.CDLL(None) + libc.malloc.argtypes = [ctypes.c_size_t] + libc.malloc.restype = ctypes.c_void_p + + @cb_type + def _callback( + _user_data: int, + node_ptr: bytes, + arg_ptr: bytes, + state_ptr: bytes, + ) -> ctypes.c_void_p: + try: + node = node_ptr.decode() + arg = json.loads(arg_ptr.decode()) + if use_shared_state: + with shared_state_lock: + before = deepcopy(_to_jsonable(shared_state)) + result = callback(node, arg, shared_state) + state_after_call = shared_state + else: + state_raw = json.loads(state_ptr.decode()) + state_snapshot = _coerce_for_type(state_raw, state_type) + before = deepcopy(_to_jsonable(state_snapshot)) + result = callback(node, arg, state_snapshot) + state_after_call = state_snapshot + if "suspend" in result: + envelope = {"ok": True, "suspend": result["suspend"]} + else: + update = result.get("update") + if update is not None: + update_json = _to_jsonable(update) + if update_json == before: + update = None + update_json = None + else: + if use_shared_state: + _apply_update_to_state(shared_state, update) + after = _to_jsonable(state_after_call) + if update is None and after != before: + update_json = after + elif update is None: + update_json = None + else: + update_json = _to_jsonable(update) + sends = [] + for item in result.get("sends", []): + if not isinstance(item, dict): + continue + sends.append( + { + "node": item.get("node"), + "arg": _to_jsonable(item.get("arg")), + } + ) + envelope = { + "ok": True, + "payload": { + "update": update_json, + "sends": sends, + }, + } + except Exception as exc: # noqa: BLE001 + envelope = {"ok": False, "error": f"python callback failed: {exc}"} + return _malloc_c_string(json.dumps(envelope, ensure_ascii=False).encode(), libc) + + return _callback + + +def _malloc_c_string(payload: bytes, libc: Any) -> ctypes.c_void_p: + size = len(payload) + 1 + ptr = libc.malloc(size) + if not ptr: + return ctypes.c_void_p(0) + ctypes.memmove(ptr, payload, len(payload)) + ctypes.memset(ctypes.c_void_p(ptr + len(payload)), 0, 1) + return ptr + + +def _load_rust_lib() -> ctypes.CDLL: + env = Path.cwd() + root = _find_repo_root(env) + rust_core = root / "rust-core" + lib_path = _resolve_lib_path(rust_core) + if not lib_path.exists(): + subprocess.run(["cargo", "build"], cwd=rust_core, check=True) + lib = ctypes.CDLL(str(lib_path)) + _configure_signatures(lib) + return lib + + +def _find_repo_root(start: Path) -> Path: + current = start.resolve() + for candidate in [current, *current.parents]: + if (candidate / "rust-core").exists() and (candidate / "saf-python-sdk").exists(): + return candidate + here = Path(__file__).resolve() + return here.parents[3] + + +def _resolve_lib_path(rust_core: Path) -> Path: + if (rust_core / "target" / "debug" / "liblanggraph_rust_core.dylib").exists(): + return rust_core / "target" / "debug" / "liblanggraph_rust_core.dylib" + if (rust_core / "target" / "debug" / "liblanggraph_rust_core.so").exists(): + return rust_core / "target" / "debug" / "liblanggraph_rust_core.so" + if (rust_core / "target" / "debug" / "langgraph_rust_core.dll").exists(): + return rust_core / "target" / "debug" / "langgraph_rust_core.dll" + return rust_core / "target" / "debug" / "liblanggraph_rust_core.dylib" + + +def _configure_signatures(lib: ctypes.CDLL) -> None: + cb_type = ctypes.CFUNCTYPE( + ctypes.c_void_p, + ctypes.c_ulong, + ctypes.c_char_p, + ctypes.c_char_p, + ctypes.c_char_p, + ) + lib.rc_engine_new.argtypes = [] + lib.rc_engine_new.restype = ctypes.c_void_p + lib.rc_engine_free.argtypes = [ctypes.c_void_p] + lib.rc_engine_free.restype = None + lib.rc_string_free.argtypes = [ctypes.c_void_p] + lib.rc_string_free.restype = None + lib.rc_add_async_channel.argtypes = [ctypes.c_void_p, ctypes.c_char_p] + lib.rc_add_async_channel.restype = ctypes.c_void_p + lib.rc_publish_json.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_char_p] + lib.rc_publish_json.restype = ctypes.c_void_p + lib.rc_wait_any_of_json.argtypes = [ctypes.c_void_p, ctypes.c_char_p] + lib.rc_wait_any_of_json.restype = ctypes.c_void_p + lib.rc_start_stream.argtypes = [ctypes.c_void_p, ctypes.c_char_p] + lib.rc_start_stream.restype = ctypes.c_void_p + lib.rc_receive_stream_json.argtypes = [ctypes.c_void_p] + lib.rc_receive_stream_json.restype = ctypes.c_void_p + lib.rc_send_custom_stream_event.argtypes = [ctypes.c_void_p, ctypes.c_char_p] + lib.rc_send_custom_stream_event.restype = ctypes.c_void_p + lib.rc_close_stream.argtypes = [ctypes.c_void_p] + lib.rc_close_stream.restype = ctypes.c_void_p + lib.rc_run_graph_json.argtypes = [ + ctypes.c_void_p, + ctypes.c_char_p, + ctypes.c_char_p, + ctypes.c_char_p, + ctypes.c_char_p, + ctypes.c_char_p, + ctypes.c_ulong, + cb_type, + ] + lib.rc_run_graph_json.restype = ctypes.c_void_p + + +def _to_jsonable(value: Any) -> Any: + if value is None or isinstance(value, (str, int, float, bool)): + return value + if dataclasses.is_dataclass(value): + return {field.name: _to_jsonable(getattr(value, field.name)) for field in dataclasses.fields(value)} + model_dump = getattr(value, "model_dump", None) + if callable(model_dump): + return _to_jsonable(model_dump()) + if isinstance(value, dict): + return {str(k): _to_jsonable(v) for k, v in value.items()} + if isinstance(value, (list, tuple, set)): + return [_to_jsonable(v) for v in value] + return value + + +def _coerce_for_type(value: Any, typ: Any) -> Any: + if value is None: + return None + origin = get_origin(typ) + args = get_args(typ) + if origin is not None: + if origin in (list, tuple, set): + item_type = args[0] if args else Any + items = [_coerce_for_type(v, item_type) for v in value] + if origin is tuple: + return tuple(items) + if origin is set: + return set(items) + return items + if origin is dict: + value_type = args[1] if len(args) == 2 else Any + return {k: _coerce_for_type(v, value_type) for k, v in value.items()} + if isinstance(typ, type): + if dataclasses.is_dataclass(typ): + kwargs = {} + for field in dataclasses.fields(typ): + kwargs[field.name] = _coerce_for_type(value.get(field.name), field.type) + return typ(**kwargs) + model_validate = getattr(typ, "model_validate", None) + if callable(model_validate): + return model_validate(value) + return value + + +def _apply_update_to_state(state: Any, update: Any) -> None: + if update is None: + return + if isinstance(state, dict): + if isinstance(update, dict): + state.update(update) + return + if dataclasses.is_dataclass(update): + state.update(_to_jsonable(update)) + return + return + if dataclasses.is_dataclass(state): + if dataclasses.is_dataclass(update): + for field in dataclasses.fields(state): + setattr(state, field.name, getattr(update, field.name)) + return + if isinstance(update, dict): + for key, val in update.items(): + setattr(state, key, val) diff --git a/saf-python-sdk/uv.lock b/saf-python-sdk/uv.lock index ea0ad8d69..41e30f5f1 100644 --- a/saf-python-sdk/uv.lock +++ b/saf-python-sdk/uv.lock @@ -4,5 +4,5 @@ requires-python = ">=3.10" [[package]] name = "saf-python-sdk" -version = "0.1.1" +version = "0.1.2" source = { editable = "." }