Merge remote-tracking branch 'origin/main'

Author: ziyujiahao

backend/apps/chat/curd/chat.py

@@ -189,6 +189,7 @@ def create_chat(session: SessionDep, current_user: CurrentUser, create_chat_obj:
     if ds:
         chat_info.datasource_exists = True
         chat_info.datasource_name = ds.name
+        chat_info.ds_type = ds.type
 
     if require_datasource and ds:
         # generate first empty record

backend/apps/datasource/crud/datasource.py

@@ -41,7 +41,7 @@ def get_ds(session: SessionDep, id: int):
 
 
 def check_status(session: SessionDep, ds: CoreDatasource, is_raise: bool = False):
-    conn = get_engine(ds, 5)
+    conn = get_engine(ds, 10)
     try:
         with conn.connect() as connection:
             SQLBotLogUtil.info("success")

backend/apps/db/db.py

@@ -52,28 +52,24 @@ def get_uri_from_config(type: str, conf: DatasourceConf) -> str:
     return db_url
 
 
-def get_engine(ds: CoreDatasource, timeout: int = 30) -> Engine:
+def get_engine(ds: CoreDatasource, timeout: int = 0) -> Engine:
     conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration))) if ds.type != "excel" else get_engine_config()
     if conf.timeout is None:
         conf.timeout = timeout
+    if timeout > 0:
+        conf.timeout = timeout
     if ds.type == "pg" and (conf.dbSchema is not None and conf.dbSchema != ""):
         engine = create_engine(get_uri(ds),
                                connect_args={"options": f"-c search_path={urllib.parse.quote(conf.dbSchema)}",
                                              "connect_timeout": conf.timeout},
-                               pool_timeout=conf.timeout, pool_size=20, max_overflow=10)
+                               pool_timeout=conf.timeout)
     elif ds.type == 'sqlServer':
-        engine = create_engine(get_uri(ds), pool_timeout=conf.timeout,
-                               pool_size=20,
-                               max_overflow=10)
+        engine = create_engine(get_uri(ds), pool_timeout=conf.timeout)
     elif ds.type == 'oracle':
         engine = create_engine(get_uri(ds),
-                               pool_timeout=conf.timeout,
-                               pool_size=20,
-                               max_overflow=10)
+                               pool_timeout=conf.timeout)
     else:
-        engine = create_engine(get_uri(ds), connect_args={"connect_timeout": conf.timeout}, pool_timeout=conf.timeout,
-                               pool_size=20,
-                               max_overflow=10)
+        engine = create_engine(get_uri(ds), connect_args={"connect_timeout": conf.timeout}, pool_timeout=conf.timeout)
     return engine
 
 

backend/apps/db/engine.py

@@ -25,9 +25,7 @@ def get_engine_conn():
     db_url = get_engine_uri(conf)
     engine = create_engine(db_url,
                            connect_args={"options": f"-c search_path={conf.dbSchema}", "connect_timeout": conf.timeout},
-                           pool_timeout=conf.timeout,
-                           pool_size=20,
-                           max_overflow=10)
+                           pool_timeout=conf.timeout)
     return engine
 
 

backend/apps/system/crud/assistant.py

@@ -1,29 +1,33 @@
-
-
 import json
+import urllib
 from typing import Optional
-from fastapi import FastAPI
+
 import requests
+from fastapi import FastAPI
 from sqlalchemy import Engine, create_engine
 from sqlmodel import Session, select
-from apps.datasource.models.datasource import CoreDatasource, DatasourceConf
+from starlette.middleware.cors import CORSMiddleware
 
+from apps.datasource.models.datasource import CoreDatasource, DatasourceConf
 from apps.system.models.system_model import AssistantModel
 from apps.system.schemas.auth import CacheName, CacheNamespace
 from apps.system.schemas.system_schema import AssistantHeader, AssistantOutDsSchema, UserInfoDTO
-from common.core.sqlbot_cache import cache
-from common.core.db import engine
-from starlette.middleware.cors import CORSMiddleware
 from common.core.config import settings
+from common.core.db import engine
+from common.core.sqlbot_cache import cache
 from common.utils.utils import string_to_numeric_hash
 
+
 @cache(namespace=CacheNamespace.EMBEDDED_INFO, cacheName=CacheName.ASSISTANT_INFO, keyExpression="assistant_id")
 async def get_assistant_info(*, session: Session, assistant_id: int) -> AssistantModel | None:
     db_model = session.get(AssistantModel, assistant_id)
     return db_model
 
+
 def get_assistant_user(*, id: int):
-    return UserInfoDTO(id=id, account="sqlbot-inner-assistant", oid=1, name="sqlbot-inner-assistant", email="[email protected]")
+    return UserInfoDTO(id=id, account="sqlbot-inner-assistant", oid=1, name="sqlbot-inner-assistant",
+                       email="[email protected]")
+
 
 def get_assistant_ds(llm_service) -> list[dict]:
     assistant: AssistantHeader = llm_service.current_assistant
