perf: Assistant chat

Author: fit2cloud-chenyw

backend/apps/chat/api/chat.py

@@ -81,6 +81,16 @@ async def start_chat(session: SessionDep, current_user: CurrentUser, create_chat
             status_code=500,
             detail=str(e)
         )
+        
+@router.post("/assistant/start")
+async def start_chat(session: SessionDep, current_user: CurrentUser):
+    try:
+        return create_chat(session, current_user, CreateChat(), False)
+    except Exception as e:
+        raise HTTPException(
+            status_code=500,
+            detail=str(e)
+        )
 
 
 @router.post("/recommend_questions/{chat_record_id}")

backend/apps/chat/curd/chat.py

@@ -203,7 +203,8 @@ def save_question(session: SessionDep, current_user: CurrentUser, question: Chat
     if not question.question or question.question.strip() == '':
         raise Exception("Question cannot be Empty")
 
-    chat = session.query(Chat).filter(Chat.id == question.chat_id).first()
+    #chat = session.query(Chat).filter(Chat.id == question.chat_id).first()
+    chat: Chat = session.get(Chat, question.chat_id)
     if not chat:
         raise Exception(f"Chat with id {question.chat_id} not found")
 

backend/apps/chat/task/llm.py

@@ -49,13 +49,16 @@ def __init__(self, session: SessionDep, current_user: CurrentUser, chat_question
         self.session = session
         self.current_user = current_user
 
-        chat = self.session.query(Chat).filter(Chat.id == chat_question.chat_id).first()
+        #chat = self.session.query(Chat).filter(Chat.id == chat_question.chat_id).first()
+        chat_id = chat_question.chat_id
+        chat: Chat = self.session.get(Chat, chat_id)
         if not chat:
-            raise Exception(f"Chat with id {chat_question.chat_id} not found")
+            raise Exception(f"Chat with id {chat_id} not found")
         ds: CoreDatasource | None = None
         if chat.datasource:
             # Get available datasource
-            ds = self.session.query(CoreDatasource).filter(CoreDatasource.id == chat.datasource).first()
+            # ds = self.session.query(CoreDatasource).filter(CoreDatasource.id == chat.datasource).first()
+            ds = self.session.get(CoreDatasource, chat.datasource)
             if not ds:
                 raise Exception("No available datasource configuration found")
 
@@ -65,7 +68,7 @@ def __init__(self, session: SessionDep, current_user: CurrentUser, chat_question
             map(lambda x: ChatRecord(**x.model_dump()), filter(lambda r: True if r.first_chat != True else False,
                                                                list_base_records(session=self.session,
                                                                                  current_user=current_user,
-                                                                                 chart_id=chat_question.chat_id))))
+                                                                                 chart_id=chat_id))))
         # get schema
         if ds:
             chat_question.db_schema = get_table_schema(session=self.session, ds=ds)

frontend/src/api/chat.ts

