feat: chat record list

Author: ulleo

backend/apps/chat/api/chat.py

@@ -24,7 +24,7 @@ async def chats(session: SessionDep, current_user: CurrentUser):
 
 
 @router.get("/get/{chart_id}")
-async def list_chat(session: SessionDep, current_user: CurrentUser, chart_id: int):
+async def get_chat(session: SessionDep, current_user: CurrentUser, chart_id: int):
     try:
         return get_chat_with_records(chart_id=chart_id, session=session, current_user=current_user)
     except Exception as e:

backend/apps/chat/curd/chat.py

@@ -2,6 +2,7 @@
 from typing import List
 
 from sqlalchemy import and_
+from sqlalchemy.orm import load_only
 
 from apps.chat.models.chat_model import Chat, ChatRecord, CreateChat, ChatInfo, RenameChat, ChatQuestion
 from apps.datasource.models.datasource import CoreDatasource
@@ -55,7 +56,11 @@ def get_chat_with_records(session: SessionDep, chart_id: int, current_user: Curr
         chat_info.datasource_exists = True
         chat_info.datasource_name = ds.name
 
-    record_list = session.query(ChatRecord).filter(
+    record_list = session.query(ChatRecord).options(
+        load_only(ChatRecord.id, ChatRecord.chat_id, ChatRecord.create_time, ChatRecord.finish_time,
+                  ChatRecord.question, ChatRecord.sql_answer, ChatRecord.sql, ChatRecord.data,
+                  ChatRecord.chart_answer, ChatRecord.chart, ChatRecord.finish, ChatRecord.error,
+                  ChatRecord.run_time)).filter(
         and_(Chat.create_by == current_user.id, ChatRecord.chat_id == chart_id)).order_by(ChatRecord.create_time).all()
 
     chat_info.records = record_list
@@ -96,6 +101,9 @@ def create_chat(session: SessionDep, current_user: CurrentUser, create_chat_obj:
     chat_info.id = chat.id
     session.commit()
 
+    chat_info.datasource_exists = True
+    chat_info.datasource_name = ds.name
+
     return chat_info
 
 
@@ -214,6 +222,7 @@ def save_error_message(session: SessionDep, record_id: int, message: str) -> Cha
     record = session.query(ChatRecord).filter(ChatRecord.id == record_id).first()
     record.error = message
     record.finish = True
+    record.finish_time = datetime.datetime.now()
 
     result = ChatRecord(**record.model_dump())
 
@@ -242,11 +251,13 @@ def save_sql_exec_data(session: SessionDep, record_id: int, data: str) -> ChatRe
 
     return result
 
+
 def finish_record(session: SessionDep, record_id: int) -> ChatRecord:
     if not record_id:
         raise Exception("Record id cannot be None")
     record = session.query(ChatRecord).filter(ChatRecord.id == record_id).first()
     record.finish = True
+    record.finish_time = datetime.datetime.now()
 
     result = ChatRecord(**record.model_dump())
 
@@ -257,4 +268,3 @@ def finish_record(session: SessionDep, record_id: int) -> ChatRecord:
     session.commit()
 
     return result
-

backend/apps/chat/models/chat_model.py

@@ -2,7 +2,7 @@
 from sqlalchemy import Column, Text, BigInteger, DateTime, Integer, Identity, Boolean
 from datetime import datetime
 from pydantic import BaseModel
-from typing import List, Optional
+from typing import List, Optional, Any
 
 from apps.template.generate_chart.generator import get_chart_template
 from apps.template.generate_sql.generator import get_sql_template

backend/apps/datasource/crud/datasource.py

@@ -225,6 +225,8 @@ def get_table_obj_by_ds(session: SessionDep, ds: CoreDatasource) -> List[TableAn
 def get_table_schema(session: SessionDep, ds: CoreDatasource) -> str:
     schema_str = ""
     table_objs = get_table_obj_by_ds(session=session, ds=ds)
+    if len(table_objs) == 0:
+        return schema_str
     db_name = table_objs[0].schema
     schema_str += f"【DB_ID】 {db_name}\n【Schema】\n"
     for obj in table_objs:

frontend/src/api/chat.ts

@@ -20,6 +20,7 @@ export interface ChatMessage {
     role: 'user' | 'assistant'
     create_time?: Date | string
     content?: string | number
+    record?: ChatRecord
     isTyping?: boolean
     isWelcome?: boolean
 }
@@ -28,24 +29,32 @@ export class ChatRecord {
     id?: number
     chat_id?: number
     create_time?: Date | string
-    create_by?: number
-    datasource?: number
-    engine_type?: string
+    finish_time?: Date | string
     question?: string
-    answer?: string
+    sql_answer?: string
+    sql?: string
+    data?: string
+    chart_answer?: string
+    chart?: string
+    finish?: boolean = false
+    error?: string
     run_time: number = 0
 
     constructor()
-    constructor(id: number, chat_id: number, create_time: Date | string, create_by: number, datasource: number, engine_type: string, question: string, answer: string, run_time: number)
-    constructor(id?: number, chat_id?: number, create_time?: Date | string, create_by?: number, datasource?: number, engine_type?: string, question?: string, answer?: string, run_time?: number) {
+    constructor(id: number, chat_id: number, create_time: Date | string, finish_time: Date | string | undefined, question: string, sql_answer: string | undefined, sql: string | undefined, data: string | undefined, chart_answer: string | undefined, chart: string | undefined, finish: boolean, error: string | undefined, run_time: number)
+    constructor(id?: number, chat_id?: number, create_time?: Date | string, finish_time?: Date | string, question?: string, sql_answer?: string, sql?: string, data?: string, chart_answer?: string, chart?: string, finish?: boolean, error?: string, run_time?: number) {
         this.id = id
         this.chat_id = chat_id
         this.create_time = getDate(create_time)
-        this.create_by = create_by
-        this.datasource = datasource
-        this.engine_type = engine_type
+        this.finish_time = getDate(finish_time)
         this.question = question
-        this.answer = answer
+        this.sql_answer = sql_answer
+        this.sql = sql
+        this.data = data
+        this.chart_answer = chart_answer
+        this.chart = chart
+        this.finish = finish
+        this.error = error
         this.run_time = run_time ?? 0
     }
 }
@@ -111,7 +120,7 @@ function toChatRecord(data?: any): ChatRecord | undefined {
     if (!data) {
         return undefined
     }
-    return new ChatRecord(data.id, data.number, data.create_time, data.create_by, data.datasource, data.engine_type, data.question, data.answer, data.run_time)
+    return new ChatRecord(data.id, data.chat_id, data.create_time, data.finish_time, data.question, data.sql_answer, data.sql, data.data, data.chart_answer, data.chart, data.finish, data.error, data.run_time)
 }
 
 function toChatRecordList(list: any = []): ChatRecord[] {

frontend/src/views/chat/index.vue

@@ -20,13 +20,31 @@
       <el-main class="chat-record-list">
         <el-scrollbar>
           <template v-for="(message, _index) in computedMessages" :key="_index">
-            <ChatRow :current-chat="currentChat" v-model:datasource="currentChat.datasource" :msg="message"/>
+            <ChatRow :current-chat="currentChat" v-model:datasource="currentChat.datasource" :msg="message">
+              <template v-if="message.role === 'assistant'">
+                <div v-if="message.isTyping">Thinking ...</div>
+                <template v-if="message.record">
+                  <div>
+                    <div v-if="message.record.sql_answer">{{ message.record.sql_answer }}</div>
+                    <div v-if="message.record.sql">{{ message.record.sql }}</div>
+                  </div>
+                  <div>
+                    <div v-if="message.record.data">{{ message.record.data }}</div>
+                  </div>
+                  <div>
+                    <div v-if="message.record.chart_answer">{{ message.record.chart_answer }}</div>
+                    <div v-if="message.record.chart">{{ message.record.chart }}</div>
+                  </div>
+                  <div v-if="message.record.error" style="color: red">{{ message.record.error }}</div>
+                </template>
+              </template>
+            </ChatRow>
           </template>
         </el-scrollbar>
       </el-main>
       <el-footer class="chat-footer">
         <div style="height: 24px;">
-          <template v-if="currentChat.datasource">
+          <template v-if="currentChat.datasource && currentChat.datasource_name">
             使用数据源：{{ currentChat.datasource_name }}
           </template>
         </div>
@@ -55,7 +73,7 @@
 </template>
 
 <script setup lang="ts">
-import {ref, computed, nextTick, watch, onMounted} from 'vue'
+import {computed, nextTick, onMounted, ref} from 'vue'
 import {Plus, Position} from '@element-plus/icons-vue'
 import {Chat, chatApi, ChatInfo, type ChatMessage, ChatRecord, questionApi} from '@/api/chat'
 import ChatList from './ChatList.vue'
@@ -82,7 +100,6 @@ const computedMessages = computed<Array<ChatMessage>>(() => {
   if (currentChatId.value === undefined) {
     return messages
   }
-  let appendThinking = false
   for (let i = 0; i < currentChat.value.records.length; i++) {
     const record = currentChat.value.records[i]
     if (record.question !== undefined) {
@@ -92,24 +109,16 @@ const computedMessages = computed<Array<ChatMessage>>(() => {
         content: record.question,
       })
     }
-    if (record.answer !== undefined && record.answer !== '') {
-      messages.push({
-        role: 'assistant',
-        create_time: record.create_time,
-        content: record.answer,
-        isTyping: i === currentChat.value.records.length - 1 && isTyping.value
-      })
-    } else {
-      appendThinking = true
-    }
-  }
-  if (isTyping.value && appendThinking) {
     messages.push({
       role: 'assistant',
-      content: 'Thinking...',
-      isTyping: true
+      create_time: record.create_time,
+      record: record,
+      isTyping: i === currentChat.value.records.length - 1 && isTyping.value
     })
   }
+
+  console.log(messages)
+
   return messages
 })
 
@@ -138,6 +147,7 @@ function onClickHistory(chat: Chat) {
     chatApi.get(chat.id)
         .then((res) => {
           const info = chatApi.toChatInfo(res)
+          console.log(info)
           if (info) {
             currentChat.value = info
           }
@@ -174,12 +184,6 @@ onMounted(() => {
 })
 
 
-const updateMessageContent = (content: string) => {
-  if (currentChat.value) {
-    currentChat.value.records[currentChat.value.records.length - 1].answer = content
-  }
-}
-
 const sendMessage = async () => {
   if (!inputMessage.value.trim()) return
   if (computedMessages.value[0].content === undefined) return
@@ -191,14 +195,17 @@ const sendMessage = async () => {
   currentRecord.create_time = new Date()
   currentRecord.chat_id = currentChatId.value
   currentRecord.question = inputMessage.value
-  currentRecord.answer = ''
+  currentRecord.sql_answer = ''
+  currentRecord.sql = ''
+  currentRecord.chart_answer = ''
+  currentRecord.chart = ''
 
   currentChat.value.records.push(currentRecord)
   inputMessage.value = ''
 
   let error = false
   if (currentChatId.value === undefined) {
-    chatApi.startChat({question: currentRecord.question.trim(), datasource: currentChat.value.datasource})
+    await chatApi.startChat({question: currentRecord.question.trim(), datasource: currentChat.value.datasource})
         .then((res) => {
           const chat = chatApi.toChatInfo(res)
           if (chat !== undefined) {
@@ -228,50 +235,60 @@ const sendMessage = async () => {
 
     while (true) {
       const {done, value} = await reader.read()
+      console.log(done)
       if (done) {
         isTyping.value = false
         break
       }
 
       const chunk = decoder.decode(value)
-      const lines = chunk.split('\n\n')
-
-      for (const line of lines) {
-        if (line.startsWith('data: ')) {
-          const data = JSON.parse(line.replace('data:', '').trim())
-          let realContent = data.content
-
-          switch (data.type) {
-            case 'html':
-              realContent = '\n' + data.content
-              break
-
-            case 'error':
-              realContent = '\nError: ' + data.content
-              break
-
-            default:
-              if (data.content) {
-                const newContent = currentRecord.answer + realContent
-
-                updateMessageContent(newContent)
-                await nextTick()
-              }
-          }
+      console.log(chunk)
+      const data = JSON.parse(chunk)
+
+      await nextTick(() => {
+        switch (data.type) {
+          case 'info':
+            console.log(data.msg)
+            break
+          case 'error':
+            currentRecord.error = data.content
+            isTyping.value = false
+            break
+          case 'sql-result':
+            currentRecord.sql_answer = currentRecord.sql_answer + data.content
+            break
+          case 'sql':
+            currentRecord.sql = data.content
+            break
+          case 'sql-data':
+            currentRecord.data = data.content
+            break
+          case 'chart-result':
+            currentRecord.chart_answer = currentRecord.chart_answer + data.content
+            break
+          case 'chart':
+            currentRecord.chart = data.content
+            break
+          case 'finish':
+            isTyping.value = false
+            break
         }
-      }
+      })
+
     }
   } catch (error) {
-    updateMessageContent(currentRecord.answer + '\nError: Failed to get response')
+    if (!currentRecord.error) {
+      currentRecord.error = ''
+    }
+    if (currentRecord.error.trim().length !== 0) {
+      currentRecord.error = currentRecord.error + '\n'
+    }
+    currentRecord.error = currentRecord.error + 'Error:' + error
     console.error('Error:', error)
     isTyping.value = false
   }
 }
 
-watch(() => currentChat.value?.records[currentChat.value.records.length - 1]?.answer, () => {
-  //scrollToBottom()
-}, {deep: true})
-
 //@ts-ignore
 const formatMessage = (content: string) => {
   if (!content) return ''