RAG с использованием Milvus и асинхронного API LlamaIndex

В этом руководстве показано, как использовать LlamaIndex вместе с Milvus для построения конвейера асинхронной обработки документов для RAG. LlamaIndex позволяет обрабатывать документы и хранить их в векторной базе данных, такой как Milvus. Используя асинхронный API LlamaIndex и библиотеку клиента Milvus для Python, мы можем повысить пропускную способность конвейера для эффективной обработки и индексирования больших объемов данных.

В этом руководстве мы сначала в общих чертах рассмотрим использование асинхронных методов для построения RAG с помощью LlamaIndex и Milvus, а затем перейдём к низкоуровневым методам и сравнению производительности синхронного и асинхронного подходов.

Перед началом

Фрагменты кода на этой странице требуют наличия зависимостей pymilvus и llamaindex. Вы можете установить их с помощью следующих команд:

$ pip install -U pymilvus llama-index-vector-stores-milvus llama-index nest-asyncio

Мы будем использовать модели от OpenAI. Вам необходимо подготовить ключ API OPENAI_API_KEY в качестве переменной среды.

import os

os.environ["OPENAI_API_KEY"] = "sk-***********"

Если вы используете Jupyter Notebook, вам необходимо выполнить эту строку кода перед запуском асинхронного кода.

import nest_asyncio

nest_asyncio.apply()

Подготовка данных

Вы можете скачать пример данных с помощью следующих команд:

$ mkdir -p 'data/'
$ wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham_essay.txt'
$ wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/10k/uber_2021.pdf' -O 'data/uber_2021.pdf'

Создание RAG с асинхронной обработкой

В этом разделе показано, как построить систему RAG, способную обрабатывать документы асинхронно.

Импортируйте необходимые библиотеки и определите URI Milvus, а также размерность вложения.

import asyncio
import random
import time

from llama_index.core.schema import TextNode, NodeRelationship, RelatedNodeInfo
from llama_index.core.vector_stores import VectorStoreQuery
from llama_index.vector_stores.milvus import MilvusVectorStore

URI = "http://localhost:19530"
DIM = 768

Определите функцию инициализации, которую мы сможем использовать повторно для восстановления коллекции Milvus.

def init_vector_store():
    return MilvusVectorStore(
        uri=URI,
        # token=TOKEN,
        dim=DIM,
        collection_name="test_collection",
        embedding_field="embedding",
        id_field="id",
        similarity_metric="COSINE",
        consistency_level="Bounded",  # Supported values are (`"Strong"`, `"Session"`, `"Bounded"`, `"Eventually"`). See https://milvus.io/docs/tune_consistency.md#Consistency-Level for more details.
        overwrite=True,  # To overwrite the collection if it already exists
    )


vector_store = init_vector_store()

2025-01-24 20:04:39,414 [DEBUG][_create_connection]: Created new connection using: faa8be8753f74288bffc7e6d38942f8a (async_milvus_client.py:600)

Используйте SimpleDirectoryReader для обертывания объекта документа LlamaIndex из файла paul_graham_essay.txt.

from llama_index.core import SimpleDirectoryReader

# load documents
documents = SimpleDirectoryReader(
    input_files=["./data/paul_graham_essay.txt"]
).load_data()

print("Document ID:", documents[0].doc_id)

Document ID: 41a6f99c-489f-49ff-9821-14e2561140eb

Создайте экземпляр модели встраивания Hugging Face локально. Использование локальной модели позволяет избежать риска достижения лимитов скорости API во время асинхронной вставки данных, так как одновременные запросы к API могут быстро накапливаться и исчерпать ваш бюджет в публичном API. Однако, если у вас высокий лимит скорости, вы можете вместо этого использовать удалённый сервис модели.

from llama_index.embeddings.huggingface import HuggingFaceEmbedding


embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-base-en-v1.5")

Создайте индекс и вставьте документ.

Мы устанавливаем значение параметра ` use_async ` равным ` True `, чтобы включить режим асинхронной вставки.

# Create an index over the documents
from llama_index.core import VectorStoreIndex, StorageContext

storage_context = StorageContext.from_defaults(vector_store=vector_store)
index = VectorStoreIndex.from_documents(
    documents,
    storage_context=storage_context,
    embed_model=embed_model,
    use_async=True,
)

Инициализируем LLM.

from llama_index.llms.openai import OpenAI

llm = OpenAI(model="gpt-3.5-turbo")

При создании механизма запросов вы также можете установить для параметра use_async значение True, чтобы включить асинхронный поиск.

query_engine = index.as_query_engine(use_async=True, llm=llm)
response = await query_engine.aquery("What did the author learn?")

print(response)

The author learned that the field of artificial intelligence, as practiced at the time, was not as promising as initially believed. The approach of using explicit data structures to represent concepts in AI was not effective in achieving true understanding of natural language. This realization led the author to shift his focus towards Lisp and eventually towards exploring the field of art.

Ознакомьтесь с асинхронным API

В этом разделе мы расскажем об использовании низкоуровневого API и сравним производительность синхронного и асинхронного выполнения.

Асинхронное добавление

