From 103565189e9f7fd6e4dbedb4bcbf4277f0bccf24 Mon Sep 17 00:00:00 2001 From: mrh Date: Mon, 30 Dec 2024 11:51:32 +0800 Subject: [PATCH] feat(devops): remove NewStateGraph call --- devops/internal/model/container.go | 7 +++-- devops/internal/model/container_test.go | 35 +++++++++++++------------ 2 files changed, 21 insertions(+), 21 deletions(-) diff --git a/devops/internal/model/container.go b/devops/internal/model/container.go index b8eeb80..a29e496 100644 --- a/devops/internal/model/container.go +++ b/devops/internal/model/container.go @@ -73,10 +73,9 @@ func (gi GraphInfo) BuildDevGraph(fromNode string) (g *Graph, err error) { } if gi.Option.GenState != nil { - g = &Graph{StateGraph: compose.NewStateGraph[any, any, any](gi.Option.GenState)} + g = &Graph{Graph: compose.NewGraph[any, any](compose.WithGenLocalState(gi.Option.GenState))} } else { - genState := func(ctx context.Context) any { return nil } - g = &Graph{StateGraph: compose.NewStateGraph[any, any, any](genState)} + g = &Graph{Graph: compose.NewGraph[any, any]()} } var ( @@ -367,7 +366,7 @@ func (gi GraphInfo) buildSubGraphSchema() (subGraphSchema map[string]*devmodel.G } type Graph struct { - *compose.StateGraph[any, any, any] + *compose.Graph[any, any] } func (g *Graph) Compile(opts ...compose.GraphCompileOption) (Runnable, error) { diff --git a/devops/internal/model/container_test.go b/devops/internal/model/container_test.go index 427c5b9..8076487 100644 --- a/devops/internal/model/container_test.go +++ b/devops/internal/model/container_test.go @@ -25,12 +25,13 @@ import ( "strings" "testing" + "github.com/stretchr/testify/assert" + "github.com/cloudwego/eino/components" "github.com/cloudwego/eino/components/prompt" "github.com/cloudwego/eino/components/retriever" "github.com/cloudwego/eino/compose" "github.com/cloudwego/eino/schema" - "github.com/stretchr/testify/assert" "github.com/cloudwego/eino-ext/devops/internal/utils/generic" devmodel "github.com/cloudwego/eino-ext/devops/model" @@ -125,9 +126,9 @@ func Test_GraphInfo_BuildDevGraph(t *testing.T) { err = g.AddGraphNode("node_5", sc) assert.NoError(t, err) - ssg := compose.NewStateGraph[[]string, []string, string](func(ctx context.Context) (state string) { + ssg := compose.NewGraph[[]string, []string](compose.WithGenLocalState(func(ctx context.Context) (state string) { return "" - }) + })) err = ssg.AddLambdaNode("ssg_node_1", compose.InvokableLambda(func(ctx context.Context, input []string) (output []string, err error) { output = append(input, fmt.Sprintf("sub_state_graph_out_lambda_1")) return output, nil @@ -567,7 +568,7 @@ func Test_GraphInfo_BuildDevGraph(t *testing.T) { genState: genState, } - g := compose.NewStateGraph[int, []string, any](genState) + g := compose.NewGraph[int, []string](compose.WithGenLocalState(genState)) err := g.AddLambdaNode("node_1", compose.InvokableLambda(func(ctx context.Context, input int) (output []string, err error) { return []string{strconv.Itoa(input), fmt.Sprintf("out_lambda_1")}, nil })) @@ -1390,7 +1391,7 @@ func Test_Graph_addNode(t *testing.T) { } t.Run("instance not Embedding", func(t *testing.T) { - g := &Graph{StateGraph: compose.NewStateGraph[any, any, any](genState)} + g := &Graph{Graph: compose.NewGraph[any, any](compose.WithGenLocalState(genState))} gni := compose.GraphNodeInfo{ Component: components.ComponentOfEmbedding, Instance: nil, @@ -1401,7 +1402,7 @@ func Test_Graph_addNode(t *testing.T) { }) t.Run("instance not Retriever", func(t *testing.T) { - g := &Graph{StateGraph: compose.NewStateGraph[any, any, any](genState)} + g := &Graph{Graph: compose.NewGraph[any, any](compose.WithGenLocalState(genState))} gni := compose.GraphNodeInfo{ Component: components.ComponentOfRetriever, Instance: nil, @@ -1412,7 +1413,7 @@ func Test_Graph_addNode(t *testing.T) { }) t.Run("instance not LoaderSplitter", func(t *testing.T) { - g := &Graph{StateGraph: compose.NewStateGraph[any, any, any](genState)} + g := &Graph{Graph: compose.NewGraph[any, any](compose.WithGenLocalState(genState))} gni := compose.GraphNodeInfo{ Component: components.ComponentOfLoaderSplitter, Instance: nil, @@ -1422,7 +1423,7 @@ func Test_Graph_addNode(t *testing.T) { }) t.Run("instance not Indexer", func(t *testing.T) { - g := &Graph{StateGraph: compose.NewStateGraph[any, any, any](genState)} + g := &Graph{Graph: compose.NewGraph[any, any](compose.WithGenLocalState(genState))} gni := compose.GraphNodeInfo{ Component: components.ComponentOfIndexer, Instance: nil, @@ -1432,7 +1433,7 @@ func Test_Graph_addNode(t *testing.T) { }) t.Run("instance not ChatModel", func(t *testing.T) { - g := &Graph{StateGraph: compose.NewStateGraph[any, any, any](genState)} + g := &Graph{Graph: compose.NewGraph[any, any](compose.WithGenLocalState(genState))} gni := compose.GraphNodeInfo{ Component: components.ComponentOfChatModel, Instance: nil, @@ -1442,7 +1443,7 @@ func Test_Graph_addNode(t *testing.T) { }) t.Run("Prompt", func(t *testing.T) { - g := &Graph{StateGraph: compose.NewStateGraph[any, any, any](genState)} + g := &Graph{Graph: compose.NewGraph[any, any](compose.WithGenLocalState(genState))} gni := compose.GraphNodeInfo{ Component: components.ComponentOfPrompt, Instance: &prompt.DefaultChatTemplate{}, @@ -1452,7 +1453,7 @@ func Test_Graph_addNode(t *testing.T) { }) t.Run("ToolsNode", func(t *testing.T) { - g := &Graph{StateGraph: compose.NewStateGraph[any, any, any](genState)} + g := &Graph{Graph: compose.NewGraph[any, any](compose.WithGenLocalState(genState))} gni := compose.GraphNodeInfo{ Component: compose.ComponentOfToolsNode, Instance: &compose.ToolsNode{}, @@ -1462,7 +1463,7 @@ func Test_Graph_addNode(t *testing.T) { }) t.Run("Graph", func(t *testing.T) { - g := &Graph{StateGraph: compose.NewStateGraph[any, any, any](genState)} + g := &Graph{Graph: compose.NewGraph[any, any](compose.WithGenLocalState(genState))} gni := compose.GraphNodeInfo{ Component: compose.ComponentOfStateGraph, Instance: compose.NewGraph[string, string](), @@ -1472,7 +1473,7 @@ func Test_Graph_addNode(t *testing.T) { }) t.Run("Chain", func(t *testing.T) { - g := &Graph{StateGraph: compose.NewStateGraph[any, any, any](genState)} + g := &Graph{Graph: compose.NewGraph[any, any](compose.WithGenLocalState(genState))} gni := compose.GraphNodeInfo{ Component: compose.ComponentOfChain, Instance: compose.NewChain[string, string](), @@ -1481,13 +1482,13 @@ func Test_Graph_addNode(t *testing.T) { assert.NoError(t, err) }) - t.Run("StateGraph", func(t *testing.T) { - g := &Graph{StateGraph: compose.NewStateGraph[any, any, any](genState)} + t.Run("Graph", func(t *testing.T) { + g := &Graph{Graph: compose.NewGraph[any, any](compose.WithGenLocalState(genState))} gni := compose.GraphNodeInfo{ Component: compose.ComponentOfStateGraph, - Instance: compose.NewStateGraph[string, string, string](func(ctx context.Context) string { + Instance: compose.NewGraph[string, string](compose.WithGenLocalState(func(ctx context.Context) string { return "" - }), + })), } err := g.addNode("node_1", gni) assert.NoError(t, err)