mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-30 19:59:40 +02:00
Some fancy stuff
This commit is contained in:
William Fu-Hinthorn
@@ -0,0 +1,28 @@
|
||||
.PHONY: all format build test
|
||||
|
||||
# Default target executed when no arguments are given to make.
|
||||
all: help
|
||||
|
||||
format:
|
||||
go fmt ./...
|
||||
|
||||
build:
|
||||
go build ./...
|
||||
|
||||
test:
|
||||
go test ./...
|
||||
|
||||
|
||||
######################
|
||||
# HELP
|
||||
######################
|
||||
|
||||
help:
|
||||
@echo '===================='
|
||||
@echo '-- DOCUMENTATION --'
|
||||
|
||||
@echo '-- LINTING --'
|
||||
@echo 'format - run code formatters'
|
||||
@echo 'build - build the project'
|
||||
@echo 'test - run unit tests'
|
||||
|
||||
@@ -0,0 +1,10 @@
|
||||
module langchain.dev/langgraph
|
||||
|
||||
go 1.23.0
|
||||
|
||||
toolchain go1.23.9
|
||||
|
||||
require (
|
||||
github.com/google/uuid v1.6.0 // indirect
|
||||
golang.org/x/sync v0.14.0 // indirect
|
||||
)
|
||||
@@ -0,0 +1,4 @@
|
||||
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
|
||||
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
|
||||
golang.org/x/sync v0.14.0 h1:woo0S4Yywslg6hp4eUFjTVOyKt0RookbpAHG4c1HmhQ=
|
||||
golang.org/x/sync v0.14.0/go.mod h1:1dzgHSNfp02xaA81J2MS99Qcpr2w7fw1gpm99rleRqA=
|
||||
@@ -0,0 +1,522 @@
|
||||
package pregel
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
"reflect"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
func PrepareNextTasks(
|
||||
ctx context.Context,
|
||||
checkpoint Checkpoint,
|
||||
pendingWrites []interface{},
|
||||
processes map[string]PregelNode,
|
||||
channels map[string]BaseChannel,
|
||||
managed ManagedValueMapping,
|
||||
config RunnableConfig,
|
||||
step int,
|
||||
forExecution bool,
|
||||
store BaseStore,
|
||||
checkpointer BaseCheckpointSaver,
|
||||
// ─ optimisation hints (optional) ─
|
||||
triggerToNodes map[string][]string,
|
||||
updatedChannels map[string]struct{},
|
||||
) (map[string]interface{}, error) {
|
||||
|
||||
// Decode checkpoint.id (UUID/xxhash) into raw bytes for deterministic task-id hashing.
|
||||
cleanID := strings.ReplaceAll(checkpoint.ID, "-", "")
|
||||
checkpointIDBytes, err := hex.DecodeString(cleanID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
nullVersion := checkpointNullVersion(checkpoint)
|
||||
tasks := make(map[string]interface{})
|
||||
|
||||
// Consume pending sends
|
||||
for idx := range checkpoint.PendingSends {
|
||||
task, err := PrepareSingleTask(
|
||||
ctx,
|
||||
[]interface{}{PUSH, idx},
|
||||
"",
|
||||
checkpoint,
|
||||
checkpointIDBytes,
|
||||
nullVersion,
|
||||
pendingWrites,
|
||||
processes,
|
||||
channels,
|
||||
managed,
|
||||
config,
|
||||
step,
|
||||
forExecution,
|
||||
store,
|
||||
checkpointer,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if task == nil {
|
||||
continue
|
||||
}
|
||||
if id, ok := taskID(task); ok {
|
||||
tasks[id] = task
|
||||
}
|
||||
}
|
||||
|
||||
var candidateNodes []string
|
||||
|
||||
if len(updatedChannels) > 0 && len(triggerToNodes) > 0 {
|
||||
nodeSet := map[string]struct{}{}
|
||||
for ch := range updatedChannels {
|
||||
for _, n := range triggerToNodes[ch] {
|
||||
nodeSet[n] = struct{}{}
|
||||
}
|
||||
}
|
||||
for n := range nodeSet {
|
||||
candidateNodes = append(candidateNodes, n)
|
||||
}
|
||||
sort.Strings(candidateNodes) // deterministic order
|
||||
} else if len(checkpoint.ChannelVersions) == 0 {
|
||||
candidateNodes = nil
|
||||
} else {
|
||||
for n := range processes {
|
||||
candidateNodes = append(candidateNodes, n)
|
||||
}
|
||||
sort.Strings(candidateNodes)
|
||||
}
|
||||
|
||||
for _, name := range candidateNodes {
|
||||
task, err := PrepareSingleTask(
|
||||
ctx,
|
||||
[]interface{}{PULL, name},
|
||||
"", // checksum only used when resuming a partial step
|
||||
checkpoint,
|
||||
checkpointIDBytes,
|
||||
nullVersion,
|
||||
pendingWrites,
|
||||
processes,
|
||||
channels,
|
||||
managed,
|
||||
config,
|
||||
step,
|
||||
forExecution,
|
||||
store,
|
||||
checkpointer,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if task == nil {
|
||||
continue
|
||||
}
|
||||
if id, ok := taskID(task); ok {
|
||||
tasks[id] = task
|
||||
}
|
||||
}
|
||||
|
||||
return tasks, nil
|
||||
}
|
||||
|
||||
func PrepareSingleTask(
|
||||
ctx context.Context,
|
||||
taskPath []interface{}, // e.g. [PUSH, idx] OR [PULL, "node"]
|
||||
taskIDChecksum string, // optional – used when resuming
|
||||
checkpoint Checkpoint, // state captured at end of previous step
|
||||
checkpointIDBytes []byte, // checkpoint.id as bytes (uuid / xxhash)
|
||||
checkpointNullVersion interface{}, // sentinel “null” version value
|
||||
pendingWrites []interface{}, // successful writes from *this* step so far
|
||||
processes map[string]PregelNode, // graph definition
|
||||
channels map[string]BaseChannel, // live channel values
|
||||
managed ManagedValueMapping, // placeholder resolver
|
||||
config RunnableConfig, // config inherited from graph.Invoke()
|
||||
step int, // current super-step (n+1)
|
||||
forExecution bool, // false = planning pass, true = exec pass
|
||||
store BaseStore, // needed for reads/writes
|
||||
checkpointer BaseCheckpointSaver, // used only when executing
|
||||
) (interface{}, error) {
|
||||
// Ensure checkpoint.ChannelVersions is initialized
|
||||
if checkpoint.ChannelVersions == nil {
|
||||
checkpoint.ChannelVersions = make(map[string]int64)
|
||||
}
|
||||
|
||||
cfgSection := config.Configurable
|
||||
if cfgSection == nil {
|
||||
cfgSection = map[string]interface{}{}
|
||||
}
|
||||
parentNS, _ := cfgSection[CONFIG_KEY_CHECKPOINT_NS].(string)
|
||||
|
||||
emitConfig := func(base RunnableConfig, md map[string]interface{}) RunnableConfig {
|
||||
// Make a shallow copy of the struct
|
||||
out := base
|
||||
|
||||
if out.Configurable == nil {
|
||||
out.Configurable = map[string]interface{}{}
|
||||
}
|
||||
confClone := make(map[string]interface{}, len(out.Configurable))
|
||||
for k, v := range out.Configurable {
|
||||
confClone[k] = v
|
||||
}
|
||||
confClone[CONFIG_KEY_SCRATCHPAD] = createScratchpad(
|
||||
out.Configurable[CONFIG_KEY_SCRATCHPAD].(map[string]interface{}),
|
||||
pendingWrites,
|
||||
md["langgraph_checkpoint_ns"].(string),
|
||||
md["langgraph_checkpoint_ns"].(string),
|
||||
out.Configurable[CONFIG_KEY_RESUME_MAP].(map[string]interface{}),
|
||||
)
|
||||
confClone[CONFIG_KEY_CHECKPOINTER] = checkpointer
|
||||
out.Configurable = confClone
|
||||
|
||||
if out.Metadata == nil {
|
||||
out.Metadata = map[string]interface{}{}
|
||||
}
|
||||
for k, v := range md {
|
||||
out.Metadata[k] = v
|
||||
}
|
||||
|
||||
return out
|
||||
}
|
||||
|
||||
// Convenience for checksum comparison
|
||||
checkSumMatch := func(need string) error {
|
||||
if taskIDChecksum != "" && taskIDChecksum != need {
|
||||
return fmt.Errorf("%s != %s", need, taskIDChecksum)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// PUSH
|
||||
if len(taskPath) > 0 && taskPath[0] == PUSH {
|
||||
|
||||
// PUSH triggered via explicit Call (happens during node execution)
|
||||
// taskPath shape: [PUSH, parentPath, writeIdx, parentTaskID, Call]
|
||||
if len(taskPath) >= 5 {
|
||||
call, ok := taskPath[4].(Call)
|
||||
if ok {
|
||||
name, isStr := call.Func.(string)
|
||||
if !isStr {
|
||||
name = "unknown"
|
||||
}
|
||||
|
||||
// Hash-stable checkpoint namespace
|
||||
var checkpointNS string
|
||||
if parentNS == "" {
|
||||
checkpointNS = name
|
||||
} else {
|
||||
checkpointNS = parentNS + NS_SEP + name
|
||||
}
|
||||
|
||||
// Deterministic task-id
|
||||
taskID := taskIDFunc(
|
||||
checkpointIDBytes,
|
||||
checkpointNS,
|
||||
strconv.Itoa(step),
|
||||
name,
|
||||
PUSH,
|
||||
taskPathStr(taskPath[1]),
|
||||
fmt.Sprintf("%v", taskPath[2]),
|
||||
)
|
||||
if err := checkSumMatch(taskID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
taskCheckpointNS := checkpointNS + NS_END + taskID
|
||||
metadata := map[string]interface{}{
|
||||
"langgraph_step": step,
|
||||
"langgraph_node": name,
|
||||
"langgraph_triggers": []string{PUSH},
|
||||
"langgraph_path": taskPath[:3],
|
||||
"langgraph_checkpoint_ns": taskCheckpointNS,
|
||||
}
|
||||
|
||||
if forExecution {
|
||||
var node NodeRunnable
|
||||
if proc, ok := processes[name]; ok {
|
||||
node = proc.Node
|
||||
}
|
||||
|
||||
return PregelExecutableTask{
|
||||
PregelTask: PregelTask{
|
||||
ID: taskID,
|
||||
Name: name,
|
||||
Path: taskPath[:3],
|
||||
},
|
||||
Input: call.Input,
|
||||
Node: node,
|
||||
Writes: []Write{},
|
||||
Config: emitConfig(config, metadata),
|
||||
Triggers: []string{PUSH},
|
||||
}, nil
|
||||
}
|
||||
return PregelTask{ID: taskID, Name: name, Path: taskPath[:3]}, nil
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------
|
||||
// 1b. Standard pending-send packet: taskPath shape [PUSH, idx]
|
||||
// ---------------------------------------------------------------------
|
||||
if len(taskPath) == 2 {
|
||||
idx, ok := taskPath[1].(int)
|
||||
if !ok || idx >= len(checkpoint.PendingSends) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
packet := checkpoint.PendingSends[idx]
|
||||
proc, ok := processes[packet.Node]
|
||||
if !ok || proc.Node == nil {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
checkpointNS := parentNS
|
||||
if checkpointNS != "" {
|
||||
checkpointNS += NS_SEP + packet.Node
|
||||
} else {
|
||||
checkpointNS = packet.Node
|
||||
}
|
||||
|
||||
taskID := taskIDFunc(
|
||||
checkpointIDBytes,
|
||||
checkpointNS,
|
||||
strconv.Itoa(step),
|
||||
packet.Node,
|
||||
PUSH,
|
||||
strconv.Itoa(idx),
|
||||
)
|
||||
if err := checkSumMatch(taskID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
taskCheckpointNS := checkpointNS + NS_END + taskID
|
||||
metadata := map[string]interface{}{
|
||||
"langgraph_step": step,
|
||||
"langgraph_node": packet.Node,
|
||||
"langgraph_triggers": []string{PUSH},
|
||||
"langgraph_path": taskPath,
|
||||
"langgraph_checkpoint_ns": taskCheckpointNS,
|
||||
}
|
||||
|
||||
if forExecution {
|
||||
return PregelExecutableTask{
|
||||
PregelTask: PregelTask{
|
||||
ID: taskID,
|
||||
Name: packet.Node,
|
||||
Path: taskPath,
|
||||
},
|
||||
Input: packet.Arg,
|
||||
Node: proc.Node,
|
||||
Writes: nil,
|
||||
Config: emitConfig(config, metadata),
|
||||
Triggers: []string{PUSH},
|
||||
}, nil
|
||||
}
|
||||
return PregelTask{ID: taskID, Name: packet.Node, Path: taskPath}, nil
|
||||
}
|
||||
|
||||
// An ill-formed PUSH path – nothing to schedule
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// PULL branch
|
||||
if len(taskPath) > 0 && taskPath[0] == PULL {
|
||||
if len(taskPath) < 2 {
|
||||
return nil, nil
|
||||
}
|
||||
name, ok := taskPath[1].(string)
|
||||
if !ok {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
proc, ok := processes[name]
|
||||
if !ok || proc.Node == nil {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
seen := map[string]interface{}{}
|
||||
if v, _ := checkpoint.VersionsSeen[name].(map[string]interface{}); v != nil {
|
||||
for k, vv := range v { // shallow copy
|
||||
seen[k] = vv
|
||||
}
|
||||
}
|
||||
|
||||
var triggers []string
|
||||
for _, ch := range proc.Triggers {
|
||||
cv, exists := checkpoint.ChannelVersions[ch]
|
||||
if !exists {
|
||||
cv = checkpointNullVersion.(int64) // use the provided null version
|
||||
}
|
||||
sv, _ := seen[ch].(int64) // default to 0 if not exists or wrong type
|
||||
if compareVersion(cv, sv) > 0 {
|
||||
triggers = append(triggers, ch)
|
||||
}
|
||||
}
|
||||
if len(triggers) == 0 {
|
||||
return nil, nil // not ready
|
||||
}
|
||||
sort.Strings(triggers)
|
||||
|
||||
input := map[string]interface{}{}
|
||||
for _, ch := range proc.Triggers {
|
||||
if v, ok := channels[ch]; ok {
|
||||
input[ch] = v
|
||||
}
|
||||
}
|
||||
|
||||
checkpointNS := parentNS
|
||||
if checkpointNS != "" {
|
||||
checkpointNS += NS_SEP + name
|
||||
} else {
|
||||
checkpointNS = name
|
||||
}
|
||||
|
||||
taskID := taskIDFunc(
|
||||
checkpointIDBytes,
|
||||
checkpointNS,
|
||||
strconv.Itoa(step),
|
||||
name,
|
||||
PULL,
|
||||
// join triggers to guarantee deterministic id
|
||||
fmt.Sprintf("%v", triggers),
|
||||
)
|
||||
if err := checkSumMatch(taskID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
taskCheckpointNS := checkpointNS + NS_END + taskID
|
||||
metadata := map[string]interface{}{
|
||||
"langgraph_step": step,
|
||||
"langgraph_node": name,
|
||||
"langgraph_triggers": triggers,
|
||||
"langgraph_path": taskPath,
|
||||
"langgraph_checkpoint_ns": taskCheckpointNS,
|
||||
}
|
||||
|
||||
if forExecution {
|
||||
return PregelExecutableTask{
|
||||
PregelTask: PregelTask{
|
||||
ID: taskID,
|
||||
Name: name,
|
||||
Path: taskPath,
|
||||
},
|
||||
Input: input,
|
||||
Node: proc.Node,
|
||||
Writes: nil,
|
||||
Config: emitConfig(config, metadata),
|
||||
Triggers: triggers,
|
||||
}, nil
|
||||
}
|
||||
return PregelTask{ID: taskID, Name: name, Path: taskPath}, nil
|
||||
}
|
||||
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// Private / Helpers
|
||||
|
||||
// taskIDFunc deterministically hashes the checkpoint-scoped information that
|
||||
// must be unique for a task in a given super-step.
|
||||
func taskIDFunc(checkpointIDBytes []byte, parts ...string) string {
|
||||
h := sha256.New()
|
||||
_, _ = h.Write(checkpointIDBytes)
|
||||
for _, p := range parts {
|
||||
_, _ = h.Write([]byte(p))
|
||||
}
|
||||
return hex.EncodeToString(h.Sum(nil))
|
||||
}
|
||||
|
||||
// taskPathStr is only used so the path element contributes to the hash in a
|
||||
// deterministic textual form.
|
||||
func taskPathStr(path interface{}) string {
|
||||
return fmt.Sprintf("%v", path)
|
||||
}
|
||||
|
||||
// createScratchpad returns an *immutable* copy of the scratchpad that will
|
||||
// be injected into the task-local Config. We:
|
||||
//
|
||||
// 1. start from the previous scratchpad (if any),
|
||||
// 2. merge in any successful writes from earlier tasks in this super-step,
|
||||
// 3. copy-on-write so individual tasks never share interior maps.
|
||||
//
|
||||
// The logic below is intentionally simple; extend as needed.
|
||||
func createScratchpad(
|
||||
current map[string]interface{},
|
||||
pendingWrites []interface{},
|
||||
taskID string,
|
||||
checkpointHash string,
|
||||
resumeMap map[string]interface{},
|
||||
) map[string]interface{} {
|
||||
out := map[string]interface{}{}
|
||||
for k, v := range current {
|
||||
out[k] = v
|
||||
}
|
||||
if len(pendingWrites) > 0 {
|
||||
out["pending_writes"] = append([]interface{}{}, pendingWrites...)
|
||||
}
|
||||
if checkpointHash != "" {
|
||||
out["checkpoint_hash"] = checkpointHash
|
||||
}
|
||||
if resumeMap != nil {
|
||||
out["resume_map"] = resumeMap
|
||||
}
|
||||
out["task_id"] = taskID
|
||||
return out
|
||||
}
|
||||
|
||||
func checkpointNullVersion(_ Checkpoint) interface{} {
|
||||
// Return the zero value for int64 as the null version
|
||||
return int64(0)
|
||||
}
|
||||
|
||||
func taskID(t interface{}) (string, bool) {
|
||||
switch v := t.(type) {
|
||||
case PregelTask:
|
||||
return v.ID, true
|
||||
case PregelExecutableTask:
|
||||
return v.ID, true
|
||||
default:
|
||||
return "", false
|
||||
}
|
||||
}
|
||||
|
||||
func compareVersion(a, b interface{}) int {
|
||||
switch av := a.(type) {
|
||||
case int:
|
||||
bv, _ := b.(int)
|
||||
return av - bv
|
||||
case int64:
|
||||
var bv int64
|
||||
switch bvVal := b.(type) {
|
||||
case int64:
|
||||
bv = bvVal
|
||||
case int:
|
||||
bv = int64(bvVal)
|
||||
default:
|
||||
bv = 0
|
||||
}
|
||||
if av == bv {
|
||||
return 0
|
||||
}
|
||||
if av < bv {
|
||||
return -1
|
||||
}
|
||||
return 1
|
||||
case string:
|
||||
bv, _ := b.(string)
|
||||
if av == bv {
|
||||
return 0
|
||||
}
|
||||
if av < bv {
|
||||
return -1
|
||||
}
|
||||
return 1
|
||||
// Fallback to reflect.DeepEqual comparison: not perfect but safe.
|
||||
default:
|
||||
if reflect.DeepEqual(a, b) {
|
||||
return 0
|
||||
}
|
||||
return 1
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
package pregel
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
func EmptyCheckpoint() (*Checkpoint, error) {
|
||||
uid, err := uuid.NewV6()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &Checkpoint{
|
||||
Version: 1,
|
||||
ID: uid.String(),
|
||||
Timestamp: time.Now().Format(time.RFC3339),
|
||||
ChannelValues: map[string]interface{}{},
|
||||
ChannelVersions: map[string]int64{},
|
||||
VersionsSeen: map[string]interface{}{},
|
||||
PendingSends: []Send{},
|
||||
}, nil
|
||||
}
|
||||
|
||||
type Checkpointer interface {
|
||||
PutWrites(ctx context.Context, checkpoint Checkpoint, writes []Write) error
|
||||
}
|
||||
@@ -0,0 +1,138 @@
|
||||
package pregel
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
)
|
||||
|
||||
// Pregel is the top-level graph object.
|
||||
type Pregel struct {
|
||||
Name string
|
||||
Nodes map[string]PregelNode
|
||||
Channels map[string]BaseChannel
|
||||
LoopCfg RunnableConfig
|
||||
Checkptr BaseCheckpointSaver
|
||||
Store BaseStore
|
||||
Debug bool
|
||||
}
|
||||
|
||||
func (g *Pregel) Stream(
|
||||
input any,
|
||||
cfg RunnableConfig,
|
||||
opts *StreamOptions,
|
||||
) (<-chan StreamChunk, <-chan error) {
|
||||
eventCh := make(chan StreamChunk, 16)
|
||||
errCh := make(chan error, 1)
|
||||
if opts == nil {
|
||||
opts = &StreamOptions{}
|
||||
}
|
||||
ctx := opts.Context
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
mode := opts.Mode
|
||||
if mode == "" {
|
||||
mode = StreamValues
|
||||
}
|
||||
// TODO: opts.Debug
|
||||
|
||||
// ensure output channels are set / valid
|
||||
outChans := opts.OutputChannels
|
||||
if len(outChans) == 0 {
|
||||
for k := range g.Channels {
|
||||
if _, ok := g.Channels[k]; ok {
|
||||
outChans = append(outChans, k)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if opts.MaxConcurrency > 0 {
|
||||
cfg.MaxConcurrency = opts.MaxConcurrency
|
||||
}
|
||||
if cfg.MaxConcurrency == 0 {
|
||||
cfg.MaxConcurrency = 4
|
||||
}
|
||||
if cfg.RecursionLimit == 0 {
|
||||
cfg.RecursionLimit = 25
|
||||
}
|
||||
if opts.CheckpointDuring != nil {
|
||||
cfg.Configurable[CONFIG_KEY_CHECKPOINT_DURING] = *opts.CheckpointDuring
|
||||
}
|
||||
checkpoint, err := EmptyCheckpoint()
|
||||
if err != nil {
|
||||
errCh <- err
|
||||
return nil, errCh
|
||||
}
|
||||
loop := NewLoop(
|
||||
ctx,
|
||||
*checkpoint,
|
||||
g.Nodes,
|
||||
g.channelsAsConcrete(),
|
||||
nil, // managed values
|
||||
cfg,
|
||||
nil, // g.checkpointer, // may be nil
|
||||
nil, // g.store,
|
||||
)
|
||||
loop.interruptBefore = opts.InterruptBefore
|
||||
loop.interruptAfter = opts.InterruptAfter
|
||||
loop.streamCh = eventCh
|
||||
loop.streamMode = mode
|
||||
// loop.debug = debug
|
||||
|
||||
go func() {
|
||||
defer close(eventCh)
|
||||
defer close(errCh)
|
||||
|
||||
// Create a runner to execute tasks
|
||||
runner := NewPregelRunner(loop, nil)
|
||||
|
||||
// Use the tick method in a loop instead of Run()
|
||||
for {
|
||||
more, err := loop.tick(outChans)
|
||||
if err != nil {
|
||||
// Check if this is a GraphInterrupt error
|
||||
var interrupt GraphInterrupt
|
||||
if errors.As(err, &interrupt) {
|
||||
// Handle interrupt gracefully
|
||||
break
|
||||
}
|
||||
// Otherwise, it's a real error
|
||||
errCh <- err
|
||||
return
|
||||
}
|
||||
|
||||
runnerOpts := TickOptions{
|
||||
MaxConcurrency: cfg.MaxConcurrency,
|
||||
}
|
||||
|
||||
if opts.Debug != nil && *opts.Debug {
|
||||
runnerOpts.OnStepWrite = func(step int, writes []Write) {
|
||||
// TODO: Handle debugging info
|
||||
}
|
||||
}
|
||||
|
||||
if err := runner.tick(runnerOpts); err != nil {
|
||||
errCh <- err
|
||||
return
|
||||
}
|
||||
|
||||
// No more iterations needed, we're done
|
||||
if !more {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
}()
|
||||
|
||||
return eventCh, errCh
|
||||
}
|
||||
|
||||
func (g *Pregel) channelsAsConcrete() map[string]BaseChannel {
|
||||
out := make(map[string]BaseChannel, len(g.Channels))
|
||||
for k, v := range g.Channels {
|
||||
if ch, ok := v.(BaseChannel); ok {
|
||||
out[k] = ch
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,500 @@
|
||||
package pregel
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"fmt"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
type GraphInterrupt struct {
|
||||
Interrupts any
|
||||
}
|
||||
|
||||
func (e GraphInterrupt) Error() string { return "graph interrupted" }
|
||||
|
||||
type GraphDelegate struct {
|
||||
Payload map[string]any
|
||||
}
|
||||
|
||||
func (e GraphDelegate) Error() string { return "graph delegation requested" }
|
||||
|
||||
func hashID(checkpointID string, parts ...string) string {
|
||||
b, _ := hex.DecodeString(checkpointID)
|
||||
h := sha256.New()
|
||||
h.Write(b)
|
||||
for _, p := range parts {
|
||||
h.Write([]byte(p))
|
||||
}
|
||||
return hex.EncodeToString(h.Sum(nil))
|
||||
}
|
||||
|
||||
type PregelLoop struct {
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
cfg RunnableConfig
|
||||
store BaseStore
|
||||
checkpoint Checkpoint
|
||||
checkporter BaseCheckpointSaver
|
||||
|
||||
processes map[string]PregelNode
|
||||
channels map[string]BaseChannel
|
||||
managed ManagedValueMapping
|
||||
|
||||
step int
|
||||
stop int
|
||||
interruptBefore []string
|
||||
interruptAfter []string
|
||||
|
||||
pendingWrites []WriteRecord
|
||||
tasks map[string]*PregelExecutableTask
|
||||
toInterrupt []*PregelExecutableTask
|
||||
|
||||
triggerToNodes map[string][]string
|
||||
updatedChans map[string]struct{}
|
||||
|
||||
// synchronisation / workers
|
||||
workers int
|
||||
wg sync.WaitGroup
|
||||
errMu sync.Mutex
|
||||
runErr error
|
||||
// Streaming
|
||||
streamCh chan<- StreamChunk
|
||||
streamMode StreamMode
|
||||
|
||||
pendingMu sync.Mutex
|
||||
checkpointPendingWrites []PendingWrite
|
||||
|
||||
checkpointer Checkpointer // interface with PutWrites()
|
||||
checkpointConfig RunnableConfig
|
||||
|
||||
emit func(task *PregelExecutableTask, writes []Write, cached bool)
|
||||
}
|
||||
|
||||
type WriteRecord struct {
|
||||
Task string
|
||||
Chan string
|
||||
Value any
|
||||
}
|
||||
|
||||
// NewLoop initialises a fully-featured loop.
|
||||
func NewLoop(
|
||||
ctx context.Context,
|
||||
checkpoint Checkpoint,
|
||||
processes map[string]PregelNode,
|
||||
channels map[string]BaseChannel,
|
||||
managed ManagedValueMapping,
|
||||
cfg RunnableConfig,
|
||||
checkporter BaseCheckpointSaver,
|
||||
store BaseStore,
|
||||
) *PregelLoop {
|
||||
c, cancel := context.WithCancel(ctx)
|
||||
// Ensure checkpoint is properly initialized
|
||||
if checkpoint.ChannelVersions == nil {
|
||||
checkpoint = NewCheckpoint()
|
||||
}
|
||||
loop := &PregelLoop{
|
||||
ctx: c,
|
||||
cancel: cancel,
|
||||
checkpoint: checkpoint,
|
||||
processes: processes,
|
||||
channels: channels,
|
||||
managed: managed,
|
||||
cfg: cfg,
|
||||
checkporter: checkporter,
|
||||
store: store,
|
||||
step: 0,
|
||||
stop: cfg.RecursionLimit,
|
||||
workers: cfg.MaxConcurrency,
|
||||
pendingWrites: make([]WriteRecord, 0, 16),
|
||||
tasks: map[string]*PregelExecutableTask{},
|
||||
}
|
||||
if loop.workers <= 0 {
|
||||
loop.workers = 1
|
||||
}
|
||||
return loop
|
||||
}
|
||||
|
||||
// Run blocks until completion (or first error)
|
||||
func (l *PregelLoop) Run() error {
|
||||
defer l.cancel()
|
||||
|
||||
for {
|
||||
more, err := l.tick(nil)
|
||||
if err != nil {
|
||||
if errors.As(err, &GraphInterrupt{}) {
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
}
|
||||
if !more {
|
||||
break
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// tick executes a single iteration of the Pregel loop.
|
||||
// Returns true if more iterations are needed, false if done.
|
||||
func (l *PregelLoop) tick(inputKeys []string) (bool, error) {
|
||||
// TODO: Use inputKeys to get the first values.
|
||||
// Check if we need to evaluate interrupts before execution
|
||||
if err := l.evaluateInterrupt("before"); err != nil {
|
||||
return false, err
|
||||
}
|
||||
|
||||
// Build tasks
|
||||
tasks, err := PrepareNextTasks(
|
||||
l.ctx,
|
||||
l.checkpoint,
|
||||
convertPending(l.pendingWrites),
|
||||
l.processes,
|
||||
l.channels,
|
||||
l.managed,
|
||||
l.cfg,
|
||||
l.step,
|
||||
true,
|
||||
l.store,
|
||||
l.checkporter,
|
||||
l.triggerToNodes,
|
||||
l.updatedChans,
|
||||
)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
if len(tasks) == 0 {
|
||||
return false, nil // done, no more tasks
|
||||
}
|
||||
l.tasks = make(map[string]*PregelExecutableTask)
|
||||
for k, v := range tasks {
|
||||
te := v.(PregelExecutableTask)
|
||||
l.tasks[k] = &te
|
||||
}
|
||||
|
||||
// parallel execute
|
||||
workCh := make(chan *PregelExecutableTask)
|
||||
errCh := make(chan error, l.workers)
|
||||
|
||||
for i := 0; i < l.workers; i++ {
|
||||
go l.worker(workCh, errCh)
|
||||
}
|
||||
|
||||
for _, t := range l.tasks {
|
||||
if len(t.Writes) > 0 {
|
||||
continue // already satisfied
|
||||
}
|
||||
workCh <- t
|
||||
}
|
||||
close(workCh)
|
||||
|
||||
for i := 0; i < l.workers; i++ {
|
||||
if err := <-errCh; err != nil {
|
||||
return false, err
|
||||
}
|
||||
}
|
||||
|
||||
// All tasks finished; apply writes
|
||||
if err := l.applyWrites(); err != nil {
|
||||
return false, err
|
||||
}
|
||||
// checkpoint
|
||||
if err := l.saveCheckpoint(); err != nil {
|
||||
return false, err
|
||||
}
|
||||
|
||||
// Check if we need to evaluate interrupts after execution
|
||||
if err := l.evaluateInterrupt("after"); err != nil {
|
||||
return false, err
|
||||
}
|
||||
|
||||
// Check if we've exceeded the recursion limit
|
||||
l.step++
|
||||
if l.step > l.stop {
|
||||
return false, fmt.Errorf("exceeded recursion limit (%d)", l.stop)
|
||||
}
|
||||
|
||||
return true, nil
|
||||
}
|
||||
|
||||
// prepareAndExecuteStep is kept for backward compatibility
|
||||
func (l *PregelLoop) prepareAndExecuteStep() error {
|
||||
more, err := l.tick(nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !more {
|
||||
return nil
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (l *PregelLoop) worker(in <-chan *PregelExecutableTask, out chan<- error) {
|
||||
for task := range in {
|
||||
err := l.runTask(task)
|
||||
out <- err
|
||||
}
|
||||
}
|
||||
|
||||
func (l *PregelLoop) runTask(t *PregelExecutableTask) error {
|
||||
// retry loop
|
||||
attempts := 0
|
||||
max := 1
|
||||
if p, ok := l.processes[t.Name]; ok {
|
||||
max = maxAttempts(p.Retry)
|
||||
}
|
||||
for {
|
||||
attempts++
|
||||
select {
|
||||
case <-l.ctx.Done():
|
||||
return l.ctx.Err()
|
||||
default:
|
||||
}
|
||||
|
||||
writes, err := t.Node.Invoke(l.ctx, t.Input, t.Config, l)
|
||||
if err == nil {
|
||||
for _, w := range writes {
|
||||
l.recordWrite(t.ID, w.Channel, w.Value)
|
||||
}
|
||||
t.Writes = writes
|
||||
return nil
|
||||
}
|
||||
|
||||
if attempts >= max {
|
||||
return err
|
||||
}
|
||||
time.Sleep(backoffDelay(attempts))
|
||||
}
|
||||
}
|
||||
|
||||
// putWrites is called by PregelRunner (or nested tasks via the SEND helper)
|
||||
// to persist writes produced by a task *during the current super-step*.
|
||||
// It is safe for concurrent use.
|
||||
func (l *PregelLoop) putWrites(taskID string, writes []Write) {
|
||||
if len(writes) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------
|
||||
// 1. Deduplicate if every write is for a “special” indexed channel.
|
||||
// (“last one wins”, exactly like in TS / Python)
|
||||
// ---------------------------------------------------------------------
|
||||
allIndexed := true
|
||||
for _, w := range writes {
|
||||
if _, ok := WRITES_IDX_MAP[w.Channel]; !ok {
|
||||
allIndexed = false
|
||||
break
|
||||
}
|
||||
}
|
||||
if allIndexed {
|
||||
dedup := make(map[string]Write, len(writes))
|
||||
for _, w := range writes {
|
||||
dedup[w.Channel] = w
|
||||
}
|
||||
writes = make([]Write, 0, len(dedup))
|
||||
for _, w := range dedup {
|
||||
writes = append(writes, w)
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------
|
||||
// 2. Merge into l.checkpointPendingWrites.
|
||||
// We need a mutex because PregelRunner goroutines call us in parallel.
|
||||
// ---------------------------------------------------------------------
|
||||
l.pendingMu.Lock()
|
||||
for _, w := range writes {
|
||||
replaced := false
|
||||
|
||||
// If it is an indexed channel and an entry already exists for (task,channel),
|
||||
// overwrite it (=> keep only the newest write).
|
||||
if _, special := WRITES_IDX_MAP[w.Channel]; special {
|
||||
for i := range l.checkpointPendingWrites {
|
||||
pw := &l.checkpointPendingWrites[i]
|
||||
if pw.TaskID == taskID && pw.Channel == w.Channel {
|
||||
pw.Value = w.Value
|
||||
replaced = true
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Otherwise (or if not found) just append.
|
||||
if !replaced {
|
||||
l.checkpointPendingWrites = append(
|
||||
l.checkpointPendingWrites,
|
||||
PendingWrite{TaskID: taskID, Channel: w.Channel, Value: w.Value},
|
||||
)
|
||||
}
|
||||
}
|
||||
l.pendingMu.Unlock()
|
||||
|
||||
// ---------------------------------------------------------------------
|
||||
// 3. Forward the writes to the configured checkpointer (if any).
|
||||
// We don’t block the caller – a quick “fire-and-forget” goroutine
|
||||
// is fine because checkpointer.PutWrites() is thread-safe by design.
|
||||
// ---------------------------------------------------------------------
|
||||
// if l.checkpointer != nil {
|
||||
// cfg := l.checkpointConfig // shallow copy is enough – we never mutate it
|
||||
// go l.checkpointer.PutWrites(cfg, writes, taskID)
|
||||
// }
|
||||
|
||||
// ---------------------------------------------------------------------
|
||||
// 4. Emit stream/debug output if the loop is already running.
|
||||
// ---------------------------------------------------------------------
|
||||
if len(l.tasks) > 0 {
|
||||
l.outputWrites(taskID, writes, false)
|
||||
}
|
||||
}
|
||||
|
||||
// outputWrites mirrors TS _outputWrites (omits hidden tasks & handles modes).
|
||||
// This is a *minimal* version; extend if you need streaming/debug UI parity.
|
||||
func (l *PregelLoop) outputWrites(taskID string, writes []Write, cached bool) {
|
||||
task, ok := l.tasks[taskID]
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
for _, tag := range task.Config.Tags {
|
||||
if tag == TAG_HIDDEN {
|
||||
return
|
||||
}
|
||||
}
|
||||
// TODO: implement streaming
|
||||
// delegate to whatever streaming mechanism you implemented…
|
||||
// if l.emit != nil {
|
||||
// l.emit(task, writes, cached)
|
||||
// }
|
||||
}
|
||||
|
||||
func maxAttempts(r RetryPolicy) int {
|
||||
if r.MaxAttempts <= 0 {
|
||||
return 1
|
||||
}
|
||||
return r.MaxAttempts
|
||||
}
|
||||
|
||||
func backoffDelay(at int) time.Duration { return time.Duration(at) * 50 * time.Millisecond }
|
||||
|
||||
func (l *PregelLoop) Send(taskID string, writes []Write) {
|
||||
for _, w := range writes {
|
||||
l.recordWrite(taskID, w.Channel, w.Value)
|
||||
}
|
||||
}
|
||||
|
||||
// Read returns a copy of current channel values
|
||||
func (l *PregelLoop) Read(selectKeys []string) map[string]any {
|
||||
out := map[string]any{}
|
||||
for _, k := range selectKeys {
|
||||
if ch, ok := l.channels[k]; ok {
|
||||
out[k] = ch.Get()
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (l *PregelLoop) AcceptPush(origin PregelExecutableTask, writeIdx int, call *Call) (*PregelExecutableTask, error) {
|
||||
ppath := origin.Path
|
||||
newPath := []interface{}{PUSH, ppath, writeIdx, origin.ID, call}
|
||||
cpid, _ := hex.DecodeString(l.checkpoint.ID)
|
||||
nullVer := -1
|
||||
task, err := PrepareSingleTask(
|
||||
l.ctx,
|
||||
newPath,
|
||||
"",
|
||||
l.checkpoint,
|
||||
cpid,
|
||||
nullVer,
|
||||
convertPending(l.pendingWrites),
|
||||
l.processes,
|
||||
l.channels,
|
||||
l.managed,
|
||||
l.cfg,
|
||||
l.step,
|
||||
true,
|
||||
l.store,
|
||||
l.checkporter,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if task == nil {
|
||||
return nil, nil
|
||||
}
|
||||
te := task.(PregelExecutableTask)
|
||||
l.tasks[te.ID] = &te
|
||||
return &te, nil
|
||||
}
|
||||
|
||||
func (l *PregelLoop) recordWrite(taskID, ch string, val any) {
|
||||
l.pendingWrites = append(l.pendingWrites, WriteRecord{taskID, ch, val})
|
||||
}
|
||||
|
||||
func convertPending(ws []WriteRecord) []interface{} {
|
||||
out := make([]interface{}, 0, len(ws))
|
||||
for _, w := range ws {
|
||||
out = append(out, []interface{}{w.Task, w.Chan, w.Value})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (l *PregelLoop) applyWrites() error {
|
||||
if len(l.pendingWrites) == 0 {
|
||||
return nil
|
||||
}
|
||||
for _, wr := range l.pendingWrites {
|
||||
ch, ok := l.channels[wr.Chan]
|
||||
if !ok {
|
||||
ch = &simpleChan{}
|
||||
l.channels[wr.Chan] = ch
|
||||
}
|
||||
ch.Set(wr.Value)
|
||||
// TODO: Handle other version types.
|
||||
if _, exists := l.checkpoint.ChannelVersions[wr.Chan]; !exists {
|
||||
l.checkpoint.ChannelVersions[wr.Chan] = 0
|
||||
}
|
||||
l.checkpoint.ChannelVersions[wr.Chan]++
|
||||
}
|
||||
l.pendingWrites = l.pendingWrites[:0]
|
||||
return nil
|
||||
}
|
||||
|
||||
func (l *PregelLoop) saveCheckpoint() error {
|
||||
if l.checkporter == nil {
|
||||
return nil
|
||||
}
|
||||
md := map[string]any{
|
||||
"step": l.step,
|
||||
"source": "loop",
|
||||
"time": time.Now().UTC().Format(time.RFC3339Nano),
|
||||
}
|
||||
return l.checkporter.Put(l.cfg, l.checkpoint, md, nil)
|
||||
}
|
||||
|
||||
func (l *PregelLoop) evaluateInterrupt(stage string) error {
|
||||
var conditions []string
|
||||
if stage == "before" {
|
||||
conditions = l.interruptBefore
|
||||
} else {
|
||||
conditions = l.interruptAfter
|
||||
}
|
||||
if len(conditions) == 0 {
|
||||
return nil
|
||||
}
|
||||
seen := map[string]struct{}{}
|
||||
for _, t := range l.tasks {
|
||||
for _, trg := range t.Triggers {
|
||||
seen[trg] = struct{}{}
|
||||
}
|
||||
}
|
||||
for _, cond := range conditions {
|
||||
if _, ok := seen[cond]; ok || cond == "*" {
|
||||
return GraphInterrupt{}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type Result struct {
|
||||
Err error
|
||||
}
|
||||
@@ -0,0 +1,217 @@
|
||||
// runner.go
|
||||
package pregel
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"golang.org/x/sync/errgroup"
|
||||
)
|
||||
|
||||
// PregelRunner is responsible for executing the set of tasks that a
|
||||
// PregelLoop prepared for the *current super-step*. It runs them with
|
||||
// respect to retry-policy, max-concurrency, timeouts, cancellation and
|
||||
// Pregel-specific error semantics (GraphInterrupt / GraphBubbleUp).
|
||||
type PregelRunner struct {
|
||||
loop *PregelLoop
|
||||
nodeFinished func(string) // Optional user-callback
|
||||
}
|
||||
|
||||
// NewPregelRunner links the runner to its parent loop.
|
||||
func NewPregelRunner(loop *PregelLoop, nodeFinished func(string)) *PregelRunner {
|
||||
return &PregelRunner{loop: loop, nodeFinished: nodeFinished}
|
||||
}
|
||||
|
||||
// TickOptions mirrors the semantics in the TS/Python implementations.
|
||||
type TickOptions struct {
|
||||
Timeout time.Duration // Deadline for the whole super-step
|
||||
RetryPolicy RetryPolicy // Per-task retry policy
|
||||
OnStepWrite func(int, []Write) // Hook after *all* writes are committed
|
||||
MaxConcurrency int // ≤0 ⇒ unlimited
|
||||
Ctx context.Context // Root ctx (optional)
|
||||
}
|
||||
|
||||
// Tick executes every task whose Writes slice is still empty.
|
||||
// It returns when *all* tasks have completed (successfully or not) **or**
|
||||
// when the first non-interrupt error bubbles up.
|
||||
func (r *PregelRunner) tick(opt TickOptions) error {
|
||||
// Choose base context
|
||||
ctx := opt.Ctx
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
// We cancel siblings on first fatal error
|
||||
ctx, cancel := context.WithCancel(ctx)
|
||||
defer cancel()
|
||||
|
||||
// Optional global timeout
|
||||
if opt.Timeout > 0 {
|
||||
ctx, cancel = context.WithTimeout(ctx, opt.Timeout)
|
||||
defer cancel()
|
||||
}
|
||||
|
||||
// Gather tasks that still need to run in this super-step
|
||||
var pending []*PregelExecutableTask
|
||||
for _, t := range r.loop.tasks {
|
||||
if len(t.Writes) == 0 {
|
||||
pending = append(pending, t)
|
||||
}
|
||||
}
|
||||
if len(pending) == 0 {
|
||||
return nil // nothing to do
|
||||
}
|
||||
|
||||
// errgroup manages goroutines and collects the first returned error
|
||||
g, gctx := errgroup.WithContext(ctx)
|
||||
|
||||
maxConc := opt.MaxConcurrency
|
||||
if maxConc <= 0 {
|
||||
maxConc = len(pending)
|
||||
}
|
||||
sem := make(chan struct{}, maxConc)
|
||||
|
||||
var mu sync.Mutex
|
||||
|
||||
for _, task := range pending {
|
||||
task := task // capture
|
||||
sem <- struct{}{}
|
||||
g.Go(func() error {
|
||||
defer func() { <-sem }()
|
||||
|
||||
err := runWithRetry(gctx, opt.RetryPolicy, func(c context.Context) error {
|
||||
// NOTE: Node.Run must honour ctx for cancellation / deadlines.
|
||||
writes, runErr := task.Node.Invoke(c, task.Input, task.Config, r.loop)
|
||||
if runErr == nil {
|
||||
task.Writes = writes
|
||||
}
|
||||
return runErr
|
||||
})
|
||||
|
||||
r.commit(task, err)
|
||||
|
||||
switch {
|
||||
case err == nil:
|
||||
return nil
|
||||
case errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded):
|
||||
return err // propagate
|
||||
}
|
||||
|
||||
var gi GraphInterrupt
|
||||
if errors.As(err, &gi) {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
// kep track so that loop can raise combined interrupt later
|
||||
return gi
|
||||
}
|
||||
|
||||
cancel()
|
||||
return err
|
||||
})
|
||||
}
|
||||
|
||||
// Wait for all goroutines (or first fatal error)
|
||||
if err := g.Wait(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Step-level callback after *all* commits
|
||||
if opt.OnStepWrite != nil {
|
||||
var all []Write
|
||||
for _, t := range r.loop.tasks {
|
||||
all = append(all, t.Writes...)
|
||||
}
|
||||
opt.OnStepWrite(r.loop.step, all)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// commit replicates the Python/TS commit semantics.
|
||||
func (r *PregelRunner) commit(task *PregelExecutableTask, execErr error) {
|
||||
// On success ensure at least one NO_WRITES marker so loop knows it's done.
|
||||
if execErr == nil && len(task.Writes) == 0 {
|
||||
task.Writes = append(task.Writes, Write{Channel: NO_WRITES})
|
||||
}
|
||||
|
||||
// Persist writes (or error) through the loop’s thread-safe adaptor.
|
||||
switch {
|
||||
case execErr == nil:
|
||||
r.loop.putWrites(task.ID, task.Writes)
|
||||
case errors.As(execErr, new(GraphInterrupt)):
|
||||
// Interrupt carries its own writes payload
|
||||
r.loop.putWrites(task.ID, task.Writes)
|
||||
default:
|
||||
// Record generic error
|
||||
r.loop.putWrites(task.ID, []Write{{Channel: ERROR, Value: execErr}})
|
||||
}
|
||||
|
||||
// optional callback
|
||||
if execErr == nil && r.nodeFinished != nil {
|
||||
r.nodeFinished(task.Name)
|
||||
}
|
||||
}
|
||||
|
||||
// runWithRetry is a minimal exponential-back-off retry helper.
|
||||
func runWithRetry(ctx context.Context, pol RetryPolicy, fn func(context.Context) error) error {
|
||||
if pol.MaxAttempts <= 0 {
|
||||
pol.MaxAttempts = 1
|
||||
}
|
||||
// if pol.Backoff == nil {
|
||||
// // default: exponential capped at 2 s
|
||||
// pol.Backoff = func(attempt int) time.Duration {
|
||||
// d := time.Duration(math.Pow(2, float64(attempt))) * 50 * time.Millisecond
|
||||
// if d > 2*time.Second {
|
||||
// d = 2 * time.Second
|
||||
// }
|
||||
// return d
|
||||
// }
|
||||
// }
|
||||
// if pol.Retryable == nil {
|
||||
// pol.Retryable = func(error) bool { return true }
|
||||
// }
|
||||
|
||||
var err error
|
||||
for attempt := 0; attempt < pol.MaxAttempts; attempt++ {
|
||||
if err = fn(ctx); err == nil { // || !pol.Retryable(err) {
|
||||
return err
|
||||
}
|
||||
// // wait before next try
|
||||
// wait := pol.Backoff(attempt)
|
||||
// select {
|
||||
// case <-time.After(wait):
|
||||
// case <-ctx.Done():
|
||||
// return ctx.Err()
|
||||
// }
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
/* --------------------------------------------------------------------------
|
||||
Missing symbols? If your project does not yet declare the following items
|
||||
just add minimal stubs like the ones below (remove before wiring in
|
||||
real implementations to avoid duplicates).
|
||||
|
||||
// Constants that mark write types
|
||||
const (
|
||||
ERROR = "error"
|
||||
NO_WRITES = "no_writes"
|
||||
)
|
||||
|
||||
// GraphInterrupt / BubbleUp marker errors
|
||||
type GraphInterrupt struct{ Msg string }
|
||||
func (g GraphInterrupt) Error() string { return g.Msg }
|
||||
type GraphBubbleUp struct{ error }
|
||||
|
||||
// Minimal Write + RetryPolicy
|
||||
type Write struct{ Channel string; Value any }
|
||||
|
||||
type RetryPolicy struct {
|
||||
MaxAttempts int
|
||||
Backoff func(attempt int) time.Duration
|
||||
Retryable func(error) bool
|
||||
}
|
||||
|
||||
// PregelExecutableTask, PregelLoop, etc. should exist elsewhere.
|
||||
// -------------------------------------------------------------------------- */
|
||||
@@ -0,0 +1,292 @@
|
||||
package pregel
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// Constants for task types and reserved keys
|
||||
const (
|
||||
// Task types
|
||||
PUSH = "__pregel_push" // Denotes push-style tasks, ie. those created by Send objects
|
||||
PULL = "__pregel_pull" // Denotes pull-style tasks, ie. those triggered by edges
|
||||
|
||||
// Reserved write keys
|
||||
INPUT = "__input__" // For values passed as input to the graph
|
||||
INTERRUPT = "__interrupt__" // For dynamic interrupts raised by nodes
|
||||
RESUME = "__resume__" // For values passed to resume a node after an interrupt
|
||||
ERROR = "__error__" // For errors raised by nodes
|
||||
NO_WRITES = "__no_writes__" // Marker to signal node didn't write anything
|
||||
SCHEDULED = "__scheduled__" // Marker to signal node was scheduled (in distributed mode)
|
||||
TASKS = "__pregel_tasks" // For Send objects returned by nodes/edges
|
||||
RETURN = "__return__" // For writes of a task where we simply record the return value
|
||||
|
||||
// Public constants
|
||||
START = "__start__" // The first (maybe virtual) node in graph-style Pregel
|
||||
END = "__end__" // The last (maybe virtual) node in graph-style Pregel
|
||||
SELF = "__self__" // The implicit branch that handles each node's Control values
|
||||
PREVIOUS = "__previous__" // Previous value
|
||||
|
||||
// Other constants
|
||||
NS_SEP = "|" // For checkpoint_ns, separates each level (ie. graph|subgraph|subsubgraph)
|
||||
NS_END = ":" // For checkpoint_ns, for each level, separates the namespace from the task_id
|
||||
NULL_TASK_ID = "00000000-0000-0000-0000-000000000000" // The task_id to use for writes that are not associated with a task
|
||||
CONF = "configurable" // Key for the configurable dict in RunnableConfig
|
||||
|
||||
// Reserved config.configurable keys
|
||||
CONFIG_KEY_SEND = "__pregel_send" // Holds the `write` function that accepts writes to state/edges/reserved keys
|
||||
CONFIG_KEY_READ = "__pregel_read" // Holds the `read` function that returns a copy of the current state
|
||||
CONFIG_KEY_CALL = "__pregel_call" // Holds the `call` function that accepts a node/func, args and returns a future
|
||||
CONFIG_KEY_CHECKPOINTER = "__pregel_checkpointer" // Holds a `BaseCheckpointSaver` passed from parent graph to child graphs
|
||||
CONFIG_KEY_STREAM = "__pregel_stream" // Holds a `StreamProtocol` passed from parent graph to child graphs
|
||||
CONFIG_KEY_STREAM_WRITER = "__pregel_stream_writer" // Holds a `StreamWriter` for stream_mode=custom
|
||||
CONFIG_KEY_STORE = "__pregel_store" // Holds a `BaseStore` made available to managed values
|
||||
CONFIG_KEY_CACHE = "__pregel_cache" // Holds a `BaseCache` made available to subgraphs
|
||||
CONFIG_KEY_RESUMING = "__pregel_resuming" // Holds a boolean indicating if subgraphs should resume from a previous checkpoint
|
||||
CONFIG_KEY_TASK_ID = "__pregel_task_id" // Holds the task ID for the current task
|
||||
CONFIG_KEY_DEDUPE_TASKS = "__pregel_dedupe_tasks" // Holds a boolean indicating if tasks should be deduplicated (for distributed mode)
|
||||
CONFIG_KEY_ENSURE_LATEST = "__pregel_ensure_latest" // Holds a boolean indicating whether to assert the requested checkpoint is the latest
|
||||
CONFIG_KEY_DELEGATE = "__pregel_delegate" // Holds a boolean indicating whether to delegate subgraphs (for distributed mode)
|
||||
CONFIG_KEY_THREAD_ID = "thread_id" // Holds the thread ID for the current invocation
|
||||
CONFIG_KEY_CHECKPOINT_MAP = "checkpoint_map" // Holds a mapping of checkpoint_ns -> checkpoint_id for parent graphs
|
||||
CONFIG_KEY_CHECKPOINT_ID = "checkpoint_id" // Holds the current checkpoint_id, if any
|
||||
CONFIG_KEY_CHECKPOINT_NS = "checkpoint_ns" // Holds the current checkpoint_ns, "" for root graph
|
||||
CONFIG_KEY_NODE_FINISHED = "__pregel_node_finished" // Holds a callback to be called when a node is finished
|
||||
CONFIG_KEY_SCRATCHPAD = "__pregel_scratchpad" // Holds a mutable dict for temporary storage scoped to the current task
|
||||
CONFIG_KEY_PREVIOUS = "__pregel_previous" // Holds the previous return value from a stateful Pregel graph
|
||||
CONFIG_KEY_RUNNER_SUBMIT = "__pregel_runner_submit" // Holds a function that receives tasks from runner, executes them and returns results
|
||||
CONFIG_KEY_CHECKPOINT_DURING = "__pregel_checkpoint_during" // Holds a boolean indicating whether to checkpoint during the run (or only at the end)
|
||||
CONFIG_KEY_RESUME_MAP = "__pregel_resume_map" // Holds a mapping of task ns -> resume value for resuming tasks
|
||||
TAG_HIDDEN = "langsmith:hidden" // Holds a boolean indicating whether to hide a node/edge from certain tracing/streaming environments.
|
||||
)
|
||||
|
||||
// StreamMode defines how the graph streams its output
|
||||
type StreamMode string
|
||||
|
||||
// WRITES_IDX_MAP maps special channel names to negative indices
|
||||
// to avoid conflicts with regular writes.
|
||||
var WRITES_IDX_MAP = map[string]int{
|
||||
ERROR: -1,
|
||||
SCHEDULED: -2,
|
||||
INTERRUPT: -3,
|
||||
RESUME: -4,
|
||||
}
|
||||
|
||||
// TS
|
||||
// export type PendingWriteValue = unknown;
|
||||
|
||||
// export type PendingWrite<Channel = string> = [Channel, PendingWriteValue];
|
||||
|
||||
// export type CheckpointPendingWrite<TaskId = string> = [
|
||||
// TaskId,
|
||||
// ...PendingWrite<string>
|
||||
// ];
|
||||
// Py
|
||||
// PendingWrite = Tuple[str, str, Any]
|
||||
|
||||
type PendingWrite struct {
|
||||
TaskID string
|
||||
Channel string
|
||||
Value interface{}
|
||||
}
|
||||
|
||||
const (
|
||||
// StreamValues emits all values in the state after each step
|
||||
StreamValues StreamMode = "values"
|
||||
// StreamUpdates emits only the node or task names and updates
|
||||
StreamUpdates StreamMode = "updates"
|
||||
// StreamCustom emits custom data from inside nodes or tasks
|
||||
StreamCustom StreamMode = "custom"
|
||||
// StreamMessages emits LLM messages token-by-token
|
||||
StreamMessages StreamMode = "messages"
|
||||
// StreamDebug emits debug events with as much information as possible
|
||||
StreamDebug StreamMode = "debug"
|
||||
)
|
||||
|
||||
// PregelTask represents a task in the Pregel system
|
||||
|
||||
type PregelTask struct {
|
||||
ID string
|
||||
Name string
|
||||
Path []interface{}
|
||||
Error error
|
||||
Interrupts []interface{}
|
||||
Result interface{}
|
||||
}
|
||||
|
||||
// PregelExecutableTask represents a task that can be executed
|
||||
type PregelExecutableTask struct {
|
||||
PregelTask
|
||||
Input interface{}
|
||||
Node NodeRunnable
|
||||
Writes []Write
|
||||
Config RunnableConfig
|
||||
Triggers []string
|
||||
RetryPolicy interface{}
|
||||
CacheKey *CacheKey
|
||||
Writers map[string]interface{} // Flat writers
|
||||
Subgraphs map[string]interface{} // Subgraphs
|
||||
}
|
||||
|
||||
// StreamChunk is what the consumer receives.
|
||||
type StreamChunk struct {
|
||||
Namespace []string // sub-graph path (reserved for future use)
|
||||
Mode StreamMode
|
||||
Payload any
|
||||
}
|
||||
|
||||
type StreamOptions struct {
|
||||
Mode StreamMode
|
||||
OutputChannels []string // defaults to all non-context channels
|
||||
InterruptBefore []string // interrupt gate (before)
|
||||
InterruptAfter []string // interrupt gate (after)
|
||||
MaxConcurrency int // overrides config[ "max_concurrency" ]
|
||||
CheckpointDuring *bool // nil → inherit config
|
||||
Debug *bool // nil → inherit graph.debug
|
||||
Context context.Context // optional, default = context.Background()
|
||||
}
|
||||
|
||||
// CacheKey represents a key for caching
|
||||
type CacheKey struct {
|
||||
Namespace []string
|
||||
Key string
|
||||
TTL int64
|
||||
}
|
||||
|
||||
type PregelNode struct {
|
||||
Node NodeRunnable
|
||||
Triggers []string
|
||||
Metadata map[string]interface{}
|
||||
Tags []string
|
||||
CachePolicy interface{} // CachePolicy equivalent
|
||||
RetryPolicy interface{} // RetryPolicy equivalent
|
||||
FlatWriters map[string]interface{}
|
||||
Subgraphs map[string]interface{}
|
||||
Retry RetryPolicy
|
||||
}
|
||||
|
||||
type NodeRunnable interface {
|
||||
Invoke(ctx context.Context, input any, cfg RunnableConfig, loop LoopCallback) ([]Write, error)
|
||||
}
|
||||
|
||||
type Write struct {
|
||||
Channel string
|
||||
Value any
|
||||
}
|
||||
|
||||
// Checkpoint represents a checkpoint in the Pregel system
|
||||
type Checkpoint struct {
|
||||
ID string
|
||||
ChannelValues map[string]interface{} `json:"channel_values,omitempty"`
|
||||
ChannelVersions map[string]int64 `json:"channel_versions,omitempty"`
|
||||
VersionsSeen map[string]interface{} `json:"versions_seen,omitempty"`
|
||||
PendingSends []Send `json:"pending_sends,omitempty"`
|
||||
Version int `json:"version,omitempty"`
|
||||
Timestamp string `json:"timestamp,omitempty"`
|
||||
}
|
||||
|
||||
// NewCheckpoint creates a new Checkpoint with all fields properly initialized
|
||||
func NewCheckpoint() Checkpoint {
|
||||
return Checkpoint{
|
||||
ChannelValues: make(map[string]interface{}),
|
||||
ChannelVersions: make(map[string]int64),
|
||||
VersionsSeen: make(map[string]interface{}),
|
||||
PendingSends: make([]Send, 0),
|
||||
}
|
||||
}
|
||||
|
||||
// Send represents a message to be sent to a node
|
||||
type Send struct {
|
||||
Node string
|
||||
Arg interface{}
|
||||
}
|
||||
|
||||
// Call represents a function call
|
||||
type Call struct {
|
||||
Func interface{} // Function to call
|
||||
Input []interface{} // Arguments
|
||||
Callbacks interface{} // Callbacks
|
||||
CachePolicy interface{} // CachePolicy
|
||||
Retry interface{} // RetryPolicy
|
||||
}
|
||||
|
||||
// PregelTaskWrites represents writes from a task
|
||||
type PregelTaskWrites struct {
|
||||
Path []interface{}
|
||||
Name string
|
||||
Writes []interface{} // Deque in Python
|
||||
Triggers []string
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Interfaces from previous snippets (slim versions here)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
type RetryPolicy struct {
|
||||
MaxAttempts int
|
||||
BackoffMs int
|
||||
}
|
||||
|
||||
type BaseChannel interface {
|
||||
Set(v any)
|
||||
Get() any
|
||||
}
|
||||
|
||||
type simpleChan struct{ val atomicValue }
|
||||
|
||||
type atomicValue struct {
|
||||
mu sync.RWMutex
|
||||
v any
|
||||
}
|
||||
|
||||
func (a *atomicValue) Store(v any) {
|
||||
a.mu.Lock()
|
||||
a.v = v
|
||||
a.mu.Unlock()
|
||||
}
|
||||
func (a *atomicValue) Load() (v any) { a.mu.RLock(); v = a.v; a.mu.RUnlock(); return }
|
||||
|
||||
func (c *simpleChan) Set(v any) { c.val.Store(v) }
|
||||
func (c *simpleChan) Get() any { return c.val.Load() }
|
||||
|
||||
// Managed values -------------------------------------------------------------
|
||||
|
||||
type WritableManagedValue interface {
|
||||
Update([]any) error
|
||||
}
|
||||
|
||||
type ManagedValueMapping map[string]WritableManagedValue
|
||||
|
||||
type SendPacket struct {
|
||||
Node string
|
||||
Arg any
|
||||
}
|
||||
|
||||
type BaseCheckpointSaver interface {
|
||||
Put(cfg RunnableConfig, cp Checkpoint, md map[string]any, newVers map[string]int) error
|
||||
GetTuple(cfg RunnableConfig) (*Checkpoint, error)
|
||||
}
|
||||
|
||||
// Stores ---------------------------------------------------------------------
|
||||
|
||||
type BaseStore interface{}
|
||||
|
||||
// Loop callback interface passed to Nodes for localWrite / localRead
|
||||
type LoopCallback interface {
|
||||
Send(taskID string, writes []Write)
|
||||
Read(selectKeys []string) map[string]any
|
||||
AcceptPush(originTask PregelExecutableTask, writeIdx int, call *Call) (*PregelExecutableTask, error)
|
||||
}
|
||||
|
||||
// RunnableConfig represents configuration for a Runnable.
|
||||
// Fields are optional
|
||||
type RunnableConfig struct {
|
||||
Tags []string `json:"tags,omitempty"` // Tags for this call and sub-calls.
|
||||
Metadata map[string]interface{} `json:"metadata,omitempty"` // Metadata for this call and sub-calls.
|
||||
Callbacks interface{} `json:"callbacks,omitempty"` // Callbacks for this call and sub-calls.
|
||||
RunName *string `json:"run_name,omitempty"` // Name for the tracer run for this call.
|
||||
MaxConcurrency int `json:"max_concurrency,omitempty"` // Max number of parallel calls.
|
||||
RecursionLimit int `json:"recursion_limit,omitempty"` // Max recursion depth.
|
||||
Configurable map[string]interface{} `json:"configurable,omitempty"` // Runtime values for configurable attributes.
|
||||
RunID *string `json:"run_id,omitempty"` // Unique identifier for the tracer run (UUID as string).
|
||||
}
|
||||
Reference in New Issue
Block a user