Merge remote-tracking branch 'origin/main'

Author: ziyujiahao

backend/apps/chat/models/chat_model.py

@@ -147,7 +147,7 @@ def guess_user_question(self, old_questions: str = "[]"):
                                                             old_questions=old_questions)
 
     def filter_sys_question(self):
-        return get_permissions_template()['system'].format(lang=self.lang)
+        return get_permissions_template()['system'].format(lang=self.lang, engine=self.engine)
 
     def filter_user_question(self):
         return get_permissions_template()['user'].format(sql=self.sql, filter=self.filter)

backend/apps/chat/task/llm.py

@@ -386,51 +386,56 @@ def select_datasource(self):
             ]
             """ _ds_list = self.session.exec(select(CoreDatasource).options(
                 load_only(CoreDatasource.id, CoreDatasource.name, CoreDatasource.description))).all() """
-        _ds_list_dict = []
-        for _ds in _ds_list:
-            _ds_list_dict.append(_ds)
-        datasource_msg.append(
-            HumanMessage(self.chat_question.datasource_user_question(orjson.dumps(_ds_list_dict).decode())))
-
-        history_msg = []
-        if self.record.full_select_datasource_message and self.record.full_select_datasource_message.strip() != '':
-            history_msg = orjson.loads(self.record.full_select_datasource_message)
-
-        self.record = save_full_select_datasource_message_and_answer(session=self.session, record_id=self.record.id,
-                                                                     answer='',
-                                                                     full_message=orjson.dumps(history_msg +
-                                                                                               [{'type': msg.type,
-                                                                                                 'content': msg.content}
-                                                                                                for msg
-                                                                                                in
-                                                                                                datasource_msg]).decode())
-        full_thinking_text = ''
-        full_text = ''
-        token_usage = {}
-        res = self.llm.stream(datasource_msg)
-        for chunk in res:
-            SQLBotLogUtil.info(chunk)
-            reasoning_content_chunk = ''
-            if 'reasoning_content' in chunk.additional_kwargs:
-                reasoning_content_chunk = chunk.additional_kwargs.get('reasoning_content', '')
-            # else:
-            #     reasoning_content_chunk = chunk.get('reasoning_content')
-            if reasoning_content_chunk is None:
+        
+        ignore_auto_select = _ds_list and len(_ds_list) == 1
+        # ignore auto select ds
+        
+        if not ignore_auto_select:
+            _ds_list_dict = []
+            for _ds in _ds_list:
+                _ds_list_dict.append(_ds)
+            datasource_msg.append(
+                HumanMessage(self.chat_question.datasource_user_question(orjson.dumps(_ds_list_dict).decode())))
+
+            history_msg = []
+            if self.record.full_select_datasource_message and self.record.full_select_datasource_message.strip() != '':
+                history_msg = orjson.loads(self.record.full_select_datasource_message)
+
+            self.record = save_full_select_datasource_message_and_answer(session=self.session, record_id=self.record.id,
+                                                                        answer='',
+                                                                        full_message=orjson.dumps(history_msg +
+                                                                                                [{'type': msg.type,
+                                                                                                    'content': msg.content}
+                                                                                                    for msg
+                                                                                                    in
+                                                                                                    datasource_msg]).decode())
+            full_thinking_text = ''
+            full_text = ''
+            token_usage = {}
+            res = self.llm.stream(datasource_msg)
+            for chunk in res:
+                SQLBotLogUtil.info(chunk)
                 reasoning_content_chunk = ''
-            full_thinking_text += reasoning_content_chunk
+                if 'reasoning_content' in chunk.additional_kwargs:
+                    reasoning_content_chunk = chunk.additional_kwargs.get('reasoning_content', '')
+                # else:
+                #     reasoning_content_chunk = chunk.get('reasoning_content')
+                if reasoning_content_chunk is None:
+                    reasoning_content_chunk = ''
+                full_thinking_text += reasoning_content_chunk
 
-            full_text += chunk.content
-            yield {'content': chunk.content, 'reasoning_content': reasoning_content_chunk}
-            get_token_usage(chunk, token_usage)
-        datasource_msg.append(AIMessage(full_text))
+                full_text += chunk.content
+                yield {'content': chunk.content, 'reasoning_content': reasoning_content_chunk}
+                get_token_usage(chunk, token_usage)
+            datasource_msg.append(AIMessage(full_text))
 
-        json_str = extract_nested_json(full_text)
+            json_str = extract_nested_json(full_text)
 
         _error: Exception | None = None
         _datasource: int | None = None
         _engine_type: str | None = None
         try:
-            data: dict = orjson.loads(json_str)
+            data: dict = _ds_list[0] if ignore_auto_select else orjson.loads(json_str)
 
             if data.get('id') and data.get('id') != 0:
                 _datasource = data['id']
@@ -465,17 +470,18 @@ def select_datasource(self):
         except Exception as e:
             _error = e
 
