perf: Assistant dynamic certificate

Author: fit2cloud-chenyw

backend/apps/chat/task/llm.py

@@ -430,9 +430,9 @@ def select_datasource(self):
         _datasource: int | None = None
         _engine_type: str | None = None
         try:
-            data = orjson.loads(json_str)
+            data: dict = orjson.loads(json_str)
 
-            if data['id'] and data['id'] != 0:
+            if data.get('id') and data.get('id') != 0:
                 _datasource = data['id']
                 _chat = self.session.get(Chat, self.record.chat_id)
                 _chat.datasource = _datasource

frontend/src/stores/assistant.ts

@@ -5,30 +5,43 @@ import { useCache } from '@/utils/useCache'
 
 const { wsCache } = useCache()
 const flagKey = 'sqlbit-assistant-flag'
+type Resolver<T = any> = (value: T | PromiseLike<T>) => void
+type Rejecter = (reason?: any) => void
+interface PendingRequest<T = any> {
+  resolve: Resolver<T>
+  reject: Rejecter
+}
 interface AssistantState {
+  id: string
   token: string
   assistant: boolean
   flag: number
   type: number
   certificate: string
   online: boolean
+  requestPromiseMap: Map<string, PendingRequest>
 }
 
 export const AssistantStore = defineStore('assistant', {
   state: (): AssistantState => {
     return {
+      id: '',
       token: '',
       assistant: false,
       flag: 0,
       type: 0,
       certificate: '',
       online: false,
+      requestPromiseMap: new Map<string, PendingRequest>(),
     }
   },
   getters: {
     getCertificate(): string {
       return this.certificate
     },
+    getId(): string {
+      return this.id
+    },
     getToken(): string {
       return this.token
     },
@@ -46,6 +59,43 @@ export const AssistantStore = defineStore('assistant', {
     },
   },
   actions: {
+    refreshCertificate<T>() {
+      const timeout = 5000
+      return new Promise((resolve, reject) => {
+        const timeoutId = setTimeout(() => {
+          this.requestPromiseMap.delete(this.id)
+          reject(new Error(`Request ${this.id} timed out after ${timeout}ms`))
+        }, timeout)
+        this.requestPromiseMap.set(this.id, {
+          resolve: (value: T) => {
+            clearTimeout(timeoutId)
+            this.requestPromiseMap.delete(this.id)
+            resolve(value)
+          },
+          reject: (reason) => {
+            clearTimeout(timeoutId)
+            this.requestPromiseMap.delete(this.id)
+            reject(reason)
+          },
+        })
+        const readyData = {
+          eventName: 'sqlbot_assistant_event',
+          busi: 'ready',
+          ready: true,
+          messageId: this.id,
+        }
+        window.parent.postMessage(readyData, '*')
+      })
+    },
+    resolveCertificate(data?: any) {
+      const peddingRequest = this.requestPromiseMap.get(this.id)
+      if (peddingRequest) {
+        peddingRequest.resolve(data)
+      }
+    },
+    setId(id: string) {
+      this.id = id
+    },
     setCertificate(certificate: string) {
       this.certificate = certificate
     },

frontend/src/utils/request.ts

@@ -245,7 +245,7 @@ class HttpService {
     return this.request({ ...config, method: 'POST', url, data })
   }
 
-  public fetchStream(url: string, data?: any, controller?: AbortController): Promise<any> {
+  public async fetchStream(url: string, data?: any, controller?: AbortController): Promise<any> {
     const token = wsCache.get('user.token')
     const heads: any = {
       'Content-Type': 'application/json',
@@ -257,6 +257,7 @@ class HttpService {
       heads['X-SQLBOT-ASSISTANT-TOKEN'] = `Assistant ${assistantStore.getToken}`
       if (heads['X-SQLBOT-TOKEN']) delete heads['X-SQLBOT-TOKEN']
       if (assistantStore.getType && assistantStore.getCertificate) {
+        await assistantStore.refreshCertificate()
         heads['X-SQLBOT-ASSISTANT-CERTIFICATE'] = btoa(
           encodeURIComponent(assistantStore.getCertificate)
         )

frontend/src/views/embedded/index.vue

@@ -58,6 +58,7 @@ const communicationCb = async (event: any) => {
       const certificate = event.data['certificate']
       assistantStore.setType(1)
       assistantStore.setCertificate(certificate)
+      assistantStore.resolveCertificate(certificate)
     }
     if (event.data?.busi == 'setOnline') {
       setFormatOnline(event.data.online)
@@ -85,6 +86,7 @@ onBeforeMount(async () => {
   setFormatOnline(online)
   const now = Date.now()
   assistantStore.setFlag(now)
+  assistantStore.setId(assistantId?.toString() || '')
   const param = {
     id: assistantId,
     virtual: assistantStore.getFlag,