mirror of
https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git
synced 2026-11-08 01:06:59 +08:00
303 lines
12 KiB
Python
303 lines
12 KiB
Python
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="[email protected]")
|
|
|
|
|
|
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))
|