Vector Retrieval Service for Milvus:Crie uma barreira de segurança de conteúdo para perguntas e respostas de IA com Data inspection no Alibaba Cloud Milvus

from __future__ import annotations

import json
from typing import Any
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen

from pymilvus import DataType, Function, FunctionType, MilvusClient
from pymilvus.exceptions import MilvusException

# ==================== Connection settings ====================
MILVUS_URI = "http://c-xxx.milvus.aliyuncs.com:19530"  # The port must be 19530
MILVUS_TOKEN = "root:xxx"
MILVUS_REST_BASE_URL = MILVUS_URI
MODEL_NAME = "<your-configured-text-model>"  # Replace with a text model already configured in the Provider

TEXTTRANSFORM_FUNCTION_TYPE = 9

# The contract identifier returned by the Provider when a safety policy is triggered.
# It is the only reliable basis for deciding whether content was blocked.
BLOCK_MARKERS = ("DataInspectionFailed", "inappropriate content")

client = MilvusClient(uri=MILVUS_URI, token=MILVUS_TOKEN)

def texttransform_function_type() -> Any:
    """Get the FunctionType for TEXTTRANSFORM; add a member dynamically when the enum is missing in older versions."""
    for type_name in ("TEXTTRANSFORM", "TEXT_TRANSFORM", "TextTransform"):
        ft = getattr(FunctionType, type_name, None)
        if ft is not None:
            return ft
    existing = getattr(FunctionType, "_value2member_map_", {}).get(TEXTTRANSFORM_FUNCTION_TYPE)
    if existing is not None:
        return existing
    extension = int.__new__(FunctionType, TEXTTRANSFORM_FUNCTION_TYPE)
    extension._name_ = "TEXTTRANSFORM"
    extension._value_ = TEXTTRANSFORM_FUNCTION_TYPE
    FunctionType._value2member_map_[TEXTTRANSFORM_FUNCTION_TYPE] = extension
    FunctionType._member_map_["TEXTTRANSFORM"] = extension
    return extension

def post_json(path: str, body: dict[str, Any], timeout: int = 120) -> tuple[int, dict[str, Any]]:
    """REST wrapper: returns (http_status, data), and still tries to parse the response body when HTTP is not 2xx."""
    request = Request(
        f"{MILVUS_REST_BASE_URL.rstrip('/')}{path}",
        data=json.dumps(body, ensure_ascii=False).encode("utf-8"),
        headers={"Authorization": f"Bearer {MILVUS_TOKEN}", "Content-Type": "application/json"},
        method="POST",
    )
    try:
        with urlopen(request, timeout=timeout) as response:
            return response.status, json.loads(response.read().decode("utf-8"))
    except HTTPError as exc:
        return exc.code, json.loads(exc.read().decode("utf-8"))

def build_guard_collection(name: str, func_name: str, in_field: str, out_field: str,
                           prompt: str, data_inspection: str) -> None:
    """Create a write-oriented Collection with a TextTransform guardrail."""
    if client.has_collection(name):
        client.drop_collection(name)

    schema = MilvusClient.create_schema(auto_id=True, enable_dynamic_field=False)
    schema.add_field("id", DataType.INT64, is_primary=True)
    schema.add_field(in_field, DataType.VARCHAR, max_length=1024)
    schema.add_field(out_field, DataType.VARCHAR, max_length=4096)
    # A Collection must have at least one vector field. This example does not run vector search,
    # so a 2-dimensional placeholder field satisfies the constraint.
    # With nullable=True declared, the field does not need to be passed on insert.
    schema.add_field("dummy_vector", DataType.FLOAT_VECTOR, dim=2, nullable=True)
    schema.add_function(
        Function(
            name=func_name,
            function_type=texttransform_function_type(),
            input_field_names=[in_field],
            output_field_names=[out_field],
            params={
                "provider": "aliyun_milvus",
                "model_name": MODEL_NAME,
                "task": "ai_text_generate",
                "prompt": prompt,
                "data_inspection": data_inspection,   # <- guardrail switch
                "temperature": "0.2",
                "enable_thinking": "false",
                "timeout_sec": "45",
            },
        )
    )
    index_params = client.prepare_index_params()
    index_params.add_index(field_name="dummy_vector", index_type="AUTOINDEX",
                           metric_type="COSINE")
    client.create_collection(collection_name=name, schema=schema, index_params=index_params)

