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

Author: dataeaseShu

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 decorator(func):
 
         @wraps(func)
         async def wrapper(*args, **kwargs):
-            # 准备键生成器参数
-            key_builder_args = {
-                "func": func,
-                "namespace": namespace,
-                "args": args,
-                "kwargs": kwargs
-            }
-            
-            # 生成缓存键
-            cache_key = used_key_builder(**key_builder_args)
-            logger.debug(f"Generated cache key: {cache_key}")
+            cache_key = used_key_builder(
+                func=func,
+                namespace=namespace or "",
+                args=args,
+                kwargs=kwargs
+            )
 
-            # 使用原始缓存装饰器
-            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 @@ async def wrapper(*args, **kwargs):
                 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) {