Dify平台集成阿里云AI安全护栏，构建AI Runtime安全防线

from fastapi import FastAPI, Body, HTTPException, Headerfrom pydantic import BaseModelimport base64from collections.abc import Generatorfrom typing import Anyimport hmacimport hashlibfrom urllib.parse import quoteimport requestsfrom datetime import datetimefrom datetime import timezoneimport uuidimport jsonimport reimport concurrent.futures# 可以根据需要调用不同区域的服务，支持上海(cn-shanghai)、北京(cn-beijing)、杭州(cn-hangzhou)、深圳(cn-shenzhen)SERVICE_URL = "https://green-cip.cn-shanghai.aliyuncs.com"# 超过这个长度时对文本进行切分MAX_LENGTH = 2000# 调用安全护栏的输入检测和输出检测的ServiceCodeSERVICE_INPUT = "query_security_check"SERVICE_OUTPUT = "response_security_check"ENCODING = "UTF-8"ISO8601_DATE_FORMAT = "%Y-%m-%dT%H:%M:%SZ"ALGORITHM = "HmacSHA1"def format_iso8601_date():    return datetime.now(timezone.utc).strftime(ISO8601_DATE_FORMAT)def percent_encode(value):    if value is None:        return ""    return (        quote(value.encode(ENCODING), safe="~").replace("+", "%20").replace("*", "%2A")    )def create_signature(string_to_sign, secret):    secret = secret + "&"    signature = hmac.new(        secret.encode(ENCODING), string_to_sign.encode(ENCODING), hashlib.sha1    ).digest()    return base64.b64encode(signature).decode(ENCODING)def create_string_to_sign(http_method, parameters):    sorted_keys = sorted(parameters.keys())    canonicalized_query_string = ""    for key in sorted_keys:        canonicalized_query_string += (            "&" + percent_encode(key) + "=" + percent_encode(parameters[key])        )    string_to_sign = (        http_method        + "&"        + percent_encode("/")        + "&"        + percent_encode(canonicalized_query_string[1:])    )    return string_to_signdef split_text(text: str, max_length: int = 1950) -> list[str]:    """将文本按 max_length 分段，尽量保留完整句子（识别多种标点）"""    segments = []    while len(text) > max_length:        # 提取当前最大长度范围内的子串        chunk = text[:max_length]        # 使用正则查找最后一个句号、感叹号、问号等断句符号的位置        match = None        for pattern in [r"[。！？；:\.?!]+"]:  # 匹配多种结束符号            matches = list(re.finditer(pattern, chunk))            if matches:                match = matches[-1]  # 取最后一个匹配项        if match:            cut_point = match.end()  # 包含标点符号        else:            cut_point = max_length  # 找不到就强制截断        segments.append(text[:cut_point])        text = text[cut_point:]    if text:        segments.append(text)    return segmentsdef request(content_segment, type, aliyun_access_key, aliyun_access_secret):    print(datetime.now(), f" [{type} request content]-> {content_segment}")    # 3.1 构造请求参数    parameters = {        "Action": "MultiModalGuard",        "Version": "2022-03-02",        "AccessKeyId": aliyun_access_key,        "Timestamp": format_iso8601_date(),        "SignatureMethod": "HMAC-SHA1",        "SignatureVersion": "1.0",        "SignatureNonce": str(uuid.uuid4()),        "Format": "JSON",        "Service": (            SERVICE_INPUT if type == "input" else SERVICE_OUTPUT        ),        "ServiceParameters": json.dumps(            {"content": content_segment}, ensure_ascii=False        ),    }    string_to_sign = create_string_to_sign("POST", parameters)    signature = create_signature(string_to_sign, aliyun_access_secret)    parameters["Signature"] = signature    # 3.2 发送请求    response = requests.post(SERVICE_URL, data=parameters)    body = response.json()    print(datetime.now(), " [response body]-> ", body)    if response.status_code != 200:        raise Exception(            f"response http status_code not 200. status_code: {response.status_code}, body: {body}"        )    if body.get("Code") != 200:        raise Exception(            f"response code not 200. code: {body.get('Code')}, body: {body}"        )    return bodyapp = FastAPI()class InputData(BaseModel):    point: str    params: dict = {}@app.post("/api/dify/receive")async def dify_receive(data: InputData = Body(...), authorization: str = Header(None)):    """    Receive API query data from Dify.    """    #print(data)    auth_scheme, _, api_key = authorization.partition(" ")    if auth_scheme.lower() != "bearer":        raise HTTPException(status_code=401, detail="Unauthorized")    # api_key decode    try:        decoded_bytes = base64.b64decode(api_key)        decoded_str = decoded_bytes.decode("utf-8")        ak, sk = decoded_str.split(":", 1)    except Exception as e:        # 如果调用失败，抛出异常        raise HTTPException(status_code=401, detail=f"Base64 Decode AK/SK fail: {e}")    point = data.point    if point == "ping":        return {"result": "pong"}    if point == "app.moderation.input":        return handle_app_moderation_input(params=data.params, ak=ak, sk=sk)    elif point == "app.moderation.output":        return handle_app_moderation_output(params=data.params, ak=ak, sk=sk)    raise HTTPException(status_code=400, detail="Not implemented")def handle_app_moderation_input(params: dict, ak: str, sk: str):    app_id = params.get("app_id")    inputs = params.get("inputs", {})    query = params.get("query")    contents = (        [query] if len(query) <= MAX_LENGTH else split_text(query, MAX_LENGTH - 50)    )    # 并发执行    bodys = []    with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:        futures = [executor.submit(request, seg, "input", ak, sk) for seg in contents]        for future in concurrent.futures.as_completed(futures):            bodys.append(future.result())    contentModerationSuggestion=""    sensitiveDataSuggestion=""    promptAttackSuggestion=""    maliciousUrlSuggestion=""    _finalSuggestion="pass"    desensitization=""    # 遍历bodys解析出各个检测项的建议    for body in bodys:        finalSuggestion = body.get("Data", {}).get("Suggestion", "")        detailList = body.get("Data", {}).get("Detail", [])        if finalSuggestion and _finalSuggestion!="block" :            _finalSuggestion = finalSuggestion        for detail in detailList:            suggestion = detail.get("Suggestion", "")            type = detail.get("Type", "")            if type == "contentModeration":                if suggestion and contentModerationSuggestion!="block" :                    contentModerationSuggestion = suggestion            elif type == "sensitiveData":                desensitization = detail.get("Result",[])[0].get("Ext",{}).get("Desensitization","")                if suggestion and sensitiveDataSuggestion!="block" :                    sensitiveDataSuggestion = suggestion            elif type == "promptAttack":                if suggestion and promptAttackSuggestion!="block" :                    promptAttackSuggestion = suggestion            elif type == "maliciousUrl":                if suggestion and maliciousUrlSuggestion!="block" :                    maliciousUrlSuggestion = suggestion    # 可以根据不同的场景返回不同的回答内容    output_response = "Your content violates our usage policy."    if contentModerationSuggestion=="block":        output_response = "Your content involves content security."    elif sensitiveDataSuggestion=="block" or sensitiveDataSuggestion=="mask":        output_response = "Your content involves sensitive data."    elif promptAttackSuggestion=="block":        output_response = "Your content involves prompt attack."    elif maliciousUrlSuggestion=="block":        output_response = "Your content involves malicious url."    flagged = False    action = "direct_output"    if _finalSuggestion == "block" :        flagged = True    elif sensitiveDataSuggestion=="mask":        flagged = True        action = "overridden"        query = desensitization    response = {"flagged": flagged, "action": action}    if flagged:        if action == "direct_output":            response["preset_response"] = output_response        elif action == "overridden":            response["inputs"] = inputs            response["query"] = query    print(response)    return responsedef handle_app_moderation_output(params: dict, ak: str, sk: str):    app_id = params.get("app_id")    text = params.get("text", "")    print(f"handle_app_moderation_output length:{len(text)}")    # 获取最近的2000字符，大小根据需要调整，建议大于dify的窗口大小    if len(text) > MAX_LENGTH:        content = text[-MAX_LENGTH:]    else:        content = text    # 执行检测    body = request(content, "output", ak, sk)    contentModerationSuggestion=""    sensitiveDataSuggestion=""    promptAttackSuggestion=""    maliciousUrlSuggestion=""    desensitization=""    _finalSuggestion=body.get("Data", {}).get("Suggestion", "")    detailList = body.get("Data", {}).get("Detail", [])    for detail in detailList:        suggestion = detail.get("Suggestion", "")        type = detail.get("Type", "")        if type == "contentModeration":            contentModerationSuggestion = suggestion        elif type == "sensitiveData":            desensitization = detail.get("Result",[])[0].get("Ext",{}).get("Desensitization","")            sensitiveDataSuggestion = suggestion        elif type == "promptAttack":            promptAttackSuggestion = suggestion        elif type == "maliciousUrl":            maliciousUrlSuggestion = suggestion    # 可以根据不同的场景返回不同的回答内容    output_response = "Your content violates our usage policy."    if contentModerationSuggestion=="block":        output_response = "Your content involves content security."    elif sensitiveDataSuggestion=="block":        output_response = "Your content involves sensitive data."    elif promptAttackSuggestion=="block":        output_response = "Your content involves prompt attack."    elif maliciousUrlSuggestion=="block":        output_response = "Your content involves malicious url."    flagged = False    action = "direct_output"    if _finalSuggestion == "block":        flagged = True    elif sensitiveDataSuggestion=="mask":        flagged = True        action = "overridden"    response = {"flagged": flagged, "action": action}    if flagged:        if action == "direct_output":            response["preset_response"] = output_response        elif action == "overridden":            response["text"] = desensitization    print(response)    return responseif __name__ == "__main__":    import uvicorn    # 开放端口可以根据自定义选择    uvicorn.run(app, host="0.0.0.0", port=8000, reload=True)