Antigravity で構築する RAG パイプライン — ベクトル検索 × LLM で社内ナレッジを即座に活用する

# semantic_chunker.py — 意味的なまとまりを保つ高度なチャンク分割
 
from langchain.text_splitter import RecursiveCharacterTextSplitter  # 後半の親子分割で使用
from langchain_experimental.text_splitter import SemanticChunker
from langchain_google_genai import GoogleGenerativeAIEmbeddings
import os
 
def create_semantic_chunks(documents: list) -> list:
    """Embedding ベースのセマンティックチャンキング"""
    embeddings = GoogleGenerativeAIEmbeddings(
        model="gemini-embedding-001",
        google_api_key=os.getenv("GOOGLE_API_KEY"),
        output_dimensionality=1536,
    )
 
    # 文の意味的な距離が閾値を超えたところで分割
    chunker = SemanticChunker(
        embeddings,
        breakpoint_threshold_type="percentile",
        breakpoint_threshold_amount=90,  # 上位10%の距離で分割
    )
 
    chunks = []
    for doc in documents:
        split = chunker.create_documents(
            [doc.page_content],
            metadatas=[doc.metadata],
        )
        chunks.extend(split)
 
    return chunks
 
# --- 親子チャンク戦略（Parent-Child Chunking）---
class ParentChildChunker:
    """
    大きな「親チャンク」で文脈を保持し、
    小さな「子チャンク」で検索精度を上げる二段構え
    """
 
    def __init__(self, parent_size=2000, child_size=400, overlap=100):
        self.parent_splitter = RecursiveCharacterTextSplitter(
            chunk_size=parent_size,
            chunk_overlap=0,
        )
        self.child_splitter = RecursiveCharacterTextSplitter(
            chunk_size=child_size,
            chunk_overlap=overlap,
        )
 
    def split(self, documents: list) -> tuple:
        """親チャンクと子チャンクを同時に生成"""
        parent_chunks = []
        child_chunks = []
 
        for doc in documents:
            parents = self.parent_splitter.split_documents([doc])
            for idx, parent in enumerate(parents):
                parent.metadata["parent_id"] = f"{doc.metadata.get('source', 'unknown')}_{idx}"
                parent_chunks.append(parent)
 
                # 親チャンクをさらに小さく分割
                children = self.child_splitter.split_documents([parent])
                for child in children:
                    child.metadata["parent_id"] = parent.metadata["parent_id"]
                    child_chunks.append(child)
 
        return parent_chunks, child_chunks
 
# 期待される動作:
# - 子チャンクで精密な検索を行い
# - ヒットした子チャンクの親チャンクをコンテキストとして LLM に渡す
# → 検索精度と文脈の豊かさを両立

# hybrid_retriever.py — ベクトル検索 + BM25 のハイブリッド検索
 
from langchain.retrievers import EnsembleRetriever
from langchain_community.retrievers import BM25Retriever
from langchain_community.vectorstores import Chroma
from langchain_google_genai import GoogleGenerativeAIEmbeddings
import os
 
def create_hybrid_retriever(chunks: list, k: int = 5):
    """ベクトル検索と BM25 を組み合わせたハイブリッドリトリーバー"""
 
    # ここでは task_type をあえて指定しません。
    # ライブラリ側が embed_documents に RETRIEVAL_DOCUMENT、
    # embed_query に RETRIEVAL_QUERY を自動で割り当てるためです（後述）。
    embeddings = GoogleGenerativeAIEmbeddings(
        model="gemini-embedding-001",
        google_api_key=os.getenv("GOOGLE_API_KEY"),
        output_dimensionality=1536,
    )
 
    # 1. ベクトル検索リトリーバー
    vectorstore = Chroma.from_documents(
        documents=chunks,
        embedding=embeddings,
        collection_name="hybrid_search",
    )
    vector_retriever = vectorstore.as_retriever(
        search_type="mmr",          # Maximal Marginal Relevance
        search_kwargs={"k": k, "fetch_k": k * 3},
    )
 
    # 2. BM25 キーワード検索リトリーバー
    bm25_retriever = BM25Retriever.from_documents(
        chunks,
        k=k,
    )
 
    # 3. アンサンブル（重み付け統合）
    ensemble = EnsembleRetriever(
        retrievers=[vector_retriever, bm25_retriever],
        weights=[0.6, 0.4],  # ベクトル検索を少し重視
    )
 
    return ensemble
 
# --- リランキング（Cohere Reranker を使用）---
from langchain.retrievers import ContextualCompressionRetriever
from langchain_cohere import CohereRerank
 
def add_reranking(base_retriever, top_n: int = 3):
    """検索結果をリランカーで再スコアリング"""
    reranker = CohereRerank(
        model="rerank-v3.5",
        cohere_api_key=os.getenv("COHERE_API_KEY"),
        top_n=top_n,
    )
 
    return ContextualCompressionRetriever(
        base_compressor=reranker,
        base_retriever=base_retriever,
    )
 
# 使用例:
# retriever = create_hybrid_retriever(chunks, k=10)
# reranked_retriever = add_reranking(retriever, top_n=3)
# results = reranked_retriever.invoke("Stripe Webhook の署名検証方法は？")
#
# 期待される動作:
# → BM25 が「Stripe」「Webhook」「署名検証」のキーワードマッチで候補を取得
# → ベクトル検索が意味的に近い文書も取得
# → リランカーが上位3件に絞り込み、最も関連性の高い文書を返す

# rag_chain.py — RAG チェーンの構築
 
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
 
# 生成モデルは1箇所に集約しておく。廃止告知が出たときの作業がここだけで済む
CHAT_MODEL = "gemini-2.5-pro"   # 2026年10月16日以降の廃止が予告済み
 
