Startseite / Artikel

Ein Mini RAG-Pipeline in 90 Zeilen Python: Daten abrufen, erweitern und generieren.

Baue eine retrieval-augmented Generation von Grund auf mit Embeddings, Kosinusähnlichkeit und einer gerankten Liste auf und untersuche anschließend, an welcher Stelle die vereinfachte Version bei größeren Datenmengen versagt.

Verlag: REACTAPP.TOP6 Min. Lesezeit1342 Wörter

Quellen: https://medium.com/@ericwang_66031/part-7-rag-in-python-retrieve-rank-then-fill-the-context-window-before-langchain-68e889b14c39

Retrieval-augmented Generation ist keine Eigenschaft irgendeines Modells. Es handelt sich dabei um ein Verpackungsproblem: Das Kontextfenster ist klein, die Dokumente sind riesig, und etwas muss entscheiden, welche Teile hineinkommen. Dieses „Etwas“ kann in etwa 90 Zeilen reinen Python geschrieben werden – ohne LangChain und ohne Vektordatenbankdienst. Es umfasst drei Schritte:

  1. Auswählen. Die Frage wird eingebettet, die Fragmente nach Ähnlichkeit gerankt und die besten k Elemente behalten.
  2. Erweitern. Die ausgewählten Fragmente werden in die Systemnachricht eingefügt.
  3. Erzeugen. Es wird derselbe Chataufruf verwendet wie in jedem einfachen LLM-Skript.

Jede Abstraktion, die ein Framework wie LangChain später anbietet, lässt sich auf einen dieser drei Schritte zurückführen – daher macht es die Verständnis des Frameworks einfacher, wenn man zunächst die reine Version betrachtet.

Eingebettete Werte verwandeln Bedeutung in Geometrie

Ein Modell, das Embeddings erzeugt, wandelt einen Textabschnitt in einen Vektor um – also eine Abfolge von Fließkommazahlen, die so angeordnet sind, dass Texte mit ähnlicher Bedeutung nah beieinander liegen. Die Entfernung wird anhand der Kosinusähnlichkeit bestimmt: 1,0 für dieselbe Richtung, 0,0 für keinen Zusammenhang. Das ist der ganze Trick. Sobald Bedeutung zu Geometrie wird, reduziert sich das Finden relevanter Dokumente auf das Auffinden benachbarter Vektoren.

# 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}")

Die typische Ausgabe dieses Skripts ist eine Ähnlichkeitswertung von 0,531 für das Rückerstattungs-Dokument und 0,126 für das Speicher-Dokument. Schauen Sie sich an, wie die Übereinstimmung zustande kommt: Die Frage lautet „Geld zurück“, während im Dokument von „Rückerstattungen“ die Rede ist. Es gibt keine gemeinsamen Schlüsselwörter, dennoch ist die Ähnlichkeit hoch. Deshalb wird ein Embedding verwendet anstelle einer Suchfunktion.

Ein Vektorlager ist eine sortierbare Liste

„Vektorlager“ klingt nach Infrastruktur. In kleinem Maßstab handelt es sich dabei um eine Liste von (Vektor, Text)-Paaren zusammen mit einer Sortierfunktion. Pinecone, Chroma und pgvector funktionieren gut bei Millionen von Vektoren, doch entscheidend ist die Schnittstelle – und diese lässt sich in eine einzige Kategorie einordnen:

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 ist die Entscheidung darüber, was ausgewählt wird – jetzt ausführbar: Alles rangieren, k beibehalten, den Rest entfernen. Der Wert von k ist Ihre Einstellmöglichkeit, bei der zwischen Erinnerungsfähigkeit, Platz im Kontextfenster und abgerechneten Tokens abgewogen wird.

Vereinigung von Abrufen, Erweitern und Erzeugen

Es gibt hier nur eine neue Funktion, build_messages. Sie nummeriert die ausgewählten Teile, fügt sie in die Systemnachricht ein und weist das Modell an, ausschließlich darauf zu antworten. Die Benutzernachricht enthält lediglich die Frage. Die Nachrichtenliste ist dieselbe wie bei jedem Chatgespräch, sodass RAG nur ihren Inhalt ändert. Diese Datei läuft unverändert:

# 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}")

Ein Beispiellauf fragt nach dem Speicherort der Benutzerdaten und ob diese exportiert werden können. Die drei besten Ergebnisse sind 0,588 für einen Abschnitt über den verschlüsselten Speicher in Frankfurt, 0,472 für einen Abschnitt über den Export aller Daten im JSON-Format und 0,244 für die Rückerstattungspolitik. Das Modell antwortet, dass die Daten in Frankfurt gespeichert sind und dort mit AES-256 verschlüsselt werden, sowie dass die Einstellungsseite einen vollständigen JSON-Export anbietet.

Sowohl der für den Speicher notwendige Abschnitt als auch der für den Export benötigte Abschnitt wurden berücksichtigt. Der Abschnitt zur Rückerstattung belegte den dritten Platz, wurde vom Modell jedoch ignoriert. Die Ergebnisse können unvollkommen sein, solange das richtige Material an erster Stelle erscheint.

Wo die Testversion versagt

Diese 90 Zeilen weisen ehrliche Grenzen auf, und jede einzelne weist auf zusätzliche Funktionen hin, die reale Systeme bieten:

  • Chunking. Die Demo speichert ganze Sätze. Reale Dokumente müssen mit Überschneidungen aufgeteilt werden, und das ist eine Entscheidung im Vorverarbeitungsprozess, nicht des Modells.
  • Re-Embedding. Das Einbetten der Abfrage bei jedem Aufruf ist in Ordnung. Das erneute Einbetten des gesamten Korpus bei jedem Neustart hingegen nicht – daher sollten die Vektoren gespeichert werden.
  • Linearer Scan. Die Berechnung der Kosinusähnlichkeit zu jedem Vektor dauert O(n). Ab etwa 100.000 Chunks lohnt sich ein ANN-Index, und genau da rentieren sich Vektordatenbanken aus.
  • Qualität der Rangliste. Die einfache Kosinus-Methode mit Top-k ist ein grobes Werkzeug; Reranker oder hybride Schlüsselwortsuchen machen sie präziser.

LangChain umhüllt jede dieser Methoden als austauschbaren Komponenten. Sobald der Roh-Loop manuell geschrieben wurde, zeigen diese Komponenten genau das, was sie sind – praktische Hilfsmittel ohne etwas Mystisches.