mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-27 01:52:25 +02:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
d26d4050c8 | ||
|
|
8f0f4a1be7 | ||
|
|
e13004da77 | ||
|
|
5a9264d124 | ||
|
|
efa86a6b14 | ||
|
|
f9866186d9 | ||
|
|
bfdd7deb60 | ||
|
|
09c8bb7c1c | ||
|
|
84d59adcd8 | ||
|
|
841ebf0c77 |
@@ -0,0 +1,6 @@
|
||||
PYTHON_VERSION ?= 3.13
|
||||
|
||||
.PHONY: run-basic
|
||||
run-basic:
|
||||
uv run --python $(PYTHON_VERSION) python examples/basic_run.py
|
||||
|
||||
@@ -0,0 +1,14 @@
|
||||
# advanced-graph-examples
|
||||
|
||||
Examples for the published `saf-python-sdk` package.
|
||||
|
||||
## Quick start
|
||||
|
||||
```bash
|
||||
cd advanced-graph-examples
|
||||
uv sync
|
||||
uv run python examples/basic_run.py
|
||||
```
|
||||
|
||||
> Note: current published wheel is for Python 3.13 on macOS arm64.
|
||||
|
||||
@@ -0,0 +1,36 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from typing import TypedDict
|
||||
|
||||
from saf_python_sdk import Command, Send
|
||||
from saf_python_sdk.advanced_graph import AdvancedStateGraph
|
||||
|
||||
|
||||
class MyState(TypedDict):
|
||||
count: int
|
||||
logs: list[str]
|
||||
|
||||
|
||||
async def start_node(state: MyState) -> Command:
|
||||
state["logs"].append("start")
|
||||
return Command(update=state, goto=Send("finish_node", "hello"))
|
||||
|
||||
|
||||
async def finish_node(input: str, state: MyState) -> Command:
|
||||
state["logs"].append(f"finish:{input}")
|
||||
state["count"] += 1
|
||||
return Command(update=state)
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
graph = AdvancedStateGraph(MyState)
|
||||
graph.add_entry_node(start_node)
|
||||
graph.add_finish_node(finish_node)
|
||||
result = await graph.compile().ainvoke({"count": 0, "logs": []})
|
||||
print(result)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
|
||||
@@ -0,0 +1,12 @@
|
||||
[project]
|
||||
name = "advanced-graph-examples"
|
||||
version = "0.1.0"
|
||||
requires-python = ">=3.13,<3.14"
|
||||
dependencies = [
|
||||
"saf-python-sdk>=0.1.1",
|
||||
]
|
||||
|
||||
[tool.uv]
|
||||
dev-dependencies = []
|
||||
|
||||
|
||||
Generated
+25
@@ -0,0 +1,25 @@
|
||||
version = 1
|
||||
revision = 3
|
||||
requires-python = "==3.13.*"
|
||||
|
||||
[[package]]
|
||||
name = "advanced-graph-examples"
|
||||
version = "0.1.0"
|
||||
source = { virtual = "." }
|
||||
dependencies = [
|
||||
{ name = "saf-python-sdk" },
|
||||
]
|
||||
|
||||
[package.metadata]
|
||||
requires-dist = [{ name = "saf-python-sdk", specifier = ">=0.1.1" }]
|
||||
|
||||
[package.metadata.requires-dev]
|
||||
dev = []
|
||||
|
||||
[[package]]
|
||||
name = "saf-python-sdk"
|
||||
version = "0.1.2"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/c5/43/2a86cdc8a1fbb185b9ce4642f0e3676acc168d48ed2e68d4ead6ad57e95e/saf_python_sdk-0.1.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:9035ab6ede9884b5588925c47d5b8be2c1c107becf5d6f88e0ebb39f4e8fac3b", size = 517053, upload-time = "2026-03-16T23:49:25.661Z" },
|
||||
]
|
||||
@@ -111,6 +111,10 @@ func (c *Context) PublishToChannel(channel string, value any) error {
|
||||
return c.engine.Publish(channel, value)
|
||||
}
|
||||
|
||||
func (c *Context) SendCustomStreamEvent(value any) error {
|
||||
return c.engine.SendCustomStreamEvent(value)
|
||||
}
|
||||
|
||||
type Handler[StateT any] struct {
|
||||
engine *RustEngine
|
||||
done chan resultOrErr[StateT]
|
||||
@@ -130,7 +134,29 @@ func (h *Handler[StateT]) WaitForResult() (StateT, error) {
|
||||
return res.state, res.err
|
||||
}
|
||||
|
||||
func (g *CompiledGraph[StateT]) Start(initialInput any, initialState StateT) (*Handler[StateT], error) {
|
||||
func (h *Handler[StateT]) ReceiveStream() (any, error) {
|
||||
event, hasEvent, err := h.engine.ReceiveStream()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !hasEvent {
|
||||
return nil, nil
|
||||
}
|
||||
return event, nil
|
||||
}
|
||||
|
||||
func (h *Handler[StateT]) CloseStream() error {
|
||||
return h.engine.CloseStream()
|
||||
}
|
||||
|
||||
func (g *CompiledGraph[StateT]) Start(initialInput any, initialState StateT, streamMode ...string) (*Handler[StateT], error) {
|
||||
resolvedStreamMode := ""
|
||||
if len(streamMode) > 1 {
|
||||
return nil, fmt.Errorf("start accepts at most one stream mode")
|
||||
}
|
||||
if len(streamMode) == 1 {
|
||||
resolvedStreamMode = streamMode[0]
|
||||
}
|
||||
engine := NewRustEngine()
|
||||
for _, ch := range g.asyncChannels {
|
||||
if err := engine.AddAsyncChannel(ch); err != nil {
|
||||
@@ -147,6 +173,7 @@ func (g *CompiledGraph[StateT]) Start(initialInput any, initialState StateT) (*H
|
||||
rawState, err := engine.RunGraph(
|
||||
g.entryPoint,
|
||||
g.finishPoint,
|
||||
resolvedStreamMode,
|
||||
initialState,
|
||||
initialInput,
|
||||
func(node string, nodeInput any, fallbackState map[string]any) (Command, error) {
|
||||
|
||||
@@ -128,6 +128,59 @@ func (e *RustEngine) AddAsyncChannel(channel string) error {
|
||||
return parseRustStatus(resp)
|
||||
}
|
||||
|
||||
func (e *RustEngine) StartStream(streamMode string) error {
|
||||
var cmode *C.char
|
||||
if streamMode != "" {
|
||||
cmode = C.CString(streamMode)
|
||||
defer C.free(unsafe.Pointer(cmode))
|
||||
}
|
||||
resp := C.rc_start_stream(e.ptr, cmode)
|
||||
return parseRustStatus(resp)
|
||||
}
|
||||
|
||||
func (e *RustEngine) ReceiveStream() (any, bool, error) {
|
||||
resp := C.rc_receive_stream_json(e.ptr)
|
||||
defer C.rc_string_free(resp)
|
||||
|
||||
raw := C.GoString(resp)
|
||||
var status struct {
|
||||
OK bool `json:"ok"`
|
||||
Error string `json:"error"`
|
||||
HasEvent bool `json:"has_event"`
|
||||
Event json.RawMessage `json:"event"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(raw), &status); err != nil {
|
||||
return nil, false, fmt.Errorf("decode rust stream response: %w", err)
|
||||
}
|
||||
if !status.OK {
|
||||
return nil, false, fmt.Errorf("rust stream failed: %s", status.Error)
|
||||
}
|
||||
if !status.HasEvent {
|
||||
return nil, false, nil
|
||||
}
|
||||
var event any
|
||||
if err := json.Unmarshal(status.Event, &event); err != nil {
|
||||
return nil, false, fmt.Errorf("decode stream event: %w", err)
|
||||
}
|
||||
return coerceJSONValue(event), true, nil
|
||||
}
|
||||
|
||||
func (e *RustEngine) SendCustomStreamEvent(value any) error {
|
||||
payload, err := json.Marshal(value)
|
||||
if err != nil {
|
||||
return fmt.Errorf("marshal stream event: %w", err)
|
||||
}
|
||||
cval := C.CString(string(payload))
|
||||
defer C.free(unsafe.Pointer(cval))
|
||||
resp := C.rc_send_custom_stream_event(e.ptr, cval)
|
||||
return parseRustStatus(resp)
|
||||
}
|
||||
|
||||
func (e *RustEngine) CloseStream() error {
|
||||
resp := C.rc_close_stream(e.ptr)
|
||||
return parseRustStatus(resp)
|
||||
}
|
||||
|
||||
func (e *RustEngine) Publish(channel string, value any) error {
|
||||
payload, err := json.Marshal(value)
|
||||
if err != nil {
|
||||
@@ -173,6 +226,7 @@ func (e *RustEngine) WaitAnyOf(cond AnyOfCondition) (WaitEvent, error) {
|
||||
func (e *RustEngine) RunGraph(
|
||||
entryPoint string,
|
||||
finishPoint string,
|
||||
streamMode string,
|
||||
initialState any,
|
||||
initialInput any,
|
||||
exec func(node string, nodeInput any, state map[string]any) (Command, error),
|
||||
@@ -189,10 +243,17 @@ func (e *RustEngine) RunGraph(
|
||||
cfinish := C.CString(finishPoint)
|
||||
cinitial := C.CString(string(initialJSON))
|
||||
cinitialInput := C.CString(string(initialInputJSON))
|
||||
var cstreamMode *C.char
|
||||
if streamMode != "" {
|
||||
cstreamMode = C.CString(streamMode)
|
||||
}
|
||||
defer C.free(unsafe.Pointer(centry))
|
||||
defer C.free(unsafe.Pointer(cfinish))
|
||||
defer C.free(unsafe.Pointer(cinitial))
|
||||
defer C.free(unsafe.Pointer(cinitialInput))
|
||||
if cstreamMode != nil {
|
||||
defer C.free(unsafe.Pointer(cstreamMode))
|
||||
}
|
||||
|
||||
callbackID := registerRunGraphCallbackCtx(&runGraphCallbackCtx{exec: exec})
|
||||
defer unregisterRunGraphCallbackCtx(callbackID)
|
||||
@@ -203,6 +264,7 @@ func (e *RustEngine) RunGraph(
|
||||
cfinish,
|
||||
cinitial,
|
||||
cinitialInput,
|
||||
cstreamMode,
|
||||
C.ulong(callbackID),
|
||||
(C.rc_node_callback_t)(C.goNodeCallback),
|
||||
)
|
||||
|
||||
@@ -0,0 +1,102 @@
|
||||
package tests
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
ag "github.com/langchain-ai/langgraph/langgraph-go/advancedgraph"
|
||||
)
|
||||
|
||||
type streamState struct {
|
||||
Done bool `json:"done"`
|
||||
}
|
||||
|
||||
type streamWorkflow struct{}
|
||||
|
||||
func (w *streamWorkflow) startNode(ctx *ag.Context, _ any, state streamState) (ag.Command, error) {
|
||||
if err := ctx.SendCustomStreamEvent(map[string]any{"step": "start", "value": 1}); err != nil {
|
||||
return ag.Command{}, err
|
||||
}
|
||||
time.Sleep(80 * time.Millisecond)
|
||||
if err := ctx.SendCustomStreamEvent(map[string]any{"step": "start", "value": 2}); err != nil {
|
||||
return ag.Command{}, err
|
||||
}
|
||||
return ag.Command{
|
||||
Update: state,
|
||||
Goto: []ag.Send{
|
||||
{Node: w.finishNode},
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (w *streamWorkflow) finishNode(ctx *ag.Context, _ any, state streamState) (ag.Command, error) {
|
||||
state.Done = true
|
||||
return ag.Command{Update: state}, nil
|
||||
}
|
||||
|
||||
func TestCustomStreamReceiveAndClose(t *testing.T) {
|
||||
workflow := &streamWorkflow{}
|
||||
graph := ag.NewAdvancedStateGraph[streamState]()
|
||||
graph.AddEntryNode(workflow.startNode)
|
||||
graph.AddFinishNode(workflow.finishNode)
|
||||
|
||||
handler, err := graph.Compile().Start(nil, streamState{Done: false}, "custom")
|
||||
if err != nil {
|
||||
t.Fatalf("start failed: %v", err)
|
||||
}
|
||||
|
||||
event, err := handler.ReceiveStream()
|
||||
if err != nil {
|
||||
t.Fatalf("receive stream failed: %v", err)
|
||||
}
|
||||
if event == nil {
|
||||
t.Fatalf("expected first stream event, got nil")
|
||||
}
|
||||
eventMap, ok := event.(map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("unexpected event type: %T", event)
|
||||
}
|
||||
if eventMap["step"] != "start" {
|
||||
t.Fatalf("unexpected stream event payload: %#v", eventMap)
|
||||
}
|
||||
|
||||
if err := handler.CloseStream(); err != nil {
|
||||
t.Fatalf("close stream failed: %v", err)
|
||||
}
|
||||
|
||||
closedEvent, err := handler.ReceiveStream()
|
||||
if err != nil {
|
||||
t.Fatalf("receive stream after close failed: %v", err)
|
||||
}
|
||||
if closedEvent != nil {
|
||||
t.Fatalf("expected nil stream event after close, got %#v", closedEvent)
|
||||
}
|
||||
|
||||
result, err := handler.WaitForResult()
|
||||
if err != nil {
|
||||
t.Fatalf("result failed: %v", err)
|
||||
}
|
||||
if !result.Done {
|
||||
t.Fatalf("expected final state done=true, got %#v", result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOnlyCustomStreamModeSupported(t *testing.T) {
|
||||
workflow := &streamWorkflow{}
|
||||
graph := ag.NewAdvancedStateGraph[streamState]()
|
||||
graph.AddEntryNode(workflow.startNode)
|
||||
graph.AddFinishNode(workflow.finishNode)
|
||||
|
||||
handler, err := graph.Compile().Start(nil, streamState{Done: false}, "values")
|
||||
if err != nil {
|
||||
t.Fatalf("start failed: %v", err)
|
||||
}
|
||||
_, runErr := handler.WaitForResult()
|
||||
if runErr == nil {
|
||||
t.Fatalf("expected run error for unsupported stream mode")
|
||||
}
|
||||
if !strings.Contains(runErr.Error(), "only `custom` is supported") {
|
||||
t.Fatalf("unexpected error: %v", runErr)
|
||||
}
|
||||
}
|
||||
Generated
-124
@@ -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"
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -19,12 +19,17 @@ void rc_engine_free(Engine* ptr);
|
||||
char* rc_add_async_channel(Engine* ptr, const char* channel);
|
||||
char* rc_publish_json(Engine* ptr, const char* channel, const char* value_json);
|
||||
char* rc_wait_any_of_json(Engine* ptr, const char* any_of_json);
|
||||
char* rc_start_stream(Engine* ptr, const char* stream_mode);
|
||||
char* rc_receive_stream_json(Engine* ptr);
|
||||
char* rc_send_custom_stream_event(Engine* ptr, const char* value_json);
|
||||
char* rc_close_stream(Engine* ptr);
|
||||
char* rc_run_graph_json(
|
||||
Engine* ptr,
|
||||
const char* entry_point,
|
||||
const char* finish_point,
|
||||
const char* initial_state_json,
|
||||
const char* initial_input_json,
|
||||
const char* stream_mode,
|
||||
unsigned long user_data,
|
||||
rc_node_callback_t callback
|
||||
);
|
||||
|
||||
+281
-7
@@ -10,7 +10,7 @@ use std::sync::OnceLock;
|
||||
use std::thread;
|
||||
use std::time::{Duration, Instant};
|
||||
use tokio::runtime::Runtime;
|
||||
use tokio::sync::Notify;
|
||||
use tokio::sync::{mpsc as tokio_mpsc, Mutex as AsyncMutex, Notify};
|
||||
|
||||
#[derive(Clone, Debug, Serialize, Deserialize)]
|
||||
#[serde(tag = "kind")]
|
||||
@@ -62,6 +62,31 @@ pub enum NodeOutcome<U, A> {
|
||||
Suspended { wait: WaitRequest },
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct CallbackSendPayloadJson {
|
||||
node: String,
|
||||
#[serde(default)]
|
||||
arg: Value,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct CallbackNodeExecResultJsonWire {
|
||||
update: Option<Value>,
|
||||
#[serde(default)]
|
||||
sends: Vec<CallbackSendPayloadJson>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct CallbackEnvelopeIn {
|
||||
ok: bool,
|
||||
#[serde(default)]
|
||||
payload: Option<CallbackNodeExecResultJsonWire>,
|
||||
#[serde(default)]
|
||||
suspend: Option<WaitRequest>,
|
||||
#[serde(default)]
|
||||
error: Option<String>,
|
||||
}
|
||||
|
||||
type Task = Box<dyn FnOnce() + Send + 'static>;
|
||||
|
||||
fn debug_enabled() -> bool {
|
||||
@@ -279,10 +304,33 @@ 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>>>>,
|
||||
channel_notify: Arc<Notify>,
|
||||
stream: Arc<StdMutex<Option<StreamChannel>>>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct StreamChannel {
|
||||
sender: tokio_mpsc::UnboundedSender<serde_json::Value>,
|
||||
receiver: Arc<AsyncMutex<tokio_mpsc::UnboundedReceiver<serde_json::Value>>>,
|
||||
}
|
||||
|
||||
impl Engine {
|
||||
@@ -297,6 +345,60 @@ impl Engine {
|
||||
channels.entry(name.to_owned()).or_default();
|
||||
}
|
||||
|
||||
pub fn start_stream(&self, stream_mode: Option<&str>) -> Result<(), String> {
|
||||
let mode = stream_mode.and_then(|m| {
|
||||
let trimmed = m.trim();
|
||||
if trimmed.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(trimmed)
|
||||
}
|
||||
});
|
||||
let mut stream = self.stream.lock().expect("stream mutex poisoned");
|
||||
if let Some(mode) = mode {
|
||||
if mode != "custom" {
|
||||
return Err(format!(
|
||||
"unsupported stream_mode `{mode}`, only `custom` is supported"
|
||||
));
|
||||
}
|
||||
let (sender, receiver) = tokio_mpsc::unbounded_channel();
|
||||
*stream = Some(StreamChannel {
|
||||
sender,
|
||||
receiver: Arc::new(AsyncMutex::new(receiver)),
|
||||
});
|
||||
} else {
|
||||
*stream = None;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn close_stream(&self) {
|
||||
let mut stream = self.stream.lock().expect("stream mutex poisoned");
|
||||
*stream = None;
|
||||
}
|
||||
|
||||
pub fn send_custom_stream_event(&self, value: serde_json::Value) {
|
||||
let sender = {
|
||||
let stream = self.stream.lock().expect("stream mutex poisoned");
|
||||
stream.as_ref().map(|s| s.sender.clone())
|
||||
};
|
||||
if let Some(tx) = sender {
|
||||
let _ = tx.send(value);
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn receive_stream_async(&self) -> Option<serde_json::Value> {
|
||||
let receiver = {
|
||||
let stream = self.stream.lock().expect("stream mutex poisoned");
|
||||
stream.as_ref().map(|s| Arc::clone(&s.receiver))
|
||||
};
|
||||
let Some(receiver) = receiver else {
|
||||
return None;
|
||||
};
|
||||
let mut guard = receiver.lock().await;
|
||||
guard.recv().await
|
||||
}
|
||||
|
||||
pub fn publish_json(&self, channel: &str, value: serde_json::Value) -> Result<(), String> {
|
||||
debug_log(&format!("Engine::publish_json(channel={channel})"));
|
||||
let mut channels = self.channels.lock().expect("channels mutex poisoned");
|
||||
@@ -339,7 +441,10 @@ impl Engine {
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn wait_for_any_of_async(&self, any_of: &AnyOfCondition) -> Result<WaitEvent, String> {
|
||||
pub async fn wait_for_any_of_async(
|
||||
&self,
|
||||
any_of: &AnyOfCondition,
|
||||
) -> Result<WaitEvent, String> {
|
||||
debug_log(&format!(
|
||||
"Engine::wait_for_any_of_async(conditions={})",
|
||||
any_of.conditions.len()
|
||||
@@ -398,11 +503,7 @@ impl Engine {
|
||||
run_loop_block_on(self.wait_for_any_of_async(any_of))
|
||||
}
|
||||
|
||||
fn try_take_channel_event(
|
||||
&self,
|
||||
channel: &str,
|
||||
n: usize,
|
||||
) -> Result<Option<WaitEvent>, String> {
|
||||
fn try_take_channel_event(&self, channel: &str, n: usize) -> Result<Option<WaitEvent>, String> {
|
||||
let mut channels = self.channels.lock().expect("channels mutex poisoned");
|
||||
let queue = channels
|
||||
.get_mut(channel)
|
||||
@@ -431,3 +532,176 @@ 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,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn parse_callback_envelope_json(
|
||||
raw: &str,
|
||||
node_name: &str,
|
||||
) -> Result<NodeOutcome<Value, Value>, String> {
|
||||
let parsed: CallbackEnvelopeIn = serde_json::from_str(raw)
|
||||
.map_err(|e| format!("decode callback envelope for `{node_name}` failed: {e}"))?;
|
||||
if !parsed.ok {
|
||||
return Err(parsed
|
||||
.error
|
||||
.unwrap_or_else(|| format!("callback reported error for `{node_name}`")));
|
||||
}
|
||||
if let Some(wait) = parsed.suspend {
|
||||
return Ok(NodeOutcome::Suspended { wait });
|
||||
}
|
||||
let payload = parsed
|
||||
.payload
|
||||
.ok_or_else(|| format!("callback payload missing for `{node_name}`"))?;
|
||||
let sends = payload
|
||||
.sends
|
||||
.into_iter()
|
||||
.map(|s| SendPayload {
|
||||
node: s.node,
|
||||
arg: s.arg,
|
||||
})
|
||||
.collect();
|
||||
Ok(NodeOutcome::Completed(NodeExecResult {
|
||||
update: payload.update,
|
||||
sends,
|
||||
}))
|
||||
}
|
||||
|
||||
@@ -1,4 +1,2 @@
|
||||
mod engine;
|
||||
mod lib_c;
|
||||
#[cfg(feature = "python-bindings")]
|
||||
mod lib_py;
|
||||
|
||||
+133
-245
@@ -1,39 +1,11 @@
|
||||
use crate::engine::{
|
||||
merge_json_update, node_pool_execute, run_loop_block_on, run_loop_spawn, AnyOfCondition,
|
||||
Engine, NodeExecResult, NodeOutcome, SendPayload, WaitEvent, WaitRequest,
|
||||
parse_callback_envelope_json, run_graph_json_with_callback, run_loop_block_on, run_loop_spawn,
|
||||
AnyOfCondition, Engine, NodeOutcome,
|
||||
};
|
||||
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 {
|
||||
node: String,
|
||||
#[serde(default)]
|
||||
arg: Value,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct NodeExecResultJsonWire {
|
||||
update: Option<Value>,
|
||||
#[serde(default)]
|
||||
sends: Vec<SendPayloadJson>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct CallbackEnvelopeIn {
|
||||
ok: bool,
|
||||
#[serde(default)]
|
||||
payload: Option<NodeExecResultJsonWire>,
|
||||
#[serde(default)]
|
||||
suspend: Option<WaitRequest>,
|
||||
#[serde(default)]
|
||||
error: Option<String>,
|
||||
}
|
||||
|
||||
type CNodeCallback = unsafe extern "C" fn(
|
||||
user_data: libc::c_ulong,
|
||||
@@ -42,9 +14,6 @@ type CNodeCallback = unsafe extern "C" fn(
|
||||
state_json: *mut c_char,
|
||||
) -> *mut c_char;
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
struct CUserData(libc::c_ulong);
|
||||
|
||||
fn cstr_to_str<'a>(ptr: *const c_char) -> Result<&'a str, String> {
|
||||
if ptr.is_null() {
|
||||
return Err("Received null pointer".to_string());
|
||||
@@ -63,213 +32,6 @@ fn into_c_ptr(s: String) -> *mut c_char {
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_c_callback_result(
|
||||
raw: String,
|
||||
node_name: &str,
|
||||
) -> Result<NodeOutcome<Value, Value>, String> {
|
||||
let parsed: CallbackEnvelopeIn = serde_json::from_str(&raw)
|
||||
.map_err(|e| format!("decode callback envelope for `{node_name}` failed: {e}"))?;
|
||||
if !parsed.ok {
|
||||
return Err(parsed
|
||||
.error
|
||||
.unwrap_or_else(|| format!("callback reported error for `{node_name}`")));
|
||||
}
|
||||
if let Some(wait) = parsed.suspend {
|
||||
return Ok(NodeOutcome::Suspended { wait });
|
||||
}
|
||||
let payload = parsed
|
||||
.payload
|
||||
.ok_or_else(|| format!("callback payload missing for `{node_name}`"))?;
|
||||
let sends = payload
|
||||
.sends
|
||||
.into_iter()
|
||||
.map(|s| SendPayload {
|
||||
node: s.node,
|
||||
arg: s.arg,
|
||||
})
|
||||
.collect();
|
||||
Ok(NodeOutcome::Completed(NodeExecResult {
|
||||
update: payload.update,
|
||||
sends,
|
||||
}))
|
||||
}
|
||||
|
||||
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()))
|
||||
@@ -382,6 +144,86 @@ pub unsafe extern "C" fn rc_wait_any_of_json(
|
||||
}
|
||||
}
|
||||
|
||||
#[no_mangle]
|
||||
/// # Safety
|
||||
/// `ptr` must be a valid engine pointer from `rc_engine_new`.
|
||||
/// `stream_mode` must be null or a valid null-terminated UTF-8 string pointer.
|
||||
pub unsafe extern "C" fn rc_start_stream(
|
||||
ptr: *mut Engine,
|
||||
stream_mode: *const c_char,
|
||||
) -> *mut c_char {
|
||||
if ptr.is_null() {
|
||||
return into_c_ptr("{\"ok\":false,\"error\":\"null engine pointer\"}".to_string());
|
||||
}
|
||||
let mode = if stream_mode.is_null() {
|
||||
None
|
||||
} else {
|
||||
match cstr_to_str(stream_mode) {
|
||||
Ok(v) => Some(v),
|
||||
Err(e) => return into_c_ptr(format!("{{\"ok\":false,\"error\":\"{e}\"}}")),
|
||||
}
|
||||
};
|
||||
match (*ptr).start_stream(mode) {
|
||||
Ok(()) => into_c_ptr("{\"ok\":true}".to_string()),
|
||||
Err(e) => into_c_ptr(format!("{{\"ok\":false,\"error\":\"{e}\"}}")),
|
||||
}
|
||||
}
|
||||
|
||||
#[no_mangle]
|
||||
/// # Safety
|
||||
/// `ptr` must be a valid engine pointer from `rc_engine_new`.
|
||||
pub unsafe extern "C" fn rc_receive_stream_json(ptr: *mut Engine) -> *mut c_char {
|
||||
if ptr.is_null() {
|
||||
return into_c_ptr("{\"ok\":false,\"error\":\"null engine pointer\"}".to_string());
|
||||
}
|
||||
let event = run_loop_block_on((*ptr).receive_stream_async());
|
||||
match event {
|
||||
Some(value) => match serde_json::to_string(&value) {
|
||||
Ok(s) => into_c_ptr(format!("{{\"ok\":true,\"has_event\":true,\"event\":{s}}}")),
|
||||
Err(e) => into_c_ptr(format!("{{\"ok\":false,\"error\":\"{e}\"}}")),
|
||||
},
|
||||
None => into_c_ptr("{\"ok\":true,\"has_event\":false}".to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
#[no_mangle]
|
||||
/// # Safety
|
||||
/// `ptr` must be a valid engine pointer from `rc_engine_new`.
|
||||
/// `value_json` must be a valid null-terminated UTF-8 string pointer.
|
||||
pub unsafe extern "C" fn rc_send_custom_stream_event(
|
||||
ptr: *mut Engine,
|
||||
value_json: *const c_char,
|
||||
) -> *mut c_char {
|
||||
if ptr.is_null() {
|
||||
return into_c_ptr("{\"ok\":false,\"error\":\"null engine pointer\"}".to_string());
|
||||
}
|
||||
let value_json = match cstr_to_str(value_json) {
|
||||
Ok(v) => v,
|
||||
Err(e) => return into_c_ptr(format!("{{\"ok\":false,\"error\":\"{e}\"}}")),
|
||||
};
|
||||
let value: Value = match serde_json::from_str(value_json) {
|
||||
Ok(v) => v,
|
||||
Err(e) => {
|
||||
return into_c_ptr(format!(
|
||||
"{{\"ok\":false,\"error\":\"invalid JSON value: {e}\"}}"
|
||||
))
|
||||
}
|
||||
};
|
||||
(*ptr).send_custom_stream_event(value);
|
||||
into_c_ptr("{\"ok\":true}".to_string())
|
||||
}
|
||||
|
||||
#[no_mangle]
|
||||
/// # Safety
|
||||
/// `ptr` must be a valid engine pointer from `rc_engine_new`.
|
||||
pub unsafe extern "C" fn rc_close_stream(ptr: *mut Engine) -> *mut c_char {
|
||||
if ptr.is_null() {
|
||||
return into_c_ptr("{\"ok\":false,\"error\":\"null engine pointer\"}".to_string());
|
||||
}
|
||||
(*ptr).close_stream();
|
||||
into_c_ptr("{\"ok\":true}".to_string())
|
||||
}
|
||||
|
||||
#[no_mangle]
|
||||
/// # Safety
|
||||
/// `ptr` must be a valid engine pointer from `rc_engine_new`.
|
||||
@@ -393,6 +235,7 @@ pub unsafe extern "C" fn rc_run_graph_json(
|
||||
finish_point: *const c_char,
|
||||
initial_state_json: *const c_char,
|
||||
initial_input_json: *const c_char,
|
||||
stream_mode: *const c_char,
|
||||
user_data: libc::c_ulong,
|
||||
callback: Option<CNodeCallback>,
|
||||
) -> *mut c_char {
|
||||
@@ -434,21 +277,66 @@ pub unsafe extern "C" fn rc_run_graph_json(
|
||||
))
|
||||
}
|
||||
};
|
||||
let stream_mode = if stream_mode.is_null() {
|
||||
None
|
||||
} else {
|
||||
match cstr_to_str(stream_mode) {
|
||||
Ok(v) => Some(v.to_string()),
|
||||
Err(e) => return into_c_ptr(format!("{{\"ok\":false,\"error\":\"{e}\"}}")),
|
||||
}
|
||||
};
|
||||
|
||||
if let Err(e) = (*ptr).start_stream(stream_mode.as_deref()) {
|
||||
return into_c_ptr(format!("{{\"ok\":false,\"error\":\"{e}\"}}"));
|
||||
}
|
||||
|
||||
let (tx, rx) = mpsc::channel::<Result<Value, String>>();
|
||||
let user_data = CUserData(user_data);
|
||||
let user_data_bits = 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_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());
|
||||
}
|
||||
parse_callback_envelope_json(&out_raw, &node)
|
||||
};
|
||||
let out = run_graph_json_with_callback(
|
||||
entry_point,
|
||||
finish_point,
|
||||
initial_state,
|
||||
initial_input,
|
||||
run_engine,
|
||||
user_data,
|
||||
callback,
|
||||
run_engine.clone(),
|
||||
callback_wrapper,
|
||||
)
|
||||
.await;
|
||||
run_engine.close_stream();
|
||||
let _ = tx.send(out);
|
||||
});
|
||||
if let Err(e) = submit {
|
||||
|
||||
@@ -1,423 +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 run_graph_py(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
entry_point: &str,
|
||||
finish_point: &str,
|
||||
initial_state: Py<PyAny>,
|
||||
callback: Py<PyAny>,
|
||||
) -> 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();
|
||||
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)
|
||||
.await;
|
||||
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(())
|
||||
}
|
||||
Generated
+214
@@ -0,0 +1,214 @@
|
||||
# This file is automatically @generated by Cargo.
|
||||
# It is not intended for manual editing.
|
||||
version = 4
|
||||
|
||||
[[package]]
|
||||
name = "bitflags"
|
||||
version = "2.11.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "843867be96c8daad0d758b57df9392b6d8d271134fce549de6ce169ff98a92af"
|
||||
|
||||
[[package]]
|
||||
name = "cfg-if"
|
||||
version = "1.0.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801"
|
||||
|
||||
[[package]]
|
||||
name = "itoa"
|
||||
version = "1.0.17"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "92ecc6618181def0457392ccd0ee51198e065e016d1d527a7ac1b6dc7c1f09d2"
|
||||
|
||||
[[package]]
|
||||
name = "langgraph_rust_core"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"parking_lot",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "libc"
|
||||
version = "0.2.183"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b5b646652bf6661599e1da8901b3b9522896f01e736bad5f723fe7a3a27f899d"
|
||||
|
||||
[[package]]
|
||||
name = "lock_api"
|
||||
version = "0.4.14"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965"
|
||||
dependencies = [
|
||||
"scopeguard",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "memchr"
|
||||
version = "2.8.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79"
|
||||
|
||||
[[package]]
|
||||
name = "parking_lot"
|
||||
version = "0.12.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a"
|
||||
dependencies = [
|
||||
"lock_api",
|
||||
"parking_lot_core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "parking_lot_core"
|
||||
version = "0.9.12"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"libc",
|
||||
"redox_syscall",
|
||||
"smallvec",
|
||||
"windows-link",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pin-project-lite"
|
||||
version = "0.2.17"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd"
|
||||
|
||||
[[package]]
|
||||
name = "proc-macro2"
|
||||
version = "1.0.106"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934"
|
||||
dependencies = [
|
||||
"unicode-ident",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "quote"
|
||||
version = "1.0.45"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "redox_syscall"
|
||||
version = "0.5.18"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d"
|
||||
dependencies = [
|
||||
"bitflags",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "scopeguard"
|
||||
version = "1.2.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49"
|
||||
|
||||
[[package]]
|
||||
name = "serde"
|
||||
version = "1.0.228"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e"
|
||||
dependencies = [
|
||||
"serde_core",
|
||||
"serde_derive",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_core"
|
||||
version = "1.0.228"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad"
|
||||
dependencies = [
|
||||
"serde_derive",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_derive"
|
||||
version = "1.0.228"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_json"
|
||||
version = "1.0.149"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "83fc039473c5595ace860d8c4fafa220ff474b3fc6bfdb4293327f1a37e94d86"
|
||||
dependencies = [
|
||||
"itoa",
|
||||
"memchr",
|
||||
"serde",
|
||||
"serde_core",
|
||||
"zmij",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "smallvec"
|
||||
version = "1.15.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03"
|
||||
|
||||
[[package]]
|
||||
name = "syn"
|
||||
version = "2.0.117"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"unicode-ident",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tokio"
|
||||
version = "1.50.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "27ad5e34374e03cfffefc301becb44e9dc3c17584f414349ebe29ed26661822d"
|
||||
dependencies = [
|
||||
"pin-project-lite",
|
||||
"tokio-macros",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tokio-macros"
|
||||
version = "2.6.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5c55a2eff8b69ce66c84f85e1da1c233edc36ceb85a2058d11b0d6a3c7e7569c"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "unicode-ident"
|
||||
version = "1.0.24"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75"
|
||||
|
||||
[[package]]
|
||||
name = "windows-link"
|
||||
version = "0.2.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5"
|
||||
|
||||
[[package]]
|
||||
name = "zmij"
|
||||
version = "1.0.21"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa"
|
||||
@@ -0,0 +1,17 @@
|
||||
[package]
|
||||
name = "langgraph_rust_core"
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
|
||||
[lib]
|
||||
name = "langgraph_rust_core"
|
||||
path = "../rust-core/src/lib.rs"
|
||||
crate-type = ["cdylib", "rlib"]
|
||||
|
||||
[dependencies]
|
||||
serde = { version = "1.0", features = ["derive"] }
|
||||
serde_json = "1.0"
|
||||
parking_lot = "0.12"
|
||||
libc = "0.2"
|
||||
tokio = { version = "1", features = ["macros", "rt-multi-thread", "sync", "time"] }
|
||||
|
||||
@@ -0,0 +1,33 @@
|
||||
|
||||
PYPI_REPOSITORY ?= pypi
|
||||
PYPI_TOKEN ?=
|
||||
PYTHON_VERSION ?= 3.13
|
||||
TEST_BASE = uv run --python $(PYTHON_VERSION) --with pytest --with anyio --with typing_extensions --with pydantic --with langgraph pytest -q -s
|
||||
|
||||
.PHONY: publish-to-pypi-saf-python-sdk all-tests test_primitives test_sub_agents test_update_elision test_run_pool_size test-benchmark
|
||||
publish-to-pypi-saf-python-sdk:
|
||||
@test -n "$(PYPI_TOKEN)" || (echo "PYPI_TOKEN is required"; exit 1)
|
||||
cd "$(CURDIR)" && \
|
||||
MATURIN_PYPI_TOKEN="$(PYPI_TOKEN)" \
|
||||
uvx maturin publish --repository $(PYPI_REPOSITORY) --non-interactive --skip-existing --no-sdist
|
||||
|
||||
all-tests:
|
||||
$(TEST_BASE) tests/advanced-graph/test_*.py
|
||||
|
||||
test_primitives:
|
||||
$(TEST_BASE) tests/advanced-graph/test_primitives.py
|
||||
|
||||
test_sub_agents:
|
||||
$(TEST_BASE) tests/advanced-graph/test_sub_agents.py
|
||||
|
||||
test_update_elision:
|
||||
$(TEST_BASE) tests/advanced-graph/test_update_elision.py
|
||||
|
||||
test_run_pool_size:
|
||||
$(TEST_BASE) tests/advanced-graph/test_run_pool_size.py
|
||||
|
||||
test-benchmark:
|
||||
# Increase per-process file descriptor limit for high-concurrency benchmark runs.
|
||||
ulimit -n 65536 || true; \
|
||||
uv run --python $(PYTHON_VERSION) --with anyio --with typing_extensions --with pydantic --with langgraph python tests/advanced-graph/benchmark_stategraph_vs_advancedgraph.py
|
||||
|
||||
@@ -0,0 +1,16 @@
|
||||
# saf-python-sdk
|
||||
|
||||
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)
|
||||
- Rust core engine via C bindings (`ctypes`)
|
||||
|
||||
It does not package the original `langgraph` `stategraph` stack.
|
||||
|
||||
## Moved assets
|
||||
|
||||
- Original advanced graph design doc: `evolve-extend-langgraph.md`
|
||||
- Original advanced graph tests: `tests/advanced-graph/`
|
||||
|
||||
+6
-9
@@ -32,11 +32,10 @@ Sub-agents also cannot simply be modeled as subgraphs, because subgraphs today e
|
||||
|
||||
The closest workaround today is double texting, but it has a fundamental flaw: when a new audio input arrives, the previous one is interrupted and canceled rather than being allowed to gracefully complete. The workflow code itself should have the control to decide whether to stop running.
|
||||
|
||||
LangGraph is, at its core, a general-purpose workflow engine. Although we focus primarily on agent development, none of the primitives it offers are exclusive to agents or dedicated solely to agentic use cases. Conversely, there is nothing that a general-purpose workflow engine provides that we can safely assume agent development will _never_ need.
|
||||
LangGraph is, at its core, a general-purpose workflow engine. Although we focus primarily on agent development, none of the primitives it offers are exclusive to agents or dedicated solely to agentic use cases. Conversely, there is nothing that a general-purpose workflow engine provides that we can safely assume agent development will _never_ need.
|
||||
|
||||
The difference is probably only priority. For example, durable timer where a step can sleep for hours, days or months before resuming. Traditional workflow engines — those built for general microservice orchestration(which doesn't need streaming) -- they may need durable timers. In the agent development world today, most agents are still relatively simple. There are not yet many scenarios that require a step to wait for hours or days before proceeding.
|
||||
|
||||
|
||||
## Deriving What's Needed from First Principles
|
||||
|
||||
Before jumping to solutions, it is worth stepping back and asking a fundamental question: what is an orchestration engine, and what do users expect it to provide?
|
||||
@@ -47,7 +46,7 @@ Starting from the simple. A developer could write a simple `main` function — a
|
||||
|
||||
But if that machine crashes, you probably do not want the process to start over from scratch. You want it to resume from the last step that completed successfully. And if a step fails, you might want it to retry automatically before giving up.
|
||||
|
||||
LangGraph handles this case very well.
|
||||
LangGraph handles this case very well.
|
||||
|
||||
There is an important constraint worth calling out explicitly: LangGraph requires the developer to organize their code into **nodes**, which serve as the boundaries at which checkpoints can be taken. This is a constraint shared by every workflow engine — it is simply not feasible to persist a checkpoint after every single line of arbitrary code.
|
||||
|
||||
@@ -59,9 +58,9 @@ This is precisely why LangGraph's superstep restriction feels awkward in practic
|
||||
|
||||
Multiple threads and processes do, however, need to coordinate with each other. In concurrent programming, channels are an essential primitive precisely because they provide a safe, structured way for threads to communicate and synchronize without relying on shared mutable memory — avoiding data races and deadlocks. In some cases, threads may use locking for coordination, but the preferred approach is message passing through channels.
|
||||
|
||||
NOTE: "channel" is overloaded term here as it's also an internal term within current LangGraph pregel algorithm.
|
||||
NOTE: "channel" is overloaded term here as it's also an internal term within current LangGraph pregel algorithm.
|
||||
|
||||
LangGraph already has a mechanism that is closely related: `interrupt`. A run can be interrupted, and then another run can resume it. If we look at this through the lens of channels, `interrupt` is essentially a **channel with size 0** — a synchronous rendezvous point where one side blocks until the other side is ready.
|
||||
LangGraph already has a mechanism that is closely related: `interrupt`. A run can be interrupted, and then another run can resume it. If we look at this through the lens of channels, `interrupt` is essentially a **channel with size 0** — a synchronous rendezvous point where one side blocks until the other side is ready.
|
||||
|
||||
The natural extension:
|
||||
|
||||
@@ -97,7 +96,6 @@ Branch 2: a → b2 → b22 → ...
|
||||
|
||||
Each branch advances at its own pace. `b1` finishing triggers `b11` immediately, without waiting for `b2`.
|
||||
|
||||
|
||||
No API change is needed from the user's perspective — the graph definition stays the same. The change is in the execution semantics: the engine no longer forces all parallel nodes to synchronize at each step boundary. Each branch is checkpointed independently, so if `b1 → b11` completes while `b2` is still running, `b11`'s result is already persisted.
|
||||
|
||||
This is necessary for the next one -- Light-weight Interrupt: Only Block the Current Node. Because we want to let other nodes continue to run while a node is waiting on something.
|
||||
@@ -118,13 +116,12 @@ The `wait_for` call takes a channel name and optionally a count `N`, meaning "wa
|
||||
|
||||
See [test_sub_agents.py](../libs/langgraph/tests/advanced-graph/test_sub_agents.py)
|
||||
|
||||
|
||||
### P2: likely needed
|
||||
#### subGraph redesign
|
||||
#### durable timers
|
||||
#### more flexiable waiting conditions on interrupts
|
||||
#### locking on state fields
|
||||
|
||||
|
||||
### P3: future needed or nice to have
|
||||
#### RPC
|
||||
#### RPC
|
||||
|
||||
@@ -0,0 +1,26 @@
|
||||
[build-system]
|
||||
requires = ["maturin>=1.7,<2"]
|
||||
build-backend = "maturin"
|
||||
|
||||
[project]
|
||||
name = "saf-python-sdk"
|
||||
version = "0.1.2"
|
||||
description = "Standalone advanced graph runtime powered by Rust engine"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10"
|
||||
license = "MIT"
|
||||
authors = [{ name = "LangGraph Contributors" }]
|
||||
classifiers = [
|
||||
"Programming Language :: Python :: 3",
|
||||
"Programming Language :: Python :: 3 :: Only",
|
||||
"Programming Language :: Rust",
|
||||
"License :: OSI Approved :: MIT License",
|
||||
"Operating System :: OS Independent",
|
||||
]
|
||||
|
||||
[tool.maturin]
|
||||
python-source = "python"
|
||||
module-name = "saf_python_sdk.langgraph_rust_core"
|
||||
bindings = "pyo3"
|
||||
features = ["python-bindings"]
|
||||
|
||||
@@ -0,0 +1,4 @@
|
||||
from .types import Command, Send
|
||||
|
||||
__all__ = ["Command", "Send"]
|
||||
|
||||
+8
-7
@@ -1,4 +1,4 @@
|
||||
from langgraph.advanced_graph.state import (
|
||||
from .state import (
|
||||
AdvancedStateGraph,
|
||||
AnyOfCondition,
|
||||
ChannelCondition,
|
||||
@@ -11,15 +11,16 @@ from langgraph.advanced_graph.state import (
|
||||
timer_condition,
|
||||
)
|
||||
|
||||
__all__ = (
|
||||
__all__ = [
|
||||
"AdvancedStateGraph",
|
||||
"AnyOfCondition",
|
||||
"ChannelCondition",
|
||||
"Context",
|
||||
"CompiledGraphEngine",
|
||||
"Context",
|
||||
"GraphRunHandler",
|
||||
"ChannelCondition",
|
||||
"TimerCondition",
|
||||
"any_of",
|
||||
"AnyOfCondition",
|
||||
"channel_condition",
|
||||
"timer_condition",
|
||||
)
|
||||
"any_of",
|
||||
]
|
||||
|
||||
+43
-14
@@ -2,18 +2,18 @@ from __future__ import annotations
|
||||
|
||||
import atexit
|
||||
import asyncio
|
||||
import os
|
||||
import inspect
|
||||
import os
|
||||
import threading
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from collections.abc import Callable, Coroutine, Sequence
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import dataclass
|
||||
from datetime import timedelta
|
||||
from typing import Any, Generic, TypeVar, cast
|
||||
|
||||
from langgraph_rust_core import PyRustEngine # type: ignore[import-untyped]
|
||||
from saf_python_sdk.rust_core_cffi import PyRustEngine
|
||||
|
||||
from langgraph.types import Command, Send
|
||||
from saf_python_sdk.types import Command, Send
|
||||
|
||||
StateT = TypeVar("StateT")
|
||||
|
||||
@@ -53,7 +53,7 @@ def _advanced_graph_executor() -> ThreadPoolExecutor:
|
||||
worker_count = max(worker_count, 1)
|
||||
_EXECUTOR = ThreadPoolExecutor(
|
||||
max_workers=worker_count,
|
||||
thread_name_prefix="langgraph-advanced-py",
|
||||
thread_name_prefix="saf-advanced-py",
|
||||
)
|
||||
atexit.register(_shutdown_advanced_graph_executor)
|
||||
return _EXECUTOR
|
||||
@@ -162,12 +162,15 @@ class CompiledGraphEngine(Generic[StateT]):
|
||||
handler = await self.astart(initial_state)
|
||||
return await handler
|
||||
|
||||
async def astart(self, initial_state: StateT) -> GraphRunHandler[StateT]:
|
||||
async def astart(
|
||||
self, initial_state: StateT, *, stream_mode: str | None = None
|
||||
) -> GraphRunHandler[StateT]:
|
||||
run = _GraphEngineRun(
|
||||
nodes=self._nodes,
|
||||
async_channel_specs=self._async_channels,
|
||||
entry_point=self._entry_point,
|
||||
finish_point=self._finish_point,
|
||||
stream_mode=stream_mode,
|
||||
)
|
||||
task = asyncio.create_task(run.run(initial_state))
|
||||
return GraphRunHandler(run=run, task=task)
|
||||
@@ -191,6 +194,9 @@ class Context:
|
||||
async def apublish_to_channel(self, channel: str, value: Any) -> None:
|
||||
await self._run.publish(channel, value)
|
||||
|
||||
def send_custom_stream_event(self, value: Any) -> None:
|
||||
self._run.send_custom_stream_event(value)
|
||||
|
||||
|
||||
class GraphRunHandler(Generic[StateT]):
|
||||
"""Handle for an active in-memory run."""
|
||||
@@ -204,6 +210,20 @@ class GraphRunHandler(Generic[StateT]):
|
||||
raise RuntimeError("Run has already completed")
|
||||
await self._run.publish(channel, value)
|
||||
|
||||
async def receive_stream(self) -> Any | None:
|
||||
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()
|
||||
|
||||
async def aresult(self) -> StateT:
|
||||
return await self._task
|
||||
|
||||
@@ -219,10 +239,12 @@ class _GraphEngineRun:
|
||||
async_channel_specs: dict[str, _ChannelSpec],
|
||||
entry_point: str,
|
||||
finish_point: str | None,
|
||||
stream_mode: str | None,
|
||||
) -> None:
|
||||
self._nodes = nodes
|
||||
self._entry_point = entry_point
|
||||
self._finish_point = finish_point
|
||||
self._stream_mode = stream_mode
|
||||
self._rust_engine = PyRustEngine()
|
||||
for name in async_channel_specs:
|
||||
self._rust_engine.add_async_channel(name)
|
||||
@@ -242,6 +264,7 @@ class _GraphEngineRun:
|
||||
finish_point,
|
||||
initial_state,
|
||||
self._execute_node_for_rust,
|
||||
self._stream_mode,
|
||||
)
|
||||
self._state = result_obj
|
||||
return cast(StateT, self._state)
|
||||
@@ -305,6 +328,15 @@ class _GraphEngineRun:
|
||||
def _publish_sync(self, channel: str, value: Any) -> None:
|
||||
self._rust_engine.publish_obj(channel, value)
|
||||
|
||||
def send_custom_stream_event(self, value: Any) -> None:
|
||||
self._rust_engine.send_custom_stream_event_obj(value)
|
||||
|
||||
def receive_stream_sync(self) -> Any | None:
|
||||
return self._rust_engine.receive_stream_obj()
|
||||
|
||||
def close_stream_sync(self) -> None:
|
||||
self._rust_engine.close_stream()
|
||||
|
||||
def _execute_node_for_rust(
|
||||
self, node_name: str, node_input: Any, state: Any
|
||||
) -> dict[str, Any]:
|
||||
@@ -350,14 +382,10 @@ class _GraphEngineRun:
|
||||
return event
|
||||
|
||||
def _run_awaitable_in_worker(self, awaitable: Coroutine[Any, Any, Any]) -> Any:
|
||||
loop = cast(
|
||||
asyncio.AbstractEventLoop | None,
|
||||
getattr(self._local, "worker_loop", None),
|
||||
)
|
||||
if loop is None or loop.is_closed():
|
||||
loop = asyncio.new_event_loop()
|
||||
self._local.worker_loop = loop
|
||||
return loop.run_until_complete(awaitable)
|
||||
# Create and close a dedicated loop per execution to avoid
|
||||
# interpreter-shutdown warnings from lingering thread-local loops.
|
||||
return asyncio.run(awaitable)
|
||||
|
||||
|
||||
def _normalize_result_to_sends(result: Any, *, default_input: Any) -> list[Send]:
|
||||
if result is None:
|
||||
@@ -542,3 +570,4 @@ def _invoke_node(node: Callable[..., Any], ctx: Context, node_input: Any, state:
|
||||
return node(node_input, state)
|
||||
|
||||
return node(ctx, node_input, state)
|
||||
|
||||
BIN
Binary file not shown.
@@ -0,0 +1 @@
|
||||
|
||||
@@ -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)
|
||||
@@ -0,0 +1,19 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Generic, TypeVar
|
||||
|
||||
N = TypeVar("N")
|
||||
|
||||
|
||||
@dataclass
|
||||
class Send(Generic[N]):
|
||||
node: N
|
||||
arg: Any = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class Command:
|
||||
update: Any = None
|
||||
goto: Any = field(default=None)
|
||||
|
||||
+10
-2
@@ -9,13 +9,13 @@ from typing import Any
|
||||
os.environ["LANGGRAPH_RUN_POOL_SIZE"] = "10"
|
||||
os.environ["LANGGRAPH_NODE_POOL_SIZE"] = "1000"
|
||||
|
||||
from langgraph.advanced_graph import (
|
||||
from saf_python_sdk.advanced_graph import (
|
||||
AdvancedStateGraph,
|
||||
channel_condition,
|
||||
timer_condition,
|
||||
)
|
||||
from saf_python_sdk.types import Command, Send
|
||||
from langgraph.graph import END, START, StateGraph
|
||||
from langgraph.types import Command, Send
|
||||
|
||||
RUNS = 100
|
||||
MIDDLE_COUNT = 10
|
||||
@@ -47,6 +47,7 @@ def build_advanced_parallel() -> Any:
|
||||
|
||||
graph.add_entry_node(start_node)
|
||||
for i in range(MIDDLE_COUNT):
|
||||
|
||||
async def middle_node(ctx: Any, state: dict[str, Any], idx: int = i) -> None:
|
||||
_ = idx
|
||||
_ = state
|
||||
@@ -71,6 +72,7 @@ def build_advanced_sequential() -> Any:
|
||||
return out
|
||||
|
||||
graph.add_entry_node(start_node)
|
||||
|
||||
def make_middle(target: str):
|
||||
async def middle_node(ctx: Any, state: dict[str, Any]) -> Command:
|
||||
_ = state
|
||||
@@ -99,6 +101,7 @@ def build_stategraph_parallel() -> Any:
|
||||
|
||||
graph.add_node("start_node", start_node)
|
||||
for i in range(MIDDLE_COUNT):
|
||||
|
||||
async def middle_node(state: dict[str, Any], idx: int = i) -> None:
|
||||
_ = idx
|
||||
_ = state
|
||||
@@ -128,6 +131,7 @@ def build_stategraph_sequential() -> Any:
|
||||
|
||||
graph.add_node("start_node", start_node)
|
||||
for i in range(MIDDLE_COUNT):
|
||||
|
||||
async def middle_node(state: dict[str, Any], idx: int = i) -> None:
|
||||
_ = idx
|
||||
_ = state
|
||||
@@ -164,6 +168,7 @@ def build_advanced_parallel_blocking() -> Any:
|
||||
|
||||
graph.add_entry_node(start_node)
|
||||
for i in range(MIDDLE_COUNT):
|
||||
|
||||
async def middle_blocking(
|
||||
ctx: Any, state: dict[str, Any], idx: int = i
|
||||
) -> None:
|
||||
@@ -223,6 +228,7 @@ def build_stategraph_parallel_blocking() -> Any:
|
||||
|
||||
graph.add_node("start_node", start_node)
|
||||
for i in range(MIDDLE_COUNT):
|
||||
|
||||
async def middle_blocking(state: dict[str, Any], idx: int = i) -> None:
|
||||
_ = idx
|
||||
_ = state
|
||||
@@ -252,6 +258,7 @@ def build_stategraph_sequential_blocking() -> Any:
|
||||
|
||||
graph.add_node("start_node", start_node)
|
||||
for i in range(MIDDLE_COUNT):
|
||||
|
||||
async def middle_blocking_seq(state: dict[str, Any], idx: int = i) -> None:
|
||||
_ = idx
|
||||
_ = state
|
||||
@@ -304,3 +311,4 @@ async def main() -> None:
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
|
||||
+3
-2
@@ -1,8 +1,8 @@
|
||||
import pytest
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.advanced_graph import AdvancedStateGraph, Context
|
||||
from langgraph.types import Command, Send
|
||||
from saf_python_sdk.advanced_graph import AdvancedStateGraph, Context
|
||||
from saf_python_sdk.types import Command, Send
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
@@ -64,3 +64,4 @@ async def test_run_ends_without_finish_node() -> None:
|
||||
assert result["counter"] == 8
|
||||
assert result["done"] == "stopped"
|
||||
assert result["logs"] == ["start", "middle:from_start"]
|
||||
|
||||
+3
-2
@@ -8,8 +8,8 @@ def test_run_pool_size_one_still_allows_parallel_runs() -> None:
|
||||
import asyncio
|
||||
import time
|
||||
from typing_extensions import TypedDict
|
||||
from langgraph.advanced_graph import AdvancedStateGraph, Context, timer_condition
|
||||
from langgraph.types import Command, Send
|
||||
from saf_python_sdk.advanced_graph import AdvancedStateGraph, Context, timer_condition
|
||||
from saf_python_sdk.types import Command, Send
|
||||
|
||||
|
||||
class RunState(TypedDict):
|
||||
@@ -53,3 +53,4 @@ asyncio.run(main())
|
||||
)
|
||||
elapsed = float(completed.stdout.strip().splitlines()[-1])
|
||||
assert elapsed < 0.35, completed.stdout
|
||||
|
||||
@@ -0,0 +1,56 @@
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from saf_python_sdk.advanced_graph import AdvancedStateGraph, Context
|
||||
from saf_python_sdk.types import Command, Send
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
|
||||
class StreamState(TypedDict):
|
||||
done: bool
|
||||
|
||||
|
||||
async def test_custom_stream_receive_and_close() -> None:
|
||||
graph: AdvancedStateGraph[StreamState] = AdvancedStateGraph(StreamState)
|
||||
|
||||
async def start_node(ctx: Context, state: StreamState) -> Command:
|
||||
ctx.send_custom_stream_event({"step": "start", "value": 1})
|
||||
await asyncio.sleep(0.08)
|
||||
ctx.send_custom_stream_event({"step": "start", "value": 2})
|
||||
return Command(update=state, goto=Send("finish_node", None))
|
||||
|
||||
async def finish_node(state: StreamState) -> dict[str, bool]:
|
||||
return {"done": True}
|
||||
|
||||
graph.add_entry_node(start_node)
|
||||
graph.add_finish_node(finish_node)
|
||||
|
||||
handler = await graph.compile().astart({"done": False}, stream_mode="custom")
|
||||
|
||||
event = await handler.receive_stream()
|
||||
assert isinstance(event, dict)
|
||||
assert event["step"] == "start"
|
||||
assert event["value"] == 1
|
||||
|
||||
handler.close_stream()
|
||||
assert await handler.receive_stream() is None
|
||||
|
||||
result = await handler.aresult()
|
||||
assert result["done"] is True
|
||||
|
||||
|
||||
async def test_only_custom_stream_mode_supported() -> None:
|
||||
graph: AdvancedStateGraph[StreamState] = AdvancedStateGraph(StreamState)
|
||||
|
||||
async def start_node(ctx: Context, state: StreamState) -> Command:
|
||||
ctx.send_custom_stream_event({"hello": "world"})
|
||||
return Command(update=state)
|
||||
|
||||
graph.add_entry_node(start_node)
|
||||
|
||||
handler = await graph.compile().astart({"done": False}, stream_mode="values")
|
||||
with pytest.raises(Exception, match="only `custom` is supported"):
|
||||
await handler.aresult()
|
||||
+4
-3
@@ -1,11 +1,12 @@
|
||||
import asyncio
|
||||
import json
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Literal
|
||||
|
||||
import pytest
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.advanced_graph import (
|
||||
from saf_python_sdk.advanced_graph import (
|
||||
AdvancedStateGraph,
|
||||
Context,
|
||||
any_of,
|
||||
@@ -14,7 +15,7 @@ from langgraph.advanced_graph import (
|
||||
)
|
||||
from langgraph.constants import END, START
|
||||
from langgraph.graph import StateGraph
|
||||
from langgraph.types import Command, Send
|
||||
from saf_python_sdk.types import Command, Send
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
@@ -212,6 +213,6 @@ async def test_async_sub_graph() -> None:
|
||||
order_food_idx = output.index("order_food: order submitted")
|
||||
assert first_sub_idx < second_sub_idx < order_food_idx
|
||||
assert llm._idx == len(llm.responses)
|
||||
import json
|
||||
|
||||
print(json.dumps(result, ensure_ascii=False, indent=2))
|
||||
|
||||
+5
-7
@@ -1,12 +1,12 @@
|
||||
import asyncio
|
||||
from dataclasses import dataclass
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
import pytest
|
||||
|
||||
from langgraph.advanced_graph import AdvancedStateGraph, CompiledGraphEngine
|
||||
from langgraph.types import Command, Send
|
||||
from saf_python_sdk.advanced_graph import AdvancedStateGraph, CompiledGraphEngine
|
||||
from saf_python_sdk.types import Command, Send
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
@@ -46,7 +46,6 @@ def _initial_state() -> UpdateElisionState:
|
||||
)
|
||||
|
||||
|
||||
|
||||
async def test_noop_slow_update_does_not_override_fast_update() -> None:
|
||||
graph: AdvancedStateGraph[UpdateElisionState] = AdvancedStateGraph(UpdateElisionState)
|
||||
|
||||
@@ -65,7 +64,6 @@ async def test_noop_slow_update_does_not_override_fast_update() -> None:
|
||||
|
||||
async def slow_node(state: UpdateElisionState) -> UpdateElisionState:
|
||||
await asyncio.sleep(0.1)
|
||||
# Returns the same values as the initial snapshot.
|
||||
return state
|
||||
|
||||
graph.add_entry_node(start_node)
|
||||
@@ -101,7 +99,6 @@ async def test_changed_slow_update_overrides_fast_update() -> None:
|
||||
|
||||
async def slow_node(state: UpdateElisionState) -> UpdateElisionState:
|
||||
await asyncio.sleep(0.1)
|
||||
# Slow node makes real changes for all field types.
|
||||
state.x = 2
|
||||
state.dc.value = 2
|
||||
state.model.value = 2
|
||||
@@ -124,3 +121,4 @@ async def test_changed_slow_update_overrides_fast_update() -> None:
|
||||
assert result.td == {"flag": False, "n": 2}
|
||||
assert result.obj == {"n": 2}
|
||||
assert result.items == [0, 1, 2]
|
||||
|
||||
Generated
+8
@@ -0,0 +1,8 @@
|
||||
version = 1
|
||||
revision = 3
|
||||
requires-python = ">=3.10"
|
||||
|
||||
[[package]]
|
||||
name = "saf-python-sdk"
|
||||
version = "0.1.2"
|
||||
source = { editable = "." }
|
||||
Reference in New Issue
Block a user