From 6bf12df847664afa5ae70c55586b24cc56a670a6 Mon Sep 17 00:00:00 2001
From: Zhang Minghan <zmh13054618081@dingtalk.com>
Date: Sat, 9 Mar 2024 18:48:09 +0800
Subject: [PATCH] feat: update addition feature storage and frontend features

---
 .dockerignore                   |  1 +
 README.md                       |  3 ++-
 addition/article/api.go         |  4 ++--
 addition/article/generate.go    |  4 ++--
 addition/article/utils.go       |  2 +-
 addition/generation/api.go      |  4 ++--
 addition/generation/build.go    |  2 +-
 addition/generation/generate.go |  2 +-
 app/src/components/Markdown.tsx |  5 ++++-
 app/src/routes/Article.tsx      | 16 +++++++++++++---
 app/src/routes/Generation.tsx   |  2 +-
 docker-compose.stable.yaml      |  1 +
 docker-compose.watch.yaml       |  1 +
 docker-compose.yaml             |  1 +
 14 files changed, 33 insertions(+), 15 deletions(-)

diff --git a/.dockerignore b/.dockerignore
index 537aa4c..97e82ae 100644
--- a/.dockerignore
+++ b/.dockerignore
@@ -15,6 +15,7 @@ screenshot
 config.yaml
 config.dev.yaml
 
