Vector Retrieval Service for Milvus:Build a content safety guardrail for AI question answering with Data inspection in Alibaba Cloud Milvus

Step 1: Prepare the shared code

The following code contains the connection settings, a REST wrapper, a compatibility fallback for the TEXTTRANSFORM type, and a helper function that creates a Collection with a guardrail. The guardrail switch is the single data_inspection line inside params.

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)

A Collection must contain at least one vector field, otherwise creation fails with schema does not contain vector field. This example does not run vector search, so a 2-dimensional placeholder vector field satisfies the constraint and is declared nullable=True to avoid passing a value on every write.

Step 2: Converge interception results into three dispositions

This is the engineering difficulty that is most often underestimated when a guardrail goes into production. When a safety policy is triggered, the server may return an HTTP error, a non-zero business code, or an explicit refusal as text inside HTTP 200. The client therefore needs to converge the observed signals into three mutually exclusive dispositions:

# ==================== 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"

classify_rest is deliberately conservative. When the response is HTTP 200 with a normal structure, the protocol layer alone cannot certify that the body is safe, because the model may have embedded an explicit refusal in it. The function therefore returns manual_review for this undecidable case instead of silently passing it. Pair classify_rest with your own answer validation so that a response you can confirm as a genuine business answer is treated as a pass, and only the responses that remain undecided reach the review queue. This is what keeps the scenario table's "Passed through" for normal input consistent with the goal of routing only a small number of ambiguous samples to manual review.

Why the decision must be based on the Provider contract identifier rather than the HTTP status code alone. Testing verified that "a non-compliant input blocked by the guardrail" and "a service fault caused by a non-existent model name" return exactly the same status code combination:

The HTTP status code and business code in the first two rows are identical, so they cannot distinguish "content blocked" from "service failed". The operational meaning of the two is completely different: misclassifying a policy block as a service fault turns every content safety interception into a false service fault alert, which over time buries real faults. The decision must therefore rely on the DataInspectionFailed contract identifier in the response body. For the same reason, do not depend on specific HTTP status code values, because they can change across gateway versions.

Do not use keywords such as "risk", "reject", or "cannot" to decide whether content was blocked, because normal business text can also contain these words and cause false blocks. Keywords are at most an auxiliary signal and must not change the final disposition.

Step 3: Inspect compliant content in input and output modes

The input mode inspects user input before the model is called, and the output mode inspects model output before the result is returned. Compliant content passes through normally.

# ==================== 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.

The following results were verified in testing:

When the input mode is triggered, the model is not called, so the risk is blocked and the model call cost is saved at the same time. The output mode suits scenarios where the input is trusted and the only concern is the model improvising.

Step 4: Block non-compliant input in both mode

Use both when neither end is trusted. The following example verifies that the guardrail actually works by sending one clearly non-compliant test input. Run this guardrail effectiveness check before launch.

Replace <a clearly non-compliant test input> with a real non-compliant test string (a canary) before you run this step. The placeholder shown below is inert and does not trigger the guardrail, so you cannot reproduce the interception result until you substitute a real value.

# ==================== 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

Testing verified that both paths blocked the input successfully, and the server returned a consistent contract identifier:

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

The guardrail was triggered on the input side, the write was blocked, and the model generated no content. After the classification functions process the signals, both the gRPC path and the REST path reach the same policy_blocked conclusion.

Testing also confirmed that the both mode does not falsely block normal business input: sending a normal customer service question to the same interface returned HTTP 200 with a complete answer.

Step 5: Log dispositions safely

After the guardrail is in production, logs must locate problems without becoming a new data leak surface. Follow these recommendations:

• Record only machine-readable, redacted fields: trace id, inspection stage, HTTP status, provider code, disposition enum, and disposition action.

• Trim or redact raw text and personal sensitive information. During review, retrieve the data through an authorized channel using the trace id instead of keeping raw text in business logs.

• Stop and alert when a policy is triggered or operational_error occurs, and do not retry indefinitely.

• Combine with AI_PII_MASK: redact before logging to keep the sensitive data exposure surface as small as possible.

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