{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Local Embeddings with HuggingFace\n",
    "\n",
    "LlamaIndex has support for HuggingFace embedding models, including BGE, Instructor, and more.\n",
    "\n",
    "Furthermore, we provide utilities to create and use ONNX models using the [Optimum library](https://huggingface.co/docs/transformers/serialization#exporting-a-transformers-model-to-onnx-with-optimumonnxruntime) from HuggingFace."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## HuggingFaceEmbedding\n",
    "\n",
    "The base `HuggingFaceEmbedding` class is a generic wrapper around any HuggingFace model for embeddings. All [embedding models](https://huggingface.co/models?library=sentence-transformers) on Hugging Face should work. You can refer to the [embeddings leaderboard](https://huggingface.co/spaces/mteb/leaderboard) for more recommendations.\n",
    "\n",
    "This class depends on the sentence-transformers package, which you can install with `pip install sentence-transformers`.\n",
    "\n",
    "NOTE: if you were previously using a `HuggingFaceEmbeddings` from LangChain, this should give equivalent results."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-embeddings-huggingface\n",
    "%pip install llama-index-embeddings-instructor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install llama-index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/loganm/miniconda3/envs/llama-index/lib/python3.11/site-packages/torch/cuda/__init__.py:546: UserWarning: Can't initialize NVML\n",
      "  warnings.warn(\"Can't initialize NVML\")\n"
     ]
    }
   ],
   "source": [
    "from llama_index.embeddings.huggingface import HuggingFaceEmbedding\n",
    "\n",
    "# loads BAAI/bge-small-en\n",
    "# embed_model = HuggingFaceEmbedding()\n",
    "\n",
    "# loads BAAI/bge-small-en-v1.5\n",
    "embed_model = HuggingFaceEmbedding(model_name=\"BAAI/bge-small-en-v1.5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Hello World!\n",
      "384\n",
      "[-0.030880315229296684, -0.11021008342504501, 0.3917851448059082, -0.35962796211242676, 0.22797748446464539]\n"
     ]
    }
   ],
   "source": [
    "embeddings = embed_model.get_text_embedding(\"Hello World!\")\n",
    "print(len(embeddings))\n",
    "print(embeddings[:5])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## InstructorEmbedding\n",
    "\n",
    "Instructor Embeddings are a class of embeddings specifically trained to augment their embeddings according to an instruction. By default, queries are given `query_instruction=\"Represent the question for retrieving supporting documents: \"` and text is given `text_instruction=\"Represent the document for retrieval: \"`.\n",
    "\n",
    "They rely on the `Instructor` and `SentenceTransformers` (version 2.2.2) pip package, which you can install with `pip install InstructorEmbedding` and `pip install -U sentence-transformers==2.2.2`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/loganm/miniconda3/envs/llama-index/lib/python3.11/site-packages/InstructorEmbedding/instructor.py:7: TqdmExperimentalWarning: Using `tqdm.autonotebook.tqdm` in notebook mode. Use `tqdm.tqdm` instead to force console mode (e.g. in jupyter console)\n",
      "  from tqdm.autonotebook import trange\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "load INSTRUCTOR_Transformer\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/loganm/miniconda3/envs/llama-index/lib/python3.11/site-packages/torch/cuda/__init__.py:546: UserWarning: Can't initialize NVML\n",
      "  warnings.warn(\"Can't initialize NVML\")\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "max_seq_length  512\n"
     ]
    }
   ],
   "source": [
    "from llama_index.embeddings.instructor import InstructorEmbedding\n",
    "\n",
    "embed_model = InstructorEmbedding(model_name=\"hkunlp/instructor-base\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "768\n",
      "[ 0.02155361 -0.06098218  0.01796207  0.05490903  0.01526906]\n"
     ]
    }
   ],
   "source": [
    "embeddings = embed_model.get_text_embedding(\"Hello World!\")\n",
    "print(len(embeddings))\n",
    "print(embeddings[:5])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## OptimumEmbedding\n",
    "\n",
    "Optimum in a HuggingFace library for exporting and running HuggingFace models in the ONNX format.\n",
    "\n",
    "You can install the dependencies with `pip install transformers optimum[exporters]`.\n",
    "\n",
    "First, we need to create the ONNX model. ONNX models provide improved inference speeds, and can be used across platforms (i.e. in TransformersJS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/loganm/miniconda3/envs/llama-index/lib/python3.11/site-packages/torch/cuda/__init__.py:546: UserWarning: Can't initialize NVML\n",
      "  warnings.warn(\"Can't initialize NVML\")\n",
      "Framework not specified. Using pt to export to ONNX.\n",
      "Using the export variant default. Available variants are:\n",
      "\t- default: The default ONNX variant.\n",
      "Using framework PyTorch: 2.0.1+cu117\n",
      "Overriding 1 configuration item(s)\n",
      "\t- use_cache -> False\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "============= Diagnostic Run torch.onnx.export version 2.0.1+cu117 =============\n",
      "verbose: False, log level: Level.ERROR\n",
      "======================= 0 NONE 0 NOTE 0 WARNING 0 ERROR ========================\n",
      "\n",


This guide shows you how to use embedding models from [LangChain](https://python.langchain.com/docs/integrations/text_embedding/).

If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙.

from langchain.embeddings import HuggingFaceEmbeddings

from llama_index.embeddings.langchain import LangchainEmbedding

stdout

768 [-0.005906202830374241, 0.04911914840340614, -0.04757878929376602, -0.04320324584841728, 0.02837090566754341, -0.017371710389852524, -0.04422023147344589, -0.019035547971725464, 0.04941621795296669, -0.03839121758937836]


import os
os.environ["OPENAI_API_KEY"] = "sk-..."

from llama_index.embeddings.openai import OpenAIEmbedding

## Using OpenAI `text-embedding-3-large` and `text-embedding-3-small`

Note, you may have to update your openai client: `pip install -U openai`

"Open AI new Embeddings models is great."

[-0.011500772088766098, 0.02457442320883274, -0.01760469563305378, -0.017763426527380943, 0.029841400682926178]

"Open AI new Embeddings models is awesome."

## Change the dimension of output embeddings

Note: Make sure you have the latest OpenAI client

"Open AI new Embeddings models with different dimensions is awesome."


{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c62c1447-0afb-4fad-8dc6-389949c3496e",
   "metadata": {},
   "source": [
    "# Chat Summary Memory Buffer\n",
    "In this demo, we use the new *ChatSummaryMemoryBuffer* to limit the chat history to a certain token length, and iteratively summarize all messages that do not fit in the memory buffer. This can be useful if you want to limit costs and latency (assuming the summarization prompt uses and generates fewer tokens than including the entire history). \n",
    "\n",
    "The original *ChatMemoryBuffer* gives you the option to truncate the history after a certain number of tokens, which is useful to limit costs and latency, but also removes potentially relevant information from the chat history. \n",
    "\n",
    "The newer *ChatSummaryMemoryBuffer* aims to makes this a bit more flexible, so the user has more control over which chat_history is retained."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c00a753b-df2c-4164-90c3-76b8a15f74c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-llms-openai\n",
    "%pip install llama-index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1f24186-b86e-4580-b7b4-072e719d424f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"YOUR_API_KEY\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db66e3bc-9791-497b-9e7d-386765dccf74",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.memory import ChatSummaryMemoryBuffer\n",
    "from llama_index.core.llms import ChatMessage, MessageRole\n",
    "from llama_index.llms.openai import OpenAI as OpenAiLlm\n",
    "import tiktoken"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68e26f76-f819-4e1a-bc47-f2ea855ee189",
   "metadata": {},
   "source": [
    "First, we simulate some chat history that will not fit in the memory buffer in its entirety."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6402621a-4131-465c-b92d-1d9a8e7ee985",
   "metadata": {},
   "outputs": [],
   "source": [
    "chat_history = [\n",
    "    ChatMessage(role=\"user\", content=\"What is LlamaIndex?\"),\n",
    "    ChatMessage(\n",
    "        role=\"assistant\",\n",
    "        content=\"LlamaaIndex is the leading data framework for building LLM applications\",\n",
    "    ),\n",
    "    ChatMessage(role=\"user\", content=\"Can you give me some more details?\"),\n",
    "    ChatMessage(\n",
    "        role=\"assistant\",\n",
    "        content=\"\"\"LlamaIndex is a framework for building context-augmented LLM applications. Context augmentation refers to any use case that applies LLMs on top of your private or domain-specific data. Some popular use cases include the following: \n",
    "        Question-Answering Chatbots (commonly referred to as RAG systems, which stands for \"Retrieval-Augmented Generation\"), Document Understanding and Extraction, Autonomous Agents that can perform research and take actions\n",
    "        LlamaIndex provides the tools to build any of these above use cases from prototype to production. The tools allow you to both ingest/process this data and implement complex query workflows combining data access with LLM prompting.\"\"\",\n",
    "    ),\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f057a791-9d8e-43e5-b40a-6675b28f6fd0",
   "metadata": {},
   "source": [
    "By supplying an *llm* and *token_limit* for summarization, we create a *ChatSummaryMemoryBuffer* instance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67dbb654-4a66-43cf-9f6a-daead87a1084",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = \"gpt-4-0125-preview\"\n",
    "summarizer_llm = OpenAiLlm(model_name=model, max_tokens=256)\n",
    "tokenizer_fn = tiktoken.encoding_for_model(model).encode\n",
    "memory = ChatSummaryMemoryBuffer.from_defaults(\n",
    "    chat_history=chat_history,\n",
    "    llm=summarizer_llm,\n",
    "    token_limit=2,\n",
    "    tokenizer_fn=tokenizer_fn,\n",
    ")\n",
    "\n",
    "history = memory.get()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e4e333c-6c33-4e01-b1a8-750d21800076",
   "metadata": {},
   "source": [
    "When printing the history, we can observe that older messages have been summarized."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0821eb89-4164-4a06-b66c-ea2632706e11",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ChatMessage(role=<MessageRole.SYSTEM: 'system'>, content='The user inquired about LlamaIndex, a leading data framework for developing LLM applications. The assistant explained that LlamaIndex is used for building context-augmented LLM applications, giving examples such as Question-Answering Chatbots, Document Understanding and Extraction, and Autonomous Agents. It was mentioned that LlamaIndex provides tools for ingesting and processing data, as well as implementing complex query workflows combining data access with LLM prompting.', additional_kwargs={})]\n"
     ]
    }
   ],
   "source": [
    "print(history)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fae3efe0-6889-49f7-9f21-6ced186f0609",
   "metadata": {},
   "source": [
    "Let's add some new chat history."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ddb295c-c5c0-4faf-b0e5-a451d1d26d60",
   "metadata": {},
   "outputs": [],
   "source": [
    "new_chat_history = [\n",
    "    ChatMessage(role=\"user\", content=\"Why context augmentation?\"),\n",
    "    ChatMessage(\n",
    "        role=\"assistant\",\n",
    "        content=\"LLMs offer a natural language interface between humans and data. Widely available models come pre-trained on huge amounts of publicly available data. However, they are not trained on your data, which may be private or specific to the problem you're trying to solve. It's behind APIs, in SQL databases, or trapped in PDFs and slide decks. LlamaIndex provides tooling to enable context augmentation. A popular example is Retrieval-Augmented Generation (RAG) which combines context with LLMs at inference time. Another is finetuning.\",\n",
    "    ),\n",
    "    ChatMessage(role=\"user\", content=\"Who is LlamaIndex for?\"),\n",
    "    ChatMessage(\n",
    "        role=\"assistant\",\n",
    "        content=\"LlamaIndex provides tools for beginners, advanced users, and everyone in between. Our high-level API allows beginner users to use LlamaIndex to ingest and query their data in 5 lines of code. For more complex applications, our lower-level APIs allow advanced users to customize and extend any module—data connectors, indices, retrievers, query engines, reranking modules—to fit their needs.\",\n",
    "    ),\n",
    "]\n",
    "memory.put(new_chat_history[0])\n",
    "memory.put(new_chat_history[1])\n",
    "memory.put(new_chat_history[2])\n",
    "memory.put(new_chat_history[3])\n",
    "history = memory.get()"
   ]
  },
  {
   "cell_type": "markdown",


{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7ae43f8b",
   "metadata": {},
   "source": [
    ""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "9c48213d-6e6a-4c10-838a-2a7c710c3a05",
   "metadata": {},
   "source": [
    "# ChatGPT"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf6108b0",
   "metadata": {},
   "source": [
    "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "590874bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-llms-openai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14790912",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install llama-index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "690a6918-7c75-4f95-9ccc-d2c4a1fe00d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import sys\n",
    "\n",
    "logging.basicConfig(stream=sys.stdout, level=logging.INFO)\n",
    "logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))\n",
    "\n",
    "from llama_index.core import VectorStoreIndex, SimpleDirectoryReader\n",
    "from llama_index.core import Settings\n",
    "from llama_index.llms.openai import OpenAI\n",
    "from IPython.display import Markdown, display"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9bc1c58d",
   "metadata": {},
   "source": [
    "#### Download Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18c8d851",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p 'data/paul_graham/'\n",
    "!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "50d3b817-b70e-4667-be4f-d3a0fe4bd119",
   "metadata": {},
   "source": [
    "#### Load documents, build the VectorStoreIndex"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03d1691e-544b-454f-825b-5ee12f7faa8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load documents\n",
    "documents = SimpleDirectoryReader(\"./data/paul_graham\").load_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6cc980e8-f4e1-4fad-93f8-ab1bbaa874f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# set global settings config\n",
    "llm = OpenAI(temperature=0, model=\"gpt-3.5-turbo\")\n",
    "Settings.llm = llm\n",
    "Settings.chunk_size = 512"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad144ee7-96da-4dd6-be00-fd6cf0c78e58",
   "metadata": {},
   "outputs": [],
   "source": [
    "index = VectorStoreIndex.from_documents(documents)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "b6caf93b-6345-4c65-a346-a95b0f1746c4",
   "metadata": {},
   "source": [
    "#### Query Index"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "83e2905e-3789-4793-82b9-0ac488246824",
   "metadata": {},
   "source": [
    "By default, with the help of langchain's PromptSelector abstraction, we use \n",
    "a modified refine prompt tailored for ChatGPT-use if the ChatGPT model is used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85466fdf-93f3-4cb1-a5f9-0056a8245a6f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:llama_index.token_counter.token_counter:> [retrieve] Total LLM token usage: 0 tokens\n",
      "> [retrieve] Total LLM token usage: 0 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [retrieve] Total embedding token usage: 8 tokens\n",
      "> [retrieve] Total embedding token usage: 8 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [get_response] Total LLM token usage: 0 tokens\n",
      "> [get_response] Total LLM token usage: 0 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [get_response] Total embedding token usage: 0 tokens\n",
      "> [get_response] Total embedding token usage: 0 tokens\n"
     ]
    }
   ],
   "source": [
    "query_engine = index.as_query_engine(\n",
    "    similarity_top_k=3,\n",
    "    streaming=True,\n",
    ")\n",
    "response = query_engine.query(\n",
    "    \"What did the author do growing up?\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0262da8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Before college, the author worked on writing short stories and programming on an IBM 1401 using an early version of Fortran. They also worked on programming with microcomputers and eventually created a new dialect of Lisp called Arc. They later realized the potential of publishing essays on the web and began writing and publishing them. The author also worked on spam filters, painting, and cooking for groups."
     ]
    }
   ],
   "source": [
    "response.print_response_stream()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec88df57",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:llama_index.token_counter.token_counter:> [retrieve] Total LLM token usage: 0 tokens\n",
      "> [retrieve] Total LLM token usage: 0 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [retrieve] Total embedding token usage: 12 tokens\n",
      "> [retrieve] Total embedding token usage: 12 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [get_response] Total LLM token usage: 0 tokens\n",
      "> [get_response] Total LLM token usage: 0 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [get_response] Total embedding token usage: 0 tokens\n",
      "> [get_response] Total embedding token usage: 0 tokens\n"
     ]
    }
   ],
   "source": [
    "query_engine = index.as_query_engine(\n",
    "    similarity_top_k=5,\n",
    "    streaming=True,\n",
    ")\n",
    "response = query_engine.query(\n",
    "    \"What did the author do during his time at RISD?\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,


{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9af12a30",
   "metadata": {},
   "source": [
    ""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "ec51f656",
   "metadata": {},
   "source": [
    "# Azure OpenAI"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef3d26db",
   "metadata": {},
   "source": [
    "Azure openAI resources unfortunately differ from standard openAI resources as you can't generate embeddings unless you use an embedding model. The regions where these models are available can be found here: https://learn.microsoft.com/en-us/azure/cognitive-services/openai/concepts/models#embeddings-models\n",
    "\n",
    "Furthermore the regions that support embedding models unfortunately don't support the latest versions (<*>-003) of openAI models, so we are forced to use one region for embeddings and another for the text generation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "710d6708",
   "metadata": {},
   "source": [
    "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94613811",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-embeddings-azure-openai\n",
    "%pip install llama-index-llms-azure-openai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de32bc14",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install llama-index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b05e71d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.azure_openai import AzureOpenAI\n",
    "from llama_index.embeddings.azure_openai import AzureOpenAIEmbedding\n",
    "from llama_index.core import VectorStoreIndex, SimpleDirectoryReader\n",
    "import logging\n",
    "import sys\n",
    "\n",
    "logging.basicConfig(\n",
    "    stream=sys.stdout, level=logging.INFO\n",
    ")  # logging.DEBUG for more verbose output\n",
    "logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fcaaafdf",
   "metadata": {},
   "source": [
    "Here, we setup the embedding model (for retrieval) and llm (for text generation).\n",
    "Note that you need not only model names (e.g. \"text-embedding-ada-002\"), but also model deployment names (the one you chose when deploying the model in Azure.\n",
    "You must pass the deployment name as a parameter when you initialize `AzureOpenAI` and `OpenAIEmbedding`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2569cb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "api_key = \"<api-key>\"\n",
    "azure_endpoint = \"https://<your-resource-name>.openai.azure.com/\"\n",
    "api_version = \"2023-07-01-preview\"\n",
    "\n",
    "llm = AzureOpenAI(\n",
    "    model=\"gpt-35-turbo-16k\",\n",
    "    deployment_name=\"my-custom-llm\",\n",
    "    api_key=api_key,\n",
    "    azure_endpoint=azure_endpoint,\n",
    "    api_version=api_version,\n",
    ")\n",
    "\n",
    "# You need to deploy your own embedding model as well as your own chat completion model\n",
    "embed_model = AzureOpenAIEmbedding(\n",
    "    model=\"text-embedding-ada-002\",\n",
    "    deployment_name=\"my-custom-embedding\",\n",
    "    api_key=api_key,\n",
    "    azure_endpoint=azure_endpoint,\n",
    "    api_version=api_version,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72aac5a6-495e-40d2-82d3-bda8688ae919",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import Settings\n",
    "\n",
    "Settings.llm = llm\n",
    "Settings.embed_model = embed_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1cf0e9c9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:httpx:HTTP Request: POST https://test-simon.openai.azure.com//openai/deployments/my-custom-embedding/embeddings?api-version=2023-07-01-preview \"HTTP/1.1 200 OK\"\n",
      "HTTP Request: POST https://test-simon.openai.azure.com//openai/deployments/my-custom-embedding/embeddings?api-version=2023-07-01-preview \"HTTP/1.1 200 OK\"\n",
      "HTTP Request: POST https://test-simon.openai.azure.com//openai/deployments/my-custom-embedding/embeddings?api-version=2023-07-01-preview \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: POST https://test-simon.openai.azure.com//openai/deployments/my-custom-embedding/embeddings?api-version=2023-07-01-preview \"HTTP/1.1 200 OK\"\n",
      "HTTP Request: POST https://test-simon.openai.azure.com//openai/deployments/my-custom-embedding/embeddings?api-version=2023-07-01-preview \"HTTP/1.1 200 OK\"\n",
      "HTTP Request: POST https://test-simon.openai.azure.com//openai/deployments/my-custom-embedding/embeddings?api-version=2023-07-01-preview \"HTTP/1.1 200 OK\"\n"
     ]
    }
   ],
   "source": [
    "documents = SimpleDirectoryReader(\n",
    "    input_files=[\"../../data/paul_graham/paul_graham_essay.txt\"]\n",
    ").load_data()\n",
    "index = VectorStoreIndex.from_documents(documents)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98d9d3fd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:httpx:HTTP Request: POST https://test-simon.openai.azure.com//openai/deployments/my-custom-embedding/embeddings?api-version=2023-07-01-preview \"HTTP/1.1 200 OK\"\n",
      "HTTP Request: POST https://test-simon.openai.azure.com//openai/deployments/my-custom-embedding/embeddings?api-version=2023-07-01-preview \"HTTP/1.1 200 OK\"\n",
      "HTTP Request: POST https://test-simon.openai.azure.com//openai/deployments/my-custom-embedding/embeddings?api-version=2023-07-01-preview \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: POST https://test-simon.openai.azure.com//openai/deployments/my-custom-llm/chat/completions?api-version=2023-07-01-preview \"HTTP/1.1 200 OK\"\n",
      "HTTP Request: POST https://test-simon.openai.azure.com//openai/deployments/my-custom-llm/chat/completions?api-version=2023-07-01-preview \"HTTP/1.1 200 OK\"\n",
      "HTTP Request: POST https://test-simon.openai.azure.com//openai/deployments/my-custom-llm/chat/completions?api-version=2023-07-01-preview \"HTTP/1.1 200 OK\"\n",


{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "30815a85",
   "metadata": {},
   "source": [
    "# LLM Pydantic Program - NVIDIA"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "311e16cb",
   "metadata": {},
   "source": [
    "This guide shows you how to generate structured data with our `LLMTextCompletionProgram`. Given an LLM as well as an output Pydantic class, generate a structured Pydantic object.\n",
    "\n",
    "In terms of the target object, you can choose to directly specify `output_cls`, or specify a `PydanticOutputParser` or any other BaseOutputParser that generates a Pydantic object.\n",
    "\n",
    "in the examples below, we show you different ways of extracting into the `Album` object (which can contain a list of Song objects)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0611198",
   "metadata": {},
   "source": [
    "## Extract into `Album` class\n",
    "\n",
    "This is a simple example of parsing an output into an `Album` schema, which can contain multiple songs.\n",
    "\n",
    "Just pass `Album` into the `output_cls` property on initialization of the `LLMTextCompletionProgram`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "511a8171",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-readers-file llama-index-embeddings-nvidia llama-index-llms-nvidia"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b029b7e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "# del os.environ['NVIDIA_API_KEY']  ## delete key and reset\n",
    "if os.environ.get(\"NVIDIA_API_KEY\", \"\").startswith(\"nvapi-\"):\n",
    "    print(\"Valid NVIDIA_API_KEY already in environment. Delete to reset\")\n",
    "else:\n",
    "    nvapi_key = getpass.getpass(\"NVAPI Key (starts with nvapi-): \")\n",
    "    assert nvapi_key.startswith(\n",
    "        \"nvapi-\"\n",
    "    ), f\"{nvapi_key[:5]}... is not a valid key\"\n",
    "    os.environ[\"NVIDIA_API_KEY\"] = nvapi_key"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7a83b49-5c34-45d5-8cf4-62f348fb1299",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydantic import BaseModel\n",
    "from typing import List\n",
    "from llama_index.core import Settings\n",
    "from llama_index.llms.nvidia import NVIDIA\n",
    "from llama_index.embeddings.nvidia import NVIDIAEmbedding\n",
    "from llama_index.core.program import LLMTextCompletionProgram\n",
    "from llama_index.core.program import FunctionCallingProgram"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e4fc4b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = NVIDIA()\n",
    "\n",
    "embedder = NVIDIAEmbedding(model=\"NV-Embed-QA\", truncate=\"END\")\n",
    "Settings.embed_model = embedder\n",
    "Settings.llm = llm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d92e739",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Song(BaseModel):\n",
    "    \"\"\"Data model for a song.\"\"\"\n",
    "\n",
    "    title: str\n",
    "    length_seconds: int\n",
    "\n",
    "\n",
    "class Album(BaseModel):\n",
    "    \"\"\"Data model for an album.\"\"\"\n",
    "\n",
    "    name: str\n",
    "    artist: str\n",
    "    songs: List[Song]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46c2d509",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_template_str = \"\"\"\\\n",
    "Generate an example album, with an artist and a list of songs. \\\n",
    "Using the movie {movie_name} as inspiration.\\\n",
    "\"\"\"\n",
    "program = LLMTextCompletionProgram.from_defaults(\n",
    "    output_cls=Album,\n",
    "    prompt_template_str=prompt_template_str,\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "498370f4",
   "metadata": {},
   "source": [
    "Run program to get structured output.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca490bf8",
   "metadata": {},
   "outputs": [],
   "source": [
    "output = program(movie_name=\"The Shining\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40cce83f",
   "metadata": {},
   "source": [
    "The output is a valid Pydantic object that we can then use to call functions/APIs. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53934d3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6401ab8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.output_parsers import PydanticOutputParser\n",
    "\n",
    "program = LLMTextCompletionProgram.from_defaults(\n",
    "    output_parser=PydanticOutputParser(output_cls=Album),\n",
    "    prompt_template_str=prompt_template_str,\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1adc5b2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "output = program(movie_name=\"Lord of the Rings\")\n",
    "output"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a41391d9",
   "metadata": {},
   "source": [
    "## Define a Custom Output Parser\n",
    "\n",
    "Sometimes you may want to parse an output your own way into a JSON object. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60b7b669",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.output_parsers import ChainableOutputParser\n",
    "\n",
    "\n",
    "class CustomAlbumOutputParser(ChainableOutputParser):\n",
    "    \"\"\"Custom Album output parser.\n",
    "\n",
    "    Assume first line is name and artist.\n",
    "\n",
    "    Assume each subsequent line is the song.\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, verbose: bool = False):\n",
    "        self.verbose = verbose\n",
    "\n",
    "    def parse(self, output: str) -> Album:\n",
    "        \"\"\"Parse output.\"\"\"\n",
    "        if self.verbose:\n",
    "            print(f\"> Raw output: {output}\")\n",
    "        lines = output.split(\"\\n\")\n",
    "        lines = list(filter(None, (line.strip() for line in lines)))\n",
    "        name, artist = lines[1].split(\",\")\n",
    "        songs = []\n",
    "        for i in range(2, len(lines)):\n",
    "            title, length_seconds = lines[i].split(\",\")\n",
    "            songs.append(Song(title=title, length_seconds=length_seconds))\n",
    "\n",
    "        return Album(name=name, artist=artist, songs=songs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,


"Cell \u001b[0;32mIn[31], line 5\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mllama_index\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mvector_stores\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mchroma\u001b[39;00m \u001b[39mimport\u001b[39;00m ChromaVectorStore\n\u001b[1;32m      3\u001b[0m \u001b[39mimport\u001b[39;00m \u001b[39mchromadb\u001b[39;00m\n\u001b[0;32m----> 5\u001b[0m db \u001b[39m=\u001b[39m chromadb\u001b[39m.\u001b[39;49mPersistentClient(path\u001b[39m=\u001b[39;49m\u001b[39m\"\u001b[39;49m\u001b[39m./chroma_db2\u001b[39;49m\u001b[39m\"\u001b[39;49m)\n\u001b[1;32m      6\u001b[0m chroma_collection \u001b[39m=\u001b[39m db\u001b[39m.\u001b[39mget_or_create_collection(\u001b[39m\"\u001b[39m\u001b[39mquickstart2\u001b[39m\u001b[39m\"\u001b[39m)\n\u001b[1;32m      7\u001b[0m vector_store \u001b[39m=\u001b[39m ChromaVectorStore(chroma_collection\u001b[39m=\u001b[39mchroma_collection)\n",
      "File \u001b[0;32m~/giant_change/llama_index/venv/lib/python3.10/site-packages/chromadb/__init__.py:146\u001b[0m, in \u001b[0;36mPersistentClient\u001b[0;34m(path, settings, tenant, database)\u001b[0m\n\u001b[1;32m    143\u001b[0m tenant \u001b[39m=\u001b[39m \u001b[39mstr\u001b[39m(tenant)\n\u001b[1;32m    144\u001b[0m database \u001b[39m=\u001b[39m \u001b[39mstr\u001b[39m(database)\n\u001b[0;32m--> 146\u001b[0m \u001b[39mreturn\u001b[39;00m ClientCreator(tenant\u001b[39m=\u001b[39;49mtenant, database\u001b[39m=\u001b[39;49mdatabase, settings\u001b[39m=\u001b[39;49msettings)\n",
      "File \u001b[0;32m~/giant_change/llama_index/venv/lib/python3.10/site-packages/chromadb/api/client.py:139\u001b[0m, in \u001b[0;36mClient.__init__\u001b[0;34m(self, tenant, database, settings)\u001b[0m\n\u001b[1;32m    133\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39m__init__\u001b[39m(\n\u001b[1;32m    134\u001b[0m     \u001b[39mself\u001b[39m,\n\u001b[1;32m    135\u001b[0m     tenant: \u001b[39mstr\u001b[39m \u001b[39m=\u001b[39m DEFAULT_TENANT,\n\u001b[1;32m    136\u001b[0m     database: \u001b[39mstr\u001b[39m \u001b[39m=\u001b[39m DEFAULT_DATABASE,\n\u001b[1;32m    137\u001b[0m     settings: Settings \u001b[39m=\u001b[39m Settings(),\n\u001b[1;32m    138\u001b[0m ) \u001b[39m-\u001b[39m\u001b[39m>\u001b[39m \u001b[39mNone\u001b[39;00m:\n\u001b[0;32m--> 139\u001b[0m     \u001b[39msuper\u001b[39;49m()\u001b[39m.\u001b[39;49m\u001b[39m__init__\u001b[39;49m(settings\u001b[39m=\u001b[39;49msettings)\n\u001b[1;32m    140\u001b[0m     \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mtenant \u001b[39m=\u001b[39m tenant\n\u001b[1;32m    141\u001b[0m     \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mdatabase \u001b[39m=\u001b[39m database\n",
      "File \u001b[0;32m~/giant_change/llama_index/venv/lib/python3.10/site-packages/chromadb/api/client.py:43\u001b[0m, in \u001b[0;36mSharedSystemClient.__init__\u001b[0;34m(self, settings)\u001b[0m\n\u001b[1;32m     38\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39m__init__\u001b[39m(\n\u001b[1;32m     39\u001b[0m     \u001b[39mself\u001b[39m,\n\u001b[1;32m     40\u001b[0m     settings: Settings \u001b[39m=\u001b[39m Settings(),\n\u001b[1;32m     41\u001b[0m ) \u001b[39m-\u001b[39m\u001b[39m>\u001b[39m \u001b[39mNone\u001b[39;00m:\n\u001b[1;32m     42\u001b[0m     \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_identifier \u001b[39m=\u001b[39m SharedSystemClient\u001b[39m.\u001b[39m_get_identifier_from_settings(settings)\n\u001b[0;32m---> 43\u001b[0m     SharedSystemClient\u001b[39m.\u001b[39;49m_create_system_if_not_exists(\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_identifier, settings)\n",


"File \u001b[0;32m~/giant_change/llama_index/venv/lib/python3.10/site-packages/chromadb/config.py:382\u001b[0m, in \u001b[0;36mSystem.instance\u001b[0;34m(self, type)\u001b[0m\n\u001b[1;32m    379\u001b[0m     \u001b[39mtype\u001b[39m \u001b[39m=\u001b[39m get_class(fqn, \u001b[39mtype\u001b[39m)\n\u001b[1;32m    381\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mtype\u001b[39m \u001b[39mnot\u001b[39;00m \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_instances:\n\u001b[0;32m--> 382\u001b[0m     impl \u001b[39m=\u001b[39m \u001b[39mtype\u001b[39;49m(\u001b[39mself\u001b[39;49m)\n\u001b[1;32m    383\u001b[0m     \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_instances[\u001b[39mtype\u001b[39m] \u001b[39m=\u001b[39m impl\n\u001b[1;32m    384\u001b[0m     \u001b[39mif\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_running:\n",
      "File \u001b[0;32m~/giant_change/llama_index/venv/lib/python3.10/site-packages/chromadb/db/impl/sqlite.py:88\u001b[0m, in \u001b[0;36mSqliteDB.__init__\u001b[0;34m(self, system)\u001b[0m\n\u001b[1;32m     84\u001b[0m     \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_db_file \u001b[39m=\u001b[39m (\n\u001b[1;32m     85\u001b[0m         \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_settings\u001b[39m.\u001b[39mrequire(\u001b[39m\"\u001b[39m\u001b[39mpersist_directory\u001b[39m\u001b[39m\"\u001b[39m) \u001b[39m+\u001b[39m \u001b[39m\"\u001b[39m\u001b[39m/chroma.sqlite3\u001b[39m\u001b[39m\"\u001b[39m\n\u001b[1;32m     86\u001b[0m     )\n\u001b[1;32m     87\u001b[0m     \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m os\u001b[39m.\u001b[39mpath\u001b[39m.\u001b[39mexists(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_db_file):\n\u001b[0;32m---> 88\u001b[0m         os\u001b[39m.\u001b[39;49mmakedirs(os\u001b[39m.\u001b[39;49mpath\u001b[39m.\u001b[39;49mdirname(\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_db_file), exist_ok\u001b[39m=\u001b[39;49m\u001b[39mTrue\u001b[39;49;00m)\n\u001b[1;32m     89\u001b[0m     \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_conn_pool \u001b[39m=\u001b[39m PerThreadPool(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_db_file)\n\u001b[1;32m     90\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_tx_stack \u001b[39m=\u001b[39m local()\n",
      "File \u001b[0;32m~/miniforge3/lib/python3.10/os.py:225\u001b[0m, in \u001b[0;36mmakedirs\u001b[0;34m(name, mode, exist_ok)\u001b[0m\n\u001b[1;32m    223\u001b[0m         \u001b[39mreturn\u001b[39;00m\n\u001b[1;32m    224\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[0;32m--> 225\u001b[0m     mkdir(name, mode)\n\u001b[1;32m    226\u001b[0m \u001b[39mexcept\u001b[39;00m \u001b[39mOSError\u001b[39;00m:\n\u001b[1;32m    227\u001b[0m     \u001b[39m# Cannot rely on checking for EEXIST, since the operating system\u001b[39;00m\n\u001b[1;32m    228\u001b[0m     \u001b[39m# could give priority to other errors like EACCES or EROFS\u001b[39;00m\n\u001b[1;32m    229\u001b[0m     \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m exist_ok \u001b[39mor\u001b[39;00m \u001b[39mnot\u001b[39;00m path\u001b[39m.\u001b[39misdir(name):\n",
      "\u001b[0;31mFileNotFoundError\u001b[0m: [Errno 2] No such file or directory: './chroma_db2'"
     ]
    }
   ],
   "source": [
    "from llama_index.core import StorageContext, VectorStoreIndex\n",
    "from llama_index.vector_stores.chroma import ChromaVectorStore\n",
    "import chromadb\n",
    "\n",
    "db = chromadb.PersistentClient(path=\"./chroma_db\")\n",
    "chroma_collection = db.get_or_create_collection(\"quickstart2\")\n",
    "vector_store = ChromaVectorStore(chroma_collection=chroma_collection)\n",
    "storage_context = StorageContext.from_defaults(vector_store=vector_store)\n",
    "\n",
    "object_index = ObjectIndex.from_objects(\n",
    "    arbitrary_objects,\n",
    "    index_cls=VectorStoreIndex,\n",
    "    storage_context=storage_context,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28cda697",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['llamaindex is an awesome library!']"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "object_retriever = object_index.as_retriever(similarity_top_k=1)\n",
    "object_retriever.retrieve(\"llamaindex\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "358994af",
   "metadata": {},
   "source": [
    "Now, lets \"reload\" the index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61134380",


{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Controlling Agent Reasoning Loop with Return Direct Tools\n",
    "\n",
    "All tools have an option for `return_direct` -- if this is set to `True`, and the associated tool is called (without any other tools being called), the agent reasoning loop is ended and the tool output is returned directly.\n",
    "\n",
    "This can be useful for speeding up response times when you know the tool output is good enough, to avoid the agent re-writing the response, and for ending the reasoning loop."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook walks through a notebook where an agent needs to gather information from a user in order to make a restaurant booking."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install llama-index-core llama-index-llms-anthropic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"ANTHROPIC_API_KEY\"] = \"sk-ant-...\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tools setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Optional\n",
    "\n",
    "from llama_index.core.tools import FunctionTool\n",
    "from llama_index.core.bridge.pydantic import BaseModel\n",
    "\n",
    "# we will store booking under random IDs\n",
    "bookings = {}\n",
    "\n",
    "\n",
    "# we will represent and track the state of a booking as a Pydantic model\n",
    "class Booking(BaseModel):\n",
    "    name: Optional[str] = None\n",
    "    email: Optional[str] = None\n",
    "    phone: Optional[str] = None\n",
    "    date: Optional[str] = None\n",
    "    time: Optional[str] = None\n",
    "\n",
    "\n",
    "def get_booking_state(user_id: str) -> str:\n",
    "    \"\"\"Get the current state of a booking for a given booking ID.\"\"\"\n",
    "    try:\n",
    "        return str(bookings[user_id].dict())\n",
    "    except:\n",
    "        return f\"Booking ID {user_id} not found\"\n",
    "\n",
    "\n",
    "def update_booking(user_id: str, property: str, value: str) -> str:\n",
    "    \"\"\"Update a property of a booking for a given booking ID. Only enter details that are explicitly provided.\"\"\"\n",
    "    booking = bookings[user_id]\n",
    "    setattr(booking, property, value)\n",
    "    return f\"Booking ID {user_id} updated with {property} = {value}\"\n",
    "\n",
    "\n",
    "def create_booking(user_id: str) -> str:\n",
    "    \"\"\"Create a new booking and return the booking ID.\"\"\"\n",
    "    bookings[user_id] = Booking()\n",
    "    return \"Booking created, but not yet confirmed. Please provide your name, email, phone, date, and time.\"\n",
    "\n",
    "\n",
    "def confirm_booking(user_id: str) -> str:\n",
    "    \"\"\"Confirm a booking for a given booking ID.\"\"\"\n",
    "    booking = bookings[user_id]\n",
    "\n",
    "    if booking.name is None:\n",
    "        raise ValueError(\"Please provide your name.\")\n",
    "\n",
    "    if booking.email is None:\n",
    "        raise ValueError(\"Please provide your email.\")\n",
    "\n",
    "    if booking.phone is None:\n",
    "        raise ValueError(\"Please provide your phone number.\")\n",
    "\n",
    "    if booking.date is None:\n",
    "        raise ValueError(\"Please provide the date of your booking.\")\n",
    "\n",
    "    if booking.time is None:\n",
    "        raise ValueError(\"Please provide the time of your booking.\")\n",
    "\n",
    "    return f\"Booking ID {user_id} confirmed!\"\n",
    "\n",
    "\n",
    "# create tools for each function\n",
    "get_booking_state_tool = FunctionTool.from_defaults(fn=get_booking_state)\n",
    "update_booking_tool = FunctionTool.from_defaults(fn=update_booking)\n",
    "create_booking_tool = FunctionTool.from_defaults(\n",
    "    fn=create_booking, return_direct=True\n",
    ")\n",
    "confirm_booking_tool = FunctionTool.from_defaults(\n",
    "    fn=confirm_booking, return_direct=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A user has walked in! Let's help them make a booking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.anthropic import Anthropic\n",
    "from llama_index.core.llms import ChatMessage\n",
    "from llama_index.core.agent import FunctionCallingAgent\n",
    "\n",
    "llm = Anthropic(model=\"claude-3-sonnet-20240229\", temperature=0.1)\n",
    "\n",
    "user = \"user123\"\n",
    "prefix_messages = [\n",
    "    ChatMessage(\n",
    "        role=\"system\",\n",
    "        content=(\n",
    "            f\"You are now connected to the booking system and helping {user} with making a booking. \"\n",
    "            \"Only enter details that the user has explicitly provided. \"\n",
    "            \"Do not make up any details.\"\n",
    "        ),\n",
    "    )\n",
    "]\n",
    "\n",
    "agent = FunctionCallingAgent.from_tools(\n",
    "    tools=[\n",
    "        get_booking_state_tool,\n",
    "        update_booking_tool,\n",
    "        create_booking_tool,\n",
    "        confirm_booking_tool,\n",
    "    ],\n",
    "    llm=llm,\n",
    "    prefix_messages=prefix_messages,\n",
    "    max_function_calls=10,\n",
    "    allow_parallel_tool_calls=False,\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Added user message to memory: Hello! I would like to make a booking, around 5pm?\n",
      "=== LLM Response ===\n",
      "Okay, let's create a new booking for you. To do that, I'll invoke the create_booking tool with your user ID:\n",
      "=== Calling Function ===\n",
      "Calling function: create_booking with args: {\"user_id\": \"user123\"}\n",
      "=== Function Output ===\n",
      "Booking created, but not yet confirmed. Please provide your name, email, phone, date, and time.\n"
     ]
    }
   ],
   "source": [
    "response = agent.chat(\"Hello! I would like to make a booking, around 5pm?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Booking created, but not yet confirmed. Please provide your name, email, phone, date, and time.\n"
     ]
    }
   ],
   "source": [
    "print(str(response))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Perfect, we can see the function output was retruned directly, with no modification or final LLM call!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Added user message to memory: Sure! My name is Logan, and my email is test@gmail.com\n",
      "=== LLM Response ===\n",


{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# NebulaGraph Property Graph Index\n",
    "\n",
    "NebulaGraph is an open-source distributed graph database built for super large-scale graphs with milliseconds of latency.\n",
    "\n",
    "If you already have an existing graph, please skip to the end of this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index llama-index-graph-stores-nebula jupyter-nebulagraph"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Docker Setup\n",
    "\n",
    "To launch NebulaGraph locally, first ensure you have docker installed. Then, you can launch the database with the following docker command.\n",
    "\n",
    "```bash\n",
    "mkdir nebula-docker-compose\n",
    "cd nebula-docker-compose\n",
    "curl --output docker-compose.yaml https://raw.githubusercontent.com/vesoft-inc/nebula-docker-compose/master/docker-compose-lite.yaml\n",
    "docker compose up \n",
    "```\n",
    "\n",
    "After this, you are ready to create your first property graph!\n",
    "\n",
    "> Other options/details for deploying NebulaGraph can be found in the [docs](https://docs.nebula-graph.io/):\n",
    ">\n",
    "> - [ad-hoc cluster in Google Colab](https://docs.nebula-graph.io/master/4.deployment-and-installation/2.compile-and-install-nebula-graph/8.deploy-nebula-graph-with-lite/).\n",
    "> - [Docker Desktop Extension](https://docs.nebula-graph.io/master/2.quick-start/1.quick-start-workflow/).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;3;38;2;47;75;124mConnection Pool Created\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "
\n",
       "\n",
       "
\n",
       "
\n",
       "
\n",
       "      	\n",
       "
\n",
       "
\n",
       "
\n",
       "
\n",
       "
\n",
       "
"
      ],
      "text/plain": [
       "Empty DataFrame\n",
       "Columns: []\n",
       "Index: []"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# load NebulaGraph Jupyter extension to enable %ngql magic\n",
    "%load_ext ngql\n",
    "# connect to NebulaGraph service\n",
    "%ngql --address 127.0.0.1 --port 9669 --user root --password nebula\n",
    "# create a graph space(think of a Database Instance) named: llamaindex_nebula_property_graph\n",
    "%ngql CREATE SPACE IF NOT EXISTS llamaindex_nebula_property_graph(vid_type=FIXED_STRING(256));"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "
\n",
       "\n",
       "
\n",
       "
\n",
       "
\n",
       "      	\n",
       "
\n",
       "
\n",
       "
\n",
       "
\n",
       "
\n",
       "
"
      ],
      "text/plain": [
       "Empty DataFrame\n",
       "Columns: []\n",
       "Index: []"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# use the graph space, which is similar to \"use database\" in MySQL\n",
    "# The space was created in async way, so we need to wait for a while before using it, retry it if failed\n",
    "%ngql USE llamaindex_nebula_property_graph;"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Env Setup\n",
    "\n",
    "We need just a few environment setups to get started."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"sk-proj-...\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p 'data/paul_graham/'\n",
    "!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import SimpleDirectoryReader\n",
    "\n",
    "documents = SimpleDirectoryReader(\"./data/paul_graham/\").load_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We choose using gpt-4o and local embedding model intfloat/multilingual-e5-large . You can change to what you like, by editing the following lines:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-embeddings-huggingface"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import Settings\n",
    "from llama_index.llms.openai import OpenAI\n",
    "from llama_index.embeddings.openai import OpenAIEmbedding\n",
    "from llama_index.embeddings.huggingface import HuggingFaceEmbedding\n",
    "\n",
    "Settings.llm = OpenAI(model=\"gpt-4o\", temperature=0.3)\n",
    "Settings.embed_model = HuggingFaceEmbedding(\n",
    "    model_name=\"intfloat/multilingual-e5-large\"\n",
    ")\n",
    "# Settings.embed_model = OpenAIEmbedding(model_name=\"text-embedding-3-small\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Index Construction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Prepare property graph store"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [


"[NodeWithScore(node=TextNode(id_='d18bb1f1-466a-443d-98d9-6217bf71ee5a', embedding=None, metadata={'filename': 'README.md', 'category': 'codebase'}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={<NodeRelationship.SOURCE: '1'>: RelatedNodeInfo(node_id='e4c638ce-6757-482e-baed-096574550602', node_type=<ObjectType.DOCUMENT: '4'>, metadata={'filename': 'README.md', 'category': 'codebase'}, hash='3183371414f6a23e9a61e11b45ec45f808b148f9973166cfed62226e3505eb05')}, text='Context\\nLLMs are a phenomenal piece of technology for knowledge generation and reasoning.\\nThey are pre-trained on large amounts of publicly available data.\\nHow do we best augment LLMs with our own private data?\\nWe need a comprehensive toolkit to help perform this data augmentation for LLMs.\\n\\nProposed Solution\\nThat\\'s where LlamaIndex comes in. LlamaIndex is a \"data framework\" to help\\nyou build LLM  apps. It provides the following tools:\\n\\nOffers data connectors to ingest your existing data sources and data formats\\n(APIs, PDFs, docs, SQL, etc.)\\nProvides ways to structure your data (indices, graphs) so that this data can be\\neasily used with LLMs.\\nProvides an advanced retrieval/query interface over your data:\\nFeed in any LLM input prompt, get back retrieved context and knowledge-augmented output.\\nAllows easy integrations with your outer application framework\\n(e.g. with LangChain, Flask, Docker, ChatGPT, anything else).\\nLlamaIndex provides tools for both beginner users and advanced users.\\nOur high-level API allows beginner users to use LlamaIndex to ingest and\\nquery their data in 5 lines of code. Our lower-level APIs allow advanced users to\\ncustomize and extend any module (data connectors, indices, retrievers, query engines,\\nreranking modules), to fit their needs.', mimetype='text/plain', start_char_idx=1, end_char_idx=1279, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n'), score=0.850998849877966)]\n"
     ]
    }
   ],
   "source": [
    "print(\n",
    "    \"------------- Example Document used to Enrich LLM Context -------------\"\n",
    ")\n",
    "llama_index_example_document = Document.example()\n",
    "print(llama_index_example_document)\n",
    "\n",
    "index = VectorStoreIndex.from_documents([llama_index_example_document])\n",
    "\n",
    "print(\"\\n------------- Example Query Engine -------------\")\n",
    "query_response = index.as_query_engine().query(\"What is llama_index?\")\n",
    "print(query_response)\n",
    "\n",
    "print(\"\\n------------- Example Chat Engine  -------------\")\n",
    "chat_response = index.as_chat_engine().chat(\n",
    "    \"What is llama_index?\",\n",
    "    chat_history=[\n",
    "        ChatMessage(role=\"system\", content=\"You are an expert on RAG!\")\n",
    "    ],\n",
    ")\n",
    "print(chat_response)\n",
    "\n",
    "\n",
    "print(\"\\n------------- Example Retriever   -------------\")\n",
    "retriever_response = index.as_retriever().retrieve(\"What is llama_index?\")\n",
    "print(retriever_response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Log the Index with MLflow\n",
    "\n",
    "The below code logs a LlamaIndex model with MLflow, allowing you to persist and manage it across different environments. By using MLflow, you can track, version, and reproduce your model reliably. The script logs parameters, an example input, and registers the model under a specific name. The `model_uri` provides a unique identifier for retrieving the model later. This persistence is essential for ensuring consistency and reproducibility in development, testing, and production. Managing the model with MLflow simplifies loading, deployment, and sharing, maintaining an organized workflow.\n",
    "\n",
    "Key Parameters\n",
    "\n",
    "* ``engine_type``: defines the pyfunc and spark_udf inference type\n",
    "* ``input_example``: defines the the input signature and infers the output signature via a prediction\n",
    "* ``registered_model_name``: defines the name of the model in the MLflow model registry"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024/07/24 17:58:27 INFO mlflow.llama_index.serialize_objects: API key(s) will be removed from the global Settings object during serialization to protect against key leakage. At inference time, the key(s) must be passed as environment variables.\n",
      "/Users/michael.berk/opt/anaconda3/envs/mlflow-dev/lib/python3.8/site-packages/_distutils_hack/__init__.py:26: UserWarning: Setuptools is replacing distutils.\n",
      "  warnings.warn(\"Setuptools is replacing distutils.\")\n",
      "Successfully registered model 'my_llama_index_vector_store'.\n",
      "Created version '1' of model 'my_llama_index_vector_store'.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "643e7b6936674e469f98d94004f3424a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading artifacts:   0%|          | 0/12 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unique identifier for the model location for loading: runs:/036936a7ac964f0cb6ab99fa908d6421/llama_index\n"
     ]
    }
   ],
   "source": [
    "mlflow.llama_index.autolog()  # This is for enabling tracing\n",
    "\n",
    "with mlflow.start_run() as run:\n",
    "    mlflow.llama_index.log_model(\n",
    "        index,\n",
    "        artifact_path=\"llama_index\",\n",
    "        engine_type=\"query\",  # Defines the pyfunc and spark_udf inference type\n",
    "        input_example=\"hi\",  # Infers signature\n",
    "        registered_model_name=\"my_llama_index_vector_store\",  # Stores an instance in the model registry\n",
    "    )\n",
    "\n",
    "    run_id = run.info.run_id\n",
    "    model_uri = f\"runs:/{run_id}/llama_index\"\n",
    "    print(f\"Unique identifier for the model location for loading: {model_uri}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load the Index and Perform Inference\n",
    "\n",
    "The below code demonstrates three core types of inference that can be done with the loaded model.\n",
    "\n",
    "1. **Load and Perform Inference via LlamaIndex:** This method loads the model using `mlflow.llama_index.load_model` and performs direct querying, chat, or retrieval. It is ideal when you want to leverage the full capabilities of the underlying llama index object.\n",
    "2. **Load and Perform Inference via MLflow PyFunc:** This method loads the model using `mlflow.pyfunc.load_model`, enabling model predictions in a generic PyFunc format, with the engine type specified at logging time. It is useful for evaluating the model with `mlflow.evaluate` or deploying the model for serving. \n",
    "3. **Load and Perform Inference via MLflow Spark UDF:** This method uses `mlflow.pyfunc.spark_udf` to load the model as a Spark UDF, facilitating distributed inference across large datasets in a Spark DataFrame. It is ideal for handling large-scale data processing and, like with PyFunc inference, only supports the engine type defined when logging.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [


{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "2ffc6a2b",
   "metadata": {},
   "source": [
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dff7db9e-fbf9-4394-9958-35323799a4e3",
   "metadata": {},
   "source": [
    "# Building RAG from Scratch (Open-source only!) \n",
    "\n",
    "In this tutorial, we show you how to build a data ingestion pipeline into a vector database, and then build a retrieval pipeline from that vector database, from scratch.\n",
    "\n",
    "Notably, we use a fully open-source stack:\n",
    "\n",
    "- Sentence Transformers as the embedding model\n",
    "- Postgres as the vector store (we support many other [vector stores](https://gpt-index.readthedocs.io/en/stable/module_guides/storing/vector_stores.html) too!)\n",
    "- Llama 2 as the LLM (through [llama.cpp](https://github.com/ggerganov/llama.cpp))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25764729-40ba-400f-b0f8-08fb9e8bb74a",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "We setup our open-source components.\n",
    "1. Sentence Transformers\n",
    "2. Llama 2\n",
    "3. We initialize postgres and wrap it with our wrappers/abstractions."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63935557-a11c-4a22-9248-9c746cc89c4c",
   "metadata": {},
   "source": [
    "#### Sentence Transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7108d3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-readers-file pymupdf\n",
    "%pip install llama-index-vector-stores-postgres\n",
    "%pip install llama-index-embeddings-huggingface\n",
    "%pip install llama-index-llms-llama-cpp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c08162e-5a48-424c-921f-c9e84a59c72f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# sentence transformers\n",
    "from llama_index.embeddings.huggingface import HuggingFaceEmbedding\n",
    "\n",
    "embed_model = HuggingFaceEmbedding(model_name=\"BAAI/bge-small-en\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df10089c-917e-4191-a718-0ef7149a6a1e",
   "metadata": {},
   "source": [
    "#### Llama CPP\n",
    "\n",
    "In this notebook, we use the [`llama-2-chat-13b-ggml`](https://huggingface.co/TheBloke/Llama-2-13B-chat-GGML) model, along with the proper prompt formatting. \n",
    "\n",
    "Check out our [Llama CPP guide](https://gpt-index.readthedocs.io/en/stable/examples/llm/llama_2_llama_cpp.html) for full setup instructions/details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85f8a556-9f37-42a3-a88a-f688ad355ee5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: llama-cpp-python in /Users/jerryliu/Programming/gpt_index/.venv/lib/python3.10/site-packages (0.2.7)\n",
      "Requirement already satisfied: numpy>=1.20.0 in /Users/jerryliu/Programming/gpt_index/.venv/lib/python3.10/site-packages (from llama-cpp-python) (1.23.5)\n",
      "Requirement already satisfied: typing-extensions>=4.5.0 in /Users/jerryliu/Programming/gpt_index/.venv/lib/python3.10/site-packages (from llama-cpp-python) (4.7.1)\n",
      "Requirement already satisfied: diskcache>=5.6.1 in /Users/jerryliu/Programming/gpt_index/.venv/lib/python3.10/site-packages (from llama-cpp-python) (5.6.3)\n",
      "\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip available: \u001b[0m\u001b[31;49m22.3.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m23.2.1\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "!pip install llama-cpp-python"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3cb975f6-c192-4a26-ae50-e9a319d2a66b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.llama_cpp import LlamaCPP\n",
    "\n",
    "# model_url = \"https://huggingface.co/TheBloke/Llama-2-13B-chat-GGML/resolve/main/llama-2-13b-chat.ggmlv3.q4_0.bin\"\n",
    "model_url = \"https://huggingface.co/TheBloke/Llama-2-13B-chat-GGUF/resolve/main/llama-2-13b-chat.Q4_0.gguf\"\n",
    "\n",
    "llm = LlamaCPP(\n",
    "    # You can pass in the URL to a GGML model to download it automatically\n",
    "    model_url=model_url,\n",
    "    # optionally, you can set the path to a pre-downloaded model instead of model_url\n",
    "    model_path=None,\n",
    "    temperature=0.1,\n",
    "    max_new_tokens=256,\n",
    "    # llama2 has a context window of 4096 tokens, but we set it lower to allow for some wiggle room\n",
    "    context_window=3900,\n",
    "    # kwargs to pass to __call__()\n",
    "    generate_kwargs={},\n",
    "    # kwargs to pass to __init__()\n",
    "    # set to at least 1 to use GPU\n",
    "    model_kwargs={\"n_gpu_layers\": 1},\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba02cfe2-8b51-4e01-a840-d6508c76ade3",
   "metadata": {},
   "source": [
    "#### Initialize Postgres\n",
    "\n",
    "Using an existing postgres running at localhost, create the database we'll be using.\n",
    "\n",
    "**NOTE**: Of course there are plenty of other open-source/self-hosted databases you can use! e.g. Chroma, Qdrant, Weaviate, and many more. Take a look at our [vector store guide](https://gpt-index.readthedocs.io/en/stable/module_guides/storing/vector_stores.html).\n",
    "\n",
    "**NOTE**: You will need to setup postgres on your local system. Here's an example of how to set it up on OSX: https://www.sqlshack.com/setting-up-a-postgresql-database-on-mac/.\n",
    "\n",


{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "57c676db",
   "metadata": {},
   "source": [
    ""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c919f307-07b1-41bd-bc5d-51edd8677983",
   "metadata": {},
   "source": [
    "# Building Data Ingestion from Scratch\n",
    "\n",
    "In this tutorial, we show you how to build a data ingestion pipeline into a vector database.\n",
    "\n",
    "We use Pinecone as the vector database.\n",
    "\n",
    "We will show how to do the following:\n",
    "1. How to load in documents.\n",
    "2. How to use a text splitter to split documents.\n",
    "3. How to **manually** construct nodes from each text chunk.\n",
    "4. [Optional] Add metadata to each Node.\n",
    "5. How to generate embeddings for each text chunk.\n",
    "6. How to insert into a vector database."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "tsHaUeqRpflK",
   "metadata": {},
   "source": [
    "## Pinecone\n",
    "\n",
    "You will need a [pinecone.io](https://www.pinecone.io/) api key for this tutorial. You can [sign up for free](https://app.pinecone.io/?sessionType=signup) to get a Starter account.\n",
    "\n",
    "If you create a Starter account, you can name your application anything you like.\n",
    "\n",
    "Once you have an account, navigate to 'API Keys' in the Pinecone console. You can use the default key or create a new one for this tutorial.\n",
    "\n",
    "Save your api key and its environment (`gcp_starter` for free accounts). You will need them below."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92b20306",
   "metadata": {},
   "source": [
    "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ae74e61",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-embeddings-openai\n",
    "%pip install llama-index-vector-stores-pinecone\n",
    "%pip install llama-index-llms-openai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b60e707a",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install llama-index"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22fb9e0a-566b-4f34-b9cf-72193cb51adb",
   "metadata": {},
   "source": [
    "## OpenAI\n",
    "\n",
    "You will need an [OpenAI](https://openai.com/) api key for this tutorial. Login to your [platform.openai.com](https://platform.openai.com/) account, click on your profile picture in the upper right corner, and choose 'API Keys' from the menu. Create an API key for this tutorial and save it. You will need it below."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "HPwWNeZwgwE8",
   "metadata": {},
   "source": [
    "## Environment\n",
    "\n",
    "First we add our dependencies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "CyTVgLfMgmIZ",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip -q install python-dotenv pinecone-client llama-index pymupdf"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bCwZFn6_iAR1",
   "metadata": {},
   "source": [
    "#### Set Environment Variables\n",
    "\n",
    "We create a file for our environment variables. Do not commit this file or share it!\n",
    "\n",
    "Note: Google Colabs will let you create but not open a .env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "M1l2emfWgjgE",
   "metadata": {},
   "outputs": [],
   "source": [
    "dotenv_path = (\n",
    "    \"env\"  # Google Colabs will not let you open a .env, but you can set\n",
    ")\n",
    "with open(dotenv_path, \"w\") as f:\n",
    "    f.write('PINECONE_API_KEY=\"<your api key>\"\\n')\n",
    "    f.write('OPENAI_API_KEY=\"<your api key>\"\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "PWMbn7GooMm5",
   "metadata": {},
   "source": [
    "Set your OpenAI api key, and Pinecone api key and environment in the file we created."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "QOyfIoXAoVGX",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from dotenv import load_dotenv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "leZkMBXYiTl-",
   "metadata": {},
   "outputs": [],
   "source": [
    "load_dotenv(dotenv_path=dotenv_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bcb486eb-c0b8-40e2-9038-da97aef63139",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "We build an empty Pinecone Index, and define the necessary LlamaIndex wrappers/abstractions so that we can start loading data into Pinecone.\n",
    "\n",
    "\n",
    "Note: Do not save your API keys in the code or add pinecone_env to your repo!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0Izxlt0XkMII",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pinecone import Pinecone, Index, ServerlessSpec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc739b4d-491f-406d-a0e6-f6b1e8c126dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "api_key = os.environ[\"PINECONE_API_KEY\"]\n",
    "pc = Pinecone(api_key=api_key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "Whwu7HqqswIq",
   "metadata": {},
   "outputs": [],
   "source": [
    "index_name = \"llamaindex-rag-fs\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "yRKkO4g1sBMl",
   "metadata": {},
   "outputs": [],
   "source": [
    "# [Optional] Delete the index before re-running the tutorial.\n",
    "# pinecone.delete_index(index_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20ba2f76-29d8-4dc5-b25c-64dcfe9e8d23",
   "metadata": {},
   "outputs": [],
   "source": [
    "# dimensions are for text-embedding-ada-002\n",
    "if index_name not in pc.list_indexes().names():\n",
    "    pc.create_index(\n",
    "        index_name,\n",
    "        dimension=1536,\n",
    "        metric=\"euclidean\",\n",
    "        spec=ServerlessSpec(cloud=\"aws\", region=\"us-east-1\"),\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",


"7182e98f-1b8a-4aba-af18-3982b862c794\n",
      "2024-05-06 14:00:35.931813\n",
      "BaseEmbedding.get_text_embedding_batch-632972aa-3345-49cb-ae2f-46f3166e3afc\n",
      "Event type: EmbeddingStartEvent\n",
      "{'model_name': 'text-embedding-ada-002', 'embed_batch_size': 100, 'num_workers': None, 'additional_kwargs': {}, 'api_base': 'https://api.openai.com/v1', 'api_version': '', 'max_retries': 10, 'timeout': 60.0, 'default_headers': None, 'reuse_client': True, 'dimensions': None, 'class_name': 'OpenAIEmbedding'}\n",
      "-----------------------\n",
      "-----------------------\n",
      "ba86e41f-cadf-4f1f-8908-8ee90404d668\n",
      "2024-05-06 14:00:36.256237\n",
      "BaseEmbedding.get_text_embedding_batch-632972aa-3345-49cb-ae2f-46f3166e3afc\n",
      "Event type: EmbeddingEndEvent\n",
      "['filename: README.md\\ncategory: codebase\\n\\nContext\\nLLMs are a phenomenal piece of technology for knowledge generation and reasoning.\\nThey are pre-trained on large amounts of publicly available data.\\nHow do we best augment LLMs with our own private data?\\nWe need a comprehensive toolkit to help perform this data augmentation for LLMs.\\n\\nProposed Solution\\nThat\\'s where LlamaIndex comes in. LlamaIndex is a \"data framework\" to help\\nyou build LLM  apps. It provides the following tools:\\n\\nOffers data connectors to ingest your existing data sources and data formats\\n(APIs, PDFs, docs, SQL, etc.)\\nProvides ways to structure your data (indices, graphs) so that this data can be\\neasily used with LLMs.\\nProvides an advanced retrieval/query interface over your data:\\nFeed in any LLM input prompt, get back retrieved context and knowledge-augmented output.\\nAllows easy integrations with your outer application framework\\n(e.g. with LangChain, Flask, Docker, ChatGPT, anything else).\\nLlamaIndex provides tools for both beginner users and advanced users.\\nOur high-level API allows beginner users to use LlamaIndex to ingest and\\nquery their data in 5 lines of code. Our lower-level APIs allow advanced users to\\ncustomize and extend any module (data connectors, indices, retrievers, query engines,\\nreranking modules), to fit their needs.']\n",
      "[-0.005768016912043095, 0.02242799662053585, -0.020438531413674355, -0.040361806750297546, -0.01757599227130413]\n",
      "-----------------------\n",
      "-----------------------\n",
      "06935377-f1e4-4fb9-b866-86f7520dfe2b\n",
      "2024-05-06 14:00:36.305798\n",
      "BaseQueryEngine.query-a766ae6c-6445-43b4-b1fc-9c29bae99556\n",
      "Event type: QueryStartEvent\n",
      "Tell me about LLMs?\n",
      "-----------------------\n",
      "-----------------------\n",
      "62608f4f-67a1-4e2c-a653-24a4430529bb\n",
      "2024-05-06 14:00:36.305998\n",
      "BaseRetriever.retrieve-4e25a2a3-43a9-45e3-a7b9-59f4d54e8f00\n",
      "Event type: RetrievalStartEvent\n",
      "Tell me about LLMs?\n",
      "-----------------------\n",
      "-----------------------\n",
      "e984c840-919b-4dc7-943d-5c49fbff48b8\n",
      "2024-05-06 14:00:36.306265\n",
      "BaseEmbedding.get_query_embedding-d30934f4-7bd2-4425-beda-12b5f55bc38b\n",
      "Event type: EmbeddingStartEvent\n",
      "{'model_name': 'text-embedding-ada-002', 'embed_batch_size': 100, 'num_workers': None, 'additional_kwargs': {}, 'api_base': 'https://api.openai.com/v1', 'api_version': '', 'max_retries': 10, 'timeout': 60.0, 'default_headers': None, 'reuse_client': True, 'dimensions': None, 'class_name': 'OpenAIEmbedding'}\n",
      "-----------------------\n",
      "-----------------------\n",
      "c09fa993-a892-4efe-9f1b-7238ff6e5c62\n",
      "2024-05-06 14:00:36.481459\n",
      "BaseEmbedding.get_query_embedding-d30934f4-7bd2-4425-beda-12b5f55bc38b\n",
      "Event type: EmbeddingEndEvent\n",
      "['Tell me about LLMs?']\n",
      "[0.00793155562132597, 0.011421983130276203, -0.010342259891331196, -0.03294854983687401, -0.03647972270846367]\n",
      "-----------------------\n",
      "-----------------------\n",
      "b076d239-628d-4b4c-94ed-25aa2ca4b02b\n",
      "2024-05-06 14:00:36.484080\n",
      "BaseRetriever.retrieve-4e25a2a3-43a9-45e3-a7b9-59f4d54e8f00\n",
      "Event type: RetrievalEndEvent\n",
      "Tell me about LLMs?\n",
      "[NodeWithScore(node=TextNode(id_='8de2b6b2-3fda-4f9b-95a8-a3ced6cfb0e5', embedding=None, metadata={'filename': 'README.md', 'category': 'codebase'}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={<NodeRelationship.SOURCE: '1'>: RelatedNodeInfo(node_id='29e2bc8f-b62c-4752-b5eb-11346c5cbe50', node_type=<ObjectType.DOCUMENT: '4'>, metadata={'filename': 'README.md', 'category': 'codebase'}, hash='3183371414f6a23e9a61e11b45ec45f808b148f9973166cfed62226e3505eb05')}, text='Context\\nLLMs are a phenomenal piece of technology for knowledge generation and reasoning.\\nThey are pre-trained on large amounts of publicly available data.\\nHow do we best augment LLMs with our own private data?\\nWe need a comprehensive toolkit to help perform this data augmentation for LLMs.\\n\\nProposed Solution\\nThat\\'s where LlamaIndex comes in. LlamaIndex is a \"data framework\" to help\\nyou build LLM  apps. It provides the following tools:\\n\\nOffers data connectors to ingest your existing data sources and data formats\\n(APIs, PDFs, docs, SQL, etc.)\\nProvides ways to structure your data (indices, graphs) so that this data can be\\neasily used with LLMs.\\nProvides an advanced retrieval/query interface over your data:\\nFeed in any LLM input prompt, get back retrieved context and knowledge-augmented output.\\nAllows easy integrations with your outer application framework\\n(e.g. with LangChain, Flask, Docker, ChatGPT, anything else).\\nLlamaIndex provides tools for both beginner users and advanced users.\\nOur high-level API allows beginner users to use LlamaIndex to ingest and\\nquery their data in 5 lines of code. Our lower-level APIs allow advanced users to\\ncustomize and extend any module (data connectors, indices, retrievers, query engines,\\nreranking modules), to fit their needs.', start_char_idx=1, end_char_idx=1279, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n'), score=0.807312731672428)]\n",
      "-----------------------\n",
      "-----------------------\n",
      "5e3289be-c597-48e7-ad3f-787722b766ea\n",
      "2024-05-06 14:00:36.484436\n",


{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "dd006f66",
   "metadata": {},
   "source": [
    ""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d617ade9-796f-431f-86ff-6b865e0eb007",
   "metadata": {},
   "source": [
    "# TopicNodeParser\n",
    "\n",
    "[MedGraphRAG](https://arxiv.org/html/2408.04187) aims to improve the capabilities of LLMs in the medical domain by generating evidence-based results through a novel graph-based Retrieval-Augmented Generation framework, improving safety and reliability in handling private medical data.\n",
    "\n",
    "`TopicNodeParser` implements an approximate version of the chunking technique described in the paper.\n",
    "\n",
    "Here is the technique as outlined in the paper:\n",
    "\n",
    "```\n",
    "Large medical documents often contain multiple themes or diverse content. To process these effectively, we first segment the document into data chunks that conform to the context limitations of Large Language Models (LLMs). Traditional methods such as chunking based on token size or fixed characters typically fail to detect subtle shifts in topics accurately. Consequently, these chunks may not fully capture the intended context, leading to a loss in the richness of meaning.\n",
    "\n",
    "To enhance accuracy, we adopt a mixed method of character separation coupled with topic-based segmentation. Specifically, we utilize static characters (line break symbols) to isolate individual paragraphs within the document. Following this, we apply a derived form of the text for semantic chunking. Our approach includes the use of proposition transfer, which extracts standalone statements from a raw text Chen et al. (2023). Through proposition transfer, each paragraph is transformed into self-sustaining statements. We then conduct a sequential analysis of the document to assess each proposition, deciding whether it should merge with an existing chunk or initiate a new one. This decision is made via a zero-shot approach by an LLM. To reduce noise generated by sequential processing, we implement a sliding window technique, managing five paragraphs at a time. We continuously adjust the window by removing the first paragraph and adding the next, maintaining focus on topic consistency. We set a hard threshold that the longest chunk cannot excess the context length limitation of LLM. After chunking the document, we construct graph on each individual of the data chunk.\n",
    "```\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d1c5118",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index llama-index-node-parser-topic"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12dcc784-f2c6-4c37-8771-57a921ff2eab",
   "metadata": {},
   "source": [
    "## Setup Data\n",
    "\n",
    "Here we consider a sample text.\n",
    "\n",
    "Note: The propositions were created by an LLM, which might lead to longer processing times when creating nodes. Exercise caution while experimenting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7fdcd874",
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"\"\"In this paper, we introduce a novel graph RAG method for applying LLMs to the medical domain, which we refer to as Medical Graph RAG (MedRAG). This technique improves LLM performance in the medical domain by response queries with grounded source citations and clear interpretations of medical terminology, boosting the transparency and interpretability of the results. This approach involves a three-tier hierarchical graph construction method. Initially, we use documents provided by users as our top-level source to extract entities. These entities are then linked to a second level consisting of more basic entities previously abstracted from credible medical books and papers. Subsequently, these entities are connected to a third level—the fundamental medical dictionary graph—that provides detailed explanations of each medical term and their semantic relationships. We then construct a comprehensive graph at the highest level by linking entities based on their content and hierarchical connections. This method ensures that the knowledge can be traced back to its sources and the results are factually accurate.\n",
    "\n",
    "To respond to user queries, we implement a U-retrieve strategy that combines top-down retrieval with bottom-up response generation. The process begins by structuring the query using predefined medical tags and indexing them through the graphs in a top-down manner. The system then generates responses based on these queries, pulling from meta-graphs—nodes retrieved along with their TopK related nodes and relationships—and summarizing the information into a detailed response. This technique maintains a balance between global context awareness and the contextual limitations inherent in LLMs.\n",
    "\n",
    "Our medical graph RAG provides Intrinsic source citation can enhance LLM transparency, interpretability, and verifiability. The results provides the provenance, or source grounding information, as it generates each response, and demonstrates that an answer is grounded in the dataset. Having the cited source for each assertion readily available also enables a human user to quickly and accurately audit the LLM’s output directly against the original source material. It is super useful in the field of medicine that security is very important, and each of the reasoning should be evidence-based. By using such a method, we construct an evidence-based Medical LLM that the clinician could easiely check the source of the reasoning and calibrate the model response to ensure the safty usage of llm in the clinical senarios.\n",
    "\n",
    "To evaluate our medical graph RAG, we implemented the method on several popular open and closed-source LLMs, including ChatGPT OpenAI (2023a) and LLaMA Touvron et al. (2023), testing them across mainstream medical Q&A benchmarks such as PubMedQA Jin et al. (2019), MedMCQA Pal et al. (2022), and USMLE Kung et al. (2023). For the RAG process, we supplied a comprehensive medical dictionary as the foundational knowledge layer, the UMLS medical knowledge graph Lindberg et al. (1993) as the foundamental layer detailing semantic relationships, and a curated MedC-K dataset Wu et al. (2023) —comprising the latest medical papers and books—as the intermediate level of data to simulate user-provided private data. Our experiments demonstrate that our model significantly enhances the performance of general-purpose LLMs on medical questions. Remarkably, it even surpasses many fine-tuned or specially trained LLMs on medical corpora, solely using the RAG approach without additional training.\n",
    "\"\"\"\n",
    "\n",
    "from llama_index.core import Document\n",
    "\n",
    "documents = [Document(text=text)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "717bd52c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "In this paper, we introduce a novel graph RAG method for applying LLMs to the medical domain, which we refer to as Medical Graph RAG (MedRAG). This technique improves LLM performance in the medical domain by response queries with grounded source citations and clear interpretations of medical terminology, boosting the transparency and interpretability of the results. This approach involves a three-tier hierarchical graph construction method. Initially, we use documents provided by users as our top-level source to extract entities. These entities are then linked to a second level consisting of more basic entities previously abstracted from credible medical books and papers. Subsequently, these entities are connected to a third level—the fundamental medical dictionary graph—that provides detailed explanations of each medical term and their semantic relationships. We then construct a comprehensive graph at the highest level by linking entities based on their content and hierarchical connections. This method ensures that the knowledge can be traced back to its sources and the results are factually accurate.\n",
      "\n",
      "To respond to user queries, we implement a U-retrieve strategy that combines top-down retrieval with bottom-up response generation. The process begins by structuring the query using predefined medical tags and indexing them through the graphs in a top-down manner. The system then generates responses based on these queries, pulling from meta-graphs—nodes retrieved along with their TopK related nodes and relationships—and summarizing the information into a detailed response. This technique maintains a balance between global context awareness and the contextual limitations inherent in LLMs.\n",
      "\n",


{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Router Query Engine\n",
    "\n",
    "`RouterQueryEngine` chooses the most appropriate query engine from multiple options to process a given query.\n",
    "\n",
    "This notebook walks through implementation of Router Query Engine, using workflows.\n",
    "\n",
    "Specifically we will implement [RouterQueryEngine](https://docs.llamaindex.ai/en/stable/examples/query_engine/RouterQueryEngine/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -U llama-index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"sk-..\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since workflows are async first, this all runs fine in a notebook. If you were running in your own code, you would want to use `asyncio.run()` to start an async event loop if one isn't already running.\n",
    "\n",
    "```python\n",
    "async def main():\n",
    "    <async code>\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    import asyncio\n",
    "    asyncio.run(main())\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define Events"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.workflow import Event\n",
    "from llama_index.core.base.base_selector import SelectorResult\n",
    "from typing import Dict, List, Any\n",
    "from llama_index.core.base.response.schema import RESPONSE_TYPE\n",
    "\n",
    "\n",
    "class QueryEngineSelectionEvent(Event):\n",
    "    \"\"\"Result of selecting the query engine tools.\"\"\"\n",
    "\n",
    "    selected_query_engines: SelectorResult\n",
    "\n",
    "\n",
    "class SynthesizeEvent(Event):\n",
    "    \"\"\"Event for synthesizing the response from different query engines.\"\"\"\n",
    "\n",
    "    result: List[RESPONSE_TYPE]\n",
    "    selected_query_engines: SelectorResult"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Workflow\n",
    "\n",
    "`selector:`\n",
    "\n",
    "1. It takes a StartEvent as input and returns a QueryEngineSelectionEvent.\n",
    "2. The `LLMSingleSelector`/ `PydanticSingleSelector`/ `PydanticMultiSelector` will select one/ multiple query engine tools.\n",
    "\n",
    "`generate_responses:`\n",
    "\n",
    "This function uses the selected query engines to generate responses and returns SynthesizeEvent.\n",
    "\n",
    "`synthesize_responses:`\n",
    "\n",
    "This function combines the generated responses and synthesizes the final response if multiple query engines are selected otherwise returns the single generated response.\n",
    "\n",
    "\n",
    "The steps will use the built-in `StartEvent` and `StopEvent` events.\n",
    "\n",
    "With our events defined, we can construct our workflow and steps. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.workflow import (\n",
    "    Context,\n",
    "    Workflow,\n",
    "    StartEvent,\n",
    "    StopEvent,\n",
    "    step,\n",
    ")\n",
    "\n",
    "from llama_index.llms.openai import OpenAI\n",
    "from llama_index.core.selectors.utils import get_selector_from_llm\n",
    "from llama_index.core.base.response.schema import (\n",
    "    PydanticResponse,\n",
    "    Response,\n",
    "    AsyncStreamingResponse,\n",
    ")\n",
    "from llama_index.core.bridge.pydantic import BaseModel\n",
    "from llama_index.core.response_synthesizers import TreeSummarize\n",
    "from llama_index.core.schema import QueryBundle\n",
    "from llama_index.core import Settings\n",
    "\n",
    "from IPython.display import Markdown, display\n",
    "import asyncio\n",
    "\n",
    "\n",
    "class RouterQueryEngineWorkflow(Workflow):\n",
    "    @step\n",
    "    async def selector(\n",
    "        self, ctx: Context, ev: StartEvent\n",
    "    ) -> QueryEngineSelectionEvent:\n",
    "        \"\"\"\n",
    "        Selects a single/ multiple query engines based on the query.\n",
    "        \"\"\"\n",
    "\n",
    "        await ctx.set(\"query\", ev.get(\"query\"))\n",
    "        await ctx.set(\"llm\", ev.get(\"llm\"))\n",
    "        await ctx.set(\"query_engine_tools\", ev.get(\"query_engine_tools\"))\n",
    "        await ctx.set(\"summarizer\", ev.get(\"summarizer\"))\n",
    "\n",
    "        llm = Settings.llm\n",
    "        select_multiple_query_engines = ev.get(\"select_multi\")\n",
    "        query = ev.get(\"query\")\n",
    "        query_engine_tools = ev.get(\"query_engine_tools\")\n",
    "\n",
    "        selector = get_selector_from_llm(\n",
    "            llm, is_multi=select_multiple_query_engines\n",
    "        )\n",
    "\n",
    "        query_engines_metadata = [\n",
    "            query_engine.metadata for query_engine in query_engine_tools\n",
    "        ]\n",
    "\n",
    "        selected_query_engines = await selector.aselect(\n",
    "            query_engines_metadata, query\n",
    "        )\n",
    "\n",
    "        return QueryEngineSelectionEvent(\n",
    "            selected_query_engines=selected_query_engines\n",
    "        )\n",
    "\n",
    "    @step\n",
    "    async def generate_responses(\n",
    "        self, ctx: Context, ev: QueryEngineSelectionEvent\n",
    "    ) -> SynthesizeEvent:\n",
    "        \"\"\"Generate the responses from the selected query engines.\"\"\"\n",
    "\n",
    "        query = await ctx.get(\"query\", default=None)\n",
    "        selected_query_engines = ev.selected_query_engines\n",
    "        query_engine_tools = await ctx.get(\"query_engine_tools\")\n",
    "\n",
    "        query_engines = [engine.query_engine for engine in query_engine_tools]\n",
    "\n",
    "        print(\n",
    "            f\"number of selected query engines: {len(selected_query_engines.selections)}\"\n",
    "        )\n",
    "\n",
    "        if len(selected_query_engines.selections) > 1:\n",
    "            tasks = []\n",
    "            for selected_query_engine in selected_query_engines.selections:\n",
    "                print(\n",
    "                    f\"Selected query engine: {selected_query_engine.index}: {selected_query_engine.reason}\"\n",
    "                )\n",
    "                query_engine = query_engines[selected_query_engine.index]\n",
    "                tasks.append(query_engine.aquery(query))\n",
    "\n",
    "            response_generated = await asyncio.gather(*tasks)\n",
    "\n",
    "        else:\n",
    "            query_engine = query_engines[\n",
    "                selected_query_engines.selections[0].index\n",
    "            ]\n",
    "\n",
    "            print(\n",
    "                f\"Selected query engine: {selected_query_engines.ind}: {selected_query_engines.reason}\"\n",
    "            )\n",
    "\n",
    "            response_generated = [await query_engine.aquery(query)]\n",
    "\n",
    "        return SynthesizeEvent(\n",
    "            result=response_generated,\n",