Un pipeline Mini RAG en 90 lignes de Python : récupération, augmentation, génération
Construire une génération améliorée par la récupération de données à partir de zéro en utilisant des embeddings, une similarité cosinus et une liste classée, puis observer où la version simplifiée échoue à grande échelle.
La génération augmentée par récupération n’est pas une fonctionnalité de tout modèle. C’est un problème d’emballage : la fenêtre de contexte est petite, les documents sont énormes, et il faut décider quels éléments y entreront. Cette solution peut être écrite en environ 90 lignes de Python pur, sans LangChain ni service de base de données vectorielle. Elle comporte trois étapes :
- Récupérer. Incarner la question, classer les fragments par similarité et conserver les k meilleurs.
- Augmenter. Intégrer ces élus dans le message du système.
- Générer. Utiliser la même requête de chat que dans n’importe quel script LLM basique.
Toute abstraction offerte par un framework comme LangChain correspond finalement à l’une de ces trois étapes ; donc en voir d’abord la version simplifiée rend le framework plus facile à comprendre.
Les incarneurs transforment le sens en géométrie
# pip install openai
# export OPENAI_API_KEY=sk-...
import math
from openai import OpenAI
client = OpenAI()
def embed(texts):
resp = client.embeddings.create(
model="text-embedding-3-small",
input=texts,
)
return [d.embedding for d in resp.data]
def cosine(a, b):
dot = sum(x * y for x, y in zip(a, b))
na = math.sqrt(sum(x * x for x in a))
nb = math.sqrt(sum(y * y for y in b))
return dot / (na * nb)
q, hit, miss = embed([
"How do I get my money back?",
"Refunds are possible within 14 days "
"of purchase.",
"Data is stored in Frankfurt and "
"encrypted at rest.",
])
print(f"refund doc: {cosine(q, hit):.3f}")
print(f"storage doc: {cosine(q, miss):.3f}")
La sortie typique de ce script est un score de 0,531 pour le document de remboursement et de 0,126 pour le document de stockage. Observez comment la correspondance s’effectue : la question indique « money back » tandis que le document mentionne « refunds ». Aucun mot-clé n’est commun, pourtant la similarité est élevée. C’est pourquoi on utilise l’incorporation plutôt que la recherche par motifs.
Un stockage vectoriel est une liste que l’on peut classer
« Stockage vectoriel » évoque une infrastructure. À petite échelle, il s’agit d’une liste de paires (vecteur, texte) accompagnée d’un tri. Pinecone, Chroma et pgvector fonctionnent bien avec des millions de vecteurs, mais ce qui compte, c’est l’interface, et celle-ci se résume en une seule catégorie :
class VectorStore:
def __init__(self):
self.texts = []
self.vectors = []
def add(self, texts):
self.texts.extend(texts)
self.vectors.extend(embed(texts))
def top_k(self, query, k=3):
q = embed([query])[0]
scored = [
(cosine(q, v), t)
for v, t in zip(
self.vectors, self.texts
)
]
scored.sort(
key=lambda s: s[0], reverse=True
)
return scored[:k]
store = VectorStore()
store.add([
"Refunds are possible within 14 days.",
"The Pro plan costs 12 EUR per month.",
"Data is stored in Frankfurt.",
])
for score, text in store.top_k(
"can I cancel and get a refund?", k=2
):
print(f"{score:.3f} {text}")
top_k représente la décision concernant les éléments à retenir, et elle est maintenant exécutable : classer tout le contenu, conserver k éléments et en exclure le reste. La valeur de k correspond au paramètre de contrôle, permettant d’équilibrer la capacité de rappel avec l’espace alloué à la fenêtre contextuelle et le nombre de tokens utilisés.
Intégrer la récupération, l’enrichissement et la génération
Une seule fonction est nouvelle ici, build_messages. Elle numérote les fragments sélectionnés, les insère dans le message système et indique au modèle de répondre uniquement à partir d’eux. Le message de l’utilisateur ne contient que la question. La liste des messages est identique à celle utilisée dans toute conversation, de sorte que RAG ne modifie que son contenu. Ce fichier fonctionne tel quel :
# examples/part07_rag.py
# Mini RAG in ~90 lines: embed -> store -> top-k
# -> stitch into messages -> generate.
#
# Setup:
# pip install openai
# export OPENAI_API_KEY=sk-...
#
# Run:
# python part07_rag.py
import math
from openai import OpenAI
client = OpenAI()
EMBED_MODEL = "text-embedding-3-small"
CHAT_MODEL = "gpt-4o-mini"
# A tiny corpus. In real life these are chunks
# of your docs, not whole documents.
DOCS = [
"Support hours are 9am to 6pm CET, "
"Monday to Friday.",
"Refunds are possible within 14 days "
"of purchase, no questions asked.",
"The Pro plan costs 12 EUR per month "
"and includes API access.",
"The Free plan allows 3 projects and "
"community support only.",
"Data is stored in Frankfurt and "
"encrypted at rest with AES-256.",
"You can export all your data as "
"JSON from the settings page.",
]
def embed(texts):
resp = client.embeddings.create(
model=EMBED_MODEL,
input=texts,
)
return [d.embedding for d in resp.data]
def cosine(a, b):
dot = sum(x * y for x, y in zip(a, b))
na = math.sqrt(sum(x * x for x in a))
nb = math.sqrt(sum(y * y for y in b))
return dot / (na * nb)
class VectorStore:
"""A list of (vector, text) pairs. That is
all a vector store is, before you need
scale."""
def __init__(self):
self.texts = []
self.vectors = []
def add(self, texts):
self.texts.extend(texts)
self.vectors.extend(embed(texts))
def top_k(self, query, k=3):
q = embed([query])[0]
scored = [
(cosine(q, v), t)
for v, t in zip(
self.vectors, self.texts
)
]
scored.sort(key=lambda s: s[0], reverse=True)
return scored[:k]
def build_messages(question, hits):
context = "\n\n".join(
f"[{i + 1}] {text}"
for i, (_, text) in enumerate(hits)
)
system = (
"Answer using only the context below. "
"If the answer is not in the context, "
"say you do not know.\n\n"
"Context:\n" + context
)
return [
{"role": "system", "content": system},
{"role": "user", "content": question},
]
def ask(store, question):
hits = store.top_k(question, k=3)
print(f"Q: {question}")
for score, text in hits:
print(f" {score:.3f} {text[:44]}...")
messages = build_messages(question, hits)
resp = client.chat.completions.create(
model=CHAT_MODEL,
messages=messages,
)
return resp.choices[0].message.content
if __name__ == "__main__":
store = VectorStore()
store.add(DOCS)
answer = ask(
store,
"Where is my data stored, and can "
"I get it out?",
)
print(f"A: {answer}")
Un essai de test demande où sont stockés les données de l’utilisateur et s’il est possible de les exporter. Les trois meilleurs résultats sont de 0,588 pour un extrait concernant le stockage à Francfort avec chiffrement, de 0,472 pour un extrait sur l’exportation de toutes les données au format JSON, et de 0,244 pour la politique de remboursement. Le modèle indique que les données se trouvent à Francfort, chiffrées en repos avec AES-256, et que la page des paramètres permet une exportation complète au format JSON.
Les deux extraits nécessaires, celui concernant le stockage et celui relatif à l’exportation, ont été retenus. L’extrait sur les remboursements a terminé à la troisième place et le modèle l’a ignoré. La récupération peut être imparfaite tant que le matériel adéquat se trouve en tête de liste.
Où la version simplifiée échoue
Ces 90 lignes présentent des limites réelles, et chacune souligne quelque chose que les systèmes complets ajoutent :
- Chunking. La démonstration stocke des phrases complètes. Les documents réels nécessitent une segmentation avec chevauchement, ce qui relève d’une décision de prétraitement et non du modèle.
- Re-embedding. Il est acceptable d’incorporer la requête à chaque appel. En revanche, il n’est pas judicieux de réincorporer l’ensemble du corpus à chaque redémarrage ; il convient donc de conserver les vecteurs.
- Scan linéaire. La similarité cosinus par rapport à chaque vecteur est de complexité O(n). Au-delà d’environ 100 000 fragments, un index ANN devient utile, et c’est là que les bases de données vectorielles se révèlent rentables.
- Qualité du classement. La méthode top-k basique sur la similarité cosinus est un outil peu précis ; l’utilisation de systèmes de réclassement ou de recherches hybrides à mots-clés améliore sa précision.
LangChain encapsule chacun de ces éléments en composants interchangeables. Une fois le boucle brute écrite manuellement, ces composants restent ce qu’ils sont : de simples outils pratiques sans rien de mystérieux.