feat 存档

backend/app.go

@@ -4,9 +4,14 @@ import (
 	"context"
 	"encoding/json"
 	"fiction/backend/ai"
+	"fiction/backend/doubao"
+	"fiction/db/model"
+	"fiction/db/sqlite"
 	"fmt"
 	"os"
 	"path/filepath"
+	"strings"
+	"time"
 )
 
 // App struct
@@ -21,6 +26,14 @@ func NewApp() *App {
 
 func (a *App) Startup(ctx context.Context) {
 	a.ctx = ctx
+	if err := sqlite.Init("./db/db.db"); err != nil {
+		panic(err)
+	}
+}
+
+func (a *App) Shutdown(ctx context.Context) {
+	a.ctx = nil
+	sqlite.Stop()
 }
 
 // GetNovelList returns the list of txt files in the source directory
@@ -57,33 +70,113 @@ func (a *App) GetNovelContent(filename string) string {
 	return string(content)
 }
 
-func (a *App) WriteNovelContent(filename string) string {
-	filePath := filepath.Join("./config", "config.json")
+// GetFangWenContent 获取仿文的内容
+func (a *App) GetFangWenContent(filename string) string {
+	first := &model.List{}
+	sqlite.DB.Model(&model.List{}).
+		Where("name = ?", filename).
+		First(&first)
+	if first.ID > 0 {
+		if len(first.Imitation) > 0 {
+			return first.Imitation
+		}
+		if first.Status == 1 {
+			return "正在仿写中..."
+		}
+	}
+	return ""
+}
 
-	content, err := os.ReadFile(filePath)
+// GetAIContent 获取ai对话记录
+func (a *App) GetAIContent(filename string) []ai.Message {
+	first := &model.List{}
+	sqlite.DB.Model(&model.List{}).
+		Where("name = ?", filename).
+		First(&first)
+	if first.ID > 0 {
+		if len(first.AI) > 0 {
+			resp := make([]ai.Message, 0)
+			_ = json.Unmarshal([]byte(first.AI), &resp)
+			return resp
+		}
+	}
+	return make([]ai.Message, 0)
+}
+
+// GetToken 获取ai token消耗
+func (a *App) GetToken(filename string) string {
+	first := &model.List{}
+	sqlite.DB.Model(&model.List{}).
+		Where("name = ?", filename).
+		First(&first)
+	return first.Token
+}
+
+func (a *App) WriteNovelContent(filename string) string {
+	content, err := os.ReadFile(filepath.Join("./config", "config.json"))
 	if err != nil {
 		fmt.Printf("读取提示词文件错误: %v\n", err)
 		return "Error reading file"
 	}
-
 	var str []string
 	if err = json.Unmarshal(content, &str); err != nil {
 		fmt.Printf("读取提示词错误: %v\n", err)
 		return "Error Unmarshal JSON"
 	}
 
-	resp := ""
-	err = ai.NewAIClient(false, "", 0.9).
+	first := &model.List{}
+	sqlite.DB.Model(&model.List{}).
+		Where("name = ?", filename).
+		First(&first)
+	if first.ID == 0 {
+		first.Name = filename
+		first.Status = 1
+		sqlite.DB.Create(&first)
+	}
+
+	resp := make([]string, 0)
+	msgList, in, out, err := doubao.NewAIClient(false, "", 0).
 		RequestAI(
-			a.GetNovelContent(filename),
+			a.GetModelConfig().Model,
+			a.GetModelConfig().Key,
+			[]string{
+				"本次对话是通过API的方式对话，所以请你不要在回答中掺杂各种图片、修正格式的符号等等，仅以纯文本+换行符的方式回答",
+				a.GetNovelContent(filename),
+			},
 			str,
 			func(content string, i int) error {
-				resp = content
+				processedContent := strings.ReplaceAll(content, "\n\n", "\n")
+				resp = append(resp, processedContent)
 				return nil
 			},
 		)
 	if err != nil {
 		return err.Error()
 	}
-	return resp
+	marshal, _ := json.Marshal(msgList)
+
+	first.Name = filename
+	first.Imitation = resp[len(resp)-1]
+	first.AI = string(marshal)
+	first.Status = 2
+	first.Token = fmt.Sprintf("模型：%v，输入 %d Token，输出 %d Token，预计费用：%.4f元", a.GetModelConfig().Model, in, out, a.GetModelConfig().In/1000000*float64(in)+a.GetModelConfig().Out/1000000*float64(out))
+	sqlite.DB.Save(first)
+
+	saveToFile(filename, resp[len(resp)-1])
+
+	return resp[len(resp)-1]
+}
+
+func saveToFile(name string, content string) {
+	// 获取当前时间戳（精确到秒）
+	timestamp := time.Now().Format("01-02 15-04-05") // 格式化为 YYYYMMDDHHMMSS
+
+	// 构建输出文件名：原文件名 + 时间戳 + .txt
+	baseName := strings.TrimSuffix(name, filepath.Ext(name))
+	outputFileName := fmt.Sprintf("%s_%s.txt", baseName, timestamp)
+
+	outputPath := filepath.Join("./generate", outputFileName)
+	if err := os.WriteFile(outputPath, []byte(content), 0644); err != nil {
+		fmt.Printf("写入文件失败: %v\n", err)
+	}
 }