From c0fac567cc61deb521fc2a4cf297273d2684f7c7 Mon Sep 17 00:00:00 2001
From: Onelevenvy <qingxin1114@163.com>
Date: Sat, 28 Sep 2024 08:45:47 +0800
Subject: [PATCH] feat:change the img output format

---
 .../app/core/tools/siliconflow/siliconflow.py |  2 +-
 backend/app/core/tools/spark/spark.py         | 50 +++++++++++--------
 web/src/app/(applayout)/dashboard/page.tsx    |  8 ++-
 web/src/components/Playground/MessageBox.tsx  | 24 ++++-----
 4 files changed, 45 insertions(+), 39 deletions(-)

diff --git a/backend/app/core/tools/siliconflow/siliconflow.py b/backend/app/core/tools/siliconflow/siliconflow.py
index 2d9f0547..04ce08bf 100644
--- a/backend/app/core/tools/siliconflow/siliconflow.py
+++ b/backend/app/core/tools/siliconflow/siliconflow.py
@@ -36,7 +36,7 @@ def text2img(
 
         response = requests.post(url, json=payload, headers=headers)
 
-        return response.json()
+        return response.json()["images"][0]["url"]
 
     except Exception as e:
         return json.dumps(f"There is a error occured . {e}")
diff --git a/backend/app/core/tools/spark/spark.py b/backend/app/core/tools/spark/spark.py
index db83e33a..d8572bb4 100644
--- a/backend/app/core/tools/spark/spark.py
+++ b/backend/app/core/tools/spark/spark.py
@@ -103,28 +103,34 @@ def spark_response(text, appid, apisecret, apikey):
 
 
 def img_generation(prompt):
-    response = spark_response(
-        text=prompt,
-        appid=os.environ.get("SPARK_APPID"),
-        apisecret=os.environ.get("SPARK_APISECRET"),
-        apikey=os.environ.get("SPARK_APIKEY"),
-    )
-    try:
-        data = json.loads(response)
-        code = data["header"]["code"]
-        if code != 0:
-            return f"error: {code}, {data}"
-        else:
-            text = data["payload"]["choices"]["text"]
-            image_content = text[0]
-            image_base = image_content["content"]
-            bs64data = "data:image/jpeg;base64," + image_base
-
-            return bs64data
-        # return aaa
-
-    except Exception as e:
-        return json.dumps(f"There is a error occured . {e}")
+    appid = os.environ.get("SPARK_APPID", "")
+    apisecret = os.environ.get("SPARK_APISECRET", "")
+    apikey = os.environ.get("SPARK_APIKEY", "")
+    if not appid or not apisecret or not apikey:
+        return "api key is not set or not correct"
+    else:
+        response = spark_response(
+            text=prompt,
+            appid=appid,
+            apisecret=apisecret,
+            apikey=apikey,
+        )
+        try:
+            data = json.loads(response)
+            code = data["header"]["code"]
+            if code != 0:
+                return f"error: {code}, {data}"
+            else:
+                text = data["payload"]["choices"]["text"]
+                image_content = text[0]
+                image_base = image_content["content"]
+                bs64data = "data:image/jpeg;base64," + image_base
+
+                return bs64data
+            # return aaa
+
+        except Exception as e:
+            return json.dumps(f"There is a error occured . {e}")
 
 
 spark = StructuredTool.from_function(
diff --git a/web/src/app/(applayout)/dashboard/page.tsx b/web/src/app/(applayout)/dashboard/page.tsx
index 70579eda..00d5010f 100644
--- a/web/src/app/(applayout)/dashboard/page.tsx
+++ b/web/src/app/(applayout)/dashboard/page.tsx
@@ -1,6 +1,6 @@
 "use client";
 import useAuth from "@/hooks/useAuth";
-import { Box, Container, Text } from "@chakra-ui/react";
+import { Box, Container, Image, Text } from "@chakra-ui/react";
 import React from "react";
 
 function Dashboard() {
@@ -14,6 +14,12 @@ function Dashboard() {
             Hi, {currentUser?.full_name || currentUser?.email} 👋🏼
           </Text>
           <Text>Welcome back, nice to see you again!</Text>
+          <Image
+            src={'https://sf-maas-uat-prod.oss-cn-shanghai.aliyuncs.com/outputs/f967b5c9-c3c8-4f4b-9f1f-708e2dfda39c_0.png'}
+            alt="aaa"
+            width={"auto"}
+            height={"auto"}
+          />
         </Box>
       </Container>
     </>
diff --git a/web/src/components/Playground/MessageBox.tsx b/web/src/components/Playground/MessageBox.tsx
index 87956e86..4e95ffc9 100644
--- a/web/src/components/Playground/MessageBox.tsx
+++ b/web/src/components/Playground/MessageBox.tsx
@@ -89,21 +89,15 @@ const MessageBox = ({ message, onResume, isPlayground }: MessageBoxProps) => {
   };
 
   function isImag(content: any): boolean {
-    if (content) {
-      // 检查是否为字符串并以 "data:image/" 开头
-      if (typeof content === "string" && content.startsWith("data:image/")) {
-        return true;
-      }
-
-      // 检查是否为对象并且包含 'images' 键
-      if (typeof content === "object" && content.images) {
-        const images = content.images;
-        // 确保 'images' 是一个非空数组，并检查第一个元素是否包含 'url' 键
-        if (Array.isArray(images) && images.length > 0 && "url" in images[0]) {
-          return true;
-        }
-      }
+    if (typeof content === "string") {
+      // 检查是否为 data URL 或有效的图像 URL
+      return (
+        content.startsWith("data:image/") ||
+        content.startsWith("http://") ||
+        content.startsWith("https://")
+      );
     }
+
     return false;
   }
 
@@ -277,7 +271,7 @@ const MessageBox = ({ message, onResume, isPlayground }: MessageBoxProps) => {
                           <AccordionPanel pb={4}>
                             {isImag(content) ? (
                               <Image
-                                src="content"
+                                src={content!}
                                 alt="img"
                                 width={"100%"}
                                 height={"100%"}