Construire RAG avec Milvus et EmbedAnything

EmbedAnything est un pipeline d'intégration ultrarapide et léger construit en Rust qui prend en charge le texte, les PDF, les images, l'audio et bien plus encore.

Dans ce tutoriel, nous allons montrer comment construire un pipeline de Génération Assistée par Récupération (RAG) en utilisant EmbedAnything avec Milvus. Plutôt que de se coupler étroitement à une base de données spécifique, EmbedAnything utilise un système d'adaptateurs enfichables - les adaptateurs servent d'enveloppes qui définissent comment les embeddings sont formatés, indexés et stockés dans le magasin vectoriel cible.

En associant EmbedAnything à un adaptateur Milvus, vous pouvez générer des embeddings à partir de divers types de fichiers et les stocker efficacement dans Milvus en quelques lignes de code seulement.

⚠️ Note : Bien que l'adaptateur de EmbedAnything gère l'insertion dans Milvus, il ne prend pas en charge la recherche dès le départ. Pour construire un pipeline RAG complet, vous devrez également instancier un MilvusClient séparément et mettre en œuvre la logique d'extraction (par exemple, la recherche de similarité sur les vecteurs) dans le cadre de votre application.

Préparation

Dépendances et environnement

$ pip install -qU pymilvus milvus-lite openai embed_anything

Cloner le référentiel et charger l'adaptateur

Ensuite, nous allons cloner le repo EmbedAnything et ajouter le répertoire examples/adapters au chemin Python. C'est là que nous stockons l'implémentation de l'adaptateur Milvus personnalisé, qui permet à EmbedAnything de communiquer avec Milvus pour l'insertion de vecteurs.

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.

Nous utiliserons OpenAI comme LLM dans ce pipeline RAG. Vous devez préparer la clé api OPENAI_API_KEY en tant que variable d'environnement.

import os
from openai import OpenAI

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

openai_client = OpenAI()

Construire RAG

Initialiser Milvus

Avant d'intégrer des fichiers, nous devons préparer deux composants qui interagissent avec Milvus :

1. MilvusVectorAdapter - Il s'agit de l'adaptateur Milvus pour EmbedAnything, utilisé uniquement pour l'ingestion de vecteurs (c'est-à-dire l'insertion d'embeddings et la création d'index). Il ne prend actuellement pas en charge les opérations de recherche.
2. MilvusClient - Il s'agit du client officiel de pymilvus, qui permet un accès complet aux fonctionnalités de Milvus, notamment la recherche vectorielle, le filtrage et la gestion des collections.

Pour éviter toute confusion :

• Considérez MilvusVectorAdapter comme votre outil de stockage de vecteurs en "écriture seule".
• Pensez à MilvusClient comme votre moteur de "lecture et de recherche" pour effectuer des requêtes et récupérer des documents pour le 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.

Initialisation du modèle d'intégration et intégration du document PDF

Nous allons maintenant initialiser le modèle d'intégration. Nous utiliserons le modèle all-MiniLM-L12-v2 model de la bibliothèque sentence-transformers, qui est un modèle léger mais puissant pour générer des embeddings de texte. Il produit des encastrements à 384 dimensions, ce qui correspond à la dimension de notre collection Milvus, fixée à 384. Cet alignement est crucial et assure la compatibilité entre les dimensions vectorielles stockées dans Milvus et celles générées par le modèle.

EmbedAnything prend en charge beaucoup plus de modèles d'intégration. Pour plus de détails, veuillez vous référer à la documentation officielle.

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

Maintenant, incorporons un fichier PDF. EmbedAnything facilite le traitement des documents PDF (et bien d'autres) et le stockage de leurs incorporations directement dans 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.

Récupérer et générer une réponse

Encore une fois, le site MilvusVectorAdapter d'EmbedAnything est actuellement une abstraction légère pour l'ingestion et l'indexation de vecteurs uniquement. Elle ne prend pas en charge les requêtes de recherche ou d'extraction. Par conséquent, pour rechercher des documents pertinents afin de construire notre pipeline RAG, nous devons utiliser directement l'instance MilvusClient (milvus_client) pour interroger notre magasin de vecteurs Milvus.

Définir une fonction pour récupérer les documents pertinents de 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

Définir une fonction pour générer une réponse en utilisant les documents récupérés dans le 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

Testons le pipeline RAG avec un exemple de question.

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.