Gemini Embeddings × リランカーで本番RAGの精度を底上げする — Vertex AI Ranking と LLM-as-judge の使い分け

# rerank_with_vertex.py
# 役割: Embeddings で取得した候補を Vertex AI Ranking API でリランクする
# 期待出力: スコア付き上位N件のドキュメント
import os
from google.cloud import discoveryengine_v1
 
PROJECT_ID = os.environ["GOOGLE_CLOUD_PROJECT"]
LOCATION = "global"
RANKING_CONFIG = "default_ranking_config"
 
def rerank(query: str, candidates: list[dict], top_k: int = 10) -> list[dict]:
    """
    candidates: [{"id": "doc-1", "title": "...", "content": "..."}, ...]
    """
    client = discoveryengine_v1.RankServiceClient()
    parent = (
        f"projects/{PROJECT_ID}/locations/{LOCATION}/"
        f"rankingConfigs/{RANKING_CONFIG}"
    )
 
    records = [
        discoveryengine_v1.RankingRecord(
            id=c["id"],
            title=c.get("title", ""),
            content=c["content"][:4000],  # 1レコード上限に注意
        )
        for c in candidates
    ]
 
    request = discoveryengine_v1.RankRequest(
        ranking_config=parent,
        model="semantic-ranker-default@latest",
        top_n=top_k,
        query=query,
        records=records,
    )
    try:
        response = client.rank(request=request)
    except Exception as e:
        # リランカーが落ちても元の候補で動かし続ける
        print(f"[rerank] fallback to embeddings order: {e}")
        return candidates[:top_k]
 
    # スコア降順で並んでいる
    ranked = []
    by_id = {c["id"]: c for c in candidates}
    for r in response.records:
        doc = by_id[r.id]
        doc["rerank_score"] = r.score
        ranked.append(doc)
    return ranked

# pipeline.py
from rerank_with_vertex import rerank
from your_embeddings_search import search_with_gemini_embeddings
 
def retrieve(query: str) -> list[dict]:
    # Stage 1: 50件をリコール
    candidates = search_with_gemini_embeddings(query, top_k=50)
    # Stage 2: 上位10件にリランク
    top10 = rerank(query, candidates, top_k=10)
    return top10

# rerank_with_gemini.py
# 役割: Gemini 2.5 Flash をリランカーとして使い、候補を一括採点する
# 期待出力: スコア付き上位N件
import os
import json
import google.generativeai as genai
from pydantic import BaseModel
 
genai.configure(api_key=os.environ["GEMINI_API_KEY"])
 
RERANK_PROMPT = """あなたは検索エンジンの関連度評価者です。
以下のクエリに対して、各ドキュメントが「直接的な答え」になっているかを 0.0〜1.0 で採点してください。
 
採点基準:
- 1.0: クエリの問いに直接かつ明確に答えている
- 0.7: 答えに近いが部分的、または周辺情報を含む
- 0.4: 関連はあるが答えではない
- 0.0: 関連性が低い
 
【クエリ】
{query}
 
【ドキュメント一覧】
{documents}
 
JSON形式で `{{"scores": [{{"id": "doc-1", "score": 0.92, "reason": "簡潔な理由"}}]}}` のように返してください。
"""
 
class RankItem(BaseModel):
    id: str
    score: float
    reason: str
 
class RankResult(BaseModel):
    scores: list[RankItem]
 
def rerank_with_gemini(
    query: str, candidates: list[dict], top_k: int = 10
) -> list[dict]:
    docs_text = "\n\n".join(
        f"[{c['id']}]\n{c['content'][:1500]}"  # 1500字に切り詰めて Token 暴発を防ぐ
        for c in candidates
    )
    prompt = RERANK_PROMPT.format(query=query, documents=docs_text)
 
    model = genai.GenerativeModel(
        "gemini-2.5-flash",
        generation_config={
            "response_mime_type": "application/json",
            "response_schema": RankResult,
            "temperature": 0.0,  # 採点は決定論にする
        },
    )
 
    try:
        response = model.generate_content(prompt)
        result = RankResult.model_validate_json(response.text)
    except Exception as e:
        print(f"[rerank_gemini] fallback: {e}")
        return candidates[:top_k]
 
    # スコア降順で並べ替え
    score_by_id = {item.id: item.score for item in result.scores}
    ranked = sorted(
        candidates,
        key=lambda c: score_by_id.get(c["id"], 0.0),
        reverse=True,
    )[:top_k]
 
    for c in ranked:
        c["rerank_score"] = score_by_id.get(c["id"], 0.0)
    return ranked

# hybrid_pipeline.py
# 役割: BM25 + Embeddings の候補をマージし、Vertex Ranking API でリランク
from rank_bm25 import BM25Okapi
from your_embeddings_search import search_with_gemini_embeddings
from rerank_with_vertex import rerank
 
