wehub/run-llama--llama_index
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
main
run-llama--llama_index/docs/examples/vector_stores/MilvusAsyncAPIDemo.ipynb
T
wehub-resource-sync a0c8464e58
Build Package / build (ubuntu-latest) (push) Failing after 1s
Details
CodeQL / Analyze (python) (push) Failing after 1s
Details
Core Typecheck / core-typecheck (push) Failing after 1s
Details
Linting / lint (push) Failing after 1s
Details
llama-dev tests / test-llama-dev (push) Failing after 1s
Details
Publish Sub-Package to PyPI if Needed / publish_subpackage_if_needed (push) Has been skipped
Details
Sync Docs to Developer Hub / sync-docs (push) Failing after 0s
Details
Build Package / build (windows-latest) (push) Has been cancelled
Details
chore: import upstream snapshot with attribution
2026-07-13 12:26:52 +08:00

698 lines
19 KiB
Plaintext
Raw Permalink Blame History

{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/vector_stores/MilvusAsyncAPIDemo.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Milvus Vector Store with Async API\n",
"\n",
"This tutorial demonstrates how to use [LlamaIndex](https://www.llamaindex.ai/) with [Milvus](https://milvus.io/) to build asynchronous document processing pipeline for RAG. LlamaIndex provides a way to process documents and store in vector db like Milvus. By leveraging the async API of LlamaIndex and Milvus Python client library, we can increase the throughput of the pipeline to efficiently process and index large volumes of data.\n",
"\n",
" \n",
"In this tutorial, we will first introduce the use of asynchronous methods to build a RAG with LlamaIndex and Milvus from a high level, and then introduce the use of low level methods and the performance comparison of synchronous and asynchronous.\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Before you begin\n",
"\n",
"Code snippets on this page require pymilvus and llamaindex dependencies. You can install them using the following commands:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"! pip install -U pymilvus llama-index-vector-stores-milvus llama-index nest-asyncio"
]
},
{
"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": [
"We will use the models from OpenAI. You should prepare the [api key](https://platform.openai.com/docs/quickstart) `OPENAI_API_KEY` as an environment variable."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"\n",
"os.environ[\"OPENAI_API_KEY\"] = \"sk-***********\""
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"If you are using Jupyter Notebook, you need to run this line of code before running the asynchronous code."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import nest_asyncio\n",
"\n",
"nest_asyncio.apply()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Prepare data\n",
"\n",
"You can download sample data with the following commands:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"! mkdir -p 'data/'\n",
"! wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham_essay.txt'\n",
"! wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/10k/uber_2021.pdf' -O 'data/uber_2021.pdf'"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Build RAG with Asynchronous Processing\n",
"This section show how to build a RAG system that can process docs in asynchronous manner.\n",
"\n",
"Import the necessary libraries and define Milvus URI and the dimension of the embedding. \n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import asyncio\n",
"import random\n",
"import time\n",
"\n",
"from llama_index.core.schema import TextNode, NodeRelationship, RelatedNodeInfo\n",
"from llama_index.core.vector_stores import VectorStoreQuery\n",
"from llama_index.vector_stores.milvus import MilvusVectorStore\n",
"\n",
"URI = \"http://localhost:19530\"\n",
"DIM = 768"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"> - If you have large scale of data, you can set up a 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.\n",
"> - In the case of complex systems (such as network communication), asynchronous processing can bring performance improvement compared to synchronization. So we think Milvus-Lite is not suitable for using asynchronous interfaces because the scenarios used are not suitable.\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Define an initialization function that we can use again to rebuild the Milvus collection."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"2025-01-24 20:04:39,414 [DEBUG][_create_connection]: Created new connection using: faa8be8753f74288bffc7e6d38942f8a (async_milvus_client.py:600)\n"
]
}
],
"source": [
"def init_vector_store():\n",
" return MilvusVectorStore(\n",
" uri=URI,\n",
" # token=TOKEN,\n",
" dim=DIM,\n",
" collection_name=\"test_collection\",\n",
" embedding_field=\"embedding\",\n",
" id_field=\"id\",\n",
" similarity_metric=\"COSINE\",\n",
" consistency_level=\"Strong\",\n",
" overwrite=True, # To overwrite the collection if it already exists\n",
" )\n",
"\n",
"\n",
"vector_store = init_vector_store()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Use SimpleDirectoryReader to wrap a LlamaIndex document object from the file `paul_graham_essay.txt`."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Document ID: 41a6f99c-489f-49ff-9821-14e2561140eb\n"
]
}
],
"source": [
"from llama_index.core import SimpleDirectoryReader\n",
"\n",
"# load documents\n",
"documents = SimpleDirectoryReader(\n",
" input_files=[\"./data/paul_graham_essay.txt\"]\n",
").load_data()\n",
"\n",
"print(\"Document ID:\", documents[0].doc_id)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Instantiate a Hugging Face embedding model locally. Using a local model avoids the risk of reaching API rate limits during asynchronous data insertion, as concurrent API requests can quickly add up and use up your budget in public API. However, if you have a high rate limit, you may opt to use a remote model service instead."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from llama_index.embeddings.huggingface import HuggingFaceEmbedding\n",
"\n",
"\n",
"embed_model = HuggingFaceEmbedding(model_name=\"BAAI/bge-base-en-v1.5\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Create an index and insert the document.\n",
"\n",
"We set the `use_async` to `True` to enable async insert mode."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Create an index over the documents\n",
"from llama_index.core import VectorStoreIndex, StorageContext\n",
"\n",
"storage_context = StorageContext.from_defaults(vector_store=vector_store)\n",
"index = VectorStoreIndex.from_documents(\n",
" documents,\n",
" storage_context=storage_context,\n",
" embed_model=embed_model,\n",
" use_async=True,\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Initialize the LLM."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from llama_index.llms.openai import OpenAI\n",
"\n",
"llm = OpenAI(model=\"gpt-3.5-turbo\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"When building the query engine, you can also set the `use_async` parameter to `True` to enable asynchronous search."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"query_engine = index.as_query_engine(use_async=True, llm=llm)\n",
"response = await query_engine.aquery(\"What did the author learn?\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"The author learned that the field of artificial intelligence, as practiced at the time, was not as promising as initially believed. The approach of using explicit data structures to represent concepts in AI was not effective in achieving true understanding of natural language. This realization led the author to shift his focus towards Lisp and eventually towards exploring the field of art.\n"
]
}
],
"source": [
"print(response)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Explore the Async API\n",
"\n",
"In this section, we'll introduce lower level API usage and compare the performance of synchronous and asynchronous runs.\n",
"\n",
"### Async add\n",
"Re-initialize the vector store."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"2025-01-24 20:07:38,727 [DEBUG][_create_connection]: Created new connection using: 5e0d130f3b644555ad7ea6b8df5f1fc2 (async_milvus_client.py:600)\n"
]
}
],
"source": [
"vector_store = init_vector_store()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Let's define a node producing function, which will be used to generate large number of test nodes for the index."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def random_id():\n",
" random_num_str = \"\"\n",
" for _ in range(16):\n",
" random_digit = str(random.randint(0, 9))\n",
" random_num_str += random_digit\n",
" return random_num_str\n",
"\n",
"\n",
"def produce_nodes(num_adding):\n",
" node_list = []\n",
" for i in range(num_adding):\n",
" node = TextNode(\n",
" id_=random_id(),\n",
" text=f\"n{i}_text\",\n",
" embedding=[0.5] * (DIM - 1) + [random.random()],\n",
" relationships={\n",
" NodeRelationship.SOURCE: RelatedNodeInfo(node_id=f\"n{i+1}\")\n",
" },\n",
" )\n",
" node_list.append(node)\n",
" return node_list"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Define a aync function to add documents to the vector store. We use the `async_add()` function in Milvus vector store instance."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"async def async_add(num_adding):\n",
" node_list = produce_nodes(num_adding)\n",
" start_time = time.time()\n",
" tasks = []\n",
" for i in range(num_adding):\n",
" sub_nodes = node_list[i]\n",
" task = vector_store.async_add([sub_nodes]) # use async_add()\n",
" tasks.append(task)\n",
" results = await asyncio.gather(*tasks)\n",
" end_time = time.time()\n",
" return end_time - start_time"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"add_counts = [10, 100, 1000]"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Get the event loop."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"loop = asyncio.get_event_loop()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Asynchronously add documents to the vector store."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Async add for 10 took 0.19 seconds\n",
"Async add for 100 took 0.48 seconds\n",
"Async add for 1000 took 3.22 seconds\n"
]
}
],
"source": [
"for count in add_counts:\n",
"\n",
" async def measure_async_add():\n",
" async_time = await async_add(count)\n",
" print(f\"Async add for {count} took {async_time:.2f} seconds\")\n",
" return async_time\n",
"\n",
" loop.run_until_complete(measure_async_add())"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"2025-01-24 20:07:45,554 [DEBUG][_create_connection]: Created new connection using: b14dde8d6d24489bba26a907593f692d (async_milvus_client.py:600)\n"
]
}
],
"source": [
"vector_store = init_vector_store()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### Compare with synchronous add\n",
"Define a sync add function. Then measure the running time under the same condition."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def sync_add(num_adding):\n",
" node_list = produce_nodes(num_adding)\n",
" start_time = time.time()\n",
" for node in node_list:\n",
" result = vector_store.add([node])\n",
" end_time = time.time()\n",
" return end_time - start_time"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Sync add for 10 took 0.56 seconds\n",
"Sync add for 100 took 5.85 seconds\n",
"Sync add for 1000 took 62.91 seconds\n"
]
}
],
"source": [
"for count in add_counts:\n",
" sync_time = sync_add(count)\n",
" print(f\"Sync add for {count} took {sync_time:.2f} seconds\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The result shows that the sync adding process is much slower than the async one.\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Async search\n",
"\n",
"Re-initialize the vector store and add some documents before running the search."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"2025-01-24 20:08:57,982 [DEBUG][_create_connection]: Created new connection using: 351dc7ea4fb14d4386cfab02621ab7d1 (async_milvus_client.py:600)\n"
]
}
],
"source": [
"vector_store = init_vector_store()\n",
"node_list = produce_nodes(num_adding=1000)\n",
"inserted_ids = vector_store.add(node_list)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Define an async search function. We use the `aquery()` function in Milvus vector store instance."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"async def async_search(num_queries):\n",
" start_time = time.time()\n",
" tasks = []\n",
" for _ in range(num_queries):\n",
" query = VectorStoreQuery(\n",
" query_embedding=[0.5] * (DIM - 1) + [0.6], similarity_top_k=3\n",
" )\n",
" task = vector_store.aquery(query=query) # use aquery()\n",
" tasks.append(task)\n",
" results = await asyncio.gather(*tasks)\n",
" end_time = time.time()\n",
" return end_time - start_time"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"query_counts = [10, 100, 1000]"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Asynchronously search from Milvus store."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Async search for 10 queries took 0.55 seconds\n",
"Async search for 100 queries took 1.39 seconds\n",
"Async search for 1000 queries took 8.81 seconds\n"
]
}
],
"source": [
"for count in query_counts:\n",
"\n",
" async def measure_async_search():\n",
" async_time = await async_search(count)\n",
" print(\n",
" f\"Async search for {count} queries took {async_time:.2f} seconds\"\n",
" )\n",
" return async_time\n",
"\n",
" loop.run_until_complete(measure_async_search())"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### Compare with synchronous search\n",
"Define a sync search function. Then measure the running time under the same condition."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def sync_search(num_queries):\n",
" start_time = time.time()\n",
" for _ in range(num_queries):\n",
" query = VectorStoreQuery(\n",
" query_embedding=[0.5] * (DIM - 1) + [0.6], similarity_top_k=3\n",
" )\n",
" result = vector_store.query(query=query)\n",
" end_time = time.time()\n",
" return end_time - start_time"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Sync search for 10 queries took 3.29 seconds\n",
"Sync search for 100 queries took 30.80 seconds\n",
"Sync search for 1000 queries took 308.80 seconds\n"
]
}
],
"source": [
"for count in query_counts:\n",
" sync_time = sync_search(count)\n",
" print(f\"Sync search for {count} queries took {sync_time:.2f} seconds\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The result shows that the sync search process is much slower than the async one."
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"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"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
Reference in New Issue View Git Blame Copy Permalink