diff --git a/langgraph-go/advancedgraph/graph.go b/langgraph-go/advancedgraph/graph.go new file mode 100644 index 000000000..84ade22e8 --- /dev/null +++ b/langgraph-go/advancedgraph/graph.go @@ -0,0 +1,115 @@ +package advancedgraph + +import ( + "fmt" +) + +type NodeFunc func(ctx *Context, arg any) (Command, error) + +type AdvancedStateGraph struct { + nodes map[string]NodeFunc + asyncChannels []string + entryPoint string + finishPoint string +} + +func NewAdvancedStateGraph() *AdvancedStateGraph { + return &AdvancedStateGraph{ + nodes: make(map[string]NodeFunc), + } +} + +func (g *AdvancedStateGraph) AddNode(name string, fn NodeFunc) { + g.nodes[name] = fn +} + +func (g *AdvancedStateGraph) AddAsyncChannel(name string) { + g.asyncChannels = append(g.asyncChannels, name) +} + +func (g *AdvancedStateGraph) SetEntryPoint(name string) { + g.entryPoint = name +} + +func (g *AdvancedStateGraph) SetFinishPoint(name string) { + g.finishPoint = name +} + +func (g *AdvancedStateGraph) Compile() *CompiledGraph { + return &CompiledGraph{ + nodes: g.nodes, + asyncChannels: g.asyncChannels, + entryPoint: g.entryPoint, + finishPoint: g.finishPoint, + } +} + +type CompiledGraph struct { + nodes map[string]NodeFunc + asyncChannels []string + entryPoint string + finishPoint string +} + +type Context struct { + engine *RustEngine +} + +func (c *Context) WaitFor(cond AnyOfCondition) (WaitEvent, error) { + return c.engine.WaitAnyOf(cond) +} + +func (c *Context) PublishToChannel(channel string, value any) error { + return c.engine.Publish(channel, value) +} + +type Handler struct { + engine *RustEngine + done chan resultOrErr +} + +type resultOrErr struct { + state map[string]any + err error +} + +func (h *Handler) APublishToChannel(channel string, value any) error { + return h.engine.Publish(channel, value) +} + +func (h *Handler) AResult() (map[string]any, error) { + res := <-h.done + return res.state, res.err +} + +func (g *CompiledGraph) AStart(initialState map[string]any) (*Handler, error) { + engine := NewRustEngine() + for _, ch := range g.asyncChannels { + if err := engine.AddAsyncChannel(ch); err != nil { + return nil, err + } + } + + handler := &Handler{ + engine: engine, + done: make(chan resultOrErr, 1), + } + go func() { + defer engine.Close() + state, err := engine.RunGraph( + g.entryPoint, + g.finishPoint, + initialState, + func(node string, arg any, _ map[string]any) (Command, error) { + fn, ok := g.nodes[node] + if !ok { + return Command{}, fmt.Errorf("unknown node `%s`", node) + } + return fn(&Context{engine: engine}, arg) + }, + ) + handler.done <- resultOrErr{state: state, err: err} + close(handler.done) + }() + return handler, nil +} diff --git a/langgraph-go/advancedgraph/rust_engine.go b/langgraph-go/advancedgraph/rust_engine.go new file mode 100644 index 000000000..cccd67bf5 --- /dev/null +++ b/langgraph-go/advancedgraph/rust_engine.go @@ -0,0 +1,241 @@ +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 +extern char* goNodeCallback(void* user_data, char* node, char* arg_json, char* state_json); +*/ +import "C" + +import ( + "encoding/json" + "fmt" + "runtime/cgo" + "unsafe" +) + +type RustEngine struct { + ptr *C.Engine +} + +type runGraphCallbackCtx struct { + exec func(node string, arg any, state map[string]any) (Command, error) +} + +//export goNodeCallback +func goNodeCallback(userData unsafe.Pointer, node *C.char, argJSON *C.char, stateJSON *C.char) *C.char { + handle := cgo.Handle(uintptr(userData)) + ctx, ok := handle.Value().(*runGraphCallbackCtx) + if !ok { + return cCallbackEnvelopeError("invalid callback context") + } + + nodeName := C.GoString(node) + + var arg any + if err := json.Unmarshal([]byte(C.GoString(argJSON)), &arg); 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)) + } + arg = coerceJSONValue(arg) + 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, arg, state) + if err != nil { + return cCallbackEnvelopeError(err.Error()) + } + + sends := make([]map[string]any, 0, len(cmd.Goto)) + for _, send := range cmd.Goto { + sends = append(sends, map[string]any{ + "node": send.Node, + "arg": send.Arg, + }) + } + 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) 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, + initialState map[string]any, + exec func(node string, arg 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) + } + centry := C.CString(entryPoint) + cfinish := C.CString(finishPoint) + cinitial := C.CString(string(initialJSON)) + defer C.free(unsafe.Pointer(centry)) + defer C.free(unsafe.Pointer(cfinish)) + defer C.free(unsafe.Pointer(cinitial)) + + handle := cgo.NewHandle(&runGraphCallbackCtx{exec: exec}) + defer handle.Delete() + + resp := C.rc_run_graph_json( + e.ptr, + centry, + cfinish, + cinitial, + unsafe.Pointer(uintptr(handle)), + (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 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 +} diff --git a/langgraph-go/advancedgraph/types.go b/langgraph-go/advancedgraph/types.go new file mode 100644 index 000000000..8748b7be5 --- /dev/null +++ b/langgraph-go/advancedgraph/types.go @@ -0,0 +1,70 @@ +package advancedgraph + +import "encoding/json" + +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 string + Arg any +} + +type Command struct { + Update map[string]any + Goto []Send +} + +func DecodeString(raw json.RawMessage) string { + var s string + _ = json.Unmarshal(raw, &s) + return s +} diff --git a/langgraph-go/go.mod b/langgraph-go/go.mod new file mode 100644 index 000000000..d297fbf70 --- /dev/null +++ b/langgraph-go/go.mod @@ -0,0 +1,4 @@ +module github.com/langchain-ai/langgraph/langgraph-go + +go 1.25 + diff --git a/langgraph-go/tests/sub_agents_test.go b/langgraph-go/tests/sub_agents_test.go new file mode 100644 index 000000000..ed1a72362 --- /dev/null +++ b/langgraph-go/tests/sub_agents_test.go @@ -0,0 +1,195 @@ +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 mockPlanner struct { + responses [][]decision + i int +} + +func (m *mockPlanner) ainvoke() []decision { + if m.i >= len(m.responses) { + return []decision{} + } + resp := m.responses[m.i] + m.i++ + return resp +} + +func TestSubAgentsEquivalentFlow(t *testing.T) { + planner := &mockPlanner{ + 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"}}, + }, + } + + graph := ag.NewAdvancedStateGraph() + graph.AddAsyncChannel("tool_completion_channel") + graph.AddAsyncChannel("subagent_completion_channel") + graph.AddAsyncChannel("user_input_channel") + + graph.AddNode("llm_node", func(ctx *ag.Context, arg any) (ag.Command, error) { + state := arg.(map[string]any) + decisions := planner.ainvoke() + sends := make([]ag.Send, 0, 4) + for _, d := range decisions { + if d.Type == "end" { + return ag.Command{ + Goto: []ag.Send{ + { + Node: "order_food_node", + Arg: map[string]any{ + "state": state, + "complete": d.Complete, + }, + }, + }, + }, nil + } + if d.Type == "sub_agent" { + sends = append(sends, ag.Send{Node: "sub_agent_node", Arg: d.SubAgent}) + } + if d.Type == "tool" { + sends = append(sends, ag.Send{Node: "tool_node", Arg: d.Tool}) + } + } + sends = append(sends, ag.Send{Node: "wait_node", Arg: state}) + return ag.Command{Goto: sends}, nil + }) + + graph.AddNode("wait_node", func(ctx *ag.Context, arg any) (ag.Command, error) { + state := arg.(map[string]any) + 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 := state["output"].([]string) + 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: "llm_node", Arg: state}}}, nil + } + + output = append(output, "timer: no updates yet") + state["output"] = output + return ag.Command{Goto: []ag.Send{{Node: "wait_node", Arg: state}}}, nil + }) + + graph.AddNode("tool_node", func(ctx *ag.Context, arg any) (ag.Command, error) { + toolInput := arg.(string) + time.Sleep(100 * time.Millisecond) + err := ctx.PublishToChannel("tool_completion_channel", "tool completed for: "+toolInput) + return ag.Command{}, err + }) + + graph.AddNode("sub_agent_node", func(ctx *ag.Context, arg any) (ag.Command, error) { + subInput := arg.(string) + time.Sleep(5 * time.Second) + err := ctx.PublishToChannel( + "subagent_completion_channel", + "research sub agent completed for: "+subInput, + ) + return ag.Command{}, err + }) + + graph.AddNode("order_food_node", func(ctx *ag.Context, arg any) (ag.Command, error) { + payload := arg.(map[string]any) + state := payload["state"].(map[string]any) + complete := payload["complete"].(string) + output := state["output"].([]string) + output = append(output, "order_food: "+complete) + state["output"] = output + state["done"] = complete + return ag.Command{Update: state}, nil + }) + + graph.SetEntryPoint("llm_node") + graph.SetFinishPoint("order_food_node") + + handler, err := graph.Compile().AStart( + map[string]any{ + "input": "help me get something for lunch", + "output": []string{}, + "done": nil, + }, + ) + if err != nil { + t.Fatalf("start failed: %v", err) + } + + time.Sleep(10 * time.Millisecond) + if err := handler.APublishToChannel("user_input_channel", "No spicy food please"); err != nil { + t.Fatalf("publish failed: %v", err) + } + + result, err := handler.AResult() + if err != nil { + t.Fatalf("result failed: %v", err) + } + + output := result["output"].([]string) + 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]) + } +} + diff --git a/libs/langgraph/langgraph/advanced_graph/state.py b/libs/langgraph/langgraph/advanced_graph/state.py index 859032fc2..f192a6833 100644 --- a/libs/langgraph/langgraph/advanced_graph/state.py +++ b/libs/langgraph/langgraph/advanced_graph/state.py @@ -2,11 +2,13 @@ from __future__ import annotations import asyncio import inspect -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Callable, Coroutine, Sequence from dataclasses import dataclass from datetime import timedelta from typing import Any, Generic, TypeVar, cast +from langgraph_rust_core import PyRustEngine # type: ignore[import-untyped] + from langgraph.types import Command, Send StateT = TypeVar("StateT") @@ -191,44 +193,30 @@ class _GraphEngineRun: self._nodes = nodes self._entry_point = entry_point self._finish_point = finish_point - self._async_channels: dict[str, asyncio.Queue[Any]] = { - name: asyncio.Queue() for name, _spec in async_channel_specs.items() - } + 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.context = Context(self) async def run(self, initial_state: StateT) -> StateT: - self._state = initial_state - self._schedule(Send(self._entry_point, initial_state)) - try: - while self._tasks and not self._finished: - done, _ = await asyncio.wait( - self._tasks, return_when=asyncio.FIRST_COMPLETED - ) - for task in done: - self._tasks.remove(task) - exc = task.exception() - if exc is not None: - await self._cancel_all_tasks() - raise exc - sends = task.result() - for send in sends: - self._schedule(send) - if self._finished: - await self._cancel_all_tasks() - return cast(StateT, self._state) - finally: - await self._cancel_all_tasks() + result_obj = await asyncio.to_thread( + self._rust_engine.run_graph_py, + self._entry_point, + self._finish_point, + initial_state, + self._execute_node_for_rust, + ) + self._state = result_obj + return cast(StateT, self._state) async def publish(self, channel: str, value: Any) -> None: - queue = self._get_async_channel(channel) - await queue.put(value) + await asyncio.to_thread(self._publish_sync, channel, value) def publish_nowait(self, channel: str, value: Any) -> None: - queue = self._get_async_channel(channel) - queue.put_nowait(value) + self._publish_sync(channel, value) async def wait_for(self, target: WaitCondition | AnyOfCondition) -> Any: if isinstance(target, ChannelCondition): @@ -239,8 +227,7 @@ class _GraphEngineRun: "value": value, } if isinstance(target, TimerCondition): - await asyncio.sleep(target.seconds) - return {"condition": "timer", "seconds": target.seconds} + return await asyncio.to_thread(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}") @@ -248,91 +235,47 @@ class _GraphEngineRun: async def _wait_for_channel_values(self, channel: str, n: int) -> Any: if n < 1: raise ValueError("wait_for count `n` must be >= 1") - queue = self._get_async_channel(channel) - if n == 1: - return await queue.get() - values: list[Any] = [] - for _ in range(n): - values.append(await queue.get()) - return values + event = await asyncio.to_thread(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] + } + return await asyncio.to_thread(self._rust_engine.wait_any_of_obj, payload) - tasks = [ - asyncio.create_task(self.wait_for(inner_condition)) - for inner_condition in condition.conditions - ] - done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) - for task in pending: - task.cancel() - await asyncio.gather(*pending, return_exceptions=True) - first = done.pop() - return first.result() + def _publish_sync(self, channel: str, value: Any) -> None: + self._rust_engine.publish_obj(channel, value) - def _get_async_channel(self, channel: str) -> asyncio.Queue[Any]: - if channel not in self._async_channels: - raise ValueError(f"Unknown channel `{channel}`") - return self._async_channels[channel] + def _execute_node_for_rust(self, node_name: str, arg: Any, state: Any) -> dict[str, Any]: - def _schedule(self, send: Send) -> None: - if self._finished: - return - task: asyncio.Task[list[Send]] = asyncio.create_task(self._execute_send(send)) - self._tasks.add(task) - - async def _cancel_all_tasks(self) -> None: - if not self._tasks: - return - to_cancel = list(self._tasks) - for task in to_cancel: - task.cancel() - await asyncio.gather(*to_cancel, return_exceptions=True) - self._tasks.clear() - - async def _execute_send(self, send: Send) -> list[Send]: - node_name = _resolve_target_name(send.node) if node_name not in self._nodes: raise ValueError(f"Unknown node `{node_name}`") node = self._nodes[node_name] - - result = _invoke_node(node, self.context, send.arg) + result = _invoke_node(node, self.context, arg) if inspect.isawaitable(result): - result = await result + result = asyncio.run(cast(Coroutine[Any, Any, Any], result)) if isinstance(result, Command): - self._apply_update(result.update) - next_sends = _normalize_goto(result.goto, default_arg=self._state) + update = result.update + sends = _normalize_goto(result.goto, default_arg=state) else: - self._apply_update(result) - next_sends = _normalize_result_to_sends(result, default_arg=self._state) + update = result + sends = _normalize_result_to_sends(result, default_arg=state) - if node_name == self._finish_point: - self._finished = True - return [] - return next_sends - - def _apply_update(self, update: Any) -> None: - if update is None: - return - if isinstance(update, Mapping): - if isinstance(self._state, Mapping): - # Keep semantics simple: in-place update for mapping-like state. - cast(dict[str, Any], self._state).update(update) - return - if isinstance(update, Sequence) and not isinstance(update, (str, bytes)): - pairs = list(update) - if all( - isinstance(item, tuple) and len(item) == 2 and isinstance(item[0], str) - for item in pairs - ): - if isinstance(self._state, Mapping): - cast(dict[str, Any], self._state).update( - cast(dict[str, Any], pairs) - ) - return + if update is None and isinstance(arg, dict): + # Preserve in-place state mutations for prototype nodes like wait_node. + update = arg + return { + "update": update, + "sends": [ + {"node": _resolve_target_name(send.node), "arg": send.arg} + for send in sends + ], + } def _normalize_result_to_sends(result: Any, *, default_arg: Any) -> list[Send]: if result is None: @@ -417,6 +360,14 @@ def any_of(*conditions: WaitCondition) -> AnyOfCondition: 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 _infer_node_name(node: Callable[..., Any]) -> str: node_name = getattr(node, "__name__", "") if not node_name or node_name == "": diff --git a/libs/langgraph/tests/advanced-graph/test_sub_agents.py b/libs/langgraph/tests/advanced-graph/test_sub_agents.py index 4f575fb3b..7e6107f9c 100644 --- a/libs/langgraph/tests/advanced-graph/test_sub_agents.py +++ b/libs/langgraph/tests/advanced-graph/test_sub_agents.py @@ -202,7 +202,8 @@ async def test_async_sub_graph() -> None: "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 + "sub_agent: research sub agent completed for: find vegetarian fallback" + in output ) assert output[-1] == "order_food: order submitted" @@ -216,4 +217,5 @@ async def test_async_sub_graph() -> None: assert first_sub_idx < second_sub_idx < order_food_idx assert planner._idx == len(planner.responses) import json + print(json.dumps(result, ensure_ascii=False, indent=2)) diff --git a/rust-core/Cargo.lock b/rust-core/Cargo.lock new file mode 100644 index 000000000..10132774b --- /dev/null +++ b/rust-core/Cargo.lock @@ -0,0 +1,310 @@ +# 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", +] + +[[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 = "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 = "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" diff --git a/rust-core/Cargo.toml b/rust-core/Cargo.toml new file mode 100644 index 000000000..79c508e66 --- /dev/null +++ b/rust-core/Cargo.toml @@ -0,0 +1,20 @@ +[package] +name = "langgraph_rust_core" +version = "0.1.0" +edition = "2021" + +[lib] +name = "langgraph_rust_core" +crate-type = ["cdylib", "rlib"] + +[features] +default = ["python-bindings"] +python-bindings = ["dep:pyo3"] + +[dependencies] +pyo3 = { version = "0.23.5", features = ["extension-module"], optional = true } +serde = { version = "1.0", features = ["derive"] } +serde_json = "1.0" +parking_lot = "0.12" +libc = "0.2" + diff --git a/rust-core/include/langgraph_rust_core.h b/rust-core/include/langgraph_rust_core.h new file mode 100644 index 000000000..81465a338 --- /dev/null +++ b/rust-core/include/langgraph_rust_core.h @@ -0,0 +1,38 @@ +#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)( + void* 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_run_graph_json( + Engine* ptr, + const char* entry_point, + const char* finish_point, + const char* initial_state_json, + void* user_data, + rc_node_callback_t callback +); + +void rc_string_free(char* ptr); + +#ifdef __cplusplus +} +#endif + +#endif + diff --git a/rust-core/src/engine.rs b/rust-core/src/engine.rs new file mode 100644 index 000000000..c6314528e --- /dev/null +++ b/rust-core/src/engine.rs @@ -0,0 +1,314 @@ +use parking_lot::{Condvar, Mutex}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use std::collections::{HashMap, VecDeque}; +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}; + +#[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, +} + +#[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 }, +} + +pub struct SendPayload { + pub node: String, + pub arg: A, +} + +pub struct NodeExecResult { + pub update: Option, + pub sends: Vec>, +} + +type Task = Box; + +struct ThreadPool { + tx: mpsc::Sender, + _workers: Vec>, +} + +impl ThreadPool { + fn new(size: usize, label: &str) -> Self { + let (tx, rx) = mpsc::channel::(); + 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(&self, task: F) -> Result<(), String> + where + F: FnOnce() + Send + 'static, + { + self.tx + .send(Box::new(task)) + .map_err(|e| format!("thread-pool send failed: {e}")) + } +} + +pub fn node_pool_execute(task: F) -> Result<(), String> +where + F: FnOnce() + Send + 'static, +{ + static NODE_POOL: OnceLock = OnceLock::new(); + let pool = NODE_POOL.get_or_init(|| { + let size = thread::available_parallelism() + .map(|n| n.get().max(2)) + .unwrap_or(4); + ThreadPool::new(size, "langgraph-node") + }); + pool.execute(task) +} + +pub fn run_loop_pool_execute(task: F) -> Result<(), String> +where + F: FnOnce() + Send + 'static, +{ + static RUN_LOOP_POOL: OnceLock = OnceLock::new(); + let pool = RUN_LOOP_POOL.get_or_init(|| ThreadPool::new(2, "langgraph-runloop")); + pool.execute(task) +} + +pub fn run_scheduler_loop( + entry_point: String, + finish_point: &str, + initial_arg: A, + mut spawn: FSpawn, + mut merge: FMerge, + rx: mpsc::Receiver), String>>, +) -> Result<(), String> +where + FSpawn: FnMut(String, A) -> Result<(), String>, + FMerge: FnMut(&str, Option) -> Result<(), String>, +{ + let mut active: usize = 1; + spawn(entry_point, initial_arg)?; + + while active > 0 { + 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}"))?; + + merge(&node_name, node_result.update)?; + + if node_name == finish_point { + break; + } + + for send in node_result.sends { + active += 1; + spawn(send.node, send.arg)?; + } + } + + Ok(()) +} + +pub fn merge_json_update(state: &mut Value, update: Option) { + let Some(update_value) = update else { + return; + }; + match (&mut *state, update_value) { + (Value::Object(state_obj), Value::Object(update_obj)) => { + for (k, v) in update_obj { + state_obj.insert(k, v); + } + } + (Value::Object(state_obj), Value::Array(entries)) => { + 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()); + } + } + } + } + } + _ => {} + } +} + +#[derive(Clone, Default)] +pub struct Engine { + channels: Arc>>>, + channel_notify: Arc, +} + +impl Engine { + pub fn new() -> Self { + Self::default() + } + + pub fn add_async_channel(&self, name: &str) { + let mut channels = self.channels.lock(); + channels.entry(name.to_owned()).or_default(); + } + + pub fn publish_json(&self, channel: &str, value: serde_json::Value) -> Result<(), String> { + let mut channels = self.channels.lock(); + let queue = channels + .get_mut(channel) + .ok_or_else(|| format!("Unknown channel `{channel}`"))?; + queue.push_back(value); + // Wake up waiters blocked on channel conditions/any_of. + self.channel_notify.notify_all(); + Ok(()) + } + + pub fn wait_for(&self, cond: &WaitCondition) -> Result { + match cond { + WaitCondition::Channel { channel, n } => { + if *n < 1 { + return Err("channel condition n must be >= 1".to_string()); + } + let mut channels = self.channels.lock(); + loop { + let queue = channels + .get_mut(channel) + .ok_or_else(|| format!("Unknown channel `{channel}`"))?; + if queue.len() >= *n { + if *n == 1 { + if let Some(value) = queue.pop_front() { + return Ok(WaitEvent::Channel { + channel: channel.clone(), + value, + }); + } + } else { + let mut values = Vec::with_capacity(*n); + for _ in 0..*n { + if let Some(v) = queue.pop_front() { + values.push(v); + } + } + return Ok(WaitEvent::Channel { + channel: channel.clone(), + value: serde_json::Value::Array(values), + }); + } + } + self.channel_notify.wait(&mut channels); + } + } + WaitCondition::Timer { seconds } => { + if *seconds <= 0.0 { + return Err("timer condition must be > 0".to_string()); + } + std::thread::sleep(Duration::from_secs_f64(*seconds)); + Ok(WaitEvent::Timer { seconds: *seconds }) + } + } + } + + pub fn wait_for_any_of(&self, any_of: &AnyOfCondition) -> Result { + 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 = 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))); + } + } + + let mut channels = self.channels.lock(); + 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()); + } + let queue = channels + .get_mut(channel) + .ok_or_else(|| format!("Unknown channel `{channel}`"))?; + if queue.len() >= *n { + if *n == 1 { + if let Some(value) = queue.pop_front() { + return Ok(WaitEvent::Channel { + channel: channel.clone(), + value, + }); + } + } else { + let mut values = Vec::with_capacity(*n); + for _ in 0..*n { + if let Some(v) = queue.pop_front() { + values.push(v); + } + } + return Ok(WaitEvent::Channel { + channel: channel.clone(), + value: serde_json::Value::Array(values), + }); + } + } + } + } + + 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); + self.channel_notify.wait_for(&mut channels, remaining); + } else { + self.channel_notify.wait(&mut channels); + } + } + } +} diff --git a/rust-core/src/lib.rs b/rust-core/src/lib.rs new file mode 100644 index 000000000..7bffe5eac --- /dev/null +++ b/rust-core/src/lib.rs @@ -0,0 +1,4 @@ +mod engine; +mod lib_c; +#[cfg(feature = "python-bindings")] +mod lib_py; diff --git a/rust-core/src/lib_c.rs b/rust-core/src/lib_c.rs new file mode 100644 index 000000000..c721402b0 --- /dev/null +++ b/rust-core/src/lib_c.rs @@ -0,0 +1,364 @@ +use crate::engine::{ + merge_json_update, node_pool_execute, run_loop_pool_execute, run_scheduler_loop, + AnyOfCondition, Engine, NodeExecResult, SendPayload, +}; +use serde::Deserialize; +use serde_json::Value; +use std::ffi::{CStr, CString}; +use std::os::raw::{c_char, c_void}; +use std::sync::mpsc; +use std::sync::{Arc, Mutex}; + +#[derive(Debug, Deserialize)] +struct SendPayloadJson { + node: String, + #[serde(default)] + arg: Value, +} + +#[derive(Debug, Deserialize)] +struct NodeExecResultJsonWire { + update: Option, + #[serde(default)] + sends: Vec, +} + +#[derive(Debug, Deserialize)] +struct CallbackEnvelopeIn { + ok: bool, + #[serde(default)] + payload: Option, + #[serde(default)] + error: Option, +} + +type CNodeCallback = unsafe extern "C" fn( + user_data: *mut c_void, + node: *mut c_char, + arg_json: *mut c_char, + state_json: *mut c_char, +) -> *mut c_char; + +#[derive(Clone, Copy)] +struct CUserData(*mut c_void); +unsafe impl Send for CUserData {} +unsafe impl Sync for CUserData {} + +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(), + } +} + +fn parse_c_callback_result( + raw: String, + node_name: &str, +) -> Result, 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}`"))); + } + 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(NodeExecResult { + update: payload.update, + sends, + }) +} + +fn spawn_json_node_task( + node: String, + arg: Value, + state_snapshot: Value, + tx: mpsc::Sender), String>>, + user_data_bits: usize, + callback: CNodeCallback, +) -> Result<(), String> { + node_pool_execute(move || { + let result = (|| -> Result<(String, NodeExecResult), 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 as *mut c_void, + node_c.as_ptr() as *mut c_char, + arg_c.as_ptr() as *mut c_char, + state_c.as_ptr() as *mut c_char, + ) + }; + if out_ptr.is_null() { + return Err(format!("callback returned null for `{node}`")); + } + let out_raw = unsafe { CStr::from_ptr(out_ptr) } + .to_string_lossy() + .into_owned(); + unsafe { + libc::free(out_ptr.cast()); + } + let payload = parse_c_callback_result(out_raw, &node)?; + Ok((node, payload)) + })(); + let _ = tx.send(result); + }) +} + +fn run_graph_scheduler_json( + entry_point: String, + finish_point: String, + initial_state: Value, + user_data: CUserData, + callback: CNodeCallback, +) -> Result { + let (tx, rx) = mpsc::channel::), String>>(); + let state = Arc::new(Mutex::new(initial_state)); + let user_data_bits = user_data.0 as usize; + let tx_for_spawn = tx.clone(); + let state_for_spawn = Arc::clone(&state); + let state_for_merge = Arc::clone(&state); + run_scheduler_loop( + entry_point, + &finish_point, + state.lock().expect("state mutex poisoned").clone(), + move |node, arg| { + let snapshot = state_for_spawn + .lock() + .expect("state mutex poisoned") + .clone(); + spawn_json_node_task( + node, + arg, + snapshot, + tx_for_spawn.clone(), + user_data_bits, + callback, + ) + }, + move |_node_name, update| { + let mut guard = state_for_merge.lock().expect("state mutex poisoned"); + merge_json_update(&mut guard, update); + Ok(()) + }, + rx, + )?; + let final_state = state.lock().expect("state mutex poisoned").clone(); + Ok(final_state) +} + +#[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 = (*ptr).wait_for_any_of(&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`. +/// `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, + user_data: *mut c_void, + callback: Option, +) -> *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 (tx, rx) = mpsc::channel::>(); + let user_data = CUserData(user_data); + let submit = run_loop_pool_execute(move || { + let out = run_graph_scheduler_json( + entry_point, + finish_point, + initial_state, + user_data, + callback, + ); + 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}\"}}")), + } +} diff --git a/rust-core/src/lib_py.rs b/rust-core/src/lib_py.rs new file mode 100644 index 000000000..384f96121 --- /dev/null +++ b/rust-core/src/lib_py.rs @@ -0,0 +1,312 @@ +#[cfg(feature = "python-bindings")] +use crate::engine::{ + node_pool_execute, run_loop_pool_execute, run_scheduler_loop, AnyOfCondition, Engine, + NodeExecResult, SendPayload, 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 pyo3::types::{PyDict, PyList, PyTuple}; +#[cfg(feature = "python-bindings")] +use serde_json::Value; +#[cfg(feature = "python-bindings")] +use std::sync::mpsc; +#[cfg(feature = "python-bindings")] +use std::sync::Arc; + +#[cfg(feature = "python-bindings")] +#[pyclass] +struct PyRustEngine { + inner: Engine, +} + +#[cfg(feature = "python-bindings")] +#[pymethods] +impl PyRustEngine { + #[new] + fn new() -> Self { + Self { + inner: Engine::new(), + } + } + + fn add_async_channel(&self, name: &str) { + self.inner.add_async_channel(name); + } + + fn publish_json(&self, channel: &str, value_json: &str) -> PyResult<()> { + let value: Value = serde_json::from_str(value_json) + .map_err(|e| PyValueError::new_err(format!("Invalid JSON value: {e}")))?; + self.inner + .publish_json(channel, value) + .map_err(PyValueError::new_err) + } + + fn publish_obj(&self, py: Python<'_>, channel: &str, value: Py) -> PyResult<()> { + let value_json = py_obj_to_json_string(py, &value.bind(py))?; + let parsed: Value = serde_json::from_str(&value_json) + .map_err(|e| PyValueError::new_err(format!("Invalid Python JSON value: {e}")))?; + self.inner + .publish_json(channel, parsed) + .map_err(PyValueError::new_err) + } + + fn wait_any_of_json(&self, any_of_json: &str) -> PyResult { + let any_of: AnyOfCondition = serde_json::from_str(any_of_json) + .map_err(|e| PyValueError::new_err(format!("Invalid any_of JSON: {e}")))?; + let event = self + .inner + .wait_for_any_of(&any_of) + .map_err(PyValueError::new_err)?; + serde_json::to_string(&event) + .map_err(|e| PyValueError::new_err(format!("Serialize event failed: {e}"))) + } + + fn wait_channel(&self, py: Python<'_>, channel: &str, n: usize) -> PyResult> { + let cond = WaitCondition::Channel { + channel: channel.to_string(), + n, + }; + let event = self.inner.wait_for(&cond).map_err(PyValueError::new_err)?; + let event_json = serde_json::to_string(&event) + .map_err(|e| PyValueError::new_err(format!("Serialize event failed: {e}")))?; + json_string_to_py_obj(py, &event_json) + } + + fn wait_timer(&self, py: Python<'_>, seconds: f64) -> PyResult> { + let cond = WaitCondition::Timer { seconds }; + let event = self.inner.wait_for(&cond).map_err(PyValueError::new_err)?; + let event_json = serde_json::to_string(&event) + .map_err(|e| PyValueError::new_err(format!("Serialize event failed: {e}")))?; + json_string_to_py_obj(py, &event_json) + } + + fn wait_any_of_obj(&self, py: Python<'_>, any_of_payload: Py) -> PyResult> { + let payload_json = py_obj_to_json_string(py, &any_of_payload.bind(py))?; + let any_of: AnyOfCondition = serde_json::from_str(&payload_json) + .map_err(|e| PyValueError::new_err(format!("Invalid any_of payload: {e}")))?; + let event = self + .inner + .wait_for_any_of(&any_of) + .map_err(PyValueError::new_err)?; + let event_json = serde_json::to_string(&event) + .map_err(|e| PyValueError::new_err(format!("Serialize event failed: {e}")))?; + json_string_to_py_obj(py, &event_json) + } + + fn wait_condition_json(&self, cond_json: &str) -> PyResult { + let cond: WaitCondition = serde_json::from_str(cond_json) + .map_err(|e| PyValueError::new_err(format!("Invalid condition JSON: {e}")))?; + let event = self.inner.wait_for(&cond).map_err(PyValueError::new_err)?; + serde_json::to_string(&event) + .map_err(|e| PyValueError::new_err(format!("Serialize event failed: {e}"))) + } + + fn run_graph_py( + &self, + py: Python<'_>, + entry_point: &str, + finish_point: &str, + initial_state: Py, + callback: Py, + ) -> PyResult> { + let state = Arc::new(initial_state); + let callback = Arc::new(callback); + let entry_point = entry_point.to_string(); + let finish_point = finish_point.to_string(); + let (done_tx, done_rx) = mpsc::channel::>(); + let state_for_run = Arc::clone(&state); + run_loop_pool_execute(move || { + let run_result = + run_graph_scheduler(entry_point, finish_point, callback, state_for_run); + let _ = done_tx.send(run_result); + }) + .map_err(PyValueError::new_err)?; + let run_result = py + .allow_threads(move || done_rx.recv()) + .map_err(|e| PyValueError::new_err(format!("run-loop recv failed: {e}")))?; + run_result.map_err(PyValueError::new_err)?; + Ok((*state).clone_ref(py)) + } +} + +#[cfg(feature = "python-bindings")] +fn run_graph_scheduler( + entry_point: String, + finish_point: String, + callback: Arc>, + state: Arc>, +) -> Result<(), String> { + let (tx, rx) = + mpsc::channel::, Py>), String>>(); + let initial_arg = Python::with_gil(|py| (*state).clone_ref(py)); + let callback_for_spawn = Arc::clone(&callback); + let state_for_spawn = Arc::clone(&state); + let tx_for_spawn = tx.clone(); + let state_for_merge = Arc::clone(&state); + run_scheduler_loop( + entry_point, + &finish_point, + initial_arg, + move |node, arg| { + spawn_node_task( + node, + arg, + tx_for_spawn.clone(), + Arc::clone(&callback_for_spawn), + Arc::clone(&state_for_spawn), + ) + }, + move |node_name, update| { + Python::with_gil(|py| -> Result<(), String> { + if let Some(update) = update { + apply_update_to_state(py, state_for_merge.as_ref(), &update) + .map_err(|e| format!("state merge failed for `{node_name}`: {e}"))?; + } + Ok(()) + }) + }, + rx, + ) +} + +#[cfg(feature = "python-bindings")] +fn spawn_node_task( + node: String, + arg: Py, + tx: mpsc::Sender, Py>), String>>, + callback: Arc>, + state_for_task: Arc>, +) -> Result<(), String> { + node_pool_execute(move || { + let outcome = Python::with_gil( + |py| -> Result<(String, NodeExecResult, Py>), String> { + let callback_bound = callback.as_ref().bind(py); + let payload_obj = callback_bound + .call1((node.as_str(), arg, (*state_for_task).clone_ref(py))) + .map_err(|e| format!("callback failed for node `{node}`: {e}"))?; + let payload = parse_node_exec_result(&payload_obj) + .map_err(|e| format!("invalid callback payload for `{node}`: {e}"))?; + Ok((node, payload)) + }, + ); + let _ = tx.send(outcome); + }) +} + +#[cfg(feature = "python-bindings")] +fn parse_node_exec_result( + payload_obj: &Bound<'_, PyAny>, +) -> Result, Py>, String> { + let payload_dict = payload_obj + .downcast::() + .map_err(|_| "payload must be a dict".to_string())?; + + let update_item = payload_dict + .get_item("update") + .map_err(|e| format!("failed to read update: {e}"))?; + let update = match update_item { + Some(v) if !v.is_none() => Some(v.unbind()), + _ => None, + }; + + let sends_obj = payload_dict + .get_item("sends") + .map_err(|e| format!("failed to read sends: {e}"))? + .ok_or_else(|| "missing sends".to_string())?; + let sends_list = sends_obj + .downcast::() + .map_err(|_| "sends must be a list".to_string())?; + + let mut sends = Vec::with_capacity(sends_list.len()); + for item in sends_list.iter() { + let send_dict = item + .downcast::() + .map_err(|_| "send item must be a dict".to_string())?; + let node_obj = send_dict + .get_item("node") + .map_err(|e| format!("failed to read send.node: {e}"))? + .ok_or_else(|| "send.node is required".to_string())?; + let node = node_obj + .extract::() + .map_err(|e| format!("send.node must be string: {e}"))?; + let arg: Py = match send_dict.get_item("arg") { + Ok(Some(v)) => v.unbind(), + Ok(None) => Python::with_gil(|py| py.None()), + Err(e) => return Err(format!("failed to read send.arg: {e}")), + }; + sends.push(SendPayload { node, arg }); + } + + Ok(NodeExecResult { update, sends }) +} + +#[cfg(feature = "python-bindings")] +fn apply_update_to_state(py: Python<'_>, state: &Py, update: &Py) -> PyResult<()> { + let state_obj = state.bind(py); + let update_obj = update.bind(py); + + if update_obj.is_none() { + return Ok(()); + } + if state_obj.is_instance_of::() && update_obj.is_instance_of::() { + let state_dict = state_obj.downcast::()?; + let update_dict = update_obj.downcast::()?; + state_dict.call_method1("update", (update_dict,))?; + return Ok(()); + } + if let Ok(tuple_like) = update_obj.downcast::() { + apply_pair_updates(state_obj, tuple_like)?; + return Ok(()); + } + if let Ok(tuple_like) = update_obj.downcast::() { + let list = PyList::new(py, tuple_like)?; + apply_pair_updates(state_obj, &list)?; + } + Ok(()) +} + +#[cfg(feature = "python-bindings")] +fn apply_pair_updates(state_obj: &Bound<'_, PyAny>, entries: &Bound<'_, PyList>) -> PyResult<()> { + if !state_obj.is_instance_of::() { + return Ok(()); + } + let state_dict = state_obj.downcast::()?; + for entry in entries.iter() { + if let Ok(pair) = entry.downcast::() { + if pair.len() == 2 { + let key_obj = pair.get_item(0)?; + if let Ok(key) = key_obj.extract::() { + let value_obj = pair.get_item(1)?; + state_dict.set_item(key, value_obj)?; + } + } + } + } + Ok(()) +} + +#[cfg(feature = "python-bindings")] +fn py_obj_to_json_string(py: Python<'_>, obj: &Bound<'_, PyAny>) -> PyResult { + let json_mod = py.import("json")?; + let dumped = json_mod.call_method1("dumps", (obj,))?; + dumped.extract::() +} + +#[cfg(feature = "python-bindings")] +fn json_string_to_py_obj(py: Python<'_>, value: &str) -> PyResult> { + let json_mod = py.import("json")?; + let loaded = json_mod.call_method1("loads", (value,))?; + Ok(loaded.unbind()) +} + +#[cfg(feature = "python-bindings")] +#[pymodule] +fn langgraph_rust_core(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_class::()?; + Ok(()) +}