✨ feat(精修): 支持对单个文件进行精修

Author: dingdinglz

bind.go

@@ -12,4 +12,5 @@ func BindRoutes(server *fiber.App) {
 	apiGroup.Get("/file", route.FileRoute)
 	apiGroup.Post("/save", route.FileUpdateRoute)
 	apiGroup.Get("/load", route.FileListRoute)
+	apiGroup.Get("/change", route.ChangeFileRoute)
 }

route/file.go

@@ -1,13 +1,15 @@
 package route
 
 import (
+	"encoding/json"
 	"io/fs"
 	"os"
 	"path/filepath"
 	"strings"
 	"ui/global"
 	"ui/tool"
 
+	"github.com/dingdinglz/openai"
 	"github.com/gofiber/fiber/v2"
 )
 
@@ -32,6 +34,7 @@ func FileUpdateRoute(c *fiber.Ctx) error {
 		return c.SendStatus(403)
 	}
 	os.WriteFile(filepath.Join(global.RootPath, "data", id, file), []byte(content), os.ModePerm)
+	updateHistoryMessage(id, file, content)
 	return c.JSON(map[string]interface{}{
 		"message": "ok",
 	})
@@ -67,3 +70,21 @@ func FileListRoute(c *fiber.Ctx) error {
 	})
 	return c.JSON(res)
 }
+
+func updateHistoryMessage(id string, name string, content string) {
+	workFile := filepath.Join(global.RootPath, "data", id, "messages.json")
+	if !tool.FileExist(workFile) {
+		return
+	}
+	historyMessages := []openai.Message{}
+	j, _ := os.ReadFile(workFile)
+	json.Unmarshal(j, &historyMessages)
+	for i := 0; i < len(historyMessages); i++ {
+		if historyMessages[i].Role == "user" && historyMessages[i].Content == name {
+			historyMessages[i+1].Content = content
+			break
+		}
+	}
+	res, _ := json.Marshal(historyMessages)
+	os.WriteFile(workFile, res, os.ModePerm)
+}

route/generate.go

@@ -40,6 +40,7 @@ func SendSSEMessage(w *bufio.Writer, message string) {
 	fmt.Fprintf(w, "data: %s\n\n", message)
 	w.Flush()
 }
+
 func GenerateRoute(ctx *fiber.Ctx) error {
 	ctx.Set("Content-Type", "text/event-stream")
 	ctx.Set("Cache-Control", "no-cache")
@@ -121,3 +122,46 @@ func GenerateRoute(ctx *fiber.Ctx) error {
 	}))
 	return nil
 }
+
+func ChangeFileRoute(ctx *fiber.Ctx) error {
+	ctx.Set("Content-Type", "text/event-stream")
+	ctx.Set("Cache-Control", "no-cache")
+	ctx.Set("Connection", "keep-alive")
+	ctx.Set("Transfer-Encoding", "chunked")
+	id := ctx.Query("id")
+	name := ctx.Query("name")
+	prompt := ctx.Query("prompt")
+	ctx.Status(fiber.StatusOK).Context().SetBodyStreamWriter(fasthttp.StreamWriter(func(w *bufio.Writer) {
+		if id == "" || prompt == "" || name == "" {
+			SendSSEMessage(w, MakeGenerateMessage("error", "不能为空！"))
+			return
+		}
+		workPath := filepath.Join(global.RootPath, "data", id)
+		if !tool.FileExist(filepath.Join(workPath, "messages.json")) {
+			SendSSEMessage(w, MakeGenerateMessage("error", "历史记录缺失！"))
+			return
+		}
+		historyContent, _ := os.ReadFile(filepath.Join(workPath, "messages.json"))
+		historyMessages := []openai.Message{}
+		json.Unmarshal(historyContent, &historyMessages)
+		historyMessages = append(historyMessages, openai.Message{
+			Role:    "user",
+			Content: "下面你需要重新生成" + name + "，要求是" + prompt + "\n注意，你需要和之前的生成规则保持一致，即仅返回html的内容即可",
+		})
+		fileContent := ""
+		e := global.OpenaiClient.ChatStream(config.ConfigVar.Model.Model, historyMessages, func(s string) {
+			fileContent += s
+			SendSSEMessage(w, MakeGenerateMessage("update", fileContent))
+		})
+		if e != nil {
+			SendSSEMessage(w, MakeGenerateMessage("error", e.Error()))
+			return
+		}
+		fileContent = tool.StringBetween(fileContent, "```html", "```")
+		os.WriteFile(filepath.Join(workPath, name), []byte(fileContent), os.ModePerm)
+		updateHistoryMessage(id, name, fileContent)
+		SendSSEMessage(w, MakeGenerateMessage("update", "生成完成！如遇任何问题请联系dinglz"))
+		SendSSEMessage(w, MakeGenerateMessage("end", "生成完成"))
+	}))
+	return nil
+}