knowledge-adapter-service/deploy/nginx/nginx_adapter_helper.py

#!/usr/bin/env python3
import hashlib
import json
import os
import pathlib
import sys
import tempfile
import urllib.error
import urllib.request
import uuid

ENV_PATH = pathlib.Path("/opt/knowledge-adapter/adapter.env")
NEW_BASE = "https://weknora-api-test.dobot.cn"
API_HOST = "weknora-api-test.dobot.cn"
UI_HOST = "weknora-test.dobot.cn"
DIFY_KEY_PATH = pathlib.Path("/opt/knowledge-adapter/secrets/dify_test_key")


def parse_env(path: pathlib.Path) -> dict[str, str]:
    result: dict[str, str] = {}
    for raw in path.read_text().splitlines():
        line = raw.strip()
        if line and not line.startswith("#") and "=" in line:
            key, value = line.split("=", 1)
            result[key] = value
    return result


def update_env() -> None:
    st = ENV_PATH.stat()
    lines = ENV_PATH.read_text().splitlines(keepends=True)
    found = 0
    output: list[str] = []
    for line in lines:
        if line.startswith("ASSET_PUBLIC_BASE_URL="):
            output.append(f"ASSET_PUBLIC_BASE_URL={NEW_BASE}\n")
            found += 1
        else:
            output.append(line)
    if found != 1:
        raise RuntimeError(f"expected exactly one ASSET_PUBLIC_BASE_URL, got {found}")
    fd, tmp_name = tempfile.mkstemp(prefix=".adapter.env.", dir=str(ENV_PATH.parent))
    try:
        with os.fdopen(fd, "w") as handle:
            handle.writelines(output)
            handle.flush()
            os.fsync(handle.fileno())
        os.chmod(tmp_name, st.st_mode & 0o777)
        os.chown(tmp_name, st.st_uid, st.st_gid)
        os.replace(tmp_name, ENV_PATH)
    finally:
        if os.path.exists(tmp_name):
            os.unlink(tmp_name)
    if parse_env(ENV_PATH).get("ASSET_PUBLIC_BASE_URL") != NEW_BASE:
        raise RuntimeError("ASSET_PUBLIC_BASE_URL readback mismatch")
    print("adapter_env_update=ok")


def request(method: str, path: str, host: str, *, data: bytes | None = None, headers: dict[str, str] | None = None) -> tuple[int, dict[str, str], bytes]:
    merged = {"Host": host, "Accept": "application/json"}
    if headers:
        merged.update(headers)
    req = urllib.request.Request(
        f"http://127.0.0.1{path}", data=data, method=method, headers=merged
    )
    try:
        with urllib.request.urlopen(req, timeout=25) as response:
            return response.status, dict(response.headers.items()), response.read(2_000_000)
    except urllib.error.HTTPError as exc:
        return exc.code, dict(exc.headers.items()), exc.read(2_000_000)


def accept() -> None:
    env = parse_env(ENV_PATH)
    allowed = [x.strip() for x in env["WEKNORA_ALLOWED_KB_IDS"].split(",") if x.strip()]
    if len(allowed) != 1:
        raise RuntimeError("acceptance expects exactly one allowlisted KB")
    raw_key = DIFY_KEY_PATH.read_text().strip()
    if not raw_key:
        raise RuntimeError("Dify key file is empty")
    configured_hash = env["DIFY_KEY_HASH_CURRENT"].lower()
    if hashlib.sha256(raw_key.encode()).hexdigest() != configured_hash:
        raise RuntimeError("Dify key file does not match configured hash")

    status, _, _ = request("GET", "/", UI_HOST)
    if status != 200:
        raise RuntimeError(f"UI Host route expected 200, got {status}")
    print("ui_host_route=200")

    status, _, _ = request("GET", "/", API_HOST)
    if status != 404:
        raise RuntimeError(f"API root expected 404, got {status}")
    print("api_root=404")

    status, _, body = request(
        "POST", "/retrieval", API_HOST, data=b"{}", headers={"Content-Type": "application/json"}
    )
    if status != 401 or json.loads(body) != {"detail": "invalid API key"}:
        raise RuntimeError("missing-auth boundary mismatch")
    print("retrieval_missing_key=401_adapter")

    payload = json.dumps(
        {
            "knowledge_id": allowed[0],
            "query": "越疆机器人技术资料",
            "retrieval_setting": {"top_k": 3, "score_threshold": 0.0},
        }
    ).encode()
    status, _, body = request(
        "POST",
        "/retrieval",
        API_HOST,
        data=payload,
        headers={"Content-Type": "application/json", "Authorization": f"Bearer {raw_key}"},
    )
    if status != 200:
        raise RuntimeError(f"allowlisted retrieval expected 200, got {status}")
    result = json.loads(body)
    records = result.get("records")
    if not isinstance(records, list):
        raise RuntimeError("retrieval response has no records list")
    print(f"retrieval_allowed=200 records={len(records)}")

    denied_payload = json.dumps(
        {
            "knowledge_id": str(uuid.UUID("00000000-0000-4000-8000-000000000001")),
            "query": "boundary probe",
            "retrieval_setting": {"top_k": 1, "score_threshold": 0.0},
        }
    ).encode()
    status, _, body = request(
        "POST",
        "/retrieval",
        API_HOST,
        data=denied_payload,
        headers={"Content-Type": "application/json", "Authorization": f"Bearer {raw_key}"},
    )
    if status != 400 or json.loads(body) != {"detail": "invalid request"}:
        raise RuntimeError("denied-KB boundary mismatch")
    print("retrieval_denied_kb=400")


if __name__ == "__main__":
    if len(sys.argv) != 2 or sys.argv[1] not in {"update-env", "accept"}:
        raise SystemExit("usage: helper.py update-env|accept")
    if sys.argv[1] == "update-env":
        update_env()
    else:
        accept()