WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-27 12:04:58 +02:00
Code Issues Packages Projects Releases Wiki Activity

refactor-go

Browse Source
This commit is contained in:
Quanzheng Long
2026-03-12 22:41:32 -07:00
parent 03ad7a011c
commit 13847a32c0
4 changed files with 262 additions and 103 deletions
Download Patch File Download Diff File Expand all files Collapse all files
langgraph-go/advancedgraph/graph.go
+47 -8
View File
@@ -2,9 +2,12 @@ package advancedgraph
import (
"fmt"
"reflect"
"runtime"
"strings"
)
type NodeFunc func(ctx *Context, arg any) (Command, error)
type NodeFunc func(ctx *Context, state map[string]any) (Command, error)
type AdvancedStateGraph struct {
nodes map[string]NodeFunc
@@ -19,20 +22,25 @@ func NewAdvancedStateGraph() *AdvancedStateGraph {
}
}
func (g *AdvancedStateGraph) AddNode(name string, fn NodeFunc) {
func (g *AdvancedStateGraph) AddNode(fn NodeFunc) string {
name := NodeName(fn)
if _, exists := g.nodes[name]; exists {
panic(fmt.Sprintf("node `%s` already exists", name))
}
g.nodes[name] = fn
return name
}
func (g *AdvancedStateGraph) AddAsyncChannel(name string) {
g.asyncChannels = append(g.asyncChannels, name)
}
func (g *AdvancedStateGraph) SetEntryPoint(name string) {
g.entryPoint = name
func (g *AdvancedStateGraph) SetEntryNode(fn NodeFunc) {
g.entryPoint = NodeName(fn)
}
func (g *AdvancedStateGraph) SetFinishPoint(name string) {
g.finishPoint = name
func (g *AdvancedStateGraph) SetFinishNode(fn NodeFunc) {
g.finishPoint = NodeName(fn)
}
func (g *AdvancedStateGraph) Compile() *CompiledGraph {
@@ -100,12 +108,19 @@ func (g *CompiledGraph) Start(initialState map[string]any) (*Handler, error) {
g.entryPoint,
g.finishPoint,
initialState,
func(node string, arg any, _ map[string]any) (Command, error) {
func(node string, arg any, fallbackState 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)
stateArg, ok := arg.(map[string]any)
if !ok {
stateArg = fallbackState
}
if stateArg == nil {
return Command{}, fmt.Errorf("node `%s` expected map state argument", node)
}
return fn(&Context{engine: engine}, stateArg)
},
)
handler.done <- resultOrErr{state: state, err: err}
@@ -113,3 +128,27 @@ func (g *CompiledGraph) Start(initialState map[string]any) (*Handler, error) {
}()
return handler, nil
}
func NodeName(fn NodeFunc) string {
pc := reflect.ValueOf(fn).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
}
langgraph-go/tests/sub_agents_test.go
+137 -93
View File
@@ -29,6 +29,128 @@ func (m *mockPlanner) invoke() []decision {
return resp
}
type lunchWorkflow struct {
planner *mockPlanner
names map[string]string
}
func cloneState(state map[string]any) map[string]any {
out := make(map[string]any, len(state)+2)
for k, v := range state {
out[k] = v
}
out["output"] = append([]string(nil), outputSlice(state)...)
return out
}
func outputSlice(state map[string]any) []string {
raw, ok := state["output"]
if !ok || raw == nil {
return []string{}
}
switch v := raw.(type) {
case []string:
return v
case []any:
out := make([]string, 0, len(v))
for _, item := range v {
if s, ok := item.(string); ok {
out = append(out, s)
}
}
return out
default:
return []string{}
}
}
func (w *lunchWorkflow) llmNode(ctx *ag.Context, state map[string]any) (ag.Command, error) {
decisions := w.planner.invoke()
sends := make([]ag.Send, 0, 4)
for _, d := range decisions {
if d.Type == "end" {
next := cloneState(state)
next["complete"] = d.Complete
return ag.Command{
Goto: []ag.Send{
{Node: w.names["order"], Arg: next},
},
}, nil
}
if d.Type == "sub_agent" {
next := cloneState(state)
next["sub_agent_input"] = d.SubAgent
sends = append(sends, ag.Send{Node: w.names["sub"], Arg: next})
}
if d.Type == "tool" {
next := cloneState(state)
next["tool_input"] = d.Tool
sends = append(sends, ag.Send{Node: w.names["tool"], Arg: next})
}
}
sends = append(sends, ag.Send{Node: w.names["wait"], Arg: cloneState(state)})
return ag.Command{Goto: sends}, nil
}
func (w *lunchWorkflow) waitNode(ctx *ag.Context, state map[string]any) (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 := outputSlice(state)
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.names["llm"], Arg: cloneState(state)}}}, nil
}
output = append(output, "timer: no updates yet")
state["output"] = output
return ag.Command{Goto: []ag.Send{{Node: w.names["wait"], Arg: cloneState(state)}}}, nil
}
func (w *lunchWorkflow) toolNode(ctx *ag.Context, state map[string]any) (ag.Command, error) {
toolInput, _ := state["tool_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, state map[string]any) (ag.Command, error) {
subInput, _ := state["sub_agent_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, state map[string]any) (ag.Command, error) {
complete, _ := state["complete"].(string)
output := outputSlice(state)
output = append(output, "order_food: "+complete)
state["output"] = output
state["done"] = complete
return ag.Command{Update: state}, nil
}
func TestSubAgentsEquivalentFlow(t *testing.T) {
planner := &mockPlanner{
responses: [][]decision{
@@ -42,105 +164,24 @@ func TestSubAgentsEquivalentFlow(t *testing.T) {
{{Type: "end", Complete: "order submitted"}},
},
}
workflow := &lunchWorkflow{
planner: planner,
names: make(map[string]string),
}
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.invoke()
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
})
workflow.names["llm"] = graph.AddNode(workflow.llmNode)
workflow.names["wait"] = graph.AddNode(workflow.waitNode)
workflow.names["tool"] = graph.AddNode(workflow.toolNode)
workflow.names["sub"] = graph.AddNode(workflow.subAgentNode)
workflow.names["order"] = graph.AddNode(workflow.orderFoodNode)
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")
graph.SetEntryNode(workflow.llmNode)
graph.SetFinishNode(workflow.orderFoodNode)
handler, err := graph.Compile().Start(
map[string]any{
@@ -163,7 +204,10 @@ func TestSubAgentsEquivalentFlow(t *testing.T) {
t.Fatalf("result failed: %v", err)
}
output := result["output"].([]string)
output := outputSlice(result)
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"])
}
rust-core/src/engine.rs
+76 -1
View File
@@ -2,6 +2,7 @@ use parking_lot::{Condvar, Mutex};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::{HashMap, VecDeque};
use std::env;
use std::sync::mpsc;
use std::sync::Arc;
use std::sync::Mutex as StdMutex;
@@ -47,6 +48,28 @@ pub struct NodeExecResult<U, A> {
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}");
}
}
struct ThreadPool {
tx: mpsc::Sender<Task>,
_workers: Vec<thread::JoinHandle<()>>,
@@ -85,6 +108,7 @@ impl ThreadPool {
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}"))
@@ -95,6 +119,7 @@ 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 size = thread::available_parallelism()
@@ -109,6 +134,7 @@ pub fn run_loop_pool_execute<F>(task: F) -> Result<(), String>
where
F: FnOnce() + Send + 'static,
{
debug_log("run_loop_pool_execute() called");
static RUN_LOOP_POOL: OnceLock<ThreadPool> = OnceLock::new();
let pool = RUN_LOOP_POOL.get_or_init(|| ThreadPool::new(2, "langgraph-runloop"));
pool.execute(task)
@@ -126,42 +152,65 @@ 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 {
@@ -172,7 +221,7 @@ pub fn merge_json_update(state: &mut Value, update: Option<Value>) {
}
}
}
_ => {}
_ => debug_log("merge_json_update(): unsupported update shape, ignored"),
}
}
@@ -184,15 +233,18 @@ pub struct Engine {
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();
channels.entry(name.to_owned()).or_default();
}
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();
let queue = channels
.get_mut(channel)
@@ -204,6 +256,7 @@ impl Engine {
}
pub fn wait_for(&self, cond: &WaitCondition) -> Result<WaitEvent, String> {
debug_log(&format!("Engine::wait_for(cond={cond:?})"));
match cond {
WaitCondition::Channel { channel, n } => {
if *n < 1 {
@@ -215,6 +268,9 @@ impl Engine {
.get_mut(channel)
.ok_or_else(|| format!("Unknown channel `{channel}`"))?;
if queue.len() >= *n {
debug_log(&format!(
"Engine::wait_for channel ready (channel={channel}, n={n})"
));
if *n == 1 {
if let Some(value) = queue.pop_front() {
return Ok(WaitEvent::Channel {
@@ -235,6 +291,9 @@ impl Engine {
});
}
}
debug_log(&format!(
"Engine::wait_for waiting on channel condvar (channel={channel}, n={n})"
));
self.channel_notify.wait(&mut channels);
}
}
@@ -243,12 +302,17 @@ impl Engine {
return Err("timer condition must be > 0".to_string());
}
std::thread::sleep(Duration::from_secs_f64(*seconds));
debug_log(&format!("Engine::wait_for timer fired (seconds={seconds})"));
Ok(WaitEvent::Timer { seconds: *seconds })
}
}
}
pub fn wait_for_any_of(&self, any_of: &AnyOfCondition) -> Result<WaitEvent, String> {
debug_log(&format!(
"Engine::wait_for_any_of(conditions={})",
any_of.conditions.len()
));
if any_of.conditions.is_empty() {
return Err("any_of requires at least one condition".to_string());
}
@@ -275,6 +339,9 @@ impl Engine {
.get_mut(channel)
.ok_or_else(|| format!("Unknown channel `{channel}`"))?;
if queue.len() >= *n {
debug_log(&format!(
"Engine::wait_for_any_of channel hit (channel={channel}, n={n})"
));
if *n == 1 {
if let Some(value) = queue.pop_front() {
return Ok(WaitEvent::Channel {
@@ -302,11 +369,19 @@ impl Engine {
let timeout = Duration::from_secs_f64(seconds);
let elapsed = started.elapsed();
if elapsed >= timeout {
debug_log(&format!(
"Engine::wait_for_any_of timer hit (seconds={seconds})"
));
return Ok(WaitEvent::Timer { seconds });
}
let remaining = timeout.saturating_sub(elapsed);
debug_log(&format!(
"Engine::wait_for_any_of waiting on condvar with timeout {:?}",
remaining
));
self.channel_notify.wait_for(&mut channels, remaining);
} else {
debug_log("Engine::wait_for_any_of waiting on condvar without timeout");
self.channel_notify.wait(&mut channels);
}
}
rust-core/src/lib_c.rs
+2 -1
View File
@@ -147,10 +147,11 @@ fn run_graph_scheduler_json(
let tx_for_spawn = tx.clone();
let state_for_spawn = Arc::clone(&state);
let state_for_merge = Arc::clone(&state);
let initial_arg = state.lock().expect("state mutex poisoned").clone();
run_scheduler_loop(
entry_point,
&finish_point,
state.lock().expect("state mutex poisoned").clone(),
initial_arg,
move |node, arg| {
let snapshot = state_for_spawn
.lock()