From 8141700746aa365cea14c56320b980efededc9c4 Mon Sep 17 00:00:00 2001 From: junjun Date: Tue, 29 Jul 2025 11:47:12 +0800 Subject: [PATCH 1/8] feat: assistant connect --- backend/apps/datasource/crud/datasource.py | 2 +- backend/apps/system/crud/assistant.py | 86 ++++++++++++++-------- 2 files changed, 57 insertions(+), 31 deletions(-) diff --git a/backend/apps/datasource/crud/datasource.py b/backend/apps/datasource/crud/datasource.py index a73a22bf..ed52dfa4 100644 --- a/backend/apps/datasource/crud/datasource.py +++ b/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") diff --git a/backend/apps/system/crud/assistant.py b/backend/apps/system/crud/assistant.py index 9b1579fc..72bdfa57 100644 --- a/backend/apps/system/crud/assistant.py +++ b/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="sqlbot-inner-assistant@sqlbot.com") + return UserInfoDTO(id=id, account="sqlbot-inner-assistant", oid=1, name="sqlbot-inner-assistant", + email="sqlbot-inner-assistant@sqlbot.com") + 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 @@ class AssistantOutDs: 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 @@ class AssistantOutDs: 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 @@ class AssistantOutDs: 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 @@ class AssistantOutDs: 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,26 @@ 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, pool_size=20, max_overflow=10) + elif ds.type == 'sqlServer': + engine = create_engine(uri, pool_timeout=timeout, + pool_size=20, + max_overflow=10) + elif ds.type == 'oracle': + engine = create_engine(uri, + pool_timeout=timeout, + pool_size=20, + max_overflow=10) + else: + engine = create_engine(uri, connect_args={"connect_timeout": timeout}, pool_timeout=timeout, + pool_size=20, + max_overflow=10) return engine - \ No newline at end of file From f6253cdb0ce6c2d129d0db532895965507ab84f1 Mon Sep 17 00:00:00 2001 From: junjun Date: Tue, 29 Jul 2025 12:00:03 +0800 Subject: [PATCH 2/8] feat: connect --- backend/apps/db/db.py | 14 ++++---------- backend/apps/db/engine.py | 4 +--- backend/apps/system/crud/assistant.py | 14 ++++---------- 3 files changed, 9 insertions(+), 23 deletions(-) diff --git a/backend/apps/db/db.py b/backend/apps/db/db.py index 83be9c75..ede3d277 100644 --- a/backend/apps/db/db.py +++ b/backend/apps/db/db.py @@ -60,20 +60,14 @@ def get_engine(ds: CoreDatasource, timeout: int = 30) -> Engine: 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 diff --git a/backend/apps/db/engine.py b/backend/apps/db/engine.py index e86756f5..f4c95274 100644 --- a/backend/apps/db/engine.py +++ b/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 diff --git a/backend/apps/system/crud/assistant.py b/backend/apps/system/crud/assistant.py index 72bdfa57..090c4f1b 100644 --- a/backend/apps/system/crud/assistant.py +++ b/backend/apps/system/crud/assistant.py @@ -213,18 +213,12 @@ def get_ds_engine(ds: AssistantOutDsSchema) -> Engine: engine = create_engine(uri, connect_args={"options": f"-c search_path={urllib.parse.quote(ds.db_schema)}", "connect_timeout": timeout}, - pool_timeout=timeout, pool_size=20, max_overflow=10) + pool_timeout=timeout) elif ds.type == 'sqlServer': - engine = create_engine(uri, pool_timeout=timeout, - pool_size=20, - max_overflow=10) + engine = create_engine(uri, pool_timeout=timeout) elif ds.type == 'oracle': engine = create_engine(uri, - pool_timeout=timeout, - pool_size=20, - max_overflow=10) + pool_timeout=timeout) else: - engine = create_engine(uri, connect_args={"connect_timeout": timeout}, pool_timeout=timeout, - pool_size=20, - max_overflow=10) + engine = create_engine(uri, connect_args={"connect_timeout": timeout}, pool_timeout=timeout) return engine From 91cd477e3563f69ee8f6071f69887be7a5804f40 Mon Sep 17 00:00:00 2001 From: junjun Date: Tue, 29 Jul 2025 12:14:39 +0800 Subject: [PATCH 3/8] feat: connect --- backend/apps/db/db.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/backend/apps/db/db.py b/backend/apps/db/db.py index ede3d277..607d7cb0 100644 --- a/backend/apps/db/db.py +++ b/backend/apps/db/db.py @@ -52,10 +52,12 @@ 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)}", From e5384668386018a4bb41b59fd6eb4c1561560adf Mon Sep 17 00:00:00 2001 From: dataeaseShu Date: Tue, 29 Jul 2025 14:00:06 +0800 Subject: [PATCH 4/8] fix: bug fix --- frontend/src/views/chat/chat-block/ChartBlock.vue | 5 ++++- frontend/src/views/embedded/index.vue | 7 +++++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/frontend/src/views/chat/chat-block/ChartBlock.vue b/frontend/src/views/chat/chat-block/ChartBlock.vue index 52e59969..36c016f3 100644 --- a/frontend/src/views/chat/chat-block/ChartBlock.vue +++ b/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' const props = withDefaults( defineProps<{ @@ -50,6 +51,8 @@ const dataObject = computed<{ } return {} }) +const assistantStore = useAssistantStore() +const isAssistant = computed(() => assistantStore.getAssistant) const data = computed(() => { if (props.isPredict) { @@ -321,7 +324,7 @@ function copy() { { } } + + From d4832da6af045c03e6a20165ec3c4e1078cb87b1 Mon Sep 17 00:00:00 2001 From: ulleo Date: Tue, 29 Jul 2025 14:08:45 +0800 Subject: [PATCH 5/8] fix: show engin type icon after chat created --- backend/apps/chat/curd/chat.py | 1 + frontend/src/views/chat/index.vue | 1 + 2 files changed, 2 insertions(+) diff --git a/backend/apps/chat/curd/chat.py b/backend/apps/chat/curd/chat.py index 2d573070..3f10f678 100644 --- a/backend/apps/chat/curd/chat.py +++ b/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 diff --git a/frontend/src/views/chat/index.vue b/frontend/src/views/chat/index.vue index 5a964f20..fde1e808 100644 --- a/frontend/src/views/chat/index.vue +++ b/frontend/src/views/chat/index.vue @@ -485,6 +485,7 @@ function onChatCreatedQuick(chat: ChatInfo) { chatList.value.unshift(chat) currentChatId.value = chat.id currentChat.value = chat + console.log(chat) onChatCreated(chat) } From 1e5628b3b4ec0856ab11878ea9d7032810486fc0 Mon Sep 17 00:00:00 2001 From: dataeaseShu Date: Tue, 29 Jul 2025 14:34:43 +0800 Subject: [PATCH 6/8] fix: bug fix --- frontend/src/views/chat/index.vue | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/frontend/src/views/chat/index.vue b/frontend/src/views/chat/index.vue index 5a964f20..8993c3b8 100644 --- a/frontend/src/views/chat/index.vue +++ b/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 From 32b1748357cc5f94343361135e177058f9d7067f Mon Sep 17 00:00:00 2001 From: dataeaseShu Date: Tue, 29 Jul 2025 14:46:59 +0800 Subject: [PATCH 7/8] fix: bug fix --- frontend/public/assistant.js | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/frontend/public/assistant.js b/frontend/public/assistant.js index 3485b3d4..0e3d2bff 100644 --- a/frontend/public/assistant.js +++ b/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 `
- +
From 27d745460e6d4e77dc55874544a6d5da90047741 Mon Sep 17 00:00:00 2001 From: ulleo Date: Tue, 29 Jul 2025 15:02:29 +0800 Subject: [PATCH 8/8] feat: modify template --- backend/template.yaml | 1 + frontend/src/views/chat/index.vue | 1 - 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/backend/template.yaml b/backend/template.yaml index 13f902df..e7bad6bc 100644 --- a/backend/template.yaml +++ b/backend/template.yaml @@ -7,6 +7,7 @@ template: 根据表结构和问题生成符合{engine}数据库引擎规范的sql语句，以及sql中所用到的表名（不要包含schema和database，用数组返回）。 你必须遵守以下规则: - 生成的SQL必须符合{engine}的规范。 + - 提问中如果有涉及数据源名称或数据源描述的内容，则忽略数据源的信息，直接根据剩余内容生成SQL - 根据表结构生成SQL语句，需给每个表名生成一个别名（不要加AS）。 - SQL查询中不能使用星号(*)，必须明确指定字段名. - SQL查询的字段名不要自动翻译，别名必须为英文。 diff --git a/frontend/src/views/chat/index.vue b/frontend/src/views/chat/index.vue index e35d811f..8993c3b8 100644 --- a/frontend/src/views/chat/index.vue +++ b/frontend/src/views/chat/index.vue @@ -485,7 +485,6 @@ function onChatCreatedQuick(chat: ChatInfo) { chatList.value.unshift(chat) currentChatId.value = chat.id currentChat.value = chat - console.log(chat) onChatCreated(chat) }