From 022b0e59ac1f2117770820f148c5a3aec2da0c5b Mon Sep 17 00:00:00 2001 From: shentongmartin Date: Mon, 20 Jan 2025 21:57:54 +0800 Subject: [PATCH] fix: add some examples (#12) Change-Id: Ia163daffa1364c7a0cfc5f665f4166e6cf5c687d --- compose/graph/simple/graph.go | 100 ++++++++++++++++++++++++ compose/graph/simple/type_match_test.go | 68 ++++++++++++++++ flow/agent/react/react.go | 10 +-- go.mod | 2 + 4 files changed, 175 insertions(+), 5 deletions(-) create mode 100644 compose/graph/simple/graph.go create mode 100644 compose/graph/simple/type_match_test.go diff --git a/compose/graph/simple/graph.go b/compose/graph/simple/graph.go new file mode 100644 index 0000000..9ebafc9 --- /dev/null +++ b/compose/graph/simple/graph.go @@ -0,0 +1,100 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package main + +import ( + "context" + "fmt" + "io" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/prompt" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +const ( + nodeOfModel = "model" + nodeOfPrompt = "prompt" +) + +func main() { + ctx := context.Background() + g := compose.NewGraph[map[string]any, *schema.Message]() + + pt := prompt.FromMessages( + schema.FString, + schema.UserMessage("what's the weather in {location}?"), + ) + + _ = g.AddChatTemplateNode(nodeOfPrompt, pt) // error check can be skipped here because in Compile, we will check the error happened here too + _ = g.AddChatModelNode(nodeOfModel, &mockChatModel{}, compose.WithNodeName("ChatModel")) + _ = g.AddEdge(compose.START, nodeOfPrompt) + _ = g.AddEdge(nodeOfPrompt, nodeOfModel) + _ = g.AddEdge(nodeOfModel, compose.END) + + r, err := g.Compile(ctx, compose.WithMaxRunSteps(10)) + if err != nil { + panic(err) + } + + in := map[string]any{"location": "beijing"} + ret, err := r.Invoke(ctx, in) + if err != nil { + panic(err) + } + fmt.Println("invoke result: ", ret) + + // stream + s, err := r.Stream(ctx, in) + if err != nil { + panic(err) + } + + defer s.Close() + for { + chunk, err := s.Recv() + if err != nil { + if err == io.EOF { + break + } + panic(err) + } + + fmt.Println("stream chunk: ", chunk) + } +} + +type mockChatModel struct{} + +func (m *mockChatModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + return schema.AssistantMessage("the weather is good", nil), nil +} + +func (m *mockChatModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + sr, sw := schema.Pipe[*schema.Message](0) + go func() { + defer sw.Close() + sw.Send(schema.AssistantMessage("the weather is", nil), nil) + sw.Send(schema.AssistantMessage(" good", nil), nil) + }() + return sr, nil +} + +func (m *mockChatModel) BindTools(tools []*schema.ToolInfo) error { + panic("implement me") +} diff --git a/compose/graph/simple/type_match_test.go b/compose/graph/simple/type_match_test.go new file mode 100644 index 0000000..eb799fa --- /dev/null +++ b/compose/graph/simple/type_match_test.go @@ -0,0 +1,68 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package main + +import ( + "context" + "io" + "testing" + + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" + "github.com/stretchr/testify/assert" +) + +func TestTypeMatch(t *testing.T) { + ctx := context.Background() + + g1 := compose.NewGraph[[]*schema.Message, string]() + _ = g1.AddChatModelNode("model", &mockChatModel{}) + _ = g1.AddLambdaNode("lambda", compose.InvokableLambda(func(_ context.Context, msg *schema.Message) (string, error) { + return msg.Content, nil + })) + _ = g1.AddEdge(compose.START, "model") + _ = g1.AddEdge("model", "lambda") + _ = g1.AddEdge("lambda", compose.END) + + runner, err := g1.Compile(ctx) + assert.NoError(t, err) + + c, err := runner.Invoke(ctx, []*schema.Message{ + schema.UserMessage("what's the weather in beijing?"), + }) + assert.NoError(t, err) + assert.Equal(t, "the weather is good", c) + + s, err := runner.Stream(ctx, []*schema.Message{ + schema.UserMessage("what's the weather in beijing?"), + }) + assert.NoError(t, err) + + var fullStr string + for { + chunk, err := s.Recv() + if err != nil { + if err == io.EOF { + break + } + panic(err) + } + + fullStr += chunk + } + assert.Equal(t, c, fullStr) +} diff --git a/flow/agent/react/react.go b/flow/agent/react/react.go index c8f7439..0681c3b 100644 --- a/flow/agent/react/react.go +++ b/flow/agent/react/react.go @@ -38,17 +38,17 @@ import ( func main() { openAIAPIKey := os.Getenv("OPENAI_API_KEY") - openAIBaseURL := os.Getenv("OPENAI_BASE_URL") + // openAIBaseURL := os.Getenv("OPENAI_BASE_URL") openAIModelName := os.Getenv("OPENAI_MODEL_NAME") ctx := context.Background() // prepare chat model chatModel, err := openai.NewChatModel(ctx, &openai.ChatModelConfig{ - BaseURL: openAIBaseURL, - APIKey: openAIAPIKey, - Model: openAIModelName, - ByAzure: true, + // BaseURL: openAIBaseURL, // if using Azure or other models conforming to OpenAI protocol, set this to your model's base URL + APIKey: openAIAPIKey, + Model: openAIModelName, + // ByAzure: true, // if using Azure's OpenAI API, set this to true APIVersion: "2024-06-01", }) if err != nil { diff --git a/go.mod b/go.mod index 4d15bad..590a0a0 100644 --- a/go.mod +++ b/go.mod @@ -14,6 +14,7 @@ require ( github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc github.com/getkin/kin-openapi v0.118.0 github.com/ollama/ollama v0.3.0 + github.com/stretchr/testify v1.9.0 ) require ( @@ -43,6 +44,7 @@ require ( github.com/pelletier/go-toml/v2 v2.2.2 // indirect github.com/perimeterx/marshmallow v1.1.5 // indirect github.com/pkg/errors v0.9.1 // indirect + github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect github.com/rogpeppe/go-internal v1.12.0 // indirect github.com/sashabaranov/go-openai v1.32.5 // indirect github.com/sirupsen/logrus v1.9.3 // indirect