From 8dbac2a1011975e3b94870eef982fa43a683e9de Mon Sep 17 00:00:00 2001 From: fit2cloud-chenyw Date: Tue, 22 Jul 2025 10:16:58 +0800 Subject: [PATCH 1/2] perf: Assistant dynamic ds --- backend/apps/chat/task/llm.py | 2 +- backend/apps/db/db.py | 22 +++++++---- backend/apps/system/crud/assistant.py | 37 ++++++++++++++----- backend/apps/system/schemas/system_schema.py | 21 ++--------- backend/common/core/deps.py | 4 +- .../src/views/chat/answer/ChartAnswer.vue | 10 ++++- 6 files changed, 55 insertions(+), 41 deletions(-) diff --git a/backend/apps/chat/task/llm.py b/backend/apps/chat/task/llm.py index 78d5d222..02746c15 100644 --- a/backend/apps/chat/task/llm.py +++ b/backend/apps/chat/task/llm.py @@ -748,7 +748,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 or llm_service.ds.type, 'type': 'datasource'}).decode() + '\n\n' - llm_service.chat_question.db_schema = llm_service.out_ds_instance.get_db_schema() if llm_service.out_ds_instance else get_table_schema(session=llm_service.session, current_user=llm_service.current_user, ds=llm_service.ds) + llm_service.chat_question.db_schema = llm_service.out_ds_instance.get_db_schema(llm_service.ds.id) if llm_service.out_ds_instance else 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() diff --git a/backend/apps/db/db.py b/backend/apps/db/db.py index d096eeb3..5d74ed90 100644 --- a/backend/apps/db/db.py +++ b/backend/apps/db/db.py @@ -11,26 +11,32 @@ from apps.datasource.utils.utils import aes_decrypt from apps.db.engine import get_engine_config from decimal import Decimal +from apps.system.crud.assistant import get_ds_engine +from apps.system.schemas.system_schema import AssistantOutDsSchema -def get_uri(ds: CoreDatasource): + +def get_uri(ds: CoreDatasource) -> str: conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration))) if ds.type != "excel" else get_engine_config() + return get_uri_from_config(ds.type, conf) + +def get_uri_from_config(type: str,conf: DatasourceConf) -> str: db_url: str - if ds.type == "mysql": + if type == "mysql": if conf.extraJdbc is not None and conf.extraJdbc != '': db_url = f"mysql+pymysql://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}/{urllib.parse.quote(conf.database)}?{conf.extraJdbc}" else: db_url = f"mysql+pymysql://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}/{urllib.parse.quote(conf.database)}" - elif ds.type == "sqlServer": + elif type == "sqlServer": if conf.extraJdbc is not None and conf.extraJdbc != '': db_url = f"mssql+pymssql://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}/{urllib.parse.quote(conf.database)}?{conf.extraJdbc}" else: db_url = f"mssql+pymssql://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}/{urllib.parse.quote(conf.database)}" - elif ds.type == "pg" or ds.type == "excel": + elif type == "pg" or type == "excel": if conf.extraJdbc is not None and conf.extraJdbc != '': db_url = f"postgresql+psycopg2://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}/{urllib.parse.quote(conf.database)}?{conf.extraJdbc}" else: db_url = f"postgresql+psycopg2://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}/{urllib.parse.quote(conf.database)}" - elif ds.type == "oracle": + elif type == "oracle": if conf.mode == "service_name": if conf.extraJdbc is not None and conf.extraJdbc != '': db_url = f"oracle+oracledb://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}?service_name={urllib.parse.quote(conf.database)}&{conf.extraJdbc}" @@ -62,8 +68,8 @@ def get_engine(ds: CoreDatasource) -> Engine: return engine -def get_session(ds: CoreDatasource): - engine = get_engine(ds) +def get_session(ds: CoreDatasource | AssistantOutDsSchema): + engine = get_engine(ds) if isinstance(ds, CoreDatasource) else get_ds_engine(ds) session_maker = sessionmaker(bind=engine) session = session_maker() return session @@ -238,7 +244,7 @@ def get_fields(ds: CoreDatasource, table_name: str = None): session.close() -def exec_sql(ds: CoreDatasource, sql: str): +def exec_sql(ds: CoreDatasource | AssistantOutDsSchema, sql: str): session = get_session(ds) result = session.execute(text(sql)) try: diff --git a/backend/apps/system/crud/assistant.py b/backend/apps/system/crud/assistant.py index 66b674a9..11934f94 100644 --- a/backend/apps/system/crud/assistant.py +++ b/backend/apps/system/crud/assistant.py @@ -4,8 +4,10 @@ import json from typing import Optional from fastapi import FastAPI import requests +from sqlalchemy import Engine, create_engine from sqlmodel import Session, select -from apps.datasource.models.datasource import CoreDatasource +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 AssistantOutDsSchema, UserInfoDTO @@ -92,7 +94,7 @@ class AssistantOutDs: self.get_ds_from_api(certificate) #@cache(namespace=CacheNamespace.EMBEDDED_INFO, cacheName=CacheName.ASSISTANT_DS, keyExpression="current_user.id") - async def get_ds_from_api(self, certificate: Optional[str] = None): + def get_ds_from_api(self, certificate: Optional[str] = None): config: dict[any] = json.loads(self.assistant.configuration) endpoint: str = config['endpoint'] certificateList: list[any] = json.loads(certificate) @@ -106,7 +108,7 @@ class AssistantOutDs: res = requests.get(url=endpoint, headers=header, cookies=cookies, timeout=10) if res.status_code == 200: - result_json: dict[any] = json.loads(res.json()) + result_json: dict[any] = json.loads(res.text) if result_json.get('code') == 0: temp_list = result_json.get('data', []) self.ds_list = [ @@ -144,20 +146,35 @@ class AssistantOutDs: def get_ds(self, ds_id: int): if self.ds_list: for ds in self.ds_list: - if ds['id'] == ds_id: + if ds.id == ds_id: return ds else: raise Exception("Datasource list is not found.") raise Exception(f"Datasource with id {ds_id} not found.") - def get_ds_engine(self, ds_id: int): - ds = self.get_ds(ds_id) - ds_type = ds.get('type') if ds else None - if not ds_type: - raise Exception(f"Datasource with id {ds_id} not found or type is not defined.") - return ds_type + class AssistantOutDsFactory: @staticmethod def get_instance(assistant: AssistantModel, certificate: Optional[str] = None) -> AssistantOutDs: return AssistantOutDs(assistant, certificate) + +def get_ds_engine(ds: AssistantOutDsSchema) -> Engine: + timeout: int = 30 + connect_args = {"connect_timeout": timeout} + conf = DatasourceConf( + host=ds.host, + port=ds.port, + username=ds.user, + password=ds.password, + database=ds.dataBase, + driver='', + extraJdbc=ds.extraParams, + dbSchema=ds.schema or '' + ) + from apps.db.db import get_uri_from_config + uri = get_uri_from_config(ds.type, conf) + if ds.type == "pg" and ds.schema: + connect_args.update({"options": f"-c search_path={ds.schema}"}) + engine = create_engine(uri, connect_args=connect_args, pool_timeout=timeout, pool_size=20, max_overflow=10) + return engine \ No newline at end of file diff --git a/backend/apps/system/schemas/system_schema.py b/backend/apps/system/schemas/system_schema.py index d94bffd8..2183db38 100644 --- a/backend/apps/system/schemas/system_schema.py +++ b/backend/apps/system/schemas/system_schema.py @@ -125,29 +125,14 @@ class AssistantOutDsBase(BaseModel): comment: Optional[str] = None description: Optional[str] = None - def __init__(self, id: Optional[int] = None, name: str = '', type: Optional[str] = None, - type_name: Optional[str] = None, comment: Optional[str] = None): - super().__init__(id=id, name=name, type=type, type_name=type_name, comment=comment) class AssistantOutDsSchema(AssistantOutDsBase): host: Optional[str] = None port: Optional[int] = None + dataBase: Optional[str] = None user: Optional[str] = None password: Optional[str] = None schema: Optional[str] = None + extraParams: Optional[str] = None tables: Optional[list[AssistantTableSchema]] = None - - def __init__(self, id: int, name: str, comment: Optional[str] = None, type: Optional[str] = None, - type_name: Optional[str] = None, host: Optional[str] = None, port: Optional[int] = None, - username: Optional[str] = None, password: Optional[str] = None, schema: Optional[str] = None): - self.id = id - self.name = name - self.comment = comment - self.type = type - self.type_name = type_name - self.host = host - self.port = port - self.username = username - self.password = password - self.schema = schema - self.description = comment \ No newline at end of file + \ No newline at end of file diff --git a/backend/common/core/deps.py b/backend/common/core/deps.py index 558ad646..be8d1efb 100644 --- a/backend/common/core/deps.py +++ b/backend/common/core/deps.py @@ -19,8 +19,8 @@ async def get_current_user(request: Request) -> UserInfoDTO: CurrentUser = Annotated[UserInfoDTO, Depends(get_current_user)] -async def get_current_assistant(request: Request) -> AssistantModel: - return request.state.assistant +async def get_current_assistant(request: Request) -> AssistantModel | None: + return request.state.assistant if hasattr(request.state, "assistant") else None CurrentAssistant = Annotated[AssistantModel, Depends(get_current_assistant)] diff --git a/frontend/src/views/chat/answer/ChartAnswer.vue b/frontend/src/views/chat/answer/ChartAnswer.vue index c177eb27..3d9f3743 100644 --- a/frontend/src/views/chat/answer/ChartAnswer.vue +++ b/frontend/src/views/chat/answer/ChartAnswer.vue @@ -1,7 +1,9 @@