diff --git a/bootcamp/tutorials/integration/build_RAG_with_milvus_and_gemini.ipynb b/bootcamp/tutorials/integration/build_RAG_with_milvus_and_gemini.ipynb new file mode 100644 index 000000000..22bca6613 --- /dev/null +++ b/bootcamp/tutorials/integration/build_RAG_with_milvus_and_gemini.ipynb @@ -0,0 +1,500 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "\"Open \n", + " \"GitHub" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Build RAG with Milvus and Gemini\n", + "\n", + "The [Gemini API](https://ai.google.dev/gemini-api/docs) and [Google AI Studio](https://ai.google.dev/aistudio) help you start working with Google's latest models and turn your ideas into applications that scale. Gemini provides access to powerful language models like `Gemini-1.5-Flash`, `Gemini-1.5-Flash-8B`, and `Gemini-1.5-Pro` for tasks such as text generation, document processing, vision, audio analysis, and more. The API allows you to input long context with millions of tokens, fine-tune models for specific tasks, generate structured outputs like JSON, and leverage capabilities like semantic retrieval and code execution.\n", + "\n", + "In this tutorial, we will show you how to build a RAG (Retrieval-Augmented Generation) pipeline with Milvus and Gemini. We will use the Gemini model to generate text based on a given query. We will also use Milvus to store and retrieve the generated text." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "\n", + "\n", + "## Preparation\n", + "### Dependencies and Environment" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "! pip install --upgrade pymilvus google-generativeai requests tqdm" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "> If you are using Google Colab, to enable dependencies just installed, you may need to **restart the runtime** (click on the \"Runtime\" menu at the top of the screen, and select \"Restart session\" from the dropdown menu)." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "You should first log in to the Google AI Studio platform and prepare the [api key](https://aistudio.google.com/apikey) `GEMINI_API_KEY` as an environment variable." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "\n", + "os.environ[\"GEMINI_API_KEY\"] = \"***********\"" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Prepare the data\n", + "\n", + "We use the FAQ pages from the [Milvus Documentation 2.4.x](https://github.com/milvus-io/milvus-docs/releases/download/v2.4.6-preview/milvus_docs_2.4.x_en.zip) as the private knowledge in our RAG, which is a good data source for a simple RAG pipeline.\n", + "\n", + "Download the zip file and extract documents to the folder `milvus_docs`." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": { + "vscode": { + "languageId": "shellscript" + } + }, + "outputs": [], + "source": [ + "! wget https://github.com/milvus-io/milvus-docs/releases/download/v2.4.6-preview/milvus_docs_2.4.x_en.zip\n", + "! unzip -q milvus_docs_2.4.x_en.zip -d milvus_docs" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We load all markdown files from the folder `milvus_docs/en/faq`. For each document, we just simply use \"# \" to separate the content in the file, which can roughly separate the content of each main part of the markdown file." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [], + "source": [ + "from glob import glob\n", + "\n", + "text_lines = []\n", + "\n", + "for file_path in glob(\"milvus_docs/en/faq/*.md\", recursive=True):\n", + " with open(file_path, \"r\") as file:\n", + " file_text = file.read()\n", + "\n", + " text_lines += file_text.split(\"# \")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Prepare the LLM and Embedding Model\n", + "\n", + "We use the `gemini-1.5-flash` as LLM, and the `text-embedding-004` as embedding model.\n", + "\n", + "Let's try to generate a test response from the LLM:" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "I am a large language model, trained by Google. I am an AI and don't have a personal identity or consciousness. My purpose is to process information and respond to a wide range of prompts and questions in a helpful and informative way.\n", + "\n" + ] + } + ], + "source": [ + "import google.generativeai as genai\n", + "\n", + "genai.configure(api_key=os.environ[\"GEMINI_API_KEY\"])\n", + "\n", + "gemini_model = genai.GenerativeModel(\"gemini-1.5-flash\")\n", + "\n", + "response = gemini_model.generate_content(\"who are you\")\n", + "print(response.text)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Generate a test embedding and print its dimension and first few elements." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "768\n", + "[0.013588584, -0.004361838, -0.08481652, -0.039724775, 0.04723794, -0.0051557426, 0.026071774, 0.045514572, -0.016867816, 0.039378334]\n" + ] + } + ], + "source": [ + "test_embeddings = genai.embed_content(\n", + " model=\"models/text-embedding-004\", content=[\"This is a test1\", \"This is a test2\"]\n", + ")[\"embedding\"]\n", + "\n", + "embedding_dim = len(test_embeddings[0])\n", + "print(embedding_dim)\n", + "print(test_embeddings[0][:10])" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Load data into Milvus" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Create the Collection" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [], + "source": [ + "from pymilvus import MilvusClient\n", + "\n", + "milvus_client = MilvusClient(uri=\"./milvus_demo.db\")\n", + "\n", + "collection_name = \"my_rag_collection\"" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "> As for the argument of `MilvusClient`:\n", + "> - Setting the `uri` as a local file, e.g.`./milvus.db`, is the most convenient method, as it automatically utilizes [Milvus Lite](https://milvus.io/docs/milvus_lite.md) to store all data in this file.\n", + "> - If you have large scale of data, you can set up a more performant Milvus server on [docker or kubernetes](https://milvus.io/docs/quickstart.md). In this setup, please use the server uri, e.g.`http://localhost:19530`, as your `uri`.\n", + "> - If you want to use [Zilliz Cloud](https://zilliz.com/cloud), the fully managed cloud service for Milvus, adjust the `uri` and `token`, which correspond to the [Public Endpoint and Api key](https://docs.zilliz.com/docs/on-zilliz-cloud-console#free-cluster-details) in Zilliz Cloud." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Check if the collection already exists and drop it if it does." + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "outputs": [], + "source": [ + "if milvus_client.has_collection(collection_name):\n", + " milvus_client.drop_collection(collection_name)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Create a new collection with specified parameters. \n", + "\n", + "If we don't specify any field information, Milvus will automatically create a default `id` field for primary key, and a `vector` field to store the vector data. A reserved JSON field is used to store non-schema-defined fields and their values." + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "outputs": [], + "source": [ + "milvus_client.create_collection(\n", + " collection_name=collection_name,\n", + " dimension=embedding_dim,\n", + " metric_type=\"IP\", # Inner product distance\n", + " consistency_level=\"Strong\", # Strong consistency level\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Insert data\n", + "Iterate through the text lines, create embeddings, and then insert the data into Milvus.\n", + "\n", + "Here is a new field `text`, which is a non-defined field in the collection schema. It will be automatically added to the reserved JSON dynamic field, which can be treated as a normal field at a high level." + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Creating embeddings: 100%|██████████| 72/72 [00:00<00:00, 468201.38it/s]\n" + ] + }, + { + "data": { + "text/plain": [ + "{'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}" + ] + }, + "execution_count": 10, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from tqdm import tqdm\n", + "\n", + "data = []\n", + "\n", + "doc_embeddings = genai.embed_content(\n", + " model=\"models/text-embedding-004\", content=text_lines\n", + ")[\"embedding\"]\n", + "\n", + "for i, line in enumerate(tqdm(text_lines, desc=\"Creating embeddings\")):\n", + " data.append({\"id\": i, \"vector\": doc_embeddings[i], \"text\": line})\n", + "\n", + "milvus_client.insert(collection_name=collection_name, data=data)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Build RAG\n", + "\n", + "### Retrieve data for a query\n", + "\n", + "Let's specify a frequent question about Milvus." + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "metadata": {}, + "outputs": [], + "source": [ + "question = \"How is data stored in milvus?\"" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Search for the question in the collection and retrieve the semantic top-3 matches." + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "metadata": {}, + "outputs": [], + "source": [ + "question_embedding = genai.embed_content(\n", + " model=\"models/text-embedding-004\", content=question\n", + ")[\"embedding\"]\n", + "\n", + "search_res = milvus_client.search(\n", + " collection_name=collection_name,\n", + " data=[question_embedding],\n", + " limit=3, # Return top 3 results\n", + " search_params={\"metric_type\": \"IP\", \"params\": {}}, # Inner product distance\n", + " output_fields=[\"text\"], # Return the text field\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Let's take a look at the search results of the query\n" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[\n", + " [\n", + " \" 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###\",\n", + " 0.8048275113105774\n", + " ],\n", + " [\n", + " \"Does the query perform in memory? What are incremental data and historical data?\\n\\nYes. When a query request comes, Milvus searches both incremental data and historical data by loading them into memory. Incremental data are in the growing segments, which are buffered in memory before they reach the threshold to be persisted in storage engine, while historical data are from the sealed segments that are stored in the object storage. Incremental data and historical data together constitute the whole dataset to search.\\n\\n###\",\n", + " 0.7574886679649353\n", + " ],\n", + " [\n", + " \"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###\",\n", + " 0.7453608512878418\n", + " ]\n", + "]\n" + ] + } + ], + "source": [ + "import json\n", + "\n", + "retrieved_lines_with_distances = [\n", + " (res[\"entity\"][\"text\"], res[\"distance\"]) for res in search_res[0]\n", + "]\n", + "print(json.dumps(retrieved_lines_with_distances, indent=4))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Use LLM to get a RAG response\n", + "\n", + "Convert the retrieved documents into a string format." + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "metadata": {}, + "outputs": [], + "source": [ + "context = \"\\n\".join(\n", + " [line_with_distance[0] for line_with_distance in retrieved_lines_with_distances]\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Define system and user prompts for the Lanage Model. This prompt is assembled with the retrieved documents from Milvus." + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "metadata": {}, + "outputs": [], + "source": [ + "SYSTEM_PROMPT = \"\"\"\n", + "Human: You are an AI assistant. You are able to find answers to the questions from the contextual passage snippets provided.\n", + "\"\"\"\n", + "USER_PROMPT = f\"\"\"\n", + "Use the following pieces of information enclosed in tags to provide an answer to the question enclosed in tags.\n", + "\n", + "{context}\n", + "\n", + "\n", + "{question}\n", + "\n", + "\"\"\"" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Use the Gemini to generate a response based on the prompts." + ] + }, + { + "cell_type": "code", + "execution_count": 17, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Milvus stores data in two ways: Inserted data (vector data, scalar data, and collection-specific schema) is stored as an incremental log in persistent storage using object storage backends such as MinIO, AWS S3, Google Cloud Storage, Azure Blob Storage, Alibaba Cloud OSS, and Tencent Cloud Object Storage. Metadata, generated by each Milvus module, is stored in etcd.\n", + "\n" + ] + } + ], + "source": [ + "gemini_model = genai.GenerativeModel(\n", + " \"gemini-1.5-flash\", system_instruction=SYSTEM_PROMPT\n", + ")\n", + "response = gemini_model.generate_content(USER_PROMPT)\n", + "print(response.text)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Great! We have successfully built a RAG pipeline with Milvus and Gemini." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "py310", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.13" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +}