milvus-io_bootcamp

1
{
2
 "cells": [
3
  {
4
   "cell_type": "markdown",
5
   "id": "91c0b9fd-213a-4da8-b84b-c766b424716c",
6
   "metadata": {},
7
   "source": [
8
    "# Milvus + Llama-Index: Enhancing OpenAI Assistant Agent with a Custom Retriever\n",
9
    "\n",
10
    "This shows how to enhance [Llama-Index](https://www.llamaindex.ai/)'s agent built on top of the [OpenAI Assistant API](https://platform.openai.com/docs/assistants/overview) with retriever tool customized by [Milvus](https://zilliz.com).\n",
11
    "\n",
12
    "## Preparation\n",
13
    "\n",
14
    "### 1. Install dependencies"
15
   ]
16
  },
17
  {
18
   "cell_type": "code",
19
   "execution_count": 1,
20
   "id": "b363b294",
21
   "metadata": {},
22
   "outputs": [],
23
   "source": [
24
    "! pip install -q llama-index 'milvus[client]' 'openai>=1.2.0' transformers"
25
   ]
26
  },
27
  {
28
   "cell_type": "markdown",
29
   "id": "958680bb",
30
   "metadata": {},
31
   "source": [
32
    "### 2. Start Milvus Service\n",
33
    "\n",
34
    "There are 2 options to start a Milvus service:\n",
35
    "\n",
36
    "- [Zilliz Cloud](https://zilliz.com/cloud): Zilliz provides cloud-native service for Milvus. It simplifies the process of deploying and scaling vector search applications by eliminating the need to create and maintain complex data infrastructure. [Get Started Free!](https://cloud.zilliz.com/signup)\n",
37
    "- [Open Source Milvus](https://milvus.io): You can install the open source Milvus using either Docker Compose or on Kubernetes.\n",
38
    "\n",
39
    "Here, we use [Milvus Lite](https://milvus.io/docs/milvus_lite.md) to start with a lightweight version of Milvus, which works seamlessly with Google Colab and Jupyter Notebook."
40
   ]
41
  },
42
  {
43
   "cell_type": "code",
44
   "execution_count": 2,
45
   "id": "2aaa83de",
46
   "metadata": {},
47
   "outputs": [],
48
   "source": [
49
    "from milvus import default_server\n",
50
    "\n",
51
    "# default_server.cleanup()  # Optional, run this line if you want to cleanup previous data\n",
52
    "default_server.start()"
53
   ]
54
  },
55
  {
56
   "cell_type": "markdown",
57
   "id": "b695901d",
58
   "metadata": {},
59
   "source": [
60
    "### 3. Download example data\n",
61
    "\n",
62
    "You can use any file(s) to build the knowledge base.\n",
63
    "We will use a SEC file [uber_2021.pdf](https://github.com/run-llama/llama_index/blob/main/docs/examples/data/10k/uber_2021.pdf) as an example."
64
   ]
65
  },
66
  {
67
   "cell_type": "code",
68
   "execution_count": 3,
69
   "id": "dbecdfda",
70
   "metadata": {},
71
   "outputs": [],
72
   "source": [
73
    "! wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/10k/uber_2021.pdf' -O 'uber_2021.pdf'"
74
   ]
75
  },
76
  {
77
   "cell_type": "markdown",
78
   "id": "1661a1c6",
79
   "metadata": {},
80
   "source": [
81
    "## Getting Started\n",
82
    "\n",
83
    "### 1. Set up Environment\n",
84
    "\n",
85
    "You need to set up some environment variables, for example, passing your [OpenAI API Key](https://beta.openai.com/account/api-keys).\n",
86
    "Please note that your OpenAI account should have the accessibility and enough quota available for the model [GPT-4 Turbo](https://platform.openai.com/docs/models/gpt-4)."
87
   ]
88
  },
89
  {
90
   "cell_type": "code",
91
   "execution_count": 4,
92
   "id": "08081fca-0274-42ae-a777-4b66d36d691c",
93
   "metadata": {},
94
   "outputs": [],
95
   "source": [
96
    "import os\n",
97
    "from getpass import getpass\n",
98
    "\n",
99
    "os.environ['TOKENIZERS_PARALLELISM'] = 'false'\n",
100
    "os.environ['OPENAI_API_KEY'] = getpass('Enter OpenAI API Key:')"
101
   ]
102
  },
103
  {
104
   "cell_type": "markdown",
105
   "id": "124e19e5",
106
   "metadata": {},
107
   "source": [
108
    "### 2. Customize Strategies\n",
109
    "\n",
110
    "In this step, we will define some strategies to be used:\n",
111
    "\n",
112
    "- Chunking: configure the text splitter (e.g. `chunk_size`)\n",
113
    "- Embedding: choose embedding model (e.g. [`BAAI/bge-small-en`](https://huggingface.co/BAAI/bge-small-en)) and its provider (e.g. [HuggingFace](https://huggingface.co/models), [OpenAI](https://platform.openai.com/docs/guides/embeddings)).\n",
114
    "- LLM: select LLM model (e.g. `gpt-4-1106-preview`) and set up model parameters (e.g. `temperature`)."
115
   ]
116
  },
117
  {
118
   "cell_type": "code",
119
   "execution_count": 5,
120
   "id": "4712df4c",
121
   "metadata": {},
122
   "outputs": [
123
    {
124
     "name": "stderr",
125
     "output_type": "stream",
126
     "text": [
127
      "/home/mengjia.gu/anaconda3/envs/develop/lib/python3.8/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
128
      "  from .autonotebook import tqdm as notebook_tqdm\n"
129
     ]
130
    }
131
   ],
132
   "source": [
133
    "from llama_index.embeddings import HuggingFaceEmbedding\n",
134
    "from llama_index.llms import OpenAI\n",
135
    "from llama_index.vector_stores import MilvusVectorStore\n",
136
    "from llama_index import StorageContext, ServiceContext\n",
137
    "\n",
138
    "\n",
139
    "llm = OpenAI(model='gpt-4-1106-preview')\n",
140
    "embed_model = HuggingFaceEmbedding(model_name='BAAI/bge-small-en', cache_folder='./tmp/models', device='cpu')\n",
141
    "service_context = ServiceContext.from_defaults(llm=llm, embed_model=embed_model, chunk_size=350)\n",
142
    "\n",
143
    "vector_store = MilvusVectorStore(\n",
144
    "    uri=f'http://localhost:{default_server.listen_port}',\n",
145
    "    # token='',  # required for Zilliz Cloud\n",
146
    "    dim=384,  # the value changes with embedding model\n",
147
    "    overwrite=True  # drop table if exist and then create\n",
148
    "    )\n",
149
    "storage_context = StorageContext.from_defaults(vector_store=vector_store)"
150
   ]
151
  },
152
  {
153
   "cell_type": "markdown",
154
   "id": "8f1147b1",
155
   "metadata": {},
156
   "source": [
157
    "### 3. Ingest Document(s)"
158
   ]
159
  },
160
  {
161
   "cell_type": "code",
162
   "execution_count": 6,
163
   "id": "c306c18a",
164
   "metadata": {},
165
   "outputs": [],
166
   "source": [
167
    "from llama_index import SimpleDirectoryReader, VectorStoreIndex\n",
168
    "\n",
169
    "# Load document(s)\n",
170
    "docs = SimpleDirectoryReader(input_files=['./uber_2021.pdf']).load_data()\n",
171
    "\n",
172
    "# Build index\n",
173
    "vector_index = VectorStoreIndex.from_documents(docs, storage_context=storage_context, service_context=service_context)"
174
   ]
175
  },
176
  {
177
   "cell_type": "markdown",
178
   "id": "1ee425a1",
179
   "metadata": {},
180
   "source": [
181
    "### 4. Define Agent & Tool(s)\n",
182
    "\n",
183
    "In order integrate the vector store index with agent, we need to define the index as a Retriever Tool.\n",
184
    "The agent will be able to recognize the retriever via the tool's name and description in metadata."
185
   ]
186
  },
187
  {
188
   "cell_type": "code",
189
   "execution_count": 7,
190
   "id": "612254fd",
191
   "metadata": {},
192
   "outputs": [],
193
   "source": [
194
    "from llama_index.tools import RetrieverTool, ToolMetadata\n",
195
    "\n",
196
    "milvus_tool = RetrieverTool(\n",
197
    "    retriever=vector_index.as_retriever(similarity_top_k=3),  # retrieve top_k results\n",
198
    "    metadata=ToolMetadata(\n",
199
    "        name=\"CustomRetriever\",\n",
200
    "        description='Retrieve relevant information from provided documents.'\n",
201
    "    ),\n",
202
    ")"
203
   ]
204
  },
205
  {
206
   "cell_type": "markdown",
207
   "id": "7d16e580",
208
   "metadata": {},
209
   "source": [
210
    "Then let's define the agent powered by OpenAI's Assistants API.\n",
211
    "To create a agent, we will define its role, give instructions, and provide tools.\n",
212
    "Here we will make LLM thinking itself a SEC analyst, with Milvus retriever as an available tool."
213
   ]
214
  },
215
  {
216
   "cell_type": "code",
217
   "execution_count": 8,
218
   "id": "dd27eee6",
219
   "metadata": {},
220
   "outputs": [],
221
   "source": [
222
    "from llama_index.agent import OpenAIAssistantAgent\n",
223
    "\n",
224
    "agent = OpenAIAssistantAgent.from_new(\n",
225
    "    name='SEC Analyst',\n",
226
    "    instructions='You are a QA assistant designed to analyze sec filings.',\n",
227
    "    tools=[milvus_tool],\n",
228
    "    verbose=True,\n",
229
    "    run_retrieve_sleep_time=1.0\n",
230
    ")"
231
   ]
232
  },
233
  {
234
   "cell_type": "markdown",
235
   "id": "22a35004",
236
   "metadata": {},
237
   "source": [
238
    "## Try it out!\n",
239
    "\n",
240
    "Now the agent is ready as a SEC analyst. It is able to respond to users based off documents loaded into Milvus.\n",
241
    "\n",
242
    "With `verbose=True`, you are able to what information are retrieved when the agent's answering your question."
243
   ]
244
  },
245
  {
246
   "cell_type": "code",
247
   "execution_count": 9,
248
   "id": "8f88c602",
249
   "metadata": {},
250
   "outputs": [
251
    {
252
     "name": "stdout",
253
     "output_type": "stream",
254
     "text": [
255
      "=== Calling Function ===\n",
256
      "Calling function: CustomRetriever with args: {\"input\":\"Uber Technologies, Inc. annual revenue growth 2021\"}\n",
257
      "Got output: page_label = 77\n",
258
      "file_name = uber_2021.pdf\n",
259
      "file_path = uber_2021.pdf\n",
260
      "UBER TECHNOLOGIES, INC.CONSOLIDATED STATEMENTS OF\n",
261
      " OPERATIONS(In millions, except share amounts which are ref\n",
262
      "lected in thousands, and per share amounts)Year Ended December 31,\n",
263
      "2019\n",
264
      "2020 2021 Revenue\n",
265
      "$ 13,000 $ 11,139 $ 17,455 Costs and expenses\n",
266
      "Cost of revenue, exclusive of dep\n",
267
      "reciation and amortization shown separately below6,061 5,154 9,351 Operations and support\n",
268
      "2,302 1,819 1,877 Sales and marketing\n",
269
      "4,626 3,583 4,789 Research and development\n",
270
      "4,836 2,205 2,054 General and administrative\n",
271
      "3,299 2,666 2,316 Depreciation and amortization\n",
272
      "472 575 902 Total costs and expenses\n",
273
      "21,596 16,002 21,289 Loss from operations\n",
274
      "(8,596) (4,863) (3,834) Interest expense\n",
275
      "(559) (458) (483) Other income (expense), net\n",
276
      "722 (1,625) 3,292 Loss before income taxes and loss from equity me\n",
277
      "thod investments(8,433) (6,946) (1,025) Provision for (benefit fro\n",
278
      "m) income taxes45 (192) (492) Loss from equity method invest\n",
279
      "ments(34) (34) (37) Net loss including non-controlling interests\n",
280
      "(8,512) (6,\n",
281
      "\n",
282
      "page_label = 57\n",
283
      "file_name = uber_2021.pdf\n",
284
      "file_path = uber_2021.pdf\n",
285
      "(61) %(3) % Totals of percentage of \n",
286
      "revenues may not foot due to rounding.Comparison of the Years Ended December 31, 2020 and 2021\n",
287
      "Revenue\n",
288
      "Year Ended December 31,\n",
289
      "2020 to 2021 % Change\n",
290
      "(In millions, except percentages) 2020 2021 Revenue\n",
291
      "$ 11,139 $ 17,455 57 %2021\n",
292
      " Compared to 2020Revenue\n",
293
      " increased $6.3 billion, or 57%, primarily attributable to an increase in Gross Bookings of 56%, or 53% on a constant currency basis. The increase inGross\n",
294
      " Bookings was primarily driven by an increase in Delivery Gross Bookings of 71%, or 66% on a constant currency basis, due to an increase in food deliveryorders\n",
295
      " and  higher  basket  sizes  as  a  result  of  stay-at-home  order  demand  related  to  COVID-19,  as  well  as  continued  expansion  across  U.S.  and  internationalmarkets.\n",
296
      " The increase was also driven by Mobility Gross Bookings growth of 38%, or 36% on a constant currency basis, due to increases in Trip volumes as thebusiness\n",
297
      " recovers from the impacts of COVID-19.\n",
298
      "\n",
299
      "page_label = 82\n",
300
      "file_name = uber_2021.pdf\n",
301
      "file_path = uber_2021.pdf\n",
302
      "UBER TECHNOLOGIES, INC.CONSOLIDATED STATEMENTS OF\n",
303
      " CASH FLOWS(In millions)\n",
304
      "Year Ended December 31,\n",
305
      "2019\n",
306
      "2020 2021 Cash flows from operating activities\n",
307
      "Net loss including non-controll\n",
308
      "ing interests$ (8,512) $ (6,788) $ (570) Adjustments to reconcile n\n",
309
      "et loss to net cash used in operating activities:Depreciation and amortization\n",
310
      "472 575 902 Bad debt expense\n",
311
      "92 76 109 Stock-based compensation\n",
312
      "4,596 827 1,168 Gain on extinguishment of conver\n",
313
      "tible notes and settlement of derivatives(444) — — Gain from sale of investm\n",
314
      "ents— — (413) Gain on business divestitures, net\n",
315
      "— (204) (1,684) Deferred income taxes\n",
316
      "(88) (266) (692) Impairment of debt and \n",
317
      "equity securities— 1,690 — Impairments of goodwill, long\n",
318
      "-lived assets and other assets— 404 116 Loss from equity method invest\n",
319
      "ments34 34 37 Unrealized (gain) loss on deb\n",
320
      "t and equity securities, net(2) 125 (1,142) Unrealized foreign cur\n",
321
      "rency transactions16 48 38 Other\n",
322
      "15 2 4 Change in assets and liabili\n",
323
      "ties, net of impact of business acquisitions and disposals:Accounts receivable\n",
324
      "(407) 142 (597) Prepaid expenses and other asse\n",
325
      "ts(478) 94 (236) Collateral held by insurer\n",
326
      "(1,\n",
327
      "\n",
328
      "\n",
329
      "========================\n"
330
     ]
331
    }
332
   ],
333
   "source": [
334
    "# print('Thread id:', agent.thread_id)\n",
335
    "response = agent.chat('''What was Uber's revenue growth in 2021?''')"
336
   ]
337
  },
338
  {
339
   "cell_type": "markdown",
340
   "id": "725e09bf",
341
   "metadata": {},
342
   "source": [
343
    "Check the agent's answer:"
344
   ]
345
  },
346
  {
347
   "cell_type": "code",
348
   "execution_count": 10,
349
   "id": "06dcf5dd",
350
   "metadata": {},
351
   "outputs": [
352
    {
353
     "name": "stdout",
354
     "output_type": "stream",
355
     "text": [
356
      "Uber's revenue growth in 2021 was 57 percent compared to 2020, with revenue increasing from $11,139 million in 2020 to $17,455 million in 2021.\n"
357
     ]
358
    }
359
   ],
360
   "source": [
361
    "print(str(response))"
362
   ]
363
  },
364
  {
365
   "cell_type": "markdown",
366
   "id": "fd183e6f",
367
   "metadata": {},
368
   "source": [
369
    "## Optional\n",
370
    "\n",
371
    "For Milvus-Lite, stop the service at the end."
372
   ]
373
  },
374
  {
375
   "cell_type": "code",
376
   "execution_count": 11,
377
   "id": "4a32cd16",
378
   "metadata": {},
379
   "outputs": [],
380
   "source": [
381
    "default_server.stop()"
382
   ]
383
  }
384
 ],
385
 "metadata": {
386
  "kernelspec": {
387
   "display_name": "develop",
388
   "language": "python",
389
   "name": "python3"
390
  },
391
  "language_info": {
392
   "codemirror_mode": {
393
    "name": "ipython",
394
    "version": 3
395
   },
396
   "file_extension": ".py",
397
   "mimetype": "text/x-python",
398
   "name": "python",
399
   "nbconvert_exporter": "python",
400
   "pygments_lexer": "ipython3",
401
   "version": "3.8.18"
402
  }
403
 },
404
 "nbformat": 4,
405
 "nbformat_minor": 5
406
}
407