Mini pipeline RAG w 90 liniach kodu Python: wyszukiwanie, wzbogacanie, generowanie

Retrieval-augmented generation nie jest funkcją żadnego modelu. To problem pakowania. Okno kontekstowe jest małe, dokumenty ogromne, więc coś musi zdecydować, które fragmenty trafią do środka. To rozwiązanie można napisać w około 90 liniach prostego Pythona, bez użycia LangChain ani usług bazy danych wektorowych. Składa się ono z trzech kroków:

1. Pobierz. Wbuduj pytanie, uporządkuj fragmenty według podobieństwa i zachowaj najlepsze k elementów.
2. Zwielokrotnij. Połącz te najlepsze fragmenty z wiadomością systemową.
3. Stwórz. Wykonaj tę samą próbę rozmowy, co w każdym prostym skrypcie LLM.

Każda abstrakcja oferowana przez frameworki takie jak LangChain odpowiada jednej z tych trzech kategorii, dlatego zapoznanie się najpierw z wersją podstawową ułatwia zrozumienie całego frameworka.

Embeddingi przekształcają znaczenie w geometrię

Modell, która tworzy embeddingi, przekształca fragment tekstu w wektor, czyli długą sekwencję liczb zmiennoprzecinkowych ułożonych tak, aby teksty o podobnym znaczeniu znajdowały się blisko siebie. Odległość ocenia się za pomocą podobieństwa kosinowego: 1,0 dla identycznego kierunku, 0,0 przy braku związku. W tym cały sekret. Gdy znaczenie staje się geometrią, znajdowanie istotnych dokumentów przekształca się w wyszukiwanie wektorów znajdujących się w pobliżu.

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

Zwykłym wynikiem tego skryptu jest ocena 0,531 dla dokumentu zwrotu pieniędzy i 0,126 dla dokumentu przechowywania. Zobaczmy, jak zachodzi dopasowanie: pytanie mówi „zwrot pieniędzy”, podczas gdy w dokumencie jest napisane „refunds”. Nie ma żadnych wspólnych słów kluczowych, a mimo to podobieństwo jest wysokie. Dlatego należy używać mechanizmu embedding zamiast greppingu.

Baza wektorowa to lista, którą można posortować

„Baza wektorowa” brzmi jak element infrastruktury. W małej skali jest to lista par (wektor, tekst) wraz z funkcją sortowania. Pinecone, Chroma i pgvector działają efektywnie przy milionach wektorów, ale najważniejsza jest interfejs, który mieści się w jednej klasie:

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 to decyzja o tym, co zostanie wybrane – teraz jest realizowalna: wszystko jest rankowane, zachowuje się k najlepszych elementów, a reszta jest odrzucana. Wartość k stanowi parametr regulujący budżet, pozwalający na kompromis pomiędzy dokładnością wyników a ilością dostępnego miejsca w oknie kontekstowym oraz liczbą tokenów zużywanych do obliczeń.

Łączenie funkcji pobierania, uzupełniania i generowania

Tutaj nowa jest tylko jedna funkcja, build_messages. Numeruje ona najlepsze fragmenty tekstu, umieszcza je w wiadomości systemowej i instruuje model, aby odpowiadał wyłącznie na podstawie tych fragmentów. Wiadomość od użytkownika zawiera jedynie pytanie. Lista wiadomości jest taka sama jak ta używana w każdej rozmowie, więc RAG zmienia tylko jej zawartość. Ten plik działa bez zmian:

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

Przykładowe uruchomienie pyta, gdzie przechowywane są dane użytkownika oraz czy można je eksportować. Trzy najlepsze wyniki to 0,588 dla fragmentu dotyczącego przechowywania w Frankfurcie z szyfrowaniem, 0,472 dla fragmentu o eksportowaniu wszystkich danych w formacie JSON oraz 0,244 dla polityki zwrotów. Model odpowiada, że dane znajdują się we Frankfurcie, są szyfrowane w stanie spoczynku za pomocą AES-256, a strona ustawień umożliwia pełny eksport w formacie JSON.

Oba potrzebne fragmenty – dotyczące przechowywania i eksportu – trafiły do wyników. Fragment o polityce zwrotów zajął trzecie miejsce, ale model go zignorował. Pobieranie informacji może być niedoskonałe, o ile odpowiednie materiały znajdą się na górze listy.

Gdzie zawodzi wersja próbna

Te 90 linii ma szczerze określone ograniczenia, a każda z nich wskazuje na coś, co dodają rzeczywiste systemy:

• Dzielenie na fragmenty. Dema przechowuje całe zdania. Prawdziwe dokumenty wymagają podziału z nakładaniem się fragmentów, a to jest decyzja dotycząca przetwarzania wstępnego, a nie modelu.
• Ponowne wplecenie wektorów. Wplecanie zapytania przy każdym wywołaniu jest w porządku. Ponowne wplecanie całego korpusu przy każdym uruchomieniu nie jest dobrym rozwiązaniem, dlatego należy przechowywać wektory.
• Skanowanie liniowe. Porównanie podobieństwa kosinowego z każdym wektorem ma złożoność O(n). Po przekroczeniu około 100 tysięcy fragmentów indeks ANN staje się przydatny, i właśnie wtedy bazy danych wektorowych pokrywają swoje koszty.
• Jakość rankingu. Proste metody oparte na podobieństwie kosinowym to nieskuteczne narzędzia, natomiast mechanizmy ponownego rankingu lub hybrydowe wyszukiwanie słów kluczowych sprawiają, że są one skuteczniejsze.

LangChain umożliwia wykorzystanie każdego z tych elementów jako wymiennego komponentu. Gdy surowy algorytm zostanie napisany ręcznie, te komponenty pokazują się takimi, jakimi są naprawdę – to tylko udogodnienia, a nie coś tajemniczego.