Создайте RAG с помощью Milvus и Gemini

API Gemini и Google AI Studio помогут вам начать работу с новейшими моделями Google и воплотить свои идеи в масштабируемых приложениях. Gemini предоставляет доступ к таким мощным языковым моделям, как Gemini-2.5-Flash и Gemini-2.5-Pro, для таких задач, как генерация текста, обработка документов, зрение, анализ аудио и многое другое. Он также предлагает Gemini Embedding 2, мультимодальную модель встраивания, поддерживающую текст, изображения, видео, аудио и PDF-документы с гибкими размерами вывода с помощью Matryoshka Representation Learning. API позволяет вводить длинные контексты с миллионами лексем, настраивать модели для конкретных задач, генерировать структурированные выходные данные, такие как JSON, и использовать такие возможности, как семантический поиск и выполнение кода.

В этом руководстве мы покажем вам, как построить конвейер RAG (Retrieval-Augmented Generation) с помощью Milvus и Gemini. Мы будем использовать модель Gemini для генерации ответов на основе заданного запроса, дополненного релевантной информацией, полученной из Milvus.

Подготовка

Зависимости и окружение

Сначала установите необходимые пакеты:

$ pip install --upgrade pymilvus milvus-lite google-genai requests tqdm

Сначала необходимо войти в платформу Google AI Studio и подготовить api ключ GEMINI_API_KEY в качестве переменной окружения.

import os

os.environ["GEMINI_API_KEY"] = "***********"

Подготовьте данные

Мы используем страницы FAQ из Milvus Documentation 2.4.x в качестве приватных знаний в нашем RAG, что является хорошим источником данных для простого RAG-конвейера.

Скачайте zip-файл и распакуйте документы в папку milvus_docs.

$ wget https://github.com/milvus-io/milvus-docs/releases/download/v2.4.6-preview/milvus_docs_2.4.x_en.zip
$ unzip -q milvus_docs_2.4.x_en.zip -d milvus_docs

Мы загружаем все файлы разметки из папки milvus_docs/en/faq. Для каждого документа мы просто используем "# " для разделения содержимого в файле, что позволяет примерно разделить содержимое каждой основной части файла разметки.

from glob import glob

text_lines = []

for file_path in glob("milvus_docs/en/faq/*.md", recursive=True):
    with open(file_path, "r") as file:
        file_text = file.read()

    text_lines += file_text.split("# ")

Подготовка LLM и модели встраивания

Мы используем gemini-2.5-flash в качестве LLM и gemini-embedding-2-preview в качестве модели встраивания. gemini-embedding-2-preview - это новейшая мультимодальная модель встраивания Google, поддерживающая текст, изображения, видео, аудио и PDF-документы с гибкими размерами вывода (128-3,072) с помощью Matryoshka Representation Learning.

Давайте попробуем сгенерировать тестовый ответ с помощью LLM:

from google import genai

client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])

response = client.models.generate_content(
    model="gemini-2.5-flash", contents="who are you"
)
print(response.text)

I am a large language model, trained by Google.

I'm designed to process and generate human-like text based on the vast amount of data I was trained on. This allows me to:

*   Answer questions
*   Provide summaries
*   Generate creative content
*   Translate languages
*   And much more

I don't have personal experiences, feelings, or consciousness. I'm a tool designed to be helpful and informative.

Сгенерируйте тестовое вложение и выведите его размерность и первые несколько элементов.

test_embeddings = client.models.embed_content(
    model="gemini-embedding-2-preview", contents=["This is a test1", "This is a test2"]
)

embedding_dim = len(test_embeddings.embeddings[0].values)
print(embedding_dim)
print(test_embeddings.embeddings[0].values[:10])

3072
[-0.016769307, 0.013630492, 0.020277105, 0.0035285393, 0.003968259, -0.013498845, 0.028525498, 0.025498547, -0.021553498, 0.015233516]

Загрузка данных в Milvus

Создайте коллекцию

Давайте инициализируем клиент Milvus и создадим нашу коллекцию:

from pymilvus import MilvusClient

milvus_client = MilvusClient(uri="./milvus_demo.db")

collection_name = "my_rag_collection"

Проверьте, не существует ли уже коллекция, и удалите ее, если она существует.

if milvus_client.has_collection(collection_name):
    milvus_client.drop_collection(collection_name)

Создайте новую коллекцию с указанными параметрами.

Если мы не укажем информацию о полях, Milvus автоматически создаст поле по умолчанию id для первичного ключа и поле vector для хранения векторных данных. Зарезервированное поле JSON используется для хранения не определенных схемой полей и их значений.

