RAG mit Milvus und der asynchronen LlamaIndex-API

Dieses Tutorial zeigt, wie man LlamaIndex zusammen mit Milvus nutzt, um eine asynchrone Pipeline zur Dokumentverarbeitung für RAG aufzubauen. LlamaIndex bietet die Möglichkeit, Dokumente zu verarbeiten und in einer Vektordatenbank wie Milvus zu speichern. Durch die Nutzung der asynchronen API von LlamaIndex und der Milvus-Python-Client-Bibliothek können wir den Durchsatz der Pipeline steigern, um große Datenmengen effizient zu verarbeiten und zu indizieren.

In diesem Tutorial stellen wir zunächst auf allgemeiner Ebene die Verwendung asynchroner Methoden zum Aufbau eines RAG mit LlamaIndex und Milvus vor und gehen anschließend auf die Verwendung von Low-Level-Methoden sowie den Leistungsvergleich zwischen synchroner und asynchroner Verarbeitung ein.

Bevor Sie beginnen

Die Codeausschnitte auf dieser Seite erfordern die Abhängigkeiten „pymilvus“ und „llamaindex“. Sie können diese mit den folgenden Befehlen installieren:

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

Wir werden die Modelle von OpenAI verwenden. Du solltest den API-Schlüssel OPENAI_API_KEY als Umgebungsvariable vorbereiten.

import os

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

Wenn Sie Jupyter Notebook verwenden, müssen Sie diese Codezeile ausführen, bevor Sie den asynchronen Code ausführen.

import nest_asyncio

nest_asyncio.apply()

Daten vorbereiten

Mit den folgenden Befehlen können Sie Beispieldaten herunterladen:

$ 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 mit asynchroner Verarbeitung erstellen

In diesem Abschnitt wird gezeigt, wie ein RAG-System erstellt wird, das Dokumente asynchron verarbeiten kann.

Importieren Sie die erforderlichen Bibliotheken und definieren Sie die Milvus-URI sowie die Dimension der Einbettung.

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

Definieren Sie eine Initialisierungsfunktion, die wir wiederverwenden können, um die Milvus-Sammlung neu aufzubauen.

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)

Verwenden Sie „SimpleDirectoryReader“, um ein LlamaIndex-Dokumentobjekt aus der Datei „ paul_graham_essay.txt “ zu verpacken.

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

Instanziieren Sie ein Hugging-Face-Embedding-Modell lokal. Durch die Verwendung eines lokalen Modells vermeiden Sie das Risiko, während der asynchronen Dateneingabe API-Ratenlimits zu erreichen, da sich gleichzeitige API-Anfragen schnell summieren und Ihr Budget bei öffentlichen APIs aufbrauchen können. Wenn Sie jedoch über ein hohes Ratenlimit verfügen, können Sie stattdessen einen Remote-Modelldienst verwenden.

from llama_index.embeddings.huggingface import HuggingFaceEmbedding


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

Erstellen Sie einen Index und fügen Sie das Dokument ein.

Wir setzen „ use_async “ auf „ True “, um den asynchronen Einfügemodus zu aktivieren.

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

Initialisieren Sie das LLM.

from llama_index.llms.openai import OpenAI

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

Beim Erstellen der Abfrage-Engine können Sie den Parameter „ use_async “ ebenfalls auf „ True “ setzen, um die asynchrone Suche zu aktivieren.

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.

Entdecken Sie die Async-API

In diesem Abschnitt stellen wir die Nutzung der Low-Level-API vor und vergleichen die Leistung synchroner und asynchroner Ausführungen.

Asynchrones Hinzufügen

Initialisieren Sie den Vektorspeicher erneut.

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)

Definieren wir eine Funktion zur Erzeugung von Knoten, die dazu dient, eine große Anzahl von Testknoten für den Index zu generieren.

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

Definieren Sie eine asynchrone Funktion, um Dokumente zum Vektorspeicher hinzuzufügen. Wir verwenden die Funktion „ async_add() “ in der Milvus-Vektorspeicher-Instanz.

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]

Rufen Sie die Ereignisschleife ab.

loop = asyncio.get_event_loop()

Fügen Sie Dokumente asynchron zum Vektorspeicher hinzu.

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)

Vergleichen Sie dies mit dem synchronen Hinzufügen.

Definieren Sie eine Funktion zum synchronen Hinzufügen. Messen Sie anschließend die Laufzeit unter denselben Bedingungen.

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

Das Ergebnis zeigt, dass der synchrone Hinzufügungsvorgang deutlich langsamer ist als der asynchrone.

Asynchrone Suche

Initialisieren Sie den Vektorspeicher neu und fügen Sie einige Dokumente hinzu, bevor Sie die Suche ausführen.

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)

Definieren Sie eine asynchrone Suchfunktion. Wir verwenden die Funktion „ aquery() “ in der Milvus-Vektorspeicherinstanz.

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]

Führen Sie eine asynchrone Suche im Milvus-Speicher durch.

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

Vergleich mit der synchronen Suche

Definieren Sie eine Funktion für die synchrone Suche. Messen Sie anschließend die Laufzeit unter denselben Bedingungen.

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

Das Ergebnis zeigt, dass der synchrone Suchvorgang deutlich langsamer ist als der asynchrone.