実践における 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):
    """Batch embed. In production, document vectors live in a vector DB."""
    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):
    """Rerank. documents must be a plain string array."""
    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':...}, ...]"""
    # ---- Stage 1: vector recall ----
    doc_vecs = embed([c["text"] for c in corpus])          # from your vector DB in production
    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]

    # ---- Stage 2: rerank ----
    results = rerank(query, [c["text"] for c in candidates], top_n=final_k)

    # Key step: map back by index so you keep ids and metadata
    return [{**candidates[r["index"]], "score": r["relevance_score"]}
            for r in results]


corpus = [
    {"id": "kb-001", "text": "A 429 means you hit a rate limit. Reduce concurrency, implement "
                             "exponential backoff on the client, or ask for a higher RPM quota."},
    {"id": "kb-002", "text": "HTTP 4xx status codes indicate client errors: 400 bad parameters, "
                             "401 unauthorized, 403 forbidden, 404 not found."},
    {"id": "kb-003", "text": "To check your balance, call /v1/dashboard/billing/subscription."},
]

for hit in search("My API requests keep returning 429 — how do I fix it?",
                  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':...}] — score by chunk, aggregate to best chunk per document."""
    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 = ("not ", "n't", "without", "except", "exclude", "free of", "non-")


def verify_if_negated(query, hits):
    """When the query carries negation, confirm each hit actually satisfies it."""
    q = query.lower()
    if not any(k in q 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"Question: {query}\nDocument: {h['text']}\n"
                       f"Does this document positively satisfy the question's condition? "
                       f"Answer only YES or NO."}],
            max_tokens=4,
        )
        if "YES" in resp.choices[0].message.content.upper():
            kept.append(h)
    return kept

#1–#7   [200]  cumulative tokens climbing to 16,093
#8      [429]  ← first rejection (quota 20,000)
#9–#16  [429]  nine consecutive failures
#17–#20 [200]  ← window slides, recovers