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

.typos.toml

@@ -0,0 +1,5 @@
+[default.extend-words]
+[default.extend-identifiers]
+maintain_column_froms = "maintain_column_froms"
+[files]
+extend-exclude = ["frontend/public", "backend/alembic/versions"]

backend/apps/chat/task/llm.py

@@ -296,7 +296,7 @@ class LLMService:
 
         # get schema
         if self.ds and not self.chat_question.db_schema:
-            self.chat_question.db_schema = get_table_schema(session=self.session, ds=self.ds)
+            self.chat_question.db_schema = get_table_schema(session=self.session, current_user=self.current_user, ds=self.ds)
 
         guess_msg: List[Union[BaseMessage, dict[str, Any]]] = []
         guess_msg.append(SystemMessage(content=self.chat_question.guess_sys_question()))
@@ -726,7 +726,7 @@ def run_task(llm_service: LLMService, in_chat: bool = True):
                 yield orjson.dumps({'id': llm_service.ds.id, 'datasource_name': llm_service.ds.name,
                                     'engine_type': llm_service.ds.type_name, 'type': 'datasource'}).decode() + '\n\n'
 
-            llm_service.chat_question.db_schema = get_table_schema(session=llm_service.session, ds=llm_service.ds)
+            llm_service.chat_question.db_schema = get_table_schema(session=llm_service.session, current_user=llm_service.current_user, ds=llm_service.ds)
 
         # generate sql
         sql_res = llm_service.generate_sql()

backend/apps/datasource/api/datasource.py

@@ -9,7 +9,7 @@ from fastapi import APIRouter, File, UploadFile, HTTPException
 from apps.db.engine import create_table, get_data_engine, insert_data
 from common.core.deps import SessionDep, CurrentUser
 from ..crud.datasource import get_datasource_list, check_status, create_ds, update_ds, delete_ds, getTables, getFields, \
-    execSql, update_table_and_fields, getTablesByDs, chooseTables, preview, updateTable, updateField, get_ds
+    execSql, update_table_and_fields, getTablesByDs, chooseTables, preview, updateTable, updateField, get_ds, fieldEnum
 from ..crud.field import get_fields_by_table_id
 from ..crud.table import get_tables_by_ds_id
 from ..models.datasource import CoreDatasource, CreateDatasource, TableObj, CoreTable, CoreField
@@ -103,6 +103,11 @@ async def edit_local(session: SessionDep, id: int, data: TableObj):
     return preview(session, id, data)
 
 
+@router.post("/fieldEnum/{id}")
+async def field_enum(session: SessionDep, id: int):
+    return fieldEnum(session, id)
+
+
 @router.post("/uploadExcel")
 async def upload_excel(session: SessionDep, file: UploadFile = File(...)):
     ALLOWED_EXTENSIONS = {"xlsx", "xls", "csv"}

backend/apps/datasource/crud/datasource.py

@@ -9,9 +9,9 @@ from sqlbot_xpack.permissions.models.ds_rules import DsRules
 from sqlmodel import select
 
 from apps.datasource.utils.utils import aes_decrypt
+from apps.db.constant import DB
 from apps.db.db import get_engine, get_tables, get_fields, exec_sql
-from apps.db.engine import get_engine_config
-from apps.db.engine import get_engine_conn
+from apps.db.engine import get_engine_config, get_engine_conn
 from apps.db.type import db_type_relation
 from common.core.deps import SessionDep, CurrentUser
 from common.utils.utils import deepcopy_ignore_extra
@@ -245,6 +245,23 @@ def preview(session: SessionDep, id: int, data: TableObj):
     return exec_sql(ds, sql)
 
 
+def fieldEnum(session: SessionDep, id: int):
+    field = session.query(CoreField).filter(CoreField.id == id).first()
+    if field is None:
+        return []
+    table = session.query(CoreTable).filter(CoreTable.id == field.table_id).first()
+    if table is None:
+        return []
+    ds = session.query(CoreDatasource).filter(CoreDatasource.id == table.ds_id).first()
+    if ds is None:
+        return []
+
+    db = DB.get_db(ds.type)
+    sql = f"""SELECT DISTINCT {db.prefix}{field.field_name}{db.suffix} FROM {db.prefix}{table.table_name}{db.suffix}"""
+    res = exec_sql(ds, sql)
+    return [item.get(res.get('fields')[0]) for item in res.get('data')]
+
+
 def updateNum(session: SessionDep, ds: CoreDatasource):
     all_tables = get_tables(ds)
     selected_tables = get_tables_by_ds_id(session, ds.id)

backend/apps/db/constant.py

