Merge branch 'main' of https://github.com/dataease/SQLBot

backend/apps/chat/task/llm.py

@@ -846,6 +846,8 @@ def run_analysis_or_predict_task(llm_service: LLMService, action_type: str, base
     try:
         llm_service.set_record(save_analysis_predict_record(llm_service.session, base_record, action_type))
 
+        yield orjson.dumps({'type': 'id', 'id': llm_service.get_record().id}).decode() + '\n\n'
+
         if action_type == 'analysis':
             # generate analysis
             analysis_res = llm_service.generate_analysis()

backend/apps/system/api/user.py

@@ -1,7 +1,7 @@
 from typing import Optional
 from fastapi import APIRouter, HTTPException, Query
 from sqlmodel import func, or_, select, delete as sqlmodel_delete
-from apps.system.crud.user import get_db_user, single_delete, user_ws_options
+from apps.system.crud.user import clean_user_cache, get_db_user, single_delete, user_ws_options
 from apps.system.models.system_model import UserWsModel
 from apps.system.models.user import UserModel
 from apps.system.schemas.auth import CacheName, CacheNamespace
@@ -95,13 +95,14 @@ async def ws_options(session: SessionDep, current_user: CurrentUser, trans: Tran
     return await user_ws_options(session, current_user.id, trans)
 
 @router.put("/ws/{oid}")
-@clear_cache(namespace=CacheNamespace.AUTH_INFO, cacheName=CacheName.USER_INFO, keyExpression="current_user.id")
+# @clear_cache(namespace=CacheNamespace.AUTH_INFO, cacheName=CacheName.USER_INFO, keyExpression="current_user.id")
 async def ws_change(session: SessionDep, current_user: CurrentUser, oid: int):
     ws_list: list[UserWs] = await user_ws_options(session, current_user.id)
     if not any(x.id == oid for x in ws_list):
         raise HTTPException(f"oid [{oid}] is invalid!")
     user_model: UserModel = get_db_user(session = session, user_id = current_user.id)
     user_model.oid = oid
+    await clean_user_cache(user_model.id)
     session.add(user_model)
     session.commit()
 
@@ -173,13 +174,14 @@ async def batch_del(session: SessionDep, id_list: list[int]):
         await single_delete(session, id)
     
 @router.put("/language")
-@clear_cache(namespace=CacheNamespace.AUTH_INFO, cacheName=CacheName.USER_INFO, keyExpression="current_user.id")
+#@clear_cache(namespace=CacheNamespace.AUTH_INFO, cacheName=CacheName.USER_INFO, keyExpression="current_user.id")
 async def langChange(session: SessionDep, current_user: CurrentUser, language: UserLanguage):
     lang = language.language
     if lang not in ["zh-CN", "en"]:
         return {"message": "Language not supported"}
     db_user: UserModel = get_db_user(session=session, user_id=current_user.id)
     db_user.language = lang
+    await clean_user_cache(db_user.id)
     session.add(db_user)
     session.commit()
     
@@ -194,11 +196,12 @@ async def pwdReset(session: SessionDep, current_user: CurrentUser, id: int):
     session.commit()
 
 @router.put("/pwd")
-@clear_cache(namespace=CacheNamespace.AUTH_INFO, cacheName=CacheName.USER_INFO, keyExpression="current_user.id")
+#@clear_cache(namespace=CacheNamespace.AUTH_INFO, cacheName=CacheName.USER_INFO, keyExpression="current_user.id")
 async def pwdUpdate(session: SessionDep, current_user: CurrentUser, editor: PwdEditor):
     db_user: UserModel = get_db_user(session=session, user_id=current_user.id)
     if not verify_md5pwd(editor.pwd, db_user.password):
         raise HTTPException("pwd error")
     db_user.password = md5pwd(editor.new_pwd)
+    await clean_user_cache(db_user.id)
     session.add(db_user)
     session.commit()

backend/common/core/sqlbot_cache.py

@@ -1,12 +1,27 @@
 from fastapi_cache import FastAPICache
 from fastapi_cache.decorator import cache as original_cache
 from functools import partial, wraps
-from typing import Optional, Set, Any, Dict, Tuple
+from typing import Optional, Set, Any, Dict, Tuple, Callable
 from inspect import Parameter, signature
 import logging
+from contextlib import asynccontextmanager
+import asyncio
 
 logger = logging.getLogger(__name__)
 
+# 锁管理
+_cache_locks = {}
+
+@asynccontextmanager
+async def _get_cache_lock(key: str):
+    lock = _cache_locks.setdefault(key, asyncio.Lock())
+    async with lock:
+        try:
+            yield
+        finally:
+            if key in _cache_locks and not lock.locked():
+                del _cache_locks[key]
+
 def should_skip_param(param: Parameter) -> bool:
     """判断参数是否应该被忽略(依赖注入参数)"""
     return (
@@ -26,9 +41,7 @@ def custom_key_builder(
     cacheName: Optional[str] = None,
     keyExpression: Optional[str] = None,
 ) -> str:
-    """
-    完全兼容FastAPICache的键生成器
-    """
+    """完全兼容FastAPICache的键生成器"""
     if cacheName: 
         base_key = f"{namespace}:{cacheName}:"
         
@@ -38,7 +51,6 @@ def custom_key_builder(
                 bound_args = sig.bind_partial(*args, **kwargs)
                 bound_args.apply_defaults()
                 
-                
                 if keyExpression.startswith("args["):
                     import re
                     match = re.match(r"args\[(\d+)\]", keyExpression)
@@ -47,7 +59,6 @@ def custom_key_builder(
                         value = bound_args.args[index]
                         base_key += f"{value}:"
                 else:
-                    
                     parts = keyExpression.split('.')
                     value = bound_args.arguments[parts[0]]
                     for part in parts[1:]:
@@ -58,37 +69,27 @@ def custom_key_builder(
                 logger.warning(f"Failed to evaluate keyExpression '{keyExpression}': {str(e)}")
         
         return base_key
-    # 获取函数签名
-    sig = signature(func)
     
-    # 自动识别要跳过的参数
+    sig = signature(func)
     auto_skip_args = {
         name for name, param in sig.parameters.items()
         if should_skip_param(param)
     }
-    
-    # 合并用户指定的额外跳过参数
     skip_args = auto_skip_args.union(additional_skip_args or set())
-    
-    # 过滤kwargs
     filtered_kwargs = {
         k: v for k, v in kwargs.items() if k not in skip_args
     }
     
-    # 过滤args - 将位置参数映射到它们的参数名
     bound_args = sig.bind_partial(*args, **kwargs)
     bound_args.apply_defaults()
     
     filtered_args = []
     for i, (name, value) in enumerate(bound_args.arguments.items()):
-        # 只处理位置参数 (在args中的参数)
         if i < len(args) and name not in skip_args:
             filtered_args.append(value)
     filtered_args = tuple(filtered_args)
     
-    # 获取默认键生成器
     default_key_builder = FastAPICache.get_key_builder()
-    # 调用默认键生成器（严格按照其要求的参数格式）
     return default_key_builder(
         func=func,
         namespace=namespace,
@@ -105,9 +106,7 @@ def cache(
     cacheName: Optional[str] = None,
     keyExpression: Optional[str] = None,
 ):
-    """
-    完全兼容的缓存装饰器
-    """
+    """完全兼容的缓存装饰器"""
     def decorator(func):
         if key_builder is None:
             used_key_builder = partial(
@@ -121,24 +120,23 @@ def cache(
             
         @wraps(func)
         async def wrapper(*args, **kwargs):
-            # 准备键生成器参数
-            key_builder_args = {
-                "func": func,
-                "namespace": namespace,
-                "args": args,
-                "kwargs": kwargs
-            }
+            cache_key = used_key_builder(
+                func=func,
+                namespace=namespace or "",
+                args=args,
+                kwargs=kwargs
+            )
             
-            # 生成缓存键
-            cache_key = used_key_builder(**key_builder_args)
-            logger.debug(f"Generated cache key: {cache_key}")
-            
-            # 使用原始缓存装饰器
-            return await original_cache(
-                expire=expire,
-                namespace=namespace,
-                key_builder=lambda *_, **__: cache_key  # 直接使用预生成的key
-            )(func)(*args, **kwargs)
+            async with _get_cache_lock(cache_key):
+                backend = FastAPICache.get_backend()
+                cached_value = await backend.get(cache_key)
+                if cached_value is not None:
+                    return cached_value
+                
+                result = await func(*args, **kwargs)
+                await backend.set(cache_key, result, expire)
+                return result
+                
         return wrapper
     return decorator
 
@@ -147,17 +145,10 @@ def clear_cache(
     cacheName: Optional[str] = None,
     keyExpression: Optional[str] = None,
 ):
-    """
-    清除缓存的装饰器，参数与 @cache 保持一致
-    使用方式：
-    @clear_cache(namespace="user", cacheName="info", keyExpression="user_id")
-    async def update_user(user_id: int):
-        ...
-    """
+    """清除缓存的装饰器"""
     def decorator(func):
         @wraps(func)
         async def wrapper(*args, **kwargs):
-            # 1. 生成缓存键（复用 custom_key_builder 逻辑）
             cache_key = custom_key_builder(
                 func=func,
                 namespace=namespace or "",
@@ -167,13 +158,10 @@ def clear_cache(
                 keyExpression=keyExpression,
             )
             
-            logger.debug(f"Clearing cache for key: {cache_key}")
-            # 2. 清除缓存
-            if await FastAPICache.get_backend().get(cache_key):
+            async with _get_cache_lock(cache_key):
                 await FastAPICache.clear(key=cache_key)
-            
-            # 3. 执行原函数
-            return await func(*args, **kwargs)
+                result = await func(*args, **kwargs)
+                return result
         
         return wrapper
     return decorator

backend/template.yaml

@@ -129,7 +129,7 @@ template:
   predict:
     system: |
       ### 说明：
-      你是一个数据分析师，你的任务是根据给定的数据进行数据预测，我将以json格式给你一组数据，你帮我预测之后的数据（一段可以展示趋势的数据，至少2个周期），用json格式返回。
+      你是一个数据分析师，你的任务是根据给定的数据进行数据预测，我将以json格式给你一组数据，你帮我预测之后的数据（一段可以展示趋势的数据，至少2个周期），用json格式返回，返回的格式需要与传入的数据格式保持一致。
       ```json
       
       无法预测或者不支持预测的数据请直接返回(不需要返回JSON格式)："抱歉，该数据无法进行预测。(有原因则返回无法预测的原因)"

frontend/public/assistant.js

@@ -40,7 +40,7 @@
   const getChatContainerHtml = (data) => {
     return `
 <div id="sqlbot-assistant-chat-container">
-  <iframe id="sqlbot-assistant-chat" allow="microphone" src="${data.domain_url}/#/assistant?id=${data.id}"></iframe>
+  <iframe id="sqlbot-assistant-chat-iframe-${data.id}" allow="microphone" src="${data.domain_url}/#/assistant?id=${data.id}"></iframe>
 <div class="sqlbot-assistant-operate">
   <div class="sqlbot-assistant-closeviewport sqlbot-assistant-viewportnone">
     <svg xmlns="http://www.w3.org/2000/svg" width="20" height="20" viewBox="0 0 20 20" fill="none">
@@ -311,7 +311,7 @@
     #sqlbot-assistant #sqlbot-assistant-chat-container .sqlbot-assistant-viewportnone{
       display:none;
     }
-    #sqlbot-assistant #sqlbot-assistant-chat-container #sqlbot-assistant-chat{
+    #sqlbot-assistant #sqlbot-assistant-chat-container #sqlbot-assistant-chat-iframe-${data.id}{
      height:100%;
      width:100%;
      border: none;
@@ -334,7 +334,79 @@
     const url = new URL(src)
     return url.searchParams.get(key)
   }
+  function parsrCertificate(config) {
+    const certificateList = config.certificate
+    if (!certificateList?.length) {
+      return null
+    }
+    const list = certificateList.map((item) => formatCertificate(item)).filter((item) => !!item)
+    return JSON.stringify(list)
+  }
+  function isEmpty(obj) {
+    return obj == null || typeof obj == 'undefined'
+  }
+  function formatCertificate(item) {
+    const { type, source, target, target_key, target_val } = item
+    let source_val = null
+    if (type.toLocaleLowerCase() == 'localstorage') {
+      source_val = localStorage.getItem(source)
+    }
+    if (type.toLocaleLowerCase() == 'sessionstorage') {
+      source_val = sessionStorage.getItem(source)
+    }
+    if (type.toLocaleLowerCase() == 'cookie') {
+      source_val = getCookie(source)
+    }
+    if (type.toLocaleLowerCase() == 'custom') {
+      source_val = source
+    }
+    if (isEmpty(source_val)) {
+      return null
+    }
+    return {
+      target,
+      key: target_key || source,
+      value: (target_val && eval(target_val)) || source_val,
+    }
+  }
+  function getCookie(key) {
+    if (!key || !document.cookie) {
+      return null
+    }
+    const cookies = document.cookie.split(';')
+    for (let i = 0; i < cookies.length; i++) {
+      const cookie = cookies[i].trim()
 
+      if (cookie.startsWith(key + '=')) {
+        return decodeURIComponent(cookie.substring(key.length + 1))
+      }
+    }
+    return null
+  }
+  function registerMessageEvent(id, data) {
+    const iframe = document.getElementById(`sqlbot-assistant-chat-iframe-${id}`)
+    const url = iframe.src
+    const eventName = 'sqlbot_assistant_event'
+    window.addEventListener('message', (event) => {
+      if (event.data?.eventName === eventName) {
+        if (event.data?.messageId !== id) {
+          return
+        }
+        if (event.data?.busi == 'ready' && event.data?.ready) {
+          const certificate = parsrCertificate(data)
+          console.log(certificate)
+          params = {
+            busi: 'certificate',
+            certificate,
+            eventName,
+            messageId: id,
+          }
+          const contentWindow = iframe.contentWindow
+          contentWindow.postMessage(params, url)
+        }
+      }
+    })
+  }
   function loadScript(src, id) {
     const domain_url = getDomain(src)
     let url = `${domain_url}/api/v1/system/assistant/info/${id}`
@@ -353,6 +425,10 @@
           tempData = Object.assign(tempData, config)
         }
         initsqlbot_assistant(tempData)
+        if (data.type == 1) {
+          registerMessageEvent(id, tempData)
+          // postMessage the certificate to iframe
+        }
       })
   }
   function getDomain(src) {

frontend/src/stores/assistant.ts

@@ -9,6 +9,8 @@ interface AssistantState {
   token: string
   assistant: boolean
   flag: number
+  type: number
+  certificate: string
 }
 
 export const AssistantStore = defineStore('assistant', {
@@ -17,9 +19,14 @@ export const AssistantStore = defineStore('assistant', {
       token: '',
       assistant: false,
       flag: 0,
+      type: 0,
+      certificate: '',
     }
   },
   getters: {
+    getCertificate(): string {
+      return this.certificate
+    },
     getToken(): string {
       return this.token
     },
@@ -29,8 +36,17 @@ export const AssistantStore = defineStore('assistant', {
     getFlag(): number {
       return this.flag
     },
+    getType(): number {
+      return this.type
+    },
   },
   actions: {
+    setCertificate(certificate: string) {
+      this.certificate = certificate
+    },
+    setType(type: number) {
+      this.type = type
+    },
     setToken(token: string) {
       this.token = token
     },

frontend/src/views/chat/RecommendQuestion.vue

@@ -68,7 +68,7 @@ function clickQuestion(question: string): void {
   .question-grid {
     display: grid;
     grid-gap: 12px;
-    grid-template-columns: repeat(2, 50%);
+    grid-template-columns: repeat(2, calc(50% - 6px));
   }
 
   .question {

frontend/src/views/chat/chat-block/UserChat.vue

@@ -45,6 +45,7 @@ function clickPredict() {
 
   .prefix-title {
     color: rgba(28, 186, 144, 1);
+    white-space: nowrap;
   }
 }
 </style>

frontend/src/views/chat/index.vue

@@ -116,7 +116,7 @@
                     @finish="onChartAnswerFinish"
                     @error="onChartAnswerError"
                   >
-                    <ChartBlock :message="message" />
+                    <ChartBlock style="margin-top: 12px" :message="message" />
                     <div
                       v-if="message.record?.error && message.record?.error?.trim().length > 0"
                       class="error-container"
@@ -220,7 +220,7 @@
                     @finish="onPredictAnswerFinish"
                     @error="onPredictAnswerError"
                   >
-                    <ChartBlock :message="message" is-predict />
+                    <ChartBlock style="margin-top: 12px" :message="message" is-predict />
                     <div
                       v-if="message.record?.error && message.record?.error?.trim().length > 0"
                       class="error-container"

frontend/src/views/embedded/assistant.vue

@@ -57,7 +57,7 @@
   </div>
 </template>
 <script setup lang="ts">
-import { onBeforeMount, ref } from 'vue'
+import { onBeforeMount, onBeforeUnmount, ref } from 'vue'
 import ChatComponent from '@/views/chat/index.vue'
 import AssistantGif from '@/assets/img/assistant.gif'
 import history from '@/assets/svg/chart/history.svg'
@@ -101,6 +101,21 @@ const validator = ref({
   token: '',
 })
 const loading = ref(true)
+const eventName = 'sqlbot_assistant_event'
+const communicationCb = async (event: any) => {
+  if (event.data?.eventName === eventName) {
+    if (event.data?.messageId !== route.query.id) {
+      return
+    }
+    if (event.data?.busi == 'certificate') {
+      const certificate = event.data['certificate']
+      console.log(certificate)
+      assistantStore.setType(1)
+      assistantStore.setCertificate(certificate)
+      // store certificate to pinia
+    }
+  }
+}
 onBeforeMount(async () => {
   const assistantId = route.query.id
   const online = route.query.online
@@ -116,6 +131,19 @@ onBeforeMount(async () => {
   assistantStore.setToken(validator.value.token)
   assistantStore.setAssistant(true)
   loading.value = false
+
+  window.addEventListener('message', communicationCb)
+  const readyData = {
+    eventName: 'sqlbot_assistant_event',
+    busi: 'ready',
+    ready: true,
+    messageId: assistantId,
+  }
+  window.parent.postMessage(readyData, '*')
+})
+
+onBeforeUnmount(() => {
+  window.removeEventListener('message', communicationCb)
 })
 </script>