Merge pull request #1 from wangle201210/feat/loop

Feat/loop

Author: wangle201210

.gitignore

@@ -1 +1,2 @@
-.idea
+.idea
+/cmd/text2sql/go.work.sum

README.md

@@ -38,6 +38,7 @@ func main() {
         DbLink:    "root:password@tcp(127.0.0.1:3306)/database?charset=utf8mb4&parseTime=True&loc=Local",
         Try:       5,      // 失败时的重试次数
         ShouldRun: true,   // 是否执行生成的SQL
+        Times:     3,      // 同时生成3个SQL，选择最合适的一个
     }
 
     // 将中文问题转换为SQL并执行
@@ -58,6 +59,10 @@ func main() {
 - `ShouldRun`: 是否执行生成的SQL语句
     - 设置为`true`时会执行SQL并返回结果
     - 设置为`false`时只返回生成的SQL语句
+- `Times`: 同时生成SQL的次数（可选，默认1）
+    - 取值范围：1-10
+    - 数值越大，生成的SQL候选数越多，选择最合适的SQL的准确率越高
+    - 注意：数值越大消耗的token也越多
 
 ## 命令行工具
 [命令行工具](./cmd/text2sql/README.md)

cmd/text2sql/README.md

@@ -7,11 +7,15 @@ text2sql还提供了命令行工具，可以直接在终端中使用：
 go install github.com/wangle201210/text2sql/cmd/text2sql@latest
 
 # 使用示例
-text2sql -l "root:password@tcp(127.0.0.1:3306)/database?charset=utf8mb4&parseTime=True&loc=Local" -q "王五的openid" -t 5 -r
+text2sql -l "root:password@tcp(127.0.0.1:3306)/database" -q "王五的openid" -t 5 -r -n 3
 ```
 
 参数说明：
 - `-l`: MySQL数据库连接信息（必填）
 - `-q`: 查询语句描述（必填）
 - `-t`: 失败后重试次数（可选，默认5次）
 - `-r`: 是否执行生成的SQL（可选，默认false）
+- `-n`: 同时生成SQL的次数（可选，默认1）
+    - 取值范围：1-10
+    - 数值越大，生成的SQL候选数越多，选择最合适的SQL的准确率越高
+    - 注意：数值越大消耗的token也越多

cmd/text2sql/go.work

@@ -0,0 +1,5 @@
+go 1.23.1
+
+use .
+
+replace github.com/wangle201210/text2sql => ../../

cmd/text2sql/main.go

@@ -16,12 +16,14 @@ func main() {
 		question string
 		try      int
 		run      bool
+		times    int
 	)
 
 	flag.StringVar(&link, "l", "", "mysql链接信息")
 	flag.StringVar(&question, "q", "", "查询语句描述")
 	flag.IntVar(&try, "t", 5, "失败后重试次数")
 	flag.BoolVar(&run, "r", false, "获取到sql后是否执行")
+	flag.IntVar(&times, "n", 1, "同时生成SQL的次数(1-10之间)，数值越大准确率越高，但消耗的token也越多")
 	flag.Parse()
 	if link == "" || question == "" {
 		fmt.Println("请输入正确参数")
@@ -32,6 +34,7 @@ func main() {
 		DbLink:    link,
 		Try:       try,
 		ShouldRun: run,
+		Times:     times,
 	}
 	sql, result, err := ts.Do(question)
 	if err != nil {

eino/eino.go

@@ -2,13 +2,41 @@ package eino
 
 import (
 	"context"
+	"fmt"
 )
 
-func GetSQL(ddl, question string) (sql string) {
+func GetSQL(ddl, question string) (sql string, err error) {
 	ctx := context.Background()
-	messages := createMessagesFromTemplate(ddl, question)
-	cm := createOpenAIChatModel(ctx)
-	result := generate(ctx, cm, messages)
-	sql = result.Content
-	return
+	messages, err := ddl2sqlMessages(ddl, question)
+	if err != nil {
+		return "", err
+	}
+
+	cm, err := createOpenAIChatModel(ctx)
+	if err != nil {
+		return "", err
+	}
+	result, err := generate(ctx, cm, messages)
+	if err != nil {
+		return "", fmt.Errorf("生成SQL失败: %w", err)
+	}
+	return result.Content, nil
+}
+
+func 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)
+	if err != nil {
+		return "", fmt.Errorf("选择SQL失败: %w", err)
+	}
+	return result.Content, nil
 }

eino/generate.go

@@ -2,16 +2,17 @@ package eino
 
 import (
 	"context"
-	"log"
+	"fmt"
 
 	"github.com/cloudwego/eino/components/model"
 	"github.com/cloudwego/eino/schema"
 )
 
-func generate(ctx context.Context, llm model.ChatModel, in []*schema.Message) *schema.Message {
-	result, err := llm.Generate(ctx, in)
+func generate(ctx context.Context, llm model.ChatModel, in []*schema.Message) (message *schema.Message, err error) {
+	message, err = llm.Generate(ctx, in)
 	if err != nil {
-		log.Fatalf("llm generate failed: %v", err)
+		err = fmt.Errorf("llm generate failed: %v", err)
+		return
 	}
-	return result
+	return
 }

eino/message.go

@@ -2,12 +2,18 @@ package eino
 
 import (
 	"context"
-	"log"
+	"fmt"
 
 	"github.com/cloudwego/eino/components/prompt"
 	"github.com/cloudwego/eino/schema"
 )
 
+const (
+	limit = 1000
+	role  = "You are a MySQL expert."
+)
+
+// createTemplate 创建并返回一个配置好的聊天模板
 func createTemplate() prompt.ChatTemplate {
 	// 创建模板，使用 FString 格式
 	return prompt.FromMessages(schema.FString,
@@ -33,18 +39,48 @@ func createTemplate() prompt.ChatTemplate {
 	)
 }
 
-func createMessagesFromTemplate(ddl, question string) []*schema.Message {
+// formatMessages 格式化消息并处理错误
+func formatMessages(template prompt.ChatTemplate, data map[string]any) ([]*schema.Message, error) {
+	messages, err := template.Format(context.Background(), data)
+	if err != nil {
+		return nil, fmt.Errorf("格式化模板失败: %w", err)
+	}
+	return messages, nil
+}
+
+// ddl2sqlMessages 将DDL和问题转换为消息列表
+func ddl2sqlMessages(ddl, question string) ([]*schema.Message, error) {
 	template := createTemplate()
-	// 使用模板生成消息
-	messages, err := template.Format(context.Background(), map[string]any{
-		"role":         "You are a MySQL expert.",
+	data := map[string]any{
+		"role":         role,
 		"question":     question,
 		"ddl":          ddl,
-		"limit":        10,
+		"limit":        limit,
 		"chat_history": []*schema.Message{},
-	})
+	}
+	messages, err := formatMessages(template, data)
+	if err != nil {
+		return nil, err
+	}
+	return messages, nil
+}
+
+// choiceSqlMessages 生成SQL选择消息列表
+func choiceSqlMessages(sqls, ddl, question string) ([]*schema.Message, error) {
+	template := createTemplate()
+	data := map[string]any{
+		"role":     role,
+		"question": "Select the most suitable SQL output from the above SQL statements",
+		"ddl":      ddl,
+		"limit":    limit,
+		"chat_history": []*schema.Message{
+			schema.UserMessage(question),
+			schema.AssistantMessage(sqls, nil),
+		},
+	}
+	messages, err := formatMessages(template, data)
 	if err != nil {
-		log.Fatalf("format template failed: %v\n", err)
+		return nil, err
 	}
-	return messages
+	return messages, nil
 }

eino/model.go

@@ -2,29 +2,42 @@ package eino
 
 import (
 	"context"
+	"fmt"
 	"log"
 	"os"
 
 	"github.com/cloudwego/eino-ext/components/model/openai"
 	"github.com/cloudwego/eino/components/model"
 )
 
-func createOpenAIChatModel(ctx context.Context) model.ChatModel {
-	m := os.Getenv("OPENAI_MODEL_NAME")
-	if m == "" {
-		log.Fatalf("请在环境变量中设置你的 OPENAI_MODEL_NAME")
+// createOpenAIChatModel 创建OpenAI聊天模型实例
+func createOpenAIChatModel(ctx context.Context) (model.ChatModel, error) {
+	// 验证必需的环境变量
+	modelName := os.Getenv("OPENAI_MODEL_NAME")
+	if modelName == "" {
+		return nil, fmt.Errorf("环境变量OPENAI_MODEL_NAME未设置")
 	}
-	apikey := os.Getenv("OPENAI_API_KEY")
-	if apikey == "" {
-		log.Fatalf("请在环境变量中设置你的 OPENAI_API_KEY")
+
+	apiKey := os.Getenv("OPENAI_API_KEY")
+	if apiKey == "" {
+		return nil, fmt.Errorf("环境变量OPENAI_API_KEY未设置")
+	}
+
+	// 获取可选的baseURL
+	baseURL := os.Getenv("OPENAI_BASE_URL")
+	if baseURL == "" {
+		log.Println("未设置OPENAI_BASE_URL，将使用默认API地址")
 	}
+
+	// 创建聊天模型
 	chatModel, err := openai.NewChatModel(ctx, &openai.ChatModelConfig{
-		Model:   os.Getenv("OPENAI_MODEL_NAME"),
-		APIKey:  os.Getenv("OPENAI_API_KEY"),
-		BaseURL: os.Getenv("OPENAI_BASE_URL"),
+		Model:   modelName,
+		APIKey:  apiKey,
+		BaseURL: baseURL,
 	})
 	if err != nil {
-		log.Fatalf("create openai chat model failed, err=%v", err)
+		return nil, fmt.Errorf("创建OpenAI聊天模型失败: %w", err)
 	}
-	return chatModel
+
+	return chatModel, nil
 }