Gemini × LlamaIndex で構築する RAG エージェント — ドキュメント検索からマルチステップ推論まで

from llama_index.core import Settings
from llama_index.llms.gemini import Gemini
from llama_index.embeddings.gemini import GeminiEmbedding
 
# LLM の設定（Gemini 3.1 Pro を使用）
Settings.llm = Gemini(
    model="models/gemini-3.1-pro",
    temperature=0.3,
    max_tokens=4096,
)
 
# エンベディングモデルの設定
Settings.embed_model = GeminiEmbedding(
    model_name="models/text-embedding-004",
)
 
print("Gemini LLM and embedding model configured successfully")
# 出力: Gemini LLM and embedding model configured successfully

from llama_index.core import SimpleDirectoryReader
 
# ディレクトリ内のドキュメントを一括読み込み
documents = SimpleDirectoryReader(
    input_dir="./docs",
    recursive=True,
    required_exts=[".txt", ".pdf", ".md"],
).load_data()
 
print(f"Loaded {len(documents)} documents")
# 出力例: Loaded 42 documents

from llama_index.core import VectorStoreIndex
 
# インデックスの構築（チャンク分割 → エンベディング生成 → インデックス化）
index = VectorStoreIndex.from_documents(
    documents,
    show_progress=True,
)
 
print("Vector index built successfully")
# 出力: Vector index built successfully

# クエリエンジンの作成
query_engine = index.as_query_engine(
    similarity_top_k=5,        # 上位5件の類似チャンクを取得
    response_mode="compact",   # コンパクトな回答生成
)
 
# 質問を投げる
response = query_engine.query(
    "プロジェクトのデプロイ手順を教えてください"
)
 
print(response)
# 出力例: デプロイ手順は以下の通りです。まず...（ドキュメントの内容に基づいた回答）
 
# ソースドキュメントの確認
for node in response.source_nodes:
    print(f"  Source: {node.metadata.get('file_name', 'unknown')} "
          f"(score: {node.score:.3f})")
# 出力例:
#   Source: deployment-guide.md (score: 0.892)
#   Source: ci-cd-pipeline.txt (score: 0.847)

from llama_index.core.node_parser import SentenceSplitter
 
Settings.text_splitter = SentenceSplitter(
    chunk_size=512,
    chunk_overlap=64,
    paragraph_separator="\n\n",
)
 
# 設定を変えたら必ずインデックスを作り直す
index = VectorStoreIndex.from_documents(documents, show_progress=True)

from llama_index.core.agent import ReActAgent
from llama_index.core.tools import QueryEngineTool, ToolMetadata
 
# 複数のインデックスからツールを作成
api_docs_tool = QueryEngineTool(
    query_engine=api_index.as_query_engine(similarity_top_k=3),
    metadata=ToolMetadata(
        name="api_docs",
        description="API のリファレンスドキュメントを検索します。"
                    "エンドポイント、パラメータ、レスポンス形式について回答できます。",
    ),
)
 
tutorial_tool = QueryEngineTool(
    query_engine=tutorial_index.as_query_engine(similarity_top_k=3),
    metadata=ToolMetadata(
        name="tutorials",
        description="チュートリアルとハウツーガイドを検索します。"
                    "実装手順やベストプラクティスについて回答できます。",
    ),
)
 
# ReAct エージェントの構築
agent = ReActAgent.from_tools(
    tools=[api_docs_tool, tutorial_tool],
    llm=Settings.llm,
    verbose=True,
    max_iterations=5,
)
 
# エージェントに質問
response = agent.chat(
    "ユーザー認証APIの仕様を確認して、"
    "それを使ったログイン機能の実装手順を教えてください"
)
print(response)
# 出力例:
# Thought: まずAPI仕様を確認し、次にチュートリアルで実装手順を調べます
# Action: api_docs → 認証エンドポイントの仕様を取得
# Action: tutorials → ログイン実装のチュートリアルを取得
# Answer: ユーザー認証APIは POST /api/auth/login で...

