RAG with Milvus and Gemini page is outdated, here's updated code using Google.genai

[Paxtian769]

I was trying to follow the tutorial on https://milvus.io/docs/build_RAG_with_milvus_and_gemini.md, but that tutorial uses the google.generativeai package, which is now out of date. The new package is google.genai. I was able to create the full pipeline and wanted to share the code here, in case it's helpful to anyone and if anyone wanted to update that webpage. Thanks!

import os
from glob import glob
from google import genai
from google.genai import types
from pymilvus import MilvusClient
from tqdm import tqdm
import json


# In Linux, edit your ~/.bashrc file to include the following line:
# export GEMINI_API_KEY=<your Gemini API key>
# Then the following code will simply grab your key from the os environment
client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])

# Get the text to be inserted into the database
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("# ")

doc = client.models.embed_content(
    model="text-embedding-004", # NOTE: Although gemini-embedding-exp-03-07 is most current, it doesn't appear to work for me
    contents=text_lines
)

# Get the embedding dimensions, per Google Gemini in this case
embedding_dim = len(doc.embeddings[0].values)

# Create the database collection and load data into Milvus
milvus_client = MilvusClient("./milvus_demo.db")

collection_name = "my_rag_collection"

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

milvus_client.create_collection(
    collection_name=collection_name,
    dimension=embedding_dim,
    metric_type="IP",  # Inner product distance
    consistency_level="Strong",  # Supported values are (`"Strong"`, `"Session"`, `"Bounded"`, `"Eventually"`). See https://milvus.io/docs/consistency.md#Consistency-Level for more details.
)

# Insert data into the database
data = []

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
)

# Example query to be used to both query the database and then be augmented by DB results
question = "How is data stored in milvus?"

quest_embed = client.models.embed_content(
    model="text-embedding-004",
    contents=question
).embeddings

search_res = milvus_client.search(
    collection_name=collection_name,
    data=[quest_embed[0].values],
    limit=3,
    search_params={"metric_type": "IP", "params": {}},
    output_fields=["text"],
)

retrieved_lines_with_distances = [
    (res["entity"]["text"], res["distance"]) for res in search_res[0]
]

# Join all the DB results for the mega-prompt
context = "\n".join(
    [line_with_distance[0] for line_with_distance in retrieved_lines_with_distances]
)

# Mega-prompt outline
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>
"""

# Get a response from Gemini using the mega-prompt
response = client.models.generate_content(
    model="gemini-2.0-flash",
    config=types.GenerateContentConfig(
        system_instruction=SYSTEM_PROMPT
    ),
    contents=USER_PROMPT
)

# See the results
print(response.text)

[yhmo]

Thank you for your efforts, I will notify the doc team to update the page.

[zc277584121]

@Paxtian769 thanks, I have updated this notebook, and it will sync to the milvus.io page within few days
milvus-io/bootcamp#1535

[zc277584121]

@Paxtian769 FYI, it's updated https://milvus.io/docs/build_RAG_with_milvus_and_gemini.md