@@ -34,13 +38,14 @@ def get_assistant_ds(llm_service) -> list[dict]:
         if configuration:
             config: dict[any] = json.loads(configuration)
             oid: int = int(config['oid'])
-            stmt = select(CoreDatasource.id, CoreDatasource.name, CoreDatasource.description).where(CoreDatasource.oid == oid)
+            stmt = select(CoreDatasource.id, CoreDatasource.name, CoreDatasource.description).where(
+                CoreDatasource.oid == oid)
             if not assistant.online:
-                private_list:list[int] = config.get('private_list') or None
+                private_list: list[int] = config.get('private_list') or None
                 if private_list:
                     stmt = stmt.where(~CoreDatasource.id.in_(private_list))
         db_ds_list = session.exec(stmt)
-        
+
         result_list = [
             {
                 "id": ds.id,
@@ -49,7 +54,7 @@ def get_assistant_ds(llm_service) -> list[dict]:
             }
             for ds in db_ds_list
         ]
-    
+
         # filter private ds if offline
         return result_list
     out_ds_instance: AssistantOutDs = AssistantOutDsFactory.get_instance(assistant)
@@ -58,8 +63,9 @@ def get_assistant_ds(llm_service) -> list[dict]:
     # format?
     return dslist
 
+
 def init_dynamic_cors(app: FastAPI):
-    try: 
+    try:
         with Session(engine) as session:
             list_result = session.exec(select(AssistantModel).order_by(AssistantModel.create_time)).all()
             seen = set()
@@ -81,20 +87,20 @@ def init_dynamic_cors(app: FastAPI):
                 cors_middleware.kwargs['allow_origins'] = updated_origins
     except Exception as e:
         return False, e
-    
-    
+
 
 class AssistantOutDs:
     assistant: AssistantHeader
     ds_list: Optional[list[AssistantOutDsSchema]] = None
     certificate: Optional[str] = None
+
     def __init__(self, assistant: AssistantHeader):
         self.assistant = assistant
         self.ds_list = None
         self.certificate = assistant.certificate
         self.get_ds_from_api()
-        
-    #@cache(namespace=CacheNamespace.EMBEDDED_INFO, cacheName=CacheName.ASSISTANT_DS, keyExpression="current_user.id")    
+
+    # @cache(namespace=CacheNamespace.EMBEDDED_INFO, cacheName=CacheName.ASSISTANT_DS, keyExpression="current_user.id")
     def get_ds_from_api(self):
         config: dict[any] = json.loads(self.assistant.configuration)
         endpoint: str = config['endpoint']
@@ -118,23 +124,23 @@ def get_ds_from_api(self):
                     self.convert2schema(item)
                     for item in temp_list
                 ]
-                
+
                 return self.ds_list
             else:
                 raise Exception(f"Failed to get datasource list from {endpoint}, error: {result_json.get('message')}")
         else:
             raise Exception(f"Failed to get datasource list from {endpoint}, status code: {res.status_code}")
-    
+
     def get_simple_ds_list(self):
         if self.ds_list:
             return [{'id': ds.id, 'name': ds.name, 'description': ds.comment} for ds in self.ds_list]
         else:
             raise Exception("Datasource list is not found.")
-       
+
     def get_db_schema(self, ds_id: int) -> str:
         ds = self.get_ds(ds_id)
         schema_str = ""
-        #db_name = ds.db_schema
+        # db_name = ds.db_schema
         db_name = ds.db_schema if ds.db_schema is not None and ds.db_schema != "" else ds.dataBase
         schema_str += f"【DB_ID】 {db_name}\n【Schema】\n"
         for table in ds.tables:
@@ -144,7 +150,7 @@ def get_db_schema(self, ds_id: int) -> str:
                 schema_str += '\n[\n'
             else:
                 schema_str += f", {table_comment}\n[\n"
-                
+
             field_list = []
             for field in table.fields:
                 field_comment = field.comment
@@ -155,7 +161,7 @@ def get_db_schema(self, ds_id: int) -> str:
             schema_str += ",\n".join(field_list)
             schema_str += '\n]\n'
         return schema_str
-    
+
     def get_ds(self, ds_id: int):
         if self.ds_list:
             for ds in self.ds_list:
@@ -175,20 +181,22 @@ def convert2schema(self, ds_dict: dict) -> AssistantOutDsSchema:
         db_schema = ds_dict.get('schema', ds_dict.get('db_schema', ''))
         ds_dict.pop("schema", None)
         return AssistantOutDsSchema(**{**ds_dict, "id": id, "db_schema": db_schema})
-    
+
+
 class AssistantOutDsFactory:
     @staticmethod
     def get_instance(assistant: AssistantHeader) -> AssistantOutDs:
         return AssistantOutDs(assistant)
 
