Underground-Digital/Workflow-Engine

1import json2from copy import deepcopy3from datetime import datetime, timezone4from typing import Any, Optional, Union5 6import httpx7import validators8 9from constants import HIDDEN_VALUE10from core.helper import ssrf_proxy11from extensions.ext_database import db12from models.dataset import (13    Dataset,14    ExternalKnowledgeApis,15    ExternalKnowledgeBindings,16)17from services.entities.external_knowledge_entities.external_knowledge_entities import (18    Authorization,19    ExternalKnowledgeApiSetting,20)21from services.errors.dataset import DatasetNameDuplicateError22 23 24class ExternalDatasetService:25    @staticmethod26    def get_external_knowledge_apis(page, per_page, tenant_id, search=None) -> tuple[list[ExternalKnowledgeApis], int]:27        query = ExternalKnowledgeApis.query.filter(ExternalKnowledgeApis.tenant_id == tenant_id).order_by(28            ExternalKnowledgeApis.created_at.desc()29        )30        if search:31            query = query.filter(ExternalKnowledgeApis.name.ilike(f"%{search}%"))32 33        external_knowledge_apis = query.paginate(page=page, per_page=per_page, max_per_page=100, error_out=False)34 35        return external_knowledge_apis.items, external_knowledge_apis.total36 37    @classmethod38    def validate_api_list(cls, api_settings: dict):39        if not api_settings:40            raise ValueError("api list is empty")41        if "endpoint" not in api_settings and not api_settings["endpoint"]:42            raise ValueError("endpoint is required")43        if "api_key" not in api_settings and not api_settings["api_key"]:44            raise ValueError("api_key is required")45 46    @staticmethod47    def create_external_knowledge_api(tenant_id: str, user_id: str, args: dict) -> ExternalKnowledgeApis:48        ExternalDatasetService.check_endpoint_and_api_key(args.get("settings"))49        external_knowledge_api = ExternalKnowledgeApis(50            tenant_id=tenant_id,51            created_by=user_id,52            updated_by=user_id,53            name=args.get("name"),54            description=args.get("description", ""),55            settings=json.dumps(args.get("settings"), ensure_ascii=False),56        )57 58        db.session.add(external_knowledge_api)59        db.session.commit()60        return external_knowledge_api61 62    @staticmethod63    def check_endpoint_and_api_key(settings: dict):64        if "endpoint" not in settings or not settings["endpoint"]:65            raise ValueError("endpoint is required")66        if "api_key" not in settings or not settings["api_key"]:67            raise ValueError("api_key is required")68 69        endpoint = f"{settings['endpoint']}/retrieval"70        api_key = settings["api_key"]71        if not validators.url(endpoint, simple_host=True):72            raise ValueError(f"invalid endpoint: {endpoint}")73        try:74            response = httpx.post(endpoint, headers={"Authorization": f"Bearer {api_key}"})75        except Exception as e:76            raise ValueError(f"failed to connect to the endpoint: {endpoint}")77        if response.status_code == 502:78            raise ValueError(f"Bad Gateway: failed to connect to the endpoint: {endpoint}")79        if response.status_code == 404:80            raise ValueError(f"Not Found: failed to connect to the endpoint: {endpoint}")81        if response.status_code == 403:82            raise ValueError(f"Forbidden: Authorization failed with api_key: {api_key}")83 84    @staticmethod85    def get_external_knowledge_api(external_knowledge_api_id: str) -> ExternalKnowledgeApis:86        return ExternalKnowledgeApis.query.filter_by(id=external_knowledge_api_id).first()87 88    @staticmethod89    def update_external_knowledge_api(tenant_id, user_id, external_knowledge_api_id, args) -> ExternalKnowledgeApis:90        external_knowledge_api = ExternalKnowledgeApis.query.filter_by(91            id=external_knowledge_api_id, tenant_id=tenant_id92        ).first()93        if external_knowledge_api is None:94            raise ValueError("api template not found")95        if args.get("settings") and args.get("settings").get("api_key") == HIDDEN_VALUE:96            args.get("settings")["api_key"] = external_knowledge_api.settings_dict.get("api_key")97 98        external_knowledge_api.name = args.get("name")99        external_knowledge_api.description = args.get("description", "")100        external_knowledge_api.settings = json.dumps(args.get("settings"), ensure_ascii=False)101        external_knowledge_api.updated_by = user_id102        external_knowledge_api.updated_at = datetime.now(timezone.utc).replace(tzinfo=None)103        db.session.commit()104 105        return external_knowledge_api106 107    @staticmethod108    def delete_external_knowledge_api(tenant_id: str, external_knowledge_api_id: str):109        external_knowledge_api = ExternalKnowledgeApis.query.filter_by(110            id=external_knowledge_api_id, tenant_id=tenant_id111        ).first()112        if external_knowledge_api is None:113            raise ValueError("api template not found")114 115        db.session.delete(external_knowledge_api)116        db.session.commit()117 118    @staticmethod119    def external_knowledge_api_use_check(external_knowledge_api_id: str) -> tuple[bool, int]:120        count = ExternalKnowledgeBindings.query.filter_by(external_knowledge_api_id=external_knowledge_api_id).count()121        if count > 0:122            return True, count123        return False, 0124 125    @staticmethod126    def get_external_knowledge_binding_with_dataset_id(tenant_id: str, dataset_id: str) -> ExternalKnowledgeBindings:127        external_knowledge_binding = ExternalKnowledgeBindings.query.filter_by(128            dataset_id=dataset_id, tenant_id=tenant_id129        ).first()130        if not external_knowledge_binding:131            raise ValueError("external knowledge binding not found")132        return external_knowledge_binding133 134    @staticmethod135    def document_create_args_validate(tenant_id: str, external_knowledge_api_id: str, process_parameter: dict):136        external_knowledge_api = ExternalKnowledgeApis.query.filter_by(137            id=external_knowledge_api_id, tenant_id=tenant_id138        ).first()139        if external_knowledge_api is None:140            raise ValueError("api template not found")141        settings = json.loads(external_knowledge_api.settings)142        for setting in settings:143            custom_parameters = setting.get("document_process_setting")144            if custom_parameters:145                for parameter in custom_parameters:146                    if parameter.get("required", False) and not process_parameter.get(parameter.get("name")):147                        raise ValueError(f'{parameter.get("name")} is required')148 149    @staticmethod150    def process_external_api(151        settings: ExternalKnowledgeApiSetting, files: Union[None, dict[str, Any]]152    ) -> httpx.Response:153        """154        do http request depending on api bundle155        """156 157        kwargs = {158            "url": settings.url,159            "headers": settings.headers,160            "follow_redirects": True,161        }162 163        response = getattr(ssrf_proxy, settings.request_method)(data=json.dumps(settings.params), files=files, **kwargs)164 165        return response166 167    @staticmethod168    def assembling_headers(authorization: Authorization, headers: Optional[dict] = None) -> dict[str, Any]:169        authorization = deepcopy(authorization)170        if headers:171            headers = deepcopy(headers)172        else:173            headers = {}174        if authorization.type == "api-key":175            if authorization.config is None:176                raise ValueError("authorization config is required")177 178            if authorization.config.api_key is None:179                raise ValueError("api_key is required")180 181            if not authorization.config.header:182                authorization.config.header = "Authorization"183 184            if authorization.config.type == "bearer":185                headers[authorization.config.header] = f"Bearer {authorization.config.api_key}"186            elif authorization.config.type == "basic":187                headers[authorization.config.header] = f"Basic {authorization.config.api_key}"188            elif authorization.config.type == "custom":189                headers[authorization.config.header] = authorization.config.api_key190 191        return headers192 193    @staticmethod194    def get_external_knowledge_api_settings(settings: dict) -> ExternalKnowledgeApiSetting:195        return ExternalKnowledgeApiSetting.parse_obj(settings)196 197    @staticmethod198    def create_external_dataset(tenant_id: str, user_id: str, args: dict) -> Dataset:199        # check if dataset name already exists200        if Dataset.query.filter_by(name=args.get("name"), tenant_id=tenant_id).first():201            raise DatasetNameDuplicateError(f"Dataset with name {args.get('name')} already exists.")202        external_knowledge_api = ExternalKnowledgeApis.query.filter_by(203            id=args.get("external_knowledge_api_id"), tenant_id=tenant_id204        ).first()205 206        if external_knowledge_api is None:207            raise ValueError("api template not found")208 209        dataset = Dataset(210            tenant_id=tenant_id,211            name=args.get("name"),212            description=args.get("description", ""),213            provider="external",214            retrieval_model=args.get("external_retrieval_model"),215            created_by=user_id,216        )217 218        db.session.add(dataset)219        db.session.flush()220 221        external_knowledge_binding = ExternalKnowledgeBindings(222            tenant_id=tenant_id,223            dataset_id=dataset.id,224            external_knowledge_api_id=args.get("external_knowledge_api_id"),225            external_knowledge_id=args.get("external_knowledge_id"),226            created_by=user_id,227        )228        db.session.add(external_knowledge_binding)229 230        db.session.commit()231 232        return dataset233 234    @staticmethod235    def fetch_external_knowledge_retrieval(236        tenant_id: str, dataset_id: str, query: str, external_retrieval_parameters: dict237    ) -> list:238        external_knowledge_binding = ExternalKnowledgeBindings.query.filter_by(239            dataset_id=dataset_id, tenant_id=tenant_id240        ).first()241        if not external_knowledge_binding:242            raise ValueError("external knowledge binding not found")243 244        external_knowledge_api = ExternalKnowledgeApis.query.filter_by(245            id=external_knowledge_binding.external_knowledge_api_id246        ).first()247        if not external_knowledge_api:248            raise ValueError("external api template not found")249 250        settings = json.loads(external_knowledge_api.settings)251        headers = {"Content-Type": "application/json"}252        if settings.get("api_key"):253            headers["Authorization"] = f"Bearer {settings.get('api_key')}"254        score_threshold_enabled = external_retrieval_parameters.get("score_threshold_enabled") or False255        score_threshold = external_retrieval_parameters.get("score_threshold", 0.0) if score_threshold_enabled else 0.0256        request_params = {257            "retrieval_setting": {258                "top_k": external_retrieval_parameters.get("top_k"),259                "score_threshold": score_threshold,260            },261            "query": query,262            "knowledge_id": external_knowledge_binding.external_knowledge_id,263        }264 265        external_knowledge_api_setting = {266            "url": f"{settings.get('endpoint')}/retrieval",267            "request_method": "post",268            "headers": headers,269            "params": request_params,270        }271        response = ExternalDatasetService.process_external_api(272            ExternalKnowledgeApiSetting(**external_knowledge_api_setting), None273        )274        if response.status_code == 200:275            return response.json().get("records", [])276        return []277