Retrieval-Augmented Generation (RAG) dengan Milvus dan Haystack

Panduan ini menunjukkan cara membangun sistem Retrieval-Augmented Generation (RAG) menggunakan Haystack dan Milvus.

Sistem RAG menggabungkan sistem pencarian dengan model generatif untuk menghasilkan teks baru berdasarkan prompt yang diberikan. Sistem ini pertama-tama mengambil dokumen yang relevan dari korpus menggunakan Milvus, kemudian menggunakan model generatif untuk menghasilkan teks baru berdasarkan dokumen yang telah diambil.

Haystack adalah kerangka kerja Python sumber terbuka dari deepset untuk membangun aplikasi khusus dengan model bahasa besar (LLM). Milvus adalah basis data vektor sumber terbuka tercanggih di dunia, yang dibangun untuk mendukung pencarian kesamaan embedding dan aplikasi AI.

Persyaratan

Tutorial ini memerlukan Python 3.10 atau yang lebih baru. Instal integrasi Haystack untuk Milvus bersama dengan paket PyMilvus dan Milvus Lite terbaru:

pip install --upgrade --quiet "milvus-haystack>=0.1.0" "pymilvus>=3.0.2" "milvus-lite>=3.2.2" "nltk>=3.9.1" markdown-it-py mdit_plain

Kami akan menggunakan model dari OpenAI. Anda harus menyiapkan kunci API OPENAI_API_KEY sebagai variabel lingkungan.

import os

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

Siapkan data

Kami menggunakan konten daring tentang Leonardo Da Vinci sebagai penyimpanan pengetahuan pribadi untuk pipa RAG kami, yang merupakan sumber data yang baik untuk pipa RAG sederhana.

Unduh dan simpan sebagai berkas teks lokal.

import os
import urllib.request

url = "https://www.gutenberg.org/cache/epub/7785/pg7785.txt"
file_path = "./davinci.txt"

if not os.path.exists(file_path):
    urllib.request.urlretrieve(url, file_path)

Buat pipa pengindeksan

Buat pipa pengindeksan yang mengubah teks menjadi dokumen, memecahnya menjadi kalimat, dan menyematkannya. Dokumen-dokumen tersebut kemudian ditulis ke penyimpanan dokumen Milvus.

from haystack import Pipeline
from haystack.components.converters import MarkdownToDocument
from haystack.components.embedders import OpenAIDocumentEmbedder, OpenAITextEmbedder
from haystack.components.preprocessors import DocumentSplitter
from haystack.components.writers import DocumentWriter
from haystack.utils import Secret

from milvus_haystack import MilvusDocumentStore, MilvusEmbeddingRetriever

document_store = MilvusDocumentStore(
    collection_name="haystack_rag_demo",
    connection_args={"uri": "./milvus.db"},
    # connection_args={"uri": "http://localhost:19530"},
    # connection_args={"uri": os.environ["MILVUS_URI"], "token": Secret.from_env_var("MILVUS_TOKEN")},
    drop_old=True,
)

indexing_pipeline = Pipeline()
indexing_pipeline.add_component("converter", MarkdownToDocument())
indexing_pipeline.add_component(
    "splitter", DocumentSplitter(split_by="sentence", split_length=2)
)
indexing_pipeline.add_component("embedder", OpenAIDocumentEmbedder())
indexing_pipeline.add_component("writer", DocumentWriter(document_store))
indexing_pipeline.connect("converter", "splitter")
indexing_pipeline.connect("splitter", "embedder")
indexing_pipeline.connect("embedder", "writer")
indexing_pipeline.run({"converter": {"sources": [file_path]}})

print("Number of documents:", document_store.count_documents())

Number of documents: 134

Buat pipa pengambilan

Buat pipa pengambilan yang mengambil dokumen dari penyimpanan dokumen Milvus menggunakan mesin pencari kesamaan vektor.

question = "Where is Leonardo's earliest dated work kept?"

retrieval_pipeline = Pipeline()
retrieval_pipeline.add_component("embedder", OpenAITextEmbedder())
retrieval_pipeline.add_component(
    "retriever", MilvusEmbeddingRetriever(document_store=document_store, top_k=3)
)
retrieval_pipeline.connect("embedder.embedding", "retriever.query_embedding")

retrieval_results = retrieval_pipeline.run({"embedder": {"text": question}})

print(retrieval_results["retriever"]["documents"][0].content)

This
small picture, which has been considerably repainted, and is perhaps
by Andrea del Verrocchio, Leonardo's master, is the subject of Plate
II.

To January 1473 belongs Leonardo's earliest dated work, a pen-and-ink
drawing--"A Wide View over a Plain," now in the Uffizi. The
inscription together with the date in the top left-hand corner is
reversed, and proves a remarkable characteristic of Leonardo's
handwriting--viz., that he wrote from right to left; indeed, it has
been suggested that he did this in order to make it difficult for any
one else to read the words, which were frequently committed to paper
by the aid of peculiar abbreviations.

Leonardo continued to work in his master's studio till about 1477.

Buat pipa RAG

Buat pipa RAG yang menggabungkan MilvusEmbeddingRetriever dan OpenAIChatGenerator untuk menjawab pertanyaan menggunakan dokumen yang telah diambil.

from haystack.components.builders import ChatPromptBuilder
from haystack.components.generators.chat import OpenAIChatGenerator
from haystack.dataclasses import ChatMessage

prompt_template = [
    ChatMessage.from_user(
        """Answer the following query based on the provided context. If the context does
        not include an answer, reply with 'I don't know'.

        Query: {{query}}
        Documents:
        {% for doc in documents %}
            {{ doc.content }}
        {% endfor %}
        Answer:
        """
    )
]

rag_pipeline = Pipeline()
rag_pipeline.add_component("text_embedder", OpenAITextEmbedder())
rag_pipeline.add_component(
    "retriever", MilvusEmbeddingRetriever(document_store=document_store, top_k=3)
)
rag_pipeline.add_component(
    "prompt_builder", ChatPromptBuilder(template=prompt_template)
)
rag_pipeline.add_component(
    "generator",
    OpenAIChatGenerator(
        api_key=Secret.from_env_var("OPENAI_API_KEY"),
    ),
)
rag_pipeline.connect("text_embedder.embedding", "retriever.query_embedding")
rag_pipeline.connect("retriever.documents", "prompt_builder.documents")
rag_pipeline.connect("prompt_builder.prompt", "generator.messages")

results = rag_pipeline.run(
    {
        "text_embedder": {"text": question},
        "prompt_builder": {"query": question},
    }
)
print("RAG answer:", results["generator"]["replies"][0].text)

RAG answer: Leonardo's earliest dated work is kept in the Uffizi Gallery (the Uffizi).

Kesimpulan

Anda telah membangun pipa pengindeksan, pipa pencarian semantik, dan pipa RAG dengan Haystack dan Milvus. Untuk opsi integrasi lainnya, lihat berkas README milvus-haystack.