-        self.record = save_full_select_datasource_message_and_answer(session=self.session, record_id=self.record.id,
-                                                                     answer=orjson.dumps({'content': full_text,
-                                                                                          'reasoning_content': full_thinking_text}).decode(),
-                                                                     datasource=_datasource,
-                                                                     engine_type=_engine_type,
-                                                                     full_message=orjson.dumps(history_msg +
-                                                                                               [{'type': msg.type,
-                                                                                                 'content': msg.content}
-                                                                                                for msg
-                                                                                                in
-                                                                                                datasource_msg]).decode())
+        if not ignore_auto_select:
+            self.record = save_full_select_datasource_message_and_answer(session=self.session, record_id=self.record.id,
+                                                                        answer=orjson.dumps({'content': full_text,
+                                                                                            'reasoning_content': full_thinking_text}).decode(),
+                                                                        datasource=_datasource,
+                                                                        engine_type=_engine_type,
+                                                                        full_message=orjson.dumps(history_msg +
+                                                                                                [{'type': msg.type,
+                                                                                                    'content': msg.content}
+                                                                                                    for msg
+                                                                                                    in
+                                                                                                    datasource_msg]).decode())
         self.init_messages()
 
         if _error:
@@ -815,7 +821,7 @@ def run_task(self, in_chat: bool = True):
             SQLBotLogUtil.info(full_sql_text)
 
             # todo row permission
-            if is_normal_user(self.current_user) or (self.current_assistant and self.current_assistant.type == 1):
+            if (not self.current_assistant and is_normal_user(self.current_user)) or (self.current_assistant and self.current_assistant.type == 1):
                 sql_json_str = extract_nested_json(full_sql_text)
                 data = orjson.loads(sql_json_str)
 

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)
+    conn = get_engine(ds, 5)
     try:
         with conn.connect() as connection:
             SQLBotLogUtil.info("success")

backend/apps/db/db.py

@@ -52,10 +52,10 @@ def get_uri_from_config(type: str, conf: DatasourceConf) -> str:
     return db_url
 
 
-def get_engine(ds: CoreDatasource) -> Engine:
+def get_engine(ds: CoreDatasource, timeout: int = 30) -> Engine:
     conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration))) if ds.type != "excel" else get_engine_config()
     if conf.timeout is None:
-        conf.timeout = 30
+        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)}",

backend/apps/system/crud/assistant.py

@@ -101,13 +101,15 @@ def get_ds_from_api(self):
         certificateList: list[any] = json.loads(self.certificate)
         header = {}
         cookies = {}
+        param = {}
         for item in certificateList:
             if item['target'] == 'header':
                 header[item['key']] = item['value']
             if item['target'] == 'cookie':
                 cookies[item['key']] = item['value']
-        
-        res = requests.get(url=endpoint, headers=header, cookies=cookies, timeout=10)
+            if item['target'] == 'param':
+                param[item['key']] = item['value']
+        res = requests.get(url=endpoint, params=param, headers=header, cookies=cookies, timeout=10)
         if res.status_code == 200:
             result_json: dict[any] = json.loads(res.text)
             if result_json.get('code') == 0:

frontend/src/i18n/en.json

@@ -472,7 +472,8 @@
     "duplicate_email": "Duplicate email",
     "interface_credentials": "Interface credentials",
     "no_credentials_yet": "No credentials yet",
-    "intelligent_customer_service": "SQLBot Intelligent Customer Service"
+    "intelligent_customer_service": "SQLBot Intelligent Customer Service",
+    "origin_format_error": "Format invalid, starts with http or https and cannot end with /"
   },
   "chat": {
     "type": "Chart Type",
@@ -518,4 +519,4 @@
     "back_community": "Restore to Community Edition",
     "confirm_tips": "Are you sure to restore to Community Edition?"
   }
-}
+}

frontend/src/i18n/zh-CN.json

@@ -487,7 +487,8 @@
     "stop_replying": "停止回答",
     "i_am_sqlbot": "你好，我是 SQLBot",
     "predict_data_etc": "我可以查询数据、生成图表、检测数据异常、预测数据等",
-    "intelligent_data_query": "赶快开启智能问数吧～"
+    "intelligent_data_query": "赶快开启智能问数吧～",
+    "origin_format_error": "格式错误，http或https开头，不能以 / 结尾"
   },
   "chat": {
     "type": "图表类型",
@@ -533,4 +534,4 @@
     "back_community": "还原至社区版",
     "confirm_tips": "确定还原至社区版？"
   }
-}
+}

frontend/src/views/system/embedded/index.vue

@@ -245,6 +245,11 @@ const rules = {
         t('datasource.please_enter') + t('common.empty') + t('embedded.cross_domain_settings'),
       trigger: 'blur',
     },
+    {
+      pattern: /^(https?:\/\/)?([\w-]+\.)*[\w-]+(:\d+)?$/,
+      message: t('embedded.origin_format_error'),
+      trigger: 'blur',
+    },
   ],
 }