@@ -302,6 +302,9 @@ export const chatApi = {
   startChat: (data: any): Promise<ChatInfo> => {
     return request.post('/chat/start', data)
   },
+  startAssistantChat: (): Promise<ChatInfo> => {
+    return request.post('/chat/assistant/start')
+  },
   renameChat: (chat_id: number | undefined, brief: string): Promise<string> => {
     return request.post('/chat/rename', { id: chat_id, brief: brief })
   },

frontend/src/stores/assistant.ts

@@ -1,25 +1,42 @@
 import { defineStore } from 'pinia'
 import { store } from './index'
+import { chatApi, ChatInfo } from '@/api/chat'
 
 interface AssistantState {
   token: string
+  assistant: boolean
 }
 
 export const AssistantStore = defineStore('assistant', {
   state: (): AssistantState => {
     return {
       token: '',
+      assistant: false,
     }
   },
   getters: {
     getToken(): string {
       return this.token
     },
+    getAssistant(): boolean {
+      return this.assistant
+    },
   },
   actions: {
     setToken(token: string) {
       this.token = token
     },
+    setAssistant(assistant: boolean) {
+      this.assistant = assistant
+    },
+    async setChat() {
+      if (!this.assistant) {
+        return null
+      }
+      const res = await chatApi.startAssistantChat()
+      const chat: ChatInfo | undefined = chatApi.toChatInfo(res)
+      return chat
+    },
     clear() {
       this.$reset()
     },

frontend/src/utils/request.ts

@@ -180,7 +180,9 @@ class HttpService {
       }
       if (error?.response?.data) {
         const msgData: any = error.response.data
-        msgData.msg && (errorMessage = msgData.msg)
+        if (msgData?.msg) {
+          errorMessage = msgData.msg
+        }
       }
     } else if (error.request) {
       errorMessage = 'No response from server'
@@ -234,6 +236,11 @@ class HttpService {
     if (token) {
       heads['X-SQLBOT-TOKEN'] = `Bearer ${token}`
     }
+    if (assistantStore.getToken) {
+      heads['X-SQLBOT-ASSISTANT-TOKEN'] = `Assistant ${assistantStore.getToken}`
+      if (heads['X-SQLBOT-TOKEN']) heads.delete('X-SQLBOT-TOKEN')
+    }
+    // eslint-disable-next-line @typescript-eslint/ban-ts-comment
     // @ts-ignore
     const request_key = LicenseGenerator.generate()
     heads['X-SQLBOT-KEY'] = request_key

frontend/src/views/chat/index.vue

@@ -1,6 +1,6 @@
 <template>
   <el-container class="chat-container">
-    <el-aside class="chat-container-left">
+    <el-aside v-if="!isAssistant" class="chat-container-left">
       <el-container class="chat-container-right-container">
         <el-header class="chat-list-header">
           <el-button type="primary" @click="createNewChat">
@@ -114,7 +114,10 @@
           </template>
         </el-scrollbar>
       </el-main>
-      <el-footer v-if="computedMessages.length > 0" class="chat-footer">
+      <el-footer
+        v-if="computedMessages.length > 0 || (isAssistant && currentChatId)"
+        class="chat-footer"
+      >
         <div style="height: 24px">
           <template v-if="currentChat.datasource && currentChat.datasource_name">
             {{ t('ds.title') }}：{{ currentChat.datasource_name }}
@@ -147,7 +150,7 @@
       </el-footer>
     </el-container>
 
-    <ChatCreator ref="chatCreatorRef" @on-chat-created="onChatCreated" />
+    <ChatCreator v-if="!isAssistant" ref="chatCreatorRef" @on-chat-created="onChatCreated" />
   </el-container>
 </template>
 
@@ -165,6 +168,11 @@ import ChatCreator from './ChatCreator.vue'
 import { useI18n } from 'vue-i18n'
 import { endsWith, find, startsWith } from 'lodash-es'
 
+import { useAssistantStore } from '@/stores/assistant'
+const assistantStore = useAssistantStore()
+
+const isAssistant = computed(() => assistantStore.getAssistant)
+
 const { t } = useI18n()
 
 const inputMessage = ref('')
@@ -223,8 +231,15 @@ const goEmpty = () => {
   inputMessage.value = ''
 }
 
-const createNewChat = () => {
+const createNewChat = async () => {
   goEmpty()
+  if (isAssistant.value) {
+    const assistantChat = await assistantStore.setChat()
+    if (assistantChat) {
+      onChatCreated(assistantChat as any)
+    }
+    return
+  }
   chatCreatorRef.value?.showDs()
 }
 
@@ -391,6 +406,11 @@ const sendMessage = async () => {
   loading.value = true
   isTyping.value = true
 
+  /* const assistantChat = await assistantStore.setChat()
+  if (assistantChat) {
+    onChatCreated(assistantChat as any)
+  } */
+
   const currentRecord = new ChatRecord()
   currentRecord.create_time = new Date()
   currentRecord.chat_id = currentChatId.value

frontend/src/views/embedded/assistant.vue

@@ -98,6 +98,7 @@ onBeforeMount(async () => {
   const assistantId = route.params.id
   validator.value = await assistantApi.validate(assistantId)
   assistantStore.setToken(validator.value.token)
+  assistantStore.setAssistant(true)
   loading.value = false
 })
 </script>