feat: image generate and download

.eslintrc-auto-import.json

@@ -7,13 +7,15 @@
     "DEFAULT_SHORTCUT_KEY": true,
     "EffectScope": true,
     "HOST_URL": true,
+    "ImageCost": true,
     "InjectionKey": true,
     "Message": true,
     "OPEN_AI_MODEL": true,
     "PropType": true,
     "Ref": true,
     "VNode": true,
     "acceptHMRUpdate": true,
+    "calcImageSize": true,
     "checkVersion": true,
     "computed": true,
     "copyCode": true,
@@ -84,6 +86,7 @@
     "request": true,
     "resolveComponent": true,
     "saveImage": true,
+    "saveImageFromBase64": true,
     "saveMarkdown": true,
     "selectSQL": true,
     "setActivePinia": true,

src/api/openAi.ts

@@ -37,7 +37,7 @@ export const getOpenAIResultStreamApi = async (messages: MessageData[]) => {
   const apiKey = getOpenAIKey()
   if (!apiKey) return
 
-  const { updateSessionData, changeLastSessionContent } = useSessionStore()
+  const { updateSessionData } = useSessionStore()
   const { sessionDataList, chatController } = storeToRefs(useSessionStore())
   const {
     proxy: { bypass, url: proxyURL },

src/components/Function/components/AddSession.vue

@@ -5,9 +5,8 @@ const sessionStore = useSessionStore()
 const { switchSession } = sessionStore
 const { isThinking, currentSession } = storeToRefs(sessionStore)
 
-const handleSelect = async (value: MessageType) => {
+const handleSelect = async (value: MessageType | any) => {
   await switchSession()
-
   currentSession.value!.type = value
 }
 </script>

src/components/Function/components/ToggleAnswer.vue

@@ -1,18 +1,16 @@
 <script setup lang="ts">
 const sessionStore = useSessionStore()
-const { updateSessionData, changeLastSessionContent } = sessionStore
-const { isThinking, sessionDataList, chatController } =
+const { updateSessionData } = sessionStore
+const { isThinking, sessionDataList, chatController, currentSession } =
   storeToRefs(sessionStore)
 
 const handleClick = () => {
   if (isThinking.value) {
     chatController.value?.abort()
 
-    changeLastSessionContent('终止了回答')
-
     updateSessionData(sessionDataList.value.at(-1)!)
   } else {
-    getAiMessage()
+    currentSession.value?.type === 'text' ? getAiMessage() : getAiIamge()
   }
 }
 </script>

src/components/Input/index.vue

@@ -81,7 +81,7 @@ onMounted(() => {
           <a-option>1024x1024</a-option>
         </a-select>
         <a-select placeholder="选择生成的数量" v-model="imageValue.number">
-          <a-option v-for="i in 10" :key="i">{{ i + '' }}</a-option>
+          <a-option v-for="i in 10" :key="i">{{ i }}</a-option>
         </a-select>
       </div>
     </div>

src/components/Session/components/SessionAvatar.vue

@@ -0,0 +1,71 @@
+<script lang="ts" setup>
+import type { SessionData } from '@/types'
+import { estimateTokens } from '@/utils'
+import { ImageCost } from '@/constants'
+
+const props = defineProps<{
+  data: SessionData
+}>()
+const { currentRole } = storeToRefs(useRoleStore())
+const { uuid, isTokenUsage } = storeToRefs(useSettingsStore())
+const { sessionDataList, currentSession } = storeToRefs(useSessionStore())
+
+/**
+ * 计算单条消息消耗的token
+ * @param data 单条消息
+ */
+const calcToken = (data: SessionData) => {
+  if (!currentRole.value) return ''
+
+  if (currentSession.value?.type === 'text') {
+    // 角色描述字符数
+    const roleToken = estimateTokens(currentRole.value!.description)
+    // 消息内容字符数
+    const contentToken = estimateTokens(data.message.content)
+    // 记忆模式下额外消耗的字符数
+    let memoryToken = 0
+    if (data.is_memory) {
+      // 获取sessionDataList中的此条之前的最后5条消息
+      const memoryList = sessionDataList.value
+        .filter((data) => data.id! < data.id!)
+        .slice(-5)
+      memoryToken = estimateTokens(
+        memoryList.map((data) => data.message.content).join('')
+      )
+    }
+
+    if (data.is_ask) {
+      return `${roleToken + contentToken + memoryToken}TK${
+        data.is_memory ? '*' : ''
+      }`
+    }
+
+    return `${contentToken}TK`
+  }
+
+  if (currentSession.value?.type === 'image') {
+    return '0TK'
+    // const { message, message_type } = props.data
+
+    // if (message_type !== 'image') return '0TK'
+
+    // // 计算图片的 token * 数量
+    // const num = message.content.length
+
+    // TODO 通过base64获得图片的大小
+    // const size = calcImageSize(message.content[0].b64_json)
+    // if (!size) return '0TK'
+
+    // return `${num * ImageCost[size]}TK`
+  }
+}
+</script>
+
+<template>
+  <div class="flex flex-col items-center gap-1">
+    <Avatar class="w-12!" :value="data.is_ask ? uuid : currentRole?.name" />
+    <span class="text-gray text-xs" v-if="isTokenUsage">
+      {{ calcToken(data) }}
+    </span>
+  </div>
+</template>

src/components/Session/index.vue

@@ -2,8 +2,7 @@
 import { listen } from '@tauri-apps/api/event'
 import MarkdownIt from 'markdown-it'
 import MarkdownItHighlight from 'markdown-it-highlightjs'
-import type { SessionData } from '@/types'
-import { estimateTokens } from '@/utils'
+
 import dayjs from 'dayjs'
 import utc from 'dayjs/plugin/utc'
 
@@ -15,9 +14,7 @@ const marked = new MarkdownIt({
   .use(MarkdownItHighlight)
   .use(copyCode)
 
-const { uuid, isTokenUsage } = storeToRefs(useSettingsStore())
 const { sessionDataList, currentSession } = storeToRefs(useSessionStore())
-const { currentRole } = storeToRefs(useRoleStore())
 
 const getLocalTime = (time: string) =>
   dayjs.utc(time).local().format('YYYY-MM-DD HH:mm:ss')
@@ -61,36 +58,6 @@ const handleScroll = () => {
   }
 }
 
-/**
- * 计算单条消息消耗的token
- * @param item 单条消息
- */
-const calcToken = (item: SessionData) => {
-  // 角色描述字符数
-  const roleToken = estimateTokens(currentRole.value!.description)
-  // 消息内容字符数
-  const contentToken = estimateTokens(item.message.content)
-  // 记忆模式下额外消耗的字符数
-  let memoryToken = 0
-  if (item.is_memory) {
-    // 获取sessionDataList中的此条之前的最后5条消息
-    const memoryList = sessionDataList.value
-      .filter((data) => data.id! < item.id!)
-      .slice(-5)
-    memoryToken = estimateTokens(
-      memoryList.map((data) => data.message.content).join('')
-    )
-  }
-
-  if (item.is_ask) {
-    return `${roleToken + contentToken + memoryToken}TK${
-      item.is_memory ? '*' : ''
-    }`
-  }
-
-  return `${contentToken}TK`
-}
-
 onMounted(() => {
   listen('scroll-to-bottom', () => {
     isAutoScroll.value = true
@@ -123,15 +90,7 @@ watch([currentSession, sessionDataList], () => {
         :class="item.is_ask && 'flex-row-reverse'"
         :key="item.id"
       >
-        <div class="flex flex-col items-center gap-1">
-          <Avatar
-            class="w-12!"
-            :value="item.is_ask ? uuid : currentRole?.name"
-          />
-          <span class="text-gray text-xs" v-if="isTokenUsage">
-            {{ calcToken(item) }}
-          </span>
-        </div>
+        <SessionAvatar :data="item" />
         <div
           class="relative flex w-[calc(100%-8rem)] flex-col gap-2"
           :class="item.is_ask && 'items-end'"
@@ -211,8 +170,22 @@ watch([currentSession, sessionDataList], () => {
               "
             ></div>
 
-            <!--  -->
-            <!-- <img  :src="" alt="" v-else /> -->
+            <a-image-preview-group infinite v-else>
+              <a-space>
+                <a-image
+                  v-for="(img, index) in item.message.content"
+                  :key="index"
+                  :src="`data:image/png;base64,${img.b64_json}`"
+                >
+                  <template #extra>
+                    <icon-download
+                      class="text-6"
+                      @click="saveImageFromBase64(img.b64_json)"
+                    />
+                  </template>
+                </a-image>
+              </a-space>
+            </a-image-preview-group>
           </div>
         </div>
       </div>

src/constants/shared.ts

@@ -11,3 +11,9 @@ export enum HOST_URL {
   OPENAI = 'https://api.openai.com',
   GITHUB = 'https://api.github.com'
 }
+
+export enum ImageCost {
+  '256x256' = 8000,
+  '512x512' = 9000,
+  '1024x1024' = 10000
+}

src/stores/role.ts

@@ -15,7 +15,7 @@ export const useRoleStore = defineStore(
     const textAreaValue = ref('')
     // 图片的参数
     const imageValue = reactive({
-      number: '1',
+      number: 1,
       size: '256x256'
     })
     // 是否有角色正在编辑

src/stores/session.ts

@@ -2,8 +2,7 @@ import type {
   SessionPayload,
   SessionData,
   MessageData,
-  MessageType,
-  ImageData
+  MessageType
 } from '@/types'
 
 export const useSessionStore = defineStore(
@@ -19,8 +18,6 @@ export const useSessionStore = defineStore(
     const isThinking = ref(false)
     // 停止请求
     const chatController = ref<AbortController>()
-    // 类型
-    const sessionType = ref<MessageType>('text')
 
     // 获取会话列表
     const getSessionList = async () => {
@@ -48,6 +45,8 @@ export const useSessionStore = defineStore(
 
     // 创建新会话
     const createNewSession = async () => {
+      if (!currentSession.value) return
+
       const defaultRole = (
         await selectSQL('role', [{ key: 'is_default', value: true }])
       )[0]
@@ -56,7 +55,7 @@ export const useSessionStore = defineStore(
         id: crypto.randomUUID(),
         title: '',
         role_id: defaultRole.id,
-        type: 'text'
+        type: currentSession.value.type
       }
     }
 
@@ -76,7 +75,9 @@ export const useSessionStore = defineStore(
 
       if (!isExist && currentRole?.id) {
         const title =
-          messageType === 'text' ? data.content : data.content.prompt
+          currentSession.value.type === 'text'
+            ? data.content
+            : data.content.prompt
 
         currentSession.value.title = title
 
@@ -96,15 +97,15 @@ export const useSessionStore = defineStore(
           is_ask: isAsk,
           is_memory: isMemory,
           message: data,
-          message_type: messageType
+          message_type: 'text'
         })
       } else if (messageType === 'image') {
         await insertSQL('session_data', {
           session_id: currentSession.value.id,
           is_ask: isAsk,
           is_memory: false,
           message: data,
-          message_type: messageType
+          message_type: 'image'
         })
       }
 
@@ -119,8 +120,12 @@ export const useSessionStore = defineStore(
 
     // 更新会话信息
     const updateSession = async (payload: SessionPayload) => {
+      const newPayload = { ...payload }
+      delete newPayload.name
+      delete newPayload.isEdit
+
       await updateSQL('session', {
-        ...payload,
+        ...newPayload,
         update_time: Date.now().toString()
       })
     }
@@ -155,11 +160,14 @@ export const useSessionStore = defineStore(
       if (!session) await createNewSession()
       else currentSession.value = session
 
-      const { changeCurrentRole } = useRoleStore()
+      sessionDataList.value.length = 0
+      setTimeout(() => {
+        const { changeCurrentRole } = useRoleStore()
 
-      changeCurrentRole()
-      getSessionData()
-      getSessionList()
+        changeCurrentRole()
+        getSessionData()
+        getSessionList()
+      }, 50)
     }
 
     // 修改最后一个对话内容

src/types/sql.d.ts

@@ -2,10 +2,12 @@ type Role = 'system' | 'user' | 'assistant'
 
 export type TableName = 'session' | 'session_data' | 'role' | 'credit'
 
+export type ImageSize = '256x256' | '512x512' | '1024x1024'
+
 export interface ImageData {
   n: number
   prompt: string
-  size: '256x256' | '512x512' | '1024x1024'
+  size: ImageSize
   response_format: string
 }