+
 def get_ds_engine(ds: AssistantOutDsSchema) -> Engine:
     timeout: int = 30
     connect_args = {"connect_timeout": timeout}
     conf = DatasourceConf(
-        host=ds.host, 
-        port=ds.port, 
+        host=ds.host,
+        port=ds.port,
         username=ds.user,
-        password=ds.password,   
+        password=ds.password,
         database=ds.dataBase,
         driver='',
         extraJdbc=ds.extraParams,
@@ -197,8 +205,20 @@ def get_ds_engine(ds: AssistantOutDsSchema) -> Engine:
     conf.extraJdbc = ''
     from apps.db.db import get_uri_from_config
     uri = get_uri_from_config(ds.type, conf)
+    # if ds.type == "pg" and ds.db_schema:
+    #     connect_args.update({"options": f"-c search_path={ds.db_schema}"})
+    # engine = create_engine(uri, connect_args=connect_args, pool_timeout=timeout, pool_size=20, max_overflow=10)
+
     if ds.type == "pg" and ds.db_schema:
-        connect_args.update({"options": f"-c search_path={ds.db_schema}"})
-    engine = create_engine(uri, connect_args=connect_args, pool_timeout=timeout, pool_size=20, max_overflow=10)
+        engine = create_engine(uri,
+                               connect_args={"options": f"-c search_path={urllib.parse.quote(ds.db_schema)}",
+                                             "connect_timeout": timeout},
+                               pool_timeout=timeout)
+    elif ds.type == 'sqlServer':
+        engine = create_engine(uri, pool_timeout=timeout)
+    elif ds.type == 'oracle':
+        engine = create_engine(uri,
+                               pool_timeout=timeout)
+    else:
+        engine = create_engine(uri, connect_args={"connect_timeout": timeout}, pool_timeout=timeout)
     return engine
-    

backend/template.yaml

@@ -7,6 +7,7 @@ template:
       根据表结构和问题生成符合{engine}数据库引擎规范的sql语句，以及sql中所用到的表名（不要包含schema和database，用数组返回）。
       你必须遵守以下规则:
       - 生成的SQL必须符合{engine}的规范。
+      - 提问中如果有涉及数据源名称或数据源描述的内容，则忽略数据源的信息，直接根据剩余内容生成SQL
       - 根据表结构生成SQL语句，需给每个表名生成一个别名（不要加AS）。
       - SQL查询中不能使用星号(*)，必须明确指定字段名.
       - SQL查询的字段名不要自动翻译，别名必须为英文。

frontend/public/assistant.js

@@ -1,4 +1,4 @@
-;(function () {
+; (function () {
   window.sqlbot_assistant_handler = window.sqlbot_assistant_handler || {}
   const defaultData = {
     id: '1',
@@ -68,7 +68,7 @@
   const getChatContainerHtml = (data) => {
     return `
 <div id="sqlbot-assistant-chat-container">
-  <iframe id="sqlbot-assistant-chat-iframe-${data.id}" allow="microphone" src="${data.domain_url}/#/assistant?id=${data.id}&online=${!!data.online}"></iframe>
+  <iframe id="sqlbot-assistant-chat-iframe-${data.id}" allow="microphone;clipboard-read 'src'; clipboard-write 'src'" src="${data.domain_url}/#/assistant?id=${data.id}&online=${!!data.online}"></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">

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

@@ -18,6 +18,7 @@ import icon_window_mini_outlined from '@/assets/svg/icon_window-mini_outlined.sv
 import icon_copy_outlined from '@/assets/svg/icon_copy_outlined.svg'
 import { useI18n } from 'vue-i18n'
 import SQLComponent from '@/views/chat/component/SQLComponent.vue'
+import { useAssistantStore } from '@/stores/assistant'
 import AddViewDashboard from '@/views/dashboard/common/AddViewDashboard.vue'
 
 const props = withDefaults(
@@ -51,6 +52,8 @@ const dataObject = computed<{
   }
   return {}
 })
+const assistantStore = useAssistantStore()
+const isAssistant = computed(() => assistantStore.getAssistant)
 
 const data = computed(() => {
   if (props.isPredict) {
@@ -325,7 +328,7 @@ function copy() {
 
     <el-drawer
       v-model="sqlShow"
-      size="600"
+      :size="isAssistant ? '100%' : '600px'"
       :title="t('chat.show_sql')"
       direction="rtl"
       body-class="chart-sql-drawer-body"

frontend/src/views/chat/index.vue

@@ -296,7 +296,7 @@
             type="textarea"
             :autosize="{ minRows: 1, maxRows: 8.583 }"
             :placeholder="t('qa.question_placeholder')"
-            @keydown.enter.exact.prevent="sendMessage"
+            @keydown.enter.exact.prevent="($event: any) => sendMessage($event)"
             @keydown.ctrl.enter.exact.prevent="handleCtrlEnter"
           />
 
@@ -551,7 +551,10 @@ const assistantPrepareSend = async () => {
     }
   }
 }
-const sendMessage = async () => {
+const sendMessage = async ($event: any = {}) => {
+  if ($event?.isComposing) {
+    return
+  }
   if (!inputMessage.value.trim()) return
 
   loading.value = true

frontend/src/views/embedded/index.vue

@@ -191,3 +191,10 @@ onBeforeUnmount(() => {
   }
 }
 </style>
+
+<style lang="less">
+.ed-overlay-dialog,
+.ed-drawer {
+  margin-top: 50px;
+}
+</style>