1. 优化项目结构

2. 可并发请求
3. 不需要反复初始化对象

Author: wangle201210

cmd/text2sql/main.go

@@ -6,7 +6,9 @@ import (
 	"log"
 	"os"
 
+	"github.com/cloudwego/eino-ext/components/model/openai"
 	"github.com/wangle201210/text2sql"
+	"github.com/wangle201210/text2sql/eino"
 )
 
 func main() {
@@ -30,13 +32,23 @@ func main() {
 		os.Exit(1)
 	}
 
-	ts := &text2sql.Text2sql{
+	cfg := &text2sql.Config{
 		DbLink:    link,
 		Try:       try,
 		ShouldRun: run,
 		Times:     times,
 	}
-	sql, result, err := ts.Do(question)
+	newEino, err := eino.NewEino(&openai.ChatModelConfig{
+		APIKey:  os.Getenv("OPENAI_API_KEY"),
+		BaseURL: os.Getenv("OPENAI_BASE_URL"),
+		Model:   os.Getenv("OPENAI_MODEL_NAME"),
+	})
+	if err != nil {
+		log.Fatalf("NewEino err: %+v", err)
+		return
+	}
+	ts := text2sql.NewText2sql(cfg, newEino)
+	sql, result, err := ts.Pretty(question)
 	if err != nil {
 		log.Fatalf("text2sql err: %+v", err)
 	}

eino/eino.go

@@ -5,48 +5,53 @@ import (
 	"encoding/json"
 	"fmt"
 	"strings"
+
+	"github.com/cloudwego/eino-ext/components/model/openai"
+	"github.com/cloudwego/eino/components/model"
 )
 
-func GetSQL(ddl, question string) (sql string, err error) {
-	ctx := context.Background()
-	messages, err := ddl2sqlMessages(ddl, question)
+type Eino struct {
+	cm model.ChatModel
+}
+
+func NewEino(cfg *openai.ChatModelConfig) (res *Eino, err error) {
+	chatModel, err := openai.NewChatModel(context.Background(), cfg)
 	if err != nil {
-		return "", err
+		return nil, err
 	}
+	return &Eino{cm: chatModel}, nil
+}
 
-	cm, err := createOpenAIChatModel(ctx)
+func (x *Eino) GetSQL(ddl, question string) (sql string, err error) {
+	ctx := context.Background()
+	messages, err := ddl2sqlMessages(ddl, question)
 	if err != nil {
 		return "", err
 	}
-	result, err := generate(ctx, cm, messages)
+	result, err := generate(ctx, x.cm, messages)
 	if err != nil {
 		return "", fmt.Errorf("生成SQL失败: %w", err)
 	}
 	sql = result.Content
 	return trimSql(sql), nil
 }
 
-func ChoiceSQL(sqls, ddl, question string) (sql string, err error) {
+func (x *Eino) ChoiceSQL(sqls, ddl, question string) (sql string, err error) {
 	ctx := context.Background()
 	messages, err := choiceSqlMessages(sqls, ddl, question)
 	if err != nil {
 		return "", err
 	}
 
-	cm, err := createOpenAIChatModel(ctx)
-	if err != nil {
-		return "", err
-	}
-	result, err := generate(ctx, cm, messages)
+	result, err := generate(ctx, x.cm, messages)
 	if err != nil {
 		return "", fmt.Errorf("选择SQL失败: %w", err)
 	}
 	sql = result.Content
 	return trimSql(sql), nil
 }
 
-func PrettyRes(sql, question string, runResult []map[string]interface{}) (res string, err error) {
-	ctx := context.Background()
+func (x *Eino) PrettyRes(sql, question string, runResult []map[string]interface{}) (res string, err error) {
 	marshal, err := json.Marshal(runResult)
 	if err != nil {
 		return "", err
@@ -56,11 +61,7 @@ func PrettyRes(sql, question string, runResult []map[string]interface{}) (res st
 		return "", err
 	}
 
-	cm, err := createOpenAIChatModel(ctx)
-	if err != nil {
-		return "", err
-	}
-	result, err := generate(ctx, cm, messages)
+	result, err := generate(context.Background(), x.cm, messages)
 	if err != nil {
 		return "", fmt.Errorf("优化回答失败: %w", err)
 	}