Construindo RAG com Milvus e EmbedAnything

O EmbedAnything é um pipeline de incorporação leve e extremamente rápido, construído em Rust, que suporta texto, PDFs, imagens, áudio e muito mais.

Neste tutorial, demonstraremos como construir um pipeline Retrieval-Augmented Generation (RAG) usando o EmbedAnything junto com o Milvus. Em vez de se acoplar firmemente a qualquer banco de dados específico, o EmbedAnything usa um sistema de adaptador plugável - os adaptadores servem como invólucros que definem como os embeddings são formatados, indexados e armazenados no armazenamento de vetor de destino.

Ao emparelhar o EmbedAnything com um adaptador Milvus, pode gerar embeddings a partir de diversos tipos de ficheiros e armazená-los eficientemente no Milvus em apenas algumas linhas de código.

⚠️ Nota: Embora o adaptador no EmbedAnything lide com a inserção no Milvus, ele não suporta a pesquisa imediata. Para construir um pipeline RAG completo, terá também de instanciar um MilvusClient separadamente e implementar a lógica de recuperação (por exemplo, pesquisa por semelhança em vectores) como parte da sua aplicação.

Preparação

Dependências e ambiente

$ pip install -qU pymilvus milvus-lite openai embed_anything

Clonar o repositório e carregar o adaptador

Em seguida, vamos clonar o repositório EmbedAnything e adicionar o diretório examples/adapters ao caminho do Python. É aqui que armazenamos a implementação personalizada do adaptador Milvus, que permite que o EmbedAnything comunique com o Milvus para a inserção de vectores.

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.

Utilizaremos o OpenAI como LLM neste pipeline RAG. Deve preparar a chave api OPENAI_API_KEY como uma variável de ambiente.

import os
from openai import OpenAI

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

openai_client = OpenAI()

Construir o RAG

Inicializar o Milvus

Antes de incorporarmos quaisquer ficheiros, temos de preparar dois componentes que interagem com o Milvus:

1. MilvusVectorAdapter - Este é o adaptador do Milvus para o EmbedAnything, e é utilizado apenas para a ingestão de vectores (ou seja, inserir embeddings e criar índices). Atualmente, não suporta operações de pesquisa.
2. MilvusClient - Este é o cliente oficial de pymilvus, que permite o acesso total às capacidades do Milvus, incluindo pesquisa de vectores, filtragem e gestão de colecções.

Para evitar confusões:

• Pense em MilvusVectorAdapter como a sua ferramenta "só de escrita" para armazenar vectores.
• Pense em MilvusClient como o seu motor de "leitura e pesquisa" para efetuar consultas e obter documentos para o 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.

Inicializar o modelo de incorporação e incorporar o documento PDF

Agora vamos inicializar o modelo de incorporação. Usaremos o all-MiniLM-L12-v2 model da biblioteca sentence-transformers, que é um modelo leve, mas poderoso, para gerar embeddings de texto. Produz 384 dimensões de embeddings, o que se alinha com a dimensão da nossa coleção Milvus que está definida para 384. Este alinhamento é crucial e garante a compatibilidade entre as dimensões vectoriais armazenadas no Milvus e as geradas pelo modelo.

O EmbedAnything suporta muitos mais modelos de incorporação. Para mais pormenores, consulte a documentação oficial.

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

Agora, vamos incorporar um ficheiro PDF. O EmbedAnything facilita o processamento de documentos PDF (e muitos mais) e o armazenamento dos seus embeddings diretamente no 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.

Recuperar e gerar resposta

Mais uma vez, o MilvusVectorAdapter do EmbedAnything é atualmente uma abstração leve apenas para ingestão e indexação de vectores. Não suporta consultas de pesquisa ou recuperação. Por conseguinte, para pesquisar documentos relevantes para construir o nosso pipeline RAG, temos de utilizar diretamente a instância MilvusClient (milvus_client) para consultar o nosso armazenamento de vectores Milvus.

Definir uma função para obter documentos relevantes do 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

Definir uma função para gerar uma resposta utilizando os documentos recuperados no 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

Vamos testar o pipeline RAG com uma pergunta de exemplo.

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.