import json from copy import deepcopy from typing import Any, cast from urllib.parse import urlparse import httpx from sqlalchemy import func, select from constants import HIDDEN_VALUE from core.helper import ssrf_proxy from core.rag.entities import MetadataFilteringCondition from extensions.ext_database import db from graphon.nodes.http_request.exc import InvalidHttpMethodError from libs.datetime_utils import naive_utc_now from models.dataset import ( Dataset, ExternalKnowledgeApis, ExternalKnowledgeBindings, ) from services.entities.external_knowledge_entities.external_knowledge_entities import ( Authorization, ExternalKnowledgeApiSetting, ) from services.errors.dataset import DatasetNameDuplicateError class ExternalDatasetService: @staticmethod def get_external_knowledge_apis( page, per_page, tenant_id, search=None ) -> tuple[list[ExternalKnowledgeApis], int | None]: query = ( select(ExternalKnowledgeApis) .where(ExternalKnowledgeApis.tenant_id == tenant_id) .order_by(ExternalKnowledgeApis.created_at.desc()) ) if search: from libs.helper import escape_like_pattern escaped_search = escape_like_pattern(search) query = query.where(ExternalKnowledgeApis.name.ilike(f"%{escaped_search}%", escape="\\")) external_knowledge_apis = db.paginate( select=query, page=page, per_page=per_page, max_per_page=100, error_out=False ) return external_knowledge_apis.items, external_knowledge_apis.total @classmethod def validate_api_list(cls, api_settings: dict[str, Any]): if not api_settings: raise ValueError("api list is empty") if not api_settings.get("endpoint"): raise ValueError("endpoint is required") if not api_settings.get("api_key"): raise ValueError("api_key is required") @staticmethod def create_external_knowledge_api(tenant_id: str, user_id: str, args: dict[str, Any]) -> ExternalKnowledgeApis: settings = args.get("settings") if settings is None: raise ValueError("settings is required") ExternalDatasetService.check_endpoint_and_api_key(settings) external_knowledge_api = ExternalKnowledgeApis( tenant_id=tenant_id, created_by=user_id, updated_by=user_id, name=str(args.get("name")), description=args.get("description", ""), settings=json.dumps(args.get("settings"), ensure_ascii=False), ) db.session.add(external_knowledge_api) db.session.commit() return external_knowledge_api @staticmethod def check_endpoint_and_api_key(settings: dict[str, Any]): if "endpoint" not in settings or not settings["endpoint"]: raise ValueError("endpoint is required") if "api_key" not in settings or not settings["api_key"]: raise ValueError("api_key is required") endpoint = f"{settings['endpoint']}/retrieval" api_key = settings["api_key"] parsed_url = urlparse(endpoint) if not all([parsed_url.scheme, parsed_url.netloc]): if not endpoint.startswith("http://") and not endpoint.startswith("https://"): raise ValueError(f"invalid endpoint: {endpoint} must start with http:// or https://") else: raise ValueError(f"invalid endpoint: {endpoint}") try: response = ssrf_proxy.post(endpoint, headers={"Authorization": f"Bearer {api_key}"}) except Exception as e: raise ValueError(f"failed to connect to the endpoint: {endpoint}") from e if response.status_code == 502: raise ValueError(f"Bad Gateway: failed to connect to the endpoint: {endpoint}") if response.status_code == 404: raise ValueError(f"Not Found: failed to connect to the endpoint: {endpoint}") if response.status_code == 403: raise ValueError(f"Forbidden: Authorization failed with api_key: {api_key}") @staticmethod def get_external_knowledge_api(external_knowledge_api_id: str, tenant_id: str) -> ExternalKnowledgeApis: external_knowledge_api: ExternalKnowledgeApis | None = db.session.scalar( select(ExternalKnowledgeApis) .where(ExternalKnowledgeApis.id == external_knowledge_api_id, ExternalKnowledgeApis.tenant_id == tenant_id) .limit(1) ) if external_knowledge_api is None: raise ValueError("api template not found") return external_knowledge_api @staticmethod def update_external_knowledge_api(tenant_id, user_id, external_knowledge_api_id, args) -> ExternalKnowledgeApis: external_knowledge_api: ExternalKnowledgeApis | None = db.session.scalar( select(ExternalKnowledgeApis) .where(ExternalKnowledgeApis.id == external_knowledge_api_id, ExternalKnowledgeApis.tenant_id == tenant_id) .limit(1) ) if external_knowledge_api is None: raise ValueError("api template not found") settings = args.get("settings") if settings and settings.get("api_key") == HIDDEN_VALUE and external_knowledge_api.settings_dict: settings["api_key"] = external_knowledge_api.settings_dict.get("api_key") external_knowledge_api.name = args.get("name") external_knowledge_api.description = args.get("description", "") external_knowledge_api.settings = json.dumps(args.get("settings"), ensure_ascii=False) external_knowledge_api.updated_by = user_id external_knowledge_api.updated_at = naive_utc_now() db.session.commit() return external_knowledge_api @staticmethod def delete_external_knowledge_api(tenant_id: str, external_knowledge_api_id: str): external_knowledge_api = db.session.scalar( select(ExternalKnowledgeApis) .where(ExternalKnowledgeApis.id == external_knowledge_api_id, ExternalKnowledgeApis.tenant_id == tenant_id) .limit(1) ) if external_knowledge_api is None: raise ValueError("api template not found") db.session.delete(external_knowledge_api) db.session.commit() @staticmethod def external_knowledge_api_use_check(external_knowledge_api_id: str, tenant_id: str) -> tuple[bool, int]: """ Return usage for an external knowledge API within a single tenant. The caller already scopes access by tenant, so this query must do the same; otherwise the endpoint becomes a cross-tenant UUID oracle. """ count = ( db.session.scalar( select(func.count(ExternalKnowledgeBindings.id)).where( ExternalKnowledgeBindings.external_knowledge_api_id == external_knowledge_api_id, ExternalKnowledgeBindings.tenant_id == tenant_id, ) ) or 0 ) return count > 0, count @staticmethod def get_external_knowledge_binding_with_dataset_id(tenant_id: str, dataset_id: str) -> ExternalKnowledgeBindings: external_knowledge_binding: ExternalKnowledgeBindings | None = db.session.scalar( select(ExternalKnowledgeBindings) .where(ExternalKnowledgeBindings.dataset_id == dataset_id, ExternalKnowledgeBindings.tenant_id == tenant_id) .limit(1) ) if not external_knowledge_binding: raise ValueError("external knowledge binding not found") return external_knowledge_binding @staticmethod def document_create_args_validate( tenant_id: str, external_knowledge_api_id: str, process_parameter: dict[str, Any] ): external_knowledge_api = db.session.scalar( select(ExternalKnowledgeApis) .where(ExternalKnowledgeApis.id == external_knowledge_api_id, ExternalKnowledgeApis.tenant_id == tenant_id) .limit(1) ) if external_knowledge_api is None or external_knowledge_api.settings is None: raise ValueError("api template not found") settings = json.loads(external_knowledge_api.settings) for setting in settings: custom_parameters = setting.get("document_process_setting") if custom_parameters: for parameter in custom_parameters: if parameter.get("required", False) and not process_parameter.get(parameter.get("name")): raise ValueError(f"{parameter.get('name')} is required") @staticmethod def process_external_api(settings: ExternalKnowledgeApiSetting, files: dict[str, Any] | None) -> httpx.Response: """ do http request depending on api bundle """ kwargs: dict[str, Any] = { "url": settings.url, "headers": settings.headers, "follow_redirects": True, } _METHOD_MAP = { "get": ssrf_proxy.get, "head": ssrf_proxy.head, "post": ssrf_proxy.post, "put": ssrf_proxy.put, "delete": ssrf_proxy.delete, "patch": ssrf_proxy.patch, } method_lc = settings.request_method.lower() if method_lc not in _METHOD_MAP: raise InvalidHttpMethodError(f"Invalid http method {settings.request_method}") response: httpx.Response = _METHOD_MAP[method_lc](data=json.dumps(settings.params), files=files, **kwargs) return response @staticmethod def assembling_headers(authorization: Authorization, headers: dict[str, Any] | None = None) -> dict[str, Any]: authorization = deepcopy(authorization) if headers: headers = deepcopy(headers) else: headers = {} if authorization.type == "api-key": if authorization.config is None: raise ValueError("authorization config is required") if authorization.config.api_key is None: raise ValueError("api_key is required") if not authorization.config.header: authorization.config.header = "Authorization" if authorization.config.type == "bearer": headers[authorization.config.header] = f"Bearer {authorization.config.api_key}" elif authorization.config.type == "basic": headers[authorization.config.header] = f"Basic {authorization.config.api_key}" elif authorization.config.type == "custom": headers[authorization.config.header] = authorization.config.api_key return headers @staticmethod def get_external_knowledge_api_settings(settings: dict[str, Any]) -> ExternalKnowledgeApiSetting: return ExternalKnowledgeApiSetting.model_validate(settings) @staticmethod def create_external_dataset(tenant_id: str, user_id: str, args: dict[str, Any]) -> Dataset: # check if dataset name already exists if db.session.scalar( select(Dataset).where(Dataset.name == args.get("name"), Dataset.tenant_id == tenant_id).limit(1) ): raise DatasetNameDuplicateError(f"Dataset with name {args.get('name')} already exists.") external_knowledge_api = db.session.scalar( select(ExternalKnowledgeApis) .where( ExternalKnowledgeApis.id == args.get("external_knowledge_api_id"), ExternalKnowledgeApis.tenant_id == tenant_id, ) .limit(1) ) if external_knowledge_api is None: raise ValueError("api template not found") dataset = Dataset( tenant_id=tenant_id, name=args.get("name"), description=args.get("description", ""), provider="external", retrieval_model=args.get("external_retrieval_model"), created_by=user_id, ) db.session.add(dataset) db.session.flush() if args.get("external_knowledge_id") is None: raise ValueError("external_knowledge_id is required") if args.get("external_knowledge_api_id") is None: raise ValueError("external_knowledge_api_id is required") external_knowledge_binding = ExternalKnowledgeBindings( tenant_id=tenant_id, dataset_id=dataset.id, external_knowledge_api_id=args.get("external_knowledge_api_id") or "", external_knowledge_id=args.get("external_knowledge_id") or "", created_by=user_id, ) db.session.add(external_knowledge_binding) db.session.commit() return dataset @staticmethod def fetch_external_knowledge_retrieval( tenant_id: str, dataset_id: str, query: str, external_retrieval_parameters: dict[str, Any], metadata_condition: MetadataFilteringCondition | None = None, ): external_knowledge_binding = db.session.scalar( select(ExternalKnowledgeBindings) .where(ExternalKnowledgeBindings.dataset_id == dataset_id, ExternalKnowledgeBindings.tenant_id == tenant_id) .limit(1) ) if not external_knowledge_binding: raise ValueError("external knowledge binding not found") external_knowledge_api = db.session.scalar( select(ExternalKnowledgeApis) .where( ExternalKnowledgeApis.id == external_knowledge_binding.external_knowledge_api_id, ExternalKnowledgeApis.tenant_id == tenant_id, ) .limit(1) ) if external_knowledge_api is None or external_knowledge_api.settings is None: raise ValueError("external api template not found") settings = json.loads(external_knowledge_api.settings) headers = {"Content-Type": "application/json"} if settings.get("api_key"): headers["Authorization"] = f"Bearer {settings.get('api_key')}" score_threshold_enabled = external_retrieval_parameters.get("score_threshold_enabled") or False score_threshold = external_retrieval_parameters.get("score_threshold", 0.0) if score_threshold_enabled else 0.0 request_params = { "retrieval_setting": { "top_k": external_retrieval_parameters.get("top_k"), "score_threshold": score_threshold, }, "query": query, "knowledge_id": external_knowledge_binding.external_knowledge_id, "metadata_condition": metadata_condition.model_dump() if metadata_condition else None, } response = ExternalDatasetService.process_external_api( ExternalKnowledgeApiSetting( url=f"{settings.get('endpoint')}/retrieval", request_method="post", headers=headers, params=request_params, ), None, ) if response.status_code == 200: return cast(list[Any], response.json().get("records", [])) else: raise ValueError(response.text)