import asyncio import json import os from datetime import datetime, timedelta, timezone from zoneinfo import ZoneInfo from typing import List, Optional from fastapi import APIRouter, Form, HTTPException, Path, Query, Request, Response, UploadFile from fastapi.responses import StreamingResponse from sqlbot_xpack.file_utils import SQLBotFileUtils from sqlmodel import select from apps.datasource.models.datasource import CoreDatasource from apps.db.constant import DB from apps.swagger.i18n import PLACEHOLDER_PREFIX from apps.system.crud.assistant import AssistantOutDs, AssistantOutDsFactory, get_assistant_info from apps.system.crud.assistant_manage import dynamic_upgrade_cors, get_ws_assistant, save from apps.system.models.system_model import AssistantModel from apps.system.schemas.auth import CacheName, CacheNamespace from apps.system.schemas.permission import SqlbotPermission, require_permissions from apps.system.schemas.system_schema import AssistantBase, AssistantDTO, AssistantUiSchema, AssistantValidator from common.core.config import settings from common.core.deps import CurrentAssistant, SessionDep, Trans, CurrentUser from common.core.security import create_access_token from common.core.sqlbot_cache import clear_cache from common.utils.utils import get_origin_from_referer, origin_match_domain router = APIRouter(tags=["system_assistant"], prefix="/system/assistant") from common.audit.models.log_model import OperationType, OperationModules from common.audit.schemas.logger_decorator import LogConfig, system_log from sqlbot_xpack.core import decrypt_embedded_sign @router.get("/info/{id}", include_in_schema=False) async def info(request: Request, response: Response, session: SessionDep, trans: Trans, id: int, virtual: Optional[int] = Query(None)): if not id: raise Exception('miss assistant id') db_model = await get_assistant_info(session=session, assistant_id=id) if not db_model: raise RuntimeError(f"assistant application not exist") db_model = AssistantModel.model_validate(db_model) # 校验 SQLBOT-EMBEDDED-SIGN 请求头 sign_header = request.headers.get("SQLBOT-EMBEDDED-SIGN") if not sign_header: raise RuntimeError(trans('i18n_embedded.missing_sign')) try: sign_data = await decrypt_embedded_sign(sign_header) except Exception: raise RuntimeError(trans('i18n_embedded.invalid_sign')) # 校验 assistant_id 与 id 参数一致 if str(sign_data.get("assistant_id")) != str(id): raise RuntimeError(trans('i18n_embedded.assistant_id_mismatch')) # 校验 target（来源域名）是否合法 target = sign_data.get("target", "") request_origin = request.headers.get("origin") or get_origin_from_referer(request) if not request_origin: raise RuntimeError(trans('i18n_embedded.missing_origin')) request_origin = request_origin.rstrip('/') if not target or target == "null": target = request_origin elif target != request_origin: raise RuntimeError(trans('i18n_embedded.target_origin_mismatch')) if not origin_match_domain(target, db_model.domain): raise RuntimeError(trans('i18n_embedded.invalid_origin', origin=target or '')) # 校验 sign_time 是否在 10 秒内 sign_time_str = sign_data.get("sign_time", "") try: sign_time = datetime.fromisoformat(sign_time_str) except (ValueError, TypeError): raise RuntimeError(trans('i18n_embedded.sign_time_invalid')) now_utc = datetime.now(timezone.utc) sign_time_utc = sign_time.astimezone(timezone.utc) if abs((now_utc - sign_time_utc).total_seconds()) > 10: raise RuntimeError(trans('i18n_embedded.sign_time_expired')) # 校验是否为真实浏览器请求（非自动化工具） if sign_data.get("webdriver", False): raise RuntimeError(trans('i18n_embedded.automation_tool_detected')) # 校验 User-Agent 一致性（签名中的 navigator.userAgent 与请求头一致） sign_user_agent = sign_data.get("user_agent", "") request_user_agent = request.headers.get("User-Agent", "") if sign_user_agent != request_user_agent: raise RuntimeError(trans('i18n_embedded.user_agent_mismatch')) # 校验 timezone 与 sign_time 偏移一致性（防时区伪造） tz_name = sign_data.get("timezone", "") try: tz = ZoneInfo(tz_name) except Exception: raise RuntimeError(trans('i18n_embedded.timezone_mismatch')) sign_time_naive = sign_time.replace(tzinfo=None) if tz.utcoffset(sign_time_naive) != sign_time.utcoffset(): raise RuntimeError(trans('i18n_embedded.timezone_mismatch')) origin = target.rstrip('/') response.headers["Access-Control-Allow-Origin"] = origin assistant_oid = 1 if (db_model.type == 0): configuration = db_model.configuration config_obj = json.loads(configuration) if configuration else {} assistant_oid = config_obj.get('oid', 1) access_token_expires = timedelta(minutes=settings.ACCESS_TOKEN_EXPIRE_MINUTES) assistantDict = { "id": virtual, "account": 'sqlbot-inner-assistant', "oid": assistant_oid, "assistant_id": id } access_token = create_access_token( assistantDict, expires_delta=access_token_expires ) result = db_model.model_dump() result["token"] = access_token return result @router.get("/app/{appId}", include_in_schema=False) async def getApp(request: Request, response: Response, session: SessionDep, trans: Trans, appId: str) -> AssistantModel: if not appId: raise Exception('miss assistant appId') db_model = session.exec(select(AssistantModel).where(AssistantModel.app_id == appId)).first() if not db_model: raise RuntimeError(f"assistant application not exist") db_model = AssistantModel.model_validate(db_model) origin = request.headers.get("origin") or get_origin_from_referer(request) if not origin: raise RuntimeError(trans('i18n_embedded.invalid_origin', origin=origin or '')) origin = origin.rstrip('/') if not origin_match_domain(origin, db_model.domain): raise RuntimeError(trans('i18n_embedded.invalid_origin', origin=origin or '')) response.headers["Access-Control-Allow-Origin"] = origin return db_model """ @router.get("/validator", response_model=AssistantValidator, include_in_schema=False) async def validator(session: SessionDep, id: int, virtual: Optional[int] = Query(None)): if not id: raise Exception('miss assistant id') db_model = await get_assistant_info(session=session, assistant_id=id) if not db_model: return AssistantValidator() db_model = AssistantModel.model_validate(db_model) assistant_oid = 1 if (db_model.type == 0): configuration = db_model.configuration config_obj = json.loads(configuration) if configuration else {} assistant_oid = config_obj.get('oid', 1) access_token_expires = timedelta(minutes=settings.ACCESS_TOKEN_EXPIRE_MINUTES) assistantDict = { "id": virtual, "account": 'sqlbot-inner-assistant', "oid": assistant_oid, "assistant_id": id } access_token = create_access_token( assistantDict, expires_delta=access_token_expires ) return AssistantValidator(True, True, True, access_token) """ @router.get('/picture/{file_id}', summary=f"{PLACEHOLDER_PREFIX}assistant_picture_api", description=f"{PLACEHOLDER_PREFIX}assistant_picture_api") async def picture(file_id: str = Path(description="file_id")): file_path = SQLBotFileUtils.get_file_path(file_id=file_id) if not os.path.exists(file_path): raise HTTPException(status_code=404, detail="File not found") if file_id.lower().endswith(".svg"): media_type = "image/svg+xml" else: media_type = "image/jpeg" def iterfile(): with open(file_path, mode="rb") as f: yield from f return StreamingResponse(iterfile(), media_type=media_type) @router.patch('/ui', summary=f"{PLACEHOLDER_PREFIX}assistant_ui_api", description=f"{PLACEHOLDER_PREFIX}assistant_ui_api") @require_permissions(permission=SqlbotPermission(role=['ws_admin'])) @system_log(LogConfig(operation_type=OperationType.UPDATE, module=OperationModules.APPLICATION, result_id_expr="id")) async def ui(session: SessionDep, current_user: CurrentUser, trans: Trans, data: str = Form(), files: List[UploadFile] = []): json_data = json.loads(data) uiSchema = AssistantUiSchema(**json_data) id = uiSchema.id db_model = get_ws_assistant(session, id, current_user, trans) configuration = db_model.configuration config_obj = json.loads(configuration) if configuration else {} ui_schema_dict = uiSchema.model_dump(exclude_none=True, exclude_unset=True) if files: for file in files: origin_file_name = file.filename file_name, flag_name = SQLBotFileUtils.split_filename_and_flag(origin_file_name) file.filename = file_name if flag_name == 'logo' or flag_name == 'float_icon': try: SQLBotFileUtils.check_file(file=file, file_types=[".jpg", ".jpeg", ".png"], limit_file_size=(10 * 1024 * 1024)) except ValueError as e: error_msg = str(e) if '文件大小超过限制' in error_msg: raise ValueError(f"文件大小超过限制（最大 10 M）") else: raise e if config_obj.get(flag_name): SQLBotFileUtils.delete_file(config_obj.get(flag_name)) file_id = await SQLBotFileUtils.upload(file) ui_schema_dict[flag_name] = file_id else: raise ValueError(f"Unsupported file flag: {flag_name}") for flag_name in ['logo', 'float_icon']: file_val = config_obj.get(flag_name) if file_val and not ui_schema_dict.get(flag_name): config_obj[flag_name] = None SQLBotFileUtils.delete_file(file_val) for attr, value in ui_schema_dict.items(): if attr != 'id' and not attr.startswith("__"): config_obj[attr] = value db_model.configuration = json.dumps(config_obj, ensure_ascii=False) session.add(db_model) session.commit() await clear_ui_cache(db_model.id) return db_model @clear_cache(namespace=CacheNamespace.EMBEDDED_INFO, cacheName=CacheName.ASSISTANT_INFO, keyExpression="id") async def clear_ui_cache(id: int): pass @router.get("/validate/{id}", include_in_schema=False) async def validate(request: Request, response: Response, session: SessionDep, trans: Trans, id: int): if not id: raise Exception('miss assistant id') db_model = session.get(AssistantModel, id) if not db_model: raise RuntimeError(f"assistant application not exist") origin = request.headers.get("origin") or get_origin_from_referer(request) if not origin: raise RuntimeError(trans('i18n_embedded.invalid_origin', origin=origin or '')) origin = origin.rstrip('/') if not origin_match_domain(origin, db_model.domain): raise RuntimeError(trans('i18n_embedded.invalid_origin', origin=origin or '')) response.headers["Access-Control-Allow-Origin"] = origin return {"valid": True, "origin": origin} @router.get("/ds", include_in_schema=False, response_model=list[dict]) async def ds(session: SessionDep, current_assistant: CurrentAssistant): if current_assistant.type == 0: online = current_assistant.online configuration = current_assistant.configuration config: dict[any] = json.loads(configuration) oid: int = int(config['oid']) stmt = select(CoreDatasource.id, CoreDatasource.name, CoreDatasource.description, CoreDatasource.type, CoreDatasource.type_name, CoreDatasource.num).where( CoreDatasource.oid == oid) if not online: public_list: list[int] = config.get('public_list') or None if public_list: stmt = stmt.where(CoreDatasource.id.in_(public_list)) else: return [] db_ds_list = session.exec(stmt) return [ { "id": ds.id, "name": ds.name, "description": ds.description, "type": ds.type, "type_name": ds.type_name, "num": ds.num, } for ds in db_ds_list] if current_assistant.type == 1: # 构造 AssistantOutDs 会同步请求宿主 API（requests.get，默认 10s 超时）， # 必须在工作线程执行，避免阻塞事件循环导致宿主回调死锁（Issue #1288） def inner(): return AssistantOutDsFactory.get_instance(current_assistant) out_ds_instance: AssistantOutDs = await asyncio.to_thread(inner) return [ { "id": str(ds.id), "name": ds.name, "description": ds.description or ds.comment, "type": ds.type, "type_name": get_db_type(ds.type), "num": len(ds.tables) if ds.tables else 0, } for ds in out_ds_instance.ds_list if get_db_type(ds.type) ] return None def get_db_type(type): try: db = DB.get_db(type) return db.db_name except Exception: return None @router.get("", response_model=list[AssistantModel], summary=f"{PLACEHOLDER_PREFIX}assistant_grid_api", description=f"{PLACEHOLDER_PREFIX}assistant_grid_api") @require_permissions(permission=SqlbotPermission(role=['ws_admin'])) async def query(session: SessionDep, current_user: CurrentUser): list_result = session.exec(select(AssistantModel).where(AssistantModel.oid == current_user.oid, AssistantModel.type != 4).order_by(AssistantModel.name, AssistantModel.create_time)).all() for model in list_result: model.enable_custom_model = model.enable_custom_model or False return list_result @router.get("/advanced_application", response_model=list[AssistantModel], include_in_schema=False) @require_permissions(permission=SqlbotPermission(role=['ws_admin'])) async def query_advanced_application(session: SessionDep, current_user: CurrentUser): list_result = session.exec(select(AssistantModel).where(AssistantModel.type == 1, AssistantModel.oid == current_user.oid).order_by(AssistantModel.name, AssistantModel.create_time)).all() return list_result @router.post("", summary=f"{PLACEHOLDER_PREFIX}assistant_create_api", description=f"{PLACEHOLDER_PREFIX}assistant_create_api") @require_permissions(permission=SqlbotPermission(role=['ws_admin'])) @system_log(LogConfig(operation_type=OperationType.CREATE, module=OperationModules.APPLICATION, result_id_expr="id")) async def add(request: Request, session: SessionDep, current_user: CurrentUser, creator: AssistantBase): oid = current_user.oid if creator.type != 4 else 1 return await save(request, session, creator, oid) @router.put("", summary=f"{PLACEHOLDER_PREFIX}assistant_update_api", description=f"{PLACEHOLDER_PREFIX}assistant_update_api") @require_permissions(permission=SqlbotPermission(role=['ws_admin'])) @clear_cache(namespace=CacheNamespace.EMBEDDED_INFO, cacheName=CacheName.ASSISTANT_INFO, keyExpression="editor.id") @system_log(LogConfig(operation_type=OperationType.UPDATE, module=OperationModules.APPLICATION, resource_id_expr="editor.id")) async def update(request: Request, session: SessionDep, current_user: CurrentUser, trans: Trans, editor: AssistantDTO): id = editor.id db_model = get_ws_assistant(session, id, current_user, trans) update_data = AssistantModel.model_validate(editor) if not current_user.isAdmin: # 忽略请求体 oid（含默认值 1）：非 admin 不允许跨工作空间移动/接管 update_data.oid = db_model.oid # create_time 不在请求体中，model_validate 会填默认值 0，更新时保留原值 update_data.create_time = db_model.create_time db_model.sqlmodel_update(update_data) session.add(db_model) session.commit() dynamic_upgrade_cors(request=request, session=session) """ @router.get("/{id}", response_model=AssistantModel, summary=f"{PLACEHOLDER_PREFIX}assistant_query_api", description=f"{PLACEHOLDER_PREFIX}assistant_query_api") async def get_one(session: SessionDep, id: int = Path(description="ID")): db_model = await get_assistant_info(session=session, assistant_id=id) if not db_model: raise ValueError(f"AssistantModel with id {id} not found") db_model = AssistantModel.model_validate(db_model) return db_model """ @router.delete("/{id}", summary=f"{PLACEHOLDER_PREFIX}assistant_del_api", description=f"{PLACEHOLDER_PREFIX}assistant_del_api") @require_permissions(permission=SqlbotPermission(role=['ws_admin'])) @clear_cache(namespace=CacheNamespace.EMBEDDED_INFO, cacheName=CacheName.ASSISTANT_INFO, keyExpression="id") @system_log(LogConfig(operation_type=OperationType.DELETE, module=OperationModules.APPLICATION, resource_id_expr="id")) async def delete(request: Request, session: SessionDep, current_user: CurrentUser, trans: Trans, id: int = Path(description="ID")): db_model = get_ws_assistant(session, id, current_user, trans) session.delete(db_model) session.commit() dynamic_upgrade_cors(request=request, session=session)