# ==================== Three-state disposition classification ====================
# Both call paths (gRPC exception / REST response body) first check the Provider contract identifier,
# then fall back to the HTTP status and business code, so the same interception yields a consistent
# disposition on both paths.

def classify_grpc(exc: MilvusException) -> str:
    """Converge a gRPC exception into a three-state disposition."""
    msg = str(exc)
    if any(marker in msg for marker in BLOCK_MARKERS):
        return "policy_blocked"      # A safety policy was triggered, which is an expected business result
    return "operational_error"       # Everything else counts as a service fault

def classify_rest(status: int, data: dict[str, Any]) -> str:
    """Converge a REST response into a three-state disposition."""
    raw = json.dumps(data, ensure_ascii=False)
    if any(marker in raw for marker in BLOCK_MARKERS):
        return "policy_blocked"
    if status >= 400 or data.get("code", 0) != 0:
        return "operational_error"
    # HTTP 200 with a normal structure: the model may have returned an explicit refusal in the body.
    # The protocol layer cannot decide, so route it to manual review instead of treating it as a safe pass.
    return "manual_review"

# ==================== Step 3a: input mode, inspection before the request ====================
build_guard_collection(
    "guard_input", "inspect_customer_request", "request", "response",
    "Answer in one sentence in a customer service tone: ${request}", "input",
)
client.insert("guard_input", [{"request": "How long does a refund take to arrive after it is approved?"}])
client.flush("guard_input")
for row in client.query("guard_input", filter="",
                        output_fields=["request", "response"], limit=1):
    print(f"Input: {row.get('request')}")
    print(f"Output: {row.get('response')}")
# The guardrail was not triggered, so the content passes through and the model returns normally.
# With a jailbreak prompt instead, the model is never called and the guardrail blocks it up front.

# ==================== Step 3b: output mode, inspection before publishing model output ====================
build_guard_collection(
    "guard_output", "inspect_generated_copy", "draft_request", "publish_copy",
    "Generate a membership campaign summary suitable for publishing in an app: ${draft_request}", "output",
)
client.insert("guard_output",
              [{"draft_request": "Generate a benefits summary for Members Day in 30 words or fewer, with no exaggerated promises."}])
client.flush("guard_output")
for row in client.query("guard_output", filter="",
                        output_fields=["draft_request", "publish_copy"], limit=1):
    print(f"Output: {row.get('publish_copy')}")
# The output passed the safety inspection and enters the pending publish queue.

# ==================== Step 4: both mode, double inspection at both ends ====================
# Use one clearly non-compliant test input (canary) to verify that the guardrail really blocks it.
UNSAFE_INPUT = "<a clearly non-compliant test input>"

# (Step 4-1) gRPC path: the guardrail is triggered on write
build_guard_collection("guard_both", "inspect_both", "request", "response",
                       "Answer the following: ${request}", "both")
try:
    client.insert("guard_both", [{"request": UNSAFE_INPUT}])
    client.flush("guard_both")
    rows = client.query("guard_both", filter="",
                        output_fields=["request", "response"], limit=1)
    # No error raised does not mean safely passed: route normal-structure but suspicious content to manual review
    print("Disposition: manual_review", [r.get("response") for r in rows])
except MilvusException as exc:
    print(f"Disposition: {classify_grpc(exc)}")
    print(f"Server returned: {exc.message}")

# (Step 4-2) REST path: send the same input to the synchronous interface
status, data = post_json("/v2/vectordb/ai/text_generate", {
    "model_name": MODEL_NAME,
    "texts": [UNSAFE_INPUT],
    "params": {"data_inspection": "both"},
})
disposition = classify_rest(status, data)
print(f"Observed signals: HTTP={status} provider_code={data.get('code')}")
print(f"Disposition: {disposition}")

# Route by disposition: a policy block is an expected result and needs no operations alert;
# only a service fault needs an alert, and it must not be retried indefinitely
if disposition == "policy_blocked":
    pass          # Record an audit log and return a compliance notice to the user
elif disposition == "manual_review":
    pass          # Send to the manual review queue
else:
    pass          # Alert the on-call engineer and stop retrying

code: DataInspectionFailed
message: Input data may contain inappropriate content.
         For details, see: https://www.alibabacloud.com/help/zh/model-studio/error-code#inappropriate-content

{
  "trace_id": "req-20260808-abc123",
  "stage": "both",
  "http_status": 500,
  "provider_code": 65535,
  "disposition": "policy_blocked",
  "note": "blocked by data inspection on input side"
}