from llama_index.core.retrievers import QueryFusionRetriever
from llama_index.retrievers.bm25 import BM25Retriever
from llama_index.core.query_engine import RetrieverQueryEngine
 
# 1) ベクトル検索リトリーバー
vector_retriever = index.as_retriever(similarity_top_k=5)
 
# 2) BM25（キーワード）リトリーバー
bm25_retriever = BM25Retriever.from_defaults(
    docstore=index.docstore,
    similarity_top_k=5,
)
 
# 3) 2つを融合。相互ランク融合（RRF）でスコアを統合する
fusion_retriever = QueryFusionRetriever(
    retrievers=[vector_retriever, bm25_retriever],
    similarity_top_k=5,
    num_queries=1,            # クエリ書き換えを使わない場合は 1
    mode="reciprocal_rerank", # RRF による統合
    use_async=True,
)
 
hybrid_engine = RetrieverQueryEngine.from_args(fusion_retriever)
 
response = hybrid_engine.query("RFC 7231 のキャッシュ制御の扱いは？")
print(response)

from llama_index.core.postprocessor import LLMRerank
from llama_index.core.query_engine import RetrieverQueryEngine
 
reranker = LLMRerank(
    choice_batch_size=5,
    top_n=5,          # 並べ直したあと、最終的に生成へ渡す件数
)
 
# 検索は広めに 20 件、生成に渡すのは 5 件だけ
rerank_engine = RetrieverQueryEngine.from_args(
    index.as_retriever(similarity_top_k=20),
    node_postprocessors=[reranker],
)
 
response = rerank_engine.query("トークンが失効したあとの再発行手順は？")
for node in response.source_nodes:
    print(f"{node.metadata.get('file_name')}: {node.score:.3f}")

# (質問, 正解を含むはずのファイル名) のペアを用意する
eval_set = [
    ("デプロイ手順を教えて", "deployment-guide.md"),
    ("認証トークンの有効期限は？", "auth-spec.md"),
    # ... 20件ほど
]
 
def recall_at_k(engine, eval_set, k=5):
    hit = 0
    for question, expected_file in eval_set:
        nodes = engine.retrieve(question)[:k]
        files = {n.metadata.get("file_name") for n in nodes}
        if expected_file in files:
            hit += 1
    return hit / len(eval_set)
 
score = recall_at_k(index.as_retriever(similarity_top_k=5), eval_set)
print(f"Recall@5 = {score:.2f}")
# 出力例: Recall@5 = 0.86

import os
 
PERSIST_DIR = "./storage"
 
# インデックスの保存
index.storage_context.persist(persist_dir=PERSIST_DIR)
print(f"Index saved to {PERSIST_DIR}")
# 出力: Index saved to ./storage
 
# インデックスの読み込み（2回目以降）
from llama_index.core import StorageContext, load_index_from_storage
 
if os.path.exists(PERSIST_DIR):
    storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)
    index = load_index_from_storage(storage_context)
    print("Index loaded from disk")
    # 出力: Index loaded from disk

from llama_index.core import SimpleDirectoryReader
 
# ファイルパスを doc_id に固定する（読み込みのたびに同じ ID になる）
documents = SimpleDirectoryReader(
    "./data",
    filename_as_id=True,
).load_data()
 
# 新規・更新されたものだけを反映する
refreshed = index.refresh_ref_docs(documents)
print(f"再埋め込み対象: {sum(refreshed)} / {len(refreshed)} 件")
# 出力例: 再埋め込み対象: 3 / 128 件
 
index.storage_context.persist(persist_dir=PERSIST_DIR)

existing_ids = set(index.ref_doc_info.keys())
current_ids = {d.doc_id for d in documents}
 
for stale_id in existing_ids - current_ids:
    index.delete_ref_doc(stale_id, delete_from_docstore=True)
 
index.storage_context.persist(persist_dir=PERSIST_DIR)

from llama_index.core.vector_stores import MetadataFilters, ExactMatchFilter
 
filters = MetadataFilters(
    filters=[ExactMatchFilter(key="department", value="engineering")]
)
query_engine = index.as_query_engine(similarity_top_k=5, filters=filters)