+# current in ~/storage
 addition/generation/data/*
 !addition/generation/data/.gitkeep
 
diff --git a/README.md b/README.md
index 64f5e37..58d083c 100644
--- a/README.md
+++ b/README.md
@@ -145,6 +145,7 @@ _🚀 **Next Generation AI One-Stop Solution**_
       -p 8000:8094 \
       -v ~/config:/config \
       -v ~/logs:/logs \
+      -v ~/storage:/storage \
       -e MYSQL_HOST=localhost \
       -e MYSQL_PORT=3306 \
       -e MYSQL_DATABASE=chatnio \
@@ -160,7 +161,7 @@ _🚀 **Next Generation AI One-Stop Solution**_
    > - *-p 8000:8094* 指映射宿主机端口为 8000, 可自行修改冒号前的端口号
    > - SECRET: JWT 密钥, 自行生成随机字符串修改
    > - SERVE_STATIC: 是否启用静态文件服务 (正常情况下不需要更改此项, 详见下方常见问题解答)
-   > - *-v ~/config:/config* 和 *-v ~/logs:/logs* 指挂载配置文件和日志文件的宿主机目录, 可自行修改
+   > - *-v ~/config:/config* 挂载配置文件, *-v ~/logs:/logs* 挂载日志文件的宿主机目录, *-v ~/storage:/storage* 挂载附加功能的生成文件
    > - 需配置 MySQL 和 Redis 服务, 请自行参考上方信息修改环境变量
     
     版本更新 （_开启 Watchtower 后无需手动更新, 执行后按照上述步骤重新运行即可_）：
diff --git a/addition/article/api.go b/addition/article/api.go
index bccac1e..93a7ff4 100644
--- a/addition/article/api.go
+++ b/addition/article/api.go
@@ -25,13 +25,13 @@ type WebsocketArticleResponse struct {
 func ProjectTarDownloadAPI(c *gin.Context) {
 	hash := strings.TrimSpace(c.Query("hash"))
 	c.Writer.Header().Add("Content-Disposition", "attachment; filename=article.tar.gz")
-	c.File(fmt.Sprintf("addition/article/data/out/%s.tar.gz", hash))
+	c.File(fmt.Sprintf("storage/article/%s.tar.gz", hash))
 }
 
 func ProjectZipDownloadAPI(c *gin.Context) {
 	hash := strings.TrimSpace(c.Query("hash"))
 	c.Writer.Header().Add("Content-Disposition", "attachment; filename=article.zip")
-	c.File(fmt.Sprintf("addition/article/data/out/%s.zip", hash))
+	c.File(fmt.Sprintf("storage/article/%s.zip", hash))
 }
 
 func GenerateAPI(c *gin.Context) {
diff --git a/addition/article/generate.go b/addition/article/generate.go
index b35f854..f93a60a 100644
--- a/addition/article/generate.go
+++ b/addition/article/generate.go
@@ -75,8 +75,8 @@ func CreateWorker(c *gin.Context, user *auth.User, model string, prompt string,
 
 	hook(StreamProgressResponse{Current: current, Total: total, Quota: 0})
 
-	path := fmt.Sprintf("addition/article/data/%s", hash)
-	if _, _, err := utils.GenerateCompressTask(hash, "addition/article/data/out", path, path); err != nil {
+	path := fmt.Sprintf("storage/article/data/%s", hash)
+	if _, _, err := utils.GenerateCompressTask(hash, "storage/article", path, path); err != nil {
 		globals.Debug(fmt.Sprintf("[article] error during generate compress task: %s", err.Error()))
 		return ""
 	}
diff --git a/addition/article/utils.go b/addition/article/utils.go
index 8dad4d8..4352d67 100644
--- a/addition/article/utils.go
+++ b/addition/article/utils.go
@@ -30,7 +30,7 @@ func GenerateDocxFile(target, title, content string) error {
 }
 
 func CreateArticleFile(hash, title, content string) string {
-	target := fmt.Sprintf("addition/article/data/%s/%s.docx", hash, title)
+	target := fmt.Sprintf("storage/article/data/%s/%s.docx", hash, title)
 	utils.FileDirSafe(target)
 	if err := GenerateDocxFile(target, title, content); err != nil {
 		globals.Debug(fmt.Sprintf("[article] error during generate article %s: %s", title, err.Error()))
diff --git a/addition/generation/api.go b/addition/generation/api.go
index 35f666b..466768b 100644
--- a/addition/generation/api.go
+++ b/addition/generation/api.go
@@ -18,13 +18,13 @@ type WebsocketGenerationForm struct {
 func ProjectTarDownloadAPI(c *gin.Context) {
 	hash := strings.TrimSpace(c.Query("hash"))
 	c.Writer.Header().Add("Content-Disposition", "attachment; filename=code.tar.gz")
-	c.File(fmt.Sprintf("addition/generation/data/out/%s.tar.gz", hash))
+	c.File(fmt.Sprintf("storage/generation/%s.tar.gz", hash))
 }
 
 func ProjectZipDownloadAPI(c *gin.Context) {
 	hash := strings.TrimSpace(c.Query("hash"))
 	c.Writer.Header().Add("Content-Disposition", "attachment; filename=code.zip")
-	c.File(fmt.Sprintf("addition/generation/data/out/%s.zip", hash))
+	c.File(fmt.Sprintf("storage/generation/%s.zip", hash))
 }
 
 func GenerateAPI(c *gin.Context) {
diff --git a/addition/generation/build.go b/addition/generation/build.go
index 530b672..364b2ed 100644
--- a/addition/generation/build.go
+++ b/addition/generation/build.go
@@ -7,7 +7,7 @@ import (
 )
 
 func GetFolder(hash string) string {
-	return fmt.Sprintf("addition/generation/data/%s", hash)
+	return fmt.Sprintf("storage/generation/data/%s", hash)
 }
 
 func GetFolderByHash(model string, prompt string) (string, string) {
diff --git a/addition/generation/generate.go b/addition/generation/generate.go
index e8ba285..561cf56 100644
--- a/addition/generation/generate.go
+++ b/addition/generation/generate.go
@@ -15,7 +15,7 @@ func CreateGenerationWithCache(group, model, prompt string, hook func(buffer *ut
 		}
 	}
 
-	if _, _, err := utils.GenerateCompressTask(hash, "addition/generation/data/out", path, path); err != nil {
+	if _, _, err := utils.GenerateCompressTask(hash, "storage/generation", path, path); err != nil {
 		return "", fmt.Errorf("error during generate compress task: %s", err.Error())
 	}
 
diff --git a/app/src/components/Markdown.tsx b/app/src/components/Markdown.tsx
index 2b98e58..b6f1a8b 100644
--- a/app/src/components/Markdown.tsx
+++ b/app/src/components/Markdown.tsx
@@ -246,7 +246,10 @@ function MarkdownContent({
                     <DialogClose asChild>
                       <Button variant={`outline`}>{t("cancel")}</Button>
                     </DialogClose>
-                    <DialogClose onClick={async () => await sendAction(message)} asChild>
+                    <DialogClose
+                      onClick={async () => await sendAction(message)}
+                      asChild
+                    >
                       <Button variant={`default`}>{t("confirm")}</Button>
                     </DialogClose>
                   </DialogFooter>
diff --git a/app/src/routes/Article.tsx b/app/src/routes/Article.tsx
index 003a541..a9cbfe8 100644
--- a/app/src/routes/Article.tsx
+++ b/app/src/routes/Article.tsx
@@ -1,7 +1,7 @@
 import "@/assets/pages/article.less";
 import { Button } from "@/components/ui/button.tsx";
 import router from "@/router.tsx";
-import { Check, ChevronLeft, Files, Globe, Loader2 } from "lucide-react";
+import { Check, ChevronLeft, Cloud, Files, Globe, Loader2 } from "lucide-react";
 import { Textarea } from "@/components/ui/textarea.tsx";
 import { useTranslation } from "react-i18next";
 import { useDispatch, useSelector } from "react-redux";
@@ -31,7 +31,7 @@ type ProgressProps = {
   total: number;
 };
 
-function GenerateProgress({ current, total }: ProgressProps) {
+function GenerateProgress({ current, total, quota }: ProgressProps & { quota: number }) {
   const { t } = useTranslation();
 
   return (
@@ -56,6 +56,12 @@ function GenerateProgress({ current, total }: ProgressProps) {
         )}
       </p>
       <Progress value={(100 * current) / total} />
+      <div
+        className={`article-quota flex flex-row mt-4 border border-input rounded-md py-1 px-3 select-none w-max items-center mx-auto`}
+      >
+        <Cloud className={`h-4 w-4 mr-2`} />
+        <p>{quota.toFixed(2)}</p>
+      </div>
     </div>
   );
 }
@@ -72,6 +78,7 @@ function ArticleContent() {
   const [progress, setProgress] = useState(false);
 
   const [state, setState] = useState<ProgressProps>({ current: 0, total: 0 });
+  const [quota, setQuota] = useState<number>(0);
   const [hash, setHash] = useState("");
 
   function clear() {
@@ -79,6 +86,7 @@ function ArticleContent() {
     setTitle("");
     setHash("");
     setProgress(false);
+    setQuota(0);
     setState({ current: 0, total: 0 });
   }
 
@@ -100,6 +108,8 @@ function ArticleContent() {
 
     connection.onmessage = (e) => {
       const data = JSON.parse(e.data);
+
+      data.data && data.data.quota && setQuota(quota + data.data.quota);
       if (!data.hash) setState(data.data as ProgressProps);
       else {
         toast({
@@ -123,7 +133,7 @@ function ArticleContent() {
 
   return progress ? (
     <>
-      <GenerateProgress {...state} />
+      <GenerateProgress {...state} quota={quota} />
       {hash && (
         <div className={`article-action flex flex-row items-center my-4 gap-4`}>
           <Button
diff --git a/app/src/routes/Generation.tsx b/app/src/routes/Generation.tsx
index 3dad214..31de37c 100644
--- a/app/src/routes/Generation.tsx
+++ b/app/src/routes/Generation.tsx
@@ -101,7 +101,7 @@ function Wrapper({ onSend }: WrapperProps) {
           {quota > 0 && (
             <div className={`quota-box`}>
               <Cloud className={`h-4 w-4 mr-2`} />
-              {quota}
+              {quota.toFixed(2)}
             </div>
           )}
           <pre className={`message-box`}>
diff --git a/docker-compose.stable.yaml b/docker-compose.stable.yaml
index 68c143a..e57fb55 100644
--- a/docker-compose.stable.yaml
+++ b/docker-compose.stable.yaml
@@ -56,6 +56,7 @@ services:
     volumes:
       - ./config:/config
       - ./logs:/logs
+      - ./storage:/storage
     networks:
       - chatnio-network
 
diff --git a/docker-compose.watch.yaml b/docker-compose.watch.yaml
index 13e4d44..9e3e77f 100644
--- a/docker-compose.watch.yaml
+++ b/docker-compose.watch.yaml
@@ -56,6 +56,7 @@ services:
     volumes:
       - ./config:/config
       - ./logs:/logs
+      - ./storage:/storage
     networks:
       - chatnio-network
 
diff --git a/docker-compose.yaml b/docker-compose.yaml
index d83a484..c0b80d1 100644
--- a/docker-compose.yaml
+++ b/docker-compose.yaml
@@ -56,6 +56,7 @@ services:
       volumes:
         - ./config:/config
         - ./logs:/logs
+        - ./storage:/storage
       networks:
         - chatnio-network