Міні-пайплайн RAG у 90 рядках Python: пошук, підсилення, генерація

Генерація з підсиленням за допомогою отримання даних — це не функція жодної моделі. Це проблема пакування: вікно контексту мале, документи величезні, і хтось мусить вирішити, які фрагменти потраплять всередину. Цей механізм можна описати приблизно у 90 рядках звичайного Python без використання LangChain чи сервісу векторних баз даних. Він складається з трьох кроків:

1. Отримання даних. Інкорпорувати запит, ранжувати фрагменти за схожістю та зберегти кращі k з них.
2. Підсилення. Об’єднати обрані фрагменти у системне повідомлення.
3. Генерація. Виконати той самий запит до чату, що використовується у будь-якому базовому скрипті LLM.

Кожна абстракція, яку пропонує фреймворк на кшталт LangChain, у кінцевому підсумку відповідає одній із цих трьох категорій, тому спочатку розглянувши просту версію, можна легше зрозуміти фреймворк.

Інкорпорації перетворюють значення на геометрію

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

Типовий результат виконання цього скрипту — це оцінка 0,531 для документа про повернення грошей та 0,126 для документа про зберігання. Подивіться, як відбувається зіставлення: у запитанні сказано „гроші назад“, тоді як у документі — „повернення“. Ключових слів немає спільних, проте схожість висока. Саме тому використовується вбудовування замість пошуку за шаблоном.

База векторів — це список, який можна ранжувати

„База векторів“ звучить як інфраструктура. У невеликих масштабах це список пар (вектор, текст) та алгоритм сортування. Pinecone, Chroma та pgvector ефективні при роботі з мільйонами векторів, але важлива інтерфейсна частина, яка підпадає під один клас:

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 — це критерій вибору того, що потрапляє до результатів, який тепер можна застосувати: ранжувати все, зберегти k елементів та видалити решту. Значення k є параметром, який дозволяє балансувати між точністю відповідей та обсягом контекстного вікна та кількістю оброблених токенів.

Об’єднання функцій отримання, доповнення та генерації

Тут є лише одна нова функція — build_messages. Вона пронумеровує вибрані фрагменти, додає їх до системного повідомлення та вказує моделі відповідати лише на основі цих фрагментів. Повідомлення користувача містить лише запитання. Список повідомлень залишається таким самим, як у будь-якому чаті, тож RAG змінює лише його вміст. Цей файл працює без змін:

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

Під час зразкового запуску запитується, де зберігаються дані користувача та чи можна їх експортувати. Три найкращі результати: 0,588 для частини про зберігання даних у Франкфурті з шифруванням, 0,472 для частини про експорт усіх даних у форматі JSON та 0,244 для політики повернення грошей. Модель відповідає, що дані знаходяться у Франкфурті, шифруються під час зберігання за допомогою AES-256, а сторінка налаштувань дозволяє здійснити повний експорт у форматі JSON.

Обидві необхідні частини — про зберігання та експорт — потрапили до результату. Частина про повернення грошей посіла третє місце, але модель її проігнорувала. Отримання інформації може бути недосконалим, якщо лише правильний матеріал опиниться вище за інші.

Де ламається демонстраційна версія

Ці 90 рядків мають чесні обмеження, і кожен з них вказує на те, що додають справжні системи:

• Чанкування. У демо-версії зберігаються цілі речення. Справжні документи потребують розділення з перекриттям, і це є рішенням етапу попередньої обробки, а не моделі.
• Повторне вбудовування. Вбудовування запиту при кожному виклику є прийнятним. Однак повторне вбудовування всього корпусу при кожному запуску — ні, тому необхідно зберігати вектори.
• Лінійний сканування. Обчислення косинусової схожості з кожним вектором займає час O(n). Після приблизно 100 тис. чанків індекс типу ANN стає корисним, і саме тоді бази даних векторів окупають себе.
• Якість ранжування. Просте косинусове ранжування типу top-k є грубим інструментом, тоді як алгоритми переранжування чи гібридні пошукові системи на основі ключових слів роблять його більш точним.

LangChain обгортає кожен з цих елементів як компонент, який можна замінити. Як тільки первинний цикл буде написаний вручну, ці компоненти стануть такими, якими вони є насправді — зручними інструментами, без будь-яких містичних елементів.