検索の精度は二段目で決まる — 埋め込みで粗く集め、Claude で並べ直す個人ナレッジ検索

from dataclasses import dataclass
 
@dataclass
class Candidate:
    id: str
    text: str
    score: float  # 一段目の類似度（大きいほど近い）
 
def recall_candidates(query: str, top_n: int = 20) -> list[Candidate]:
    """埋め込み検索で粗く top_n を集める。
    埋め込みモデルや保存先（pgvector 等）は問わない前提のインターフェース。"""
    vec = embed(query)                          # あなたの埋め込み関数
    rows = vector_store.search(vec, k=top_n)    # cosine 上位 top_n
    return [Candidate(id=r.id, text=r.text, score=r.similarity) for r in rows]

import os
from anthropic import Anthropic
 
client = Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
 
RERANK_TOOL = {
    "name": "rank_candidates",
    "description": "各候補が質問の答えとしてどれだけ的確かを 0.0〜1.0 で採点する",
    "input_schema": {
        "type": "object",
        "properties": {
            "rankings": {
                "type": "array",
                "items": {
                    "type": "object",
                    "properties": {
                        "id": {"type": "string"},
                        "relevance": {"type": "number"},  # 0.0-1.0
                        "reason": {"type": "string"},      # なぜその点か（根拠）
                    },
                    "required": ["id", "relevance", "reason"],
                },
            },
        },
        "required": ["rankings"],
    },
}
 
def rerank(query: str, cands: list[Candidate], model: str) -> list[dict]:
    # 候補を1リクエストに束ねる。1件ずつ呼ぶと採点基準がぶれ、往復コストも嵩む。
    passages = "\n\n".join(f"[{c.id}]\n{c.text}" for c in cands)
    msg = client.messages.create(
        model=model,
        max_tokens=1024,
        tools=[RERANK_TOOL],
        tool_choice={"type": "tool", "name": "rank_candidates"},
        messages=[{
            "role": "user",
            "content": (
                f"質問: {query}\n\n"
                "次の候補を、質問の答えとしての的確さで採点してください。"
                "言い換えが似ているかではなく、質問に直接答えているかを見てください。\n\n"
                f"{passages}"
            ),
        }],
    )
    block = next(b for b in msg.content if b.type == "tool_use")
    ranked = block.input["rankings"]
    ranked.sort(key=lambda r: r["relevance"], reverse=True)
    return ranked

# Before: 埋め込みの上位3件をそのまま返す（話題は近いが答えでないことがある）
def search_naive(query: str) -> list[Candidate]:
    return recall_candidates(query, top_n=3)
 
# After: 広く集めてから Claude で答えやすさで並べ直す
def search_two_stage(query: str, model: str) -> list[dict]:
    cands = recall_candidates(query, top_n=20)
    return rerank(query, cands, model=model)

ABSTAIN = 0.55  # 最上位がこの値未満なら「答えなし」を返す
 
def search(query: str, top_k: int = 3,
           model: str = "claude-haiku-4-5-20251001") -> dict:
    cands = recall_candidates(query, top_n=20)
    if not cands:
        return {"sources": [], "reason": "候補が見つかりませんでした"}
    ranked = rerank(query, cands, model=model)
    if ranked[0]["relevance"] < ABSTAIN:
        # 意味は近くても答えになっていない。無理に上位を返さない。
        return {"sources": [], "reason": "確信の持てる候補がありませんでした"}
    by_id = {c.id: c for c in cands}
    picked = [r for r in ranked[:top_k] if r["relevance"] >= ABSTAIN]
    return {
        "sources": [
            {"id": r["id"], "text": by_id[r["id"]].text,
             "relevance": r["relevance"], "reason": r["reason"]}
            for r in picked
        ],
    }