@@ -0,0 +1,24 @@
+# Author: Junjun
+# Date: 2025/7/16
+
+from enum import Enum
+
+
+class DB(Enum):
+    mysql = ('mysql', '`', '`')
+    sqlServer = ('sqlServer', '[', ']')
+    pg = ('pg', '"', '"')
+    excel = ('excel', '"', '"')
+    oracle = ('oracle', '"', '"')
+
+    def __init__(self, type, prefix, suffix):
+        self.type = type
+        self.prefix = prefix
+        self.suffix = suffix
+
+    @classmethod
+    def get_db(cls, type):
+        for db in cls:
+            if db.type == type:
+                return db
+        raise ValueError(f"Invalid db type: {type}")

backend/apps/system/api/workspace.py

@@ -3,7 +3,7 @@ from fastapi import APIRouter, Query
 from sqlmodel import exists, or_, select    
 from apps.system.models.system_model import UserWsModel, WorkspaceBase, WorkspaceEditor, WorkspaceModel
 from apps.system.models.user import UserModel
-from apps.system.schemas.system_schema import UserWsBase, UserWsDTO, UserWsOption, WorkspaceUser
+from apps.system.schemas.system_schema import UserWsBase, UserWsDTO, UserWsEditor, UserWsOption, WorkspaceUser
 from common.core.deps import CurrentUser, SessionDep, Trans
 from common.core.pagination import Paginator
 from common.core.schemas import PaginatedResponse, PaginationParams
@@ -42,6 +42,29 @@ async def option_pager(
         stmt=stmt,
         pagination=pagination,
     )