milvus_client.create_collection(
    collection_name=collection_name,
    dimension=embedding_dim,
    metric_type="IP",  # Inner product distance
    # Strong consistency waits for all loads to complete, adding latency with large datasets
    # consistency_level="Strong",  # Strong consistency level
)

Вставка данных

Пройдитесь по текстовым строкам, создайте вкрапления, а затем вставьте данные в Milvus.

Вот новое поле text, которое является неопределенным полем в схеме коллекции. Оно будет автоматически добавлено в зарезервированное динамическое поле JSON, с которым можно обращаться как с обычным полем на высоком уровне.

from tqdm import tqdm

data = []

doc = client.models.embed_content(model="gemini-embedding-2-preview", contents=text_lines)

for i, line in enumerate(tqdm(text_lines, desc="Creating embeddings")):
    data.append({"id": i, "vector": doc.embeddings[i].values, "text": line})

milvus_client.insert(collection_name=collection_name, data=data)

Creating embeddings: 100%|██████████| 72/72 [00:00<00:00, 337796.30it/s]





{'insert_count': 72, 'ids': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71], 'cost': 0}

Построение RAG

Получение данных для запроса

Давайте зададим частый вопрос о Milvus.

question = "How is data stored in milvus?"

Найдем этот вопрос в коллекции и получим семантический топ-3 совпадений.

quest_embed = client.models.embed_content(model="gemini-embedding-2-preview", contents=question)

search_res = milvus_client.search(
    collection_name=collection_name,
    data=[quest_embed.embeddings[0].values],
    limit=3,  # Return top 3 results
    search_params={"metric_type": "IP", "params": {}},  # Inner product distance
    output_fields=["text"],  # Return the text field
)

Давайте посмотрим на результаты поиска по этому запросу.

import json

retrieved_lines_with_distances = [
    (res["entity"]["text"], res["distance"]) for res in search_res[0]
]
print(json.dumps(retrieved_lines_with_distances, indent=4))

[
    [
        " Where does Milvus store data?\n\nMilvus deals with two types of data, inserted data and metadata. \n\nInserted data, including vector data, scalar data, and collection-specific schema, are stored in persistent storage as incremental log. Milvus supports multiple object storage backends, including [MinIO](https://min.io/), [AWS S3](https://aws.amazon.com/s3/?nc1=h_ls), [Google Cloud Storage](https://cloud.google.com/storage?hl=en#object-storage-for-companies-of-all-sizes) (GCS), [Azure Blob Storage](https://azure.microsoft.com/en-us/products/storage/blobs), [Alibaba Cloud OSS](https://www.alibabacloud.com/product/object-storage-service), and [Tencent Cloud Object Storage](https://www.tencentcloud.com/products/cos) (COS).\n\nMetadata are generated within Milvus. Each Milvus module has its own metadata that are stored in etcd.\n\n###",
        0.864
    ],
    [
        "Why is there no vector data in etcd?\n\netcd stores Milvus module metadata; MinIO stores entities.",
        0.7923
    ],
    [
        "What is the maximum dataset size Milvus can handle?\n\n  \nTheoretically, the maximum dataset size Milvus can handle is determined by the hardware it is run on, specifically system memory and storage:\n\n- Milvus loads all specified collections and partitions into memory before running queries. Therefore, memory size determines the maximum amount of data Milvus can query.\n- When new entities and and collection-related schema (currently only MinIO is supported for data persistence) are added to Milvus, system storage determines the maximum allowable size of inserted data.\n\n###",
        0.7857
    ]
]

Использование LLM для получения ответа RAG

Преобразуйте полученные документы в строковый формат.

context = "\n".join(
    [line_with_distance[0] for line_with_distance in retrieved_lines_with_distances]
)

Определите системные и пользовательские подсказки для языковой модели. Эта подсказка собирается из документов, полученных из Milvus.

from google.genai import types

SYSTEM_PROMPT = """
Human: You are an AI assistant. You are able to find answers to the questions from the contextual passage snippets provided.
"""
USER_PROMPT = f"""
Use the following pieces of information enclosed in <context> tags to provide an answer to the question enclosed in <question> tags.
<context>
{context}
</context>
<question>
{question}
</question>
"""

Используйте Gemini для генерации ответа на основе подсказок.

response = client.models.generate_content(
    model="gemini-2.5-flash",
    config=types.GenerateContentConfig(system_instruction=SYSTEM_PROMPT),
    contents=USER_PROMPT,
)
print(response.text)

Milvus stores data in two main ways:

1.  **Inserted Data:** This includes vector data, scalar data, and collection-specific schema. This type of data is stored in persistent storage as an incremental log. Milvus supports various object storage backends for this, such as MinIO, AWS S3, Google Cloud Storage (GCS), Azure Blob Storage, Alibaba Cloud OSS, and Tencent Cloud Object Storage (COS).
2.  **Metadata:** Metadata is generated within Milvus by its various modules. Each module's metadata is stored in etcd.