def build_rag_chain(retriever):
    """検索結果を元に回答を生成する RAG チェーン"""
 
    # --- プロンプトテンプレート ---
    template = ChatPromptTemplate.from_messages([
        ("system", """あなたは社内ナレッジに基づいて正確に回答するアシスタントです。
 
以下のルールを厳守してください:
1. 提供されたコンテキストの情報のみに基づいて回答する
2. コンテキストに情報がない場合は「この情報は社内ナレッジベースに見つかりませんでした」と明示する
3. 推測や外部知識を混ぜない
4. 回答の根拠となったドキュメントのソースを明記する
5. 技術的な内容にはコード例を含める
 
コンテキスト:
{context}"""),
        ("human", "{question}"),
    ])
 
    # --- LLM ---
    llm = ChatGoogleGenerativeAI(
        model=CHAT_MODEL,
        temperature=0.1,        # 事実に基づく回答のため低めに設定
        max_output_tokens=4096,
    )
 
    # --- チェーン構築 ---
    def format_docs(docs):
        """検索結果を整形してコンテキスト文字列にする"""
        formatted = []
        for i, doc in enumerate(docs, 1):
            source = doc.metadata.get("source", "不明")
            formatted.append(
                f"[ソース {i}: {source}]\n{doc.page_content}"
            )
        return "\n\n---\n\n".join(formatted)
 
    chain = (
        {
            "context": retriever | format_docs,
            "question": RunnablePassthrough(),
        }
        | template
        | llm
        | StrOutputParser()
    )
 
    return chain
 
# --- ストリーミング対応版 ---
def query_with_streaming(chain, question: str):
    """ストリーミングで回答を表示"""
    print(f"\n📝 質問: {question}\n")
    print("=" * 60)
    for chunk in chain.stream(question):
        print(chunk, end="", flush=True)
    print("\n" + "=" * 60)
 
# 使用例:
# chain = build_rag_chain(reranked_retriever)
# query_with_streaming(chain, "新規プロジェクトのセットアップ手順を教えてください")
#
# 期待される出力:
# 📝 質問: 新規プロジェクトのセットアップ手順を教えてください
# ============================================================
# 社内ナレッジベースに基づくと、新規プロジェクトのセットアップは
# 以下の手順で行います:
# 1. GitHub でリポジトリを作成し...
# [ソース: docs/setup-guide.md]
# ============================================================

# rag_evaluator.py — RAG パイプラインの評価フレームワーク
 
from dataclasses import dataclass
from langchain_google_genai import ChatGoogleGenerativeAI
from pydantic import BaseModel, Field
 
@dataclass
class EvalResult:
    question: str
    expected: str
    actual: str
    faithfulness: float    # コンテキストに忠実か（0-1）
    relevance: float       # 質問に関連しているか（0-1）
 
class Verdict(BaseModel):
    """判定モデルに返させるスキーマ"""
    score: float = Field(description="0.0 から 1.0 のスコア")
    reason: str = Field(description="そのスコアにした理由")
 
class RAGEvaluator:
    """LLM-as-a-Judge パターンで RAG 回答を評価"""
 
    def __init__(self):
        # 素の応答を json.loads すると、```json フェンス付きで返ってきた瞬間に落ちます。
        # 構造化出力に寄せれば、パース処理そのものが不要になります。
        self.judge = ChatGoogleGenerativeAI(
            model="gemini-2.5-pro",
            temperature=0,
        ).with_structured_output(Verdict)
 
    def evaluate_faithfulness(
        self, context: str, answer: str
    ) -> float:
        """回答がコンテキストに忠実かを評価"""
        prompt = f"""以下の回答が、提供されたコンテキストの情報のみに基づいているか評価してください。
 
コンテキスト:
{context}
 
回答:
{answer}
 
0.0（完全にコンテキスト外）〜 1.0（完全にコンテキストに忠実）で採点してください。"""
 
        return self.judge.invoke(prompt).score
 
    def evaluate_relevance(
        self, question: str, answer: str
    ) -> float:
        """回答が質問に関連しているかを評価"""
        prompt = f"""以下の回答が質問に対して適切に回答しているか評価してください。
 
質問: {question}
回答: {answer}
 
0.0（無関係）〜 1.0（完全に的確）で採点してください。"""
 
        return self.judge.invoke(prompt).score
 
    def run_eval_suite(
        self, chain, retriever, test_cases: list[dict]
    ) -> list[EvalResult]:
        """テストケース一覧で評価を実行"""
        results = []
        for case in test_cases:
            # 検索結果を取得
            docs = retriever.invoke(case["question"])
            context = "\n".join(d.page_content for d in docs)
 
            # 回答を生成
            answer = chain.invoke(case["question"])
 
            # 評価
            faithfulness = self.evaluate_faithfulness(context, answer)
            relevance = self.evaluate_relevance(case["question"], answer)
 
            results.append(EvalResult(
                question=case["question"],
                expected=case.get("expected", ""),
                actual=answer,
                faithfulness=faithfulness,
                relevance=relevance,
            ))
 
        # サマリー出力
        avg_faith = sum(r.faithfulness for r in results) / len(results)
        avg_rel = sum(r.relevance for r in results) / len(results)
        print(f"\n📊 評価結果サマリー（{len(results)} ケース）")
        print(f"  忠実性（Faithfulness）: {avg_faith:.2f}")
        print(f"  関連性（Relevance）:    {avg_rel:.2f}")
 
        return results
 
# テストケース例:
# test_cases = [
#     {"question": "デプロイ手順は？", "expected": "wrangler deploy を実行..."},
#     {"question": "API キーの発行方法は？", "expected": "管理画面から..."},
# ]
# evaluator = RAGEvaluator()
# results = evaluator.run_eval_suite(chain, retriever, test_cases)