This commit is contained in:
Quanzheng Long
2026-03-17 16:23:08 -07:00
parent e13004da77
commit 8f0f4a1be7
12 changed files with 554 additions and 916 deletions
-124
View File
@@ -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"
-5
View File
@@ -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"
+158
View File
@@ -279,6 +279,22 @@ pub fn merge_json_update(state: &mut Value, update: Option<Value>) {
}
}
enum SchedulerEventJson {
Node(Result<NodeExecutionJson, String>),
Resume {
node: String,
arg: Value,
event: WaitEvent,
},
WaitError(String),
}
struct NodeExecutionJson {
node: String,
arg: Value,
outcome: NodeOutcome<Value, Value>,
}
#[derive(Clone, Default)]
pub struct Engine {
channels: Arc<StdMutex<HashMap<String, VecDeque<serde_json::Value>>>>,
@@ -491,3 +507,145 @@ impl Engine {
}))
}
}
fn spawn_json_node_task<F>(
node: String,
arg: Value,
state_snapshot: Value,
tx: tokio_mpsc::UnboundedSender<SchedulerEventJson>,
callback: Arc<F>,
) -> Result<(), String>
where
F: Fn(String, Value, Value) -> Result<NodeOutcome<Value, Value>, 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<F>(
entry_point: String,
finish_point: String,
initial_state: Value,
initial_input: Value,
engine: Engine,
callback: F,
) -> Result<Value, String>
where
F: Fn(String, Value, Value) -> Result<NodeOutcome<Value, Value>, String>
+ Send
+ Sync
+ 'static,
{
let callback = Arc::new(callback);
let (tx, mut rx) = tokio_mpsc::unbounded_channel::<SchedulerEventJson>();
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,
})
}
-2
View File
@@ -1,4 +1,2 @@
mod engine;
mod lib_c;
#[cfg(feature = "python-bindings")]
mod lib_py;
+37 -187
View File
@@ -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<NodeExecutionJson, String>),
Resume {
node: String,
arg: Value,
event: WaitEvent,
},
WaitError(String),
}
struct NodeExecutionJson {
node: String,
arg: Value,
outcome: NodeOutcome<Value, Value>,
}
fn spawn_json_node_task(
node: String,
arg: Value,
state_snapshot: Value,
tx: tokio_mpsc::UnboundedSender<SchedulerEventJson>,
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<NodeExecutionJson, 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_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<Value, String> {
let (tx, mut rx) = tokio_mpsc::unbounded_channel::<SchedulerEventJson>();
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<NodeOutcome<Value, Value>, 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();
-461
View File
@@ -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<PyAny>) -> 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<String> {
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<Py<PyAny>> {
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<Py<PyAny>> {
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<PyAny>) -> PyResult<Py<PyAny>> {
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<String> {
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<Py<PyAny>> {
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<PyAny>) -> 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<PyAny>,
callback: Py<PyAny>,
stream_mode: Option<&str>,
) -> PyResult<Py<PyAny>> {
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::<Result<(), String>>();
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<NodeExecutionPy, String>),
Resume {
node: String,
arg: Py<PyAny>,
event: WaitEvent,
},
WaitError(String),
}
struct NodeExecutionPy {
node: String,
arg: Py<PyAny>,
outcome: NodeOutcome<Py<PyAny>, Py<PyAny>>,
}
#[cfg(feature = "python-bindings")]
async fn run_graph_scheduler(
entry_point: String,
finish_point: String,
callback: Arc<Py<PyAny>>,
state: Arc<Py<PyAny>>,
engine: Engine,
) -> Result<(), String> {
let (tx, mut rx) = tokio_mpsc::unbounded_channel::<SchedulerEventPy>();
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<PyAny>,
tx: tokio_mpsc::UnboundedSender<SchedulerEventPy>,
callback: Arc<Py<PyAny>>,
state_for_task: Arc<Py<PyAny>>,
) -> 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<NodeExecutionPy, String> {
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<NodeOutcome<Py<PyAny>, Py<PyAny>>, String> {
let payload_dict = payload_obj
.downcast::<PyDict>()
.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::<PyList>()
.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::<PyDict>()
.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::<String>()
.map_err(|e| format!("send.node must be string: {e}"))?;
let arg: Py<PyAny> = 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<PyAny>, event: &WaitEvent) -> Result<Py<PyAny>, String> {
Python::with_gil(|py| -> Result<Py<PyAny>, 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<PyAny>, update: &Py<PyAny>) -> 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::<PyDict>() && update_obj.is_instance_of::<PyDict>() {
let state_dict = state_obj.downcast::<PyDict>()?;
let update_dict = update_obj.downcast::<PyDict>()?;
state_dict.call_method1("update", (update_dict,))?;
return Ok(());
}
if let Ok(tuple_like) = update_obj.downcast::<PyList>() {
apply_pair_updates(state_obj, tuple_like)?;
return Ok(());
}
if let Ok(tuple_like) = update_obj.downcast::<PyTuple>() {
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::<PyDict>() {
return Ok(());
}
let state_dict = state_obj.downcast::<PyDict>()?;
for entry in entries.iter() {
if let Ok(pair) = entry.downcast::<PyTuple>() {
if pair.len() == 2 {
let key_obj = pair.get_item(0)?;
if let Ok(key) = key_obj.extract::<String>() {
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<String> {
let json_mod = py.import("json")?;
let dumped = json_mod.call_method1("dumps", (obj,))?;
dumped.extract::<String>()
}
#[cfg(feature = "python-bindings")]
fn json_string_to_py_obj(py: Python<'_>, value: &str) -> PyResult<Py<PyAny>> {
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::<PyRustEngine>()?;
Ok(())
}
-124
View File
@@ -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"
-5
View File
@@ -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"
+1 -1
View File
@@ -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.
@@ -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()
@@ -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)
+1 -1
View File
@@ -4,5 +4,5 @@ requires-python = ">=3.10"
[[package]]
name = "saf-python-sdk"
version = "0.1.1"
version = "0.1.2"
source = { editable = "." }