def hybrid_retrieve(
    query: str, corpus: list[dict], bm25: BM25Okapi, top_k: int = 10
) -> list[dict]:
    # ① BM25 で 30件
    bm25_scores = bm25.get_scores(query.split())
    bm25_top = sorted(
        zip(corpus, bm25_scores), key=lambda x: x[1], reverse=True
    )[:30]
    bm25_ids = {doc["id"] for doc, _ in bm25_top}
 
    # ② Embeddings で 30件
    emb_top = search_with_gemini_embeddings(query, top_k=30)
    emb_ids = {doc["id"] for doc in emb_top}
 
    # ③ 両者の和集合（重複は除去・最大60件）
    candidates_map = {doc["id"]: doc for doc, _ in bm25_top}
    for doc in emb_top:
        candidates_map.setdefault(doc["id"], doc)
    candidates = list(candidates_map.values())
 
    # ④ 共通している候補は信頼度を上げる（簡易RRF）
    overlap = bm25_ids & emb_ids
    for c in candidates:
        c["recall_signal"] = 2 if c["id"] in overlap else 1
 
    # ⑤ Vertex AI Ranking で精度を出す
    return rerank(query, candidates, top_k=top_k)

# eval_rerank.py
# 役割: 評価セットでMRRとNDCG@5を計算する
import math
from typing import Callable
 
def mrr(ranked_lists: list[list[str]], golds: list[str]) -> float:
    total = 0.0
    for ranked, gold in zip(ranked_lists, golds):
        try:
            rank = ranked.index(gold) + 1
            total += 1.0 / rank
        except ValueError:
            total += 0.0
    return total / len(ranked_lists)
 
def ndcg_at_k(ranked_lists: list[list[str]], golds: list[str], k: int = 5) -> float:
    total = 0.0
    for ranked, gold in zip(ranked_lists, golds):
        ideal = 1.0  # 単一正解前提
        dcg = 0.0
        for i, doc_id in enumerate(ranked[:k]):
            if doc_id == gold:
                dcg = 1.0 / math.log2(i + 2)
                break
        total += dcg / ideal
    return total / len(ranked_lists)
 
def evaluate(retrieve_fn: Callable, eval_set: list[dict]) -> dict:
    queries = [item["query"] for item in eval_set]
    golds = [item["gold_id"] for item in eval_set]
    ranked_lists = [
        [doc["id"] for doc in retrieve_fn(q)] for q in queries
    ]
    return {
        "mrr": mrr(ranked_lists, golds),
        "ndcg@5": ndcg_at_k(ranked_lists, golds, k=5),
        "n_queries": len(queries),
    }

# chunker.py
# 役割: 見出し階層を保持しながら適切な粒度でチャンク分割する
import re
 
def split_with_headings(markdown: str, max_chars: int = 600) -> list[dict]:
    chunks: list[dict] = []
    current_h1, current_h2, current_h3 = "", "", ""
    buffer = []
    buffer_len = 0
 
    def flush():
        nonlocal buffer, buffer_len
        if buffer_len > 0:
            chunks.append({
                "title": " > ".join(filter(None, [current_h1, current_h2, current_h3])),
                "content": "
".join(buffer).strip(),
            })
            buffer = []
            buffer_len = 0
 
    for line in markdown.split("
"):
        if line.startswith("# "):
            flush(); current_h1 = line[2:].strip(); current_h2 = current_h3 = ""
            continue
        if line.startswith("## "):
            flush(); current_h2 = line[3:].strip(); current_h3 = ""
            continue
        if line.startswith("### "):
            flush(); current_h3 = line[4:].strip()
            continue
        buffer.append(line)
        buffer_len += len(line)
        if buffer_len >= max_chars:
            flush()
    flush()
    return chunks

# rerank_cache.py
# 役割: ペア単位スコアキャッシュで Gemini リランカー呼び出しを削減
import hashlib
import json
from typing import Optional
import redis
 
r = redis.Redis(host="localhost", port=6379, decode_responses=True)
SCORE_TTL = 24 * 60 * 60  # 24時間
 
def _key(query: str, doc_id: str) -> str:
    h = hashlib.md5(f"{query}::{doc_id}".encode()).hexdigest()
    return f"rerank:score:{h}"
 
def get_cached_scores(query: str, doc_ids: list[str]) -> dict[str, float]:
    keys = [_key(query, d) for d in doc_ids]
    values = r.mget(keys)
    return {
        d: float(v) for d, v in zip(doc_ids, values) if v is not None
    }
 
def save_scores(query: str, scored: list[dict]) -> None:
    pipe = r.pipeline()
    for item in scored:
        pipe.set(_key(query, item["id"]), str(item["rerank_score"]), ex=SCORE_TTL)
    pipe.execute()
 
def rerank_with_cache(query: str, candidates: list[dict], rerank_fn) -> list[dict]:
    cached = get_cached_scores(query, [c["id"] for c in candidates])
    uncached = [c for c in candidates if c["id"] not in cached]
 
    # 未キャッシュ分だけ採点
    if uncached:
        new_scored = rerank_fn(query, uncached, top_k=len(uncached))
        save_scores(query, new_scored)
        for c in new_scored:
            cached[c["id"]] = c["rerank_score"]
 
    # キャッシュとマージしてスコア降順
    for c in candidates:
        c["rerank_score"] = cached.get(c["id"], 0.0)
    return sorted(candidates, key=lambda x: x["rerank_score"], reverse=True)