PythonでRAGパイプラインを自作する方法

# LlamaIndex本体とOpenAI連携用のパッケージ
pip install llama-index llama-index-llms-openai llama-index-embeddings-openai

# ベクトルデータベース（Qdrant）のクライアント
pip install qdrant-client llama-index-vector-stores-qdrant

# PDFの読み込みに必要なパッケージ
pip install pypdf

import os
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, StorageContext
from llama_index.core import Settings
from llama_index.llms.openai import OpenAI
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.vector_stores.qdrant import QdrantVectorStore
import qdrant_client

# 1. APIキーの設定
# 実際の業務では .env ファイルを使用しますが、まずは動作確認のために環境変数へセットします
os.environ["OPENAI_API_KEY"] = "sk-xxxxxxxxxxxxxxxx"

# 2. モデルの定義
# 回答用LLMにはコスパに優れた gpt-4o-mini を、
# 検索用モデル（Embedding）には text-embedding-3-small を指定します。
# なぜ small かと言うと、large にしても検索精度が劇的に変わるケースは稀で、
# トークン消費を抑える方が実務上のメリットが大きいからです。
Settings.llm = OpenAI(model="gpt-4o-mini", temperature=0.1)
Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small")

# 3. データの読み込み
# SimpleDirectoryReaderは指定したフォルダ内のファイルを自動判別して読み込みます。
# これが LlamaIndex の最も強力な「武器」の一つです。
documents = SimpleDirectoryReader("./data").load_data()

# 4. ベクトルDB（Qdrant）の初期化
# メモリ上で動作させる設定（:memory:）にします。
# 実運用時は、ここをURL（localhost等）に変更するだけで済みます。
client = qdrant_client.QdrantClient(location=":memory:")
vector_store = QdrantVectorStore(client=client, collection_name="my_rag_collection")
storage_context = StorageContext.from_defaults(vector_store=vector_store)

# 5. インデックス（検索エンジン）の作成
# ここでテキストがベクトル化され、DBに保存されます。
index = VectorStoreIndex.from_documents(
    documents,
    storage_context=storage_context,
)

# 6. クエリエンジンの作成
# similarity_top_k=3 は、「質問に関連する箇所を上位3つ探してくる」という意味です。
query_engine = index.as_query_engine(similarity_top_k=3)

# 7. 質問の実行
question = "読み込ませたドキュメントに基づいて、〇〇の注意点を要約してください"
response = query_engine.query(question)

print(f"【質問】: {question}")
print("-" * 30)
print(f"【回答】: {response}")

# 8. 根拠となったソースの確認
# 「なぜその回答をしたか」を追跡できるのがRAGのメリットです。
for node in response.source_nodes:
    print(f"\n[ソース（スコア: {node.score:.4f}）]:")
    print(node.node.get_content()[:100] + "...")

【質問】: 読み込ませたドキュメントに基づいて、〇〇の注意点を要約してください
------------------------------
【回答】: ドキュメントによると、〇〇に関する主な注意点は以下の3点です...

[ソース（スコア: 0.8542）]:
〇〇を使用する際は、必ず電源を切ってから作業を行ってください。また、湿気の多い場所での保管は...

from llama_index.core.node_parser import SentenceSplitter

# チャンクサイズの最適化
# デフォルトでは1024トークンですが、これだと情報が混ざりすぎることがあります。
# 512トークン程度に細分化し、前後の文脈（オーバーラップ）を50トークン持たせるのが
# 経験上、最もバランスが良いです。
node_parser = SentenceSplitter(chunk_size=512, chunk_overlap=50)

# インデックス作成時にこのパーサーを適用
index = VectorStoreIndex.from_documents(
    documents,
    storage_context=storage_context,
    transformations=[node_parser]
)

# 検索を強化する「ハイブリッド検索」や「リランク（再ランク付け）」
# 予算があるなら、CohereのRerank APIを挟むと劇的に精度が上がります。
# 今回は無料でできる工夫として、プロンプトに「根拠が見つからない場合は知らないと答えて」と明示します。
from llama_index.core import PromptTemplate

qa_prompt_tmpl_str = (
    "コンテキスト情報は以下の通りです。\n"
    "---------------------\n"
    "{context_str}\n"
    "---------------------\n"
    "この情報のみに基づいて、質問に答えてください。答えが見つからない場合は「該当する情報がありません」と回答してください。\n"
    "質問: {query_str}\n"
    "回答: "
)
qa_prompt_tmpl = PromptTemplate(qa_prompt_tmpl_str)
query_engine.update_prompts({"response_synthesizer:text_qa_template": qa_prompt_tmpl})