import json import re import urllib from typing import Optional import requests from fastapi import FastAPI from sqlmodel import Session, select from starlette.middleware.cors import CORSMiddleware from apps.datasource.models.datasource import CoreDatasource from apps.datasource.utils.utils import aes_encrypt 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.config import settings from common.core.db import engine from common.core.sqlbot_cache import cache from common.utils.aes_crypto import simple_aes_decrypt from common.utils.utils import SQLBotLogUtil, get_domain_list, string_to_numeric_hash from common.core.deps import Trans from common.core.response_middleware import ResponseMiddleware def _update_cors_middleware_instance(app: FastAPI, updated_origins: list[str]): """遍历 middleware 栈，找到 CORSMiddleware 实例并更新其 allow_origins。 仅修改 middleware.kwargs 不会影响已构建的中间件实例， 需要直接更新实例的 allow_origins 属性。 """ stack = getattr(app, 'middleware_stack', None) while stack is not None and hasattr(stack, 'app'): if isinstance(stack, CORSMiddleware): stack.allow_origins = updated_origins return stack = stack.app @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") def get_assistant_ds(session: Session, llm_service) -> list[dict]: assistant: AssistantHeader = llm_service.current_assistant type = assistant.type if type == 0 or type == 2: configuration = assistant.configuration 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) if not assistant.online: public_list: list[int] = config.get('public_list') or None if public_list: stmt = stmt.where(CoreDatasource.id.in_(public_list)) else: return [] """ 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, "name": ds.name, "description": ds.description } for ds in db_ds_list ] # filter private ds if offline return result_list out_ds_instance: AssistantOutDs = AssistantOutDsFactory.get_instance(assistant) llm_service.out_ds_instance = out_ds_instance dslist = out_ds_instance.get_simple_ds_list() # format? return dslist def init_dynamic_cors(app: FastAPI): try: with Session(engine) as session: list_result = session.exec(select(AssistantModel).order_by(AssistantModel.create_time)).all() seen = set() unique_domains = [] for item in list_result: if item.domain: for domain in get_domain_list(item.domain): domain = domain.strip() if domain and domain not in seen: seen.add(domain) unique_domains.append(domain) cors_middleware = None response_middleware = None for middleware in app.user_middleware: if not cors_middleware and middleware.cls == CORSMiddleware: cors_middleware = middleware if not response_middleware and middleware.cls == ResponseMiddleware: response_middleware = middleware if cors_middleware and response_middleware: break updated_origins = list(set(settings.all_cors_origins + unique_domains)) if cors_middleware: cors_middleware.kwargs['allow_origins'] = updated_origins _update_cors_middleware_instance(app, updated_origins) if response_middleware: for instance in ResponseMiddleware.instances: instance.update_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 request_origin: Optional[str] = None def __init__(self, assistant: AssistantHeader): self.assistant = assistant self.ds_list = None self.certificate = assistant.certificate self.request_origin = assistant.request_origin self.get_ds_from_api() # @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'] endpoint = self.get_complete_endpoint(endpoint=endpoint) if not endpoint: raise Exception( f"Failed to get datasource list from {config['endpoint']}, error: [Assistant domain or endpoint miss]") 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'] if item['target'] == 'param': param[item['key']] = item['value'] timeout = int(config.get('timeout')) if config.get('timeout') else 10 res = requests.get(url=endpoint, params=param, headers=header, cookies=cookies, timeout=timeout) if res.status_code == 200: result_json: dict[any] = json.loads(res.text) if result_json.get('code') == 0 or result_json.get('code') == 200: temp_list = result_json.get('data', []) temp_ds_list = [ self.convert2schema(item, config) for item in temp_list ] self.ds_list = temp_ds_list return self.ds_list else: raise Exception(f"Failed to get datasource list from {endpoint}, error: {result_json.get('message')}") else: SQLBotLogUtil.error(f"Failed to get datasource list from {endpoint}, response: {res}") raise Exception(f"Failed to get datasource list from {endpoint}, response: {res}") def get_first_element(self, text: str): parts = re.split(r'[,;]', text.strip()) first_domain = parts[0].strip() return first_domain def get_complete_endpoint(self, endpoint: str) -> str | None: if endpoint.startswith("http://") or endpoint.startswith("https://"): return endpoint domain_text = self.assistant.domain if not domain_text: return None if ',' in domain_text or ';' in domain_text: return ( self.request_origin.strip('/') if self.request_origin else self.get_first_element(domain_text).strip( '/')) + endpoint else: return f"{domain_text}{endpoint}" 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, question: str = '', embedding: bool = True, table_list: list[str] = None) -> tuple[str, list]: ds = self.get_ds(ds_id) schema_str = "" 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" tables = [] table_name_list = [] i = 0 for table in ds.tables: # 如果传入了 table_list，则只处理在列表中的表 if table_list is not None and table.name not in table_list: continue i += 1 schema_table = '' schema_table += f"# Table: {db_name}.{table.name}" if ds.type != "mysql" and ds.type != "es" else f"# Table: {table.name}" table_comment = table.comment if table_comment == '': schema_table += '\n[\n' else: schema_table += f", {table_comment}\n[\n" field_list = [] for field in table.fields: field_comment = field.comment if field_comment == '': field_list.append(f"({field.name}:{field.type})") else: field_list.append(f"({field.name}:{field.type}, {field_comment})") schema_table += ",\n".join(field_list) schema_table += '\n]\n' t_obj = {"id": i, "schema_table": schema_table} tables.append(t_obj) table_name_list.append(table.name) # do table embedding # if embedding and tables and settings.TABLE_EMBEDDING_ENABLED: # tables = get_table_embedding(tables, question) if tables: for s in tables: schema_str += s.get('schema_table') return schema_str, table_name_list def get_ds(self, ds_id: int, trans: Trans = None): if self.ds_list: for ds in self.ds_list: if ds.id == ds_id: return ds else: raise Exception("Datasource list is not found.") raise Exception(f"Datasource id {ds_id} is not found." if trans is None else trans( 'i18n_data_training.datasource_id_not_found', key=ds_id)) def convert2schema(self, ds_dict: dict, config: dict[any]) -> AssistantOutDsSchema: id_marker: str = '' attr_list = ['name', 'type', 'host', 'port', 'user', 'dataBase', 'schema', 'mode', 'lowVersion'] if config.get('encrypt', False): key = config.get('aes_key', None) iv = config.get('aes_iv', None) aes_attrs = ['host', 'user', 'password', 'dataBase', 'db_schema', 'schema', 'mode', 'lowVersion'] for attr in aes_attrs: if attr in ds_dict and ds_dict[attr]: try: ds_dict[attr] = simple_aes_decrypt(ds_dict[attr], key, iv) except Exception as e: raise Exception( f"Failed to encrypt {attr} for datasource {ds_dict.get('name')}, error: {str(e)}") id = ds_dict.get('id', None) if not id: for attr in attr_list: if attr in ds_dict: id_marker += str(ds_dict.get(attr, '')) + '--sqlbot--' id = string_to_numeric_hash(id_marker) 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_out_ds_conf(ds: AssistantOutDsSchema, timeout: int = 30) -> str: conf = { "host": ds.host or '', "port": ds.port or 0, "username": ds.user or '', "password": ds.password or '', "database": ds.dataBase or '', "driver": '', "extraJdbc": ds.extraParams or '', "dbSchema": ds.db_schema or '', "timeout": timeout or 30, "mode": ds.mode or '', "lowVersion": ds.lowVersion or False, } conf["extraJdbc"] = '' return aes_encrypt(json.dumps(conf))