Повторно инициализируем векторное хранилище.

vector_store = init_vector_store()

2025-01-24 20:07:38,727 [DEBUG][_create_connection]: Created new connection using: 5e0d130f3b644555ad7ea6b8df5f1fc2 (async_milvus_client.py:600)

Давайте определим функцию, генерирующую узлы, которая будет использоваться для создания большого количества тестовых узлов для индекса.

def random_id():
    random_num_str = ""
    for _ in range(16):
        random_digit = str(random.randint(0, 9))
        random_num_str += random_digit
    return random_num_str


def produce_nodes(num_adding):
    node_list = []
    for i in range(num_adding):
        node = TextNode(
            id_=random_id(),
            text=f"n{i}_text",
            embedding=[0.5] * (DIM - 1) + [random.random()],
            relationships={NodeRelationship.SOURCE: RelatedNodeInfo(node_id=f"n{i+1}")},
        )
        node_list.append(node)
    return node_list

Определим асинхронную функцию для добавления документов в векторное хранилище. Мы используем функцию ` async_add() ` в экземпляре векторного хранилища Milvus.

async def async_add(num_adding):
    node_list = produce_nodes(num_adding)
    start_time = time.time()
    tasks = []
    for i in range(num_adding):
        sub_nodes = node_list[i]
        task = vector_store.async_add([sub_nodes])  # use async_add()
        tasks.append(task)
    results = await asyncio.gather(*tasks)
    end_time = time.time()
    return end_time - start_time

add_counts = [10, 100, 1000]

Получим цикл событий.

loop = asyncio.get_event_loop()

Асинхронно добавляем документы в векторное хранилище.

for count in add_counts:

    async def measure_async_add():
        async_time = await async_add(count)
        print(f"Async add for {count} took {async_time:.2f} seconds")
        return async_time

    loop.run_until_complete(measure_async_add())

Async add for 10 took 0.19 seconds
Async add for 100 took 0.48 seconds
Async add for 1000 took 3.22 seconds

vector_store = init_vector_store()

2025-01-24 20:07:45,554 [DEBUG][_create_connection]: Created new connection using: b14dde8d6d24489bba26a907593f692d (async_milvus_client.py:600)

Сравните с синхронным добавлением

Определим функцию синхронного добавления. Затем измерим время выполнения в тех же условиях.

def sync_add(num_adding):
    node_list = produce_nodes(num_adding)
    start_time = time.time()
    for node in node_list:
        result = vector_store.add([node])
    end_time = time.time()
    return end_time - start_time

for count in add_counts:
    sync_time = sync_add(count)
    print(f"Sync add for {count} took {sync_time:.2f} seconds")

Sync add for 10 took 0.56 seconds
Sync add for 100 took 5.85 seconds
Sync add for 1000 took 62.91 seconds

Результаты показывают, что процесс синхронного добавления значительно медленнее, чем асинхронного.

Асинхронный поиск

Перед запуском поиска переинициализируем векторное хранилище и добавим несколько документов.

vector_store = init_vector_store()
node_list = produce_nodes(num_adding=1000)
inserted_ids = vector_store.add(node_list)

2025-01-24 20:08:57,982 [DEBUG][_create_connection]: Created new connection using: 351dc7ea4fb14d4386cfab02621ab7d1 (async_milvus_client.py:600)

Определите функцию асинхронного поиска. Мы используем функцию ` aquery() ` в экземпляре векторного хранилища Milvus.

async def async_search(num_queries):
    start_time = time.time()
    tasks = []
    for _ in range(num_queries):
        query = VectorStoreQuery(
            query_embedding=[0.5] * (DIM - 1) + [0.6], similarity_top_k=3
        )
        task = vector_store.aquery(query=query)  # use aquery()
        tasks.append(task)
    results = await asyncio.gather(*tasks)
    end_time = time.time()
    return end_time - start_time

query_counts = [10, 100, 1000]

Выполните асинхронный поиск в хранилище Milvus.

for count in query_counts:

    async def measure_async_search():
        async_time = await async_search(count)
        print(f"Async search for {count} queries took {async_time:.2f} seconds")
        return async_time

    loop.run_until_complete(measure_async_search())

Async search for 10 queries took 0.55 seconds
Async search for 100 queries took 1.39 seconds
Async search for 1000 queries took 8.81 seconds

Сравните с синхронным поиском

Определите функцию синхронного поиска. Затем измерьте время выполнения в тех же условиях.

def sync_search(num_queries):
    start_time = time.time()
    for _ in range(num_queries):
        query = VectorStoreQuery(
            query_embedding=[0.5] * (DIM - 1) + [0.6], similarity_top_k=3
        )
        result = vector_store.query(query=query)
    end_time = time.time()
    return end_time - start_time

for count in query_counts:
    sync_time = sync_search(count)
    print(f"Sync search for {count} queries took {sync_time:.2f} seconds")

Sync search for 10 queries took 3.29 seconds
Sync search for 100 queries took 30.80 seconds
Sync search for 1000 queries took 308.80 seconds

Результат показывает, что процесс синхронного поиска значительно медленнее, чем асинхронного.