Мини-пайплайн 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 становится полезным, и именно тогда векторные базы данных оправдывают свою стоимость.
• Качество ранжирования. Простой метод косинусного сходства с выбором топ-k — это грубый инструмент; для улучшения качества используются алгоритмы переранжирования или гибридные методы поиска по ключевым словам.

LangChain представляет каждый из этих элементов в виде заменяемого компонента. После того как ручной цикл будет написан, эти компоненты останутся просто инструментами для удобства, без каких-либо мистических элементов.