perf: The assistant is optimized for manual data source selection

Author: fit2cloud-chenyw

backend/apps/chat/api/chat.py

@@ -172,9 +172,9 @@ async def start_chat(session: SessionDep, current_user: CurrentUser, create_chat
     module=OperationModules.CHAT,
     result_id_expr="id"
 ))
-async def start_chat(session: SessionDep, current_user: CurrentUser):
+async def start_chat(session: SessionDep, current_user: CurrentUser, current_assistant: CurrentAssistant, create_chat_obj: CreateChat = CreateChat(origin=2)):
     try:
-        return create_chat(session, current_user, CreateChat(origin=2), False)
+        return create_chat(session, current_user, create_chat_obj, create_chat_obj and create_chat_obj.datasource, current_assistant)
     except Exception as e:
         raise HTTPException(
             status_code=500,

backend/apps/chat/curd/chat.py

@@ -11,7 +11,9 @@
     TypeEnum, OperationEnum, ChatRecordResult
 from apps.datasource.crud.recommended_problem import get_datasource_recommended_chart
 from apps.datasource.models.datasource import CoreDatasource
-from apps.system.crud.assistant import AssistantOutDsFactory
+from apps.db.constant import DB
+from apps.system.crud.assistant import AssistantOutDs, AssistantOutDsFactory
+from apps.system.schemas.system_schema import AssistantOutDsSchema
 from common.core.deps import CurrentAssistant, SessionDep, CurrentUser, Trans
 from common.utils.utils import extract_nested_json
 
@@ -447,7 +449,7 @@ def list_generate_chart_logs(session: SessionDep, chart_id: int) -> List[ChatLog
 
 
 def create_chat(session: SessionDep, current_user: CurrentUser, create_chat_obj: CreateChat,
-                require_datasource: bool = True) -> ChatInfo:
+                require_datasource: bool = True, current_assistant: CurrentAssistant = None) -> ChatInfo:
     if not create_chat_obj.datasource and require_datasource:
         raise Exception("Datasource cannot be None")
 
@@ -459,10 +461,15 @@ def create_chat(session: SessionDep, current_user: CurrentUser, create_chat_obj:
                 oid=current_user.oid if current_user.oid is not None else 1,
                 brief=create_chat_obj.question.strip()[:20],
                 origin=create_chat_obj.origin if create_chat_obj.origin is not None else 0)
-    ds: CoreDatasource | None = None
+    ds: CoreDatasource | AssistantOutDsSchema | None = None
     if create_chat_obj.datasource:
         chat.datasource = create_chat_obj.datasource
-        ds = session.get(CoreDatasource, create_chat_obj.datasource)
+        if current_assistant and current_assistant.type == 1:
+            out_ds_instance: AssistantOutDs = AssistantOutDsFactory.get_instance(current_assistant)
+            ds = out_ds_instance.get_ds(chat.datasource)
+            ds.type_name = DB.get_db(ds.type)
+        else:
+            ds = session.get(CoreDatasource, create_chat_obj.datasource)
 
         if not ds:
             raise Exception(f"Datasource with id {create_chat_obj.datasource} not found")
@@ -494,7 +501,7 @@ def create_chat(session: SessionDep, current_user: CurrentUser, create_chat_obj:
         record.finish = True
         record.create_time = datetime.datetime.now()
         record.create_by = current_user.id
-        if ds.recommended_config == 2:
+        if isinstance(ds, CoreDatasource) and ds.recommended_config == 2:
             questions = get_datasource_recommended_chart(session, ds.id)
             record.recommended_question = orjson.dumps(questions).decode()
             record.recommended_question_answer = orjson.dumps({

backend/apps/system/api/assistant.py

@@ -8,14 +8,16 @@
 from sqlbot_xpack.file_utils import SQLBotFileUtils
 from sqlmodel import select
 
+from apps.datasource.models.datasource import CoreDatasource
+from apps.db.constant import DB
 from apps.swagger.i18n import PLACEHOLDER_PREFIX
-from apps.system.crud.assistant import get_assistant_info
+from apps.system.crud.assistant import AssistantOutDs, AssistantOutDsFactory, 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.config import settings
-from common.core.deps import SessionDep, Trans, CurrentUser
+from common.core.deps import CurrentAssistant, SessionDep, Trans, CurrentUser
 from common.core.security import create_access_token
 from common.core.sqlbot_cache import clear_cache
 from common.utils.utils import get_origin_from_referer, origin_match_domain
@@ -163,6 +165,56 @@ async def ui(session: SessionDep, data: str = Form(), files: List[UploadFile] =
 async def clear_ui_cache(id: int):
     pass
 
+@router.get("/ds", include_in_schema=False, response_model=list[dict])
+async def ds(session: SessionDep, current_assistant: CurrentAssistant):
+    if current_assistant.type == 0:
+        online = current_assistant.online
+        configuration = current_assistant.configuration
+        config: dict[any] = json.loads(configuration)
+        oid: int = int(config['oid'])
+        stmt = select(CoreDatasource.id, CoreDatasource.name, CoreDatasource.description, CoreDatasource.type, CoreDatasource.type_name, CoreDatasource.num).where(
+            CoreDatasource.oid == oid)
+        if not online:
+            public_list: list[int] = config.get('public_list') or None
+            if public_list:
+                stmt = stmt.where(CoreDatasource.id.in_(public_list))
+            else:
+                return []
+        db_ds_list = session.exec(stmt)
+        return [
+            {
+                "id": ds.id,
+                "name": ds.name,
+                "description": ds.description,
+                "type": ds.type,
+                "type_name": ds.type_name,
+                "num": ds.num,
+            }
+            for ds in db_ds_list]
+    if current_assistant.type == 1:
+        out_ds_instance: AssistantOutDs = AssistantOutDsFactory.get_instance(current_assistant)
+        return [
+            {
+                "id": ds.id,
+                "name": ds.name,
+                "description": ds.description or ds.comment,
+                "type": ds.type,
+                "type_name": get_db_type(ds.type),
+                "num": len(ds.tables) if ds.tables else 0,
+            }
+            for ds in out_ds_instance.ds_list
+            if get_db_type(ds.type)
+        ]
+        
+    return None
+
+def get_db_type(type):
+    try:
+        db = DB.get_db(type)
+        return db.db_name
+    except Exception:
+        return None
+
 
 @router.get("", response_model=list[AssistantModel], summary=f"{PLACEHOLDER_PREFIX}assistant_grid_api", description=f"{PLACEHOLDER_PREFIX}assistant_grid_api")
 async def query(session: SessionDep, current_user: CurrentUser):
@@ -219,3 +271,4 @@ async def delete(request: Request, session: SessionDep, id: int = Path(descripti
     session.delete(db_model)
     session.commit()
     dynamic_upgrade_cors(request=request, session=session)
+

frontend/src/api/chat.ts

@@ -330,8 +330,8 @@ export const chatApi = {
   startChat: (data: any): Promise<ChatInfo> => {
     return request.post('/chat/start', data)
   },
-  startAssistantChat: (): Promise<ChatInfo> => {
-    return request.post('/chat/assistant/start')
+  startAssistantChat: (data?: any): Promise<ChatInfo> => {
+    return request.post('/chat/assistant/start', Object.assign({ origin: 2 }, data))
   },
   renameChat: (chat_id: number | undefined, brief: string): Promise<string> => {
     return request.post('/chat/rename', { id: chat_id, brief: brief })

frontend/src/i18n/en.json

@@ -669,7 +669,8 @@
     "data_analysis_now": "I can query data, generate charts, detect data anomalies, predict data, and more! Enable Smart Data Analysis now!",
     "window_entrance_icon": "Floating window entrance icon",
     "preview_error": "The current page prohibits embedding using Iframe!",
-    "assistant_app": "Assistant App"
+    "assistant_app": "Assistant App",
+    "auto_select_ds": "Automatically select data source"
   },
   "chat": {
     "type": "Chart Type",

frontend/src/i18n/ko-KR.json

@@ -669,7 +669,8 @@
     "data_analysis_now": "저는 데이터 조회, 차트 생성, 데이터 이상 감지, 데이터 예측 등을 할 수 있습니다. 빨리 스마트 데이터 조회를 시작하세요~",
     "window_entrance_icon": "플로팅 윈도우 입구 아이콘",
     "preview_error": "현재 페이지는 Iframe 임베딩을 허용하지 않습니다!",
-    "assistant_app": "어시스턴트 앱"
+    "assistant_app": "어시스턴트 앱",
+    "auto_select_ds": "자동 데이터 소스 선택"
   },
   "chat": {
     "type": "차트 유형",

frontend/src/i18n/zh-CN.json

@@ -669,7 +669,8 @@
     "data_analysis_now": "我可以查询数据、生成图表、检测数据异常、预测数据等赶快开启智能问数吧～",
     "window_entrance_icon": "浮窗入口图标",
     "preview_error": "当前页面禁止使用 Iframe 嵌入！",
-    "assistant_app": "小助手应用"
+    "assistant_app": "小助手应用",
+    "auto_select_ds": "自动选择数据源"
   },
   "chat": {
     "type": "图表类型",

frontend/src/stores/assistant.ts

@@ -22,6 +22,7 @@ interface AssistantState {
   pageEmbedded?: boolean
   history: boolean
   hostOrigin: string
+  autoDs?: boolean
   requestPromiseMap: Map<string, PendingRequest>
 }
 
@@ -38,6 +39,7 @@ export const AssistantStore = defineStore('assistant', {
       pageEmbedded: false,
       history: true,
       hostOrigin: '',
+      autoDs: false,
       requestPromiseMap: new Map<string, PendingRequest>(),
     }
   },
@@ -75,6 +77,9 @@ export const AssistantStore = defineStore('assistant', {
     getHostOrigin(): string {
       return this.hostOrigin
     },
+    getAutoDs(): boolean {
+      return !!this.autoDs
+    },
   },
   actions: {
     refreshCertificate<T>() {
@@ -146,6 +151,9 @@ export const AssistantStore = defineStore('assistant', {
     setHostOrigin(origin: string) {
       this.hostOrigin = origin
     },
+    setAutoDs(autoDs?: boolean) {
+      this.autoDs = !!autoDs
+    },
     async setChat() {
       if (!this.assistant) {
         return null

frontend/src/views/chat/ChatCreator.vue

@@ -9,10 +9,15 @@ import { datasourceApi } from '@/api/datasource.ts'
 import Card from '@/views/ds/ChatCard.vue'
 import AddDrawer from '@/views/ds/AddDrawer.vue'
 import { useUserStore } from '@/stores/user'
-
+import { useAssistantStore } from '@/stores/assistant'
+import { request } from '@/utils/request'
+const assistantStore = useAssistantStore()
 const userStore = useUserStore()
 
 const isWsAdmin = computed(() => userStore.isAdmin || userStore.isSpaceAdmin)
+const selectAssistantDs = computed(
+  () => assistantStore.getAssistant && !assistantStore.getEmbedded && !assistantStore.getAutoDs
+)
 const props = withDefaults(
   defineProps<{
     hidden?: boolean
@@ -42,8 +47,7 @@ const emits = defineEmits(['onChatCreated'])
 
 function listDs() {
   searchLoading.value = true
-  datasourceApi
-    .list()
+  ;(selectAssistantDs.value ? request.get('/system/assistant/ds') : datasourceApi.list())
     .then((res) => {
       datasourceList.value = res
     })
@@ -73,6 +77,10 @@ function selectDsInDialog(ds: any) {
 
 function confirmSelectDs() {
   if (innerDs.value) {
+    if (assistantStore.getType == 1) {
+      createChat(innerDs.value)
+      return
+    }
     statusLoading.value = true
     //check first
     datasourceApi
@@ -90,10 +98,15 @@ function confirmSelectDs() {
 
 function createChat(datasource: number) {
   loading.value = true
-  chatApi
-    .startChat({
-      datasource: datasource,
-    })
+  const param = {
+    datasource: datasource,
+  } as any
+  let method = chatApi.startChat
+  if (assistantStore.getAssistant) {
+    param['origin'] = 2
+    method = chatApi.startAssistantChat
+  }
+  method(param)
     .then((res) => {
       const chat: ChatInfo | undefined = chatApi.toChatInfo(res)
       if (chat == undefined) {
@@ -140,11 +153,11 @@ defineExpose({
     >
       <template #header="{ close }">
         <span style="white-space: nowrap">{{ $t('qa.select_datasource') }}</span>
-        <div class="flex-center" style="width: 100%">
+        <div class="flex-center" style="width: 100%; margin-right: 32px">
           <el-input
             v-model="keywords"
             clearable
-            style="width: 320px"
+            style="width: 320px; max-width: calc(100% - 32px)"
             :placeholder="$t('datasource.search')"
           >
             <template #prefix>

frontend/src/views/chat/index.vue

@@ -147,7 +147,7 @@
             </div>
 
             <el-button
-              v-if="isCompletePage && currentChatId === undefined"
+              v-if="(isCompletePage || selectAssistantDs) && currentChatId === undefined"
               size="large"
               type="primary"
               class="greeting-btn"
@@ -360,9 +360,12 @@
           </div>
         </el-scrollbar>
       </el-main>
-      <el-footer v-if="computedMessages.length > 0 || !isCompletePage" class="chat-footer">
+      <el-footer
+        v-if="computedMessages.length > 0 || (!isCompletePage && !selectAssistantDs)"
+        class="chat-footer"
+      >
         <div class="input-wrapper" @click="clickInput">
-          <div v-if="isCompletePage" class="datasource">
+          <div v-if="isCompletePage || selectAssistantDs" class="datasource">
             <template v-if="currentChat.datasource && currentChat.datasource_name">
               {{ t('qa.selected_datasource') }}:
               <img
@@ -397,7 +400,7 @@
             :disabled="isTyping"
             clearable
             class="input-area"
-            :class="!isCompletePage && 'is-assistant'"
+            :class="!isCompletePage && !selectAssistantDs && 'is-assistant'"
             type="textarea"
             :autosize="{ minRows: 1, maxRows: 8.583 }"
             :placeholder="t('qa.question_placeholder')"
@@ -420,7 +423,11 @@
       </el-footer>
     </el-container>
 
-    <ChatCreator v-if="isCompletePage" ref="chatCreatorRef" @on-chat-created="onChatCreatedQuick" />
+    <ChatCreator
+      v-if="isCompletePage || selectAssistantDs"
+      ref="chatCreatorRef"
+      @on-chat-created="onChatCreatedQuick"
+    />
     <ChatCreator ref="hiddenChatCreatorRef" hidden @on-chat-created="onChatCreatedQuick" />
   </el-container>
 </template>
@@ -482,6 +489,10 @@ const isCompletePage = computed(() => !assistantStore.getAssistant || assistantS
 const embeddedHistoryHidden = computed(
   () => assistantStore.getAssistant && !assistantStore.getHistory
 )
+// const autoDs = computed(() => assistantStore.getAssistant && assistantStore.getAutoDs)
+const selectAssistantDs = computed(() => {
+  return assistantStore.getAssistant && !assistantStore.getAutoDs
+})
 const customName = computed(() => {
   if (!isCompletePage.value && props.pageEmbedded) return props.appName
   return ''
@@ -634,7 +645,7 @@ const createNewChat = async () => {
     return
   }
   goEmpty()
-  if (!isCompletePage.value) {
+  if (!isCompletePage.value && !selectAssistantDs.value) {
     currentChat.value = new ChatInfo()
     currentChatId.value = undefined
     return
@@ -764,6 +775,7 @@ function onChatStop() {
 const assistantPrepareSend = async () => {
   if (
     !isCompletePage.value &&
+    !selectAssistantDs.value &&
     (currentChatId.value == null || typeof currentChatId.value == 'undefined')
   ) {
     const assistantChat = await assistantStore.setChat()