Construyendo RAG con Milvus y EmbedAnything

EmbedAnything es una tubería de incrustación rápida y ligera construida en Rust que soporta texto, PDFs, imágenes, audio y más.

En este tutorial, demostraremos cómo construir un canal de Generación de Recuperación-Aumentada (RAG) usando EmbedAnything junto con Milvus. En lugar de acoplarse estrechamente con cualquier base de datos específica, EmbedAnything utiliza un sistema de adaptador conectable - los adaptadores sirven como envoltorios que definen cómo se formatean, indexan y almacenan las incrustaciones en el almacén vectorial de destino.

Al emparejar EmbedAnything con un adaptador Milvus, usted puede generar incrustaciones de diversos tipos de archivos y almacenarlos eficientemente en Milvus en sólo unas pocas líneas de código.

⚠️ Nota: Mientras que el adaptador en EmbedAnything maneja la inserción en Milvus, no soporta la búsqueda fuera de la caja. Para construir una canalización RAG completa, también necesitará instanciar un MilvusClient por separado e implementar la lógica de recuperación (por ejemplo, búsqueda de similitud sobre vectores) como parte de su aplicación.

Preparación

Dependencias y entorno

$ pip install -qU pymilvus milvus-lite openai embed_anything

Clonar el repositorio y cargar el adaptador

A continuación, clonaremos el repositorio de EmbedAnything y añadiremos el directorio examples/adapters a la ruta de Python. Aquí es donde almacenamos la implementación personalizada del adaptador Milvus, que permite a EmbedAnything comunicarse con Milvus para la inserción 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.

Usaremos OpenAI como LLM en este pipeline RAG. Debe preparar la clave api OPENAI_API_KEY como variable de entorno.

import os
from openai import OpenAI

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

openai_client = OpenAI()

Construir RAG

Inicializar Milvus

Antes de incrustar cualquier archivo, necesitamos preparar dos componentes que interactúan con Milvus:

1. MilvusVectorAdapter - Este es el adaptador Milvus para EmbedAnything, y se utiliza sólo para la ingestión de vectores (es decir, la inserción de incrustaciones y la creación de índices). Actualmente no soporta operaciones de búsqueda.
2. MilvusClient - Este es el cliente oficial de pymilvus, que permite el acceso completo a las capacidades de Milvus, incluyendo la búsqueda de vectores, el filtrado y la gestión de colecciones.

Para evitar confusiones:

• Piense en MilvusVectorAdapter como su herramienta de "sólo escritura" para almacenar vectores.
• Piense en MilvusClient como su motor de "lectura y búsqueda" para realizar consultas y recuperar documentos para 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 el modelo de incrustación e incrustar el documento PDF

Ahora inicializaremos el modelo de incrustación. Usaremos el all-MiniLM-L12-v2 model de la librería sentence-transformers, que es un modelo ligero pero potente para generar incrustaciones de texto. Produce incrustaciones de 384 dimensiones, por lo que se alinea con la dimensión de nuestra colección Milvus que se establece en 384. Esta alineación es crucial y garantiza la compatibilidad entre las dimensiones vectoriales almacenadas en Milvus y las generadas por el modelo.

EmbedAnything soporta muchos más modelos de incrustación. Para más detalles, consulte la documentación oficial.

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

Ahora, vamos a incrustar un archivo PDF. EmbedAnything hace fácil procesar documentos PDF (y muchos más) y almacenar sus incrustaciones directamente en 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 y generar respuesta

De nuevo, el MilvusVectorAdapter de EmbedAnything actualmente es una abstracción ligera para la ingestión e indexación de vectores solamente. No admite consultas de búsqueda o recuperación. Por lo tanto, para buscar documentos relevantes para construir nuestra tubería RAG, debemos utilizar directamente la instancia MilvusClient (milvus_client) para consultar nuestro almacén de vectores Milvus.

Definir una función para recuperar documentos relevantes 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

Definir una función para generar una respuesta utilizando los documentos recuperados en la canalización 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

Probemos la canalización RAG con una pregunta de ejemplo.

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.