RAG 实战调优

import os
import numpy as np
import requests

BASE = "https://api.apiyi.com/v1"
HEADERS = {"Authorization": f"Bearer {os.environ['APIYI_API_KEY']}"}


def embed(texts):
    """批量向量化。真实项目里文档向量应离线算好存进向量库。"""
    r = requests.post(f"{BASE}/embeddings", headers=HEADERS, timeout=60,
                      json={"model": "text-embedding-3-small", "input": texts})
    r.raise_for_status()
    data = sorted(r.json()["data"], key=lambda x: x["index"])
    return np.array([d["embedding"] for d in data])


def rerank(query, documents, top_n=5):
    """精排。documents 只接受字符串数组。"""
    r = requests.post(f"{BASE}/rerank", headers=HEADERS, timeout=120,
                      json={"model": "bge-reranker-v2-m3", "query": query,
                            "documents": documents, "top_n": top_n})
    r.raise_for_status()
    return r.json()["results"]


def search(query, corpus, recall_k=50, final_k=5):
    """corpus: [{'id':..., 'text':..., 'meta':...}, ...]"""
    # ---- 第一级：向量召回 ----
    doc_vecs = embed([c["text"] for c in corpus])          # 生产环境应来自向量库
    q_vec = embed([query])[0]
    doc_vecs /= np.linalg.norm(doc_vecs, axis=1, keepdims=True)
    q_vec /= np.linalg.norm(q_vec)
    recalled_idx = np.argsort(-(doc_vecs @ q_vec))[:recall_k]
    candidates = [corpus[i] for i in recalled_idx]

    # ---- 第二级：重排序精排 ----
    results = rerank(query, [c["text"] for c in candidates], top_n=final_k)

    # 关键：用 index 回填原始文档对象，保住 id 和 meta
    return [{**candidates[r["index"]], "score": r["relevance_score"]}
            for r in results]


corpus = [
    {"id": "kb-001", "text": "429 表示触发了速率限制。解决办法是降低并发、"
                             "在客户端实现指数退避重试，或联系服务商提升账号的 RPM 配额。"},
    {"id": "kb-002", "text": "HTTP 状态码 4xx 系列表示客户端错误，常见的有 400 参数错误、"
                             "401 未授权、403 禁止访问、404 资源不存在。"},
    {"id": "kb-003", "text": "要查询账户余额，请调用 /v1/dashboard/billing/subscription 接口。"},
]

for hit in search("API 请求一直返回 429 错误，怎么解决？", corpus, recall_k=3, final_k=2):
    print(f"{hit['score']:.4f}  [{hit['id']}]  {hit['text'][:40]}")

def chunk(text, size=400, overlap=60):
    step = size - overlap
    return [text[i:i + size] for i in range(0, max(len(text) - overlap, 1), step)]


def rerank_long_docs(query, docs, top_n=5):
    """docs: [{'id':..., 'text':...}]，按块打分后按文档聚合取最高分。"""
    flat, owner = [], []
    for d in docs:
        for c in chunk(d["text"]):
            flat.append(c)
            owner.append(d)

    best = {}
    for r in rerank(query, flat, top_n=len(flat)):
        d = owner[r["index"]]
        if r["relevance_score"] > best.get(d["id"], (0, None))[0]:
            best[d["id"]] = (r["relevance_score"], flat[r["index"]], d)

    ranked = sorted(best.values(), key=lambda x: -x[0])[:top_n]
    return [{**d, "score": s, "best_chunk": c} for s, c, d in ranked]

def filter_by_relative_threshold(results, alpha=0.25, max_n=5):
    if not results:
        return []
    top = results[0]["relevance_score"]
    return [r for r in results if r["relevance_score"] >= top * alpha][:max_n]

from openai import OpenAI

client = OpenAI(api_key=os.environ["APIYI_API_KEY"],
                base_url="https://api.apiyi.com/v1")

NEG_HINT = ("不", "别", "除了", "无需", "禁止", "except", "not ", "without")


def verify_if_negated(query, hits):
    """query 含否定语义时，逐条让小模型确认是否真的满足条件。"""
    if not any(k in query for k in NEG_HINT):
        return hits
    kept = []
    for h in hits:
        resp = client.chat.completions.create(
            model="gemini-3.5-flash-lite",
            messages=[{"role": "user", "content":
                       f"问题：{query}\n文档：{h['text']}\n"
                       f"这份文档是否**正面满足**问题的条件？只回答 是 或 否。"}],
            max_tokens=4,
        )
        if "是" in resp.choices[0].message.content:
            kept.append(h)
    return kept

#1–#7   [200]  累计 tokens 递增到 16,093
#8      [429]  ← 首次拒绝（配额 20,000）
#9–#16  [429]  连续 9 个被拒
#17–#20 [200]  ← 窗口滑动，恢复