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

Author: dataeaseShu

backend/alembic/versions/036_modify_assistant.py

@@ -0,0 +1,27 @@
+"""036_modify_assistant
+
+Revision ID: 646e7ca28e0e
+Revises: 29559ee607af
+Create Date: 2025-08-18 16:12:46.041413
+
+"""
+from alembic import op
+import sqlalchemy as sa
+import sqlmodel.sql.sqltypes
+
+
+# revision identifiers, used by Alembic.
+revision = '646e7ca28e0e'
+down_revision = '29559ee607af'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    op.add_column('sys_assistant', sa.Column('app_id', sa.String(255), nullable=True, comment='app_id'))
+    op.add_column('sys_assistant', sa.Column('app_secret', sa.String(255), nullable=True, comment='app_secret'))
+
+
+def downgrade():
+    op.drop_column('sys_assistant', 'app_id')
+    op.drop_column('sys_assistant', 'app_secret')

backend/apps/system/api/assistant.py

@@ -2,21 +2,23 @@
 import json
 import os
 from typing import List, Optional
-from fastapi import APIRouter, FastAPI, Form, HTTPException, Query, Request, Response, UploadFile
+from fastapi import APIRouter, Form, HTTPException, Query, Request, Response, UploadFile
 from fastapi.responses import StreamingResponse
-from sqlmodel import Session, select
+from sqlmodel import select
 from apps.system.crud.assistant import get_assistant_info
+from apps.system.crud.assistant_manage import dynamic_upgrade_cors, save
 from apps.system.models.system_model import AssistantModel
 from apps.system.schemas.auth import CacheName, CacheNamespace
 from apps.system.schemas.system_schema import AssistantBase, AssistantDTO, AssistantUiSchema, AssistantValidator
 from common.core.deps import SessionDep, Trans
 from common.core.security import create_access_token
 from common.core.sqlbot_cache import clear_cache
 from common.utils.time import get_timestamp
-from starlette.middleware.cors import CORSMiddleware
+
 from common.core.config import settings
 from common.utils.utils import get_origin_from_referer
 from sqlbot_xpack.file_utils import SQLBotFileUtils
+
 router = APIRouter(tags=["system/assistant"], prefix="/system/assistant")
 
 @router.get("/info/{id}") 
