Voyage AI RAG 教程：嵌入召回、重排与证据质量评估

python -m pip install voyageai

import math
import voyageai

MODEL = "voyage-4-lite"
DIMENSION = 1024
tenant = "alpha"  # 正式应用从已认证的会话与权限数据库推导
documents = [
    {"id": "a-refund", "tenant": "alpha",
     "text": "Alpha 客服政策：退款申请由人工审核后处理。"},
    {"id": "a-shipping", "tenant": "alpha",
     "text": "Alpha 发货政策：工作日收到订单后安排发货。"},
    {"id": "a-account", "tenant": "alpha",
     "text": "Alpha 账号政策：更换邮箱需要验证当前身份。"},
    {"id": "b-refund", "tenant": "beta",
     "text": "Beta 的内部退款政策不得向 Alpha 用户展示。"},
]
allowed = [d for d in documents if d["tenant"] == tenant]
if not allowed:
    raise RuntimeError("没有获准访问的资料")

client = voyageai.Client()  # 读取 VOYAGE_API_KEY
doc_vectors = client.embed(
    [d["text"] for d in allowed],
    model=MODEL, input_type="document",
    output_dimension=DIMENSION, truncation=False,
).embeddings
query = "退款由谁审核？"
query_vector = client.embed(
    [query], model=MODEL, input_type="query",
    output_dimension=DIMENSION, truncation=False,
).embeddings[0]

def cosine(a, b):
    if len(a) != len(b):
        raise ValueError("向量维度不一致")
    denominator = math.sqrt(sum(x*x for x in a)) * math.sqrt(
        sum(x*x for x in b)
    )
    return sum(x*y for x, y in zip(a, b)) / denominator if denominator else 0.0

scored = sorted(
    zip(allowed, doc_vectors),
    key=lambda pair: cosine(query_vector, pair[1]), reverse=True,
)
# 生产系统可取前 20/50 条；示例只有三份获准访问的资料
candidates = [doc for doc, _ in scored[:20]]
reranked = client.rerank(
    query=query, documents=[d["text"] for d in candidates],
    model="rerank-3-lite", top_k=min(2, len(candidates)),
    truncation=False,
)
for result in reranked.results:
    doc = candidates[result.index]
    print(doc["id"], result.relevance_score, doc["text"])

def recall_at_k(ranked_ids, relevant_ids, k):
    relevant = set(relevant_ids)
    if not relevant:
        return None  # 无答案问题另外检查拒答
    return len(set(ranked_ids[:k]) & relevant) / len(relevant)

def reciprocal_rank(ranked_ids, relevant_ids):
    relevant = set(relevant_ids)
    for rank, doc_id in enumerate(ranked_ids, 1):
        if doc_id in relevant:
            return 1.0 / rank
    return 0.0