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/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/db/db.py b/backend/apps/db/db.py index 83be9c75..607d7cb0 100644 --- a/backend/apps/db/db.py +++ b/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 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 9b1579fc..090c4f1b 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,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 - \ No newline at end of file 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/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 `
- +
diff --git a/frontend/src/views/chat/chat-block/ChartBlock.vue b/frontend/src/views/chat/chat-block/ChartBlock.vue index 6c428806..32de5d61 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' 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() { @@ -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 diff --git a/frontend/src/views/embedded/index.vue b/frontend/src/views/embedded/index.vue index c728f20e..e71ce1ed 100644 --- a/frontend/src/views/embedded/index.vue +++ b/frontend/src/views/embedded/index.vue @@ -191,3 +191,10 @@ onBeforeUnmount(() => { } } + +