@@ -104,16 +106,12 @@ async def ui(session: SessionDep, data: str = Form(), files: List[UploadFile] =
 
 @router.get("", response_model=list[AssistantModel])
 async def query(session: SessionDep):
-    list_result = session.exec(select(AssistantModel).order_by(AssistantModel.name, AssistantModel.create_time)).all()
+    list_result = session.exec(select(AssistantModel).where(AssistantModel.type.in_([0, 1])).order_by(AssistantModel.name, AssistantModel.create_time)).all()
     return list_result
 
 @router.post("")
 async def add(request: Request, session: SessionDep, creator: AssistantBase):
-    db_model = AssistantModel.model_validate(creator)
-    db_model.create_time = get_timestamp()
-    session.add(db_model)
-    session.commit()
-    dynamic_upgrade_cors(request=request, session=session)
+    save(request, session, creator)
 
 
 @router.put("")
@@ -147,26 +145,7 @@ async def delete(request: Request, session: SessionDep, id: int):
     session.commit()
     dynamic_upgrade_cors(request=request, session=session)
 
-def dynamic_upgrade_cors(request: Request, session: Session):
-    list_result = session.exec(select(AssistantModel).order_by(AssistantModel.create_time)).all()
-    seen = set()
-    unique_domains = []
-    for item in list_result:
-        if item.domain:
-            for domain in item.domain.split(','):
-                domain = domain.strip()
-                if domain and domain not in seen:
-                    seen.add(domain)
-                    unique_domains.append(domain)
-    app: FastAPI = request.app
-    cors_middleware = None
-    for middleware in app.user_middleware:
-        if middleware.cls == CORSMiddleware:
-            cors_middleware = middleware
-            break
-    if cors_middleware:
-        updated_origins = list(set(settings.all_cors_origins + unique_domains))
-        cors_middleware.kwargs['allow_origins'] = updated_origins
+
 
 
 

backend/apps/system/crud/assistant_manage.py

@@ -0,0 +1,39 @@
+
+
+from fastapi import FastAPI, Request
+from sqlmodel import Session, select
+from starlette.middleware.cors import CORSMiddleware
+from apps.system.schemas.system_schema import AssistantBase
+from common.core.config import settings
+from apps.system.models.system_model import AssistantModel
+from common.utils.time import get_timestamp
+
+
+def dynamic_upgrade_cors(request: Request, session: Session):
+    list_result = session.exec(select(AssistantModel).order_by(AssistantModel.create_time)).all()
+    seen = set()
+    unique_domains = []
+    for item in list_result:
+        if item.domain:
+            for domain in item.domain.split(','):
+                domain = domain.strip()
+                if domain and domain not in seen:
+                    seen.add(domain)
+                    unique_domains.append(domain)
+    app: FastAPI = request.app
+    cors_middleware = None
+    for middleware in app.user_middleware:
+        if middleware.cls == CORSMiddleware:
+            cors_middleware = middleware
+            break
+    if cors_middleware:
+        updated_origins = list(set(settings.all_cors_origins + unique_domains))
+        cors_middleware.kwargs['allow_origins'] = updated_origins
+
+
+async def save(request: Request, session: Session, creator: AssistantBase):
+    db_model = AssistantModel.model_validate(creator)
+    db_model.create_time = get_timestamp()
+    session.add(db_model)
+    session.commit()
+    dynamic_upgrade_cors(request=request, session=session)

backend/apps/system/models/system_model.py

@@ -50,6 +50,8 @@ class AssistantBaseModel(SQLModel):
     description: Optional[str] = Field(sa_type = Text(), nullable=True)
     configuration: Optional[str] = Field(sa_type = Text(), nullable=True)
     create_time: int = Field(default=0, sa_type=BigInteger())
+    app_id: Optional[str] = Field(default=None, max_length=255,  nullable=True)
+    app_secret: Optional[str] = Field(default=None, max_length=255, nullable=True)
 
 class AssistantModel(SnowflakeBase, AssistantBaseModel, table=True):
     __tablename__ = "sys_assistant"

backend/common/core/sqlbot_cache.py

@@ -16,7 +16,7 @@ def custom_key_builder(
     kwargs: Dict[str, Any],
     cacheName: str,
     keyExpression: Optional[str] = None,
-) -> str:
+) -> str | list[str]:
     try:
         base_key = f"{namespace}:{cacheName}:"
 
@@ -31,13 +31,17 @@ def custom_key_builder(
                 if match := re.match(r"args\[(\d+)\]", keyExpression):
                     index = int(match.group(1))
                     value = bound_args.args[index]
+                    if isinstance(value, list):
+                        return [f"{base_key}{v}" for v in value]
                     return f"{base_key}{value}"
 
             # 支持属性路径格式
             parts = keyExpression.split('.')
             value = bound_args.arguments[parts[0]]
             for part in parts[1:]:
                 value = getattr(value, part)
+            if isinstance(value, list):
+                return [f"{base_key}{v}" for v in value]
             return f"{base_key}{value}"
 
         # 默认使用第一个参数作为key
@@ -102,14 +106,16 @@ async def wrapper(*args, **kwargs):
                 cacheName=cacheName,
                 keyExpression=keyExpression,
             )
+            ket_list = cache_key if isinstance(cache_key, list) else [cache_key]
             backend = FastAPICache.get_backend()
