diff --git a/backend/apps/chat/task/llm.py b/backend/apps/chat/task/llm.py index f501a63c..2937a4f9 100644 --- a/backend/apps/chat/task/llm.py +++ b/backend/apps/chat/task/llm.py @@ -430,9 +430,9 @@ class LLMService: _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 diff --git a/frontend/src/stores/assistant.ts b/frontend/src/stores/assistant.ts index f40f8a72..0c7c08fd 100644 --- a/frontend/src/stores/assistant.ts +++ b/frontend/src/stores/assistant.ts @@ -5,30 +5,43 @@ import { useCache } from '@/utils/useCache' const { wsCache } = useCache() const flagKey = 'sqlbit-assistant-flag' +type Resolver = (value: T | PromiseLike) => void +type Rejecter = (reason?: any) => void +interface PendingRequest { + resolve: Resolver + reject: Rejecter +} interface AssistantState { + id: string token: string assistant: boolean flag: number type: number certificate: string online: boolean + requestPromiseMap: Map } export const AssistantStore = defineStore('assistant', { state: (): AssistantState => { return { + id: '', token: '', assistant: false, flag: 0, type: 0, certificate: '', online: false, + requestPromiseMap: new Map(), } }, 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() { + 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 }, diff --git a/frontend/src/utils/request.ts b/frontend/src/utils/request.ts index 03a7e998..6e22a9da 100644 --- a/frontend/src/utils/request.ts +++ b/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 { + public async fetchStream(url: string, data?: any, controller?: AbortController): Promise { 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) ) diff --git a/frontend/src/views/embedded/index.vue b/frontend/src/views/embedded/index.vue index dec2b7cc..c728f20e 100644 --- a/frontend/src/views/embedded/index.vue +++ b/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,