Vector Retrieval Service for Milvus:Reduza os custos de invocação de IA no Alibaba Cloud Milvus com Embedding Cache e AI Batch

import json
import uuid
from urllib.error import HTTPError
from urllib.request import Request, urlopen

from pymilvus import DataType, Function, FunctionType, MilvusClient

MILVUS_URI = "http://c-xxx.milvus.aliyuncs.com:19530"  # The port must be 19530
MILVUS_TOKEN = "root:xxx"
MODEL_NAME = "text-embedding-v4"
VECTOR_DIM = 1024

# Cache configuration: exact match, Redis backend, 24-hour TTL
CACHE_CONFIG = json.dumps({
    "enabled": True,
    "exact_cache": {
        "enabled": True,
        "backend": "redis",
        "ttl_hours": 24,
    },
})

client = MilvusClient(uri=MILVUS_URI, token=MILVUS_TOKEN)
collection_name = "ai_embedding_cache_demo"

if client.has_collection(collection_name):
    client.drop_collection(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("content", DataType.VARCHAR, max_length=4096)
schema.add_field("embedding", DataType.FLOAT_VECTOR, dim=VECTOR_DIM)
schema.add_function(
    Function(
        name="embed_content_with_cache",
        function_type=FunctionType.TEXTEMBEDDING,
        input_field_names=["content"],
        output_field_names=["embedding"],
        params={
            "provider": "aliyun_milvus",
            "model_name": MODEL_NAME,
            "dim": VECTOR_DIM,
            "cache": CACHE_CONFIG,        # ← Enable Embedding Cache
        },
    )
)

index_params = client.prepare_index_params()
index_params.add_index(field_name="embedding", index_type="AUTOINDEX", metric_type="COSINE")
client.create_collection(collection_name=collection_name, schema=schema,
                        index_params=index_params)

# Content that hits the cache during a write generates no model invocation
row = {"content": "Milvus is an open-source vector database."}
client.insert(collection_name, [row])
client.flush(collection_name)

# ==================== Verify whether the cache is actually hit ====================
# Criteria: whether usage.total_tokens drops to zero and whether request_id is empty.
# When both hold, no model invocation actually happened, which means a cache hit.

def post_json(path, body, timeout=180):
    request = Request(
        f"{MILVUS_URI.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"))

# Use a brand-new piece of text so that the first request is guaranteed to miss
text = f"cache hit verification {uuid.uuid4().hex[:12]}"
body = {
    "model_name": MODEL_NAME,
    "texts": [text],
    "params": {"dim": VECTOR_DIM, "cache": CACHE_CONFIG},
}

for i in (1, 2, 3):
    status, data = post_json("/v2/vectordb/ai/embedding", body)
    assert status == 200 and data.get("code") == 0, data
    usage = data["data"].get("usage", {})
    request_id = data["data"].get("request_id", "")
    total_tokens = usage.get("total_tokens")
    hit = (total_tokens == 0) and not request_id
    print(f"Request {i}: total_tokens={total_tokens} "
          f"request_id={'(empty)' if not request_id else request_id} "
          f"-> {'cache hit' if hit else 'cache miss, the model was invoked'}")

{"custom_id":"article-001","method":"POST","url":"/v1/chat/completions","body":{"model":"qwen3.7-max","messages":[{"role":"user","content":"Generate a one-sentence summary for this knowledge base article..."}],"enable_thinking":false}}
{"custom_id":"article-002","method":"POST","url":"/v1/chat/completions","body":{"model":"qwen3.7-max","messages":[{"role":"user","content":"Generate a one-sentence summary for this knowledge base article..."}],"enable_thinking":false}}

Etapa 1: Fazer upload do arquivo de entrada

Faça upload do input.jsonl e guarde o input_file_id retornado, que identifica a entrada da tarefa que você criará na próxima etapa.

# 1) Upload the JSONL input file
status, data = upload_input_file("input.jsonl")
input_file_id = (data.get("data") or {}).get("id")
if status != 200 or not input_file_id:
    raise SystemExit(f"Upload failed: HTTP={status} message={data.get('message')}")
print(f"input_file_id = {input_file_id}")

Etapa 2: Criar a tarefa em lote

Crie a tarefa a partir do input_file_id e declare a janela de conclusão. Em testes, tarefas desse tipo podem levar várias horas da criação à conclusão, então completion_window geralmente é definido como 24h.

# 2) Create the Batch job and declare the completion window
status, data = post_json("/v2/vectordb/ai/batch/jobs/create", {
    "provider": PROVIDER,
    "input_file_id": input_file_id,
    "endpoint": ENDPOINT,
    "completion_window": "24h",
})
batch_id = (data.get("data") or {}).get("id")
if status != 200 or not batch_id:
    raise SystemExit(f"Creation failed: HTTP={status} message={data.get('message')}")
print(f"batch_id = {batch_id}")

Etapa 3: Consultar o status da tarefa

Tarefas em lote são tarefas assíncronas em escala horária, portanto não faça consultas intensivas em primeiro plano. Nos testes, uma tarefa enviada permanece no estado in_progress por muito tempo. Em produção, registre o batch_id após o envio e consulte o status posteriormente a partir de uma tarefa agendada, em vez de bloquear um thread de negócio.

# 3) Poll the job status. Batch jobs run asynchronously on an hourly scale. In production,
#    record the batch_id and poll later from a scheduled task instead of polling intensively in the foreground.
while True:
    status, data = post_json("/v2/vectordb/ai/batch/jobs/describe",
                             {"provider": PROVIDER, "batch_id": batch_id})
    batch = data.get("data") or {}
    batch_status = batch.get("status")
    print(f"{batch_status}  {batch.get('request_counts')}")
    if batch_status in {"completed", "failed", "expired", "cancelled"}:
        break
    time.sleep(300)      # Polling once every 5 minutes is enough

Etapa 4: Baixar os resultados

Não baixe os resultados até que o status seja completed. Linhas com falha são coletadas em error_file_id e podem ser tentadas novamente separadamente.

# 4) Download the results after completion. Failed lines are in error_file_id and can be retried separately
if batch_status == "completed":
    download_batch_file(batch_id, "output", Path("output.jsonl"))
    if batch.get("error_file_id"):
        download_batch_file(batch_id, "error", Path("error.jsonl"))