Costruire RAG con Milvus ed EmbedAnything

EmbedAnything è una pipeline di incorporamento leggera e veloce, realizzata in Rust, che supporta testo, PDF, immagini, audio e altro ancora.

In questo tutorial dimostreremo come costruire una pipeline Retrieval-Augmented Generation (RAG) utilizzando EmbedAnything insieme a Milvus. Invece di accoppiarsi strettamente con un database specifico, EmbedAnything utilizza un sistema di adattatori collegabili: gli adattatori servono come involucri che definiscono il modo in cui gli embeddings vengono formattati, indicizzati e memorizzati nel vector store di destinazione.

Abbinando EmbedAnything a un adattatore Milvus, è possibile generare embeddings da diversi tipi di file e memorizzarli in modo efficiente in Milvus con poche righe di codice.

⚠️ Nota: l'adattatore di EmbedAnything gestisce l'inserimento in Milvus, ma non supporta la ricerca. Per creare una pipeline RAG completa, è necessario istanziare un MilvusClient separatamente e implementare la logica di recupero (ad esempio, la ricerca di similarità sui vettori) come parte della propria applicazione.

Preparazione

Dipendenze e ambiente

$ pip install -qU pymilvus milvus-lite openai embed_anything

Clonare il repository e caricare l'adattatore

Successivamente, cloniamo il repository EmbedAnything e aggiungiamo la cartella examples/adapters al percorso di Python. Qui si trova l'implementazione dell'adattatore personalizzato di Milvus, che consente a EmbedAnything di comunicare con Milvus per l'inserimento dei vettori.

import sys

# Clone the EmbedAnything repository if not already cloned
![ -d "EmbedAnything" ] || git clone https://github.com/StarlightSearch/EmbedAnything.git

# Add the `examples/adapters` directory to the Python path
sys.path.append("EmbedAnything/examples/adapters")
print("✅ EmbedAnything cloned and adapter path added.")

✅ EmbedAnything cloned and adapter path added.

Utilizzeremo OpenAI come LLM in questa pipeline RAG. È necessario preparare la chiave api OPENAI_API_KEY come variabile d'ambiente.

import os
from openai import OpenAI

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

openai_client = OpenAI()

Creare RAG

Inizializzare Milvus

Prima di incorporare qualsiasi file, dobbiamo preparare due componenti che interagiscono con Milvus:

1. MilvusVectorAdapter - Questo è l'adattatore Milvus per EmbedAnything ed è usato solo per l'ingestione di vettori (cioè l'inserimento di embeddings e la creazione di indici). Attualmente non supporta le operazioni di ricerca.
2. MilvusClient - Questo è il client ufficiale di pymilvus, che consente l'accesso completo alle funzionalità di Milvus, tra cui la ricerca vettoriale, il filtraggio e la gestione delle raccolte.

Per evitare confusione:

• Pensate a MilvusVectorAdapter come a uno strumento di "sola scrittura" per la memorizzazione dei vettori.
• Pensate a MilvusClient come al motore di "lettura e ricerca" per eseguire query e recuperare documenti per RAG.

import embed_anything
from embed_anything import (
    WhichModel,
    EmbeddingModel,
)
from milvus_db import MilvusVectorAdapter
from pymilvus import MilvusClient

# Official Milvus client for full operations
milvus_client = MilvusClient(uri="./milvus.db", token="")

# EmbedAnything adapter for pushing embeddings into Milvus
index_name = "embed_anything_milvus_collection"
milvus_adapter = MilvusVectorAdapter(
    uri="./milvus.db", token="", collection_name=index_name
)

# Delete existing collection if it exists
if milvus_client.has_collection(index_name):
    milvus_client.drop_collection(index_name)

# Create a new collection with dimension matching the embedding model later used
milvus_adapter.create_index(dimension=384)

Ok - Milvus DB connection established.
Collection 'embed_anything_milvus_collection' created with index.

Inizializzare il modello di incorporamento e incorporare il documento PDF

Ora inizializzeremo il modello di incorporamento. Utilizzeremo all-MiniLM-L12-v2 model dalla libreria sentence-transformers, che è un modello leggero ma potente per generare embedding di testo. Produce embeddings a 384 dimensioni, quindi si allinea con la dimensione della nostra collezione Milvus, impostata a 384. Questo allineamento è fondamentale e garantisce la compatibilità tra le dimensioni dei vettori memorizzati in Milvus e quelle generate dal modello.

EmbedAnything supporta molti altri modelli di incorporazione. Per maggiori dettagli, consultare la documentazione ufficiale.

# Initialize the embedding model
model = EmbeddingModel.from_pretrained_hf(
    WhichModel.Bert, model_id="sentence-transformers/all-MiniLM-L12-v2"
)

Ora incorporiamo un file PDF. EmbedAnything consente di elaborare facilmente i documenti PDF (e molti altri) e di memorizzarne gli incorporamenti direttamente in Milvus.

# Embed a PDF file
data = embed_anything.embed_file(
    "./pdf_files/WhatisMilvus.pdf",
    embedder=model,
    adapter=milvus_adapter,
)

Converted 12 embeddings for insertion.
Successfully inserted 12 embeddings.

Recuperare e generare la risposta

Anche in questo caso, MilvusVectorAdapter di EmbedAnything è attualmente un'astrazione leggera solo per l'ingestione e l'indicizzazione dei vettori. Non supporta query di ricerca o di recupero. Pertanto, per cercare documenti rilevanti per costruire la nostra pipeline RAG, dobbiamo usare direttamente l'istanza MilvusClient (milvus_client) per interrogare il nostro archivio vettoriale Milvus.

Definire una funzione per recuperare i documenti rilevanti da Milvus.

def retrieve_documents(question, top_k=3):
    query_vector = list(
        embed_anything.embed_query([question], embedder=model)[0].embedding
    )
    search_res = milvus_client.search(
        collection_name=index_name,
        data=[query_vector],
        limit=top_k,
        output_fields=["text"],
    )
    docs = [(res["entity"]["text"], res["distance"]) for res in search_res[0]]
    return docs

Definire una funzione per generare una risposta utilizzando i documenti recuperati nella pipeline RAG.

def generate_rag_response(question):
    retrieved_docs = retrieve_documents(question)
    context = "\n".join([f"Text: {doc[0]}\n" for doc in retrieved_docs])
    system_prompt = (
        "You are an AI assistant. Provide answers based on the given context."
    )
    user_prompt = f"""
    Use the following pieces of information to answer the question. If the information is not in the context, say you don't know.
    
    Context:
    {context}
    
    Question: {question}
    """
    response = openai_client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ],
    )
    return response.choices[0].message.content

Testiamo la pipeline RAG con una domanda di esempio.

question = "How does Milvus search for similar documents?"
answer = generate_rag_response(question)
print(f"Question: {question}")
print(f"Answer: {answer}")

Question: How does Milvus search for similar documents?
Answer: Milvus searches for similar documents primarily through Approximate Nearest Neighbor (ANN) search, which finds the top K vectors closest to a given query vector. It also supports various other types of searches, such as filtering search under specified conditions, range search within a specified radius, hybrid search based on multiple vector fields, and keyword search based on BM25. Additionally, it can perform reranking to adjust the order of search results based on additional criteria, refining the initial ANN search results.