Python: Minimal RAG pipeline in ~80 lines — FAISS + sentence-transformers + OpenAI, no framework

minimal_rag.py

"""
Minimal RAG pipeline in ~80 lines — no framework, no vector DB service.
Embeds docs, stores them in a local FAISS index, retrieves top-k, and
hands them to an LLM as context.

Install:
    pip install openai faiss-cpu sentence-transformers numpy

Env:
    export OPENAI_API_KEY=sk-...
"""

import os
import numpy as np
import faiss
from sentence_transformers import SentenceTransformer
from openai import OpenAI

EMBED_MODEL = "all-MiniLM-L6-v2"   # 384-dim, fast, free, local
LLM_MODEL = "gpt-4o-mini"
TOP_K = 3

client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])
embedder = SentenceTransformer(EMBED_MODEL)


def build_index(docs: list[str]) -> tuple[faiss.Index, list[str]]:
    vectors = embedder.encode(docs, normalize_embeddings=True)
    index = faiss.IndexFlatIP(vectors.shape[1])  # cosine via inner product
    index.add(np.asarray(vectors, dtype="float32"))
    return index, docs


def retrieve(index: faiss.Index, docs: list[str], query: str, k: int = TOP_K) -> list[str]:
    q = embedder.encode([query], normalize_embeddings=True).astype("float32")
    _, idx = index.search(q, k)
    return [docs[i] for i in idx[0] if i != -1]


def answer(index, docs, question: str) -> str:
    context = retrieve(index, docs, question)
    prompt = (
        "Answer the question using ONLY the context below. "
        "If the context does not contain the answer, say you don't know.\n\n"
        f"Context:\n" + "\n---\n".join(context) + f"\n\nQuestion: {question}"
    )
    res = client.chat.completions.create(
        model=LLM_MODEL,
        max_tokens=512,
        messages=[{"role": "user", "content": prompt}],
    )
    return res.choices[0].message.content


if __name__ == "__main__":
    docs = [
        "The Eiffel Tower was completed in 1889 for the World's Fair.",
        "Mount Everest is 8,848.86 meters tall.",
        "Python was created by Guido van Rossum and first released in 1991.",
        "The Great Wall of China is over 13,000 miles long.",
        "FAISS is a library by Facebook AI for efficient similarity search.",
    ]
    index, docs = build_index(docs)
    print(answer(index, docs, "Who created Python and when?"))
    print("---")
    print(answer(index, docs, "How long is the Great Wall?"))