-            if await backend.get(cache_key):
-                if settings.CACHE_TYPE.lower() == "redis":
-                    redis = backend.redis
-                    await redis.delete(cache_key)
-                else:
-                    await backend.clear(key=cache_key)
-                SQLBotLogUtil.debug(f"Cache cleared: {cache_key}")
+            for temp_cache_key in ket_list:
+                if await backend.get(temp_cache_key):
+                    if settings.CACHE_TYPE.lower() == "redis":
+                        redis = backend.redis
+                        await redis.delete(temp_cache_key)
+                    else:
+                        await backend.clear(key=temp_cache_key)
+                    SQLBotLogUtil.debug(f"Cache cleared: {temp_cache_key}")
             return await func(*args, **kwargs)
 
         return wrapper

backend/common/utils/whitelist.py

@@ -31,6 +31,7 @@
     "/system/assistant/validator*",
     "/system/assistant/info/*",
     "/system/assistant/picture/*",
+    "/system/embedded*",
     "/datasource/uploadExcel"
 ]
 

frontend/src/router/watch.ts

@@ -8,7 +8,7 @@ const appearanceStore = useAppearanceStoreWithOut()
 const userStore = useUserStore()
 const { wsCache } = useCache()
 const whiteList = ['/login']
-const assistantWhiteList = ['/assistant']
+const assistantWhiteList = ['/assistant', '/embeddedPage']
 export const watchRouter = (router: any) => {
   router.beforeEach(async (to: any, from: any, next: any) => {
     await loadXpackStatic()

frontend/src/views/embedded/page.vue

@@ -1,13 +1,100 @@
 <template>
   <div class="sqlbot--embedded-page">
-    <chat-component ref="chatRef" />
+    <chat-component v-if="!loading" ref="chatRef" />
   </div>
 </template>
 <script setup lang="ts">
 import ChatComponent from '@/views/chat/index.vue'
-import { ref } from 'vue'
+import { onBeforeMount, onBeforeUnmount, ref } from 'vue'
+import { useRoute } from 'vue-router'
+import { assistantApi } from '@/api/assistant'
+import { useAssistantStore } from '@/stores/assistant'
 
 const chatRef = ref()
+const assistantStore = useAssistantStore()
+const route = useRoute()
+const assistantName = ref('')
+
+const validator = ref({
+  id: '',
+  valid: false,
+  id_match: false,
+  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']
+      assistantStore.setType(1)
+      assistantStore.setCertificate(certificate)
+      assistantStore.resolveCertificate(certificate)
+    }
+    if (event.data?.busi == 'setOnline') {
+      setFormatOnline(event.data.online)
+    }
+  }
+}
+const setFormatOnline = (text?: any) => {
+  if (text === null || typeof text === 'undefined') {
+    assistantStore.setOnline(false)
+    return
+  }
+  if (typeof text === 'boolean') {
+    assistantStore.setOnline(text)
+    return
+  }
+  if (typeof text === 'string') {
+    assistantStore.setOnline(text.toLowerCase() === 'true')
+    return
+  }
+  assistantStore.setOnline(false)
+}
+
+onBeforeMount(async () => {
+  debugger
+  const assistantId = route.query.id
+  if (!assistantId) {
+    ElMessage.error('Miss embedded id, please check embedded url')
+    return
+  }
+  const online = route.query.online
+  setFormatOnline(online)
+
+  let name = route.query.name
+  if (name) {
+    assistantName.value = decodeURIComponent(name.toString())
+  }
+  const now = Date.now()
+  assistantStore.setFlag(now)
+  assistantStore.setId(assistantId?.toString() || '')
+  const param = {
+    id: assistantId,
+    virtual: assistantStore.getFlag,
+    online,
+  }
+  validator.value = await assistantApi.validate(param)
+  assistantStore.setToken(validator.value.token)
+  assistantStore.setAssistant(true)
+  loading.value = false
+
+  window.addEventListener('message', communicationCb)
+  const readyData = {
+    eventName: 'sqlbot_embedded_event',
+    busi: 'ready',
+    ready: true,
+    messageId: assistantId,
+  }
+  window.parent.postMessage(readyData, '*')
+})
+
+onBeforeUnmount(() => {
+  window.removeEventListener('message', communicationCb)
+})
 </script>
 
 <style lang="less" scoped>