RAG com o Milvus e a API assíncrona do LlamaIndex

Este tutorial demonstra como utilizar o LlamaIndex com o Milvus para criar um pipeline de processamento assíncrono de documentos para RAG. O LlamaIndex permite processar documentos e armazená-los numa base de dados vetorial, como o Milvus. Ao tirar partido da API assíncrona do LlamaIndex e da biblioteca cliente Python do Milvus, podemos aumentar a capacidade de processamento do pipeline para processar e indexar eficientemente grandes volumes de dados.

Neste tutorial, iremos primeiro apresentar, a um nível geral, a utilização de métodos assíncronos para criar um RAG com o LlamaIndex e o Milvus e, em seguida, abordar a utilização de métodos de baixo nível e a comparação de desempenho entre os métodos síncronos e assíncronos.

Antes de começar

Os trechos de código nesta página requerem as dependências pymilvus e llamaindex. Pode instalá-las utilizando os seguintes comandos:

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

Iremos utilizar os modelos da OpenAI. Deve preparar a chave API OPENAI_API_KEY como variável de ambiente.

import os

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

Se estiver a utilizar o Jupyter Notebook, terá de executar esta linha de código antes de executar o código assíncrono.

import nest_asyncio

nest_asyncio.apply()

Preparar os dados

Pode descarregar dados de exemplo com os seguintes comandos:

$ 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'

Criar um RAG com processamento assíncrono

Esta secção mostra como criar um sistema RAG capaz de processar documentos de forma assíncrona.

Importe as bibliotecas necessárias e defina o URI do Milvus e a dimensão da incorporação.

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

Defina uma função de inicialização que possamos utilizar novamente para reconstruir a coleção do 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)

Utilize o SimpleDirectoryReader para encapsular um objeto de documento LlamaIndex a partir do ficheiro 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

Instancie um modelo de embedding do Hugging Face localmente. A utilização de um modelo local evita o risco de atingir os limites de taxa da API durante a inserção assíncrona de dados, uma vez que os pedidos simultâneos à API podem acumular-se rapidamente e esgotar o seu orçamento na API pública. No entanto, se tiver um limite de taxa elevado, pode optar por utilizar um serviço de modelo remoto.

from llama_index.embeddings.huggingface import HuggingFaceEmbedding


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

Crie um índice e insira o documento.

Definimos o parâmetro ` use_async ` para ` True ` para ativar o modo de inserção assíncrona.

# 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,
)

Inicialize o LLM.

from llama_index.llms.openai import OpenAI

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

Ao criar o motor de pesquisa, também pode definir o parâmetro « use_async » como « True » para ativar a pesquisa assíncrona.

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.

Explorar a API assíncrona

Nesta secção, iremos apresentar a utilização da API de nível inferior e comparar o desempenho das execuções síncronas e assíncronas.

Adição assíncrona

Reinicialize o armazenamento de vetores.

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)

Vamos definir uma função de produção de nós, que será utilizada para gerar um grande número de nós de teste para o índice.

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

Defina uma função assíncrona para adicionar documentos ao armazenamento de vetores. Utilizamos a função ` async_add() ` na instância do armazenamento de vetores do 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]

Obter o ciclo de eventos.

loop = asyncio.get_event_loop()

Adicione documentos de forma assíncrona ao armazenamento de vetores.

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)

Compare com a adição síncrona

Defina uma função de adição síncrona. Em seguida, meça o tempo de execução nas mesmas condições.

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

O resultado mostra que o processo de adição síncrona é muito mais lento do que o assíncrono.

Pesquisa assíncrona

Reinicialize o armazenamento de vetores e adicione alguns documentos antes de executar a pesquisa.

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)

Defina uma função de pesquisa assíncrona. Utilizamos a função « aquery() » na instância do armazenamento de vetores 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]

Efetue uma pesquisa assíncrona a partir do armazenamento 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

Compare com a pesquisa síncrona

Defina uma função de pesquisa síncrona. Em seguida, meça o tempo de execução nas mesmas condições.

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

O resultado mostra que o processo de pesquisa síncrona é muito mais lento do que o assíncrono.