eino-llm/orchestration.go at master · Penryn/eino-llm · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
package main

import (
	"context"
	"github.com/joho/godotenv"
	"log"
	"os"

	"github.com/cloudwego/eino/compose"
	"github.com/cloudwego/eino/schema"
)

// 模型基本信息
var (
	ModelType   string // 模型名称
	OwnerAPIKey string // apiKey
	BaseURL     string // 模型api地址

	arkAPIKey    string // ark模型apiKey
	arkModelName string // ark模型名称
)

func buildeinoLLM(ctx context.Context) (r compose.Runnable[string, *schema.Message], err error) {
	const (
		agent          = "agent"
		ChatTemplate1  = "ChatTemplate1"
		ChatModel2     = "ChatModel2"
		SearchTemplate = "SearchTemplate"
		ConveyMap      = "ConveyMap"
		ConveyMap1     = "ConveyMap1"
	)
	err = godotenv.Load()
	if err != nil {
		log.Fatal("加载 .env 文件出错")
	}

	ModelType = os.Getenv("Model_Type")
	OwnerAPIKey = os.Getenv("Owner_API_Key")
	BaseURL = os.Getenv("Base_URL")
	arkAPIKey = os.Getenv("ARK_API_KEY")
	arkModelName = os.Getenv("ARK_MODEL_NAME")

	g := compose.NewGraph[string, *schema.Message]()
	agentKeyOfLambda, err := newLambda(ctx)
	if err != nil {
		return nil, err
	}
	_ = g.AddLambdaNode(agent, agentKeyOfLambda)
	chatTemplate1KeyOfChatTemplate, err := newChatTemplate(ctx)
	if err != nil {
		return nil, err
	}
	_ = g.AddChatTemplateNode(ChatTemplate1, chatTemplate1KeyOfChatTemplate)
	chatModel2KeyOfChatModel, err := newChatModel1(ctx)
	if err != nil {
		return nil, err
	}
	_ = g.AddChatModelNode(ChatModel2, chatModel2KeyOfChatModel)
	searchTemplateKeyOfChatTemplate, err := newChatTemplate1(ctx)
	if err != nil {
		return nil, err
	}
	_ = g.AddChatTemplateNode(SearchTemplate, searchTemplateKeyOfChatTemplate)
	_ = g.AddLambdaNode(ConveyMap, compose.InvokableLambda(newLambda1))
	_ = g.AddLambdaNode(ConveyMap1, compose.InvokableLambda(newLambda2))
	_ = g.AddEdge(compose.START, ConveyMap)
	_ = g.AddEdge(ChatModel2, compose.END)
	_ = g.AddEdge(SearchTemplate, agent)
	_ = g.AddEdge(agent, ConveyMap1)
	_ = g.AddEdge(ConveyMap1, ChatTemplate1)
	_ = g.AddEdge(ChatTemplate1, ChatModel2)
	_ = g.AddEdge(ConveyMap, SearchTemplate)
	r, err = g.Compile(ctx, compose.WithGraphName("einoLLM"), compose.WithNodeTriggerMode(compose.AllPredecessor))
	if err != nil {
		return nil, err
	}
	return r, err
}