mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-20 06:35:46 +02:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
8b2bafd481 | ||
|
|
c535521e2d | ||
|
|
d26d4050c8 | ||
|
|
8f0f4a1be7 | ||
|
|
e13004da77 | ||
|
|
5a9264d124 | ||
|
|
efa86a6b14 | ||
|
|
f9866186d9 | ||
|
|
bfdd7deb60 | ||
|
|
09c8bb7c1c | ||
|
|
84d59adcd8 | ||
|
|
841ebf0c77 | ||
|
|
4f5b775819 | ||
|
|
76bee17ec4 | ||
|
|
d4335683e6 | ||
|
|
b43107ec99 | ||
|
|
3987f9cd63 | ||
|
|
2bbc3bb1da | ||
|
|
b3474d2db1 | ||
|
|
56e9fe1b10 | ||
|
|
1b3d075dbb | ||
|
|
c2a9661d2a | ||
|
|
9e723642c0 | ||
|
|
7c0a9275f9 | ||
|
|
783b3d3435 | ||
|
|
597b3402e6 | ||
|
|
54b723775c | ||
|
|
93c4a0a2d5 | ||
|
|
a5b90fbb95 | ||
|
|
85272db354 | ||
|
|
d7614999b0 | ||
|
|
af5f74f2ad | ||
|
|
708c0dff7f | ||
|
|
d50801d871 | ||
|
|
0327a86b80 | ||
|
|
13847a32c0 | ||
|
|
03ad7a011c | ||
|
|
896c1a8054 | ||
|
|
388d6b3593 | ||
|
|
667b679694 | ||
|
|
ccb9f4c41a | ||
|
|
789be99634 | ||
|
|
e33842ff54 | ||
|
|
89d37a0f9b | ||
|
|
8f717d3874 | ||
|
|
508193272a | ||
|
|
788ae6cb72 | ||
|
|
b2dca399e8 | ||
|
|
4b2167dd25 | ||
|
|
be46e91180 | ||
|
|
bbb308259b |
@@ -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" },
|
||||
]
|
||||
@@ -0,0 +1,377 @@
|
||||
package advancedgraph
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"reflect"
|
||||
"runtime"
|
||||
"strings"
|
||||
)
|
||||
|
||||
type nodeExecutor func(ctx *Context, input any, state map[string]any) (Command, error)
|
||||
|
||||
type AdvancedStateGraph[StateT any] struct {
|
||||
nodes map[string]nodeExecutor
|
||||
asyncChannels []string
|
||||
entryPoint string
|
||||
finishPoint string
|
||||
stateType reflect.Type
|
||||
}
|
||||
|
||||
func NewAdvancedStateGraph[StateT any]() *AdvancedStateGraph[StateT] {
|
||||
stateType := mustTypeOf[StateT]()
|
||||
if stateType.Kind() != reflect.Struct {
|
||||
panic(fmt.Sprintf("StateT must be a struct, got %s", stateType.String()))
|
||||
}
|
||||
return &AdvancedStateGraph[StateT]{
|
||||
nodes: make(map[string]nodeExecutor),
|
||||
stateType: stateType,
|
||||
}
|
||||
}
|
||||
|
||||
// AddNode keeps `fn` as `any` because advanced graph nodes can have different
|
||||
// input argument types per node, while only `StateT` is globally constrained.
|
||||
// We validate and adapt node signatures at runtime in compileNodeExecutor.
|
||||
func (g *AdvancedStateGraph[StateT]) AddNode(fn any) string {
|
||||
name := NodeName(fn)
|
||||
return g.AddNodeAs(name, fn)
|
||||
}
|
||||
|
||||
func (g *AdvancedStateGraph[StateT]) AddNodeAs(name string, fn any) string {
|
||||
if _, exists := g.nodes[name]; exists {
|
||||
panic(fmt.Sprintf("node `%s` already exists", name))
|
||||
}
|
||||
exec, err := compileNodeExecutor(fn, g.stateType)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
g.nodes[name] = exec
|
||||
return name
|
||||
}
|
||||
|
||||
func (g *AdvancedStateGraph[StateT]) AddAsyncChannel(name string) {
|
||||
g.asyncChannels = append(g.asyncChannels, name)
|
||||
}
|
||||
|
||||
func (g *AdvancedStateGraph[StateT]) AddEntryNode(fn any) string {
|
||||
name := NodeName(fn)
|
||||
return g.AddEntryNodeAs(name, fn)
|
||||
}
|
||||
|
||||
func (g *AdvancedStateGraph[StateT]) AddEntryNodeAs(name string, fn any) string {
|
||||
name = g.AddNodeAs(name, fn)
|
||||
g.entryPoint = name
|
||||
return name
|
||||
}
|
||||
|
||||
func (g *AdvancedStateGraph[StateT]) AddFinishNode(fn any) string {
|
||||
name := NodeName(fn)
|
||||
return g.AddFinishNodeAs(name, fn)
|
||||
}
|
||||
|
||||
func (g *AdvancedStateGraph[StateT]) AddFinishNodeAs(name string, fn any) string {
|
||||
name = g.AddNodeAs(name, fn)
|
||||
g.finishPoint = name
|
||||
return name
|
||||
}
|
||||
|
||||
func (g *AdvancedStateGraph[StateT]) Compile() *CompiledGraph[StateT] {
|
||||
return &CompiledGraph[StateT]{
|
||||
nodes: g.nodes,
|
||||
asyncChannels: g.asyncChannels,
|
||||
entryPoint: g.entryPoint,
|
||||
finishPoint: g.finishPoint,
|
||||
stateType: g.stateType,
|
||||
}
|
||||
}
|
||||
|
||||
type CompiledGraph[StateT any] struct {
|
||||
nodes map[string]nodeExecutor
|
||||
asyncChannels []string
|
||||
entryPoint string
|
||||
finishPoint string
|
||||
stateType reflect.Type
|
||||
}
|
||||
|
||||
type Context struct {
|
||||
engine *RustEngine
|
||||
resumeEvent *WaitEvent
|
||||
}
|
||||
|
||||
func (c *Context) WaitFor(cond AnyOfCondition) (WaitEvent, error) {
|
||||
if c.resumeEvent != nil {
|
||||
event := *c.resumeEvent
|
||||
c.resumeEvent = nil
|
||||
return event, nil
|
||||
}
|
||||
return WaitEvent{}, ErrWaitRequested{Condition: cond}
|
||||
}
|
||||
|
||||
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]
|
||||
}
|
||||
|
||||
type resultOrErr[StateT any] struct {
|
||||
state StateT
|
||||
err error
|
||||
}
|
||||
|
||||
func (h *Handler[StateT]) PublishToChannel(channel string, value any) error {
|
||||
return h.engine.Publish(channel, value)
|
||||
}
|
||||
|
||||
func (h *Handler[StateT]) WaitForResult() (StateT, error) {
|
||||
res := <-h.done
|
||||
return res.state, res.err
|
||||
}
|
||||
|
||||
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 {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
handler := &Handler[StateT]{
|
||||
engine: engine,
|
||||
done: make(chan resultOrErr[StateT], 1),
|
||||
}
|
||||
go func() {
|
||||
defer engine.Close()
|
||||
rawState, err := engine.RunGraph(
|
||||
g.entryPoint,
|
||||
g.finishPoint,
|
||||
resolvedStreamMode,
|
||||
initialState,
|
||||
initialInput,
|
||||
func(node string, nodeInput any, fallbackState map[string]any) (Command, error) {
|
||||
fn, ok := g.nodes[node]
|
||||
if !ok {
|
||||
return Command{}, fmt.Errorf("unknown node `%s`", node)
|
||||
}
|
||||
if fallbackState == nil {
|
||||
return Command{}, fmt.Errorf("node `%s` expected map state argument", node)
|
||||
}
|
||||
resolvedInput, resumeEvent := unwrapResumeInput(nodeInput)
|
||||
return fn(&Context{engine: engine, resumeEvent: resumeEvent}, resolvedInput, fallbackState)
|
||||
},
|
||||
)
|
||||
if err != nil {
|
||||
handler.done <- resultOrErr[StateT]{err: err}
|
||||
close(handler.done)
|
||||
return
|
||||
}
|
||||
state, err := mapToState[StateT](rawState)
|
||||
handler.done <- resultOrErr[StateT]{state: state, err: err}
|
||||
close(handler.done)
|
||||
}()
|
||||
return handler, nil
|
||||
}
|
||||
|
||||
func NodeName(fn any) string {
|
||||
rv := reflect.ValueOf(fn)
|
||||
if !rv.IsValid() || rv.Kind() != reflect.Func {
|
||||
panic("cannot infer node name from non-function value")
|
||||
}
|
||||
pc := rv.Pointer()
|
||||
f := runtime.FuncForPC(pc)
|
||||
if f == nil {
|
||||
panic("cannot infer node name from nil function")
|
||||
}
|
||||
full := f.Name()
|
||||
if strings.Contains(full, ".func") {
|
||||
panic("anonymous functions are not allowed as nodes")
|
||||
}
|
||||
short := full
|
||||
if i := strings.LastIndex(short, "/"); i >= 0 {
|
||||
short = short[i+1:]
|
||||
}
|
||||
if i := strings.LastIndex(short, "."); i >= 0 {
|
||||
short = short[i+1:]
|
||||
}
|
||||
short = strings.TrimSuffix(short, "-fm")
|
||||
if short == "" || strings.Contains(short, "func") {
|
||||
panic(fmt.Sprintf("cannot infer stable node name from `%s`", full))
|
||||
}
|
||||
return short
|
||||
}
|
||||
|
||||
func compileNodeExecutor(fn any, expectedStateType reflect.Type) (nodeExecutor, error) {
|
||||
rv := reflect.ValueOf(fn)
|
||||
if !rv.IsValid() || rv.Kind() != reflect.Func {
|
||||
return nil, fmt.Errorf("node must be a function")
|
||||
}
|
||||
rt := rv.Type()
|
||||
if rt.NumIn() != 3 {
|
||||
return nil, fmt.Errorf("node `%s` must accept exactly 3 args: (*Context, input, state)", NodeName(fn))
|
||||
}
|
||||
ctxType := reflect.TypeOf((*Context)(nil))
|
||||
if rt.In(0) != ctxType {
|
||||
return nil, fmt.Errorf("node `%s` first arg must be *Context", NodeName(fn))
|
||||
}
|
||||
if rt.NumOut() != 2 {
|
||||
return nil, fmt.Errorf("node `%s` must return (Command, error)", NodeName(fn))
|
||||
}
|
||||
cmdType := reflect.TypeOf(Command{})
|
||||
if rt.Out(0) != cmdType {
|
||||
return nil, fmt.Errorf("node `%s` first return must be Command", NodeName(fn))
|
||||
}
|
||||
errType := reflect.TypeOf((*error)(nil)).Elem()
|
||||
if !rt.Out(1).Implements(errType) {
|
||||
return nil, fmt.Errorf("node `%s` second return must be error", NodeName(fn))
|
||||
}
|
||||
|
||||
inputType := rt.In(1)
|
||||
stateType := rt.In(2)
|
||||
if stateType != expectedStateType {
|
||||
return nil, fmt.Errorf(
|
||||
"node `%s` state type mismatch: got %s, graph expects %s",
|
||||
NodeName(fn),
|
||||
stateType.String(),
|
||||
expectedStateType.String(),
|
||||
)
|
||||
}
|
||||
return func(ctx *Context, input any, state map[string]any) (Command, error) {
|
||||
stateArg, err := convertStateArg(state, stateType)
|
||||
if err != nil {
|
||||
return Command{}, fmt.Errorf("node `%s` state decode failed: %w", NodeName(fn), err)
|
||||
}
|
||||
args := []reflect.Value{
|
||||
reflect.ValueOf(ctx),
|
||||
reflect.Zero(inputType),
|
||||
stateArg,
|
||||
}
|
||||
if input != nil {
|
||||
inVal := reflect.ValueOf(input)
|
||||
if inVal.Type().AssignableTo(inputType) {
|
||||
args[1] = inVal
|
||||
} else if inVal.Type().ConvertibleTo(inputType) {
|
||||
args[1] = inVal.Convert(inputType)
|
||||
} else {
|
||||
return Command{}, fmt.Errorf(
|
||||
"node `%s` input type mismatch: got %T, want %s",
|
||||
NodeName(fn),
|
||||
input,
|
||||
inputType.String(),
|
||||
)
|
||||
}
|
||||
}
|
||||
out := rv.Call(args)
|
||||
cmd := out[0].Interface().(Command)
|
||||
if cmd.Update != nil {
|
||||
updateType := reflect.TypeOf(cmd.Update)
|
||||
if updateType != stateType {
|
||||
return Command{}, fmt.Errorf(
|
||||
"node `%s` update type mismatch: got %s, graph expects %s",
|
||||
NodeName(fn),
|
||||
updateType.String(),
|
||||
stateType.String(),
|
||||
)
|
||||
}
|
||||
}
|
||||
if out[1].IsNil() {
|
||||
return cmd, nil
|
||||
}
|
||||
return cmd, out[1].Interface().(error)
|
||||
}, nil
|
||||
}
|
||||
|
||||
func convertStateArg(state map[string]any, stateType reflect.Type) (reflect.Value, error) {
|
||||
if stateType == reflect.TypeOf(map[string]any{}) {
|
||||
return reflect.ValueOf(state), nil
|
||||
}
|
||||
raw, err := json.Marshal(state)
|
||||
if err != nil {
|
||||
return reflect.Value{}, fmt.Errorf("marshal state: %w", err)
|
||||
}
|
||||
if stateType.Kind() == reflect.Ptr {
|
||||
target := reflect.New(stateType.Elem())
|
||||
if err := json.Unmarshal(raw, target.Interface()); err != nil {
|
||||
return reflect.Value{}, fmt.Errorf("unmarshal state into %s: %w", stateType.String(), err)
|
||||
}
|
||||
return target, nil
|
||||
}
|
||||
target := reflect.New(stateType)
|
||||
if err := json.Unmarshal(raw, target.Interface()); err != nil {
|
||||
return reflect.Value{}, fmt.Errorf("unmarshal state into %s: %w", stateType.String(), err)
|
||||
}
|
||||
return target.Elem(), nil
|
||||
}
|
||||
|
||||
func mapToState[StateT any](raw map[string]any) (StateT, error) {
|
||||
var out StateT
|
||||
if anyVal, ok := any(raw).(StateT); ok {
|
||||
return anyVal, nil
|
||||
}
|
||||
payload, err := json.Marshal(raw)
|
||||
if err != nil {
|
||||
return out, fmt.Errorf("marshal state: %w", err)
|
||||
}
|
||||
if err := json.Unmarshal(payload, &out); err != nil {
|
||||
return out, fmt.Errorf("unmarshal state: %w", err)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func mustTypeOf[T any]() reflect.Type {
|
||||
var zero T
|
||||
t := reflect.TypeOf(zero)
|
||||
if t != nil {
|
||||
return t
|
||||
}
|
||||
// Handles nil-able types where zero value has no dynamic type.
|
||||
return reflect.TypeOf((*T)(nil)).Elem()
|
||||
}
|
||||
|
||||
func unwrapResumeInput(input any) (any, *WaitEvent) {
|
||||
wrapper, ok := input.(map[string]any)
|
||||
if !ok {
|
||||
return input, nil
|
||||
}
|
||||
rawArg, hasArg := wrapper["__lg_resume_arg__"]
|
||||
rawEvent, hasEvent := wrapper["__lg_resume_event__"]
|
||||
if !hasArg || !hasEvent {
|
||||
return input, nil
|
||||
}
|
||||
eventPayload, err := json.Marshal(rawEvent)
|
||||
if err != nil {
|
||||
return rawArg, nil
|
||||
}
|
||||
var event WaitEvent
|
||||
if err := json.Unmarshal(eventPayload, &event); err != nil {
|
||||
return rawArg, nil
|
||||
}
|
||||
return rawArg, &event
|
||||
}
|
||||
@@ -0,0 +1,371 @@
|
||||
package advancedgraph
|
||||
|
||||
/*
|
||||
#cgo CFLAGS: -I${SRCDIR}/../../rust-core/include
|
||||
#cgo LDFLAGS: -L${SRCDIR}/../../rust-core/target/debug -llanggraph_rust_core
|
||||
#include "langgraph_rust_core.h"
|
||||
#include <stdlib.h>
|
||||
extern char* goNodeCallback(unsigned long user_data, char* node, char* arg_json, char* state_json);
|
||||
*/
|
||||
import "C"
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"reflect"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"unsafe"
|
||||
)
|
||||
|
||||
type RustEngine struct {
|
||||
ptr *C.Engine
|
||||
}
|
||||
|
||||
type runGraphCallbackCtx struct {
|
||||
exec func(node string, nodeInput any, state map[string]any) (Command, error)
|
||||
}
|
||||
|
||||
var (
|
||||
callbackRegistryMu sync.RWMutex
|
||||
callbackRegistry = map[uint64]*runGraphCallbackCtx{}
|
||||
callbackNextID uint64
|
||||
)
|
||||
|
||||
func registerRunGraphCallbackCtx(ctx *runGraphCallbackCtx) uint64 {
|
||||
id := atomic.AddUint64(&callbackNextID, 1)
|
||||
callbackRegistryMu.Lock()
|
||||
callbackRegistry[id] = ctx
|
||||
callbackRegistryMu.Unlock()
|
||||
return id
|
||||
}
|
||||
|
||||
func unregisterRunGraphCallbackCtx(id uint64) {
|
||||
callbackRegistryMu.Lock()
|
||||
delete(callbackRegistry, id)
|
||||
callbackRegistryMu.Unlock()
|
||||
}
|
||||
|
||||
func getRunGraphCallbackCtx(id uint64) (*runGraphCallbackCtx, bool) {
|
||||
callbackRegistryMu.RLock()
|
||||
ctx, ok := callbackRegistry[id]
|
||||
callbackRegistryMu.RUnlock()
|
||||
return ctx, ok
|
||||
}
|
||||
|
||||
//export goNodeCallback
|
||||
func goNodeCallback(userData C.ulong, node *C.char, argJSON *C.char, stateJSON *C.char) *C.char {
|
||||
ctx, ok := getRunGraphCallbackCtx(uint64(userData))
|
||||
if !ok {
|
||||
return cCallbackEnvelopeError("invalid callback context (possibly stale callback)")
|
||||
}
|
||||
|
||||
nodeName := C.GoString(node)
|
||||
|
||||
var nodeInput any
|
||||
if err := json.Unmarshal([]byte(C.GoString(argJSON)), &nodeInput); err != nil {
|
||||
return cCallbackEnvelopeError(fmt.Sprintf("decode arg failed for `%s`: %v", nodeName, err))
|
||||
}
|
||||
var state map[string]any
|
||||
if err := json.Unmarshal([]byte(C.GoString(stateJSON)), &state); err != nil {
|
||||
return cCallbackEnvelopeError(fmt.Sprintf("decode state failed for `%s`: %v", nodeName, err))
|
||||
}
|
||||
nodeInput = coerceJSONValue(nodeInput)
|
||||
stateAny := coerceJSONValue(state)
|
||||
state, ok = stateAny.(map[string]any)
|
||||
if !ok {
|
||||
return cCallbackEnvelopeError(fmt.Sprintf("decoded state has unexpected type for `%s`", nodeName))
|
||||
}
|
||||
|
||||
cmd, err := ctx.exec(nodeName, nodeInput, state)
|
||||
if err != nil {
|
||||
if waitReq, ok := AsErrWaitRequested(err); ok {
|
||||
return cCallbackEnvelopeSuspend(waitReq.Condition)
|
||||
}
|
||||
return cCallbackEnvelopeError(err.Error())
|
||||
}
|
||||
|
||||
sends := make([]map[string]any, 0, len(cmd.Goto))
|
||||
for _, send := range cmd.Goto {
|
||||
targetNode, err := resolveSendTarget(send.Node)
|
||||
if err != nil {
|
||||
return cCallbackEnvelopeError(err.Error())
|
||||
}
|
||||
sends = append(sends, map[string]any{
|
||||
"node": targetNode,
|
||||
"arg": send.NodeInput,
|
||||
})
|
||||
}
|
||||
payload := map[string]any{
|
||||
"update": cmd.Update,
|
||||
"sends": sends,
|
||||
}
|
||||
raw, err := json.Marshal(map[string]any{
|
||||
"ok": true,
|
||||
"payload": payload,
|
||||
})
|
||||
if err != nil {
|
||||
return cCallbackEnvelopeError(fmt.Sprintf("encode callback payload failed: %v", err))
|
||||
}
|
||||
return C.CString(string(raw))
|
||||
}
|
||||
|
||||
func NewRustEngine() *RustEngine {
|
||||
return &RustEngine{ptr: C.rc_engine_new()}
|
||||
}
|
||||
|
||||
func (e *RustEngine) Close() {
|
||||
if e.ptr != nil {
|
||||
C.rc_engine_free(e.ptr)
|
||||
e.ptr = nil
|
||||
}
|
||||
}
|
||||
|
||||
func (e *RustEngine) AddAsyncChannel(channel string) error {
|
||||
cch := C.CString(channel)
|
||||
defer C.free(unsafe.Pointer(cch))
|
||||
resp := C.rc_add_async_channel(e.ptr, cch)
|
||||
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 {
|
||||
return fmt.Errorf("marshal publish value: %w", err)
|
||||
}
|
||||
cch := C.CString(channel)
|
||||
cval := C.CString(string(payload))
|
||||
defer C.free(unsafe.Pointer(cch))
|
||||
defer C.free(unsafe.Pointer(cval))
|
||||
resp := C.rc_publish_json(e.ptr, cch, cval)
|
||||
return parseRustStatus(resp)
|
||||
}
|
||||
|
||||
func (e *RustEngine) WaitAnyOf(cond AnyOfCondition) (WaitEvent, error) {
|
||||
payload, err := json.Marshal(cond)
|
||||
if err != nil {
|
||||
return WaitEvent{}, fmt.Errorf("marshal any_of: %w", err)
|
||||
}
|
||||
cpayload := C.CString(string(payload))
|
||||
defer C.free(unsafe.Pointer(cpayload))
|
||||
resp := C.rc_wait_any_of_json(e.ptr, cpayload)
|
||||
defer C.rc_string_free(resp)
|
||||
|
||||
raw := C.GoString(resp)
|
||||
var status struct {
|
||||
OK bool `json:"ok"`
|
||||
Error string `json:"error"`
|
||||
Event json.RawMessage `json:"event"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(raw), &status); err != nil {
|
||||
return WaitEvent{}, fmt.Errorf("decode rust wait response: %w", err)
|
||||
}
|
||||
if !status.OK {
|
||||
return WaitEvent{}, fmt.Errorf("rust wait failed: %s", status.Error)
|
||||
}
|
||||
var event WaitEvent
|
||||
if err := json.Unmarshal(status.Event, &event); err != nil {
|
||||
return WaitEvent{}, fmt.Errorf("decode wait event: %w", err)
|
||||
}
|
||||
return event, nil
|
||||
}
|
||||
|
||||
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),
|
||||
) (map[string]any, error) {
|
||||
initialJSON, err := json.Marshal(initialState)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("marshal initial state: %w", err)
|
||||
}
|
||||
initialInputJSON, err := json.Marshal(initialInput)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("marshal initial input: %w", err)
|
||||
}
|
||||
centry := C.CString(entryPoint)
|
||||
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)
|
||||
|
||||
resp := C.rc_run_graph_json(
|
||||
e.ptr,
|
||||
centry,
|
||||
cfinish,
|
||||
cinitial,
|
||||
cinitialInput,
|
||||
cstreamMode,
|
||||
C.ulong(callbackID),
|
||||
(C.rc_node_callback_t)(C.goNodeCallback),
|
||||
)
|
||||
defer C.rc_string_free(resp)
|
||||
|
||||
raw := C.GoString(resp)
|
||||
var status struct {
|
||||
OK bool `json:"ok"`
|
||||
Error string `json:"error"`
|
||||
State map[string]any `json:"state"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(raw), &status); err != nil {
|
||||
return nil, fmt.Errorf("decode rust run response: %w", err)
|
||||
}
|
||||
if !status.OK {
|
||||
return nil, fmt.Errorf("rust run failed: %s", status.Error)
|
||||
}
|
||||
coerced := coerceJSONValue(status.State)
|
||||
typed, ok := coerced.(map[string]any)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("unexpected state type from rust run")
|
||||
}
|
||||
return typed, nil
|
||||
}
|
||||
|
||||
func cCallbackEnvelopeError(message string) *C.char {
|
||||
raw, _ := json.Marshal(map[string]any{
|
||||
"ok": false,
|
||||
"error": message,
|
||||
})
|
||||
return C.CString(string(raw))
|
||||
}
|
||||
|
||||
func cCallbackEnvelopeSuspend(cond AnyOfCondition) *C.char {
|
||||
raw, _ := json.Marshal(map[string]any{
|
||||
"ok": true,
|
||||
"suspend": map[string]any{
|
||||
"kind": "any_of",
|
||||
"any_of": cond,
|
||||
},
|
||||
})
|
||||
return C.CString(string(raw))
|
||||
}
|
||||
|
||||
func resolveSendTarget(target any) (string, error) {
|
||||
if name, ok := target.(string); ok {
|
||||
if name == "" {
|
||||
return "", fmt.Errorf("send target cannot be empty string")
|
||||
}
|
||||
return name, nil
|
||||
}
|
||||
rv := reflect.ValueOf(target)
|
||||
if rv.IsValid() && rv.Kind() == reflect.Func {
|
||||
return NodeName(target), nil
|
||||
}
|
||||
return "", fmt.Errorf("unsupported send target type %T", target)
|
||||
}
|
||||
|
||||
func coerceJSONValue(v any) any {
|
||||
switch t := v.(type) {
|
||||
case map[string]any:
|
||||
out := make(map[string]any, len(t))
|
||||
for k, val := range t {
|
||||
out[k] = coerceJSONValue(val)
|
||||
}
|
||||
return out
|
||||
case []any:
|
||||
coerced := make([]any, len(t))
|
||||
allStrings := true
|
||||
for i, val := range t {
|
||||
cv := coerceJSONValue(val)
|
||||
coerced[i] = cv
|
||||
if _, ok := cv.(string); !ok {
|
||||
allStrings = false
|
||||
}
|
||||
}
|
||||
if allStrings {
|
||||
out := make([]string, len(coerced))
|
||||
for i, item := range coerced {
|
||||
out[i] = item.(string)
|
||||
}
|
||||
return out
|
||||
}
|
||||
return coerced
|
||||
default:
|
||||
return v
|
||||
}
|
||||
}
|
||||
|
||||
func parseRustStatus(resp *C.char) error {
|
||||
defer C.rc_string_free(resp)
|
||||
raw := C.GoString(resp)
|
||||
var status struct {
|
||||
OK bool `json:"ok"`
|
||||
Error string `json:"error"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(raw), &status); err != nil {
|
||||
return fmt.Errorf("decode rust response: %w", err)
|
||||
}
|
||||
if !status.OK {
|
||||
return fmt.Errorf("rust error: %s", status.Error)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,89 @@
|
||||
package advancedgraph
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
)
|
||||
|
||||
type WaitCondition interface {
|
||||
toAny() map[string]any
|
||||
}
|
||||
|
||||
type ChannelCondition struct {
|
||||
Channel string
|
||||
N int
|
||||
}
|
||||
|
||||
func (c ChannelCondition) toAny() map[string]any {
|
||||
n := c.N
|
||||
if n <= 0 {
|
||||
n = 1
|
||||
}
|
||||
return map[string]any{
|
||||
"kind": "channel",
|
||||
"channel": c.Channel,
|
||||
"n": n,
|
||||
}
|
||||
}
|
||||
|
||||
type TimerCondition struct {
|
||||
Seconds float64
|
||||
}
|
||||
|
||||
func (t TimerCondition) toAny() map[string]any {
|
||||
return map[string]any{
|
||||
"kind": "timer",
|
||||
"seconds": t.Seconds,
|
||||
}
|
||||
}
|
||||
|
||||
type AnyOfCondition struct {
|
||||
Conditions []map[string]any `json:"conditions"`
|
||||
}
|
||||
|
||||
func AnyOf(conditions ...WaitCondition) AnyOfCondition {
|
||||
result := AnyOfCondition{Conditions: make([]map[string]any, 0, len(conditions))}
|
||||
for _, cond := range conditions {
|
||||
result.Conditions = append(result.Conditions, cond.toAny())
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
type WaitEvent struct {
|
||||
Condition string `json:"condition"`
|
||||
Channel string `json:"channel,omitempty"`
|
||||
Value json.RawMessage `json:"value,omitempty"`
|
||||
Seconds float64 `json:"seconds,omitempty"`
|
||||
}
|
||||
|
||||
type Send struct {
|
||||
Node any
|
||||
NodeInput any
|
||||
}
|
||||
|
||||
type Command struct {
|
||||
Update any
|
||||
Goto []Send
|
||||
}
|
||||
|
||||
type ErrWaitRequested struct {
|
||||
Condition AnyOfCondition
|
||||
}
|
||||
|
||||
func (e ErrWaitRequested) Error() string {
|
||||
return "wait requested"
|
||||
}
|
||||
|
||||
func AsErrWaitRequested(err error) (ErrWaitRequested, bool) {
|
||||
var target ErrWaitRequested
|
||||
if !errors.As(err, &target) {
|
||||
return target, false
|
||||
}
|
||||
return target, true
|
||||
}
|
||||
|
||||
func DecodeString(raw json.RawMessage) string {
|
||||
var s string
|
||||
_ = json.Unmarshal(raw, &s)
|
||||
return s
|
||||
}
|
||||
@@ -0,0 +1,3 @@
|
||||
module github.com/langchain-ai/langgraph/langgraph-go
|
||||
|
||||
go 1.25
|
||||
@@ -0,0 +1,320 @@
|
||||
package stategraph
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"slices"
|
||||
|
||||
ag "github.com/langchain-ai/langgraph/langgraph-go/advancedgraph"
|
||||
)
|
||||
|
||||
type StateNodeFunc[StateT any] func(ctx *Context, state StateT) (StateT, error)
|
||||
|
||||
const (
|
||||
internalBarrierChannel = "__stategraph_barrier"
|
||||
internalInterruptChannel = "__stategraph_interrupt"
|
||||
)
|
||||
|
||||
type Context struct {
|
||||
inner *ag.Context
|
||||
}
|
||||
|
||||
func (c *Context) Interrupt(name string) (any, error) {
|
||||
if name == "" {
|
||||
return nil, fmt.Errorf("interrupt name cannot be empty")
|
||||
}
|
||||
event, err := c.inner.WaitFor(ag.AnyOf(ag.ChannelCondition{
|
||||
Channel: internalInterruptChannel,
|
||||
N: 1,
|
||||
}))
|
||||
if err != nil {
|
||||
if waitReq, ok := ag.AsErrWaitRequested(err); ok {
|
||||
return nil, errInterruptRequested{
|
||||
Name: name,
|
||||
Condition: waitReq.Condition,
|
||||
}
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
if len(event.Value) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
var payload interruptPayload
|
||||
if err := json.Unmarshal(event.Value, &payload); err != nil {
|
||||
var value any
|
||||
if err := json.Unmarshal(event.Value, &value); err != nil {
|
||||
return nil, fmt.Errorf("decode interrupt `%s` value: %w", name, err)
|
||||
}
|
||||
return value, nil
|
||||
}
|
||||
if payload.Name != "" && payload.Name != name {
|
||||
return nil, fmt.Errorf("interrupt name mismatch: expected `%s`, got `%s`", name, payload.Name)
|
||||
}
|
||||
if len(payload.Value) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
var value any
|
||||
if err := json.Unmarshal(payload.Value, &value); err != nil {
|
||||
return nil, fmt.Errorf("decode interrupt `%s` payload: %w", name, err)
|
||||
}
|
||||
return value, nil
|
||||
}
|
||||
|
||||
type errInterruptRequested struct {
|
||||
Name string
|
||||
Condition ag.AnyOfCondition
|
||||
}
|
||||
|
||||
func (e errInterruptRequested) Error() string {
|
||||
if e.Name == "" {
|
||||
return "interrupt requested"
|
||||
}
|
||||
return fmt.Sprintf("interrupt requested: %s", e.Name)
|
||||
}
|
||||
|
||||
func asErrInterruptRequested(err error) (errInterruptRequested, bool) {
|
||||
var target errInterruptRequested
|
||||
if !errors.As(err, &target) {
|
||||
return target, false
|
||||
}
|
||||
return target, true
|
||||
}
|
||||
|
||||
type BasicStateGraph[StateT any] struct {
|
||||
nodes map[string]StateNodeFunc[StateT]
|
||||
edges map[string][]string
|
||||
}
|
||||
|
||||
type interruptPayload struct {
|
||||
Name string `json:"name"`
|
||||
Value json.RawMessage `json:"value"`
|
||||
}
|
||||
|
||||
func NewBasicStateGraph[StateT any]() *BasicStateGraph[StateT] {
|
||||
return &BasicStateGraph[StateT]{
|
||||
nodes: make(map[string]StateNodeFunc[StateT]),
|
||||
edges: make(map[string][]string),
|
||||
}
|
||||
}
|
||||
|
||||
func (g *BasicStateGraph[StateT]) AddNode(fn StateNodeFunc[StateT]) string {
|
||||
name := ag.NodeName(fn)
|
||||
if _, exists := g.nodes[name]; exists {
|
||||
panic(fmt.Sprintf("node `%s` already exists", name))
|
||||
}
|
||||
g.nodes[name] = fn
|
||||
return name
|
||||
}
|
||||
|
||||
func (g *BasicStateGraph[StateT]) AddEdge(from StateNodeFunc[StateT], to StateNodeFunc[StateT]) {
|
||||
fromName := ag.NodeName(from)
|
||||
toName := ag.NodeName(to)
|
||||
if _, ok := g.nodes[fromName]; !ok {
|
||||
panic(fmt.Sprintf("source node `%s` does not exist", fromName))
|
||||
}
|
||||
if _, ok := g.nodes[toName]; !ok {
|
||||
panic(fmt.Sprintf("target node `%s` does not exist", toName))
|
||||
}
|
||||
g.edges[fromName] = append(g.edges[fromName], toName)
|
||||
}
|
||||
|
||||
type CompiledBasicStateGraph[StateT any] struct {
|
||||
inner *ag.CompiledGraph[StateT]
|
||||
}
|
||||
|
||||
type Handler[StateT any] struct {
|
||||
inner *ag.Handler[StateT]
|
||||
}
|
||||
|
||||
func (h *Handler[StateT]) WaitForResult() (StateT, error) {
|
||||
return h.inner.WaitForResult()
|
||||
}
|
||||
|
||||
func (h *Handler[StateT]) Resume(name string, value any) error {
|
||||
if name == "" {
|
||||
return fmt.Errorf("interrupt name cannot be empty")
|
||||
}
|
||||
return h.inner.PublishToChannel(internalInterruptChannel, map[string]any{
|
||||
"name": name,
|
||||
"value": value,
|
||||
})
|
||||
}
|
||||
|
||||
func (g *BasicStateGraph[StateT]) Compile() *CompiledBasicStateGraph[StateT] {
|
||||
if len(g.nodes) == 0 {
|
||||
panic("graph has no nodes")
|
||||
}
|
||||
levels, err := g.computeSupersteps()
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
adv := ag.NewAdvancedStateGraph[StateT]()
|
||||
adv.AddAsyncChannel(internalBarrierChannel)
|
||||
adv.AddAsyncChannel(internalInterruptChannel)
|
||||
|
||||
const finalNodeName = "__stategraph_finish"
|
||||
finalNode := func(_ *ag.Context, _ any, state StateT) (ag.Command, error) {
|
||||
return ag.Command{Update: state}, nil
|
||||
}
|
||||
adv.AddFinishNodeAs(finalNodeName, finalNode)
|
||||
|
||||
for stepIdx, stepNodes := range levels {
|
||||
for _, nodeName := range stepNodes {
|
||||
userFn := g.nodes[nodeName]
|
||||
nextBarrier := fmt.Sprintf("__stategraph_barrier_%d", stepIdx+1)
|
||||
wrapper := func(ctx *ag.Context, _ any, state StateT) (ag.Command, error) {
|
||||
updated, err := userFn(&Context{inner: ctx}, state)
|
||||
if err != nil {
|
||||
if interruptReq, ok := asErrInterruptRequested(err); ok {
|
||||
cond := interruptReq.Condition
|
||||
if len(cond.Conditions) == 0 {
|
||||
cond = ag.AnyOf(ag.ChannelCondition{
|
||||
Channel: internalInterruptChannel,
|
||||
N: 1,
|
||||
})
|
||||
}
|
||||
return ag.Command{}, ag.ErrWaitRequested{Condition: cond}
|
||||
}
|
||||
return ag.Command{}, err
|
||||
}
|
||||
if err := ctx.PublishToChannel(internalBarrierChannel, map[string]any{
|
||||
"step": stepIdx,
|
||||
}); err != nil {
|
||||
return ag.Command{}, err
|
||||
}
|
||||
return ag.Command{
|
||||
Update: updated,
|
||||
Goto: []ag.Send{{Node: nextBarrier}},
|
||||
}, nil
|
||||
}
|
||||
adv.AddNodeAs(fmt.Sprintf("__stategraph_node_%s", nodeName), wrapper)
|
||||
}
|
||||
}
|
||||
|
||||
lastBarrier := len(levels)
|
||||
for barrierStep := 0; barrierStep <= lastBarrier; barrierStep++ {
|
||||
barrierName := fmt.Sprintf("__stategraph_barrier_%d", barrierStep)
|
||||
nextStep := barrierStep
|
||||
barrier := func(ctx *ag.Context, _ any, state StateT) (ag.Command, error) {
|
||||
if nextStep > 0 {
|
||||
needed := len(levels[nextStep-1])
|
||||
_, err := ctx.WaitFor(ag.AnyOf(ag.ChannelCondition{
|
||||
Channel: internalBarrierChannel,
|
||||
N: needed,
|
||||
}))
|
||||
if err != nil {
|
||||
return ag.Command{}, err
|
||||
}
|
||||
}
|
||||
if nextStep >= len(levels) {
|
||||
return ag.Command{
|
||||
Update: state,
|
||||
Goto: []ag.Send{{Node: finalNodeName}},
|
||||
}, nil
|
||||
}
|
||||
sends := make([]ag.Send, 0, len(levels[nextStep]))
|
||||
for _, nodeName := range levels[nextStep] {
|
||||
sends = append(sends, ag.Send{
|
||||
Node: fmt.Sprintf("__stategraph_node_%s", nodeName),
|
||||
})
|
||||
}
|
||||
return ag.Command{
|
||||
Update: state,
|
||||
Goto: sends,
|
||||
}, nil
|
||||
}
|
||||
if barrierStep == 0 {
|
||||
adv.AddEntryNodeAs(barrierName, barrier)
|
||||
} else {
|
||||
adv.AddNodeAs(barrierName, barrier)
|
||||
}
|
||||
}
|
||||
|
||||
return &CompiledBasicStateGraph[StateT]{
|
||||
inner: adv.Compile(),
|
||||
}
|
||||
}
|
||||
|
||||
func (g *CompiledBasicStateGraph[StateT]) Start(initialState StateT) (*Handler[StateT], error) {
|
||||
raw, err := g.inner.Start(nil, initialState)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &Handler[StateT]{inner: raw}, nil
|
||||
}
|
||||
|
||||
func (g *CompiledBasicStateGraph[StateT]) Invoke(initialState StateT) (StateT, error) {
|
||||
handler, err := g.Start(initialState)
|
||||
if err != nil {
|
||||
var zero StateT
|
||||
return zero, err
|
||||
}
|
||||
return handler.WaitForResult()
|
||||
}
|
||||
|
||||
func (g *BasicStateGraph[StateT]) computeSupersteps() ([][]string, error) {
|
||||
indegree := make(map[string]int, len(g.nodes))
|
||||
for name := range g.nodes {
|
||||
indegree[name] = 0
|
||||
}
|
||||
for from, tos := range g.edges {
|
||||
if _, ok := g.nodes[from]; !ok {
|
||||
return nil, fmt.Errorf("edge source `%s` does not exist", from)
|
||||
}
|
||||
for _, to := range tos {
|
||||
if _, ok := g.nodes[to]; !ok {
|
||||
return nil, fmt.Errorf("edge target `%s` does not exist", to)
|
||||
}
|
||||
indegree[to]++
|
||||
}
|
||||
}
|
||||
|
||||
queue := make([]string, 0, len(g.nodes))
|
||||
level := make(map[string]int, len(g.nodes))
|
||||
for name, deg := range indegree {
|
||||
if deg == 0 {
|
||||
queue = append(queue, name)
|
||||
}
|
||||
}
|
||||
if len(queue) == 0 {
|
||||
return nil, fmt.Errorf("graph has no entry nodes (cycle suspected)")
|
||||
}
|
||||
|
||||
processed := 0
|
||||
for len(queue) > 0 {
|
||||
curr := queue[0]
|
||||
queue = queue[1:]
|
||||
processed++
|
||||
currLevel := level[curr]
|
||||
for _, to := range g.edges[curr] {
|
||||
if level[to] < currLevel+1 {
|
||||
level[to] = currLevel + 1
|
||||
}
|
||||
indegree[to]--
|
||||
if indegree[to] == 0 {
|
||||
queue = append(queue, to)
|
||||
}
|
||||
}
|
||||
}
|
||||
if processed != len(g.nodes) {
|
||||
return nil, fmt.Errorf("graph contains a cycle")
|
||||
}
|
||||
|
||||
maxLevel := 0
|
||||
for _, lv := range level {
|
||||
if lv > maxLevel {
|
||||
maxLevel = lv
|
||||
}
|
||||
}
|
||||
levels := make([][]string, maxLevel+1)
|
||||
for nodeName := range g.nodes {
|
||||
lv := level[nodeName]
|
||||
levels[lv] = append(levels[lv], nodeName)
|
||||
}
|
||||
for i := range levels {
|
||||
slices.Sort(levels[i])
|
||||
}
|
||||
return levels, nil
|
||||
}
|
||||
@@ -0,0 +1,218 @@
|
||||
package stategraph_test
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
sg "github.com/langchain-ai/langgraph/langgraph-go/stategraph"
|
||||
)
|
||||
|
||||
type stateGraphState struct {
|
||||
Noop bool `json:"noop"`
|
||||
}
|
||||
|
||||
type orderRecorder struct {
|
||||
mu sync.Mutex
|
||||
orders map[string]int32
|
||||
seq int32
|
||||
}
|
||||
|
||||
func newOrderRecorder() *orderRecorder {
|
||||
return &orderRecorder{orders: make(map[string]int32)}
|
||||
}
|
||||
|
||||
func (f *orderRecorder) record(name string) {
|
||||
idx := atomic.AddInt32(&f.seq, 1)
|
||||
f.mu.Lock()
|
||||
f.orders[name] = idx
|
||||
f.mu.Unlock()
|
||||
}
|
||||
|
||||
type orderGraph struct {
|
||||
recorder *orderRecorder
|
||||
}
|
||||
|
||||
func newOrderGraph() *orderGraph {
|
||||
return &orderGraph{
|
||||
recorder: newOrderRecorder(),
|
||||
}
|
||||
}
|
||||
|
||||
func (g *orderGraph) A(_ *sg.Context, state stateGraphState) (stateGraphState, error) {
|
||||
g.recorder.record("A")
|
||||
return state, nil
|
||||
}
|
||||
|
||||
func (g *orderGraph) B1(_ *sg.Context, state stateGraphState) (stateGraphState, error) {
|
||||
g.recorder.record("B1")
|
||||
return state, nil
|
||||
}
|
||||
|
||||
func (g *orderGraph) B2(_ *sg.Context, state stateGraphState) (stateGraphState, error) {
|
||||
g.recorder.record("B2")
|
||||
return state, nil
|
||||
}
|
||||
|
||||
func (g *orderGraph) C1(_ *sg.Context, state stateGraphState) (stateGraphState, error) {
|
||||
g.recorder.record("C1")
|
||||
return state, nil
|
||||
}
|
||||
|
||||
func (g *orderGraph) C2(_ *sg.Context, state stateGraphState) (stateGraphState, error) {
|
||||
g.recorder.record("C2")
|
||||
return state, nil
|
||||
}
|
||||
|
||||
func (g *orderGraph) C3(_ *sg.Context, state stateGraphState) (stateGraphState, error) {
|
||||
g.recorder.record("C3")
|
||||
return state, nil
|
||||
}
|
||||
|
||||
func (g *orderGraph) D(_ *sg.Context, state stateGraphState) (stateGraphState, error) {
|
||||
g.recorder.record("D")
|
||||
return state, nil
|
||||
}
|
||||
|
||||
func TestBasicStateGraphWithoutInterrupt(t *testing.T) {
|
||||
fixture := newOrderGraph()
|
||||
graph := sg.NewBasicStateGraph[stateGraphState]()
|
||||
graph.AddNode(fixture.A)
|
||||
graph.AddNode(fixture.B1)
|
||||
graph.AddNode(fixture.B2)
|
||||
graph.AddNode(fixture.C1)
|
||||
graph.AddNode(fixture.C2)
|
||||
graph.AddNode(fixture.C3)
|
||||
graph.AddNode(fixture.D)
|
||||
|
||||
graph.AddEdge(fixture.A, fixture.B1)
|
||||
graph.AddEdge(fixture.A, fixture.B2)
|
||||
graph.AddEdge(fixture.B1, fixture.C1)
|
||||
graph.AddEdge(fixture.B1, fixture.C2)
|
||||
graph.AddEdge(fixture.B2, fixture.C3)
|
||||
graph.AddEdge(fixture.C1, fixture.D)
|
||||
graph.AddEdge(fixture.C2, fixture.D)
|
||||
graph.AddEdge(fixture.C3, fixture.D)
|
||||
|
||||
_, err := graph.Compile().Invoke(stateGraphState{})
|
||||
if err != nil {
|
||||
t.Fatalf("invoke failed: %v", err)
|
||||
}
|
||||
|
||||
fixture.recorder.mu.Lock()
|
||||
orders := make(map[string]int32, len(fixture.recorder.orders))
|
||||
for k, v := range fixture.recorder.orders {
|
||||
orders[k] = v
|
||||
}
|
||||
fixture.recorder.mu.Unlock()
|
||||
|
||||
for _, name := range []string{"A", "B1", "B2", "C1", "C2", "C3", "D"} {
|
||||
if _, ok := orders[name]; !ok {
|
||||
t.Fatalf("node %s did not execute; orders=%v", name, orders)
|
||||
}
|
||||
}
|
||||
maxB := maxInt32(orders["B1"], orders["B2"])
|
||||
minC := minInt32(orders["C1"], minInt32(orders["C2"], orders["C3"]))
|
||||
maxC := maxInt32(orders["C1"], maxInt32(orders["C2"], orders["C3"]))
|
||||
if !(orders["A"] < orders["B1"] && orders["A"] < orders["B2"]) {
|
||||
t.Fatalf("A should run before B-step, orders=%v", orders)
|
||||
}
|
||||
if !(maxB < minC) {
|
||||
t.Fatalf("B-step should finish before C-step, orders=%v", orders)
|
||||
}
|
||||
if !(maxC < orders["D"]) {
|
||||
t.Fatalf("C-step should finish before D, orders=%v", orders)
|
||||
}
|
||||
}
|
||||
|
||||
type interruptState struct {
|
||||
A bool `json:"a"`
|
||||
B bool `json:"b"`
|
||||
}
|
||||
|
||||
type interruptFixture struct{}
|
||||
|
||||
func (f *interruptFixture) A(_ *sg.Context, state interruptState) (interruptState, error) {
|
||||
state.A = true
|
||||
return state, nil
|
||||
}
|
||||
|
||||
func (f *interruptFixture) B(ctx *sg.Context, state interruptState) (interruptState, error) {
|
||||
if !state.A {
|
||||
return state, fmt.Errorf("B should observe A=true")
|
||||
}
|
||||
value, err := ctx.Interrupt("resume_channel")
|
||||
if err != nil {
|
||||
return state, err
|
||||
}
|
||||
s, ok := value.(string)
|
||||
if !ok || s != "go" {
|
||||
return state, fmt.Errorf("unexpected interrupt payload: %#v", value)
|
||||
}
|
||||
state.B = true
|
||||
return state, nil
|
||||
}
|
||||
|
||||
func TestBasicStateGraphWithInterrupt(t *testing.T) {
|
||||
fixture := &interruptFixture{}
|
||||
graph := sg.NewBasicStateGraph[interruptState]()
|
||||
graph.AddNode(fixture.A)
|
||||
graph.AddNode(fixture.B)
|
||||
graph.AddEdge(fixture.A, fixture.B)
|
||||
|
||||
handler, err := graph.Compile().Start(interruptState{})
|
||||
if err != nil {
|
||||
t.Fatalf("start failed: %v", err)
|
||||
}
|
||||
|
||||
doneCh := make(chan interruptState, 1)
|
||||
errCh := make(chan error, 1)
|
||||
go func() {
|
||||
result, runErr := handler.WaitForResult()
|
||||
if runErr != nil {
|
||||
errCh <- runErr
|
||||
return
|
||||
}
|
||||
doneCh <- result
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-doneCh:
|
||||
t.Fatalf("run should pause for interrupt, but completed early")
|
||||
case err := <-errCh:
|
||||
t.Fatalf("run should pause for interrupt, but failed early: %v", err)
|
||||
case <-time.After(120 * time.Millisecond):
|
||||
// expected: paused
|
||||
}
|
||||
|
||||
if err := handler.Resume("resume_channel", "go"); err != nil {
|
||||
t.Fatalf("resume interrupt failed: %v", err)
|
||||
}
|
||||
|
||||
select {
|
||||
case err := <-errCh:
|
||||
t.Fatalf("run failed after interrupt: %v", err)
|
||||
case result := <-doneCh:
|
||||
if !(result.A && result.B) {
|
||||
t.Fatalf("unexpected final state: %#v", result)
|
||||
}
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatalf("timeout waiting for resumed run completion")
|
||||
}
|
||||
}
|
||||
|
||||
func minInt32(a int32, b int32) int32 {
|
||||
if a < b {
|
||||
return a
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
func maxInt32(a int32, b int32) int32 {
|
||||
if a > b {
|
||||
return a
|
||||
}
|
||||
return b
|
||||
}
|
||||
@@ -0,0 +1,204 @@
|
||||
package tests
|
||||
|
||||
import (
|
||||
"slices"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
ag "github.com/langchain-ai/langgraph/langgraph-go/advancedgraph"
|
||||
)
|
||||
|
||||
type decision struct {
|
||||
Type string
|
||||
SubAgent string
|
||||
Tool string
|
||||
Complete string
|
||||
}
|
||||
|
||||
type mockLLM struct {
|
||||
responses [][]decision
|
||||
i int
|
||||
}
|
||||
|
||||
func (m *mockLLM) invoke() []decision {
|
||||
if m.i >= len(m.responses) {
|
||||
return []decision{}
|
||||
}
|
||||
resp := m.responses[m.i]
|
||||
m.i++
|
||||
return resp
|
||||
}
|
||||
|
||||
type lunchWorkflow struct {
|
||||
planner *mockLLM
|
||||
}
|
||||
|
||||
type lunchState struct {
|
||||
Input string `json:"input"`
|
||||
Output []string `json:"output"`
|
||||
Done string `json:"done"`
|
||||
}
|
||||
|
||||
func (w *lunchWorkflow) llmNode(ctx *ag.Context, _ any, _ lunchState) (ag.Command, error) {
|
||||
decisions := w.planner.invoke()
|
||||
sends := make([]ag.Send, 0, 4)
|
||||
for _, d := range decisions {
|
||||
if d.Type == "end" {
|
||||
return ag.Command{
|
||||
Goto: []ag.Send{
|
||||
{Node: w.orderFoodNode, NodeInput: d.Complete},
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
if d.Type == "sub_agent" {
|
||||
sends = append(sends, ag.Send{Node: w.subAgentNode, NodeInput: d.SubAgent})
|
||||
}
|
||||
if d.Type == "tool" {
|
||||
sends = append(sends, ag.Send{Node: w.toolNode, NodeInput: d.Tool})
|
||||
}
|
||||
}
|
||||
sends = append(sends, ag.Send{Node: w.waitNode})
|
||||
return ag.Command{Goto: sends}, nil
|
||||
}
|
||||
|
||||
func (w *lunchWorkflow) waitNode(ctx *ag.Context, _ any, state lunchState) (ag.Command, error) {
|
||||
event, err := ctx.WaitFor(
|
||||
ag.AnyOf(
|
||||
ag.ChannelCondition{Channel: "tool_completion_channel", N: 1},
|
||||
ag.ChannelCondition{Channel: "subagent_completion_channel", N: 1},
|
||||
ag.ChannelCondition{Channel: "user_input_channel", N: 1},
|
||||
ag.TimerCondition{Seconds: 1},
|
||||
),
|
||||
)
|
||||
if err != nil {
|
||||
return ag.Command{}, err
|
||||
}
|
||||
|
||||
output := append([]string(nil), state.Output...)
|
||||
if event.Condition == "channel" {
|
||||
payload := ag.DecodeString(event.Value)
|
||||
switch event.Channel {
|
||||
case "tool_completion_channel":
|
||||
output = append(output, "tool: "+payload)
|
||||
case "subagent_completion_channel":
|
||||
output = append(output, "sub_agent: "+payload)
|
||||
case "user_input_channel":
|
||||
output = append(output, "user_input: "+payload)
|
||||
}
|
||||
state.Output = output
|
||||
return ag.Command{Goto: []ag.Send{{Node: w.llmNode}}, Update: state}, nil
|
||||
}
|
||||
|
||||
output = append(output, "timer: no updates yet")
|
||||
state.Output = output
|
||||
return ag.Command{Goto: []ag.Send{{Node: w.waitNode}}, Update: state}, nil
|
||||
}
|
||||
|
||||
func (w *lunchWorkflow) toolNode(ctx *ag.Context, input any, _ lunchState) (ag.Command, error) {
|
||||
toolInput, _ := input.(string)
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
err := ctx.PublishToChannel("tool_completion_channel", "tool completed for: "+toolInput)
|
||||
return ag.Command{}, err
|
||||
}
|
||||
|
||||
func (w *lunchWorkflow) subAgentNode(ctx *ag.Context, input any, _ lunchState) (ag.Command, error) {
|
||||
subInput, _ := input.(string)
|
||||
time.Sleep(5 * time.Second)
|
||||
err := ctx.PublishToChannel(
|
||||
"subagent_completion_channel",
|
||||
"research sub agent completed for: "+subInput,
|
||||
)
|
||||
return ag.Command{}, err
|
||||
}
|
||||
|
||||
func (w *lunchWorkflow) orderFoodNode(ctx *ag.Context, input any, state lunchState) (ag.Command, error) {
|
||||
complete, _ := input.(string)
|
||||
output := append([]string(nil), state.Output...)
|
||||
output = append(output, "order_food: "+complete)
|
||||
state.Output = output
|
||||
state.Done = complete
|
||||
return ag.Command{Update: state}, nil
|
||||
}
|
||||
|
||||
func TestSubAgentsEquivalentFlow(t *testing.T) {
|
||||
planner := &mockLLM{
|
||||
responses: [][]decision{
|
||||
{
|
||||
{Type: "sub_agent", SubAgent: "research lunch options"},
|
||||
{Type: "tool", Tool: "slack_tool"},
|
||||
},
|
||||
{},
|
||||
{},
|
||||
{{Type: "sub_agent", SubAgent: "find vegetarian fallback"}},
|
||||
{{Type: "end", Complete: "order submitted"}},
|
||||
},
|
||||
}
|
||||
workflow := &lunchWorkflow{
|
||||
planner: planner,
|
||||
}
|
||||
|
||||
graph := ag.NewAdvancedStateGraph[lunchState]()
|
||||
graph.AddAsyncChannel("tool_completion_channel")
|
||||
graph.AddAsyncChannel("subagent_completion_channel")
|
||||
graph.AddAsyncChannel("user_input_channel")
|
||||
|
||||
graph.AddEntryNode(workflow.llmNode)
|
||||
graph.AddNode(workflow.waitNode)
|
||||
graph.AddNode(workflow.toolNode)
|
||||
graph.AddNode(workflow.subAgentNode)
|
||||
graph.AddFinishNode(workflow.orderFoodNode)
|
||||
|
||||
handler, err := graph.Compile().Start(
|
||||
nil,
|
||||
lunchState{
|
||||
Input: "help me get something for lunch",
|
||||
Output: []string{},
|
||||
Done: "",
|
||||
},
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("start failed: %v", err)
|
||||
}
|
||||
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
if err := handler.PublishToChannel("user_input_channel", "No spicy food please"); err != nil {
|
||||
t.Fatalf("publish failed: %v", err)
|
||||
}
|
||||
|
||||
result, err := handler.WaitForResult()
|
||||
if err != nil {
|
||||
t.Fatalf("result failed: %v", err)
|
||||
}
|
||||
|
||||
output := result.Output
|
||||
if len(output) == 0 {
|
||||
t.Fatalf("output is empty, full result=%#v", result)
|
||||
}
|
||||
if result.Done != "order submitted" {
|
||||
t.Fatalf("unexpected done: %v", result.Done)
|
||||
}
|
||||
if !slices.Contains(output, "user_input: No spicy food please") {
|
||||
t.Fatalf("missing user input output: %#v", output)
|
||||
}
|
||||
if !slices.Contains(output, "tool: tool completed for: slack_tool") {
|
||||
t.Fatalf("missing tool output: %#v", output)
|
||||
}
|
||||
if !slices.Contains(output, "sub_agent: research sub agent completed for: research lunch options") {
|
||||
t.Fatalf("missing first sub-agent output: %#v", output)
|
||||
}
|
||||
if !slices.Contains(output, "sub_agent: research sub agent completed for: find vegetarian fallback") {
|
||||
t.Fatalf("missing second sub-agent output: %#v", output)
|
||||
}
|
||||
timerCount := 0
|
||||
for _, line := range output {
|
||||
if line == "timer: no updates yet" {
|
||||
timerCount++
|
||||
}
|
||||
}
|
||||
if timerCount < 3 {
|
||||
t.Fatalf("expected >=3 timer outputs, got %d, output=%#v", timerCount, output)
|
||||
}
|
||||
if output[len(output)-1] != "order_food: order submitted" {
|
||||
t.Fatalf("unexpected last output: %#v", output[len(output)-1])
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,124 @@
|
||||
package tests
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"testing"
|
||||
|
||||
ag "github.com/langchain-ai/langgraph/langgraph-go/advancedgraph"
|
||||
)
|
||||
|
||||
type primitiveWorkflow struct {
|
||||
}
|
||||
|
||||
type primitiveState struct {
|
||||
Count int `json:"count"`
|
||||
Logs []string `json:"logs"`
|
||||
Done string `json:"done"`
|
||||
}
|
||||
|
||||
func (w *primitiveWorkflow) startNode(ctx *ag.Context, input int, state primitiveState) (ag.Command, error) {
|
||||
state.Logs = append(state.Logs, fmt.Sprintf("start:%d", input))
|
||||
return ag.Command{
|
||||
Update: state,
|
||||
Goto: []ag.Send{
|
||||
{Node: w.middleNode, NodeInput: "from_start"},
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (w *primitiveWorkflow) middleNode(ctx *ag.Context, input string, state primitiveState) (ag.Command, error) {
|
||||
state.Logs = append(state.Logs, "middle:"+input)
|
||||
return ag.Command{
|
||||
Update: state,
|
||||
Goto: []ag.Send{
|
||||
{Node: w.finishNode, NodeInput: "from_middle"},
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (w *primitiveWorkflow) finishNode(ctx *ag.Context, input string, state primitiveState) (ag.Command, error) {
|
||||
state.Logs = append(state.Logs, "finish:"+input)
|
||||
state.Done = input
|
||||
return ag.Command{Update: state}, nil
|
||||
}
|
||||
|
||||
func TestInputAndStatePrimitivesCompatible(t *testing.T) {
|
||||
workflow := &primitiveWorkflow{}
|
||||
graph := ag.NewAdvancedStateGraph[primitiveState]()
|
||||
|
||||
graph.AddEntryNode(workflow.startNode)
|
||||
graph.AddNode(workflow.middleNode)
|
||||
graph.AddFinishNode(workflow.finishNode)
|
||||
|
||||
handler, err := graph.Compile().Start(100, primitiveState{
|
||||
Count: 1,
|
||||
Logs: []string{},
|
||||
Done: "",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("start failed: %v", err)
|
||||
}
|
||||
|
||||
result, err := handler.WaitForResult()
|
||||
if err != nil {
|
||||
t.Fatalf("result failed: %v", err)
|
||||
}
|
||||
if result.Done != "from_middle" {
|
||||
t.Fatalf("unexpected done: %v", result.Done)
|
||||
}
|
||||
if result.Count != 1 {
|
||||
t.Fatalf("unexpected count: %v", result.Count)
|
||||
}
|
||||
if len(result.Logs) != 3 || result.Logs[0] != "start:100" || result.Logs[1] != "middle:from_start" || result.Logs[2] != "finish:from_middle" {
|
||||
t.Fatalf("unexpected logs: %#v", result.Logs)
|
||||
}
|
||||
}
|
||||
|
||||
func (w *primitiveWorkflow) startNoFinishNode(ctx *ag.Context, _ any, state primitiveState) (ag.Command, error) {
|
||||
state.Logs = append(state.Logs, "start")
|
||||
return ag.Command{
|
||||
Update: state,
|
||||
Goto: []ag.Send{
|
||||
{Node: w.middleNoFinishNode, NodeInput: "from_start"},
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (w *primitiveWorkflow) middleNoFinishNode(ctx *ag.Context, input string, state primitiveState) (ag.Command, error) {
|
||||
state.Logs = append(state.Logs, "middle:"+input)
|
||||
state.Count += 1
|
||||
state.Done = "stopped"
|
||||
// No goto and no finish node configured: run should end automatically.
|
||||
return ag.Command{Update: state}, nil
|
||||
}
|
||||
|
||||
func TestRunEndsWithoutFinishNode(t *testing.T) {
|
||||
workflow := &primitiveWorkflow{}
|
||||
graph := ag.NewAdvancedStateGraph[primitiveState]()
|
||||
|
||||
graph.AddEntryNode(workflow.startNoFinishNode)
|
||||
graph.AddNode(workflow.middleNoFinishNode)
|
||||
|
||||
handler, err := graph.Compile().Start(nil, primitiveState{
|
||||
Count: 7,
|
||||
Logs: []string{},
|
||||
Done: "",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("start failed: %v", err)
|
||||
}
|
||||
|
||||
result, err := handler.WaitForResult()
|
||||
if err != nil {
|
||||
t.Fatalf("result failed: %v", err)
|
||||
}
|
||||
if result.Done != "stopped" {
|
||||
t.Fatalf("unexpected done: %v", result.Done)
|
||||
}
|
||||
if result.Count != 8 {
|
||||
t.Fatalf("unexpected count: %v", result.Count)
|
||||
}
|
||||
if len(result.Logs) != 2 || result.Logs[0] != "start" || result.Logs[1] != "middle:from_start" {
|
||||
t.Fatalf("unexpected logs: %#v", result.Logs)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,54 @@
|
||||
package tests
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
ag "github.com/langchain-ai/langgraph/langgraph-go/advancedgraph"
|
||||
)
|
||||
|
||||
func TestNewAdvancedStateGraphRejectsNonStructState(t *testing.T) {
|
||||
defer func() {
|
||||
if r := recover(); r == nil {
|
||||
t.Fatalf("expected panic for non-struct StateT")
|
||||
}
|
||||
}()
|
||||
_ = ag.NewAdvancedStateGraph[map[string]any]()
|
||||
}
|
||||
|
||||
type stateTypeA struct {
|
||||
X int `json:"x"`
|
||||
}
|
||||
|
||||
type stateTypeB struct {
|
||||
X int `json:"x"`
|
||||
}
|
||||
|
||||
type wrongUpdateWorkflow struct{}
|
||||
|
||||
func (w *wrongUpdateWorkflow) startNode(ctx *ag.Context, _ any, _ stateTypeA) (ag.Command, error) {
|
||||
return ag.Command{Goto: []ag.Send{{Node: w.badNode}}}, nil
|
||||
}
|
||||
|
||||
func (w *wrongUpdateWorkflow) badNode(ctx *ag.Context, _ any, _ stateTypeA) (ag.Command, error) {
|
||||
return ag.Command{Update: stateTypeB{X: 1}}, nil
|
||||
}
|
||||
|
||||
func TestNodeUpdateTypeMustMatchGraphStateType(t *testing.T) {
|
||||
workflow := &wrongUpdateWorkflow{}
|
||||
graph := ag.NewAdvancedStateGraph[stateTypeA]()
|
||||
graph.AddEntryNode(workflow.startNode)
|
||||
graph.AddFinishNode(workflow.badNode)
|
||||
|
||||
handler, err := graph.Compile().Start(nil, stateTypeA{X: 0})
|
||||
if err != nil {
|
||||
t.Fatalf("start failed: %v", err)
|
||||
}
|
||||
_, err = handler.WaitForResult()
|
||||
if err == nil {
|
||||
t.Fatalf("expected runtime error for wrong update type")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "update type mismatch") {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,135 @@
|
||||
package tests
|
||||
|
||||
import (
|
||||
"reflect"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
ag "github.com/langchain-ai/langgraph/langgraph-go/advancedgraph"
|
||||
)
|
||||
|
||||
type updateElisionState struct {
|
||||
X int `json:"x"`
|
||||
S updateStruct `json:"s"`
|
||||
M map[string]int `json:"m"`
|
||||
L []int `json:"l"`
|
||||
PS *updateStruct `json:"ps"`
|
||||
PM *map[string]int `json:"pm"`
|
||||
PL *[]int `json:"pl"`
|
||||
}
|
||||
|
||||
type updateStruct struct {
|
||||
V int `json:"v"`
|
||||
}
|
||||
|
||||
type updateElisionWorkflow struct{}
|
||||
|
||||
func makeState(v int) updateElisionState {
|
||||
m := map[string]int{"n": v}
|
||||
l := []int{v}
|
||||
return updateElisionState{
|
||||
X: v,
|
||||
S: updateStruct{V: v},
|
||||
M: map[string]int{"n": v},
|
||||
L: []int{v},
|
||||
PS: &updateStruct{V: v},
|
||||
PM: &m,
|
||||
PL: &l,
|
||||
}
|
||||
}
|
||||
|
||||
func (w *updateElisionWorkflow) startNoopNode(ctx *ag.Context, _ any, _ updateElisionState) (ag.Command, error) {
|
||||
return ag.Command{
|
||||
Goto: []ag.Send{
|
||||
{Node: w.fastNode},
|
||||
{Node: w.slowNoopNode},
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (w *updateElisionWorkflow) startChangedNode(ctx *ag.Context, _ any, _ updateElisionState) (ag.Command, error) {
|
||||
return ag.Command{
|
||||
Goto: []ag.Send{
|
||||
{Node: w.fastNode},
|
||||
{Node: w.slowChangedNode},
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (w *updateElisionWorkflow) fastNode(ctx *ag.Context, _ any, state updateElisionState) (ag.Command, error) {
|
||||
_ = state
|
||||
return ag.Command{Update: makeState(1)}, nil
|
||||
}
|
||||
|
||||
func (w *updateElisionWorkflow) slowNoopNode(ctx *ag.Context, _ any, state updateElisionState) (ag.Command, error) {
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
// Returns same state as initial snapshot; without runtime elision this can overwrite newer updates.
|
||||
return ag.Command{Update: state}, nil
|
||||
}
|
||||
|
||||
func (w *updateElisionWorkflow) slowChangedNode(ctx *ag.Context, _ any, _ updateElisionState) (ag.Command, error) {
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
// Real change should not be elided.
|
||||
return ag.Command{Update: makeState(2)}, nil
|
||||
}
|
||||
|
||||
func assertStateEquals(t *testing.T, got updateElisionState, expected updateElisionState) {
|
||||
t.Helper()
|
||||
if got.X != expected.X {
|
||||
t.Fatalf("unexpected X: got=%d want=%d", got.X, expected.X)
|
||||
}
|
||||
if got.S != expected.S {
|
||||
t.Fatalf("unexpected S: got=%#v want=%#v", got.S, expected.S)
|
||||
}
|
||||
if !reflect.DeepEqual(got.M, expected.M) {
|
||||
t.Fatalf("unexpected M: got=%#v want=%#v", got.M, expected.M)
|
||||
}
|
||||
if !reflect.DeepEqual(got.L, expected.L) {
|
||||
t.Fatalf("unexpected L: got=%#v want=%#v", got.L, expected.L)
|
||||
}
|
||||
if got.PS == nil || expected.PS == nil || *got.PS != *expected.PS {
|
||||
t.Fatalf("unexpected PS: got=%#v want=%#v", got.PS, expected.PS)
|
||||
}
|
||||
if got.PM == nil || expected.PM == nil || !reflect.DeepEqual(*got.PM, *expected.PM) {
|
||||
t.Fatalf("unexpected PM: got=%#v want=%#v", got.PM, expected.PM)
|
||||
}
|
||||
if got.PL == nil || expected.PL == nil || !reflect.DeepEqual(*got.PL, *expected.PL) {
|
||||
t.Fatalf("unexpected PL: got=%#v want=%#v", got.PL, expected.PL)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNoopSlowUpdateCanOverrideFastUpdate(t *testing.T) {
|
||||
workflow := &updateElisionWorkflow{}
|
||||
graph := ag.NewAdvancedStateGraph[updateElisionState]()
|
||||
graph.AddEntryNode(workflow.startNoopNode)
|
||||
graph.AddNode(workflow.fastNode)
|
||||
graph.AddFinishNode(workflow.slowNoopNode)
|
||||
|
||||
handler, err := graph.Compile().Start(nil, makeState(0))
|
||||
if err != nil {
|
||||
t.Fatalf("start failed: %v", err)
|
||||
}
|
||||
result, err := handler.WaitForResult()
|
||||
if err != nil {
|
||||
t.Fatalf("result failed: %v", err)
|
||||
}
|
||||
assertStateEquals(t, result, makeState(0))
|
||||
}
|
||||
|
||||
func TestChangedSlowUpdateOverridesFastUpdate(t *testing.T) {
|
||||
workflow := &updateElisionWorkflow{}
|
||||
graph := ag.NewAdvancedStateGraph[updateElisionState]()
|
||||
graph.AddEntryNode(workflow.startChangedNode)
|
||||
graph.AddNode(workflow.fastNode)
|
||||
graph.AddFinishNode(workflow.slowChangedNode)
|
||||
|
||||
handler, err := graph.Compile().Start(nil, makeState(0))
|
||||
if err != nil {
|
||||
t.Fatalf("start failed: %v", err)
|
||||
}
|
||||
result, err := handler.WaitForResult()
|
||||
if err != nil {
|
||||
t.Fatalf("result failed: %v", err)
|
||||
}
|
||||
assertStateEquals(t, result, makeState(2))
|
||||
}
|
||||
+70
-288
@@ -29,7 +29,7 @@ from langgraph_cli.exec import Runner, subp_exec
|
||||
from langgraph_cli.host_backend import HostBackendClient, HostBackendError
|
||||
from langgraph_cli.progress import Progress
|
||||
from langgraph_cli.templates import TEMPLATE_HELP_STRING, create_new
|
||||
from langgraph_cli.util import format_deployments_table, warn_non_wolfi_distro
|
||||
from langgraph_cli.util import warn_non_wolfi_distro
|
||||
from langgraph_cli.version import __version__
|
||||
|
||||
RESERVED_ENV_VARS = frozenset(
|
||||
@@ -287,22 +287,6 @@ OPT_API_VERSION = click.option(
|
||||
help="API server version to use for the base image. If unspecified, the latest version will be used.",
|
||||
)
|
||||
|
||||
OPT_HOST_API_KEY = click.option(
|
||||
"--api-key",
|
||||
envvar="LANGGRAPH_HOST_API_KEY",
|
||||
help=(
|
||||
"API key. Can also be set via LANGGRAPH_HOST_API_KEY, "
|
||||
"LANGSMITH_API_KEY, or LANGCHAIN_API_KEY environment variable or .env file."
|
||||
),
|
||||
)
|
||||
|
||||
OPT_HOST_URL = click.option(
|
||||
"--host-url",
|
||||
envvar="LANGGRAPH_HOST_URL",
|
||||
default="https://api.host.langchain.com",
|
||||
hidden=True,
|
||||
)
|
||||
|
||||
OPT_ENGINE_RUNTIME_MODE = click.option(
|
||||
"--engine-runtime-mode",
|
||||
type=click.Choice(["combined_queue_worker", "distributed"]),
|
||||
@@ -311,67 +295,7 @@ OPT_ENGINE_RUNTIME_MODE = click.option(
|
||||
)
|
||||
|
||||
|
||||
class NestedHelpGroup(click.Group):
|
||||
"""Click group that shows one level of nested subcommands in top-level help."""
|
||||
|
||||
def format_commands(
|
||||
self, ctx: click.Context, formatter: click.HelpFormatter
|
||||
) -> None:
|
||||
command_entries: list[tuple[str, click.Command]] = []
|
||||
# Collect the top-level commands first, then append one level of nested
|
||||
# subcommands using names like "deploy list" so they show up in the
|
||||
# top-level help output.
|
||||
for command_name in self.list_commands(ctx):
|
||||
command = self.get_command(ctx, command_name)
|
||||
if command is None or command.hidden:
|
||||
continue
|
||||
command_entries.append((command_name, command))
|
||||
if isinstance(command, click.Group):
|
||||
# Build a child context so Click resolves the subcommands the same
|
||||
# way it would for the nested group itself.
|
||||
sub_ctx = click.Context(command, info_name=command_name, parent=ctx)
|
||||
for subcommand_name in command.list_commands(sub_ctx):
|
||||
subcommand = command.get_command(sub_ctx, subcommand_name)
|
||||
if subcommand is None or subcommand.hidden:
|
||||
continue
|
||||
command_entries.append(
|
||||
(f"{command_name} {subcommand_name}", subcommand)
|
||||
)
|
||||
|
||||
# Compute the available width for help text up front so we can truncate
|
||||
# descriptions before handing them to Click. That keeps each command on
|
||||
# a single line instead of allowing wrapped descriptions.
|
||||
command_width = max((len(name) for name, _ in command_entries), default=0)
|
||||
help_width = max(formatter.width - command_width - 6, 10)
|
||||
rows = [
|
||||
(name, command.get_short_help_str(help_width))
|
||||
for name, command in command_entries
|
||||
]
|
||||
|
||||
if rows:
|
||||
# Render the flattened command list using Click's standard
|
||||
# definition-list formatter so alignment stays consistent with the
|
||||
# rest of the CLI help output.
|
||||
with formatter.section("Commands"):
|
||||
formatter.write_dl(rows)
|
||||
|
||||
|
||||
class DeployGroup(NestedHelpGroup):
|
||||
"""Group that treats leading '-' args as passthrough docker flags."""
|
||||
|
||||
def parse_args(self, ctx: click.Context, args: list[str]) -> list[str]:
|
||||
result = super().parse_args(ctx, args)
|
||||
if ctx._protected_args and ctx._protected_args[0].startswith("-"):
|
||||
# Click stores the would-be subcommand in _protected_args; if it looks
|
||||
# like an option (e.g. --build-arg) treat it as passthrough docker
|
||||
# args instead of insisting on a nested command.
|
||||
ctx.args = [*ctx._protected_args, *ctx.args]
|
||||
ctx._protected_args = []
|
||||
return ctx.args
|
||||
return result
|
||||
|
||||
|
||||
@click.group(cls=NestedHelpGroup)
|
||||
@click.group()
|
||||
@click.version_option(version=__version__, prog_name="LangGraph CLI")
|
||||
def cli():
|
||||
pass
|
||||
@@ -669,89 +593,72 @@ def build(
|
||||
)
|
||||
|
||||
|
||||
def _deploy_base_options(
|
||||
func: Callable | None = None,
|
||||
*,
|
||||
include_docker_args: bool = True,
|
||||
validate_config_path: bool = True,
|
||||
):
|
||||
"""Apply shared deploy flags.
|
||||
|
||||
The group shares most options but should not consume subcommands, so the
|
||||
docker build args are only attached when requested.
|
||||
"""
|
||||
|
||||
def _apply(target: Callable) -> Callable:
|
||||
decorators = [
|
||||
OPT_HOST_API_KEY,
|
||||
click.option(
|
||||
"--name",
|
||||
envvar="LANGSMITH_DEPLOYMENT_NAME",
|
||||
help=(
|
||||
"Deployment name. Can also be set via LANGSMITH_DEPLOYMENT_NAME "
|
||||
"environment variable or .env file. Defaults to current directory name "
|
||||
"if --deployment-id is not provided."
|
||||
),
|
||||
),
|
||||
click.option(
|
||||
"--deployment-id",
|
||||
help=(
|
||||
"ID of an existing deployment to update. If omitted, "
|
||||
"--name is used to find or create the deployment."
|
||||
),
|
||||
),
|
||||
click.option(
|
||||
"--deployment-type",
|
||||
type=click.Choice(["dev", "prod"]),
|
||||
default="dev",
|
||||
show_default=True,
|
||||
help="Deployment type (used when creating a new deployment).",
|
||||
),
|
||||
click.option(
|
||||
"--no-wait",
|
||||
is_flag=True,
|
||||
default=False,
|
||||
help="Skip waiting for deployment status.",
|
||||
),
|
||||
OPT_VERBOSE,
|
||||
OPT_HOST_URL,
|
||||
click.option("--image-name", hidden=True),
|
||||
click.option("--image-tag", default="latest", hidden=True),
|
||||
click.option(
|
||||
"--config",
|
||||
"-c",
|
||||
default=DEFAULT_CONFIG,
|
||||
hidden=True,
|
||||
type=click.Path(
|
||||
exists=validate_config_path,
|
||||
file_okay=True,
|
||||
dir_okay=False,
|
||||
resolve_path=True,
|
||||
path_type=pathlib.Path,
|
||||
),
|
||||
),
|
||||
click.option("--pull/--no-pull", default=True, hidden=True),
|
||||
click.option("--base-image", hidden=True),
|
||||
click.option("--install-command", hidden=True),
|
||||
click.option("--build-command", hidden=True),
|
||||
click.option("--api-version", type=str, hidden=True),
|
||||
]
|
||||
if include_docker_args:
|
||||
# Only attach build args to the default command; on the group they
|
||||
# would capture subcommand names like `list` before Click resolves
|
||||
# them, making those subcommands unreachable.
|
||||
decorators.append(
|
||||
click.argument("docker_build_args", nargs=-1, type=click.UNPROCESSED)
|
||||
)
|
||||
for decorator in reversed(decorators):
|
||||
target = decorator(target)
|
||||
return target
|
||||
|
||||
return _apply(func) if func is not None else _apply
|
||||
|
||||
|
||||
@cli.group(
|
||||
cls=DeployGroup,
|
||||
@click.option(
|
||||
"--api-key",
|
||||
envvar="LANGGRAPH_HOST_API_KEY",
|
||||
help=(
|
||||
"API key. Can also be set via LANGGRAPH_HOST_API_KEY, "
|
||||
"LANGSMITH_API_KEY, or LANGCHAIN_API_KEY environment variable or .env file."
|
||||
),
|
||||
)
|
||||
@click.option(
|
||||
"--name",
|
||||
envvar="LANGSMITH_DEPLOYMENT_NAME",
|
||||
help=(
|
||||
"Deployment name. Can also be set via LANGSMITH_DEPLOYMENT_NAME "
|
||||
"environment variable or .env file. Defaults to current directory name "
|
||||
"if --deployment-id is not provided."
|
||||
),
|
||||
)
|
||||
@click.option(
|
||||
"--deployment-id",
|
||||
help=(
|
||||
"ID of an existing deployment to update. If omitted, "
|
||||
"--name is used to find or create the deployment."
|
||||
),
|
||||
)
|
||||
@click.option(
|
||||
"--deployment-type",
|
||||
type=click.Choice(["dev", "prod"]),
|
||||
default="dev",
|
||||
show_default=True,
|
||||
help="Deployment type (used when creating a new deployment).",
|
||||
)
|
||||
@click.option(
|
||||
"--no-wait",
|
||||
is_flag=True,
|
||||
default=False,
|
||||
help="Skip waiting for deployment status.",
|
||||
)
|
||||
@OPT_VERBOSE
|
||||
@click.option(
|
||||
"--host-url",
|
||||
envvar="LANGGRAPH_HOST_URL",
|
||||
default="https://api.host.langchain.com",
|
||||
hidden=True,
|
||||
)
|
||||
@click.option("--image-name", hidden=True)
|
||||
@click.option("--image-tag", default="latest", hidden=True)
|
||||
@click.option(
|
||||
"--config",
|
||||
"-c",
|
||||
default=DEFAULT_CONFIG,
|
||||
hidden=True,
|
||||
type=click.Path(
|
||||
exists=True,
|
||||
file_okay=True,
|
||||
dir_okay=False,
|
||||
resolve_path=True,
|
||||
path_type=pathlib.Path,
|
||||
),
|
||||
)
|
||||
@click.option("--pull/--no-pull", default=True, hidden=True)
|
||||
@click.option("--base-image", hidden=True)
|
||||
@click.option("--install-command", hidden=True)
|
||||
@click.option("--build-command", hidden=True)
|
||||
@click.option("--api-version", type=str, hidden=True)
|
||||
@click.argument("docker_build_args", nargs=-1, type=click.UNPROCESSED)
|
||||
@cli.command(
|
||||
help=(
|
||||
"[Beta] Build and deploy a LangGraph image to LangSmith Deployments.\n\n"
|
||||
"This command is in beta and under active development. "
|
||||
@@ -760,26 +667,10 @@ def _deploy_base_options(
|
||||
"is located). This command also accepts build flags (--base-image, "
|
||||
"--pull, etc.). See 'langgraph build --help' for details."
|
||||
),
|
||||
context_settings=dict(ignore_unknown_options=True, allow_extra_args=True),
|
||||
invoke_without_command=True, # allow `deploy` click group to execute without command
|
||||
context_settings=dict(ignore_unknown_options=True),
|
||||
)
|
||||
@_deploy_base_options(include_docker_args=False, validate_config_path=False)
|
||||
@click.pass_context
|
||||
@log_command
|
||||
def deploy(ctx: click.Context, **_: object):
|
||||
# We register deploy as both a group and a command here.
|
||||
# if we detect no subcommand, we run _deploy (basically run langgraph deploy as a top level command)
|
||||
# otherwise, we return None here and click will proceed to actually run the subcommand (list or delete)
|
||||
if ctx.invoked_subcommand is not None:
|
||||
return
|
||||
docker_build_args = tuple(ctx.args)
|
||||
ctx.args = [] # Prevent Click from re-processing passthrough args later.
|
||||
return ctx.forward(_deploy, docker_build_args=docker_build_args)
|
||||
|
||||
|
||||
@_deploy_base_options()
|
||||
@click.command(context_settings=dict(ignore_unknown_options=True))
|
||||
def _deploy(
|
||||
def deploy(
|
||||
config: pathlib.Path,
|
||||
pull: bool,
|
||||
verbose: bool,
|
||||
@@ -1159,115 +1050,6 @@ def _deploy(
|
||||
)
|
||||
|
||||
|
||||
def _create_host_backend_client(
|
||||
host_url: str | None,
|
||||
api_key: str | None,
|
||||
env_vars: dict[str, str] | None = None,
|
||||
) -> HostBackendClient:
|
||||
if env_vars is None:
|
||||
env_vars = _parse_env_from_config({}, pathlib.Path.cwd() / DEFAULT_CONFIG)
|
||||
resolved_api_key = api_key
|
||||
if not resolved_api_key:
|
||||
for key_name in _API_KEY_ENV_NAMES:
|
||||
val = env_vars.get(key_name)
|
||||
if val:
|
||||
resolved_api_key = val
|
||||
break
|
||||
val = os.environ.get(key_name)
|
||||
if val:
|
||||
resolved_api_key = val
|
||||
break
|
||||
if not resolved_api_key:
|
||||
resolved_api_key = click.prompt("Host API key", hide_input=True)
|
||||
return HostBackendClient(host_url, resolved_api_key)
|
||||
|
||||
|
||||
def _call_host_backend_with_optional_tenant(
|
||||
client: HostBackendClient,
|
||||
operation: Callable[[HostBackendClient], object],
|
||||
) -> object:
|
||||
try:
|
||||
return operation(client)
|
||||
except HostBackendError as err:
|
||||
if err.status_code == 403 and "requires workspace specification" in err.message:
|
||||
click.secho(
|
||||
"Your API key is org-scoped and requires a workspace ID.",
|
||||
fg="yellow",
|
||||
)
|
||||
click.secho(
|
||||
"Find your workspace ID in LangSmith under Settings > Workspaces.",
|
||||
fg="yellow",
|
||||
)
|
||||
tenant_id = click.prompt("Workspace ID")
|
||||
client = HostBackendClient(
|
||||
client._base_url, client._api_key, tenant_id=tenant_id
|
||||
)
|
||||
return operation(client)
|
||||
raise
|
||||
|
||||
|
||||
@OPT_HOST_API_KEY
|
||||
@OPT_HOST_URL
|
||||
@click.option(
|
||||
"--name-contains",
|
||||
default="",
|
||||
help="Only show deployments whose names contain this value.",
|
||||
)
|
||||
@deploy.command("list", help="[Beta] List LangSmith Deployments.")
|
||||
def deploy_list(api_key: str | None, host_url: str | None, name_contains: str) -> None:
|
||||
client = _create_host_backend_client(host_url, api_key)
|
||||
response = _call_host_backend_with_optional_tenant(
|
||||
client,
|
||||
lambda current_client: current_client.list_deployments(
|
||||
name_contains=name_contains
|
||||
),
|
||||
)
|
||||
resources = response.get("resources", []) if isinstance(response, dict) else []
|
||||
deployments = [item for item in resources if isinstance(item, dict)]
|
||||
if not deployments:
|
||||
click.echo("No deployments found.")
|
||||
return
|
||||
click.echo(format_deployments_table(deployments))
|
||||
|
||||
|
||||
@OPT_HOST_API_KEY
|
||||
@OPT_HOST_URL
|
||||
@click.option(
|
||||
"--force",
|
||||
is_flag=True,
|
||||
default=False,
|
||||
help="Delete without prompting for confirmation.",
|
||||
)
|
||||
@click.argument("deployment_id")
|
||||
@deploy.command(
|
||||
"delete",
|
||||
help=(
|
||||
"[Beta] Delete a LangSmith Deployment.\n\n"
|
||||
"Use the `deploy list` command to list deployment IDs."
|
||||
),
|
||||
)
|
||||
def deploy_delete(
|
||||
api_key: str | None, host_url: str | None, force: bool, deployment_id: str
|
||||
) -> None:
|
||||
if not force:
|
||||
response = click.prompt(
|
||||
click.style(
|
||||
f"Are you sure you want to delete deployment ID {deployment_id}? (Y/n)",
|
||||
fg="yellow",
|
||||
),
|
||||
default="Y",
|
||||
show_default=False,
|
||||
)
|
||||
if response.strip().lower() not in {"y", "yes"}:
|
||||
raise click.Abort()
|
||||
client = _create_host_backend_client(host_url, api_key)
|
||||
_call_host_backend_with_optional_tenant(
|
||||
client,
|
||||
lambda current_client: current_client.delete_deployment(deployment_id),
|
||||
)
|
||||
click.secho(f"Deleted deployment {deployment_id}.", fg="green")
|
||||
|
||||
|
||||
def _normalize_image_name(value: str | None) -> str:
|
||||
"""Sanitize a deployment/directory name into a valid Docker repository name.
|
||||
|
||||
|
||||
@@ -39,14 +39,10 @@ class HostBackendClient:
|
||||
)
|
||||
|
||||
def _request(
|
||||
self,
|
||||
method: str,
|
||||
path: str,
|
||||
payload: dict[str, Any] | None = None,
|
||||
params: dict[str, Any] | None = None,
|
||||
self, method: str, path: str, payload: dict[str, Any] | None = None
|
||||
) -> Any:
|
||||
try:
|
||||
resp = self._client.request(method, path, json=payload, params=params)
|
||||
resp = self._client.request(method, path, json=payload)
|
||||
resp.raise_for_status()
|
||||
except httpx.HTTPStatusError as err:
|
||||
detail = err.response.text or str(err.response.status_code)
|
||||
@@ -69,19 +65,12 @@ class HostBackendClient:
|
||||
def create_deployment(self, payload: dict[str, Any]) -> dict[str, Any]:
|
||||
return self._request("POST", "/v2/deployments", payload)
|
||||
|
||||
def list_deployments(self, name_contains: str = "") -> dict[str, Any]:
|
||||
return self._request(
|
||||
"GET",
|
||||
"/v2/deployments",
|
||||
params={"name_contains": name_contains},
|
||||
)
|
||||
def list_deployments(self, name_contains: str) -> dict[str, Any]:
|
||||
return self._request("GET", f"/v2/deployments?name_contains={name_contains}")
|
||||
|
||||
def get_deployment(self, deployment_id: str) -> dict[str, Any]:
|
||||
return self._request("GET", f"/v2/deployments/{deployment_id}")
|
||||
|
||||
def delete_deployment(self, deployment_id: str) -> None:
|
||||
return self._request("DELETE", f"/v2/deployments/{deployment_id}")
|
||||
|
||||
def request_push_token(self, deployment_id: str) -> dict[str, Any]:
|
||||
return self._request(
|
||||
"POST",
|
||||
|
||||
@@ -1,5 +1,3 @@
|
||||
from collections.abc import Sequence
|
||||
|
||||
import click
|
||||
|
||||
|
||||
@@ -25,35 +23,3 @@ def warn_non_wolfi_distro(config_json: dict) -> None:
|
||||
fg="yellow",
|
||||
)
|
||||
click.secho("") # Empty line for better readability
|
||||
|
||||
|
||||
def _extract_deployment_url(deployment: dict[str, object]) -> str:
|
||||
source_config = deployment.get("source_config")
|
||||
if isinstance(source_config, dict):
|
||||
custom_url = source_config.get("custom_url")
|
||||
if isinstance(custom_url, str) and custom_url:
|
||||
return custom_url
|
||||
return "-"
|
||||
|
||||
|
||||
def format_deployments_table(deployments: Sequence[dict[str, object]]) -> str:
|
||||
headers = ("Deployment ID", "Deployment Name", "Deployment URL")
|
||||
rows = [
|
||||
(
|
||||
str(deployment.get("id", "-") or "-"),
|
||||
str(deployment.get("name", "-") or "-"),
|
||||
_extract_deployment_url(deployment),
|
||||
)
|
||||
for deployment in deployments
|
||||
]
|
||||
widths = [
|
||||
max(len(headers[index]), *(len(row[index]) for row in rows))
|
||||
for index in range(len(headers))
|
||||
]
|
||||
|
||||
def format_row(row: Sequence[str]) -> str:
|
||||
return " ".join(value.ljust(widths[index]) for index, value in enumerate(row))
|
||||
|
||||
lines = [format_row(headers), format_row(tuple("-" * width for width in widths))]
|
||||
lines.extend(format_row(row) for row in rows)
|
||||
return "\n".join(lines)
|
||||
|
||||
@@ -9,7 +9,6 @@ from pathlib import Path
|
||||
|
||||
from click.testing import CliRunner
|
||||
|
||||
import langgraph_cli.cli as cli_module
|
||||
from langgraph_cli.cli import cli, prepare_args_and_stdin
|
||||
from langgraph_cli.config import Config, _get_pip_cleanup_lines, validate_config
|
||||
from langgraph_cli.docker import DEFAULT_POSTGRES_URI, DockerCapabilities, Version
|
||||
@@ -288,238 +287,6 @@ def test_version_option() -> None:
|
||||
)
|
||||
|
||||
|
||||
def test_top_level_help_shows_deploy_subcommands() -> None:
|
||||
runner = CliRunner()
|
||||
|
||||
result = runner.invoke(cli, ["--help"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "deploy" in result.output
|
||||
assert "deploy list" in result.output
|
||||
assert "deploy delete" in result.output
|
||||
assert "[Beta] List LangSmith Deployments." in result.output
|
||||
|
||||
|
||||
def test_top_level_help_truncates_command_descriptions_to_single_line() -> None:
|
||||
runner = CliRunner()
|
||||
|
||||
result = runner.invoke(cli, ["--help"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
lines = result.output.splitlines()
|
||||
deploy_line = next(line for line in lines if line.strip().startswith("deploy"))
|
||||
deploy_list_line = next(
|
||||
line for line in lines if line.strip().startswith("deploy list")
|
||||
)
|
||||
|
||||
assert not lines[lines.index(deploy_line) + 1].startswith(" ")
|
||||
assert "..." in deploy_line
|
||||
assert "[Beta] List LangSmith Deployments." in deploy_list_line
|
||||
|
||||
|
||||
def test_deploy_list_command(monkeypatch) -> None:
|
||||
runner = CliRunner()
|
||||
captured: dict[str, str] = {}
|
||||
|
||||
class FakeClient:
|
||||
def __init__(self, host_url: str, api_key: str, tenant_id: str | None = None):
|
||||
captured["host_url"] = host_url
|
||||
captured["api_key"] = api_key
|
||||
captured["tenant_id"] = tenant_id or ""
|
||||
|
||||
def list_deployments(self, name_contains: str = ""):
|
||||
captured["name_contains"] = name_contains
|
||||
return {
|
||||
"resources": [
|
||||
{
|
||||
"id": "dep-123",
|
||||
"name": "alpha",
|
||||
"source_config": {"custom_url": "https://alpha.example.com"},
|
||||
},
|
||||
{
|
||||
"id": "dep-456",
|
||||
"name": "beta",
|
||||
"source_config": {"custom_url": "https://beta.example.com"},
|
||||
},
|
||||
]
|
||||
}
|
||||
|
||||
monkeypatch.setattr(cli_module, "HostBackendClient", FakeClient)
|
||||
|
||||
result = runner.invoke(
|
||||
cli,
|
||||
[
|
||||
"deploy",
|
||||
"list",
|
||||
"--api-key",
|
||||
"test-key",
|
||||
"--host-url",
|
||||
"https://api.example.com",
|
||||
"--name-contains",
|
||||
"alp",
|
||||
],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert captured == {
|
||||
"host_url": "https://api.example.com",
|
||||
"api_key": "test-key",
|
||||
"tenant_id": "",
|
||||
"name_contains": "alp",
|
||||
}
|
||||
assert "Deployment ID" in result.output
|
||||
assert "Deployment Name" in result.output
|
||||
assert "Deployment URL" in result.output
|
||||
assert "dep-123" in result.output
|
||||
assert "https://beta.example.com" in result.output
|
||||
|
||||
|
||||
def test_deploy_list_command_no_results(monkeypatch) -> None:
|
||||
runner = CliRunner()
|
||||
|
||||
class FakeClient:
|
||||
def __init__(self, host_url: str, api_key: str, tenant_id: str | None = None):
|
||||
pass
|
||||
|
||||
def list_deployments(self, name_contains: str = ""):
|
||||
return {"resources": []}
|
||||
|
||||
monkeypatch.setattr(cli_module, "HostBackendClient", FakeClient)
|
||||
|
||||
result = runner.invoke(
|
||||
cli,
|
||||
[
|
||||
"deploy",
|
||||
"list",
|
||||
"--api-key",
|
||||
"test-key",
|
||||
"--host-url",
|
||||
"https://api.example.com",
|
||||
],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert result.output.strip() == "No deployments found."
|
||||
|
||||
|
||||
def test_deploy_delete_command(monkeypatch) -> None:
|
||||
runner = CliRunner()
|
||||
captured: dict[str, str] = {}
|
||||
|
||||
class FakeClient:
|
||||
def __init__(self, host_url: str, api_key: str, tenant_id: str | None = None):
|
||||
captured["host_url"] = host_url
|
||||
captured["api_key"] = api_key
|
||||
captured["tenant_id"] = tenant_id or ""
|
||||
|
||||
def delete_deployment(self, deployment_id: str):
|
||||
captured["deployment_id"] = deployment_id
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(cli_module, "HostBackendClient", FakeClient)
|
||||
|
||||
result = runner.invoke(
|
||||
cli,
|
||||
[
|
||||
"deploy",
|
||||
"delete",
|
||||
"--api-key",
|
||||
"test-key",
|
||||
"--host-url",
|
||||
"https://api.example.com",
|
||||
"dep-123",
|
||||
],
|
||||
input="y\n",
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert captured == {
|
||||
"host_url": "https://api.example.com",
|
||||
"api_key": "test-key",
|
||||
"tenant_id": "",
|
||||
"deployment_id": "dep-123",
|
||||
}
|
||||
assert (
|
||||
"Are you sure you want to delete deployment ID dep-123? (Y/n):" in result.output
|
||||
)
|
||||
assert result.output.strip().endswith("Deleted deployment dep-123.")
|
||||
|
||||
|
||||
def test_deploy_delete_command_cancelled(monkeypatch) -> None:
|
||||
runner = CliRunner()
|
||||
deleted = False
|
||||
|
||||
class FakeClient:
|
||||
def __init__(self, host_url: str, api_key: str, tenant_id: str | None = None):
|
||||
pass
|
||||
|
||||
def delete_deployment(self, deployment_id: str):
|
||||
nonlocal deleted
|
||||
deleted = True
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(cli_module, "HostBackendClient", FakeClient)
|
||||
|
||||
result = runner.invoke(
|
||||
cli,
|
||||
[
|
||||
"deploy",
|
||||
"delete",
|
||||
"--api-key",
|
||||
"test-key",
|
||||
"--host-url",
|
||||
"https://api.example.com",
|
||||
"dep-123",
|
||||
],
|
||||
input="n\n",
|
||||
)
|
||||
|
||||
assert result.exit_code == 1, result.output
|
||||
assert not deleted
|
||||
assert "Aborted!" in result.output
|
||||
|
||||
|
||||
def test_deploy_delete_command_force(monkeypatch) -> None:
|
||||
runner = CliRunner()
|
||||
captured: dict[str, str] = {}
|
||||
|
||||
class FakeClient:
|
||||
def __init__(self, host_url: str, api_key: str, tenant_id: str | None = None):
|
||||
captured["host_url"] = host_url
|
||||
captured["api_key"] = api_key
|
||||
captured["tenant_id"] = tenant_id or ""
|
||||
|
||||
def delete_deployment(self, deployment_id: str):
|
||||
captured["deployment_id"] = deployment_id
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(cli_module, "HostBackendClient", FakeClient)
|
||||
|
||||
result = runner.invoke(
|
||||
cli,
|
||||
[
|
||||
"deploy",
|
||||
"delete",
|
||||
"--force",
|
||||
"--api-key",
|
||||
"test-key",
|
||||
"--host-url",
|
||||
"https://api.example.com",
|
||||
"dep-123",
|
||||
],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "Are you sure you want to delete deployment ID dep-123?" not in result.output
|
||||
assert captured == {
|
||||
"host_url": "https://api.example.com",
|
||||
"api_key": "test-key",
|
||||
"tenant_id": "",
|
||||
"deployment_id": "dep-123",
|
||||
}
|
||||
assert result.output.strip() == "Deleted deployment dep-123."
|
||||
|
||||
|
||||
def test_dockerfile_command_basic() -> None:
|
||||
"""Test the 'dockerfile' command with basic configuration."""
|
||||
runner = CliRunner()
|
||||
|
||||
@@ -1784,9 +1784,7 @@ def test_config_to_compose_distributed_mode():
|
||||
# Executor service is present with correct base image
|
||||
assert "langgraph-executor:" in actual_compose_stdin
|
||||
assert "FROM langchain/langgraph-executor:3.11" in actual_compose_stdin
|
||||
assert (
|
||||
'entrypoint: ["sh", "/storage/executor_entrypoint.sh"]' in actual_compose_stdin
|
||||
)
|
||||
assert 'entrypoint: ["sh", "/storage/executor_entrypoint.sh"]' in actual_compose_stdin
|
||||
|
||||
# Executor has required environment variables
|
||||
assert "EXECUTOR_GRPC_PORT:" in actual_compose_stdin
|
||||
|
||||
@@ -135,28 +135,6 @@ def test_list_deployments(client):
|
||||
assert result == {"ok": True}
|
||||
|
||||
|
||||
def test_list_deployments_sends_query_params():
|
||||
def handler(req: httpx.Request) -> httpx.Response:
|
||||
assert req.url.path == "/v2/deployments"
|
||||
assert req.url.params["name_contains"] == "my app"
|
||||
return httpx.Response(200, json={"ok": True})
|
||||
|
||||
c = HostBackendClient("https://api.example.com", "test-key")
|
||||
c._client = httpx.Client(
|
||||
base_url="https://api.example.com",
|
||||
transport=httpx.MockTransport(handler),
|
||||
headers={"X-Api-Key": "test-key", "Accept": "application/json"},
|
||||
timeout=30,
|
||||
)
|
||||
result = c.list_deployments("my app")
|
||||
assert result == {"ok": True}
|
||||
|
||||
|
||||
def test_delete_deployment(client):
|
||||
result = client.delete_deployment("dep-123")
|
||||
assert result == {"ok": True}
|
||||
|
||||
|
||||
def test_request_push_token(client):
|
||||
result = client.request_push_token("dep-123")
|
||||
assert result == {"ok": True}
|
||||
|
||||
@@ -1,11 +1,6 @@
|
||||
from unittest.mock import patch
|
||||
|
||||
from langgraph_cli.util import (
|
||||
_extract_deployment_url,
|
||||
clean_empty_lines,
|
||||
format_deployments_table,
|
||||
warn_non_wolfi_distro,
|
||||
)
|
||||
from langgraph_cli.util import clean_empty_lines, warn_non_wolfi_distro
|
||||
|
||||
|
||||
def test_clean_empty_lines():
|
||||
@@ -191,36 +186,3 @@ def test_warn_non_wolfi_distro_does_not_modify_config():
|
||||
warn_non_wolfi_distro(config_copy)
|
||||
|
||||
assert config_copy == original_config # Config should remain unchanged
|
||||
|
||||
|
||||
def test_extract_deployment_url_uses_custom_url():
|
||||
deployment = {"source_config": {"custom_url": "https://example.com/custom"}}
|
||||
assert _extract_deployment_url(deployment) == "https://example.com/custom"
|
||||
|
||||
|
||||
def test_extract_deployment_url_defaults_to_dash():
|
||||
assert _extract_deployment_url({"id": "dep-123"}) == "-"
|
||||
|
||||
|
||||
def test_format_deployments_table():
|
||||
output = format_deployments_table(
|
||||
[
|
||||
{
|
||||
"id": "dep-123",
|
||||
"name": "alpha",
|
||||
"source_config": {"custom_url": "https://alpha.example.com"},
|
||||
},
|
||||
{
|
||||
"id": "dep-456",
|
||||
"name": "beta",
|
||||
"url": "https://beta.example.com",
|
||||
},
|
||||
]
|
||||
)
|
||||
assert "Deployment ID" in output
|
||||
assert "Deployment Name" in output
|
||||
assert "Deployment URL" in output
|
||||
assert "dep-123" in output
|
||||
assert "alpha" in output
|
||||
assert "https://alpha.example.com" in output
|
||||
assert "dep-456" in output
|
||||
|
||||
@@ -1,67 +0,0 @@
|
||||
# RESUME Writes Stripping: Complete Flow Reference
|
||||
|
||||
## Legend
|
||||
|
||||
| Column | Meaning |
|
||||
|---|---|
|
||||
| **Level** | P = Parent, S = Subgraph |
|
||||
| **`is_replaying`** | `CONFIG_KEY_CHECKPOINT_ID` key exists in `config[CONF]` (line 249) |
|
||||
| **`__enter__` via** | Which branch loads the checkpoint: **ckpt_id** (explicit checkpoint_id in checkpoint_config), **replay_state** (parent's ReplayState), **latest** (fetch most recent) |
|
||||
| **`RESUMING`** | Value of `CONFIG_KEY_RESUMING` in configurable (set by parent for subgraphs, absent for outer graph) |
|
||||
| **`is_resuming`** | Computed at line 633 — controls whether to "proceed past previous checkpoint" |
|
||||
| **`in_map`** | `replaying_from_checkpoint_map` — subgraph's ns found in checkpoint_map |
|
||||
| **Strip?** | Are RESUME pending writes stripped? (line 662-671) |
|
||||
|
||||
## Setup
|
||||
|
||||
```
|
||||
Parent: START → executor (subgraph, checkpointer=True) → END
|
||||
Subgraph: START → step_a → ask_1 (interrupt) → ask_2 (interrupt) → END
|
||||
```
|
||||
|
||||
## The Table
|
||||
|
||||
| # | Scenario | Level | User call | `__enter__` via | `is_replaying` | `RESUMING` | `is_resuming` | `in_map` | Strip? | Why correct |
|
||||
|---|---|---|---|---|---|---|---|---|---|---|
|
||||
| 1 | **Fresh run** | P | `invoke({"v":[]}, cfg)` | latest (None) | False | _(absent)_ | False | — | N/A | No checkpoint yet, no writes to strip |
|
||||
| 1 | | S | _(Send from parent)_ | latest (None) | True¹ | False | False | False | N/A | No checkpoint yet |
|
||||
| 2 | **Resume single interrupt** | P | `invoke(Cmd(resume="a"), cfg)` | latest | False | _(absent)_ | True | — | No | Resuming — keep RESUME writes for interrupt() to return answer |
|
||||
| 2 | | S | _(Send)_ | latest | True¹ | True | True | False | No | `RESUMING=True` → keep. interrupt() returns "a", node completes |
|
||||
| 3 | **Resume 1st of 2 interrupts** | P | `invoke(Cmd(resume="a1"), cfg)` | latest | False | _(absent)_ | True | — | No | Keep RESUME writes — ask_1's answer must survive |
|
||||
| 3 | | S | _(Send)_ | latest | True¹ | True | True | False | **No** | ask_1 gets "a1" from RESUME write. ask_2 has no RESUME write → interrupt() re-fires. Correct. |
|
||||
| 4 | **Replay parent ckpt** (parent was mid-subgraph) | P | `invoke(None, parent_hist_cfg)` | ckpt_id | True | _(absent)_ | True | — | **Yes** | Replaying — strip stale RESUME writes so interrupts re-fire |
|
||||
| 4 | | S | _(Send)_ | replay_state² | True¹ | _(popped)_³ | False | False | **Yes** | `is_replaying=T`, `RESUMING` absent → strip. Subgraph replays cleanly |
|
||||
| 5 | **Time-travel to subgraph ckpt** (THE BUG) | P | `invoke(None, sub_cfg)` | ckpt_id⁴ | True | _(absent)_ | True | — | **Yes** | Parent replays from historical checkpoint |
|
||||
| 5 | | S | _(Send)_ | **ckpt_id**⁵ | True¹ | **True** | **True** | **True** | **Yes** ✨ | `in_map=True` overrides `RESUMING=True` → force strip. THE FIX. |
|
||||
| 5 | | S _(without fix)_ | _(Send)_ | ckpt_id⁵ | True¹ | **True** | **True** | _(no check)_ | **No** ❌ | BUG: `RESUMING=True` prevents strip → stale RESUME values → interrupt() doesn't re-fire |
|
||||
| 6 | **Fork from subgraph ckpt** | P | `invoke(None, update_state(sub_cfg,...))` | ckpt_id | True | _(absent)_ | True | — | **Yes** | Same as case 5 — fork creates new ckpt, but checkpoint_map still resolves |
|
||||
| 6 | | S | _(Send)_ | ckpt_id⁵ | True¹ | True | True | **True** | **Yes** ✨ | Same fix applies |
|
||||
| 7 | **Resume after case 5 re-interrupts** | P | `invoke(Cmd(resume="a2"), cfg)` | latest | False | _(absent)_ | True | — | No | Normal resume — keep RESUME writes |
|
||||
| 7 | | S | _(Send)_ | latest | True¹ | True | True | False⁶ | **No** | ask_2 gets "a2" from fresh RESUME write. Correct. |
|
||||
|
||||
## Footnotes
|
||||
|
||||
**¹** `is_replaying` is always `True` for subgraphs on tick 1 because `_algo.py` sets `CONFIG_KEY_CHECKPOINT_ID: None` — the key exists (even with `None` value), so `key in dict` is `True`. After tick 1, line 563 sets `is_replaying = False`.
|
||||
|
||||
**²** `replay_state` branch: parent passed `CONFIG_KEY_REPLAY_STATE = ReplayState(parent_ckpt_id)`. The subgraph uses `replay_state.get_checkpoint()` which does `checkpointer.list(before=parent_ckpt_id, limit=1)` to find the subgraph's checkpoint from before the replay point.
|
||||
|
||||
**³** The `replay_state` branch in `__enter__` (line 1158) explicitly pops `CONFIG_KEY_RESUMING` from config. This makes `is_resuming = False` in `_first()` because for nested graphs the fallback (`self.input is None or input_is_command`) is False (input is a Send arg).
|
||||
|
||||
**⁴** Parent `__init__` clears `checkpoint_ns → ""` and `checkpoint_id → None` (line 273-277), then resolves `""` from checkpoint_map → gets `parent_checkpoint_id` onto `checkpoint_config` (line 278-290).
|
||||
|
||||
**⁵** Subgraph `__init__` resolves its namespace (e.g. `"executor:task_id"`) from checkpoint_map → gets `subgraph_checkpoint_id` onto `checkpoint_config`. This is why the new first branch in `__enter__` (line 1141) fires — `checkpoint_config` has a truthy `checkpoint_id`.
|
||||
|
||||
**⁶** After case 5 completes/re-interrupts and user resumes, the config is a normal thread config with no checkpoint_map entry for the subgraph. `in_map` is False, so normal resume logic applies.
|
||||
|
||||
## The core tension (case 5)
|
||||
|
||||
The parent **can't distinguish** these cases when propagating flags to subgraphs:
|
||||
|
||||
| Parent sees | What's actually happening | Subgraph should strip RESUME? |
|
||||
|---|---|---|
|
||||
| `input=None`, has checkpoint | Resume after interrupt | Yes (replaying) |
|
||||
| `input=None`, has checkpoint | Resume after interrupt | No (resuming) |
|
||||
| `input=Command(resume=...)` | Active resume | No (resuming) |
|
||||
| `input=None`, has checkpoint | Time-travel to subgraph | Yes (replaying) |
|
||||
|
||||
The **only** distinguishing signal at the subgraph level is whether its namespace appears in `checkpoint_map`.
|
||||
@@ -647,24 +647,13 @@ class PregelLoop:
|
||||
# writes so that interrupt() calls re-fire instead of returning
|
||||
# stale values. But if we're actively resuming, keep them —
|
||||
# multi-interrupt scenarios need previously resolved values preserved.
|
||||
if self.is_replaying and (
|
||||
# Time-travel to a subgraph checkpoint: the parent sets
|
||||
# RESUMING=True (it can't distinguish time-travel from resume),
|
||||
# so we check if this subgraph's own ns is in checkpoint_map.
|
||||
# Normally the map only has ancestor entries (_algo.py); the
|
||||
# subgraph's own entry only appears via get_state(subgraphs=True).
|
||||
(
|
||||
self.is_nested
|
||||
and configurable.get(CONFIG_KEY_CHECKPOINT_NS, "")
|
||||
in configurable.get(CONFIG_KEY_CHECKPOINT_MAP, {})
|
||||
)
|
||||
or not (
|
||||
# Outer graph: resume arrives as Command(resume=...)
|
||||
(input_is_command and cast(Command, self.input).resume is not None)
|
||||
# Subgraphs: resume arrives via config flag from parent
|
||||
# (subgraph input is a Send arg, not a Command)
|
||||
or configurable.get(CONFIG_KEY_RESUMING, False)
|
||||
)
|
||||
# We check two conditions because resume signals arrive differently:
|
||||
# - Command(resume=...): the outer graph receives resume via input
|
||||
# - CONFIG_KEY_RESUMING: child subgraphs receive it via config from
|
||||
# the parent (their input is a Send arg, not a Command)
|
||||
if self.is_replaying and not (
|
||||
(input_is_command and cast(Command, self.input).resume is not None)
|
||||
or configurable.get(CONFIG_KEY_RESUMING, False)
|
||||
):
|
||||
self.checkpoint_pending_writes = [
|
||||
w for w in self.checkpoint_pending_writes if w[1] != RESUME
|
||||
@@ -1138,15 +1127,9 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
|
||||
def __enter__(self) -> Self:
|
||||
if not self.checkpointer:
|
||||
saved = None
|
||||
elif self.checkpoint_config[CONF].get(CONFIG_KEY_CHECKPOINT_ID):
|
||||
# Explicit checkpoint_id requested — fetch that exact checkpoint.
|
||||
# This covers both normal replay and subgraphs resolved via
|
||||
# checkpoint_map during time-travel.
|
||||
saved = self.checkpointer.get_tuple(self.checkpoint_config)
|
||||
elif replay_state := self.config[CONF].get(CONFIG_KEY_REPLAY_STATE):
|
||||
# Subgraph replay: the parent is replaying and passed us a
|
||||
# replay_state with its checkpoint_id. Look up our checkpoint
|
||||
# from the parent's checkpoint_map instead of fetching latest.
|
||||
elif self.is_nested and (
|
||||
replay_state := self.config[CONF].get(CONFIG_KEY_REPLAY_STATE)
|
||||
):
|
||||
saved = replay_state.get_checkpoint(
|
||||
self.config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, ""),
|
||||
self.checkpointer,
|
||||
@@ -1158,7 +1141,9 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
|
||||
self.config[CONF].pop(CONFIG_KEY_RESUMING, None)
|
||||
else:
|
||||
# Normal case: fetch the most recent checkpoint for this
|
||||
# graph/thread. Returns None on first invocation.
|
||||
# graph/thread. If a specific checkpoint_id is in the config,
|
||||
# fetch that exact checkpoint; otherwise fetch the latest one.
|
||||
# Returns None on first invocation (no checkpoints exist yet).
|
||||
saved = self.checkpointer.get_tuple(self.checkpoint_config)
|
||||
|
||||
if saved is None:
|
||||
@@ -1337,15 +1322,9 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
|
||||
async def __aenter__(self) -> Self:
|
||||
if not self.checkpointer:
|
||||
saved = None
|
||||
elif self.checkpoint_config[CONF].get(CONFIG_KEY_CHECKPOINT_ID):
|
||||
# Explicit checkpoint_id requested — fetch that exact checkpoint.
|
||||
# This covers both normal replay and subgraphs resolved via
|
||||
# checkpoint_map during time-travel.
|
||||
saved = await self.checkpointer.aget_tuple(self.checkpoint_config)
|
||||
elif replay_state := self.config[CONF].get(CONFIG_KEY_REPLAY_STATE):
|
||||
# Subgraph replay: the parent is replaying and passed us a
|
||||
# replay_state with its checkpoint_id. Look up our checkpoint
|
||||
# from the parent's checkpoint_map instead of fetching latest.
|
||||
elif self.is_nested and (
|
||||
replay_state := self.config[CONF].get(CONFIG_KEY_REPLAY_STATE)
|
||||
):
|
||||
saved = await replay_state.aget_checkpoint(
|
||||
self.config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, ""),
|
||||
self.checkpointer,
|
||||
@@ -1357,7 +1336,9 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
|
||||
self.config[CONF].pop(CONFIG_KEY_RESUMING, None)
|
||||
else:
|
||||
# Normal case: fetch the most recent checkpoint for this
|
||||
# graph/thread. Returns None on first invocation.
|
||||
# graph/thread. If a specific checkpoint_id is in the config,
|
||||
# fetch that exact checkpoint; otherwise fetch the latest one.
|
||||
# Returns None on first invocation (no checkpoints exist yet).
|
||||
saved = await self.checkpointer.aget_tuple(self.checkpoint_config)
|
||||
|
||||
if saved is None:
|
||||
|
||||
@@ -2456,7 +2456,7 @@ class Pregel(
|
||||
debug: bool | None = None,
|
||||
version: Literal["v2"],
|
||||
**kwargs: Unpack[DeprecatedKwargs],
|
||||
) -> Iterator[StreamPart[StateT, OutputT]]: ...
|
||||
) -> Iterator[StreamPart[OutputT, StateT]]: ...
|
||||
|
||||
@overload
|
||||
def stream(
|
||||
@@ -2787,7 +2787,7 @@ class Pregel(
|
||||
debug: bool | None = None,
|
||||
version: Literal["v2"],
|
||||
**kwargs: Unpack[DeprecatedKwargs],
|
||||
) -> AsyncIterator[StreamPart[StateT, OutputT]]: ...
|
||||
) -> AsyncIterator[StreamPart[OutputT, StateT]]: ...
|
||||
|
||||
@overload
|
||||
def astream(
|
||||
@@ -3194,7 +3194,7 @@ class Pregel(
|
||||
durability: Durability | None = None,
|
||||
version: Literal["v2"],
|
||||
**kwargs: Any,
|
||||
) -> list[StreamPart[StateT, OutputT]]: ...
|
||||
) -> list[StreamPart[OutputT, StateT]]: ...
|
||||
|
||||
@overload
|
||||
def invoke(
|
||||
@@ -3364,7 +3364,7 @@ class Pregel(
|
||||
durability: Durability | None = None,
|
||||
version: Literal["v2"],
|
||||
**kwargs: Any,
|
||||
) -> list[StreamPart[StateT, OutputT]]: ...
|
||||
) -> list[StreamPart[OutputT, StateT]]: ...
|
||||
|
||||
@overload
|
||||
async def ainvoke(
|
||||
|
||||
@@ -117,7 +117,7 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, ContextT, InputT, Ou
|
||||
interrupt_after: All | Sequence[str] | None = None,
|
||||
subgraphs: bool = False,
|
||||
version: Literal["v2"],
|
||||
) -> Iterator[StreamPart[StateT, OutputT]]: ...
|
||||
) -> Iterator[StreamPart[OutputT, StateT]]: ...
|
||||
|
||||
@overload
|
||||
@abstractmethod
|
||||
@@ -161,7 +161,7 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, ContextT, InputT, Ou
|
||||
interrupt_after: All | Sequence[str] | None = None,
|
||||
subgraphs: bool = False,
|
||||
version: Literal["v2"],
|
||||
) -> AsyncIterator[StreamPart[StateT, OutputT]]: ...
|
||||
) -> AsyncIterator[StreamPart[OutputT, StateT]]: ...
|
||||
|
||||
@overload
|
||||
@abstractmethod
|
||||
|
||||
@@ -5,7 +5,6 @@ from collections.abc import AsyncIterator, Iterator, Sequence
|
||||
from dataclasses import asdict
|
||||
from typing import (
|
||||
Any,
|
||||
Generic,
|
||||
Literal,
|
||||
cast,
|
||||
overload,
|
||||
@@ -66,7 +65,6 @@ from langgraph.types import (
|
||||
StreamMode,
|
||||
StreamPart,
|
||||
)
|
||||
from langgraph.typing import ContextT, InputT, OutputT, StateT
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -110,10 +108,7 @@ class RemoteException(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class RemoteGraph(
|
||||
PregelProtocol[StateT, ContextT, InputT, OutputT],
|
||||
Generic[StateT, ContextT, InputT, OutputT],
|
||||
):
|
||||
class RemoteGraph(PregelProtocol):
|
||||
"""The `RemoteGraph` class is a client implementation for calling remote
|
||||
APIs that implement the LangGraph Server API specification.
|
||||
|
||||
@@ -693,10 +688,9 @@ class RemoteGraph(
|
||||
@overload
|
||||
def stream(
|
||||
self,
|
||||
input: InputT | Command | None,
|
||||
input: dict[str, Any] | Any,
|
||||
config: RunnableConfig | None = None,
|
||||
*,
|
||||
context: ContextT | None = None,
|
||||
stream_mode: StreamMode | list[StreamMode] | None = None,
|
||||
interrupt_before: All | Sequence[str] | None = None,
|
||||
interrupt_after: All | Sequence[str] | None = None,
|
||||
@@ -705,15 +699,14 @@ class RemoteGraph(
|
||||
params: QueryParamTypes | None = None,
|
||||
version: Literal["v2"],
|
||||
**kwargs: Any,
|
||||
) -> Iterator[StreamPart[StateT, OutputT]]: ...
|
||||
) -> Iterator[StreamPart]: ...
|
||||
|
||||
@overload
|
||||
def stream(
|
||||
self,
|
||||
input: InputT | Command | None,
|
||||
input: dict[str, Any] | Any,
|
||||
config: RunnableConfig | None = None,
|
||||
*,
|
||||
context: ContextT | None = None,
|
||||
stream_mode: StreamMode | list[StreamMode] | None = None,
|
||||
interrupt_before: All | Sequence[str] | None = None,
|
||||
interrupt_after: All | Sequence[str] | None = None,
|
||||
@@ -726,10 +719,9 @@ class RemoteGraph(
|
||||
|
||||
def stream(
|
||||
self,
|
||||
input: InputT | Command | None,
|
||||
input: dict[str, Any] | Any,
|
||||
config: RunnableConfig | None = None,
|
||||
*,
|
||||
context: ContextT | None = None,
|
||||
stream_mode: StreamMode | list[StreamMode] | None = None,
|
||||
interrupt_before: All | Sequence[str] | None = None,
|
||||
interrupt_after: All | Sequence[str] | None = None,
|
||||
@@ -777,7 +769,6 @@ class RemoteGraph(
|
||||
input=input,
|
||||
command=command,
|
||||
config=sanitized_config,
|
||||
context=context,
|
||||
stream_mode=stream_modes,
|
||||
interrupt_before=interrupt_before,
|
||||
interrupt_after=interrupt_after,
|
||||
@@ -848,10 +839,9 @@ class RemoteGraph(
|
||||
@overload
|
||||
def astream(
|
||||
self,
|
||||
input: InputT | Command | None,
|
||||
input: dict[str, Any] | Any,
|
||||
config: RunnableConfig | None = None,
|
||||
*,
|
||||
context: ContextT | None = None,
|
||||
stream_mode: StreamMode | list[StreamMode] | None = None,
|
||||
interrupt_before: All | Sequence[str] | None = None,
|
||||
interrupt_after: All | Sequence[str] | None = None,
|
||||
@@ -860,15 +850,14 @@ class RemoteGraph(
|
||||
params: QueryParamTypes | None = None,
|
||||
version: Literal["v2"],
|
||||
**kwargs: Any,
|
||||
) -> AsyncIterator[StreamPart[StateT, OutputT]]: ...
|
||||
) -> AsyncIterator[StreamPart]: ...
|
||||
|
||||
@overload
|
||||
def astream(
|
||||
self,
|
||||
input: InputT | Command | None,
|
||||
input: dict[str, Any] | Any,
|
||||
config: RunnableConfig | None = None,
|
||||
*,
|
||||
context: ContextT | None = None,
|
||||
stream_mode: StreamMode | list[StreamMode] | None = None,
|
||||
interrupt_before: All | Sequence[str] | None = None,
|
||||
interrupt_after: All | Sequence[str] | None = None,
|
||||
@@ -881,10 +870,9 @@ class RemoteGraph(
|
||||
|
||||
async def astream(
|
||||
self,
|
||||
input: InputT | Command | None,
|
||||
input: dict[str, Any] | Any,
|
||||
config: RunnableConfig | None = None,
|
||||
*,
|
||||
context: ContextT | None = None,
|
||||
stream_mode: StreamMode | list[StreamMode] | None = None,
|
||||
interrupt_before: All | Sequence[str] | None = None,
|
||||
interrupt_after: All | Sequence[str] | None = None,
|
||||
@@ -932,7 +920,6 @@ class RemoteGraph(
|
||||
input=input,
|
||||
command=command,
|
||||
config=sanitized_config,
|
||||
context=context,
|
||||
stream_mode=stream_modes,
|
||||
interrupt_before=interrupt_before,
|
||||
interrupt_after=interrupt_after,
|
||||
@@ -1019,25 +1006,23 @@ class RemoteGraph(
|
||||
@overload
|
||||
def invoke(
|
||||
self,
|
||||
input: InputT | Command | None,
|
||||
input: dict[str, Any] | Any,
|
||||
config: RunnableConfig | None = None,
|
||||
*,
|
||||
context: ContextT | None = None,
|
||||
interrupt_before: All | Sequence[str] | None = None,
|
||||
interrupt_after: All | Sequence[str] | None = None,
|
||||
headers: dict[str, str] | None = None,
|
||||
params: QueryParamTypes | None = None,
|
||||
version: Literal["v2"],
|
||||
**kwargs: Any,
|
||||
) -> GraphOutput[OutputT]: ...
|
||||
) -> GraphOutput[dict[str, Any]]: ...
|
||||
|
||||
@overload
|
||||
def invoke(
|
||||
self,
|
||||
input: InputT | Command | None,
|
||||
input: dict[str, Any] | Any,
|
||||
config: RunnableConfig | None = None,
|
||||
*,
|
||||
context: ContextT | None = None,
|
||||
interrupt_before: All | Sequence[str] | None = None,
|
||||
interrupt_after: All | Sequence[str] | None = None,
|
||||
headers: dict[str, str] | None = None,
|
||||
@@ -1048,10 +1033,9 @@ class RemoteGraph(
|
||||
|
||||
def invoke(
|
||||
self,
|
||||
input: InputT | Command | None,
|
||||
input: dict[str, Any] | Any,
|
||||
config: RunnableConfig | None = None,
|
||||
*,
|
||||
context: ContextT | None = None,
|
||||
interrupt_before: All | Sequence[str] | None = None,
|
||||
interrupt_after: All | Sequence[str] | None = None,
|
||||
headers: dict[str, str] | None = None,
|
||||
@@ -1077,7 +1061,6 @@ class RemoteGraph(
|
||||
for chunk in self.stream( # type: ignore[misc, call-overload]
|
||||
input,
|
||||
config=config,
|
||||
context=context,
|
||||
interrupt_before=interrupt_before,
|
||||
interrupt_after=interrupt_after,
|
||||
headers=headers,
|
||||
@@ -1101,25 +1084,23 @@ class RemoteGraph(
|
||||
@overload
|
||||
async def ainvoke(
|
||||
self,
|
||||
input: InputT | Command | None,
|
||||
input: dict[str, Any] | Any,
|
||||
config: RunnableConfig | None = None,
|
||||
*,
|
||||
context: ContextT | None = None,
|
||||
interrupt_before: All | Sequence[str] | None = None,
|
||||
interrupt_after: All | Sequence[str] | None = None,
|
||||
headers: dict[str, str] | None = None,
|
||||
params: QueryParamTypes | None = None,
|
||||
version: Literal["v2"],
|
||||
**kwargs: Any,
|
||||
) -> GraphOutput[OutputT]: ...
|
||||
) -> GraphOutput[dict[str, Any]]: ...
|
||||
|
||||
@overload
|
||||
async def ainvoke(
|
||||
self,
|
||||
input: InputT | Command | None,
|
||||
input: dict[str, Any] | Any,
|
||||
config: RunnableConfig | None = None,
|
||||
*,
|
||||
context: ContextT | None = None,
|
||||
interrupt_before: All | Sequence[str] | None = None,
|
||||
interrupt_after: All | Sequence[str] | None = None,
|
||||
headers: dict[str, str] | None = None,
|
||||
@@ -1130,10 +1111,9 @@ class RemoteGraph(
|
||||
|
||||
async def ainvoke(
|
||||
self,
|
||||
input: InputT | Command | None,
|
||||
input: dict[str, Any] | Any,
|
||||
config: RunnableConfig | None = None,
|
||||
*,
|
||||
context: ContextT | None = None,
|
||||
interrupt_before: All | Sequence[str] | None = None,
|
||||
interrupt_after: All | Sequence[str] | None = None,
|
||||
headers: dict[str, str] | None = None,
|
||||
@@ -1159,7 +1139,6 @@ class RemoteGraph(
|
||||
async for chunk in self.astream( # type: ignore[misc, call-overload]
|
||||
input,
|
||||
config=config,
|
||||
context=context,
|
||||
interrupt_before=interrupt_before,
|
||||
interrupt_after=interrupt_after,
|
||||
headers=headers,
|
||||
|
||||
@@ -335,7 +335,7 @@ StreamPart = TypeAliasType(
|
||||
| CheckpointStreamPart[StateT]
|
||||
| TasksStreamPart
|
||||
| DebugStreamPart[StateT],
|
||||
type_params=(StateT, OutputT),
|
||||
type_params=(OutputT, StateT),
|
||||
)
|
||||
"""A discriminated union of all v2 stream part types.
|
||||
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "langgraph"
|
||||
version = "1.1.1"
|
||||
version = "1.1.0"
|
||||
description = "Building stateful, multi-actor applications with LLMs"
|
||||
authors = []
|
||||
requires-python = ">=3.10"
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
import re
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
from typing import Annotated
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
@@ -11,7 +10,6 @@ from langchain_core.runnables import RunnableConfig
|
||||
from langchain_core.runnables.graph import Edge as DrawableEdge
|
||||
from langchain_core.runnables.graph import Node as DrawableNode
|
||||
from langgraph_sdk.schema import StreamPart
|
||||
from pydantic import BaseModel
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.errors import GraphInterrupt
|
||||
@@ -910,188 +908,6 @@ async def test_ainvoke():
|
||||
assert result == {"messages": [{"type": "human", "content": "world"}]}
|
||||
|
||||
|
||||
def test_stream_context():
|
||||
"""Test that context is passed through to the SDK client in stream."""
|
||||
mock_sync_client = MagicMock()
|
||||
mock_sync_client.runs.stream.return_value = [
|
||||
StreamPart(event="values", data={"chunk": "data1"}),
|
||||
]
|
||||
|
||||
remote_pregel = RemoteGraph(
|
||||
"test_graph_id",
|
||||
sync_client=mock_sync_client,
|
||||
)
|
||||
|
||||
config = {"configurable": {"thread_id": "thread_1"}}
|
||||
context = {"model_name": "anthropic", "user_id": "123"}
|
||||
stream_parts = list(
|
||||
remote_pregel.stream(
|
||||
{"input": "data"},
|
||||
config,
|
||||
context=context,
|
||||
stream_mode="values",
|
||||
)
|
||||
)
|
||||
|
||||
assert stream_parts == [{"chunk": "data1"}]
|
||||
_, kwargs = mock_sync_client.runs.stream.call_args
|
||||
assert kwargs["context"] == {"model_name": "anthropic", "user_id": "123"}
|
||||
|
||||
|
||||
def test_stream_context_none():
|
||||
"""Test that context defaults to None when not provided."""
|
||||
mock_sync_client = MagicMock()
|
||||
mock_sync_client.runs.stream.return_value = [
|
||||
StreamPart(event="values", data={"chunk": "data1"}),
|
||||
]
|
||||
|
||||
remote_pregel = RemoteGraph(
|
||||
"test_graph_id",
|
||||
sync_client=mock_sync_client,
|
||||
)
|
||||
|
||||
config = {"configurable": {"thread_id": "thread_1"}}
|
||||
list(remote_pregel.stream({"input": "data"}, config, stream_mode="values"))
|
||||
|
||||
_, kwargs = mock_sync_client.runs.stream.call_args
|
||||
assert kwargs["context"] is None
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_astream_context():
|
||||
"""Test that context is passed through to the SDK client in astream."""
|
||||
mock_async_client = MagicMock()
|
||||
async_iter = MagicMock()
|
||||
async_iter.__aiter__.return_value = [
|
||||
StreamPart(event="values", data={"chunk": "data1"}),
|
||||
]
|
||||
mock_async_client.runs.stream.return_value = async_iter
|
||||
|
||||
remote_pregel = RemoteGraph(
|
||||
"test_graph_id",
|
||||
client=mock_async_client,
|
||||
)
|
||||
|
||||
config = {"configurable": {"thread_id": "thread_1"}}
|
||||
context = {"model_name": "anthropic"}
|
||||
chunks = []
|
||||
async for chunk in remote_pregel.astream(
|
||||
{"input": "data"},
|
||||
config,
|
||||
context=context,
|
||||
stream_mode="values",
|
||||
):
|
||||
chunks.append(chunk)
|
||||
|
||||
assert chunks == [{"chunk": "data1"}]
|
||||
_, kwargs = mock_async_client.runs.stream.call_args
|
||||
assert kwargs["context"] == {"model_name": "anthropic"}
|
||||
|
||||
|
||||
def test_invoke_context():
|
||||
"""Test that context is passed through to the SDK client in invoke."""
|
||||
mock_sync_client = MagicMock()
|
||||
mock_sync_client.runs.stream.return_value = [
|
||||
StreamPart(event="values", data={"result": "done"}),
|
||||
]
|
||||
|
||||
remote_pregel = RemoteGraph(
|
||||
"test_graph_id",
|
||||
sync_client=mock_sync_client,
|
||||
)
|
||||
|
||||
config = {"configurable": {"thread_id": "thread_1"}}
|
||||
context = {"model_name": "openai"}
|
||||
result = remote_pregel.invoke({"input": "data"}, config, context=context)
|
||||
|
||||
assert result == {"result": "done"}
|
||||
_, kwargs = mock_sync_client.runs.stream.call_args
|
||||
assert kwargs["context"] == {"model_name": "openai"}
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_ainvoke_context():
|
||||
"""Test that context is passed through to the SDK client in ainvoke."""
|
||||
mock_async_client = MagicMock()
|
||||
async_iter = MagicMock()
|
||||
async_iter.__aiter__.return_value = [
|
||||
StreamPart(event="values", data={"result": "done"}),
|
||||
]
|
||||
mock_async_client.runs.stream.return_value = async_iter
|
||||
|
||||
remote_pregel = RemoteGraph(
|
||||
"test_graph_id",
|
||||
client=mock_async_client,
|
||||
)
|
||||
|
||||
config = {"configurable": {"thread_id": "thread_1"}}
|
||||
context = {"user_id": "456"}
|
||||
result = await remote_pregel.ainvoke({"input": "data"}, config, context=context)
|
||||
|
||||
assert result == {"result": "done"}
|
||||
_, kwargs = mock_async_client.runs.stream.call_args
|
||||
assert kwargs["context"] == {"user_id": "456"}
|
||||
|
||||
|
||||
def test_stream_context_dataclass():
|
||||
"""Test that a dataclass context is passed through to the SDK client."""
|
||||
|
||||
@dataclass
|
||||
class MyContext:
|
||||
model_name: str
|
||||
user_id: str
|
||||
|
||||
mock_sync_client = MagicMock()
|
||||
mock_sync_client.runs.stream.return_value = [
|
||||
StreamPart(event="values", data={"chunk": "data1"}),
|
||||
]
|
||||
|
||||
remote_pregel = RemoteGraph(
|
||||
"test_graph_id",
|
||||
sync_client=mock_sync_client,
|
||||
)
|
||||
|
||||
config = {"configurable": {"thread_id": "thread_1"}}
|
||||
ctx = MyContext(model_name="anthropic", user_id="123")
|
||||
list(
|
||||
remote_pregel.stream(
|
||||
{"input": "data"}, config, context=ctx, stream_mode="values"
|
||||
)
|
||||
)
|
||||
|
||||
_, kwargs = mock_sync_client.runs.stream.call_args
|
||||
assert kwargs["context"] == ctx
|
||||
|
||||
|
||||
def test_stream_context_base_model():
|
||||
"""Test that a BaseModel context is passed through to the SDK client."""
|
||||
|
||||
class MyContext(BaseModel):
|
||||
model_name: str
|
||||
user_id: str
|
||||
|
||||
mock_sync_client = MagicMock()
|
||||
mock_sync_client.runs.stream.return_value = [
|
||||
StreamPart(event="values", data={"chunk": "data1"}),
|
||||
]
|
||||
|
||||
remote_pregel = RemoteGraph(
|
||||
"test_graph_id",
|
||||
sync_client=mock_sync_client,
|
||||
)
|
||||
|
||||
config = {"configurable": {"thread_id": "thread_1"}}
|
||||
ctx = MyContext(model_name="anthropic", user_id="123")
|
||||
list(
|
||||
remote_pregel.stream(
|
||||
{"input": "data"}, config, context=ctx, stream_mode="values"
|
||||
)
|
||||
)
|
||||
|
||||
_, kwargs = mock_sync_client.runs.stream.call_args
|
||||
assert kwargs["context"] == ctx
|
||||
|
||||
|
||||
@pytest.mark.skip(
|
||||
"Unskip this test to manually test the LangSmith Deployment integration"
|
||||
)
|
||||
|
||||
@@ -1116,631 +1116,6 @@ def test_subgraph_replay_from_subgraph_checkpoint(
|
||||
assert "post" in final_result["value"]
|
||||
|
||||
|
||||
def test_subgraph_time_travel_to_first_interrupt(
|
||||
sync_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
"""Time travel to a subgraph checkpoint at the FIRST interrupt.
|
||||
|
||||
Architecture:
|
||||
Parent: START --> executor (subgraph, checkpointer=True) --> END
|
||||
Executor: START --> step_a --> ask_1 (interrupt) --> ask_2 (interrupt) --> END
|
||||
|
||||
Flow: run through both interrupts, then time travel back to the subgraph
|
||||
checkpoint captured at the first interrupt. ask_1 should re-fire,
|
||||
step_a should NOT re-run. Then resume through both interrupts with new answers.
|
||||
"""
|
||||
|
||||
called: list[str] = []
|
||||
|
||||
def step_a(state: State) -> State:
|
||||
called.append("step_a")
|
||||
return {"value": ["step_a_done"]}
|
||||
|
||||
def ask_1(state: State) -> State:
|
||||
called.append("ask_1")
|
||||
answer = interrupt("Question 1?")
|
||||
return {"value": [f"ask_1:{answer}"]}
|
||||
|
||||
def ask_2(state: State) -> State:
|
||||
called.append("ask_2")
|
||||
answer = interrupt("Question 2?")
|
||||
return {"value": [f"ask_2:{answer}"]}
|
||||
|
||||
executor = (
|
||||
StateGraph(State)
|
||||
.add_node("step_a", step_a)
|
||||
.add_node("ask_1", ask_1)
|
||||
.add_node("ask_2", ask_2)
|
||||
.add_edge(START, "step_a")
|
||||
.add_edge("step_a", "ask_1")
|
||||
.add_edge("ask_1", "ask_2")
|
||||
.add_edge("ask_2", "__end__")
|
||||
.compile(checkpointer=True)
|
||||
)
|
||||
|
||||
graph = (
|
||||
StateGraph(State)
|
||||
.add_node("executor", executor)
|
||||
.add_edge(START, "executor")
|
||||
.compile(checkpointer=sync_checkpointer)
|
||||
)
|
||||
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
# Run until first interrupt (ask_1)
|
||||
result = graph.invoke({"value": []}, config)
|
||||
assert "__interrupt__" in result
|
||||
assert result["__interrupt__"][0].value == "Question 1?"
|
||||
|
||||
# Capture subgraph state at the first interrupt
|
||||
parent_state = graph.get_state(config, subgraphs=True)
|
||||
sub_config_at_first = parent_state.tasks[0].state.config
|
||||
|
||||
# Resume first interrupt
|
||||
result = graph.invoke(Command(resume="answer_1"), config)
|
||||
assert result["__interrupt__"][0].value == "Question 2?"
|
||||
|
||||
# Resume second interrupt to complete
|
||||
result = graph.invoke(Command(resume="answer_2"), config)
|
||||
assert "__interrupt__" not in result
|
||||
|
||||
# --- Scenario 1: Replay from subgraph checkpoint at 1st interrupt ---
|
||||
called.clear()
|
||||
replay_result = graph.invoke(None, sub_config_at_first)
|
||||
assert "__interrupt__" in replay_result
|
||||
assert replay_result["__interrupt__"][0].value == "Question 1?"
|
||||
# step_a should NOT re-run — it was before this checkpoint
|
||||
assert "step_a" not in called
|
||||
# ask_1 re-fires because the interrupt replays
|
||||
assert "ask_1" in called
|
||||
|
||||
# --- Scenario 2: Fork from subgraph checkpoint at 1st interrupt ---
|
||||
called.clear()
|
||||
fork_config = graph.update_state(sub_config_at_first, {"value": ["forked"]})
|
||||
fork_result = graph.invoke(None, fork_config)
|
||||
assert "__interrupt__" in fork_result
|
||||
assert fork_result["__interrupt__"][0].value == "Question 1?"
|
||||
assert "step_a" not in called
|
||||
assert "ask_1" in called
|
||||
|
||||
|
||||
def test_subgraph_time_travel_to_second_interrupt(
|
||||
sync_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
"""Time travel to a subgraph checkpoint at the SECOND interrupt.
|
||||
|
||||
Architecture:
|
||||
Parent: START --> executor (subgraph, checkpointer=True) --> END
|
||||
Executor: START --> step_a --> ask_1 (interrupt) --> ask_2 (interrupt) --> END
|
||||
|
||||
Flow: run through both interrupts resuming each, then time travel back to the
|
||||
subgraph checkpoint at the second interrupt. Only ask_2 should re-fire.
|
||||
Then resume with a new answer and verify state.
|
||||
"""
|
||||
|
||||
called: list[str] = []
|
||||
|
||||
def step_a(state: State) -> State:
|
||||
called.append("step_a")
|
||||
return {"value": ["step_a_done"]}
|
||||
|
||||
def ask_1(state: State) -> State:
|
||||
called.append("ask_1")
|
||||
answer = interrupt("Question 1?")
|
||||
return {"value": [f"ask_1:{answer}"]}
|
||||
|
||||
def ask_2(state: State) -> State:
|
||||
called.append("ask_2")
|
||||
answer = interrupt("Question 2?")
|
||||
return {"value": [f"ask_2:{answer}"]}
|
||||
|
||||
executor = (
|
||||
StateGraph(State)
|
||||
.add_node("step_a", step_a)
|
||||
.add_node("ask_1", ask_1)
|
||||
.add_node("ask_2", ask_2)
|
||||
.add_edge(START, "step_a")
|
||||
.add_edge("step_a", "ask_1")
|
||||
.add_edge("ask_1", "ask_2")
|
||||
.add_edge("ask_2", "__end__")
|
||||
.compile(checkpointer=True)
|
||||
)
|
||||
|
||||
graph = (
|
||||
StateGraph(State)
|
||||
.add_node("executor", executor)
|
||||
.add_edge(START, "executor")
|
||||
.compile(checkpointer=sync_checkpointer)
|
||||
)
|
||||
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
# Run until first interrupt (ask_1)
|
||||
result = graph.invoke({"value": []}, config)
|
||||
assert result["__interrupt__"][0].value == "Question 1?"
|
||||
|
||||
# Resume first interrupt
|
||||
result = graph.invoke(Command(resume="answer_1"), config)
|
||||
assert result["__interrupt__"][0].value == "Question 2?"
|
||||
|
||||
# Capture subgraph state at the second interrupt
|
||||
parent_state = graph.get_state(config, subgraphs=True)
|
||||
sub_config = parent_state.tasks[0].state.config
|
||||
|
||||
# Resume second interrupt to complete the graph
|
||||
result = graph.invoke(Command(resume="answer_2"), config)
|
||||
assert "__interrupt__" not in result
|
||||
|
||||
# --- Scenario 1: Replay from subgraph checkpoint at 2nd interrupt ---
|
||||
called.clear()
|
||||
replay_result = graph.invoke(None, sub_config)
|
||||
assert "__interrupt__" in replay_result
|
||||
assert replay_result["__interrupt__"][0].value == "Question 2?"
|
||||
# step_a and ask_1 should NOT re-run — they were before this checkpoint
|
||||
assert "step_a" not in called
|
||||
assert "ask_1" not in called
|
||||
|
||||
# --- Scenario 2: Fork from subgraph checkpoint at 2nd interrupt ---
|
||||
called.clear()
|
||||
fork_config = graph.update_state(sub_config, {"value": ["forked"]})
|
||||
fork_result = graph.invoke(None, fork_config)
|
||||
assert "__interrupt__" in fork_result
|
||||
assert fork_result["__interrupt__"][0].value == "Question 2?"
|
||||
assert "step_a" not in called
|
||||
assert "ask_1" not in called
|
||||
|
||||
|
||||
def test_subgraph_time_travel_after_completion(
|
||||
sync_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
"""Time travel to a subgraph checkpoint AFTER both interrupts are resolved.
|
||||
|
||||
Architecture:
|
||||
Parent: START --> executor (subgraph, checkpointer=True) --> END
|
||||
Executor: START --> step_a --> ask_1 (interrupt) --> ask_2 (interrupt) --> END
|
||||
|
||||
After completing the full flow, capture the subgraph's final state checkpoint
|
||||
and replay from it — should be a no-op (no nodes re-run).
|
||||
"""
|
||||
|
||||
called: list[str] = []
|
||||
|
||||
def step_a(state: State) -> State:
|
||||
called.append("step_a")
|
||||
return {"value": ["step_a_done"]}
|
||||
|
||||
def ask_1(state: State) -> State:
|
||||
called.append("ask_1")
|
||||
answer = interrupt("Question 1?")
|
||||
return {"value": [f"ask_1:{answer}"]}
|
||||
|
||||
def ask_2(state: State) -> State:
|
||||
called.append("ask_2")
|
||||
answer = interrupt("Question 2?")
|
||||
return {"value": [f"ask_2:{answer}"]}
|
||||
|
||||
executor = (
|
||||
StateGraph(State)
|
||||
.add_node("step_a", step_a)
|
||||
.add_node("ask_1", ask_1)
|
||||
.add_node("ask_2", ask_2)
|
||||
.add_edge(START, "step_a")
|
||||
.add_edge("step_a", "ask_1")
|
||||
.add_edge("ask_1", "ask_2")
|
||||
.add_edge("ask_2", "__end__")
|
||||
.compile(checkpointer=True)
|
||||
)
|
||||
|
||||
graph = (
|
||||
StateGraph(State)
|
||||
.add_node("executor", executor)
|
||||
.add_edge(START, "executor")
|
||||
.compile(checkpointer=sync_checkpointer)
|
||||
)
|
||||
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
# Run through both interrupts
|
||||
graph.invoke({"value": []}, config)
|
||||
graph.invoke(Command(resume="answer_1"), config)
|
||||
|
||||
# Before resuming 2nd interrupt, get state history to find the
|
||||
# subgraph checkpoint that will exist after ask_2 completes
|
||||
graph.invoke(Command(resume="answer_2"), config)
|
||||
|
||||
# Get the final parent state — no pending tasks
|
||||
final_state = graph.get_state(config)
|
||||
assert len(final_state.tasks) == 0
|
||||
|
||||
# Replay from the final parent checkpoint — should be a no-op
|
||||
called.clear()
|
||||
replay_result = graph.invoke(None, final_state.config)
|
||||
assert "__interrupt__" not in replay_result
|
||||
assert "step_a" not in called
|
||||
assert "ask_1" not in called
|
||||
assert "ask_2" not in called
|
||||
# All values should be present
|
||||
assert "step_a_done" in replay_result["value"]
|
||||
assert "ask_1:answer_1" in replay_result["value"]
|
||||
assert "ask_2:answer_2" in replay_result["value"]
|
||||
|
||||
|
||||
def test_3_levels_deep_time_travel_to_first_interrupt(
|
||||
sync_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
"""Time travel to the innermost subgraph checkpoint at the FIRST interrupt.
|
||||
|
||||
Architecture:
|
||||
Parent: START --> outer (subgraph, checkpointer=True) --> END
|
||||
Outer: START --> inner (subgraph, checkpointer=True) --> END
|
||||
Inner: START --> step_a --> ask_1 (interrupt) --> ask_2 (interrupt) --> END
|
||||
"""
|
||||
|
||||
called: list[str] = []
|
||||
|
||||
def step_a(state: State) -> State:
|
||||
called.append("step_a")
|
||||
return {"value": ["step_a_done"]}
|
||||
|
||||
def ask_1(state: State) -> State:
|
||||
called.append("ask_1")
|
||||
answer = interrupt("Question 1?")
|
||||
return {"value": [f"ask_1:{answer}"]}
|
||||
|
||||
def ask_2(state: State) -> State:
|
||||
called.append("ask_2")
|
||||
answer = interrupt("Question 2?")
|
||||
return {"value": [f"ask_2:{answer}"]}
|
||||
|
||||
inner = (
|
||||
StateGraph(State)
|
||||
.add_node("step_a", step_a)
|
||||
.add_node("ask_1", ask_1)
|
||||
.add_node("ask_2", ask_2)
|
||||
.add_edge(START, "step_a")
|
||||
.add_edge("step_a", "ask_1")
|
||||
.add_edge("ask_1", "ask_2")
|
||||
.add_edge("ask_2", "__end__")
|
||||
.compile(checkpointer=True)
|
||||
)
|
||||
|
||||
middle = (
|
||||
StateGraph(State)
|
||||
.add_node("inner", inner)
|
||||
.add_edge(START, "inner")
|
||||
.compile(checkpointer=True)
|
||||
)
|
||||
|
||||
graph = (
|
||||
StateGraph(State)
|
||||
.add_node("outer", middle)
|
||||
.add_edge(START, "outer")
|
||||
.compile(checkpointer=sync_checkpointer)
|
||||
)
|
||||
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
# Run until first interrupt
|
||||
result = graph.invoke({"value": []}, config)
|
||||
assert result["__interrupt__"][0].value == "Question 1?"
|
||||
|
||||
# Capture innermost subgraph state at the first interrupt
|
||||
parent_state = graph.get_state(config, subgraphs=True)
|
||||
mid_state = parent_state.tasks[0].state
|
||||
inner_config = mid_state.tasks[0].state.config
|
||||
|
||||
# Resume through both interrupts to complete
|
||||
graph.invoke(Command(resume="answer_1"), config)
|
||||
graph.invoke(Command(resume="answer_2"), config)
|
||||
|
||||
# --- Scenario 1: Replay from innermost checkpoint at 1st interrupt ---
|
||||
called.clear()
|
||||
replay_result = graph.invoke(None, inner_config)
|
||||
assert "__interrupt__" in replay_result
|
||||
assert replay_result["__interrupt__"][0].value == "Question 1?"
|
||||
assert "step_a" not in called
|
||||
assert "ask_1" in called
|
||||
|
||||
# --- Scenario 2: Fork from innermost checkpoint at 1st interrupt ---
|
||||
called.clear()
|
||||
fork_config = graph.update_state(inner_config, {"value": ["forked"]})
|
||||
fork_result = graph.invoke(None, fork_config)
|
||||
assert "__interrupt__" in fork_result
|
||||
assert fork_result["__interrupt__"][0].value == "Question 1?"
|
||||
assert "step_a" not in called
|
||||
assert "ask_1" in called
|
||||
|
||||
|
||||
def test_3_levels_deep_time_travel_to_second_interrupt(
|
||||
sync_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
"""Time travel to the innermost subgraph checkpoint at the SECOND interrupt.
|
||||
|
||||
Architecture:
|
||||
Parent: START --> outer (subgraph, checkpointer=True) --> END
|
||||
Outer: START --> inner (subgraph, checkpointer=True) --> END
|
||||
Inner: START --> step_a --> ask_1 (interrupt) --> ask_2 (interrupt) --> END
|
||||
"""
|
||||
|
||||
called: list[str] = []
|
||||
|
||||
def step_a(state: State) -> State:
|
||||
called.append("step_a")
|
||||
return {"value": ["step_a_done"]}
|
||||
|
||||
def ask_1(state: State) -> State:
|
||||
called.append("ask_1")
|
||||
answer = interrupt("Question 1?")
|
||||
return {"value": [f"ask_1:{answer}"]}
|
||||
|
||||
def ask_2(state: State) -> State:
|
||||
called.append("ask_2")
|
||||
answer = interrupt("Question 2?")
|
||||
return {"value": [f"ask_2:{answer}"]}
|
||||
|
||||
inner = (
|
||||
StateGraph(State)
|
||||
.add_node("step_a", step_a)
|
||||
.add_node("ask_1", ask_1)
|
||||
.add_node("ask_2", ask_2)
|
||||
.add_edge(START, "step_a")
|
||||
.add_edge("step_a", "ask_1")
|
||||
.add_edge("ask_1", "ask_2")
|
||||
.add_edge("ask_2", "__end__")
|
||||
.compile(checkpointer=True)
|
||||
)
|
||||
|
||||
middle = (
|
||||
StateGraph(State)
|
||||
.add_node("inner", inner)
|
||||
.add_edge(START, "inner")
|
||||
.compile(checkpointer=True)
|
||||
)
|
||||
|
||||
graph = (
|
||||
StateGraph(State)
|
||||
.add_node("outer", middle)
|
||||
.add_edge(START, "outer")
|
||||
.compile(checkpointer=sync_checkpointer)
|
||||
)
|
||||
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
# Run until first interrupt
|
||||
graph.invoke({"value": []}, config)
|
||||
|
||||
# Resume first interrupt
|
||||
result = graph.invoke(Command(resume="answer_1"), config)
|
||||
assert result["__interrupt__"][0].value == "Question 2?"
|
||||
|
||||
# Capture innermost subgraph state at the second interrupt
|
||||
parent_state = graph.get_state(config, subgraphs=True)
|
||||
mid_state = parent_state.tasks[0].state
|
||||
inner_config = mid_state.tasks[0].state.config
|
||||
|
||||
# Resume second interrupt to complete
|
||||
graph.invoke(Command(resume="answer_2"), config)
|
||||
|
||||
# --- Scenario 1: Replay from innermost checkpoint at 2nd interrupt ---
|
||||
called.clear()
|
||||
replay_result = graph.invoke(None, inner_config)
|
||||
assert "__interrupt__" in replay_result
|
||||
assert replay_result["__interrupt__"][0].value == "Question 2?"
|
||||
assert "step_a" not in called
|
||||
assert "ask_1" not in called
|
||||
|
||||
# --- Scenario 2: Fork from innermost checkpoint at 2nd interrupt ---
|
||||
called.clear()
|
||||
fork_config = graph.update_state(inner_config, {"value": ["forked"]})
|
||||
fork_result = graph.invoke(None, fork_config)
|
||||
assert "__interrupt__" in fork_result
|
||||
assert fork_result["__interrupt__"][0].value == "Question 2?"
|
||||
assert "step_a" not in called
|
||||
assert "ask_1" not in called
|
||||
|
||||
|
||||
def test_3_levels_deep_time_travel_to_middle_subgraph(
|
||||
sync_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
"""Time travel to the MIDDLE-level subgraph checkpoint (not innermost).
|
||||
|
||||
Architecture:
|
||||
Parent: START --> outer (subgraph, checkpointer=True) --> END
|
||||
Outer: START --> inner (subgraph, checkpointer=True) --> END
|
||||
Inner: START --> step_a --> ask_1 (interrupt) --> ask_2 (interrupt) --> END
|
||||
|
||||
After completing the full flow, time travel back to the middle subgraph's
|
||||
checkpoint at the second interrupt. The middle subgraph should replay the
|
||||
inner subgraph from the correct point.
|
||||
"""
|
||||
|
||||
called: list[str] = []
|
||||
|
||||
def step_a(state: State) -> State:
|
||||
called.append("step_a")
|
||||
return {"value": ["step_a_done"]}
|
||||
|
||||
def ask_1(state: State) -> State:
|
||||
called.append("ask_1")
|
||||
answer = interrupt("Question 1?")
|
||||
return {"value": [f"ask_1:{answer}"]}
|
||||
|
||||
def ask_2(state: State) -> State:
|
||||
called.append("ask_2")
|
||||
answer = interrupt("Question 2?")
|
||||
return {"value": [f"ask_2:{answer}"]}
|
||||
|
||||
inner = (
|
||||
StateGraph(State)
|
||||
.add_node("step_a", step_a)
|
||||
.add_node("ask_1", ask_1)
|
||||
.add_node("ask_2", ask_2)
|
||||
.add_edge(START, "step_a")
|
||||
.add_edge("step_a", "ask_1")
|
||||
.add_edge("ask_1", "ask_2")
|
||||
.add_edge("ask_2", "__end__")
|
||||
.compile(checkpointer=True)
|
||||
)
|
||||
|
||||
middle = (
|
||||
StateGraph(State)
|
||||
.add_node("inner", inner)
|
||||
.add_edge(START, "inner")
|
||||
.compile(checkpointer=True)
|
||||
)
|
||||
|
||||
graph = (
|
||||
StateGraph(State)
|
||||
.add_node("outer", middle)
|
||||
.add_edge(START, "outer")
|
||||
.compile(checkpointer=sync_checkpointer)
|
||||
)
|
||||
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
# Run until first interrupt
|
||||
graph.invoke({"value": []}, config)
|
||||
|
||||
# Resume first, capture middle config at second interrupt
|
||||
graph.invoke(Command(resume="answer_1"), config)
|
||||
parent_state = graph.get_state(config, subgraphs=True)
|
||||
mid_config = parent_state.tasks[0].state.config
|
||||
|
||||
# Resume second to complete
|
||||
graph.invoke(Command(resume="answer_2"), config)
|
||||
|
||||
# --- Scenario 1: Replay from middle-level subgraph checkpoint ---
|
||||
# The middle subgraph's checkpoint knows about the inner subgraph's state
|
||||
# via checkpoint_map, so the inner replays from the correct point.
|
||||
called.clear()
|
||||
replay_result = graph.invoke(None, mid_config)
|
||||
assert "__interrupt__" in replay_result
|
||||
|
||||
# --- Scenario 2: Fork from middle-level subgraph checkpoint ---
|
||||
called.clear()
|
||||
fork_config = graph.update_state(mid_config, {"value": ["forked"]})
|
||||
fork_result = graph.invoke(None, fork_config)
|
||||
assert "__interrupt__" in fork_result
|
||||
|
||||
|
||||
def test_3_levels_deep_middle_has_interrupts(
|
||||
sync_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
"""Time travel when the MIDDLE subgraph itself has interrupts.
|
||||
|
||||
Architecture:
|
||||
Parent: START --> outer (subgraph, checkpointer=True) --> END
|
||||
Outer: START --> pre (interrupt) --> inner (subgraph, checkpointer=True) --> END
|
||||
Inner: START --> step_a --> ask_1 (interrupt) --> END
|
||||
|
||||
Flow: run through both interrupts (pre then ask_1), then time travel back to
|
||||
the middle subgraph checkpoint at each interrupt point.
|
||||
"""
|
||||
|
||||
called: list[str] = []
|
||||
|
||||
def pre(state: State) -> State:
|
||||
called.append("pre")
|
||||
answer = interrupt("Pre-question?")
|
||||
return {"value": [f"pre:{answer}"]}
|
||||
|
||||
def step_a(state: State) -> State:
|
||||
called.append("step_a")
|
||||
return {"value": ["step_a_done"]}
|
||||
|
||||
def ask_1(state: State) -> State:
|
||||
called.append("ask_1")
|
||||
answer = interrupt("Question 1?")
|
||||
return {"value": [f"ask_1:{answer}"]}
|
||||
|
||||
inner = (
|
||||
StateGraph(State)
|
||||
.add_node("step_a", step_a)
|
||||
.add_node("ask_1", ask_1)
|
||||
.add_edge(START, "step_a")
|
||||
.add_edge("step_a", "ask_1")
|
||||
.add_edge("ask_1", "__end__")
|
||||
.compile(checkpointer=True)
|
||||
)
|
||||
|
||||
middle = (
|
||||
StateGraph(State)
|
||||
.add_node("pre", pre)
|
||||
.add_node("inner", inner)
|
||||
.add_edge(START, "pre")
|
||||
.add_edge("pre", "inner")
|
||||
.add_edge("inner", "__end__")
|
||||
.compile(checkpointer=True)
|
||||
)
|
||||
|
||||
graph = (
|
||||
StateGraph(State)
|
||||
.add_node("outer", middle)
|
||||
.add_edge(START, "outer")
|
||||
.compile(checkpointer=sync_checkpointer)
|
||||
)
|
||||
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
# Run until first interrupt (pre in middle subgraph)
|
||||
result = graph.invoke({"value": []}, config)
|
||||
assert result["__interrupt__"][0].value == "Pre-question?"
|
||||
|
||||
# Capture middle subgraph config at the pre interrupt
|
||||
parent_state = graph.get_state(config, subgraphs=True)
|
||||
mid_config_at_pre = parent_state.tasks[0].state.config
|
||||
|
||||
# Resume pre, hits ask_1 in inner subgraph
|
||||
result = graph.invoke(Command(resume="pre_answer"), config)
|
||||
assert result["__interrupt__"][0].value == "Question 1?"
|
||||
|
||||
# Capture middle subgraph config at the ask_1 interrupt
|
||||
parent_state = graph.get_state(config, subgraphs=True)
|
||||
mid_config_at_ask1 = parent_state.tasks[0].state.config
|
||||
|
||||
# Resume ask_1 to complete
|
||||
result = graph.invoke(Command(resume="answer_1"), config)
|
||||
assert "__interrupt__" not in result
|
||||
|
||||
# --- Time travel to middle checkpoint at pre interrupt ---
|
||||
called.clear()
|
||||
replay_result = graph.invoke(None, mid_config_at_pre)
|
||||
assert "__interrupt__" in replay_result
|
||||
assert replay_result["__interrupt__"][0].value == "Pre-question?"
|
||||
# pre should re-fire (interrupt replays), but nothing else should run
|
||||
assert "pre" in called
|
||||
assert "step_a" not in called
|
||||
assert "ask_1" not in called
|
||||
|
||||
# Fork from middle checkpoint at pre interrupt
|
||||
called.clear()
|
||||
fork_config = graph.update_state(mid_config_at_pre, {"value": ["forked"]})
|
||||
fork_result = graph.invoke(None, fork_config)
|
||||
assert "__interrupt__" in fork_result
|
||||
assert fork_result["__interrupt__"][0].value == "Pre-question?"
|
||||
assert "pre" in called
|
||||
assert "step_a" not in called
|
||||
|
||||
# --- Time travel to middle checkpoint at ask_1 interrupt ---
|
||||
called.clear()
|
||||
replay_result = graph.invoke(None, mid_config_at_ask1)
|
||||
assert "__interrupt__" in replay_result
|
||||
assert replay_result["__interrupt__"][0].value == "Question 1?"
|
||||
# pre should NOT re-run (it completed before this checkpoint)
|
||||
assert "pre" not in called
|
||||
# ask_1 re-fires
|
||||
assert "ask_1" in called
|
||||
|
||||
# Fork from middle checkpoint at ask_1 interrupt
|
||||
called.clear()
|
||||
fork_config = graph.update_state(mid_config_at_ask1, {"value": ["forked"]})
|
||||
fork_result = graph.invoke(None, fork_config)
|
||||
assert "__interrupt__" in fork_result
|
||||
assert fork_result["__interrupt__"][0].value == "Question 1?"
|
||||
assert "pre" not in called
|
||||
assert "ask_1" in called
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Section 6: __copy__ / update_state(None)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -1051,552 +1051,6 @@ async def test_subgraph_interrupt_full_flow_no_sub_checkpointer(
|
||||
assert "post" in final_result["value"]
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
async def test_subgraph_time_travel_to_first_interrupt_async(
|
||||
async_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
"""Time travel to a subgraph checkpoint at the FIRST interrupt (async)."""
|
||||
|
||||
called: list[str] = []
|
||||
|
||||
async def step_a(state: State) -> State:
|
||||
called.append("step_a")
|
||||
return {"value": ["step_a_done"]}
|
||||
|
||||
async def ask_1(state: State) -> State:
|
||||
called.append("ask_1")
|
||||
answer = interrupt("Question 1?")
|
||||
return {"value": [f"ask_1:{answer}"]}
|
||||
|
||||
async def ask_2(state: State) -> State:
|
||||
called.append("ask_2")
|
||||
answer = interrupt("Question 2?")
|
||||
return {"value": [f"ask_2:{answer}"]}
|
||||
|
||||
executor = (
|
||||
StateGraph(State)
|
||||
.add_node("step_a", step_a)
|
||||
.add_node("ask_1", ask_1)
|
||||
.add_node("ask_2", ask_2)
|
||||
.add_edge(START, "step_a")
|
||||
.add_edge("step_a", "ask_1")
|
||||
.add_edge("ask_1", "ask_2")
|
||||
.add_edge("ask_2", "__end__")
|
||||
.compile(checkpointer=True)
|
||||
)
|
||||
|
||||
graph = (
|
||||
StateGraph(State)
|
||||
.add_node("executor", executor)
|
||||
.add_edge(START, "executor")
|
||||
.compile(checkpointer=async_checkpointer)
|
||||
)
|
||||
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
# Run until first interrupt (ask_1)
|
||||
result = await graph.ainvoke({"value": []}, config)
|
||||
assert result["__interrupt__"][0].value == "Question 1?"
|
||||
|
||||
# Capture subgraph state at the first interrupt
|
||||
parent_state = await graph.aget_state(config, subgraphs=True)
|
||||
sub_config_at_first = parent_state.tasks[0].state.config
|
||||
|
||||
# Resume through both interrupts to complete
|
||||
await graph.ainvoke(Command(resume="answer_1"), config)
|
||||
await graph.ainvoke(Command(resume="answer_2"), config)
|
||||
|
||||
# --- Scenario 1: Replay from subgraph checkpoint at 1st interrupt ---
|
||||
called.clear()
|
||||
replay_result = await graph.ainvoke(None, sub_config_at_first)
|
||||
assert "__interrupt__" in replay_result
|
||||
assert replay_result["__interrupt__"][0].value == "Question 1?"
|
||||
assert "step_a" not in called
|
||||
assert "ask_1" in called
|
||||
|
||||
# --- Scenario 2: Fork from subgraph checkpoint at 1st interrupt ---
|
||||
called.clear()
|
||||
fork_config = await graph.aupdate_state(sub_config_at_first, {"value": ["forked"]})
|
||||
fork_result = await graph.ainvoke(None, fork_config)
|
||||
assert "__interrupt__" in fork_result
|
||||
assert fork_result["__interrupt__"][0].value == "Question 1?"
|
||||
assert "step_a" not in called
|
||||
assert "ask_1" in called
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
async def test_subgraph_time_travel_to_second_interrupt_async(
|
||||
async_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
"""Time travel to a subgraph checkpoint at the SECOND interrupt (async)."""
|
||||
|
||||
called: list[str] = []
|
||||
|
||||
async def step_a(state: State) -> State:
|
||||
called.append("step_a")
|
||||
return {"value": ["step_a_done"]}
|
||||
|
||||
async def ask_1(state: State) -> State:
|
||||
called.append("ask_1")
|
||||
answer = interrupt("Question 1?")
|
||||
return {"value": [f"ask_1:{answer}"]}
|
||||
|
||||
async def ask_2(state: State) -> State:
|
||||
called.append("ask_2")
|
||||
answer = interrupt("Question 2?")
|
||||
return {"value": [f"ask_2:{answer}"]}
|
||||
|
||||
executor = (
|
||||
StateGraph(State)
|
||||
.add_node("step_a", step_a)
|
||||
.add_node("ask_1", ask_1)
|
||||
.add_node("ask_2", ask_2)
|
||||
.add_edge(START, "step_a")
|
||||
.add_edge("step_a", "ask_1")
|
||||
.add_edge("ask_1", "ask_2")
|
||||
.add_edge("ask_2", "__end__")
|
||||
.compile(checkpointer=True)
|
||||
)
|
||||
|
||||
graph = (
|
||||
StateGraph(State)
|
||||
.add_node("executor", executor)
|
||||
.add_edge(START, "executor")
|
||||
.compile(checkpointer=async_checkpointer)
|
||||
)
|
||||
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
# Run until first interrupt
|
||||
await graph.ainvoke({"value": []}, config)
|
||||
|
||||
# Resume first interrupt
|
||||
result = await graph.ainvoke(Command(resume="answer_1"), config)
|
||||
assert result["__interrupt__"][0].value == "Question 2?"
|
||||
|
||||
# Capture subgraph state at the second interrupt
|
||||
parent_state = await graph.aget_state(config, subgraphs=True)
|
||||
sub_config = parent_state.tasks[0].state.config
|
||||
|
||||
# Resume second interrupt to complete
|
||||
await graph.ainvoke(Command(resume="answer_2"), config)
|
||||
|
||||
# --- Scenario 1: Replay from subgraph checkpoint at 2nd interrupt ---
|
||||
called.clear()
|
||||
replay_result = await graph.ainvoke(None, sub_config)
|
||||
assert "__interrupt__" in replay_result
|
||||
assert replay_result["__interrupt__"][0].value == "Question 2?"
|
||||
assert "step_a" not in called
|
||||
assert "ask_1" not in called
|
||||
|
||||
# --- Scenario 2: Fork from subgraph checkpoint at 2nd interrupt ---
|
||||
called.clear()
|
||||
fork_config = await graph.aupdate_state(sub_config, {"value": ["forked"]})
|
||||
fork_result = await graph.ainvoke(None, fork_config)
|
||||
assert "__interrupt__" in fork_result
|
||||
assert fork_result["__interrupt__"][0].value == "Question 2?"
|
||||
assert "step_a" not in called
|
||||
assert "ask_1" not in called
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
async def test_subgraph_time_travel_after_completion_async(
|
||||
async_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
"""Time travel to a subgraph checkpoint AFTER both interrupts resolved (async)."""
|
||||
|
||||
called: list[str] = []
|
||||
|
||||
async def step_a(state: State) -> State:
|
||||
called.append("step_a")
|
||||
return {"value": ["step_a_done"]}
|
||||
|
||||
async def ask_1(state: State) -> State:
|
||||
called.append("ask_1")
|
||||
answer = interrupt("Question 1?")
|
||||
return {"value": [f"ask_1:{answer}"]}
|
||||
|
||||
async def ask_2(state: State) -> State:
|
||||
called.append("ask_2")
|
||||
answer = interrupt("Question 2?")
|
||||
return {"value": [f"ask_2:{answer}"]}
|
||||
|
||||
executor = (
|
||||
StateGraph(State)
|
||||
.add_node("step_a", step_a)
|
||||
.add_node("ask_1", ask_1)
|
||||
.add_node("ask_2", ask_2)
|
||||
.add_edge(START, "step_a")
|
||||
.add_edge("step_a", "ask_1")
|
||||
.add_edge("ask_1", "ask_2")
|
||||
.add_edge("ask_2", "__end__")
|
||||
.compile(checkpointer=True)
|
||||
)
|
||||
|
||||
graph = (
|
||||
StateGraph(State)
|
||||
.add_node("executor", executor)
|
||||
.add_edge(START, "executor")
|
||||
.compile(checkpointer=async_checkpointer)
|
||||
)
|
||||
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
await graph.ainvoke({"value": []}, config)
|
||||
await graph.ainvoke(Command(resume="answer_1"), config)
|
||||
await graph.ainvoke(Command(resume="answer_2"), config)
|
||||
|
||||
final_state = await graph.aget_state(config)
|
||||
assert len(final_state.tasks) == 0
|
||||
|
||||
# Replay from the final parent checkpoint — should be a no-op
|
||||
called.clear()
|
||||
replay_result = await graph.ainvoke(None, final_state.config)
|
||||
assert "__interrupt__" not in replay_result
|
||||
assert "step_a" not in called
|
||||
assert "ask_1" not in called
|
||||
assert "ask_2" not in called
|
||||
assert "step_a_done" in replay_result["value"]
|
||||
assert "ask_1:answer_1" in replay_result["value"]
|
||||
assert "ask_2:answer_2" in replay_result["value"]
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
async def test_3_levels_deep_time_travel_to_first_interrupt_async(
|
||||
async_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
"""Time travel to innermost subgraph checkpoint at FIRST interrupt (async, 3 levels)."""
|
||||
|
||||
called: list[str] = []
|
||||
|
||||
async def step_a(state: State) -> State:
|
||||
called.append("step_a")
|
||||
return {"value": ["step_a_done"]}
|
||||
|
||||
async def ask_1(state: State) -> State:
|
||||
called.append("ask_1")
|
||||
answer = interrupt("Question 1?")
|
||||
return {"value": [f"ask_1:{answer}"]}
|
||||
|
||||
async def ask_2(state: State) -> State:
|
||||
called.append("ask_2")
|
||||
answer = interrupt("Question 2?")
|
||||
return {"value": [f"ask_2:{answer}"]}
|
||||
|
||||
inner = (
|
||||
StateGraph(State)
|
||||
.add_node("step_a", step_a)
|
||||
.add_node("ask_1", ask_1)
|
||||
.add_node("ask_2", ask_2)
|
||||
.add_edge(START, "step_a")
|
||||
.add_edge("step_a", "ask_1")
|
||||
.add_edge("ask_1", "ask_2")
|
||||
.add_edge("ask_2", "__end__")
|
||||
.compile(checkpointer=True)
|
||||
)
|
||||
|
||||
middle = (
|
||||
StateGraph(State)
|
||||
.add_node("inner", inner)
|
||||
.add_edge(START, "inner")
|
||||
.compile(checkpointer=True)
|
||||
)
|
||||
|
||||
graph = (
|
||||
StateGraph(State)
|
||||
.add_node("outer", middle)
|
||||
.add_edge(START, "outer")
|
||||
.compile(checkpointer=async_checkpointer)
|
||||
)
|
||||
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
result = await graph.ainvoke({"value": []}, config)
|
||||
assert result["__interrupt__"][0].value == "Question 1?"
|
||||
|
||||
parent_state = await graph.aget_state(config, subgraphs=True)
|
||||
mid_state = parent_state.tasks[0].state
|
||||
inner_config = mid_state.tasks[0].state.config
|
||||
|
||||
await graph.ainvoke(Command(resume="answer_1"), config)
|
||||
await graph.ainvoke(Command(resume="answer_2"), config)
|
||||
|
||||
# --- Scenario 1: Replay from innermost checkpoint at 1st interrupt ---
|
||||
called.clear()
|
||||
replay_result = await graph.ainvoke(None, inner_config)
|
||||
assert "__interrupt__" in replay_result
|
||||
assert replay_result["__interrupt__"][0].value == "Question 1?"
|
||||
assert "step_a" not in called
|
||||
assert "ask_1" in called
|
||||
|
||||
# --- Scenario 2: Fork from innermost checkpoint at 1st interrupt ---
|
||||
called.clear()
|
||||
fork_config = await graph.aupdate_state(inner_config, {"value": ["forked"]})
|
||||
fork_result = await graph.ainvoke(None, fork_config)
|
||||
assert "__interrupt__" in fork_result
|
||||
assert fork_result["__interrupt__"][0].value == "Question 1?"
|
||||
assert "step_a" not in called
|
||||
assert "ask_1" in called
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
async def test_3_levels_deep_time_travel_to_second_interrupt_async(
|
||||
async_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
"""Time travel to innermost subgraph checkpoint at SECOND interrupt (async, 3 levels)."""
|
||||
|
||||
called: list[str] = []
|
||||
|
||||
async def step_a(state: State) -> State:
|
||||
called.append("step_a")
|
||||
return {"value": ["step_a_done"]}
|
||||
|
||||
async def ask_1(state: State) -> State:
|
||||
called.append("ask_1")
|
||||
answer = interrupt("Question 1?")
|
||||
return {"value": [f"ask_1:{answer}"]}
|
||||
|
||||
async def ask_2(state: State) -> State:
|
||||
called.append("ask_2")
|
||||
answer = interrupt("Question 2?")
|
||||
return {"value": [f"ask_2:{answer}"]}
|
||||
|
||||
inner = (
|
||||
StateGraph(State)
|
||||
.add_node("step_a", step_a)
|
||||
.add_node("ask_1", ask_1)
|
||||
.add_node("ask_2", ask_2)
|
||||
.add_edge(START, "step_a")
|
||||
.add_edge("step_a", "ask_1")
|
||||
.add_edge("ask_1", "ask_2")
|
||||
.add_edge("ask_2", "__end__")
|
||||
.compile(checkpointer=True)
|
||||
)
|
||||
|
||||
middle = (
|
||||
StateGraph(State)
|
||||
.add_node("inner", inner)
|
||||
.add_edge(START, "inner")
|
||||
.compile(checkpointer=True)
|
||||
)
|
||||
|
||||
graph = (
|
||||
StateGraph(State)
|
||||
.add_node("outer", middle)
|
||||
.add_edge(START, "outer")
|
||||
.compile(checkpointer=async_checkpointer)
|
||||
)
|
||||
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
await graph.ainvoke({"value": []}, config)
|
||||
|
||||
result = await graph.ainvoke(Command(resume="answer_1"), config)
|
||||
assert result["__interrupt__"][0].value == "Question 2?"
|
||||
|
||||
parent_state = await graph.aget_state(config, subgraphs=True)
|
||||
mid_state = parent_state.tasks[0].state
|
||||
inner_config = mid_state.tasks[0].state.config
|
||||
|
||||
await graph.ainvoke(Command(resume="answer_2"), config)
|
||||
|
||||
# --- Scenario 1: Replay ---
|
||||
called.clear()
|
||||
replay_result = await graph.ainvoke(None, inner_config)
|
||||
assert "__interrupt__" in replay_result
|
||||
assert replay_result["__interrupt__"][0].value == "Question 2?"
|
||||
assert "step_a" not in called
|
||||
assert "ask_1" not in called
|
||||
|
||||
# --- Scenario 2: Fork ---
|
||||
called.clear()
|
||||
fork_config = await graph.aupdate_state(inner_config, {"value": ["forked"]})
|
||||
fork_result = await graph.ainvoke(None, fork_config)
|
||||
assert "__interrupt__" in fork_result
|
||||
assert fork_result["__interrupt__"][0].value == "Question 2?"
|
||||
assert "step_a" not in called
|
||||
assert "ask_1" not in called
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
async def test_3_levels_deep_time_travel_to_middle_subgraph_async(
|
||||
async_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
"""Time travel to the MIDDLE-level subgraph checkpoint (async, 3 levels)."""
|
||||
|
||||
called: list[str] = []
|
||||
|
||||
async def step_a(state: State) -> State:
|
||||
called.append("step_a")
|
||||
return {"value": ["step_a_done"]}
|
||||
|
||||
async def ask_1(state: State) -> State:
|
||||
called.append("ask_1")
|
||||
answer = interrupt("Question 1?")
|
||||
return {"value": [f"ask_1:{answer}"]}
|
||||
|
||||
async def ask_2(state: State) -> State:
|
||||
called.append("ask_2")
|
||||
answer = interrupt("Question 2?")
|
||||
return {"value": [f"ask_2:{answer}"]}
|
||||
|
||||
inner = (
|
||||
StateGraph(State)
|
||||
.add_node("step_a", step_a)
|
||||
.add_node("ask_1", ask_1)
|
||||
.add_node("ask_2", ask_2)
|
||||
.add_edge(START, "step_a")
|
||||
.add_edge("step_a", "ask_1")
|
||||
.add_edge("ask_1", "ask_2")
|
||||
.add_edge("ask_2", "__end__")
|
||||
.compile(checkpointer=True)
|
||||
)
|
||||
|
||||
middle = (
|
||||
StateGraph(State)
|
||||
.add_node("inner", inner)
|
||||
.add_edge(START, "inner")
|
||||
.compile(checkpointer=True)
|
||||
)
|
||||
|
||||
graph = (
|
||||
StateGraph(State)
|
||||
.add_node("outer", middle)
|
||||
.add_edge(START, "outer")
|
||||
.compile(checkpointer=async_checkpointer)
|
||||
)
|
||||
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
await graph.ainvoke({"value": []}, config)
|
||||
await graph.ainvoke(Command(resume="answer_1"), config)
|
||||
|
||||
parent_state = await graph.aget_state(config, subgraphs=True)
|
||||
mid_config = parent_state.tasks[0].state.config
|
||||
|
||||
await graph.ainvoke(Command(resume="answer_2"), config)
|
||||
|
||||
# --- Scenario 1: Replay from middle-level subgraph checkpoint ---
|
||||
# The middle subgraph's checkpoint knows about the inner subgraph's state
|
||||
# via checkpoint_map, so the inner replays from the correct point.
|
||||
called.clear()
|
||||
replay_result = await graph.ainvoke(None, mid_config)
|
||||
assert "__interrupt__" in replay_result
|
||||
|
||||
# --- Scenario 2: Fork from middle-level subgraph checkpoint ---
|
||||
called.clear()
|
||||
fork_config = await graph.aupdate_state(mid_config, {"value": ["forked"]})
|
||||
fork_result = await graph.ainvoke(None, fork_config)
|
||||
assert "__interrupt__" in fork_result
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
async def test_3_levels_deep_middle_has_interrupts_async(
|
||||
async_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
"""Time travel when the MIDDLE subgraph itself has interrupts (async)."""
|
||||
|
||||
called: list[str] = []
|
||||
|
||||
async def pre(state: State) -> State:
|
||||
called.append("pre")
|
||||
answer = interrupt("Pre-question?")
|
||||
return {"value": [f"pre:{answer}"]}
|
||||
|
||||
async def step_a(state: State) -> State:
|
||||
called.append("step_a")
|
||||
return {"value": ["step_a_done"]}
|
||||
|
||||
async def ask_1(state: State) -> State:
|
||||
called.append("ask_1")
|
||||
answer = interrupt("Question 1?")
|
||||
return {"value": [f"ask_1:{answer}"]}
|
||||
|
||||
inner = (
|
||||
StateGraph(State)
|
||||
.add_node("step_a", step_a)
|
||||
.add_node("ask_1", ask_1)
|
||||
.add_edge(START, "step_a")
|
||||
.add_edge("step_a", "ask_1")
|
||||
.add_edge("ask_1", "__end__")
|
||||
.compile(checkpointer=True)
|
||||
)
|
||||
|
||||
middle = (
|
||||
StateGraph(State)
|
||||
.add_node("pre", pre)
|
||||
.add_node("inner", inner)
|
||||
.add_edge(START, "pre")
|
||||
.add_edge("pre", "inner")
|
||||
.add_edge("inner", "__end__")
|
||||
.compile(checkpointer=True)
|
||||
)
|
||||
|
||||
graph = (
|
||||
StateGraph(State)
|
||||
.add_node("outer", middle)
|
||||
.add_edge(START, "outer")
|
||||
.compile(checkpointer=async_checkpointer)
|
||||
)
|
||||
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
# Run until first interrupt (pre in middle subgraph)
|
||||
result = await graph.ainvoke({"value": []}, config)
|
||||
assert result["__interrupt__"][0].value == "Pre-question?"
|
||||
|
||||
# Capture middle subgraph config at the pre interrupt
|
||||
parent_state = await graph.aget_state(config, subgraphs=True)
|
||||
mid_config_at_pre = parent_state.tasks[0].state.config
|
||||
|
||||
# Resume pre, hits ask_1 in inner subgraph
|
||||
result = await graph.ainvoke(Command(resume="pre_answer"), config)
|
||||
assert result["__interrupt__"][0].value == "Question 1?"
|
||||
|
||||
# Capture middle subgraph config at the ask_1 interrupt
|
||||
parent_state = await graph.aget_state(config, subgraphs=True)
|
||||
mid_config_at_ask1 = parent_state.tasks[0].state.config
|
||||
|
||||
# Resume ask_1 to complete
|
||||
result = await graph.ainvoke(Command(resume="answer_1"), config)
|
||||
assert "__interrupt__" not in result
|
||||
|
||||
# --- Time travel to middle checkpoint at pre interrupt ---
|
||||
called.clear()
|
||||
replay_result = await graph.ainvoke(None, mid_config_at_pre)
|
||||
assert "__interrupt__" in replay_result
|
||||
assert replay_result["__interrupt__"][0].value == "Pre-question?"
|
||||
assert "pre" in called
|
||||
assert "step_a" not in called
|
||||
assert "ask_1" not in called
|
||||
|
||||
# Fork from middle checkpoint at pre interrupt
|
||||
called.clear()
|
||||
fork_config = await graph.aupdate_state(mid_config_at_pre, {"value": ["forked"]})
|
||||
fork_result = await graph.ainvoke(None, fork_config)
|
||||
assert "__interrupt__" in fork_result
|
||||
assert fork_result["__interrupt__"][0].value == "Pre-question?"
|
||||
assert "pre" in called
|
||||
assert "step_a" not in called
|
||||
|
||||
# --- Time travel to middle checkpoint at ask_1 interrupt ---
|
||||
called.clear()
|
||||
replay_result = await graph.ainvoke(None, mid_config_at_ask1)
|
||||
assert "__interrupt__" in replay_result
|
||||
assert replay_result["__interrupt__"][0].value == "Question 1?"
|
||||
assert "pre" not in called
|
||||
assert "ask_1" in called
|
||||
|
||||
# Fork from middle checkpoint at ask_1 interrupt
|
||||
called.clear()
|
||||
fork_config = await graph.aupdate_state(mid_config_at_ask1, {"value": ["forked"]})
|
||||
fork_result = await graph.ainvoke(None, fork_config)
|
||||
assert "__interrupt__" in fork_result
|
||||
assert fork_result["__interrupt__"][0].value == "Question 1?"
|
||||
assert "pre" not in called
|
||||
assert "ask_1" in called
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Section 6: __copy__ / update_state(None)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -8,8 +8,7 @@ from pydantic import BaseModel
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.graph import StateGraph
|
||||
from langgraph.pregel.remote import RemoteGraph
|
||||
from langgraph.types import Command, GraphOutput, StreamPart
|
||||
from langgraph.types import Command
|
||||
|
||||
|
||||
def test_typed_dict_state() -> None:
|
||||
@@ -160,75 +159,3 @@ def test_add_node_with_explicit_input_schema() -> None:
|
||||
# because it violates the principles of contravariance
|
||||
workflow.add_node("a_narrow", a, input_schema=ANarrow) # type: ignore[arg-type]
|
||||
workflow.add_node("b_narrow", b, input_schema=BNarrow) # type: ignore[arg-type]
|
||||
|
||||
|
||||
@pytest.mark.skip("Purely for type checking")
|
||||
def test_remote_graph_generics_typed_dict() -> None:
|
||||
"""RemoteGraph parameterized with TypedDict should propagate types."""
|
||||
|
||||
class MyState(TypedDict):
|
||||
messages: list[str]
|
||||
|
||||
rg: RemoteGraph[MyState, None, MyState, MyState] = RemoteGraph(
|
||||
"test", url="http://localhost:8123"
|
||||
)
|
||||
|
||||
# v2 invoke should return GraphOutput[MyState]
|
||||
result: GraphOutput[MyState] = rg.invoke({"messages": ["hi"]}, version="v2")
|
||||
_val: MyState = result.value
|
||||
|
||||
# v1 invoke should return dict[str, Any] | Any
|
||||
_v1_result: dict[str, Any] | Any = rg.invoke({"messages": ["hi"]})
|
||||
|
||||
# v2 stream should yield StreamPart[MyState, MyState]
|
||||
for part in rg.stream({"messages": ["hi"]}, version="v2"):
|
||||
_part: StreamPart[MyState, MyState] = part
|
||||
|
||||
# input should accept the state type
|
||||
rg.invoke({"messages": ["hi"]}, version="v2")
|
||||
|
||||
# input should also accept Command
|
||||
rg.invoke(Command(), version="v2")
|
||||
|
||||
# input should also accept None
|
||||
rg.invoke(None, version="v2")
|
||||
|
||||
|
||||
@pytest.mark.skip("Purely for type checking")
|
||||
def test_remote_graph_generics_pydantic() -> None:
|
||||
"""RemoteGraph parameterized with Pydantic model should propagate types."""
|
||||
|
||||
class PydanticState(BaseModel):
|
||||
messages: list[str]
|
||||
|
||||
rg: RemoteGraph[PydanticState, None, PydanticState, PydanticState] = RemoteGraph(
|
||||
"test", url="http://localhost:8123"
|
||||
)
|
||||
|
||||
result: GraphOutput[PydanticState] = rg.invoke(
|
||||
PydanticState(messages=["hi"]), version="v2"
|
||||
)
|
||||
_val: PydanticState = result.value
|
||||
|
||||
|
||||
@pytest.mark.skip("Purely for type checking")
|
||||
def test_remote_graph_separate_input_output() -> None:
|
||||
"""RemoteGraph with different input/output schemas."""
|
||||
|
||||
class InputState(TypedDict):
|
||||
query: str
|
||||
|
||||
class OutputState(TypedDict):
|
||||
answer: str
|
||||
|
||||
class FullState(InputState, OutputState): ...
|
||||
|
||||
rg: RemoteGraph[FullState, None, InputState, OutputState] = RemoteGraph(
|
||||
"test", url="http://localhost:8123"
|
||||
)
|
||||
|
||||
result: GraphOutput[OutputState] = rg.invoke({"query": "hi"}, version="v2")
|
||||
_val: OutputState = result.value
|
||||
|
||||
# wrong input type should fail type checking
|
||||
rg.invoke({"answer": "wrong"}, version="v2") # type: ignore[call-overload]
|
||||
|
||||
Generated
+14
-12
@@ -1367,7 +1367,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph"
|
||||
version = "1.1.1"
|
||||
version = "1.1.0"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
@@ -3615,19 +3615,21 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "tornado"
|
||||
version = "6.5.5"
|
||||
version = "6.5.4"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/f8/f1/3173dfa4a18db4a9b03e5d55325559dab51ee653763bb8745a75af491286/tornado-6.5.5.tar.gz", hash = "sha256:192b8f3ea91bd7f1f50c06955416ed76c6b72f96779b962f07f911b91e8d30e9", size = 516006, upload-time = "2026-03-10T21:31:02.067Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/37/1d/0a336abf618272d53f62ebe274f712e213f5a03c0b2339575430b8362ef2/tornado-6.5.4.tar.gz", hash = "sha256:a22fa9047405d03260b483980635f0b041989d8bcc9a313f8fe18b411d84b1d7", size = 513632, upload-time = "2025-12-15T19:21:03.836Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/59/8c/77f5097695f4dd8255ecbd08b2a1ed8ba8b953d337804dd7080f199e12bf/tornado-6.5.5-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:487dc9cc380e29f58c7ab88f9e27cdeef04b2140862e5076a66fb6bb68bb1bfa", size = 445983, upload-time = "2026-03-10T21:30:44.28Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/ab/5e/7625b76cd10f98f1516c36ce0346de62061156352353ef2da44e5c21523c/tornado-6.5.5-cp39-abi3-macosx_10_9_x86_64.whl", hash = "sha256:65a7f1d46d4bb41df1ac99f5fcb685fb25c7e61613742d5108b010975a9a6521", size = 444246, upload-time = "2026-03-10T21:30:46.571Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/b2/04/7b5705d5b3c0fab088f434f9c83edac1573830ca49ccf29fb83bf7178eec/tornado-6.5.5-cp39-abi3-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:e74c92e8e65086b338fd56333fb9a68b9f6f2fe7ad532645a290a464bcf46be5", size = 447229, upload-time = "2026-03-10T21:30:48.273Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/34/01/74e034a30ef59afb4097ef8659515e96a39d910b712a89af76f5e4e1f93c/tornado-6.5.5-cp39-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:435319e9e340276428bbdb4e7fa732c2d399386d1de5686cb331ec8eee754f07", size = 448192, upload-time = "2026-03-10T21:30:51.22Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/be/00/fe9e02c5a96429fce1a1d15a517f5d8444f9c412e0bb9eadfbe3b0fc55bf/tornado-6.5.5-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:3f54aa540bdbfee7b9eb268ead60e7d199de5021facd276819c193c0fb28ea4e", size = 448039, upload-time = "2026-03-10T21:30:53.52Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/82/9e/656ee4cec0398b1d18d0f1eb6372c41c6b889722641d84948351ae19556d/tornado-6.5.5-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:36abed1754faeb80fbd6e64db2758091e1320f6bba74a4cf8c09cd18ccce8aca", size = 447445, upload-time = "2026-03-10T21:30:55.541Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/5a/76/4921c00511f88af86a33de770d64141170f1cfd9c00311aea689949e274e/tornado-6.5.5-cp39-abi3-win32.whl", hash = "sha256:dd3eafaaeec1c7f2f8fdcd5f964e8907ad788fe8a5a32c4426fbbdda621223b7", size = 448582, upload-time = "2026-03-10T21:30:57.142Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/2c/23/f6c6112a04d28eed765e374435fb1a9198f73e1ec4b4024184f21faeb1ad/tornado-6.5.5-cp39-abi3-win_amd64.whl", hash = "sha256:6443a794ba961a9f619b1ae926a2e900ac20c34483eea67be4ed8f1e58d3ef7b", size = 448990, upload-time = "2026-03-10T21:30:58.857Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/b7/c8/876602cbc96469911f0939f703453c1157b0c826ecb05bdd32e023397d4e/tornado-6.5.5-cp39-abi3-win_arm64.whl", hash = "sha256:2c9a876e094109333f888539ddb2de4361743e5d21eece20688e3e351e4990a6", size = 448016, upload-time = "2026-03-10T21:31:00.43Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/ab/a9/e94a9d5224107d7ce3cc1fab8d5dc97f5ea351ccc6322ee4fb661da94e35/tornado-6.5.4-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:d6241c1a16b1c9e4cc28148b1cda97dd1c6cb4fb7068ac1bedc610768dff0ba9", size = 443909, upload-time = "2025-12-15T19:20:48.382Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/db/7e/f7b8d8c4453f305a51f80dbb49014257bb7d28ccb4bbb8dd328ea995ecad/tornado-6.5.4-cp39-abi3-macosx_10_9_x86_64.whl", hash = "sha256:2d50f63dda1d2cac3ae1fa23d254e16b5e38153758470e9956cbc3d813d40843", size = 442163, upload-time = "2025-12-15T19:20:49.791Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/ba/b5/206f82d51e1bfa940ba366a8d2f83904b15942c45a78dd978b599870ab44/tornado-6.5.4-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d1cf66105dc6acb5af613c054955b8137e34a03698aa53272dbda4afe252be17", size = 445746, upload-time = "2025-12-15T19:20:51.491Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/8e/9d/1a3338e0bd30ada6ad4356c13a0a6c35fbc859063fa7eddb309183364ac1/tornado-6.5.4-cp39-abi3-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:50ff0a58b0dc97939d29da29cd624da010e7f804746621c78d14b80238669335", size = 445083, upload-time = "2025-12-15T19:20:52.778Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/50/d4/e51d52047e7eb9a582da59f32125d17c0482d065afd5d3bc435ff2120dc5/tornado-6.5.4-cp39-abi3-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e5fb5e04efa54cf0baabdd10061eb4148e0be137166146fff835745f59ab9f7f", size = 445315, upload-time = "2025-12-15T19:20:53.996Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/27/07/2273972f69ca63dbc139694a3fc4684edec3ea3f9efabf77ed32483b875c/tornado-6.5.4-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:9c86b1643b33a4cd415f8d0fe53045f913bf07b4a3ef646b735a6a86047dda84", size = 446003, upload-time = "2025-12-15T19:20:56.101Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/d1/83/41c52e47502bf7260044413b6770d1a48dda2f0246f95ee1384a3cd9c44a/tornado-6.5.4-cp39-abi3-musllinux_1_2_i686.whl", hash = "sha256:6eb82872335a53dd063a4f10917b3efd28270b56a33db69009606a0312660a6f", size = 445412, upload-time = "2025-12-15T19:20:57.398Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/10/c7/bc96917f06cbee182d44735d4ecde9c432e25b84f4c2086143013e7b9e52/tornado-6.5.4-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:6076d5dda368c9328ff41ab5d9dd3608e695e8225d1cd0fd1e006f05da3635a8", size = 445392, upload-time = "2025-12-15T19:20:58.692Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/0c/1a/d7592328d037d36f2d2462f4bc1fbb383eec9278bc786c1b111cbbd44cfa/tornado-6.5.4-cp39-abi3-win32.whl", hash = "sha256:1768110f2411d5cd281bac0a090f707223ce77fd110424361092859e089b38d1", size = 446481, upload-time = "2025-12-15T19:21:00.008Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/d6/6d/c69be695a0a64fd37a97db12355a035a6d90f79067a3cf936ec2b1dc38cd/tornado-6.5.4-cp39-abi3-win_amd64.whl", hash = "sha256:fa07d31e0cd85c60713f2b995da613588aa03e1303d75705dca6af8babc18ddc", size = 446886, upload-time = "2025-12-15T19:21:01.287Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/50/49/8dc3fd90902f70084bd2cd059d576ddb4f8bb44c2c7c0e33a11422acb17e/tornado-6.5.4-cp39-abi3-win_arm64.whl", hash = "sha256:053e6e16701eb6cbe641f308f4c1a9541f91b6261991160391bfc342e8a551a1", size = 445910, upload-time = "2025-12-15T19:21:02.571Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
||||
Generated
+1
-1
@@ -268,7 +268,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph"
|
||||
version = "1.1.1"
|
||||
version = "1.1.0"
|
||||
source = { editable = "../langgraph" }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
|
||||
@@ -3,6 +3,6 @@ from langgraph_sdk.client import get_client, get_sync_client
|
||||
from langgraph_sdk.encryption import Encryption
|
||||
from langgraph_sdk.encryption.types import EncryptionContext
|
||||
|
||||
__version__ = "0.3.11"
|
||||
__version__ = "0.3.10"
|
||||
|
||||
__all__ = ["Auth", "Encryption", "EncryptionContext", "get_client", "get_sync_client"]
|
||||
|
||||
@@ -4,11 +4,10 @@ from __future__ import annotations
|
||||
|
||||
import warnings
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime, tzinfo
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from langgraph_sdk._async.http import HttpClient
|
||||
from langgraph_sdk._shared.utilities import _resolve_timezone
|
||||
from langgraph_sdk.schema import (
|
||||
All,
|
||||
Config,
|
||||
@@ -71,7 +70,6 @@ class CronClient:
|
||||
multitask_strategy: str | None = None,
|
||||
end_time: datetime | None = None,
|
||||
enabled: bool | None = None,
|
||||
timezone: str | tzinfo | None = None,
|
||||
stream_mode: StreamMode | Sequence[StreamMode] | None = None,
|
||||
stream_subgraphs: bool | None = None,
|
||||
stream_resumable: bool | None = None,
|
||||
@@ -86,7 +84,7 @@ class CronClient:
|
||||
assistant_id: The assistant ID or graph name to use for the cron job.
|
||||
If using graph name, will default to first assistant created from that graph.
|
||||
schedule: The cron schedule to execute this job on.
|
||||
Schedules are interpreted in UTC unless a timezone is specified.
|
||||
Schedules are interpreted in UTC.
|
||||
input: The input to the graph.
|
||||
metadata: Metadata to assign to the cron job runs.
|
||||
config: The configuration for the assistant.
|
||||
@@ -102,7 +100,6 @@ class CronClient:
|
||||
Must be one of 'reject', 'interrupt', 'rollback', or 'enqueue'.
|
||||
end_time: The time to stop running the cron job. If not provided, the cron job will run indefinitely.
|
||||
enabled: Whether the cron job is enabled or not.
|
||||
timezone: IANA timezone for the cron schedule. Accepts a string (e.g. 'America/New_York') or a ``datetime.tzinfo`` instance (e.g. ``ZoneInfo("America/New_York")``).
|
||||
stream_mode: The stream mode(s) to use.
|
||||
stream_subgraphs: Whether to stream output from subgraphs.
|
||||
stream_resumable: Whether to persist the stream chunks in order to resume the stream later.
|
||||
@@ -155,7 +152,6 @@ class CronClient:
|
||||
"webhook": webhook,
|
||||
"end_time": end_time.isoformat() if end_time else None,
|
||||
"enabled": enabled,
|
||||
"timezone": _resolve_timezone(timezone),
|
||||
"stream_mode": stream_mode,
|
||||
"stream_subgraphs": stream_subgraphs,
|
||||
"stream_resumable": stream_resumable,
|
||||
@@ -188,7 +184,6 @@ class CronClient:
|
||||
multitask_strategy: str | None = None,
|
||||
end_time: datetime | None = None,
|
||||
enabled: bool | None = None,
|
||||
timezone: str | tzinfo | None = None,
|
||||
stream_mode: StreamMode | Sequence[StreamMode] | None = None,
|
||||
stream_subgraphs: bool | None = None,
|
||||
stream_resumable: bool | None = None,
|
||||
@@ -202,7 +197,7 @@ class CronClient:
|
||||
assistant_id: The assistant ID or graph name to use for the cron job.
|
||||
If using graph name, will default to first assistant created from that graph.
|
||||
schedule: The cron schedule to execute this job on.
|
||||
Schedules are interpreted in UTC unless a timezone is specified.
|
||||
Schedules are interpreted in UTC.
|
||||
input: The input to the graph.
|
||||
metadata: Metadata to assign to the cron job runs.
|
||||
config: The configuration for the assistant.
|
||||
@@ -220,7 +215,6 @@ class CronClient:
|
||||
Must be one of 'reject', 'interrupt', 'rollback', or 'enqueue'.
|
||||
end_time: The time to stop running the cron job. If not provided, the cron job will run indefinitely.
|
||||
enabled: Whether the cron job is enabled or not.
|
||||
timezone: IANA timezone for the cron schedule. Accepts a string (e.g. 'America/New_York') or a ``datetime.tzinfo`` instance (e.g. ``ZoneInfo("America/New_York")``).
|
||||
stream_mode: The stream mode(s) to use.
|
||||
stream_subgraphs: Whether to stream output from subgraphs.
|
||||
stream_resumable: Whether to persist the stream chunks in order to resume the stream later.
|
||||
@@ -274,7 +268,6 @@ class CronClient:
|
||||
"on_run_completed": on_run_completed,
|
||||
"end_time": end_time.isoformat() if end_time else None,
|
||||
"enabled": enabled,
|
||||
"timezone": _resolve_timezone(timezone),
|
||||
"stream_mode": stream_mode,
|
||||
"stream_subgraphs": stream_subgraphs,
|
||||
"stream_resumable": stream_resumable,
|
||||
@@ -331,7 +324,6 @@ class CronClient:
|
||||
interrupt_after: All | list[str] | None = None,
|
||||
on_run_completed: OnCompletionBehavior | None = None,
|
||||
enabled: bool | None = None,
|
||||
timezone: str | tzinfo | None = None,
|
||||
stream_mode: StreamMode | Sequence[StreamMode] | None = None,
|
||||
stream_subgraphs: bool | None = None,
|
||||
stream_resumable: bool | None = None,
|
||||
@@ -344,7 +336,7 @@ class CronClient:
|
||||
Args:
|
||||
cron_id: The cron ID to update.
|
||||
schedule: The cron schedule to execute this job on.
|
||||
Schedules are interpreted in UTC unless a timezone is specified.
|
||||
Schedules are interpreted in UTC.
|
||||
end_time: The end date to stop running the cron.
|
||||
input: The input to the graph.
|
||||
metadata: Metadata to assign to the cron job runs.
|
||||
@@ -358,7 +350,6 @@ class CronClient:
|
||||
after execution. 'keep' creates a new thread for each execution but does not
|
||||
clean them up.
|
||||
enabled: Enable or disable the cron job.
|
||||
timezone: IANA timezone for the cron schedule. Accepts a string (e.g. 'America/New_York') or a ``datetime.tzinfo`` instance (e.g. ``ZoneInfo("America/New_York")``).
|
||||
stream_mode: The stream mode(s) to use.
|
||||
stream_subgraphs: Whether to stream output from subgraphs.
|
||||
stream_resumable: Whether to persist the stream chunks in order to resume the stream later.
|
||||
@@ -393,7 +384,6 @@ class CronClient:
|
||||
"interrupt_after": interrupt_after,
|
||||
"on_run_completed": on_run_completed,
|
||||
"enabled": enabled,
|
||||
"timezone": _resolve_timezone(timezone),
|
||||
"stream_mode": stream_mode,
|
||||
"stream_subgraphs": stream_subgraphs,
|
||||
"stream_resumable": stream_resumable,
|
||||
|
||||
@@ -6,17 +6,13 @@ import functools
|
||||
import os
|
||||
import re
|
||||
from collections.abc import Mapping
|
||||
from datetime import tzinfo
|
||||
from typing import TYPE_CHECKING, Any, cast
|
||||
from typing import Any, cast
|
||||
|
||||
import httpx
|
||||
|
||||
import langgraph_sdk
|
||||
from langgraph_sdk.schema import RunCreateMetadata
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from zoneinfo import ZoneInfo
|
||||
|
||||
RESERVED_HEADERS = ("x-api-key",)
|
||||
|
||||
NOT_PROVIDED = cast(None, object())
|
||||
@@ -129,35 +125,6 @@ def _sse_to_v2_dict(event: str, data: Any) -> dict[str, Any] | None:
|
||||
return result
|
||||
|
||||
|
||||
def _resolve_timezone(tz: str | tzinfo | ZoneInfo | None) -> str | None:
|
||||
"""Convert a timezone argument to an IANA timezone string.
|
||||
|
||||
Accepts:
|
||||
- A string (returned as-is, assumed to be an IANA timezone name)
|
||||
- A ``datetime.tzinfo`` instance (e.g. ``zoneinfo.ZoneInfo("America/New_York")``,
|
||||
``datetime.timezone.utc``). The ``key`` attribute is used if available,
|
||||
otherwise ``tzname(None)`` is used.
|
||||
- ``None`` (returned as ``None``)
|
||||
"""
|
||||
if tz is None or isinstance(tz, str):
|
||||
return tz
|
||||
if isinstance(tz, tzinfo):
|
||||
# ZoneInfo objects have a .key attribute with the IANA name
|
||||
if hasattr(tz, "key"):
|
||||
return tz.key # type: ignore[union-attr]
|
||||
# Fall back to tzname for fixed-offset timezones like datetime.timezone.utc
|
||||
name = tz.tzname(None)
|
||||
if name is not None:
|
||||
return name
|
||||
raise ValueError(
|
||||
f"Cannot determine timezone name from {tz!r}. "
|
||||
"Use a zoneinfo.ZoneInfo instance or pass a string like 'America/New_York'."
|
||||
)
|
||||
raise TypeError(
|
||||
f"Expected str, datetime.tzinfo, or None for timezone, got {type(tz).__name__}"
|
||||
)
|
||||
|
||||
|
||||
def _provided_vals(d: Mapping[str, Any]) -> dict[str, Any]:
|
||||
return {k: v for k, v in d.items() if v is not None}
|
||||
|
||||
|
||||
@@ -4,10 +4,9 @@ from __future__ import annotations
|
||||
|
||||
import warnings
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime, tzinfo
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from langgraph_sdk._shared.utilities import _resolve_timezone
|
||||
from langgraph_sdk._sync.http import SyncHttpClient
|
||||
from langgraph_sdk.schema import (
|
||||
All,
|
||||
@@ -65,7 +64,6 @@ class SyncCronClient:
|
||||
multitask_strategy: str | None = None,
|
||||
end_time: datetime | None = None,
|
||||
enabled: bool | None = None,
|
||||
timezone: str | tzinfo | None = None,
|
||||
stream_mode: StreamMode | Sequence[StreamMode] | None = None,
|
||||
stream_subgraphs: bool | None = None,
|
||||
stream_resumable: bool | None = None,
|
||||
@@ -80,7 +78,7 @@ class SyncCronClient:
|
||||
assistant_id: The assistant ID or graph name to use for the cron job.
|
||||
If using graph name, will default to first assistant created from that graph.
|
||||
schedule: The cron schedule to execute this job on.
|
||||
Schedules are interpreted in UTC unless a timezone is specified.
|
||||
Schedules are interpreted in UTC.
|
||||
input: The input to the graph.
|
||||
metadata: Metadata to assign to the cron job runs.
|
||||
config: The configuration for the assistant.
|
||||
@@ -94,7 +92,6 @@ class SyncCronClient:
|
||||
Must be one of 'reject', 'interrupt', 'rollback', or 'enqueue'.
|
||||
end_time: The time to stop running the cron job. If not provided, the cron job will run indefinitely.
|
||||
enabled: Whether the cron job is enabled. By default, it is considered enabled.
|
||||
timezone: IANA timezone for the cron schedule. Accepts a string (e.g. 'America/New_York') or a ``datetime.tzinfo`` instance (e.g. ``ZoneInfo("America/New_York")``).
|
||||
stream_mode: The stream mode(s) to use.
|
||||
stream_subgraphs: Whether to stream output from subgraphs.
|
||||
stream_resumable: Whether to persist the stream chunks in order to resume the stream later.
|
||||
@@ -147,7 +144,6 @@ class SyncCronClient:
|
||||
"multitask_strategy": multitask_strategy,
|
||||
"end_time": end_time.isoformat() if end_time else None,
|
||||
"enabled": enabled,
|
||||
"timezone": _resolve_timezone(timezone),
|
||||
"stream_mode": stream_mode,
|
||||
"stream_subgraphs": stream_subgraphs,
|
||||
"stream_resumable": stream_resumable,
|
||||
@@ -178,7 +174,6 @@ class SyncCronClient:
|
||||
multitask_strategy: str | None = None,
|
||||
end_time: datetime | None = None,
|
||||
enabled: bool | None = None,
|
||||
timezone: str | tzinfo | None = None,
|
||||
stream_mode: StreamMode | Sequence[StreamMode] | None = None,
|
||||
stream_subgraphs: bool | None = None,
|
||||
stream_resumable: bool | None = None,
|
||||
@@ -192,7 +187,7 @@ class SyncCronClient:
|
||||
assistant_id: The assistant ID or graph name to use for the cron job.
|
||||
If using graph name, will default to first assistant created from that graph.
|
||||
schedule: The cron schedule to execute this job on.
|
||||
Schedules are interpreted in UTC unless a timezone is specified.
|
||||
Schedules are interpreted in UTC.
|
||||
input: The input to the graph.
|
||||
metadata: Metadata to assign to the cron job runs.
|
||||
config: The configuration for the assistant.
|
||||
@@ -210,7 +205,6 @@ class SyncCronClient:
|
||||
Must be one of 'reject', 'interrupt', 'rollback', or 'enqueue'.
|
||||
end_time: The time to stop running the cron job. If not provided, the cron job will run indefinitely.
|
||||
enabled: Whether the cron job is enabled. By default, it is considered enabled.
|
||||
timezone: IANA timezone for the cron schedule. Accepts a string (e.g. 'America/New_York') or a ``datetime.tzinfo`` instance (e.g. ``ZoneInfo("America/New_York")``).
|
||||
stream_mode: The stream mode(s) to use.
|
||||
stream_subgraphs: Whether to stream output from subgraphs.
|
||||
stream_resumable: Whether to persist the stream chunks in order to resume the stream later.
|
||||
@@ -265,7 +259,6 @@ class SyncCronClient:
|
||||
"multitask_strategy": multitask_strategy,
|
||||
"end_time": end_time.isoformat() if end_time else None,
|
||||
"enabled": enabled,
|
||||
"timezone": _resolve_timezone(timezone),
|
||||
"stream_mode": stream_mode,
|
||||
"stream_subgraphs": stream_subgraphs,
|
||||
"stream_resumable": stream_resumable,
|
||||
@@ -320,7 +313,6 @@ class SyncCronClient:
|
||||
interrupt_after: All | list[str] | None = None,
|
||||
on_run_completed: OnCompletionBehavior | None = None,
|
||||
enabled: bool | None = None,
|
||||
timezone: str | tzinfo | None = None,
|
||||
stream_mode: StreamMode | Sequence[StreamMode] | None = None,
|
||||
stream_subgraphs: bool | None = None,
|
||||
stream_resumable: bool | None = None,
|
||||
@@ -333,7 +325,7 @@ class SyncCronClient:
|
||||
Args:
|
||||
cron_id: The cron ID to update.
|
||||
schedule: The cron schedule to execute this job on.
|
||||
Schedules are interpreted in UTC unless a timezone is specified.
|
||||
Schedules are interpreted in UTC.
|
||||
end_time: The end date to stop running the cron.
|
||||
input: The input to the graph.
|
||||
metadata: Metadata to assign to the cron job runs.
|
||||
@@ -347,7 +339,6 @@ class SyncCronClient:
|
||||
after execution. 'keep' creates a new thread for each execution but does not
|
||||
clean them up.
|
||||
enabled: Enable or disable the cron job.
|
||||
timezone: IANA timezone for the cron schedule. Accepts a string (e.g. 'America/New_York') or a ``datetime.tzinfo`` instance (e.g. ``ZoneInfo("America/New_York")``).
|
||||
stream_mode: The stream mode(s) to use.
|
||||
stream_subgraphs: Whether to stream output from subgraphs.
|
||||
stream_resumable: Whether to persist the stream chunks in order to resume the stream later.
|
||||
@@ -382,7 +373,6 @@ class SyncCronClient:
|
||||
"interrupt_after": interrupt_after,
|
||||
"on_run_completed": on_run_completed,
|
||||
"enabled": enabled,
|
||||
"timezone": _resolve_timezone(timezone),
|
||||
"stream_mode": stream_mode,
|
||||
"stream_subgraphs": stream_subgraphs,
|
||||
"stream_resumable": stream_resumable,
|
||||
|
||||
@@ -385,8 +385,6 @@ class Cron(TypedDict):
|
||||
"""The end date to stop running the cron."""
|
||||
schedule: str
|
||||
"""The schedule to run, cron format."""
|
||||
timezone: str | None
|
||||
"""IANA timezone for the cron schedule (e.g. 'America/New_York'). Defaults to null, which is treated as UTC."""
|
||||
created_at: datetime
|
||||
"""The time the cron was created."""
|
||||
updated_at: datetime
|
||||
@@ -408,8 +406,6 @@ class CronUpdate(TypedDict, total=False):
|
||||
|
||||
schedule: str
|
||||
"""The cron schedule to execute this job on."""
|
||||
timezone: str
|
||||
"""IANA timezone for the cron schedule (e.g. 'America/New_York')."""
|
||||
end_time: datetime
|
||||
"""The end date to stop running the cron."""
|
||||
input: Input
|
||||
@@ -486,7 +482,6 @@ CronSelectField = Literal[
|
||||
"thread_id",
|
||||
"end_time",
|
||||
"schedule",
|
||||
"timezone",
|
||||
"created_at",
|
||||
"updated_at",
|
||||
"user_id",
|
||||
|
||||
Generated
+1
-1
@@ -265,7 +265,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph"
|
||||
version = "1.1.1"
|
||||
version = "1.1.0"
|
||||
source = { editable = "../langgraph" }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
|
||||
Generated
+338
@@ -0,0 +1,338 @@
|
||||
# This file is automatically @generated by Cargo.
|
||||
# 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"
|
||||
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 = "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"
|
||||
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",
|
||||
"pyo3",
|
||||
"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 = "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"
|
||||
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 = "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"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934"
|
||||
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"
|
||||
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 = "rustversion"
|
||||
version = "1.0.22"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d"
|
||||
|
||||
[[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 = "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"
|
||||
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 = "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"
|
||||
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,21 @@
|
||||
[package]
|
||||
name = "langgraph_rust_core"
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
|
||||
[lib]
|
||||
name = "langgraph_rust_core"
|
||||
crate-type = ["cdylib", "rlib"]
|
||||
|
||||
[features]
|
||||
default = []
|
||||
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"
|
||||
libc = "0.2"
|
||||
tokio = { version = "1", features = ["macros", "rt-multi-thread", "sync", "time"] }
|
||||
|
||||
@@ -0,0 +1,44 @@
|
||||
#ifndef LANGGRAPH_RUST_CORE_H
|
||||
#define LANGGRAPH_RUST_CORE_H
|
||||
|
||||
#ifdef __cplusplus
|
||||
extern "C" {
|
||||
#endif
|
||||
|
||||
typedef struct Engine Engine;
|
||||
typedef char* (*rc_node_callback_t)(
|
||||
unsigned long user_data,
|
||||
char* node,
|
||||
char* arg_json,
|
||||
char* state_json
|
||||
);
|
||||
|
||||
Engine* rc_engine_new(void);
|
||||
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
|
||||
);
|
||||
|
||||
void rc_string_free(char* ptr);
|
||||
|
||||
#ifdef __cplusplus
|
||||
}
|
||||
#endif
|
||||
|
||||
#endif
|
||||
|
||||
@@ -0,0 +1,707 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use std::collections::{HashMap, VecDeque};
|
||||
use std::env;
|
||||
use std::future::Future;
|
||||
use std::sync::mpsc;
|
||||
use std::sync::Arc;
|
||||
use std::sync::Mutex as StdMutex;
|
||||
use std::sync::OnceLock;
|
||||
use std::thread;
|
||||
use std::time::{Duration, Instant};
|
||||
use tokio::runtime::Runtime;
|
||||
use tokio::sync::{mpsc as tokio_mpsc, Mutex as AsyncMutex, Notify};
|
||||
|
||||
#[derive(Clone, Debug, Serialize, Deserialize)]
|
||||
#[serde(tag = "kind")]
|
||||
pub enum WaitCondition {
|
||||
#[serde(rename = "channel")]
|
||||
Channel { channel: String, n: usize },
|
||||
#[serde(rename = "timer")]
|
||||
Timer { seconds: f64 },
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Serialize, Deserialize)]
|
||||
pub struct AnyOfCondition {
|
||||
pub conditions: Vec<WaitCondition>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Serialize, Deserialize)]
|
||||
#[serde(tag = "condition")]
|
||||
pub enum WaitEvent {
|
||||
#[serde(rename = "channel")]
|
||||
Channel {
|
||||
channel: String,
|
||||
value: serde_json::Value,
|
||||
},
|
||||
#[serde(rename = "timer")]
|
||||
Timer { seconds: f64 },
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Serialize, Deserialize)]
|
||||
#[serde(tag = "kind")]
|
||||
pub enum WaitRequest {
|
||||
#[serde(rename = "condition")]
|
||||
Condition { condition: WaitCondition },
|
||||
#[serde(rename = "any_of")]
|
||||
AnyOf { any_of: AnyOfCondition },
|
||||
}
|
||||
|
||||
pub struct SendPayload<A> {
|
||||
pub node: String,
|
||||
pub arg: A,
|
||||
}
|
||||
|
||||
pub struct NodeExecResult<U, A> {
|
||||
pub update: Option<U>,
|
||||
pub sends: Vec<SendPayload<A>>,
|
||||
}
|
||||
|
||||
pub enum NodeOutcome<U, A> {
|
||||
Completed(NodeExecResult<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 {
|
||||
static DEBUG: OnceLock<bool> = OnceLock::new();
|
||||
*DEBUG.get_or_init(|| {
|
||||
matches!(
|
||||
env::var("DEBUG")
|
||||
.unwrap_or_default()
|
||||
.trim()
|
||||
.to_ascii_lowercase()
|
||||
.as_str(),
|
||||
"1" | "true" | "yes" | "on"
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn debug_log(message: &str) {
|
||||
if debug_enabled() {
|
||||
let current = thread::current();
|
||||
let thread_name = current.name().unwrap_or("unnamed");
|
||||
println!("[advanced-graph][{thread_name}] {message}");
|
||||
}
|
||||
}
|
||||
|
||||
fn pool_size_from_env(var_name: &str, default: usize, min: usize) -> usize {
|
||||
let parsed = env::var(var_name)
|
||||
.ok()
|
||||
.and_then(|raw| raw.trim().parse::<usize>().ok());
|
||||
parsed.unwrap_or(default).max(min)
|
||||
}
|
||||
|
||||
struct ThreadPool {
|
||||
tx: mpsc::Sender<Task>,
|
||||
_workers: Vec<thread::JoinHandle<()>>,
|
||||
}
|
||||
|
||||
impl ThreadPool {
|
||||
fn new(size: usize, label: &str) -> Self {
|
||||
let (tx, rx) = mpsc::channel::<Task>();
|
||||
let rx = Arc::new(StdMutex::new(rx));
|
||||
let mut workers = Vec::with_capacity(size);
|
||||
for idx in 0..size {
|
||||
let thread_name = format!("{label}-{idx}");
|
||||
let rx = Arc::clone(&rx);
|
||||
let handle = thread::Builder::new()
|
||||
.name(thread_name)
|
||||
.spawn(move || loop {
|
||||
let task = {
|
||||
let guard = rx.lock().expect("thread-pool receiver mutex poisoned");
|
||||
guard.recv()
|
||||
};
|
||||
match task {
|
||||
Ok(task) => task(),
|
||||
Err(_) => break,
|
||||
}
|
||||
})
|
||||
.expect("failed to spawn thread-pool worker");
|
||||
workers.push(handle);
|
||||
}
|
||||
Self {
|
||||
tx,
|
||||
_workers: workers,
|
||||
}
|
||||
}
|
||||
|
||||
fn execute<F>(&self, task: F) -> Result<(), String>
|
||||
where
|
||||
F: FnOnce() + Send + 'static,
|
||||
{
|
||||
debug_log("thread-pool execute() called");
|
||||
self.tx
|
||||
.send(Box::new(task))
|
||||
.map_err(|e| format!("thread-pool send failed: {e}"))
|
||||
}
|
||||
}
|
||||
|
||||
pub fn node_pool_execute<F>(task: F) -> Result<(), String>
|
||||
where
|
||||
F: FnOnce() + Send + 'static,
|
||||
{
|
||||
debug_log("node_pool_execute() called");
|
||||
static NODE_POOL: OnceLock<ThreadPool> = OnceLock::new();
|
||||
let pool = NODE_POOL.get_or_init(|| {
|
||||
let default_size = thread::available_parallelism()
|
||||
.map(|n| n.get().max(2))
|
||||
.unwrap_or(4);
|
||||
let size = pool_size_from_env("LANGGRAPH_NODE_POOL_SIZE", default_size, 1);
|
||||
ThreadPool::new(size, "langgraph-node")
|
||||
});
|
||||
pool.execute(task)
|
||||
}
|
||||
|
||||
pub fn run_loop_pool_execute<F>(task: F) -> Result<(), String>
|
||||
where
|
||||
F: FnOnce() + Send + 'static,
|
||||
{
|
||||
debug_log("run_loop_pool_execute() called");
|
||||
run_runtime().spawn_blocking(task);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn run_loop_spawn<F>(future: F) -> Result<(), String>
|
||||
where
|
||||
F: Future<Output = ()> + Send + 'static,
|
||||
{
|
||||
run_runtime().spawn(future);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn run_loop_block_on<F>(future: F) -> F::Output
|
||||
where
|
||||
F: Future,
|
||||
{
|
||||
run_runtime().block_on(future)
|
||||
}
|
||||
|
||||
fn run_runtime() -> &'static Runtime {
|
||||
static RUNTIME: OnceLock<Runtime> = OnceLock::new();
|
||||
RUNTIME.get_or_init(|| {
|
||||
let default_size = thread::available_parallelism()
|
||||
.map(|n| n.get().max(2))
|
||||
.unwrap_or(2);
|
||||
let worker_threads = pool_size_from_env("LANGGRAPH_RUN_POOL_SIZE", default_size, 1);
|
||||
tokio::runtime::Builder::new_multi_thread()
|
||||
.worker_threads(worker_threads)
|
||||
.thread_name("langgraph-runloop")
|
||||
.enable_all()
|
||||
.build()
|
||||
.expect("failed to build tokio runtime")
|
||||
})
|
||||
}
|
||||
|
||||
pub fn run_scheduler_loop<U: Send + 'static, A: Send + 'static, FSpawn, FMerge>(
|
||||
entry_point: String,
|
||||
finish_point: &str,
|
||||
initial_arg: A,
|
||||
mut spawn: FSpawn,
|
||||
mut merge: FMerge,
|
||||
rx: mpsc::Receiver<Result<(String, NodeExecResult<U, A>), String>>,
|
||||
) -> Result<(), String>
|
||||
where
|
||||
FSpawn: FnMut(String, A) -> Result<(), String>,
|
||||
FMerge: FnMut(&str, Option<U>) -> Result<(), String>,
|
||||
{
|
||||
debug_log("run_scheduler_loop() started");
|
||||
let mut active: usize = 1;
|
||||
debug_log("scheduling initial entry node");
|
||||
spawn(entry_point, initial_arg)?;
|
||||
|
||||
while active > 0 {
|
||||
debug_log(&format!(
|
||||
"scheduler waiting for node result (active={active})"
|
||||
));
|
||||
let item = rx
|
||||
.recv()
|
||||
.map_err(|e| format!("scheduler recv failed: {e}"))?;
|
||||
active = active.saturating_sub(1);
|
||||
let (node_name, node_result) = item.map_err(|e| format!("node execution failed: {e}"))?;
|
||||
debug_log(&format!("scheduler received result from node={node_name}"));
|
||||
|
||||
merge(&node_name, node_result.update)?;
|
||||
debug_log(&format!("merged update from node={node_name}"));
|
||||
|
||||
if node_name == finish_point {
|
||||
debug_log("finish node reached, stopping scheduler loop");
|
||||
break;
|
||||
}
|
||||
|
||||
for send in node_result.sends {
|
||||
active += 1;
|
||||
debug_log(&format!(
|
||||
"scheduling next node={} (active={active})",
|
||||
send.node
|
||||
));
|
||||
spawn(send.node, send.arg)?;
|
||||
}
|
||||
}
|
||||
|
||||
debug_log("run_scheduler_loop() finished");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn merge_json_update(state: &mut Value, update: Option<Value>) {
|
||||
debug_log("merge_json_update() called");
|
||||
let Some(update_value) = update else {
|
||||
debug_log("merge_json_update(): no update payload");
|
||||
return;
|
||||
};
|
||||
match (&mut *state, update_value) {
|
||||
(Value::Object(state_obj), Value::Object(update_obj)) => {
|
||||
debug_log(&format!(
|
||||
"merge_json_update(): object merge with {} keys",
|
||||
update_obj.len()
|
||||
));
|
||||
for (k, v) in update_obj {
|
||||
state_obj.insert(k, v);
|
||||
}
|
||||
}
|
||||
(Value::Object(state_obj), Value::Array(entries)) => {
|
||||
debug_log(&format!(
|
||||
"merge_json_update(): tuple-list merge with {} entries",
|
||||
entries.len()
|
||||
));
|
||||
for entry in entries {
|
||||
if let Value::Array(pair) = entry {
|
||||
if pair.len() == 2 {
|
||||
if let Value::String(key) = &pair[0] {
|
||||
state_obj.insert(key.clone(), pair[1].clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => debug_log("merge_json_update(): unsupported update shape, ignored"),
|
||||
}
|
||||
}
|
||||
|
||||
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 {
|
||||
pub fn new() -> Self {
|
||||
debug_log("Engine::new()");
|
||||
Self::default()
|
||||
}
|
||||
|
||||
pub fn add_async_channel(&self, name: &str) {
|
||||
debug_log(&format!("Engine::add_async_channel(name={name})"));
|
||||
let mut channels = self.channels.lock().expect("channels mutex poisoned");
|
||||
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");
|
||||
let queue = channels
|
||||
.get_mut(channel)
|
||||
.ok_or_else(|| format!("Unknown channel `{channel}`"))?;
|
||||
queue.push_back(value);
|
||||
self.channel_notify.notify_waiters();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn wait_request_async(&self, wait: &WaitRequest) -> Result<WaitEvent, String> {
|
||||
match wait {
|
||||
WaitRequest::Condition { condition } => self.wait_for_async(condition).await,
|
||||
WaitRequest::AnyOf { any_of } => self.wait_for_any_of_async(any_of).await,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn wait_for_async(&self, cond: &WaitCondition) -> Result<WaitEvent, String> {
|
||||
debug_log(&format!("Engine::wait_for_async(cond={cond:?})"));
|
||||
match cond {
|
||||
WaitCondition::Channel { channel, n } => {
|
||||
if *n < 1 {
|
||||
return Err("channel condition n must be >= 1".to_string());
|
||||
}
|
||||
loop {
|
||||
if let Some(event) = self.try_take_channel_event(channel, *n)? {
|
||||
return Ok(event);
|
||||
}
|
||||
self.channel_notify.notified().await;
|
||||
}
|
||||
}
|
||||
WaitCondition::Timer { seconds } => {
|
||||
if *seconds <= 0.0 {
|
||||
return Err("timer condition must be > 0".to_string());
|
||||
}
|
||||
tokio::time::sleep(Duration::from_secs_f64(*seconds)).await;
|
||||
Ok(WaitEvent::Timer { seconds: *seconds })
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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()
|
||||
));
|
||||
if any_of.conditions.is_empty() {
|
||||
return Err("any_of requires at least one condition".to_string());
|
||||
}
|
||||
|
||||
let started = Instant::now();
|
||||
let mut min_timer: Option<f64> = None;
|
||||
for cond in &any_of.conditions {
|
||||
if let WaitCondition::Timer { seconds } = cond {
|
||||
if *seconds <= 0.0 {
|
||||
return Err("timer condition must be > 0".to_string());
|
||||
}
|
||||
min_timer = Some(min_timer.map_or(*seconds, |x| x.min(*seconds)));
|
||||
}
|
||||
}
|
||||
|
||||
loop {
|
||||
for cond in &any_of.conditions {
|
||||
if let WaitCondition::Channel { channel, n } = cond {
|
||||
if *n < 1 {
|
||||
return Err("channel condition n must be >= 1".to_string());
|
||||
}
|
||||
if let Some(event) = self.try_take_channel_event(channel, *n)? {
|
||||
return Ok(event);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(seconds) = min_timer {
|
||||
let timeout = Duration::from_secs_f64(seconds);
|
||||
let elapsed = started.elapsed();
|
||||
if elapsed >= timeout {
|
||||
return Ok(WaitEvent::Timer { seconds });
|
||||
}
|
||||
let remaining = timeout.saturating_sub(elapsed);
|
||||
tokio::select! {
|
||||
_ = self.channel_notify.notified() => {}
|
||||
_ = tokio::time::sleep(remaining) => {
|
||||
return Ok(WaitEvent::Timer { seconds });
|
||||
}
|
||||
}
|
||||
} else {
|
||||
self.channel_notify.notified().await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn wait_for(&self, cond: &WaitCondition) -> Result<WaitEvent, String> {
|
||||
run_loop_block_on(self.wait_for_async(cond))
|
||||
}
|
||||
|
||||
pub fn wait_for_any_of(&self, any_of: &AnyOfCondition) -> Result<WaitEvent, String> {
|
||||
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> {
|
||||
let mut channels = self.channels.lock().expect("channels mutex poisoned");
|
||||
let queue = channels
|
||||
.get_mut(channel)
|
||||
.ok_or_else(|| format!("Unknown channel `{channel}`"))?;
|
||||
if queue.len() < n {
|
||||
return Ok(None);
|
||||
}
|
||||
if n == 1 {
|
||||
if let Some(value) = queue.pop_front() {
|
||||
return Ok(Some(WaitEvent::Channel {
|
||||
channel: channel.to_string(),
|
||||
value,
|
||||
}));
|
||||
}
|
||||
return Ok(None);
|
||||
}
|
||||
let mut values = Vec::with_capacity(n);
|
||||
for _ in 0..n {
|
||||
if let Some(v) = queue.pop_front() {
|
||||
values.push(v);
|
||||
}
|
||||
}
|
||||
Ok(Some(WaitEvent::Channel {
|
||||
channel: channel.to_string(),
|
||||
value: serde_json::Value::Array(values),
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
||||
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,
|
||||
}))
|
||||
}
|
||||
@@ -0,0 +1,4 @@
|
||||
mod engine;
|
||||
mod lib_c;
|
||||
#[cfg(feature = "python-bindings")]
|
||||
mod lib_py;
|
||||
@@ -0,0 +1,363 @@
|
||||
use crate::engine::{
|
||||
parse_callback_envelope_json, run_graph_json_with_callback, run_loop_block_on, run_loop_spawn,
|
||||
AnyOfCondition, Engine, NodeOutcome,
|
||||
};
|
||||
use serde_json::Value;
|
||||
use std::ffi::{CStr, CString};
|
||||
use std::os::raw::c_char;
|
||||
use std::sync::mpsc;
|
||||
|
||||
type CNodeCallback = unsafe extern "C" fn(
|
||||
user_data: libc::c_ulong,
|
||||
node: *mut c_char,
|
||||
arg_json: *mut c_char,
|
||||
state_json: *mut c_char,
|
||||
) -> *mut c_char;
|
||||
|
||||
fn cstr_to_str<'a>(ptr: *const c_char) -> Result<&'a str, String> {
|
||||
if ptr.is_null() {
|
||||
return Err("Received null pointer".to_string());
|
||||
}
|
||||
let cstr = unsafe { CStr::from_ptr(ptr) };
|
||||
cstr.to_str()
|
||||
.map_err(|e| format!("Invalid UTF-8 input string: {e}"))
|
||||
}
|
||||
|
||||
fn into_c_ptr(s: String) -> *mut c_char {
|
||||
match CString::new(s) {
|
||||
Ok(c) => c.into_raw(),
|
||||
Err(_) => CString::new("{\"error\":\"NUL byte in output\"}")
|
||||
.expect("static string is valid")
|
||||
.into_raw(),
|
||||
}
|
||||
}
|
||||
|
||||
#[no_mangle]
|
||||
pub extern "C" fn rc_engine_new() -> *mut Engine {
|
||||
Box::into_raw(Box::new(Engine::new()))
|
||||
}
|
||||
|
||||
#[no_mangle]
|
||||
/// # Safety
|
||||
/// `ptr` must be either null or a valid pointer returned by `rc_engine_new`.
|
||||
pub unsafe extern "C" fn rc_engine_free(ptr: *mut Engine) {
|
||||
if ptr.is_null() {
|
||||
return;
|
||||
}
|
||||
drop(Box::from_raw(ptr));
|
||||
}
|
||||
|
||||
#[no_mangle]
|
||||
/// # Safety
|
||||
/// `ptr` must be either null or a valid pointer returned by this library.
|
||||
pub unsafe extern "C" fn rc_string_free(ptr: *mut c_char) {
|
||||
if ptr.is_null() {
|
||||
return;
|
||||
}
|
||||
drop(CString::from_raw(ptr));
|
||||
}
|
||||
|
||||
#[no_mangle]
|
||||
/// # Safety
|
||||
/// `ptr` must be a valid engine pointer from `rc_engine_new`.
|
||||
/// `channel` must be a valid null-terminated UTF-8 string pointer.
|
||||
pub unsafe extern "C" fn rc_add_async_channel(
|
||||
ptr: *mut Engine,
|
||||
channel: *const c_char,
|
||||
) -> *mut c_char {
|
||||
if ptr.is_null() {
|
||||
return into_c_ptr("{\"ok\":false,\"error\":\"null engine pointer\"}".to_string());
|
||||
}
|
||||
let channel = match cstr_to_str(channel) {
|
||||
Ok(v) => v,
|
||||
Err(e) => return into_c_ptr(format!("{{\"ok\":false,\"error\":\"{e}\"}}")),
|
||||
};
|
||||
(*ptr).add_async_channel(channel);
|
||||
into_c_ptr("{\"ok\":true}".to_string())
|
||||
}
|
||||
|
||||
#[no_mangle]
|
||||
/// # Safety
|
||||
/// `ptr` must be a valid engine pointer from `rc_engine_new`.
|
||||
/// `channel` and `value_json` must be valid null-terminated UTF-8 string pointers.
|
||||
pub unsafe extern "C" fn rc_publish_json(
|
||||
ptr: *mut Engine,
|
||||
channel: *const c_char,
|
||||
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 channel = match cstr_to_str(channel) {
|
||||
Ok(v) => v,
|
||||
Err(e) => return into_c_ptr(format!("{{\"ok\":false,\"error\":\"{e}\"}}")),
|
||||
};
|
||||
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}\"}}"
|
||||
))
|
||||
}
|
||||
};
|
||||
let result = (*ptr).publish_json(channel, value);
|
||||
match result {
|
||||
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`.
|
||||
/// `any_of_json` must be a valid null-terminated UTF-8 string pointer.
|
||||
pub unsafe extern "C" fn rc_wait_any_of_json(
|
||||
ptr: *mut Engine,
|
||||
any_of_json: *const c_char,
|
||||
) -> *mut c_char {
|
||||
if ptr.is_null() {
|
||||
return into_c_ptr("{\"ok\":false,\"error\":\"null engine pointer\"}".to_string());
|
||||
}
|
||||
let any_of_json = match cstr_to_str(any_of_json) {
|
||||
Ok(v) => v,
|
||||
Err(e) => return into_c_ptr(format!("{{\"ok\":false,\"error\":\"{e}\"}}")),
|
||||
};
|
||||
let any_of: AnyOfCondition = match serde_json::from_str(any_of_json) {
|
||||
Ok(v) => v,
|
||||
Err(e) => {
|
||||
return into_c_ptr(format!(
|
||||
"{{\"ok\":false,\"error\":\"invalid any_of JSON: {e}\"}}"
|
||||
))
|
||||
}
|
||||
};
|
||||
let result = run_loop_block_on((*ptr).wait_for_any_of_async(&any_of));
|
||||
match result {
|
||||
Ok(event) => match serde_json::to_string(&event) {
|
||||
Ok(s) => into_c_ptr(format!("{{\"ok\":true,\"event\":{s}}}")),
|
||||
Err(e) => into_c_ptr(format!("{{\"ok\":false,\"error\":\"{e}\"}}")),
|
||||
},
|
||||
Err(e) => into_c_ptr(format!("{{\"ok\":false,\"error\":\"{e}\"}}")),
|
||||
}
|
||||
}
|
||||
|
||||
#[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`.
|
||||
/// `entry_point`, `finish_point`, and `initial_state_json` must be valid null-terminated UTF-8 pointers.
|
||||
/// `callback` must be a valid function pointer that returns a malloc-allocated C string.
|
||||
pub unsafe extern "C" fn rc_run_graph_json(
|
||||
ptr: *mut Engine,
|
||||
entry_point: *const c_char,
|
||||
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 {
|
||||
if ptr.is_null() {
|
||||
return into_c_ptr("{\"ok\":false,\"error\":\"null engine pointer\"}".to_string());
|
||||
}
|
||||
let Some(callback) = callback else {
|
||||
return into_c_ptr("{\"ok\":false,\"error\":\"null callback pointer\"}".to_string());
|
||||
};
|
||||
let entry_point = match cstr_to_str(entry_point) {
|
||||
Ok(v) => v.to_string(),
|
||||
Err(e) => return into_c_ptr(format!("{{\"ok\":false,\"error\":\"{e}\"}}")),
|
||||
};
|
||||
let finish_point = match cstr_to_str(finish_point) {
|
||||
Ok(v) => v.to_string(),
|
||||
Err(e) => return into_c_ptr(format!("{{\"ok\":false,\"error\":\"{e}\"}}")),
|
||||
};
|
||||
let initial_state_json = match cstr_to_str(initial_state_json) {
|
||||
Ok(v) => v,
|
||||
Err(e) => return into_c_ptr(format!("{{\"ok\":false,\"error\":\"{e}\"}}")),
|
||||
};
|
||||
let initial_state: Value = match serde_json::from_str(initial_state_json) {
|
||||
Ok(v) => v,
|
||||
Err(e) => {
|
||||
return into_c_ptr(format!(
|
||||
"{{\"ok\":false,\"error\":\"invalid initial_state JSON: {e}\"}}"
|
||||
))
|
||||
}
|
||||
};
|
||||
let initial_input_json = match cstr_to_str(initial_input_json) {
|
||||
Ok(v) => v,
|
||||
Err(e) => return into_c_ptr(format!("{{\"ok\":false,\"error\":\"{e}\"}}")),
|
||||
};
|
||||
let initial_input: Value = match serde_json::from_str(initial_input_json) {
|
||||
Ok(v) => v,
|
||||
Err(e) => {
|
||||
return into_c_ptr(format!(
|
||||
"{{\"ok\":false,\"error\":\"invalid initial_input JSON: {e}\"}}"
|
||||
))
|
||||
}
|
||||
};
|
||||
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_bits = user_data;
|
||||
let run_engine = (*ptr).clone();
|
||||
let submit = run_loop_spawn(async move {
|
||||
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.clone(),
|
||||
callback_wrapper,
|
||||
)
|
||||
.await;
|
||||
run_engine.close_stream();
|
||||
let _ = tx.send(out);
|
||||
});
|
||||
if let Err(e) = submit {
|
||||
return into_c_ptr(format!("{{\"ok\":false,\"error\":\"{e}\"}}"));
|
||||
}
|
||||
|
||||
let run_result = match rx.recv() {
|
||||
Ok(v) => v,
|
||||
Err(e) => {
|
||||
return into_c_ptr(format!(
|
||||
"{{\"ok\":false,\"error\":\"run-loop recv failed: {e}\"}}"
|
||||
))
|
||||
}
|
||||
};
|
||||
match run_result {
|
||||
Ok(state) => match serde_json::to_string(&state) {
|
||||
Ok(s) => into_c_ptr(format!("{{\"ok\":true,\"state\":{s}}}")),
|
||||
Err(e) => into_c_ptr(format!(
|
||||
"{{\"ok\":false,\"error\":\"serialize state failed: {e}\"}}"
|
||||
)),
|
||||
},
|
||||
Err(e) => into_c_ptr(format!("{{\"ok\":false,\"error\":\"{e}\"}}")),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,215 @@
|
||||
#[cfg(feature = "python-bindings")]
|
||||
use crate::engine::{
|
||||
parse_callback_envelope_json, run_graph_json_with_callback, run_loop_block_on, AnyOfCondition,
|
||||
Engine, NodeOutcome, WaitCondition,
|
||||
};
|
||||
#[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 serde_json::Value;
|
||||
|
||||
#[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))?;
|
||||
self.publish_json(channel, &value_json)
|
||||
}
|
||||
|
||||
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 event_json = self.wait_any_of_json(&payload_json)?;
|
||||
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}")))
|
||||
}
|
||||
|
||||
#[pyo3(signature = (stream_mode=None))]
|
||||
fn start_stream(&self, stream_mode: Option<&str>) -> PyResult<()> {
|
||||
self.inner
|
||||
.start_stream(stream_mode)
|
||||
.map_err(PyValueError::new_err)
|
||||
}
|
||||
|
||||
fn receive_stream_obj(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
|
||||
match run_loop_block_on(self.inner.receive_stream_async()) {
|
||||
Some(value) => {
|
||||
let event_json = serde_json::to_string(&value)
|
||||
.map_err(|e| PyValueError::new_err(format!("Serialize event failed: {e}")))?;
|
||||
json_string_to_py_obj(py, &event_json)
|
||||
}
|
||||
None => Ok(py.None()),
|
||||
}
|
||||
}
|
||||
|
||||
fn send_custom_stream_event_obj(&self, py: Python<'_>, value: Py<PyAny>) -> PyResult<()> {
|
||||
let value_json = py_obj_to_json_string(py, &value.bind(py))?;
|
||||
let parsed: Value = serde_json::from_str(&value_json)
|
||||
.map_err(|e| PyValueError::new_err(format!("Invalid Python JSON value: {e}")))?;
|
||||
self.inner.send_custom_stream_event(parsed);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn close_stream(&self) {
|
||||
self.inner.close_stream();
|
||||
}
|
||||
|
||||
#[pyo3(signature = (entry_point, finish_point, initial_state, callback, stream_mode=None))]
|
||||
fn run_graph_py(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
entry_point: &str,
|
||||
finish_point: &str,
|
||||
initial_state: Py<PyAny>,
|
||||
callback: Py<PyAny>,
|
||||
stream_mode: Option<&str>,
|
||||
) -> PyResult<Py<PyAny>> {
|
||||
let state_json = py_obj_to_json_string(py, &initial_state.bind(py))?;
|
||||
let initial_state_value: Value = serde_json::from_str(&state_json)
|
||||
.map_err(|e| PyValueError::new_err(format!("Invalid initial state: {e}")))?;
|
||||
|
||||
self.inner
|
||||
.start_stream(stream_mode)
|
||||
.map_err(PyValueError::new_err)?;
|
||||
|
||||
let callback_arc = std::sync::Arc::new(callback);
|
||||
let run_result = run_loop_block_on(run_graph_json_with_callback(
|
||||
entry_point.to_string(),
|
||||
finish_point.to_string(),
|
||||
initial_state_value.clone(),
|
||||
initial_state_value,
|
||||
self.inner.clone(),
|
||||
move |node: String,
|
||||
arg: Value,
|
||||
state_snapshot: Value|
|
||||
-> Result<NodeOutcome<Value, Value>, String> {
|
||||
Python::with_gil(|py| -> Result<NodeOutcome<Value, Value>, String> {
|
||||
let callback_bound = callback_arc.as_ref().bind(py);
|
||||
let arg_json = serde_json::to_string(&arg)
|
||||
.map_err(|e| format!("serialize arg failed: {e}"))?;
|
||||
let state_json = serde_json::to_string(&state_snapshot)
|
||||
.map_err(|e| format!("serialize state failed: {e}"))?;
|
||||
let arg_obj = json_string_to_py_obj(py, &arg_json)
|
||||
.map_err(|e| format!("decode arg failed: {e}"))?;
|
||||
let state_obj = json_string_to_py_obj(py, &state_json)
|
||||
.map_err(|e| format!("decode state failed: {e}"))?;
|
||||
let payload_obj = callback_bound
|
||||
.call1((node.as_str(), arg_obj, state_obj))
|
||||
.map_err(|e| format!("callback failed for `{node}`: {e}"))?;
|
||||
let payload_json = py_obj_to_json_string(py, &payload_obj)
|
||||
.map_err(|e| format!("serialize callback payload failed: {e}"))?;
|
||||
let payload_value: Value = serde_json::from_str(&payload_json)
|
||||
.map_err(|e| format!("decode callback payload failed: {e}"))?;
|
||||
let envelope = serde_json::json!({
|
||||
"ok": true,
|
||||
"payload": {
|
||||
"update": payload_value
|
||||
.get("update")
|
||||
.cloned()
|
||||
.unwrap_or(Value::Null),
|
||||
"sends": payload_value
|
||||
.get("sends")
|
||||
.cloned()
|
||||
.unwrap_or(Value::Array(vec![])),
|
||||
},
|
||||
"suspend": payload_value.get("suspend").cloned(),
|
||||
});
|
||||
parse_callback_envelope_json(&envelope.to_string(), &node)
|
||||
})
|
||||
},
|
||||
));
|
||||
self.inner.close_stream();
|
||||
|
||||
let out = run_result.map_err(PyValueError::new_err)?;
|
||||
let out_json = serde_json::to_string(&out)
|
||||
.map_err(|e| PyValueError::new_err(format!("Serialize state failed: {e}")))?;
|
||||
json_string_to_py_obj(py, &out_json)
|
||||
}
|
||||
}
|
||||
|
||||
#[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
+338
@@ -0,0 +1,338 @@
|
||||
# This file is automatically @generated by Cargo.
|
||||
# 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"
|
||||
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 = "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"
|
||||
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",
|
||||
"pyo3",
|
||||
"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 = "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"
|
||||
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 = "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"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934"
|
||||
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"
|
||||
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 = "rustversion"
|
||||
version = "1.0.22"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d"
|
||||
|
||||
[[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 = "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"
|
||||
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 = "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"
|
||||
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,22 @@
|
||||
[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"]
|
||||
|
||||
[features]
|
||||
default = []
|
||||
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"
|
||||
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)
|
||||
- `langgraph_rust_core` (Rust execution engine via PyO3)
|
||||
|
||||
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/`
|
||||
|
||||
@@ -0,0 +1,127 @@
|
||||
# Evolve/Extend LangGraph with next level of orchestration
|
||||
|
||||
## LangGraph Today: A Strong Foundation with Creative Innovation
|
||||
|
||||
LangGraph is already an exceptional orchestration framework. It has introduced a number of creative features that no other workflow engine on the market has even attempted.
|
||||
|
||||
**First-class streaming.** No workflow engine has ever integrated streaming as seamlessly as LangGraph. Streaming is not an afterthought bolted on top — it is woven into the core execution model, allowing every node, every tool call, and every LLM interaction to emit incremental output naturally.
|
||||
|
||||
**Flexible durability modes.** LangGraph defaults to asynchronous execution and supports sync and "exit" modes as well. This is a significant departure from traditional workflow engines, which typically only offer synchronous execution. The ability to choose a durability mode gives developers fine-grained control over the trade-off between persistence guarantees and execution speed.
|
||||
|
||||
**Reusable checkpoints.** The checkpoint system allows state to be captured at any point during graph execution and freely replayed, forked, or resumed later. This enables powerful patterns like time-travel debugging, human-in-the-loop approval flows, and long-running conversations that can be picked up exactly where they left off.
|
||||
|
||||
**Double texting.** LangGraph natively handles the real-world scenario where a user sends a new message while a previous one is still being processed — a problem most orchestration frameworks simply ignore.
|
||||
|
||||
Beyond these innovative features, LangGraph provides solid support for the foundational workflow execution patterns that developers rely on daily. Sequential execution, or loops and conditional branching. Basic parallelism is also well supported: when multiple LLM calls or tool invocations are independent of each other, they can run concurrently to avoid the latency cost of sequential execution, and their results are merged back into the shared state for downstream processing.
|
||||
|
||||
LangGraph also offers a simple and intuitive mechanism for human-in-the-loop interactions, allowing a graph to pause execution and wait for user input before continuing.
|
||||
|
||||
Combined with the broader LangChain ecosystem, these have made LangGraph a significant success in the market.
|
||||
|
||||
## Emerging Gaps: What LangGraph Struggles to Support
|
||||
|
||||
As adoption has grown and use cases have become more sophisticated, we have discovered an increasing number of scenarios and design patterns that LangGraph cannot support well today.
|
||||
|
||||
**Complex sub-agent coordination.** A main agent often needs to manage multiple sub-agents, but the coordination involved is far more nuanced than simply launching a batch of sub-agents, waiting for all of them to finish, and then moving on. In practice, a main agent may launch a sub-agent, continue doing other work, spawn additional sub-agents later, wait selectively for certain results, retry with a different strategy if one sub-agent fails, or dynamically decide what to do next based on partial results that arrive at unpredictable times.
|
||||
|
||||
LangGraph today lacks the coordination primitives to express this. The current parallelism model groups multiple nodes into a single superstep — all of them execute concurrently, but _all_ must complete before the graph can advance to the next step. There is no way for one node to proceed independently while others are still running, and no built-in mechanism for selective waiting, partial result handling, or dynamic task spawning mid-execution.
|
||||
|
||||
Sub-agents also cannot simply be modeled as subgraphs, because subgraphs today execute within the same run. They cannot be scaled up independently — if a sub-agent is resource-intensive, there is no straightforward way to run it on a separate machine. Ideally, launching a sub-agent should be(or opt in) as simple as dispatching it for distributed execution across multiple machines.
|
||||
|
||||
**Concurrent input and output (e.g. audio agents).** Audio agents also present a particularly clear example of a pattern LangGraph cannot express today. In a voice interaction, speech input and speech output may happen simultaneously — the agent should be able to process a previous utterance, continue receiving new audio input, and produce output all at the same time. These three activities should not be mutually exclusive.
|
||||
|
||||
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.
|
||||
|
||||
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?
|
||||
|
||||
At its most fundamental level, a workflow engine's value proposition is making a long-running process execute reliably. If a machine crashes, execution should smoothly fail over to another machine and resume from the last point where it was interrupted — not start over from the beginning. So we can reason about what is needed by asking: what would a developer do if they had to build a long-running process _without_ a workflow engine?
|
||||
|
||||
Starting from the simple. A developer could write a simple `main` function — a single-threaded program, just like everyone writes when they first learn to code. It would have `if/else` branches, `for` loops, and maybe it would wait for command-line input. Many early agent use cases look exactly like this: execute a sequence of steps, make decisions along the way, loop when necessary.
|
||||
|
||||
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.
|
||||
|
||||
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.
|
||||
|
||||
### From Single-Threaded to Concurrent: Where the Model Breaks Down
|
||||
|
||||
But as product requirements grow more complex, a single-threaded program is no longer sufficient. The process becomes multi-threaded or multi-process. And in a multi-threaded program, each thread executes independently — when one thread finishes a step and moves on to its next step, it does not need to wait for another thread to finish _its_ current step first.
|
||||
|
||||
This is precisely why LangGraph's superstep restriction feels awkward in practice. In the superstep model, all concurrently executing nodes must complete before any of them can advance. But that is not how independent threads work. Each thread should be able to progress at its own pace, checkpoint its own state, and move to its next step without being blocked by unrelated work happening in parallel.
|
||||
|
||||
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.
|
||||
|
||||
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:
|
||||
|
||||
1. **Variable-size channels.** The channel buffer size should be configurable — size 0 for synchronous handoff (like `interrupt` today), size N for buffered communication where the sender can proceed without waiting, and unbounded for fully asynchronous fire-and-forget messaging.
|
||||
2. **Channels across boundaries.** Channels should not be limited to communication between separate runs. Nodes within the same graph should also be able to send and receive through channels — mirroring the way both multi-process communication (between runs) and multi-thread communication (between nodes within a run) work in ordinary concurrent programs.
|
||||
3. **Node-level blocking, not run-level pausing.** When a node waits on a channel (i.e. `interrupt`), only that node should block — the rest of the graph should continue executing. Today, `interrupt` pauses the entire run. In a concurrent program, when one thread blocks on a channel read, the other threads keep running. The same should be true: an interrupt should suspend the individual node, not halt the whole run.
|
||||
|
||||
## Summmary of all extension opportunity
|
||||
|
||||
### P1: urgently needed
|
||||
#### Remove the Superstep Restriction
|
||||
|
||||
Today, when multiple nodes execute in parallel, they are grouped into a superstep. All nodes in a superstep must complete before any downstream node can begin. This means that even if `b1` finishes quickly and its successor `b11` is ready to run, it must wait for `b2` to finish first.
|
||||
|
||||
With the superstep restriction removed, each parallel branch progresses independently. As soon as a node completes, its downstream successor can begin immediately — regardless of what is happening in other branches.
|
||||
|
||||
**Current behavior (superstep model):**
|
||||
|
||||
```
|
||||
Step 1: a
|
||||
Step 2: b1, b2 ← both must finish before step 3
|
||||
Step 3: b11, b22 ← both start together
|
||||
```
|
||||
|
||||
Even if `b1` finishes in 1 second and `b2` takes 30 seconds, `b11` cannot start until `b2` is done.
|
||||
|
||||
**Proposed behavior (independent branches):**
|
||||
|
||||
```
|
||||
Branch 1: a → b1 → b11 → ...
|
||||
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.
|
||||
|
||||
#### Light-weight Interrupt -- wait_for API: Only Block the Current Node Until Channel Has Enough Messages
|
||||
|
||||
Today, `interrupt` pauses the entire run. Every node stops, and nothing can proceed until the interrupt is resolved externally. This is the right behavior for a simple single-threaded workflow, but it breaks down when multiple branches are executing concurrently — one branch needing input should not freeze all the others.
|
||||
|
||||
The proposed change has three parts:
|
||||
|
||||
1. **Named channels.** A graph can declare named channels as coordination points. These are distinct from the graph's state — they are message-passing primitives, not shared memory.
|
||||
2. **`wait_for` blocks only the current node.** When a node calls `wait_for`, it suspends itself and waits for messages on the specified channel. All other nodes in the graph continue executing normally.
|
||||
3. A channel can be published from both external and internal
|
||||
|
||||
The `wait_for` call takes a channel name and optionally a count `N`, meaning "wait until N messages have arrived on this channel before resuming."
|
||||
|
||||
**Prototype:**
|
||||
|
||||
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
|
||||
|
||||
@@ -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"]
|
||||
|
||||
@@ -0,0 +1,26 @@
|
||||
from .state import (
|
||||
AdvancedStateGraph,
|
||||
AnyOfCondition,
|
||||
ChannelCondition,
|
||||
CompiledGraphEngine,
|
||||
Context,
|
||||
GraphRunHandler,
|
||||
TimerCondition,
|
||||
any_of,
|
||||
channel_condition,
|
||||
timer_condition,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"AdvancedStateGraph",
|
||||
"CompiledGraphEngine",
|
||||
"Context",
|
||||
"GraphRunHandler",
|
||||
"ChannelCondition",
|
||||
"TimerCondition",
|
||||
"AnyOfCondition",
|
||||
"channel_condition",
|
||||
"timer_condition",
|
||||
"any_of",
|
||||
]
|
||||
|
||||
@@ -0,0 +1,573 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import atexit
|
||||
import asyncio
|
||||
import inspect
|
||||
import os
|
||||
import threading
|
||||
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 saf_python_sdk.langgraph_rust_core import PyRustEngine # type: ignore[import-untyped]
|
||||
|
||||
from saf_python_sdk.types import Command, Send
|
||||
|
||||
StateT = TypeVar("StateT")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _ChannelSpec:
|
||||
typ: Any
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ChannelCondition:
|
||||
channel: str
|
||||
n: int = 1
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TimerCondition:
|
||||
seconds: float
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class AnyOfCondition:
|
||||
conditions: tuple[WaitCondition, ...]
|
||||
|
||||
|
||||
WaitCondition = ChannelCondition | TimerCondition
|
||||
|
||||
_EXECUTOR_LOCK = threading.Lock()
|
||||
_EXECUTOR: ThreadPoolExecutor | None = None
|
||||
|
||||
|
||||
def _advanced_graph_executor() -> ThreadPoolExecutor:
|
||||
global _EXECUTOR
|
||||
with _EXECUTOR_LOCK:
|
||||
if _EXECUTOR is None:
|
||||
worker_count = int(os.getenv("LANGGRAPH_ADVANCED_GRAPH_PY_THREADS", "256"))
|
||||
worker_count = max(worker_count, 1)
|
||||
_EXECUTOR = ThreadPoolExecutor(
|
||||
max_workers=worker_count,
|
||||
thread_name_prefix="saf-advanced-py",
|
||||
)
|
||||
atexit.register(_shutdown_advanced_graph_executor)
|
||||
return _EXECUTOR
|
||||
|
||||
|
||||
def _shutdown_advanced_graph_executor() -> None:
|
||||
global _EXECUTOR
|
||||
with _EXECUTOR_LOCK:
|
||||
if _EXECUTOR is not None:
|
||||
_EXECUTOR.shutdown(wait=False, cancel_futures=False)
|
||||
_EXECUTOR = None
|
||||
|
||||
|
||||
class WaitRequested(Exception):
|
||||
def __init__(self, payload: dict[str, Any]) -> None:
|
||||
super().__init__("wait requested")
|
||||
self.payload = payload
|
||||
|
||||
|
||||
class AdvancedStateGraph(Generic[StateT]):
|
||||
"""Experimental in-memory graph engine with async channels."""
|
||||
|
||||
def __init__(self, state_schema: type[StateT]) -> None:
|
||||
self.state_schema = state_schema
|
||||
self._nodes: dict[str, Callable[..., Any]] = {}
|
||||
self._async_channels: dict[str, _ChannelSpec] = {}
|
||||
self._entry_point: str | None = None
|
||||
self._finish_point: str | None = None
|
||||
|
||||
def add_node(
|
||||
self,
|
||||
name_or_node: str | Callable[..., Any],
|
||||
node: Callable[..., Any] | None = None,
|
||||
) -> str:
|
||||
if node is None:
|
||||
if not callable(name_or_node):
|
||||
raise TypeError("add_node() expects a callable when name is omitted")
|
||||
node_name = _infer_node_name(name_or_node)
|
||||
node_fn = name_or_node
|
||||
else:
|
||||
if not isinstance(name_or_node, str):
|
||||
raise TypeError("add_node() expects a string node name")
|
||||
node_name = name_or_node
|
||||
node_fn = node
|
||||
|
||||
if node_name in self._nodes:
|
||||
raise ValueError(f"Node `{node_name}` already exists")
|
||||
self._nodes[node_name] = node_fn
|
||||
return node_name
|
||||
|
||||
def add_async_channel(self, name: str, typ: Any) -> None:
|
||||
if name in self._async_channels:
|
||||
raise ValueError(f"Channel `{name}` already exists")
|
||||
self._async_channels[name] = _ChannelSpec(typ=typ)
|
||||
|
||||
def add_entry_node(self, node: Callable[..., Any]) -> str:
|
||||
node_name = self.add_node(node)
|
||||
self._entry_point = self._resolve_node_name(node_name)
|
||||
return node_name
|
||||
|
||||
def add_finish_node(self, node: Callable[..., Any]) -> str:
|
||||
node_name = self.add_node(node)
|
||||
self._finish_point = self._resolve_node_name(node_name)
|
||||
return node_name
|
||||
|
||||
def _resolve_node_name(self, name_or_node: str | Callable[..., Any]) -> str:
|
||||
if isinstance(name_or_node, str):
|
||||
return name_or_node
|
||||
node_name = _infer_node_name(name_or_node)
|
||||
if node_name not in self._nodes:
|
||||
self._nodes[node_name] = name_or_node
|
||||
return node_name
|
||||
|
||||
def compile(self) -> CompiledGraphEngine[StateT]:
|
||||
if self._entry_point is None:
|
||||
raise ValueError("Entry point is not set")
|
||||
if self._entry_point not in self._nodes:
|
||||
raise ValueError(f"Entry point node `{self._entry_point}` does not exist")
|
||||
if self._finish_point is not None and self._finish_point not in self._nodes:
|
||||
raise ValueError(f"Finish point node `{self._finish_point}` does not exist")
|
||||
return CompiledGraphEngine(
|
||||
nodes=dict(self._nodes),
|
||||
async_channels=dict(self._async_channels),
|
||||
entry_point=self._entry_point,
|
||||
finish_point=self._finish_point,
|
||||
)
|
||||
|
||||
|
||||
class CompiledGraphEngine(Generic[StateT]):
|
||||
"""Executable runtime for `AdvancedStateGraph`."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
nodes: dict[str, Callable[..., Any]],
|
||||
async_channels: dict[str, _ChannelSpec],
|
||||
entry_point: str,
|
||||
finish_point: str | None,
|
||||
) -> None:
|
||||
self._nodes = nodes
|
||||
self._async_channels = async_channels
|
||||
self._entry_point = entry_point
|
||||
self._finish_point = finish_point
|
||||
|
||||
async def ainvoke(self, initial_state: StateT) -> StateT:
|
||||
handler = await self.astart(initial_state)
|
||||
return await handler
|
||||
|
||||
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)
|
||||
|
||||
|
||||
class Context:
|
||||
"""Per-run context injected into advanced graph nodes."""
|
||||
|
||||
def __init__(self, run: _GraphEngineRun) -> None:
|
||||
self._run = run
|
||||
|
||||
async def wait_for(self, target: WaitCondition | AnyOfCondition) -> Any:
|
||||
resumed = self._run._consume_resume_event(target)
|
||||
if resumed is not None:
|
||||
return resumed
|
||||
raise WaitRequested(_target_to_suspend_payload(target))
|
||||
|
||||
def publish_to_channel(self, channel: str, value: Any) -> None:
|
||||
self._run.publish_nowait(channel, value)
|
||||
|
||||
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."""
|
||||
|
||||
def __init__(self, *, run: _GraphEngineRun, task: asyncio.Task[StateT]) -> None:
|
||||
self._run = run
|
||||
self._task = task
|
||||
|
||||
async def apublish_to_channel(self, channel: str, value: Any) -> None:
|
||||
if self._task.done():
|
||||
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
|
||||
|
||||
def __await__(self) -> Any:
|
||||
return self._task.__await__()
|
||||
|
||||
|
||||
class _GraphEngineRun:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
nodes: dict[str, Callable[..., Any]],
|
||||
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)
|
||||
self._tasks: set[asyncio.Task[list[Send]]] = set()
|
||||
self._finished = False
|
||||
self._state: Any = None
|
||||
self._local = threading.local()
|
||||
self.context = Context(self)
|
||||
|
||||
async def run(self, initial_state: StateT) -> StateT:
|
||||
finish_point = self._finish_point or ""
|
||||
loop = asyncio.get_running_loop()
|
||||
result_obj = await loop.run_in_executor(
|
||||
_advanced_graph_executor(),
|
||||
self._rust_engine.run_graph_py,
|
||||
self._entry_point,
|
||||
finish_point,
|
||||
initial_state,
|
||||
self._execute_node_for_rust,
|
||||
self._stream_mode,
|
||||
)
|
||||
self._state = result_obj
|
||||
return cast(StateT, self._state)
|
||||
|
||||
async def publish(self, channel: str, value: Any) -> None:
|
||||
loop = asyncio.get_running_loop()
|
||||
await loop.run_in_executor(
|
||||
_advanced_graph_executor(),
|
||||
self._publish_sync,
|
||||
channel,
|
||||
value,
|
||||
)
|
||||
|
||||
def publish_nowait(self, channel: str, value: Any) -> None:
|
||||
self._publish_sync(channel, value)
|
||||
|
||||
async def wait_for(self, target: WaitCondition | AnyOfCondition) -> Any:
|
||||
if isinstance(target, ChannelCondition):
|
||||
value = await self._wait_for_channel_values(target.channel, n=target.n)
|
||||
return {
|
||||
"condition": "channel",
|
||||
"channel": target.channel,
|
||||
"value": value,
|
||||
}
|
||||
if isinstance(target, TimerCondition):
|
||||
loop = asyncio.get_running_loop()
|
||||
return await loop.run_in_executor(
|
||||
_advanced_graph_executor(),
|
||||
self._rust_engine.wait_timer,
|
||||
target.seconds,
|
||||
)
|
||||
if isinstance(target, AnyOfCondition):
|
||||
return await self._wait_for_any_of(target)
|
||||
raise ValueError(f"Unsupported wait condition type: {type(target)!r}")
|
||||
|
||||
async def _wait_for_channel_values(self, channel: str, n: int) -> Any:
|
||||
if n < 1:
|
||||
raise ValueError("wait_for count `n` must be >= 1")
|
||||
loop = asyncio.get_running_loop()
|
||||
event = await loop.run_in_executor(
|
||||
_advanced_graph_executor(),
|
||||
self._rust_engine.wait_channel,
|
||||
channel,
|
||||
n,
|
||||
)
|
||||
return event["value"]
|
||||
|
||||
async def _wait_for_any_of(self, condition: AnyOfCondition) -> Any:
|
||||
if not condition.conditions:
|
||||
raise ValueError("any_of() requires at least one condition")
|
||||
payload = {
|
||||
"conditions": [_condition_to_rust(cond) for cond in condition.conditions]
|
||||
}
|
||||
loop = asyncio.get_running_loop()
|
||||
return await loop.run_in_executor(
|
||||
_advanced_graph_executor(),
|
||||
self._rust_engine.wait_any_of_obj,
|
||||
payload,
|
||||
)
|
||||
|
||||
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]:
|
||||
node_input, resume_event = _unwrap_resume_input(node_input)
|
||||
self._set_resume_event(resume_event)
|
||||
if node_name not in self._nodes:
|
||||
raise ValueError(f"Unknown node `{node_name}`")
|
||||
node = self._nodes[node_name]
|
||||
try:
|
||||
result = _invoke_node(node, self.context, node_input, state)
|
||||
if inspect.isawaitable(result):
|
||||
result = self._run_awaitable_in_worker(
|
||||
cast(Coroutine[Any, Any, Any], result)
|
||||
)
|
||||
except WaitRequested as suspend:
|
||||
return {"suspend": suspend.payload}
|
||||
finally:
|
||||
self._set_resume_event(None)
|
||||
|
||||
if isinstance(result, Command):
|
||||
update = result.update
|
||||
sends = _normalize_goto(result.goto, default_input=node_input)
|
||||
else:
|
||||
update = result
|
||||
sends = _normalize_result_to_sends(result, default_input=node_input)
|
||||
|
||||
return {
|
||||
"update": update,
|
||||
"sends": [
|
||||
{"node": _resolve_target_name(send.node), "arg": send.arg}
|
||||
for send in sends
|
||||
],
|
||||
}
|
||||
|
||||
def _set_resume_event(self, event: dict[str, Any] | None) -> None:
|
||||
self._local.resume_event = event
|
||||
|
||||
def _consume_resume_event(self, target: WaitCondition | AnyOfCondition) -> Any | None:
|
||||
event = cast(dict[str, Any] | None, getattr(self._local, "resume_event", None))
|
||||
if event is None:
|
||||
return None
|
||||
self._local.resume_event = None
|
||||
return event
|
||||
|
||||
def _run_awaitable_in_worker(self, awaitable: Coroutine[Any, Any, Any]) -> Any:
|
||||
# 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:
|
||||
return []
|
||||
if isinstance(result, Send):
|
||||
return [result]
|
||||
if callable(result):
|
||||
return [Send(_infer_node_name(result), default_input)]
|
||||
if isinstance(result, str):
|
||||
return [Send(result, default_input)]
|
||||
if isinstance(result, Sequence) and not isinstance(result, (str, bytes)):
|
||||
sends: list[Send] = []
|
||||
for item in result:
|
||||
if isinstance(item, Send):
|
||||
sends.append(item)
|
||||
elif callable(item):
|
||||
sends.append(Send(_infer_node_name(item), default_input))
|
||||
elif isinstance(item, str):
|
||||
sends.append(Send(item, default_input))
|
||||
return sends
|
||||
return []
|
||||
|
||||
|
||||
def _normalize_goto(goto: Any, *, default_input: Any) -> list[Send]:
|
||||
if not goto:
|
||||
return []
|
||||
if isinstance(goto, Send):
|
||||
return [goto]
|
||||
if callable(goto):
|
||||
return [Send(_infer_node_name(goto), default_input)]
|
||||
if isinstance(goto, str):
|
||||
return [Send(goto, default_input)]
|
||||
if isinstance(goto, Sequence):
|
||||
sends: list[Send] = []
|
||||
for item in goto:
|
||||
if isinstance(item, Send):
|
||||
sends.append(item)
|
||||
elif callable(item):
|
||||
sends.append(Send(_infer_node_name(item), default_input))
|
||||
elif isinstance(item, str):
|
||||
sends.append(Send(item, default_input))
|
||||
return sends
|
||||
return []
|
||||
|
||||
|
||||
def channel_condition(channel: str, n: int = 1) -> ChannelCondition:
|
||||
if n < 1:
|
||||
raise ValueError("channel_condition `n` must be >= 1")
|
||||
return ChannelCondition(channel=channel, n=n)
|
||||
|
||||
|
||||
def timer_condition(
|
||||
timeout: float | timedelta | None = None,
|
||||
*,
|
||||
seconds: float | None = None,
|
||||
minutes: float | None = None,
|
||||
) -> TimerCondition:
|
||||
if timeout is not None and (seconds is not None or minutes is not None):
|
||||
raise ValueError(
|
||||
"Provide either `timeout` or named `seconds`/`minutes`, not both"
|
||||
)
|
||||
|
||||
if isinstance(timeout, timedelta):
|
||||
resolved_seconds = timeout.total_seconds()
|
||||
elif isinstance(timeout, (int, float)):
|
||||
resolved_seconds = float(timeout)
|
||||
else:
|
||||
resolved_seconds = 0.0
|
||||
if seconds is not None:
|
||||
resolved_seconds += float(seconds)
|
||||
if minutes is not None:
|
||||
resolved_seconds += float(minutes) * 60.0
|
||||
|
||||
if resolved_seconds <= 0:
|
||||
raise ValueError("timer_condition must be greater than 0 seconds")
|
||||
return TimerCondition(seconds=resolved_seconds)
|
||||
|
||||
|
||||
def any_of(*conditions: WaitCondition) -> AnyOfCondition:
|
||||
if not conditions:
|
||||
raise ValueError("any_of() requires at least one condition")
|
||||
return AnyOfCondition(conditions=tuple(conditions))
|
||||
|
||||
|
||||
def _condition_to_rust(condition: WaitCondition) -> dict[str, Any]:
|
||||
if isinstance(condition, ChannelCondition):
|
||||
return {"kind": "channel", "channel": condition.channel, "n": condition.n}
|
||||
if isinstance(condition, TimerCondition):
|
||||
return {"kind": "timer", "seconds": condition.seconds}
|
||||
raise TypeError(f"Unsupported condition type: {type(condition)!r}")
|
||||
|
||||
|
||||
def _target_to_suspend_payload(target: WaitCondition | AnyOfCondition) -> dict[str, Any]:
|
||||
if isinstance(target, AnyOfCondition):
|
||||
return {
|
||||
"kind": "any_of",
|
||||
"any_of": {
|
||||
"conditions": [_condition_to_rust(cond) for cond in target.conditions]
|
||||
},
|
||||
}
|
||||
return {"kind": "condition", "condition": _condition_to_rust(target)}
|
||||
|
||||
|
||||
def _unwrap_resume_input(node_input: Any) -> tuple[Any, dict[str, Any] | None]:
|
||||
if not isinstance(node_input, dict):
|
||||
return node_input, None
|
||||
if "__lg_resume_arg__" not in node_input or "__lg_resume_event__" not in node_input:
|
||||
return node_input, None
|
||||
resume_arg = node_input["__lg_resume_arg__"]
|
||||
resume_event = node_input["__lg_resume_event__"]
|
||||
if isinstance(resume_event, dict):
|
||||
return resume_arg, resume_event
|
||||
return resume_arg, None
|
||||
|
||||
|
||||
def _infer_node_name(node: Callable[..., Any]) -> str:
|
||||
node_name = getattr(node, "__name__", "")
|
||||
if not node_name or node_name == "<lambda>":
|
||||
raise ValueError("Cannot infer node name from anonymous callable")
|
||||
return node_name
|
||||
|
||||
|
||||
def _resolve_target_name(target: Any) -> str:
|
||||
if isinstance(target, str):
|
||||
return target
|
||||
if callable(target):
|
||||
return _infer_node_name(target)
|
||||
raise ValueError(f"Unsupported node target type: {type(target)!r}")
|
||||
|
||||
|
||||
def _invoke_node(node: Callable[..., Any], ctx: Context, node_input: Any, state: Any) -> Any:
|
||||
try:
|
||||
params = list(inspect.signature(node).parameters.values())
|
||||
except (TypeError, ValueError):
|
||||
params = []
|
||||
|
||||
if not params:
|
||||
return node()
|
||||
|
||||
names = [param.name.lower() for param in params]
|
||||
has_ctx = [("ctx" in name or "context" in name) for name in names]
|
||||
has_state = [("state" in name) for name in names]
|
||||
has_input = [("input" in name) for name in names]
|
||||
|
||||
kwargs: dict[str, Any] = {}
|
||||
unresolved = False
|
||||
for idx, param in enumerate(params):
|
||||
if has_ctx[idx]:
|
||||
kwargs[param.name] = ctx
|
||||
elif has_state[idx]:
|
||||
kwargs[param.name] = state
|
||||
elif has_input[idx]:
|
||||
kwargs[param.name] = node_input
|
||||
else:
|
||||
unresolved = True
|
||||
|
||||
if kwargs and not unresolved:
|
||||
return node(**kwargs)
|
||||
|
||||
if len(params) == 1:
|
||||
if has_ctx[0]:
|
||||
return node(ctx)
|
||||
if has_state[0]:
|
||||
return node(state)
|
||||
return node(node_input)
|
||||
|
||||
if len(params) == 2:
|
||||
if has_ctx[0] and has_state[1]:
|
||||
return node(ctx, state)
|
||||
if has_ctx[0] and has_input[1]:
|
||||
return node(ctx, node_input)
|
||||
if has_input[0] and has_state[1]:
|
||||
return node(node_input, state)
|
||||
if has_state[0] and has_input[1]:
|
||||
return node(state, node_input)
|
||||
if has_state[0]:
|
||||
return node(state, node_input)
|
||||
if has_state[1]:
|
||||
return node(node_input, state)
|
||||
if has_ctx[0]:
|
||||
return node(ctx, node_input)
|
||||
return node(node_input, state)
|
||||
|
||||
return node(ctx, node_input, state)
|
||||
|
||||
BIN
Binary file not shown.
@@ -0,0 +1 @@
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -0,0 +1,314 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
# Configure advanced-graph runtime pools for this benchmark run.
|
||||
os.environ["LANGGRAPH_RUN_POOL_SIZE"] = "10"
|
||||
os.environ["LANGGRAPH_NODE_POOL_SIZE"] = "1000"
|
||||
|
||||
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
|
||||
|
||||
RUNS = 100
|
||||
MIDDLE_COUNT = 10
|
||||
SLEEP_SECONDS = 2.0
|
||||
BLOCKING_SECONDS = 0.1
|
||||
STATE_BYTES = 10 * 1024
|
||||
|
||||
|
||||
def make_initial_state() -> dict[str, Any]:
|
||||
return {"payload": "x" * STATE_BYTES, "done": False}
|
||||
|
||||
|
||||
def build_advanced_parallel() -> Any:
|
||||
graph: AdvancedStateGraph[dict[str, Any]] = AdvancedStateGraph(dict)
|
||||
done_channel = "__bench_done_channel"
|
||||
graph.add_async_channel(done_channel, str)
|
||||
|
||||
async def start_node(state: dict[str, Any]) -> Command:
|
||||
_ = state
|
||||
sends = [Send(f"middle_{i}", None) for i in range(MIDDLE_COUNT)]
|
||||
sends.append(Send("end_node", None))
|
||||
return Command(goto=sends)
|
||||
|
||||
async def end_node(ctx: Any, state: dict[str, Any]) -> dict[str, Any]:
|
||||
await ctx.wait_for(channel_condition(done_channel, n=MIDDLE_COUNT))
|
||||
out = dict(state)
|
||||
out["done"] = True
|
||||
return out
|
||||
|
||||
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
|
||||
await ctx.wait_for(timer_condition(seconds=SLEEP_SECONDS))
|
||||
ctx.publish_to_channel(done_channel, "done")
|
||||
|
||||
graph.add_node(f"middle_{i}", middle_node)
|
||||
graph.add_finish_node(end_node)
|
||||
return graph.compile()
|
||||
|
||||
|
||||
def build_advanced_sequential() -> Any:
|
||||
graph: AdvancedStateGraph[dict[str, Any]] = AdvancedStateGraph(dict)
|
||||
|
||||
async def start_node(state: dict[str, Any]) -> Command:
|
||||
_ = state
|
||||
return Command(goto=Send("middle_0", None))
|
||||
|
||||
async def end_node(state: dict[str, Any]) -> dict[str, Any]:
|
||||
out = dict(state)
|
||||
out["done"] = True
|
||||
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
|
||||
await ctx.wait_for(timer_condition(seconds=SLEEP_SECONDS))
|
||||
return Command(goto=Send(target, None))
|
||||
|
||||
return middle_node
|
||||
|
||||
for i in range(MIDDLE_COUNT):
|
||||
next_name = "end_node" if i == MIDDLE_COUNT - 1 else f"middle_{i+1}"
|
||||
graph.add_node(f"middle_{i}", make_middle(next_name))
|
||||
graph.add_finish_node(end_node)
|
||||
return graph.compile()
|
||||
|
||||
|
||||
def build_stategraph_parallel() -> Any:
|
||||
graph = StateGraph(dict)
|
||||
|
||||
async def start_node(state: dict[str, Any]) -> None:
|
||||
_ = state
|
||||
|
||||
async def end_node(state: dict[str, Any]) -> dict[str, Any]:
|
||||
out = dict(state)
|
||||
out["done"] = True
|
||||
return out
|
||||
|
||||
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
|
||||
await asyncio.sleep(SLEEP_SECONDS)
|
||||
|
||||
graph.add_node(f"middle_{i}", middle_node)
|
||||
graph.add_node("end_node", end_node)
|
||||
|
||||
graph.add_edge(START, "start_node")
|
||||
for i in range(MIDDLE_COUNT):
|
||||
graph.add_edge("start_node", f"middle_{i}")
|
||||
graph.add_edge(f"middle_{i}", "end_node")
|
||||
graph.add_edge("end_node", END)
|
||||
return graph.compile()
|
||||
|
||||
|
||||
def build_stategraph_sequential() -> Any:
|
||||
graph = StateGraph(dict)
|
||||
|
||||
async def start_node(state: dict[str, Any]) -> None:
|
||||
_ = state
|
||||
|
||||
async def end_node(state: dict[str, Any]) -> dict[str, Any]:
|
||||
out = dict(state)
|
||||
out["done"] = True
|
||||
return out
|
||||
|
||||
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
|
||||
await asyncio.sleep(SLEEP_SECONDS)
|
||||
|
||||
graph.add_node(f"middle_{i}", middle_node)
|
||||
graph.add_node("end_node", end_node)
|
||||
|
||||
graph.add_edge(START, "start_node")
|
||||
graph.add_edge("start_node", "middle_0")
|
||||
for i in range(MIDDLE_COUNT - 1):
|
||||
graph.add_edge(f"middle_{i}", f"middle_{i+1}")
|
||||
graph.add_edge(f"middle_{MIDDLE_COUNT - 1}", "end_node")
|
||||
graph.add_edge("end_node", END)
|
||||
return graph.compile()
|
||||
|
||||
|
||||
def build_advanced_parallel_blocking() -> Any:
|
||||
graph: AdvancedStateGraph[dict[str, Any]] = AdvancedStateGraph(dict)
|
||||
done_channel = "__bench_done_channel_blocking"
|
||||
graph.add_async_channel(done_channel, str)
|
||||
|
||||
async def start_node(state: dict[str, Any]) -> Command:
|
||||
_ = state
|
||||
sends = [Send(f"middle_blocking_{i}", None) for i in range(MIDDLE_COUNT)]
|
||||
sends.append(Send("end_node_blocking", None))
|
||||
return Command(goto=sends)
|
||||
|
||||
async def end_node_blocking(ctx: Any, state: dict[str, Any]) -> dict[str, Any]:
|
||||
await ctx.wait_for(channel_condition(done_channel, n=MIDDLE_COUNT))
|
||||
out = dict(state)
|
||||
out["done"] = True
|
||||
return out
|
||||
|
||||
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:
|
||||
_ = idx
|
||||
_ = state
|
||||
time.sleep(BLOCKING_SECONDS)
|
||||
ctx.publish_to_channel(done_channel, "done")
|
||||
|
||||
graph.add_node(f"middle_blocking_{i}", middle_blocking)
|
||||
graph.add_finish_node(end_node_blocking)
|
||||
return graph.compile()
|
||||
|
||||
|
||||
def build_advanced_sequential_blocking() -> Any:
|
||||
graph: AdvancedStateGraph[dict[str, Any]] = AdvancedStateGraph(dict)
|
||||
|
||||
async def start_node(state: dict[str, Any]) -> Command:
|
||||
_ = state
|
||||
return Command(goto=Send("middle_blocking_seq_0", None))
|
||||
|
||||
async def end_node_blocking_seq(state: dict[str, Any]) -> dict[str, Any]:
|
||||
out = dict(state)
|
||||
out["done"] = True
|
||||
return out
|
||||
|
||||
graph.add_entry_node(start_node)
|
||||
|
||||
def make_middle(target: str):
|
||||
async def middle_blocking_seq(state: dict[str, Any]) -> Command:
|
||||
_ = state
|
||||
time.sleep(BLOCKING_SECONDS)
|
||||
return Command(goto=Send(target, None))
|
||||
|
||||
return middle_blocking_seq
|
||||
|
||||
for i in range(MIDDLE_COUNT):
|
||||
next_name = (
|
||||
"end_node_blocking_seq"
|
||||
if i == MIDDLE_COUNT - 1
|
||||
else f"middle_blocking_seq_{i+1}"
|
||||
)
|
||||
graph.add_node(f"middle_blocking_seq_{i}", make_middle(next_name))
|
||||
graph.add_finish_node(end_node_blocking_seq)
|
||||
return graph.compile()
|
||||
|
||||
|
||||
def build_stategraph_parallel_blocking() -> Any:
|
||||
graph = StateGraph(dict)
|
||||
|
||||
async def start_node(state: dict[str, Any]) -> None:
|
||||
_ = state
|
||||
|
||||
async def end_node(state: dict[str, Any]) -> dict[str, Any]:
|
||||
out = dict(state)
|
||||
out["done"] = True
|
||||
return out
|
||||
|
||||
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
|
||||
time.sleep(BLOCKING_SECONDS)
|
||||
|
||||
graph.add_node(f"middle_blocking_{i}", middle_blocking)
|
||||
graph.add_node("end_node", end_node)
|
||||
|
||||
graph.add_edge(START, "start_node")
|
||||
for i in range(MIDDLE_COUNT):
|
||||
graph.add_edge("start_node", f"middle_blocking_{i}")
|
||||
graph.add_edge(f"middle_blocking_{i}", "end_node")
|
||||
graph.add_edge("end_node", END)
|
||||
return graph.compile()
|
||||
|
||||
|
||||
def build_stategraph_sequential_blocking() -> Any:
|
||||
graph = StateGraph(dict)
|
||||
|
||||
async def start_node(state: dict[str, Any]) -> None:
|
||||
_ = state
|
||||
|
||||
async def end_node(state: dict[str, Any]) -> dict[str, Any]:
|
||||
out = dict(state)
|
||||
out["done"] = True
|
||||
return out
|
||||
|
||||
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
|
||||
time.sleep(BLOCKING_SECONDS)
|
||||
|
||||
graph.add_node(f"middle_blocking_seq_{i}", middle_blocking_seq)
|
||||
graph.add_node("end_node", end_node)
|
||||
|
||||
graph.add_edge(START, "start_node")
|
||||
graph.add_edge("start_node", "middle_blocking_seq_0")
|
||||
for i in range(MIDDLE_COUNT - 1):
|
||||
graph.add_edge(
|
||||
f"middle_blocking_seq_{i}",
|
||||
f"middle_blocking_seq_{i+1}",
|
||||
)
|
||||
graph.add_edge(f"middle_blocking_seq_{MIDDLE_COUNT - 1}", "end_node")
|
||||
graph.add_edge("end_node", END)
|
||||
return graph.compile()
|
||||
|
||||
|
||||
async def run_benchmark(name: str, compiled: Any) -> float:
|
||||
started = time.perf_counter()
|
||||
tasks = [asyncio.create_task(compiled.ainvoke(make_initial_state())) for _ in range(RUNS)]
|
||||
results = await asyncio.gather(*tasks)
|
||||
elapsed = time.perf_counter() - started
|
||||
if not all(item.get("done") is True for item in results):
|
||||
raise RuntimeError(f"{name} produced unfinished runs")
|
||||
return elapsed
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
suites = [
|
||||
("advanced-graph-parallel", build_advanced_parallel()),
|
||||
("advanced-graph-sequential", build_advanced_sequential()),
|
||||
("state-graph-parallel", build_stategraph_parallel()),
|
||||
("state-graph-sequential", build_stategraph_sequential()),
|
||||
("advanced-graph-parallel-blocking", build_advanced_parallel_blocking()),
|
||||
("advanced-graph-sequential-blocking", build_advanced_sequential_blocking()),
|
||||
("state-graph-parallel-blocking", build_stategraph_parallel_blocking()),
|
||||
("state-graph-sequential-blocking", build_stategraph_sequential_blocking()),
|
||||
]
|
||||
print(
|
||||
f"runs={RUNS}, middle_nodes={MIDDLE_COUNT}, sleep={SLEEP_SECONDS}s, "
|
||||
f"blocking_sleep={BLOCKING_SECONDS}s, state_bytes={STATE_BYTES}"
|
||||
)
|
||||
for name, compiled in suites:
|
||||
elapsed = await run_benchmark(name, compiled)
|
||||
print(f"{name}: {elapsed:.3f}s")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
|
||||
@@ -0,0 +1,67 @@
|
||||
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 PrimitiveState(TypedDict):
|
||||
counter: int
|
||||
logs: list[str]
|
||||
done: str | None
|
||||
|
||||
|
||||
async def test_input_and_state_primitives_are_compatible() -> None:
|
||||
graph = AdvancedStateGraph(PrimitiveState)
|
||||
|
||||
async def start_node(state: PrimitiveState) -> Command:
|
||||
state["logs"].append(f"start:counter={state['counter']}")
|
||||
return Command(goto=Send("middle_node", "from_start"))
|
||||
|
||||
async def middle_node(ctx: Context, tool_input: str, state: PrimitiveState) -> Command:
|
||||
state["logs"].append(f"middle:input={tool_input}")
|
||||
return Command(update=state, goto=Send("finish_node", "from_middle"))
|
||||
|
||||
async def finish_node(payload: str, state: PrimitiveState) -> dict[str, object]:
|
||||
state["logs"].append(f"finish:input={payload}")
|
||||
return {
|
||||
"logs": state["logs"],
|
||||
"counter": state["counter"],
|
||||
"done": payload,
|
||||
}
|
||||
|
||||
graph.add_entry_node(start_node)
|
||||
graph.add_node(middle_node)
|
||||
graph.add_finish_node(finish_node)
|
||||
|
||||
result = await graph.compile().ainvoke({"counter": 7, "logs": [], "done": None})
|
||||
assert result["counter"] == 7
|
||||
assert result["done"] == "from_middle"
|
||||
assert result["logs"] == [
|
||||
"start:counter=7",
|
||||
"middle:input=from_start",
|
||||
"finish:input=from_middle",
|
||||
]
|
||||
|
||||
|
||||
async def test_run_ends_without_finish_node() -> None:
|
||||
graph = AdvancedStateGraph(PrimitiveState)
|
||||
|
||||
async def start_node(state: PrimitiveState) -> Command:
|
||||
state["logs"].append("start")
|
||||
return Command(update=state, goto=Send("middle_node", "from_start"))
|
||||
|
||||
async def middle_node(input: str, state: PrimitiveState) -> dict[str, object]:
|
||||
state["logs"].append(f"middle:{input}")
|
||||
return {"counter": state["counter"] + 1, "logs": state["logs"], "done": "stopped"}
|
||||
|
||||
graph.add_entry_node(start_node)
|
||||
graph.add_node(middle_node)
|
||||
|
||||
result = await graph.compile().ainvoke({"counter": 7, "logs": [], "done": None})
|
||||
assert result["counter"] == 8
|
||||
assert result["done"] == "stopped"
|
||||
assert result["logs"] == ["start", "middle:from_start"]
|
||||
|
||||
@@ -0,0 +1,56 @@
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
|
||||
def test_run_pool_size_one_still_allows_parallel_runs() -> None:
|
||||
script = r"""
|
||||
import asyncio
|
||||
import time
|
||||
from typing_extensions import TypedDict
|
||||
from saf_python_sdk.advanced_graph import AdvancedStateGraph, Context, timer_condition
|
||||
from saf_python_sdk.types import Command, Send
|
||||
|
||||
|
||||
class RunState(TypedDict):
|
||||
done: bool
|
||||
|
||||
|
||||
async def wait_node(ctx: Context, _: object, state: RunState) -> Command:
|
||||
await ctx.wait_for(timer_condition(seconds=0.2))
|
||||
return Command(goto=Send("finish_node", None), update=state)
|
||||
|
||||
|
||||
async def finish_node(_: object, state: RunState) -> dict[str, bool]:
|
||||
return {"done": True}
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
graph = AdvancedStateGraph(RunState)
|
||||
graph.add_entry_node(wait_node)
|
||||
graph.add_finish_node(finish_node)
|
||||
compiled = graph.compile()
|
||||
started = time.perf_counter()
|
||||
await asyncio.gather(
|
||||
compiled.ainvoke({"done": False}),
|
||||
compiled.ainvoke({"done": False}),
|
||||
)
|
||||
elapsed = time.perf_counter() - started
|
||||
print(f"{elapsed:.6f}")
|
||||
|
||||
|
||||
asyncio.run(main())
|
||||
"""
|
||||
env = os.environ.copy()
|
||||
env["LANGGRAPH_RUN_POOL_SIZE"] = "1"
|
||||
env.setdefault("LANGGRAPH_NODE_POOL_SIZE", "2")
|
||||
completed = subprocess.run(
|
||||
[sys.executable, "-c", script],
|
||||
env=env,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=True,
|
||||
)
|
||||
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()
|
||||
@@ -0,0 +1,218 @@
|
||||
import asyncio
|
||||
import json
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Literal
|
||||
|
||||
import pytest
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from saf_python_sdk.advanced_graph import (
|
||||
AdvancedStateGraph,
|
||||
Context,
|
||||
any_of,
|
||||
channel_condition,
|
||||
timer_condition,
|
||||
)
|
||||
from langgraph.constants import END, START
|
||||
from langgraph.graph import StateGraph
|
||||
from saf_python_sdk.types import Command, Send
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
|
||||
class MainAgentState(TypedDict):
|
||||
input: str
|
||||
output: list[str]
|
||||
done: str | None
|
||||
|
||||
|
||||
class SubAgentState(TypedDict):
|
||||
input: str
|
||||
output: str
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Decision:
|
||||
type: Literal["end", "sub_agent", "tool"]
|
||||
sub_agent: str | None = None
|
||||
tool: str | None = None
|
||||
complete: str | None = None
|
||||
|
||||
|
||||
class MockLLM:
|
||||
def __init__(self) -> None:
|
||||
self.responses: list[list[Decision]] = []
|
||||
self._idx = 0
|
||||
|
||||
async def ainvoke(self, _: MainAgentState) -> list[Decision]:
|
||||
if self._idx >= len(self.responses):
|
||||
return []
|
||||
response = self.responses[self._idx]
|
||||
self._idx += 1
|
||||
return response
|
||||
|
||||
|
||||
def build_sub_agent() -> Any:
|
||||
# Sub-agent uses the regular/simple StateGraph API.
|
||||
sub_agent = StateGraph(SubAgentState)
|
||||
|
||||
async def research_node(state: SubAgentState) -> dict[str, str]:
|
||||
# Intentionally slower than timer_condition(seconds=1) to validate timer path.
|
||||
await asyncio.sleep(5)
|
||||
return {"output": f"research sub agent completed for: {state['input']}"}
|
||||
|
||||
sub_agent.add_node("research_node", research_node)
|
||||
sub_agent.add_edge(START, "research_node")
|
||||
sub_agent.add_edge("research_node", END)
|
||||
return sub_agent.compile()
|
||||
|
||||
|
||||
def build_main_agent(planner: MockLLM, sub_agent: Any) -> Any:
|
||||
async def llm_node(state: MainAgentState) -> Command:
|
||||
# Planner decides whether to call a tool, spawn a sub-agent, or finish.
|
||||
decisions = await planner.ainvoke(state)
|
||||
sends: list[Send] = []
|
||||
for decision in decisions:
|
||||
if decision.type == "end":
|
||||
# NOTE: this can be simplified further in the future with a dedicated
|
||||
# complete primitive, instead of routing to a finish node manually.
|
||||
return Command(
|
||||
goto=Send(
|
||||
order_food_node,
|
||||
decision.complete or "order flow completed",
|
||||
)
|
||||
)
|
||||
if decision.type == "sub_agent" and decision.sub_agent:
|
||||
sends.append(Send("sub_agent_node", decision.sub_agent))
|
||||
if decision.type == "tool" and decision.tool:
|
||||
sends.append(Send("tool_node", decision.tool))
|
||||
# Keep the main loop responsive: wait for one inbound message and continue.
|
||||
sends.append(Send("wait_node", None))
|
||||
return Command(goto=sends)
|
||||
|
||||
async def wait_node(ctx: Context, state: MainAgentState) -> Command:
|
||||
# Lightweight interrupt: only this node blocks for the next relevant signal.
|
||||
event = await ctx.wait_for(
|
||||
any_of(
|
||||
channel_condition("tool_completion_channel"),
|
||||
channel_condition("subagent_completion_channel"),
|
||||
channel_condition("user_input_channel"),
|
||||
timer_condition(seconds=1),
|
||||
)
|
||||
)
|
||||
if event["condition"] == "channel":
|
||||
channel = event["channel"]
|
||||
payload = event["value"]
|
||||
if channel == "tool_completion_channel":
|
||||
state["output"].append(f"tool: {payload}")
|
||||
elif channel == "subagent_completion_channel":
|
||||
state["output"].append(f"sub_agent: {payload}")
|
||||
elif channel == "user_input_channel":
|
||||
state["output"].append(f"user_input: {payload}")
|
||||
# State changed -> ask planner what to do next.
|
||||
return Command(update=state, goto=Send("llm_node", None))
|
||||
else:
|
||||
state["output"].append("timer: no updates yet")
|
||||
# No meaningful state change -> keep waiting without calling planner.
|
||||
return Command(update=state, goto=Send("wait_node", None))
|
||||
|
||||
async def tool_node(ctx: Context, tool_input: str) -> None:
|
||||
await asyncio.sleep(0.1)
|
||||
# Fire-and-forget style completion: publish result to inbox and exit.
|
||||
# (i.e., just complete without explicitly going to a next node)
|
||||
ctx.publish_to_channel(
|
||||
"tool_completion_channel",
|
||||
f"tool completed for: {tool_input}",
|
||||
)
|
||||
|
||||
async def sub_agent_node(ctx: Context, sub_agent_input: str) -> None:
|
||||
# Sub-agent remains a regular StateGraph, compiled independently.
|
||||
sub_agent_output = await sub_agent.ainvoke(
|
||||
{"input": sub_agent_input, "output": ""}
|
||||
)
|
||||
# Same pattern as tool node: publish result and complete current node.
|
||||
ctx.publish_to_channel(
|
||||
"subagent_completion_channel",
|
||||
sub_agent_output["output"],
|
||||
)
|
||||
|
||||
async def order_food_node(input: str, state: MainAgentState) -> dict[str, Any]:
|
||||
complete_message = input
|
||||
return {
|
||||
"done": complete_message,
|
||||
"output": [*state["output"], f"order_food: {complete_message}"],
|
||||
}
|
||||
|
||||
advanced_flow = AdvancedStateGraph(MainAgentState)
|
||||
# Default behavior is an unbounded async channel like Rust channel
|
||||
advanced_flow.add_async_channel("tool_completion_channel", str)
|
||||
advanced_flow.add_async_channel("subagent_completion_channel", str)
|
||||
advanced_flow.add_async_channel("user_input_channel", str)
|
||||
# nodes are the same as in the regular StateGraph API
|
||||
advanced_flow.add_entry_node(llm_node)
|
||||
advanced_flow.add_node(wait_node)
|
||||
advanced_flow.add_node(tool_node)
|
||||
advanced_flow.add_node(sub_agent_node)
|
||||
advanced_flow.add_finish_node(order_food_node)
|
||||
|
||||
return advanced_flow.compile()
|
||||
|
||||
|
||||
async def test_async_sub_graph() -> None:
|
||||
llm = MockLLM()
|
||||
sub_agent = build_sub_agent()
|
||||
main_agent = build_main_agent(llm, sub_agent)
|
||||
|
||||
llm.responses = [
|
||||
[
|
||||
# First planner pass triggers one slow sub-agent.
|
||||
Decision(type="sub_agent", sub_agent="research lunch options"),
|
||||
Decision(type="tool", tool="slack_tool"),
|
||||
],
|
||||
# After user input.
|
||||
[],
|
||||
# After tool completion.
|
||||
[],
|
||||
# After first sub-agent completion, planner decides to run second research.
|
||||
[Decision(type="sub_agent", sub_agent="find vegetarian fallback")],
|
||||
# After second sub-agent completion, planner decides to end.
|
||||
[Decision(type="end", complete="order submitted")],
|
||||
]
|
||||
|
||||
handler = await main_agent.astart(
|
||||
{"input": "help me get something for lunch", "output": [], "done": None}
|
||||
)
|
||||
|
||||
# External input can be injected while graph execution is in progress.
|
||||
await asyncio.sleep(0.01)
|
||||
await handler.apublish_to_channel("user_input_channel", "No spicy food please")
|
||||
result = await handler.aresult()
|
||||
|
||||
assert result["input"] == "help me get something for lunch"
|
||||
assert result["done"] == "order submitted"
|
||||
|
||||
output = result["output"]
|
||||
assert output.count("timer: no updates yet") >= 3
|
||||
assert "user_input: No spicy food please" in output
|
||||
assert "tool: tool completed for: slack_tool" in output
|
||||
assert (
|
||||
"sub_agent: research sub agent completed for: research lunch options" in output
|
||||
)
|
||||
assert (
|
||||
"sub_agent: research sub agent completed for: find vegetarian fallback"
|
||||
in output
|
||||
)
|
||||
assert output[-1] == "order_food: order submitted"
|
||||
|
||||
first_sub_idx = output.index(
|
||||
"sub_agent: research sub agent completed for: research lunch options"
|
||||
)
|
||||
second_sub_idx = output.index(
|
||||
"sub_agent: research sub agent completed for: find vegetarian fallback"
|
||||
)
|
||||
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)
|
||||
|
||||
print(json.dumps(result, ensure_ascii=False, indent=2))
|
||||
|
||||
@@ -0,0 +1,124 @@
|
||||
import asyncio
|
||||
from dataclasses import dataclass
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from saf_python_sdk.advanced_graph import AdvancedStateGraph, CompiledGraphEngine
|
||||
from saf_python_sdk.types import Command, Send
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
|
||||
@dataclass
|
||||
class DataClassPayload:
|
||||
value: int
|
||||
|
||||
|
||||
class PydanticPayload(BaseModel):
|
||||
value: int
|
||||
|
||||
|
||||
class InnerTypedDict(TypedDict):
|
||||
flag: bool
|
||||
n: int
|
||||
|
||||
|
||||
@dataclass
|
||||
class UpdateElisionState:
|
||||
x: int
|
||||
dc: DataClassPayload
|
||||
model: PydanticPayload
|
||||
td: InnerTypedDict
|
||||
obj: dict[str, int]
|
||||
items: list[int]
|
||||
|
||||
|
||||
def _initial_state() -> UpdateElisionState:
|
||||
return UpdateElisionState(
|
||||
x=0,
|
||||
dc=DataClassPayload(0),
|
||||
model=PydanticPayload(value=0),
|
||||
td={"flag": False, "n": 0},
|
||||
obj={"n": 0},
|
||||
items=[0],
|
||||
)
|
||||
|
||||
|
||||
async def test_noop_slow_update_does_not_override_fast_update() -> None:
|
||||
graph: AdvancedStateGraph[UpdateElisionState] = AdvancedStateGraph(UpdateElisionState)
|
||||
|
||||
async def start_node(state: UpdateElisionState) -> Command:
|
||||
return Command(goto=[Send("fast_node", None), Send("slow_node", None)])
|
||||
|
||||
async def fast_node(state: UpdateElisionState) -> UpdateElisionState:
|
||||
state.x = 1
|
||||
state.dc.value = 1
|
||||
state.model.value = 1
|
||||
state.td["flag"] = True
|
||||
state.td["n"] = 1
|
||||
state.obj["n"] = 1
|
||||
state.items.append(1)
|
||||
return state
|
||||
|
||||
async def slow_node(state: UpdateElisionState) -> UpdateElisionState:
|
||||
await asyncio.sleep(0.1)
|
||||
return state
|
||||
|
||||
graph.add_entry_node(start_node)
|
||||
graph.add_node(fast_node)
|
||||
graph.add_finish_node(slow_node)
|
||||
|
||||
compiled: CompiledGraphEngine[UpdateElisionState] = graph.compile()
|
||||
initial_state: UpdateElisionState = _initial_state()
|
||||
result: UpdateElisionState = await compiled.ainvoke(initial_state)
|
||||
assert result.x == 1
|
||||
assert result.dc == DataClassPayload(1)
|
||||
assert result.model.value == 1
|
||||
assert result.td == {"flag": True, "n": 1}
|
||||
assert result.obj == {"n": 1}
|
||||
assert result.items == [0, 1]
|
||||
|
||||
|
||||
async def test_changed_slow_update_overrides_fast_update() -> None:
|
||||
graph: AdvancedStateGraph[UpdateElisionState] = AdvancedStateGraph(UpdateElisionState)
|
||||
|
||||
async def start_node(state: UpdateElisionState) -> Command:
|
||||
return Command(goto=[Send("fast_node", None), Send("slow_node", None)])
|
||||
|
||||
async def fast_node(state: UpdateElisionState) -> UpdateElisionState:
|
||||
state.x = 1
|
||||
state.dc.value = 1
|
||||
state.model.value = 1
|
||||
state.td["flag"] = True
|
||||
state.td["n"] = 1
|
||||
state.obj["n"] = 1
|
||||
state.items.append(1)
|
||||
return state
|
||||
|
||||
async def slow_node(state: UpdateElisionState) -> UpdateElisionState:
|
||||
await asyncio.sleep(0.1)
|
||||
state.x = 2
|
||||
state.dc.value = 2
|
||||
state.model.value = 2
|
||||
state.td["flag"] = False
|
||||
state.td["n"] = 2
|
||||
state.obj["n"] = 2
|
||||
state.items.append(2)
|
||||
return state
|
||||
|
||||
graph.add_entry_node(start_node)
|
||||
graph.add_node(fast_node)
|
||||
graph.add_finish_node(slow_node)
|
||||
|
||||
compiled: CompiledGraphEngine[UpdateElisionState] = graph.compile()
|
||||
initial_state: UpdateElisionState = _initial_state()
|
||||
result: UpdateElisionState = await compiled.ainvoke(initial_state)
|
||||
assert result.x == 2
|
||||
assert result.dc == DataClassPayload(2)
|
||||
assert result.model.value == 2
|
||||
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