RAGアーキテクチャ：インデックス作成と検索

RAGには2つの異なるフェーズがある

インデックス作成フェーズ：ステップ1 — Load

from pypdf import PdfReader

def load_pdf(path):
    reader = PdfReader(path)
    pages = []
    for i, page in enumerate(reader.pages):
        text = page.extract_text()
        pages.append({'text': text, 'page': i + 1, 'source': path})
    return pages

docs = load_pdf('company_policy.pdf')
print(f'Loaded {len(docs)} pages')

インデックス作成フェーズ：ステップ2 — Chunk

from langchain_text_splitters import RecursiveCharacterTextSplitter

splitter = RecursiveCharacterTextSplitter(
    chunk_size=500,      # characters per chunk
    chunk_overlap=50,    # overlap between chunks
    separators=['\n\n', '\n', '. ', ' ']
)

for page in docs:
    chunks = splitter.split_text(page['text'])
    for chunk in chunks:
        # Each chunk carries metadata from its source page
        print(f'Chunk ({len(chunk)} chars): {chunk[:80]}...')

インデックス作成フェーズ：ステップ3 — Embed

from openai import OpenAI

client = OpenAI()

def embed_chunks(chunks):
    texts = [c['text'] for c in chunks]
    response = client.embeddings.create(
        model='text-embedding-3-small',
        input=texts
    )
    for i, chunk in enumerate(chunks):
        chunk['embedding'] = response.data[i].embedding
    return chunks

# Batch up to 2048 texts per API call
embedded = embed_chunks(all_chunks)

インデックス作成フェーズ：ステップ4 — Store

import pinecone

pc = pinecone.Pinecone(api_key='YOUR_KEY')
index = pc.Index('rag-documents')

# Upsert vectors with metadata
vectors_to_upsert = [
    (
        chunk['id'],
        chunk['embedding'],
        {'text': chunk['text'], 'source': chunk['source'], 'page': chunk['page']}
    )
    for chunk in embedded_chunks
]

# Upsert in batches of 100
for i in range(0, len(vectors_to_upsert), 100):
    index.upsert(vectors=vectors_to_upsert[i:i+100])
print('Indexing complete')

検索フェーズ：クエリのembedding化

def embed_query(question):
    response = client.embeddings.create(
        model='text-embedding-3-small',  # MUST match indexing model
        input=[question]
    )
    return response.data[0].embedding

user_question = 'What is our parental leave policy?'
query_vector = embed_query(user_question)
print(f'Query embedded: {len(query_vector)}-dim vector')

検索フェーズ：ANN検索

results = index.query(
    vector=query_vector,
    top_k=5,
    include_metadata=True
)

print(f'Retrieved {len(results.matches)} chunks:')
for match in results.matches:
    print(f'  Score: {match.score:.3f} | Source: {match.metadata["source"]}')
    print(f'  Text: {match.metadata["text"][:100]}...')
    print()

2つのフェーズをつなぐ

検索時のメタデータフィルタリング

# Retrieve only from HR department documents
results = index.query(
    vector=query_vector,
    top_k=5,
    filter={'department': {'$eq': 'HR'}},
    include_metadata=True
)

# Or filter by date range
results = index.query(
    vector=query_vector,
    top_k=5,
    filter={
        'upload_year': {'$gte': 2024},
        'doc_type': {'$eq': 'policy'}
    },
    include_metadata=True
)

更新に対応する増分インデックス作成

import hashlib

def make_chunk_id(source_path, chunk_index):
    # Deterministic, stable ID for each chunk
    key = f'{source_path}::chunk_{chunk_index}'
    return hashlib.md5(key.encode()).hexdigest()

def update_document(source_path, index):
    # Delete old vectors for this document
    index.delete(filter={'source': source_path})
    # Re-index the updated document
    new_chunks = load_and_chunk(source_path)
    new_embedded = embed_chunks(new_chunks)
    index.upsert(vectors=new_embedded)
    print(f'Updated {source_path}: {len(new_chunks)} chunks')

RAGパイプライン全体の概要

理解度チェック

レッスンのまとめ