+    
+@router.get("/uws/option", response_model=UserWsOption)
+async def option_user(
+    session: SessionDep, 
+    current_user: CurrentUser,
+    keyword: str = Query(description="搜索关键字")
+    ):
+    if (not current_user.isAdmin) and current_user.weight == 0:
+        raise RuntimeError("no permission to execute this api")
+    oid = current_user.oid
+    
+    stmt = select(UserModel.id, UserModel.account, UserModel.name).where(
+        ~exists().where(UserWsModel.uid == UserModel.id, UserWsModel.oid == oid)
+    )
+    
+    if keyword:
+        stmt = stmt.where(
+            or_(
+                UserModel.account == int(keyword),
+                UserModel.name == keyword,
+            )
+        )
+    return session.exec(stmt).first()
 
 @router.get("/uws/pager/{pageNum}/{pageSize}", response_model=PaginatedResponse[WorkspaceUser])
 async def pager(
@@ -95,6 +118,20 @@ async def create(session: SessionDep, creator: UserWsDTO):
     session.add_all(db_model_list)
     session.commit()
 
+@router.put("/uws")     
+async def edit(session: SessionDep, editor: UserWsEditor):
+    if not editor.oid or not editor.uid:
+        raise RuntimeError("param [oid, uid] miss")
+    db_model = session.exec(select(UserWsModel).where(UserWsModel.uid == editor.uid, UserWsModel.oid == editor.oid)).first()
+    if not db_model:
+        raise RuntimeError("uws not exist")
+    if editor.weight == db_model.weight:
+        return
+    
+    db_model.weight = editor.weight
+    session.add(db_model)
+    session.commit()
+
 @router.delete("/uws")     
 async def delete(session: SessionDep, dto: UserWsBase):
     db_model_list: list[UserWsModel] = session.exec(select(UserWsModel).where(UserWsModel.uid.in_(dto.uid_list), UserWsModel.oid == dto.oid)).all()

backend/apps/system/schemas/system_schema.py

@@ -50,7 +50,11 @@ class UserWsBase(BaseModel):
     oid: int
 class UserWsDTO(UserWsBase):
     weight: int = 0
-    
+
+class UserWsEditor(BaseModel):
+    uid: int
+    oid: int    
+    weight: int = 0
 
 class UserInfoDTO(UserEditor):
     language: str = "zh-CN"

frontend/src/assets/svg/icon_new_chat_outlined.svg

@@ -1,4 +1 @@
-<svg width="18" height="18" viewBox="0 0 18 18" fill="none" xmlns="http://www.w3.org/2000/svg">
-<path d="M9 1.125C13.3492 1.125 16.875 4.65076 16.875 9C16.875 13.3492 13.3492 16.875 9 16.875H1.875C1.46079 16.875 1.125 16.5392 1.125 16.125V9C1.125 4.65076 4.65076 1.125 9 1.125ZM2.625 15.375H9C12.5208 15.375 15.375 12.5208 15.375 8.99998C15.375 5.47917 12.5208 2.62498 9 2.62498C5.47918 2.62498 2.625 5.47917 2.625 8.99998V15.375Z" fill="currentColor"/>
-<path d="M8.625 5.625C8.41789 5.625 8.25 5.79289 8.25 6V8.25H6C5.79289 8.25 5.625 8.41789 5.625 8.625V9.375C5.625 9.58211 5.79289 9.75 6 9.75H8.25V12C8.25 12.2071 8.41789 12.375 8.625 12.375H9.375C9.58211 12.375 9.75 12.2071 9.75 12V9.75H12C12.2071 9.75 12.375 9.58211 12.375 9.375V8.625C12.375 8.41789 12.2071 8.25 12 8.25H9.75V6C9.75 5.79289 9.58211 5.625 9.375 5.625H8.625Z" fill="currentColor"/>
-</svg>
+<?xml version="1.0" standalone="no"?><!DOCTYPE svg PUBLIC "-//W3C//DTD SVG 1.1//EN" "http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd"><svg t="1752641977067" class="icon" viewBox="0 0 1024 1024" version="1.1" xmlns="http://www.w3.org/2000/svg" p-id="4591" xmlns:xlink="http://www.w3.org/1999/xlink" width="256" height="256"><path d="M512 983.04H143.36c-25.6 0-51.2-10.24-71.68-30.72-20.48-20.48-30.72-46.08-30.72-71.68V512c0-122.88 51.2-245.76 138.24-332.8C266.24 92.16 384 40.96 512 40.96c122.88 0 245.76 51.2 332.8 138.24C931.84 266.24 983.04 384 983.04 512c0 122.88-51.2 245.76-138.24 332.8-87.04 87.04-209.92 138.24-332.8 138.24z m0-849.92c-102.4 0-194.56 40.96-266.24 112.64S133.12 409.6 133.12 512v368.64c0 5.12 0 5.12 5.12 5.12l5.12 5.12H512c102.4 0 194.56-40.96 266.24-112.64S890.88 614.4 890.88 512s-40.96-194.56-112.64-266.24S614.4 133.12 512 133.12z" fill="" p-id="4592"></path><path d="M322.56 471.04h373.76c15.36 0 25.6 10.24 25.6 25.6v35.84c0 15.36-10.24 25.6-25.6 25.6H322.56c-15.36 0-25.6-10.24-25.6-25.6v-35.84c0-15.36 15.36-25.6 25.6-25.6z" fill="" p-id="4593"></path><path d="M552.96 322.56v373.76c0 15.36-10.24 25.6-25.6 25.6h-30.72c-15.36 0-25.6-10.24-25.6-25.6V322.56c0-15.36 10.24-25.6 25.6-25.6h35.84c10.24 0 20.48 15.36 20.48 25.6z" fill="" p-id="4594"></path></svg>

frontend/src/views/chat/index.vue

@@ -72,11 +72,14 @@
               v-if="currentChatId === undefined"
               size="large"
               type="primary"
+              class="greeting-btn"
               @click="createNewChat"
             >
-              <el-icon>
-                <icon_new_chat_outlined />
-              </el-icon>
+              <span class="inner-icon">
+                <el-icon>
+                  <icon_new_chat_outlined />
+                </el-icon>
+              </span>
               {{ t('qa.start_sqlbot') }}
             </el-button>
           </div>
@@ -292,7 +295,7 @@ const createNewChat = async () => {
   if (isAssistant.value) {
     const assistantChat = await assistantStore.setChat()
     if (assistantChat) {
-      onChatCreated(assistantChat as any)
+      onChatCreatedQuick(assistantChat as any)
     }
     return
   }
@@ -317,6 +320,26 @@ function onClickHistory(chat: Chat) {
   scrollToBottom()
 }
 
+function toAssistantHistory(chat: Chat) {
+  currentChat.value = new ChatInfo(chat)
+  if (chat !== undefined && chat.id !== undefined && !loading.value) {
+    currentChatId.value = chat.id
+    loading.value = true
+    chatApi
+      .get(chat.id)
+      .then((res) => {
+        const info = chatApi.toChatInfo(res)
+        if (info) {
+          currentChat.value = info
+          onClickHistory(info)
+        }
+      })
+      .finally(() => {
+        loading.value = false
+      })
+  }
+}
+
 function onChatDeleted(id: number) {
   console.log('deleted', id)
 }
@@ -434,11 +457,6 @@ const sendMessage = async () => {
   loading.value = true
   isTyping.value = true
 
-  /* const assistantChat = await assistantStore.setChat()
-  if (assistantChat) {
-    onChatCreated(assistantChat as any)
-  } */
-
   const currentRecord = new ChatRecord()
   currentRecord.create_time = new Date()
   currentRecord.chat_id = currentChatId.value
@@ -833,10 +851,9 @@ const getCurrentChatId = () => {
 }
 defineExpose({
   getHistoryList,
-  onClickHistory,
-  onChatDeleted,
-  onChatRenamed,
+  toAssistantHistory,
   getCurrentChatId,
+  createNewChat,
 })
 </script>
 
@@ -879,7 +896,6 @@ defineExpose({
 
   .chat-record-list {
     padding: 0 0 20px 0;
-    background: rgba(255, 255, 255, 1);
     border-radius: 0 12px 12px 0;
 
     &.hide-sidebar {
@@ -974,6 +990,35 @@ defineExpose({
       font-size: 16px;
       line-height: 24px;
     }
+
+    .greeting-btn {
+      width: 100%;
+      height: 88px;
+
+      border-style: dashed;
+
+      .inner-icon {
+        display: flex;
+        flex-direction: row;
+        align-items: center;
+
+        margin-right: 6px;
+      }
+
+      font-size: 16px;
+      line-height: 24px;
+      font-weight: 500;
+
+      --ed-button-text-color: rgba(28, 186, 144, 1);
+      --ed-button-hover-text-color: rgba(28, 186, 144, 1);
+      --ed-button-active-text-color: rgba(28, 186, 144, 1);
+      --ed-button-bg-color: rgba(248, 249, 250, 1);
+      --ed-button-hover-bg-color: rgba(28, 186, 144, 0.1);
+      --ed-button-border-color: rgba(217, 220, 223, 1);
+      --ed-button-hover-border-color: rgba(28, 186, 144, 1);
+      --ed-button-active-bg-color: rgba(28, 186, 144, 0.2);
+      --ed-button-active-border-color: rgba(28, 186, 144, 1);
+    }
   }
 }
 </style>

frontend/src/views/embedded/assistant.vue

@@ -11,6 +11,9 @@
     <chat-component v-if="!loading" ref="chatRef" class="sqlbot-chat-container" />
   </div>
   <div class="sqlbot-top-btn">
+    <el-icon style="cursor: pointer" @click="createChat">
+      <icon_new_chat_outlined />
+    </el-icon>
     <el-icon style="cursor: pointer" @click="openHistory">
       <history></history>
     </el-icon>
@@ -41,10 +44,10 @@
                 @click="onClickHistory(chat)"
               >
                 <span class="title">{{ chat.brief ?? 'Untitled' }}</span>
-                <div class="history-operate">
+                <!-- <div class="history-operate">
                   <el-icon @click="EditPen(chat)"><IconOpeEdit /></el-icon>
                   <el-icon @click="Delete(chat)"><IconOpeDelete /></el-icon>
-                </div>
+                </div> -->
               </div>
             </template>
           </el-scrollbar>
@@ -58,8 +61,9 @@ import { onBeforeMount, 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'
-import IconOpeEdit from '@/assets/svg/operate/ope-edit.svg'
-import IconOpeDelete from '@/assets/svg/operate/ope-delete.svg'
+import icon_new_chat_outlined from '@/assets/svg/icon_new_chat_outlined.svg'
+/* import IconOpeEdit from '@/assets/svg/operate/ope-edit.svg'
+import IconOpeDelete from '@/assets/svg/operate/ope-delete.svg' */
 import { useRoute } from 'vue-router'
 import { assistantApi } from '@/api/assistant'
 import { useAssistantStore } from '@/stores/assistant'
@@ -72,21 +76,24 @@ const chatList = ref<Array<any>>([])
 const drawer = ref(false)
 const currentId = ref()
 
+const createChat = () => {
+  chatRef.value?.createNewChat()
+}
 const openHistory = () => {
   chatList.value = chatRef.value?.getHistoryList()
   currentId.value = chatRef.value?.getCurrentChatId()
   drawer.value = true
 }
 const onClickHistory = (chat: any) => {
-  chatRef.value?.onClickHistory(chat)
+  chatRef.value?.toAssistantHistory(chat)
 }
 
-const EditPen = (chat: any) => {
+/* const EditPen = (chat: any) => {
   chatRef.value?.onChatRenamed(chat)
 }
 const Delete = (chat: any) => {
   chatRef.value?.onChatDeleted(chat.id)
-}
+} */
 const validator = ref({
   id: '',
   valid: false,
@@ -153,10 +160,15 @@ onBeforeMount(async () => {
 .sqlbot-top-btn {
   right: 85px;
   z-index: 2009;
-  position: absolute;
+  position: fixed;
   top: 16px;
+  display: flex;
+  column-gap: 16px;
   // right: 16px;
   font-size: 22px;
+  i:first-child {
+    color: var(--primary-color);
+  }
 }
 .sqlbot-history-container {
   width: 100%;