FAISSまたはChromaで素朴なRAGを構築する

40行のRAG

Step 1: Install

# pip install openai chromadb

Step 2: Sample Corpus

docs = [
    'Python is a high-level interpreted programming language created by Guido van Rossum in 1991.',
    'Pizza is a savory dish of Italian origin consisting of a round flat base of dough.',
    'The Eiffel Tower is a wrought-iron lattice tower in Paris, France, completed in 1889.',
    'NumPy is a Python library used for working with arrays and numerical computing.',
    'The Great Wall of China is a series of fortifications built across northern China.'
]
print(f"Loaded {len(docs)} documents")
for i, d in enumerate(docs):
    print(f"{i+1}. {d[:50]}...")

Step 3: Set Up Chroma

import chromadb
from openai import OpenAI

client = chromadb.Client()
collection = client.create_collection('demo')
openai_client = OpenAI()

Step 4: Embed and Store

embed_response = openai_client.embeddings.create(
    model='text-embedding-3-small',
    input=docs
)
vectors = [d.embedding for d in embed_response.data]

collection.add(
    ids=[f'd{i}' for i in range(len(docs))],
    embeddings=vectors,
    documents=docs
)

Step 5: Embed the Query

query = 'Tell me about Python programming'
query_vec = openai_client.embeddings.create(
    model='text-embedding-3-small',
    input=query
).data[0].embedding

Step 6: Retrieve Top-K

results = collection.query(
    query_embeddings=[query_vec],
    n_results=2
)
retrieved = results['documents'][0]
print(retrieved)
# ['Python is a high-level...', 'NumPy is a Python library...']

Step 7: Build the RAG Prompt

context = '\n'.join(f'- {chunk}' for chunk in retrieved)
prompt = f'''
Using only the context below, answer the question. If unknown, say so.

Context:
{context}

Question: {query}
'''

Step 8: Call the LLM

response = openai_client.chat.completions.create(
    model='gpt-4o-mini',
    messages=[{'role': 'user', 'content': prompt}],
    temperature=0,
)
print(response.choices[0].message.content)

FAISSという選択肢

import faiss
import numpy as np

index = faiss.IndexFlatIP(1536)  # inner product, normalised vectors
index.add(np.array(vectors).astype('float32'))

D, I = index.search(np.array([query_vec]).astype('float32'), k=2)
retrieved = [docs[i] for i in I[0]]

この単純なRAGの制限

クエリあたりのコスト

リリースして反復改善する

素朴なRAGのTop-K

まとめ