Мультимодальный поиск

Поскольку gemini-embedding-2-preview отображает текст, изображения и другие модальности в одном и том же пространстве встраивания, мы можем выполнять кросс-модальный поиск - например, использовать текстовый запрос для поиска наиболее релевантных изображений.

Подготовка данных изображений

Мы загружаем набор архитектурных диаграмм RAG из репозитория Milvus Bootcamp, чтобы использовать его в качестве набора данных изображений.

import urllib.request
from pathlib import Path

image_dir = Path("images")
image_dir.mkdir(exist_ok=True)

image_files = [
    "vanilla_rag.png",
    "hyde.png",
    "query_routing.png",
    "self_reflection.png",
    "hybrid_and_rerank.png",
    "hierarchical_index.png",
]

base_url = "https://raw.githubusercontent.com/milvus-io/bootcamp/master/pics/advanced_rag/"

for fname in image_files:
    path = image_dir / fname
    if not path.exists():
        urllib.request.urlretrieve(base_url + fname, path)
        print(f"Downloaded {fname}")
    else:
        print(f"Already exists {fname}")

print(f"\nTotal images: {len(image_files)}")

Downloaded vanilla_rag.png
Downloaded hyde.png
Downloaded query_routing.png
Downloaded self_reflection.png
Downloaded hybrid_and_rerank.png
Downloaded hierarchical_index.png

Total images: 6

Встраивание изображений и хранение в Milvus

Мы считываем каждое изображение как байт и передаем его на gemini-embedding-2-preview для генерации вкраплений, а затем сохраняем их в новой коллекции Milvus.

from google.genai import types

image_data = []

for fname in image_files:
    path = image_dir / fname
    with open(path, "rb") as f:
        image_bytes = f.read()

    result = client.models.embed_content(
        model="gemini-embedding-2-preview",
        contents=types.Part.from_bytes(data=image_bytes, mime_type="image/png"),
    )
    image_data.append(
        {
            "id": len(image_data),
            "vector": result.embeddings[0].values,
            "filename": fname,
        }
    )
    print(f"Embedded {fname}")

# Create a new collection for images
image_collection = "image_collection"
if milvus_client.has_collection(image_collection):
    milvus_client.drop_collection(image_collection)

milvus_client.create_collection(
    collection_name=image_collection,
    dimension=len(image_data[0]["vector"]),
    metric_type="IP",
)

milvus_client.insert(collection_name=image_collection, data=image_data)
print(f"\nInserted {len(image_data)} image embeddings (dim={len(image_data[0]['vector'])})")

Embedded vanilla_rag.png
Embedded hyde.png
Embedded query_routing.png
Embedded self_reflection.png
Embedded hybrid_and_rerank.png
Embedded hierarchical_index.png

Inserted 6 image embeddings (dim=3072)

Кросс-модальный поиск: Текстовый запрос → результаты в виде изображений

Теперь давайте воспользуемся текстовым запросом для поиска по вкраплениям изображений. Поскольку и текст, и изображения отображаются в одно и то же пространство вкраплений, мы можем напрямую сравнивать их.

from IPython.display import display, Image

text_queries = [
    "How does a basic RAG pipeline work?",
    "What is the hypothetical document embedding approach?",
    "How to combine hybrid search with reranking?",
]

for query in text_queries:
    query_embed = client.models.embed_content(
        model="gemini-embedding-2-preview", contents=query
    )

    results = milvus_client.search(
        collection_name=image_collection,
        data=[query_embed.embeddings[0].values],
        limit=1,
        search_params={"metric_type": "IP", "params": {}},
        output_fields=["filename"],
    )

    best = results[0][0]
    print(f"\nQuery: {query}")
    print(f"Match: {best['entity']['filename']} (score: {best['distance']:.4f})")
    display(Image(filename=str(image_dir / best["entity"]["filename"]), width=600))

Query: How does a basic RAG pipeline work?
Match: vanilla_rag.png (score: 0.5132)

Ванильный конвейер RAG

Query: What is the hypothetical document embedding approach?
Match: hyde.png (score: 0.4756)

HyDE

Query: How to combine hybrid search with reranking?
Match: hybrid_and_rerank.png (score: 0.5271)

Гибридный поиск и реранжирование

Отлично! Мы успешно построили конвейер RAG с Milvus и Gemini, а также продемонстрировали кросс-модальный поиск с использованием текстовых запросов для получения релевантных изображений - и все это на основе единого пространства встраивания gemini-embedding-2-preview.