"**For that, importing Chroma vector database from langchain. Also, importing open source GPT4All for embedding models**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "9ff99dd7-9d47-4239-ba0a-d775792334ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_chroma import Chroma\n",
    "from langchain_community.embeddings import GPT4AllEmbeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5d1f4dd-dd8d-4a20-95d1-2dbdd204375a",
   "metadata": {},
   "source": [
    "**In next step, we will download one of the most popular embedding model \"all-MiniLM-L6-v2\". Find more details of the model at this link https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "05db3494-5d8e-4a13-9941-26330a86f5e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = \"all-MiniLM-L6-v2.gguf2.f16.gguf\"\n",
    "gpt4all_kwargs = {\"allow_download\": \"True\"}\n",
    "embeddings = GPT4AllEmbeddings(model_name=model_name, gpt4all_kwargs=gpt4all_kwargs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e53999e-1983-46ac-8039-2783e194c3ae",
   "metadata": {},
   "source": [
    "**Store all the embeddings in the Chroma database**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "0922951a-9ddf-4761-973d-8e9a86f61284",
   "metadata": {},
   "outputs": [],
   "source": [
    "vectorstore = Chroma.from_documents(documents=all_splits, embedding=embeddings)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29f94fa0-6c75-4a65-a1a3-debc75422479",
   "metadata": {},
   "source": [
    "**Now, let's find relevant splits from the documents related to the question**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "88c8152d-ec7a-4f0b-9d86-877789407537",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4\n"
     ]
    }
   ],
   "source": [
    "question = \"What is Intel CCG revenue in Q1 2024\"\n",
    "docs = vectorstore.similarity_search(question)\n",
    "print(len(docs))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53330c6b-cb0f-43f9-b379-2e57ac1e5335",
   "metadata": {},
   "source": [
    "**Look at the first retrieved document from the vector database**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "43a6d94f-b5c4-47b0-a353-2db4c3d24d9c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Document(metadata={'page': 1, 'source': 'intel_q1_2024_earnings.pdf'}, page_content='Client Computing Group (CCG) $7.5 billion up31%\\nData Center and AI (DCAI) $3.0 billion up5%\\nNetwork and Edge (NEX) $1.4 billion down 8%\\nTotal Intel Products revenue $11.9 billion up17%\\nIntel Foundry $4.4 billion down 10%\\nAll other:\\nAltera $342 million down 58%\\nMobileye $239 million down 48%\\nOther $194 million up17%\\nTotal all other revenue $775 million down 46%\\nIntersegment eliminations $(4.4) billion\\nTotal net revenue $12.7 billion up9%\\nIntel Products Highlights')"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "docs[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64ba074f-4b36-442e-b7e2-b26d6e2815c3",
   "metadata": {},
   "source": [
    "**Download Lllama-2 model from Huggingface and store locally**
\n",
    "**You can download different quantization variant of Lllama-2 model from the link below. We are using Q8 version here (7.16GB).**
\n",
    "https://huggingface.co/TheBloke/Llama-2-7B-Chat-GGUF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8dd0811-6f43-4bc6-b854-2ab377639c9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "!huggingface-cli download TheBloke/Llama-2-7b-Chat-GGUF llama-2-7b-chat.Q8_0.gguf --local-dir . --local-dir-use-symlinks False"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3895b1f5-f51d-4539-abf0-af33d7ca48ea",
   "metadata": {},
   "source": [
    "**Import langchain components required to load downloaded LLMs model**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "fb087088-aa62-44c0-8356-061e9b9f1186",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.callbacks.manager import CallbackManager\n",
    "from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler\n",
    "from langchain_community.llms import LlamaCpp"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a8a111e-2614-4b70-b034-85cd3e7304cb",
   "metadata": {},
   "source": [
    "**Loading the local Lllama-2 model using Llama-cpp library**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "fb917da2-c0d7-4995-b56d-26254276e0da",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "llama_model_loader: loaded meta data with 19 key-value pairs and 291 tensors from llama-2-7b-chat.Q8_0.gguf (version GGUF V2)\n",
      "llama_model_loader: Dumping metadata keys/values. Note: KV overrides do not apply in this output.\n",
      "llama_model_loader: - kv   0:                       general.architecture str              = llama\n",
      "llama_model_loader: - kv   1:                               general.name str              = LLaMA v2\n",
      "llama_model_loader: - kv   2:                       llama.context_length u32              = 4096\n",
      "llama_model_loader: - kv   3:                     llama.embedding_length u32              = 4096\n",
      "llama_model_loader: - kv   4:                          llama.block_count u32              = 32\n",
      "llama_model_loader: - kv   5:                  llama.feed_forward_length u32              = 11008\n",
      "llama_model_loader: - kv   6:                 llama.rope.dimension_count u32              = 128\n",
      "llama_model_loader: - kv   7:                 llama.attention.head_count u32              = 32\n",
      "llama_model_loader: - kv   8:              llama.attention.head_count_kv u32              = 32\n",
      "llama_model_loader: - kv   9:     llama.attention.layer_norm_rms_epsilon f32              = 0.000001\n",


]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "4654e5b7-635f-4767-8b31-4c430164cdd5",
   "metadata": {},
   "outputs": [],
   "source": [
    "retriever = vectorstore.as_retriever()\n",
    "qa_chain = (\n",
    "    {\"context\": retriever | format_docs, \"question\": RunnablePassthrough()}\n",
    "    | rag_prompt\n",
    "    | llm\n",
    "    | StrOutputParser()\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0979f393-fd0a-4e82-b844-68371c6ad68f",
   "metadata": {},
   "source": [
    "**Now we only need to pass the question to the chain and it will fetch the contexts directly from the vector database to generate the answer**\n",
    "<br>\n",
    "**Let's try with another question**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "3ea07b82-e6ec-4084-85f4-191373530172",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Llama.generate: prefix-match hit\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " According to the provided context, Intel DCAI revenue in Q1 2024 was $3.0 billion up 5%."
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "llama_print_timings:        load time =     131.20 ms\n",
      "llama_print_timings:      sample time =       6.28 ms /    31 runs   (    0.20 ms per token,  4937.88 tokens per second)\n",
      "llama_print_timings: prompt eval time =    2681.93 ms /   730 tokens (    3.67 ms per token,   272.19 tokens per second)\n",
      "llama_print_timings:        eval time =    1471.07 ms /    30 runs   (   49.04 ms per token,    20.39 tokens per second)\n",
      "llama_print_timings:       total time =    4206.77 ms /   760 tokens\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "' According to the provided context, Intel DCAI revenue in Q1 2024 was $3.0 billion up 5%.'"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "qa_chain.invoke(\"what is Intel DCAI revenue in Q1 2024?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9407f2a0-4a35-4315-8e96-02fcb80f210c",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.11.1 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.1"
  },
  "vscode": {
   "interpreter": {
    "hash": "1a1af0ee75eeea9e2e1ee996c87e7a2b11a0bebd85af04bb136d915cefc0abce"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}


"Created a chunk of size 1405, which is longer than the specified 1000\n",
      "Created a chunk of size 2221, which is longer than the specified 1000\n",
      "Created a chunk of size 1128, which is longer than the specified 1000\n",
      "Created a chunk of size 1021, which is longer than the specified 1000\n",
      "Created a chunk of size 1532, which is longer than the specified 1000\n",
      "Created a chunk of size 1535, which is longer than the specified 1000\n",
      "Created a chunk of size 1230, which is longer than the specified 1000\n",
      "Created a chunk of size 2456, which is longer than the specified 1000\n",
      "Created a chunk of size 1047, which is longer than the specified 1000\n",
      "Created a chunk of size 1320, which is longer than the specified 1000\n",
      "Created a chunk of size 1144, which is longer than the specified 1000\n",
      "Created a chunk of size 1509, which is longer than the specified 1000\n",
      "Created a chunk of size 1003, which is longer than the specified 1000\n",
      "Created a chunk of size 1025, which is longer than the specified 1000\n",
      "Created a chunk of size 1197, which is longer than the specified 1000\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "8244\n"
     ]
    }
   ],
   "source": [
    "from langchain_text_splitters import CharacterTextSplitter\n",
    "\n",
    "text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)\n",
    "texts = text_splitter.split_documents(docs)\n",
    "print(f\"{len(texts)}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then embed chunks and upload them to the DeepLake.\n",
    "\n",
    "This can take several minutes. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "OpenAIEmbeddings(client=<class 'openai.api_resources.embedding.Embedding'>, model='text-embedding-ada-002', deployment='text-embedding-ada-002', openai_api_version='', openai_api_base='', openai_api_type='', openai_proxy='', embedding_ctx_length=8191, openai_api_key='', openai_organization='', allowed_special=set(), disallowed_special='all', chunk_size=1000, max_retries=6, request_timeout=None, headers=None, tiktoken_model_name=None, show_progress_bar=False, model_kwargs={})"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain_openai import OpenAIEmbeddings\n",
    "\n",
    "embeddings = OpenAIEmbeddings()\n",
    "embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Your Deep Lake dataset has been successfully created!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset(path='hub://adilkhan/langchain-code', tensors=['embedding', 'id', 'metadata', 'text'])\n",
      "\n",
      "  tensor      htype       shape       dtype  compression\n",
      "  -------    -------     -------     -------  ------- \n",
      " embedding  embedding  (8244, 1536)  float32   None   \n",
      "    id        text      (8244, 1)      str     None   \n",
      " metadata     json      (8244, 1)      str     None   \n",
      "   text       text      (8244, 1)      str     None   \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": []
    },
    {
     "data": {
      "text/plain": [
       "<langchain_community.vectorstores.deeplake.DeepLake at 0x7fe1b67d7a30>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain_community.vectorstores import DeepLake\n",
    "\n",
    "username = \"<USERNAME_OR_ORG>\"\n",
    "\n",
    "\n",
    "db = DeepLake.from_documents(\n",
    "    texts, embeddings, dataset_path=f\"hub://{username}/langchain-code\", overwrite=True\n",
    ")\n",
    "db"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`Optional`: You can also use Deep Lake's Managed Tensor Database as a hosting service and run queries there. In order to do so, it is necessary to specify the runtime parameter as {'tensor_db': True} during the creation of the vector store. This configuration enables the execution of queries on the Managed Tensor Database, rather than on the client side. It should be noted that this functionality is not applicable to datasets stored locally or in-memory. In the event that a vector store has already been created outside of the Managed Tensor Database, it is possible to transfer it to the Managed Tensor Database by following the prescribed steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from langchain_community.vectorstores import DeepLake\n",
    "\n",
    "# db = DeepLake.from_documents(\n",
    "#     texts, embeddings, dataset_path=f\"hub://{<org_id>}/langchain-code\", runtime={\"tensor_db\": True}\n",
    "# )\n",
    "# db"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Question Answering\n",
    "First load the dataset, construct the retriever, then construct the Conversational Chain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Deep Lake Dataset in hub://adilkhan/langchain-code already exists, loading from the storage\n"
     ]
    }
   ],
   "source": [
    "db = DeepLake(\n",
    "    dataset_path=f\"hub://{username}/langchain-code\",\n",
    "    read_only=True,\n",
    "    embedding=embeddings,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "retriever = db.as_retriever()\n",
    "retriever.search_kwargs[\"distance_metric\"] = \"cos\"\n",
    "retriever.search_kwargs[\"fetch_k\"] = 20\n",
    "retriever.search_kwargs[\"maximal_marginal_relevance\"] = True\n",
    "retriever.search_kwargs[\"k\"] = 20"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can also specify user defined functions using [Deep Lake filters](https://docs.deeplake.ai/en/latest/deeplake.core.dataset.html#deeplake.core.dataset.Dataset.filter)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def filter(x):\n",
    "    # filter based on source code\n",
    "    if \"something\" in x[\"text\"].data()[\"value\"]:\n",
    "        return False\n",
    "\n",
    "    # filter based on path e.g. extension\n",
    "    metadata = x[\"metadata\"].data()[\"value\"]\n",
    "    return \"only_this\" in metadata[\"source\"] or \"also_that\" in metadata[\"source\"]\n",
    "\n",
    "\n",


"### turn on below for custom filtering\n",
    "# retriever.search_kwargs['filter'] = filter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from langchain.chains import ConversationalRetrievalChain\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "model = ChatOpenAI(\n",
    "    model_name=\"gpt-3.5-turbo-0613\"\n",
    ")  # 'ada' 'gpt-3.5-turbo-0613' 'gpt-4',\n",
    "qa = ConversationalRetrievalChain.from_llm(model, retriever=retriever)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-> **Question**: What is the class hierarchy? \n",
      "\n",
      "**Answer**: The class hierarchy for Memory is as follows:\n",
      "\n",
      "    BaseMemory --> BaseChatMemory --> <name>Memory  # Examples: ZepMemory, MotorheadMemory\n",
      "\n",
      "The class hierarchy for ChatMessageHistory is as follows:\n",
      "\n",
      "    BaseChatMessageHistory --> <name>ChatMessageHistory  # Example: ZepChatMessageHistory\n",
      "\n",
      "The class hierarchy for Prompt is as follows:\n",
      "\n",
      "    BasePromptTemplate --> PipelinePromptTemplate\n",
      "                           StringPromptTemplate --> PromptTemplate\n",
      "                                                    FewShotPromptTemplate\n",
      "                                                    FewShotPromptWithTemplates\n",
      "                           BaseChatPromptTemplate --> AutoGPTPrompt\n",
      "                                                      ChatPromptTemplate --> AgentScratchPadChatPromptTemplate\n",
      " \n",
      "\n",
      "-> **Question**: What classes are derived from the Chain class? \n",
      "\n",
      "**Answer**: The classes derived from the Chain class are:\n",
      "\n",
      "- APIChain\n",
      "- OpenAPIEndpointChain\n",
      "- AnalyzeDocumentChain\n",
      "- MapReduceDocumentsChain\n",
      "- MapRerankDocumentsChain\n",
      "- ReduceDocumentsChain\n",
      "- RefineDocumentsChain\n",
      "- StuffDocumentsChain\n",
      "- ConstitutionalChain\n",
      "- ConversationChain\n",
      "- ChatVectorDBChain\n",
      "- ConversationalRetrievalChain\n",
      "- FalkorDBQAChain\n",
      "- FlareChain\n",
      "- ArangoGraphQAChain\n",
      "- GraphQAChain\n",
      "- GraphCypherQAChain\n",
      "- HugeGraphQAChain\n",
      "- KuzuQAChain\n",
      "- NebulaGraphQAChain\n",
      "- NeptuneOpenCypherQAChain\n",
      "- GraphSparqlQAChain\n",
      "- HypotheticalDocumentEmbedder\n",
      "- LLMChain\n",
      "- LLMBashChain\n",
      "- LLMCheckerChain\n",
      "- LLMMathChain\n",
      "- LLMRequestsChain\n",
      "- LLMSummarizationCheckerChain\n",
      "- MapReduceChain\n",
      "- OpenAIModerationChain\n",
      "- NatBotChain\n",
      "- QAGenerationChain\n",
      "- QAWithSourcesChain\n",
      "- RetrievalQAWithSourcesChain\n",
      "- VectorDBQAWithSourcesChain\n",
      "- RetrievalQA\n",
      "- VectorDBQA\n",
      "- LLMRouterChain\n",
      "- MultiPromptChain\n",
      "- MultiRetrievalQAChain\n",
      "- MultiRouteChain\n",
      "- RouterChain\n",
      "- SequentialChain\n",
      "- SimpleSequentialChain\n",
      "- TransformChain\n",
      "- TaskPlaningChain\n",
      "- QueryChain\n",
      "- CPALChain\n",
      " \n",
      "\n",
      "-> **Question**: What kind of retrievers does LangChain have? \n",
      "\n",
      "**Answer**: The LangChain class includes various types of retrievers such as:\n",
      "\n",
      "- ArxivRetriever\n",
      "- AzureAISearchRetriever\n",
      "- BM25Retriever\n",
      "- ChaindeskRetriever\n",
      "- ChatGPTPluginRetriever\n",
      "- ContextualCompressionRetriever\n",
      "- DocArrayRetriever\n",
      "- ElasticSearchBM25Retriever\n",
      "- EnsembleRetriever\n",
      "- GoogleVertexAISearchRetriever\n",
      "- AmazonKendraRetriever\n",
      "- KNNRetriever\n",
      "- LlamaIndexGraphRetriever and LlamaIndexRetriever\n",
      "- MergerRetriever\n",
      "- MetalRetriever\n",
      "- MilvusRetriever\n",
      "- MultiQueryRetriever\n",
      "- ParentDocumentRetriever\n",
      "- PineconeHybridSearchRetriever\n",
      "- PubMedRetriever\n",
      "- RePhraseQueryRetriever\n",
      "- RemoteLangChainRetriever\n",
      "- SelfQueryRetriever\n",
      "- SVMRetriever\n",
      "- TFIDFRetriever\n",
      "- TimeWeightedVectorStoreRetriever\n",
      "- VespaRetriever\n",
      "- WeaviateHybridSearchRetriever\n",
      "- WebResearchRetriever\n",
      "- WikipediaRetriever\n",
      "- ZepRetriever\n",
      "- ZillizRetriever \n",
      "\n"
     ]
    }
   ],
   "source": [
    "questions = [\n",
    "    \"What is the class hierarchy?\",\n",
    "    \"What classes are derived from the Chain class?\",\n",
    "    \"What kind of retrievers does LangChain have?\",\n",
    "]\n",
    "chat_history = []\n",
    "qa_dict = {}\n",
    "\n",
    "for question in questions:\n",
    "    result = qa({\"question\": question, \"chat_history\": chat_history})\n",
    "    chat_history.append((question, result[\"answer\"]))\n",
    "    qa_dict[question] = result[\"answer\"]\n",
    "    print(f\"-> **Question**: {question} \\n\")\n",
    "    print(f\"**Answer**: {result['answer']} \\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'question': 'LangChain possesses a variety of retrievers including:\\n\\n1. ArxivRetriever\\n2. AzureAISearchRetriever\\n3. BM25Retriever\\n4. ChaindeskRetriever\\n5. ChatGPTPluginRetriever\\n6. ContextualCompressionRetriever\\n7. DocArrayRetriever\\n8. ElasticSearchBM25Retriever\\n9. EnsembleRetriever\\n10. GoogleVertexAISearchRetriever\\n11. AmazonKendraRetriever\\n12. KNNRetriever\\n13. LlamaIndexGraphRetriever\\n14. LlamaIndexRetriever\\n15. MergerRetriever\\n16. MetalRetriever\\n17. MilvusRetriever\\n18. MultiQueryRetriever\\n19. ParentDocumentRetriever\\n20. PineconeHybridSearchRetriever\\n21. PubMedRetriever\\n22. RePhraseQueryRetriever\\n23. RemoteLangChainRetriever\\n24. SelfQueryRetriever\\n25. SVMRetriever\\n26. TFIDFRetriever\\n27. TimeWeightedVectorStoreRetriever\\n28. VespaRetriever\\n29. WeaviateHybridSearchRetriever\\n30. WebResearchRetriever\\n31. WikipediaRetriever\\n32. ZepRetriever\\n33. ZillizRetriever\\n\\nIt also includes self query translators like:\\n\\n1. ChromaTranslator\\n2. DeepLakeTranslator\\n3. MyScaleTranslator\\n4. PineconeTranslator\\n5. QdrantTranslator\\n6. WeaviateTranslator\\n\\nAnd remote retrievers like:\\n\\n1. RemoteLangChainRetriever'}"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "qa_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The class hierarchy for Memory is as follows:\n",
      "\n",


"from langchain_core.pydantic_v1 import BaseModel, Field\n",
    "from langchain_core.runnables import RunnablePassthrough\n",
    "from langchain_core.utils.function_calling import convert_to_openai_tool\n",
    "from langchain_openai import ChatOpenAI, OpenAIEmbeddings\n",
    "from langgraph.prebuilt import ToolInvocation\n",
    "\n",
    "### Nodes ###\n",
    "\n",
    "\n",
    "def retrieve(state):\n",
    "    \"\"\"\n",
    "    Retrieve documents\n",
    "\n",
    "    Args:\n",
    "        state (dict): The current state of the agent, including all keys.\n",
    "\n",
    "    Returns:\n",
    "        dict: New key added to state, documents, that contains documents.\n",
    "    \"\"\"\n",
    "    print(\"---RETRIEVE---\")\n",
    "    state_dict = state[\"keys\"]\n",
    "    question = state_dict[\"question\"]\n",
    "    documents = retriever.invoke(question)\n",
    "    return {\"keys\": {\"documents\": documents, \"question\": question}}\n",
    "\n",
    "\n",
    "def generate(state):\n",
    "    \"\"\"\n",
    "    Generate answer\n",
    "\n",
    "    Args:\n",
    "        state (dict): The current state of the agent, including all keys.\n",
    "\n",
    "    Returns:\n",
    "        dict: New key added to state, generation, that contains generation.\n",
    "    \"\"\"\n",
    "    print(\"---GENERATE---\")\n",
    "    state_dict = state[\"keys\"]\n",
    "    question = state_dict[\"question\"]\n",
    "    documents = state_dict[\"documents\"]\n",
    "\n",
    "    # Prompt\n",
    "    prompt = hub.pull(\"rlm/rag-prompt\")\n",
    "\n",
    "    # LLM\n",
    "    llm = ChatOpenAI(model=\"gpt-3.5-turbo\", temperature=0)\n",
    "\n",
    "    # Post-processing\n",
    "    def format_docs(docs):\n",
    "        return \"\\n\\n\".join(doc.page_content for doc in docs)\n",
    "\n",
    "    # Chain\n",
    "    rag_chain = prompt | llm | StrOutputParser()\n",
    "\n",
    "    # Run\n",
    "    generation = rag_chain.invoke({\"context\": documents, \"question\": question})\n",
    "    return {\n",
    "        \"keys\": {\"documents\": documents, \"question\": question, \"generation\": generation}\n",
    "    }\n",
    "\n",
    "\n",
    "def grade_documents(state):\n",
    "    \"\"\"\n",
    "    Determines whether the retrieved documents are relevant to the question.\n",
    "\n",
    "    Args:\n",
    "        state (dict): The current state of the agent, including all keys.\n",
    "\n",
    "    Returns:\n",
    "        dict: New key added to state, filtered_documents, that contains relevant documents.\n",
    "    \"\"\"\n",
    "\n",
    "    print(\"---CHECK RELEVANCE---\")\n",
    "    state_dict = state[\"keys\"]\n",
    "    question = state_dict[\"question\"]\n",
    "    documents = state_dict[\"documents\"]\n",
    "\n",
    "    # Data model\n",
    "    class grade(BaseModel):\n",
    "        \"\"\"Binary score for relevance check.\"\"\"\n",
    "\n",
    "        binary_score: str = Field(description=\"Relevance score 'yes' or 'no'\")\n",
    "\n",
    "    # LLM\n",
    "    model = ChatOpenAI(temperature=0, model=\"gpt-4-0125-preview\", streaming=True)\n",
    "\n",
    "    # Tool\n",
    "    grade_tool_oai = convert_to_openai_tool(grade)\n",
    "\n",
    "    # LLM with tool and enforce invocation\n",
    "    llm_with_tool = model.bind(\n",
    "        tools=[convert_to_openai_tool(grade_tool_oai)],\n",
    "        tool_choice={\"type\": \"function\", \"function\": {\"name\": \"grade\"}},\n",
    "    )\n",
    "\n",
    "    # Parser\n",
    "    parser_tool = PydanticToolsParser(tools=[grade])\n",
    "\n",
    "    # Prompt\n",
    "    prompt = PromptTemplate(\n",
    "        template=\"\"\"You are a grader assessing relevance of a retrieved document to a user question. \\n \n",
    "        Here is the retrieved document: \\n\\n {context} \\n\\n\n",
    "        Here is the user question: {question} \\n\n",
    "        If the document contains keyword(s) or semantic meaning related to the user question, grade it as relevant. \\n\n",
    "        Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question.\"\"\",\n",
    "        input_variables=[\"context\", \"question\"],\n",
    "    )\n",
    "\n",
    "    # Chain\n",
    "    chain = prompt | llm_with_tool | parser_tool\n",
    "\n",
    "    # Score\n",
    "    filtered_docs = []\n",
    "    for d in documents:\n",
    "        score = chain.invoke({\"question\": question, \"context\": d.page_content})\n",
    "        grade = score[0].binary_score\n",
    "        if grade == \"yes\":\n",
    "            print(\"---GRADE: DOCUMENT RELEVANT---\")\n",
    "            filtered_docs.append(d)\n",
    "        else:\n",
    "            print(\"---GRADE: DOCUMENT NOT RELEVANT---\")\n",
    "            continue\n",
    "\n",
    "    return {\"keys\": {\"documents\": filtered_docs, \"question\": question}}\n",
    "\n",
    "\n",
    "def transform_query(state):\n",
    "    \"\"\"\n",
    "    Transform the query to produce a better question.\n",
    "\n",
    "    Args:\n",
    "        state (dict): The current state of the agent, including all keys.\n",
    "\n",
    "    Returns:\n",
    "        dict: New value saved to question.\n",
    "    \"\"\"\n",
    "\n",
    "    print(\"---TRANSFORM QUERY---\")\n",
    "    state_dict = state[\"keys\"]\n",
    "    question = state_dict[\"question\"]\n",
    "    documents = state_dict[\"documents\"]\n",
    "\n",
    "    # Create a prompt template with format instructions and the query\n",
    "    prompt = PromptTemplate(\n",
    "        template=\"\"\"You are generating questions that is well optimized for retrieval. \\n \n",
    "        Look at the input and try to reason about the underlying semantic intent / meaning. \\n \n",
    "        Here is the initial question:\n",
    "        \\n ------- \\n\n",
    "        {question} \n",
    "        \\n ------- \\n\n",
    "        Formulate an improved question: \"\"\",\n",
    "        input_variables=[\"question\"],\n",
    "    )\n",
    "\n",
    "    # Grader\n",
    "    model = ChatOpenAI(temperature=0, model=\"gpt-4-0125-preview\", streaming=True)\n",
    "\n",
    "    # Prompt\n",
    "    chain = prompt | model | StrOutputParser()\n",
    "    better_question = chain.invoke({\"question\": question})\n",
    "\n",
    "    return {\"keys\": {\"documents\": documents, \"question\": better_question}}\n",
    "\n",
    "\n",
    "def prepare_for_final_grade(state):\n",
    "    \"\"\"\n",
    "    Stage for final grade, passthrough state.\n",
    "\n",
    "    Args:\n",
    "        state (dict): The current state of the agent, including all keys.\n",
    "\n",
    "    Returns:\n",
    "        state (dict): The current state of the agent, including all keys.\n",
    "    \"\"\"\n",
    "\n",
    "    print(\"---FINAL GRADE---\")\n",
    "    state_dict = state[\"keys\"]\n",
    "    question = state_dict[\"question\"]\n",
    "    documents = state_dict[\"documents\"]\n",
    "    generation = state_dict[\"generation\"]\n",
    "\n",


"workflow.add_node(\"generate\", generate)  # generatae\n",
    "workflow.add_node(\"transform_query\", transform_query)  # transform_query\n",
    "workflow.add_node(\"prepare_for_final_grade\", prepare_for_final_grade)  # passthrough\n",
    "\n",
    "# Build graph\n",
    "workflow.set_entry_point(\"retrieve\")\n",
    "workflow.add_edge(\"retrieve\", \"grade_documents\")\n",
    "workflow.add_conditional_edges(\n",
    "    \"grade_documents\",\n",
    "    decide_to_generate,\n",
    "    {\n",
    "        \"transform_query\": \"transform_query\",\n",
    "        \"generate\": \"generate\",\n",
    "    },\n",
    ")\n",
    "workflow.add_edge(\"transform_query\", \"retrieve\")\n",
    "workflow.add_conditional_edges(\n",
    "    \"generate\",\n",
    "    grade_generation_v_documents,\n",
    "    {\n",
    "        \"supported\": \"prepare_for_final_grade\",\n",
    "        \"not supported\": \"generate\",\n",
    "    },\n",
    ")\n",
    "workflow.add_conditional_edges(\n",
    "    \"prepare_for_final_grade\",\n",
    "    grade_generation_v_question,\n",
    "    {\n",
    "        \"useful\": END,\n",
    "        \"not useful\": \"transform_query\",\n",
    "    },\n",
    ")\n",
    "\n",
    "# Compile\n",
    "app = workflow.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb69dbb9-91ee-4868-8c3c-93af3cd885be",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run\n",
    "inputs = {\"keys\": {\"question\": \"Explain how the different types of agent memory work?\"}}\n",
    "for output in app.stream(inputs):\n",
    "    for key, value in output.items():\n",
    "        pprint.pprint(f\"Output from node '{key}':\")\n",
    "        pprint.pprint(\"---\")\n",
    "        pprint.pprint(value[\"keys\"], indent=2, width=80, depth=None)\n",
    "    pprint.pprint(\"\\n---\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4138bc51-8c84-4b8a-8d24-f7f470721f6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = {\"keys\": {\"question\": \"Explain how chain of thought prompting works?\"}}\n",
    "for output in app.stream(inputs):\n",
    "    for key, value in output.items():\n",
    "        pprint.pprint(f\"Output from node '{key}':\")\n",
    "        pprint.pprint(\"---\")\n",
    "        pprint.pprint(value[\"keys\"], indent=2, width=80, depth=None)\n",
    "    pprint.pprint(\"\\n---\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "548f1c5b-4108-4aae-8abb-ec171b511b92",
   "metadata": {},
   "source": [
    "Trace - \n",
    " \n",
    "* https://smith.langchain.com/public/55d6180f-aab8-42bc-8799-dadce6247d9b/r\n",
    "* https://smith.langchain.com/public/f85ebc95-81d9-47fc-91c6-b54e5b78f359/r"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.11.1 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.1"
  },
  "vscode": {
   "interpreter": {
    "hash": "1a1af0ee75eeea9e2e1ee996c87e7a2b11a0bebd85af04bb136d915cefc0abce"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}


{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "68b24990",
   "metadata": {},
   "source": [
    "# Combine agents and vector stores\n",
    "\n",
    "This notebook covers how to combine agents and vector stores. The use case for this is that you've ingested your data into a vector store and want to interact with it in an agentic manner.\n",
    "\n",
    "The recommended method for doing so is to create a `RetrievalQA` and then use that as a tool in the overall agent. Let's take a look at doing this below. You can do this with multiple different vector DBs, and use the agent as a way to route between them. There are two different ways of doing this - you can either let the agent use the vector stores as normal tools, or you can set `return_direct=True` to really just use the agent as a router."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b22020a",
   "metadata": {},
   "source": [
    "## Create the vector store"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "2e87c10a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains import RetrievalQA\n",
    "from langchain_chroma import Chroma\n",
    "from langchain_openai import OpenAI, OpenAIEmbeddings\n",
    "from langchain_text_splitters import CharacterTextSplitter\n",
    "\n",
    "llm = OpenAI(temperature=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "0b7b772b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "relevant_parts = []\n",
    "for p in Path(\".\").absolute().parts:\n",
    "    relevant_parts.append(p)\n",
    "    if relevant_parts[-3:] == [\"langchain\", \"docs\", \"modules\"]:\n",
    "        break\n",
    "doc_path = str(Path(*relevant_parts) / \"state_of_the_union.txt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "f2675861",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running Chroma using direct local API.\n",
      "Using DuckDB in-memory for database. Data will be transient.\n"
     ]
    }
   ],
   "source": [
    "from langchain_community.document_loaders import TextLoader\n",
    "\n",
    "loader = TextLoader(doc_path)\n",
    "documents = loader.load()\n",
    "text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)\n",
    "texts = text_splitter.split_documents(documents)\n",
    "\n",
    "embeddings = OpenAIEmbeddings()\n",
    "docsearch = Chroma.from_documents(texts, embeddings, collection_name=\"state-of-union\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "bc5403d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "state_of_union = RetrievalQA.from_chain_type(\n",
    "    llm=llm, chain_type=\"stuff\", retriever=docsearch.as_retriever()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "1431cded",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.document_loaders import WebBaseLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "915d3ff3",
   "metadata": {},
   "outputs": [],
   "source": [
    "loader = WebBaseLoader(\"https://beta.ruff.rs/docs/faq/\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "96a2edf8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running Chroma using direct local API.\n",
      "Using DuckDB in-memory for database. Data will be transient.\n"
     ]
    }
   ],
   "source": [
    "docs = loader.load()\n",
    "ruff_texts = text_splitter.split_documents(docs)\n",
    "ruff_db = Chroma.from_documents(ruff_texts, embeddings, collection_name=\"ruff\")\n",
    "ruff = RetrievalQA.from_chain_type(\n",
    "    llm=llm, chain_type=\"stuff\", retriever=ruff_db.as_retriever()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71ecef90",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "c0a6c031",
   "metadata": {},
   "source": [
    "## Create the Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "eb142786",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import things that are needed generically\n",
    "from langchain.agents import AgentType, Tool, initialize_agent\n",
    "from langchain_openai import OpenAI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "850bc4e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "tools = [\n",
    "    Tool(\n",
    "        name=\"State of Union QA System\",\n",
    "        func=state_of_union.run,\n",
    "        description=\"useful for when you need to answer questions about the most recent state of the union address. Input should be a fully formed question.\",\n",
    "    ),\n",
    "    Tool(\n",
    "        name=\"Ruff QA System\",\n",
    "        func=ruff.run,\n",
    "        description=\"useful for when you need to answer questions about ruff (a python linter). Input should be a fully formed question.\",\n",
    "    ),\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "fc47f230",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Construct the agent. We will use the default agent type here.\n",
    "# See documentation for a full list of options.\n",
    "agent = initialize_agent(\n",
    "    tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "10ca2db8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3m I need to find out what Biden said about Ketanji Brown Jackson in the State of the Union address.\n",
      "Action: State of Union QA System\n",
      "Action Input: What did Biden say about Ketanji Brown Jackson in the State of the Union address?\u001b[0m\n",
      "Observation: \u001b[36;1m\u001b[1;3m Biden said that Jackson is one of the nation's top legal minds and that she will continue Justice Breyer's legacy of excellence.\u001b[0m\n",
      "Thought:\u001b[32;1m\u001b[1;3m I now know the final answer\n",
      "Final Answer: Biden said that Jackson is one of the nation's top legal minds and that she will continue Justice Breyer's legacy of excellence.\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\"Biden said that Jackson is one of the nation's top legal minds and that she will continue Justice Breyer's legacy of excellence.\""
      ]
     },
     "execution_count": 46,
     "metadata": {},


{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Human input LLM\n",
    "\n",
    "Similar to the fake LLM, LangChain provides a pseudo LLM class that can be used for testing, debugging, or educational purposes. This allows you to mock out calls to the LLM and simulate how a human would respond if they received the prompts.\n",
    "\n",
    "In this notebook, we go over how to use this.\n",
    "\n",
    "We start this with using the HumanInputLLM in an agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.llms.human import HumanInputLLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.agents import AgentType, initialize_agent, load_tools"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since we will use the `WikipediaQueryRun` tool in this notebook, you might need to install the `wikipedia` package if you haven't done so already."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install wikipedia"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "tools = load_tools([\"wikipedia\"])\n",
    "llm = HumanInputLLM(\n",
    "    prompt_func=lambda prompt: print(\n",
    "        f\"\\n===PROMPT====\\n{prompt}\\n=====END OF PROMPT======\"\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "agent = initialize_agent(\n",
    "    tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\n",
      "===PROMPT====\n",
      "Answer the following questions as best you can. You have access to the following tools:\n",
      "\n",
      "Wikipedia: A wrapper around Wikipedia. Useful for when you need to answer general questions about people, places, companies, historical events, or other subjects. Input should be a search query.\n",
      "\n",
      "Use the following format:\n",
      "\n",
      "Question: the input question you must answer\n",
      "Thought: you should always think about what to do\n",
      "Action: the action to take, should be one of [Wikipedia]\n",
      "Action Input: the input to the action\n",
      "Observation: the result of the action\n",
      "... (this Thought/Action/Action Input/Observation can repeat N times)\n",
      "Thought: I now know the final answer\n",
      "Final Answer: the final answer to the original input question\n",
      "\n",
      "Begin!\n",
      "\n",
      "Question: What is 'Bocchi the Rock!'?\n",
      "Thought:\n",
      "=====END OF PROMPT======\n",
      "\u001b[32;1m\u001b[1;3mI need to use a tool.\n",
      "Action: Wikipedia\n",
      "Action Input: Bocchi the Rock!, Japanese four-panel manga and anime series.\u001b[0m\n",
      "Observation: \u001b[36;1m\u001b[1;3mPage: Bocchi the Rock!\n",
      "Summary: Bocchi the Rock! (ぼっち・ざ・ろっく!, Bocchi Za Rokku!) is a Japanese four-panel manga series written and illustrated by Aki Hamaji. It has been serialized in Houbunsha's seinen manga magazine Manga Time Kirara Max since December 2017. Its chapters have been collected in five tankōbon volumes as of November 2022.\n",
      "An anime television series adaptation produced by CloverWorks aired from October to December 2022. The series has been praised for its writing, comedy, characters, and depiction of social anxiety, with the anime's visual creativity receiving acclaim.\n",
      "\n",
      "Page: Manga Time Kirara\n",
      "Summary: Manga Time Kirara (まんがタイムきらら, Manga Taimu Kirara) is a Japanese seinen manga magazine published by Houbunsha which mainly serializes four-panel manga. The magazine is sold on the ninth of each month and was first published as a special edition of Manga Time, another Houbunsha magazine, on May 17, 2002. Characters from this magazine have appeared in a crossover role-playing game called Kirara Fantasia.\n",
      "\n",
      "Page: Manga Time Kirara Max\n",
      "Summary: Manga Time Kirara Max (まんがタイムきららMAX) is a Japanese four-panel seinen manga magazine published by Houbunsha. It is the third magazine of the \"Kirara\" series, after \"Manga Time Kirara\" and \"Manga Time Kirara Carat\". The first issue was released on September 29, 2004. Currently the magazine is released on the 19th of each month.\u001b[0m\n",
      "Thought:\n",
      "===PROMPT====\n",
      "Answer the following questions as best you can. You have access to the following tools:\n",
      "\n",
      "Wikipedia: A wrapper around Wikipedia. Useful for when you need to answer general questions about people, places, companies, historical events, or other subjects. Input should be a search query.\n",
      "\n",
      "Use the following format:\n",
      "\n",
      "Question: the input question you must answer\n",
      "Thought: you should always think about what to do\n",
      "Action: the action to take, should be one of [Wikipedia]\n",
      "Action Input: the input to the action\n",
      "Observation: the result of the action\n",
      "... (this Thought/Action/Action Input/Observation can repeat N times)\n",
      "Thought: I now know the final answer\n",
      "Final Answer: the final answer to the original input question\n",
      "\n",
      "Begin!\n",
      "\n",
      "Question: What is 'Bocchi the Rock!'?\n",
      "Thought:I need to use a tool.\n",
      "Action: Wikipedia\n",
      "Action Input: Bocchi the Rock!, Japanese four-panel manga and anime series.\n",
      "Observation: Page: Bocchi the Rock!\n",
      "Summary: Bocchi the Rock! (ぼっち・ざ・ろっく!, Bocchi Za Rokku!) is a Japanese four-panel manga series written and illustrated by Aki Hamaji. It has been serialized in Houbunsha's seinen manga magazine Manga Time Kirara Max since December 2017. Its chapters have been collected in five tankōbon volumes as of November 2022.\n",
      "An anime television series adaptation produced by CloverWorks aired from October to December 2022. The series has been praised for its writing, comedy, characters, and depiction of social anxiety, with the anime's visual creativity receiving acclaim.\n",
      "\n",
      "Page: Manga Time Kirara\n",
      "Summary: Manga Time Kirara (まんがタイムきらら, Manga Taimu Kirara) is a Japanese seinen manga magazine published by Houbunsha which mainly serializes four-panel manga. The magazine is sold on the ninth of each month and was first published as a special edition of Manga Time, another Houbunsha magazine, on May 17, 2002. Characters from this magazine have appeared in a crossover role-playing game called Kirara Fantasia.\n",
      "\n",
      "Page: Manga Time Kirara Max\n",
      "Summary: Manga Time Kirara Max (まんがタイムきららMAX) is a Japanese four-panel seinen manga magazine published by Houbunsha. It is the third magazine of the \"Kirara\" series, after \"Manga Time Kirara\" and \"Manga Time Kirara Carat\". The first issue was released on September 29, 2004. Currently the magazine is released on the 19th of each month.\n",
      "Thought:\n",


{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "517a9fd4",
   "metadata": {},
   "source": [
    "# BabyAGI with Tools\n",
    "\n",
    "This notebook builds on top of [baby agi](baby_agi.html), but shows how you can swap out the execution chain. The previous execution chain was just an LLM which made stuff up. By swapping it out with an agent that has access to tools, we can hopefully get real reliable information"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "556af556",
   "metadata": {},
   "source": [
    "## Install and Import Required Modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c8a354b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Optional\n",
    "\n",
    "from langchain.chains import LLMChain\n",
    "from langchain.prompts import PromptTemplate\n",
    "from langchain_experimental.autonomous_agents import BabyAGI\n",
    "from langchain_openai import OpenAI, OpenAIEmbeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09f70772",
   "metadata": {},
   "source": [
    "## Connect to the Vector Store\n",
    "\n",
    "Depending on what vectorstore you use, this step may look different."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "794045d4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Note: you may need to restart the kernel to use updated packages.\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip install faiss-cpu > /dev/null\n",
    "%pip install google-search-results > /dev/null\n",
    "from langchain.docstore import InMemoryDocstore\n",
    "from langchain_community.vectorstores import FAISS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6e0305eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define your embedding model\n",
    "embeddings_model = OpenAIEmbeddings()\n",
    "# Initialize the vectorstore as empty\n",
    "import faiss\n",
    "\n",
    "embedding_size = 1536\n",
    "index = faiss.IndexFlatL2(embedding_size)\n",
    "vectorstore = FAISS(embeddings_model.embed_query, index, InMemoryDocstore({}), {})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f3b72bf",
   "metadata": {},
   "source": [
    "## Define the Chains\n",
    "\n",
    "BabyAGI relies on three LLM chains:\n",
    "- Task creation chain to select new tasks to add to the list\n",
    "- Task prioritization chain to re-prioritize tasks\n",
    "- Execution Chain to execute the tasks\n",
    "\n",
    "\n",
    "NOTE: in this notebook, the Execution chain will now be an agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b43cd580",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.agents import AgentExecutor, Tool, ZeroShotAgent\n",
    "from langchain.chains import LLMChain\n",
    "from langchain_community.utilities import SerpAPIWrapper\n",
    "from langchain_openai import OpenAI\n",
    "\n",
    "todo_prompt = PromptTemplate.from_template(\n",
    "    \"You are a planner who is an expert at coming up with a todo list for a given objective. Come up with a todo list for this objective: {objective}\"\n",
    ")\n",
    "todo_chain = LLMChain(llm=OpenAI(temperature=0), prompt=todo_prompt)\n",
    "search = SerpAPIWrapper()\n",
    "tools = [\n",
    "    Tool(\n",
    "        name=\"Search\",\n",
    "        func=search.run,\n",
    "        description=\"useful for when you need to answer questions about current events\",\n",
    "    ),\n",
    "    Tool(\n",
    "        name=\"TODO\",\n",
    "        func=todo_chain.run,\n",
    "        description=\"useful for when you need to come up with todo lists. Input: an objective to create a todo list for. Output: a todo list for that objective. Please be very clear what the objective is!\",\n",
    "    ),\n",
    "]\n",
    "\n",
    "\n",
    "prefix = \"\"\"You are an AI who performs one task based on the following objective: {objective}. Take into account these previously completed tasks: {context}.\"\"\"\n",
    "suffix = \"\"\"Question: {task}\n",
    "{agent_scratchpad}\"\"\"\n",
    "prompt = ZeroShotAgent.create_prompt(\n",
    "    tools,\n",
    "    prefix=prefix,\n",
    "    suffix=suffix,\n",
    "    input_variables=[\"objective\", \"task\", \"context\", \"agent_scratchpad\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4b00ae2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = OpenAI(temperature=0)\n",
    "llm_chain = LLMChain(llm=llm, prompt=prompt)\n",
    "tool_names = [tool.name for tool in tools]\n",
    "agent = ZeroShotAgent(llm_chain=llm_chain, allowed_tools=tool_names)\n",
    "agent_executor = AgentExecutor.from_agent_and_tools(\n",
    "    agent=agent, tools=tools, verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05ba762e",
   "metadata": {},
   "source": [
    "### Run the BabyAGI\n",
    "\n",
    "Now it's time to create the BabyAGI controller and watch it try to accomplish your objective."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3d220b69",
   "metadata": {},
   "outputs": [],
   "source": [
    "OBJECTIVE = \"Write a weather report for SF today\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "3d69899b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Logging of LLMChains\n",
    "verbose = False\n",
    "# If None, will keep on going forever\n",
    "max_iterations: Optional[int] = 3\n",
    "baby_agi = BabyAGI.from_llm(\n",
    "    llm=llm,\n",
    "    vectorstore=vectorstore,\n",
    "    task_execution_chain=agent_executor,\n",
    "    verbose=verbose,\n",
    "    max_iterations=max_iterations,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "f7957b51",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[95m\u001b[1m\n",
      "*****TASK LIST*****\n",
      "\u001b[0m\u001b[0m\n",
      "1: Make a todo list\n",
      "\u001b[92m\u001b[1m\n",
      "*****NEXT TASK*****\n",
      "\u001b[0m\u001b[0m\n",
      "1: Make a todo list\n",
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3mThought: I need to come up with a todo list\n",
      "Action: TODO\n",
      "Action Input: Write a weather report for SF today\u001b[0m\u001b[33;1m\u001b[1;3m\n",
      "\n",
      "1. Research current weather conditions in San Francisco\n",
      "2. Gather data on temperature, humidity, wind speed, and other relevant weather conditions\n",


"File \u001b[0;32m~/langchain/langchain/chains/base.py:140\u001b[0m, in \u001b[0;36mChain.__call__\u001b[0;34m(self, inputs, return_only_outputs, callbacks)\u001b[0m\n\u001b[1;32m    138\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m (\u001b[38;5;167;01mKeyboardInterrupt\u001b[39;00m, \u001b[38;5;167;01mException\u001b[39;00m) \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[1;32m    139\u001b[0m     run_manager\u001b[38;5;241m.\u001b[39mon_chain_error(e)\n\u001b[0;32m--> 140\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m e\n\u001b[1;32m    141\u001b[0m run_manager\u001b[38;5;241m.\u001b[39mon_chain_end(outputs)\n\u001b[1;32m    142\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mprep_outputs(inputs, outputs, return_only_outputs)\n",
      "File \u001b[0;32m~/langchain/langchain/chains/base.py:134\u001b[0m, in \u001b[0;36mChain.__call__\u001b[0;34m(self, inputs, return_only_outputs, callbacks)\u001b[0m\n\u001b[1;32m    128\u001b[0m run_manager \u001b[38;5;241m=\u001b[39m callback_manager\u001b[38;5;241m.\u001b[39mon_chain_start(\n\u001b[1;32m    129\u001b[0m     {\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mname\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__class__\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m},\n\u001b[1;32m    130\u001b[0m     inputs,\n\u001b[1;32m    131\u001b[0m )\n\u001b[1;32m    132\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m    133\u001b[0m     outputs \u001b[38;5;241m=\u001b[39m (\n\u001b[0;32m--> 134\u001b[0m         \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call\u001b[49m\u001b[43m(\u001b[49m\u001b[43minputs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mrun_manager\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mrun_manager\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    135\u001b[0m         \u001b[38;5;28;01mif\u001b[39;00m new_arg_supported\n\u001b[1;32m    136\u001b[0m         \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_call(inputs)\n\u001b[1;32m    137\u001b[0m     )\n\u001b[1;32m    138\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m (\u001b[38;5;167;01mKeyboardInterrupt\u001b[39;00m, \u001b[38;5;167;01mException\u001b[39;00m) \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[1;32m    139\u001b[0m     run_manager\u001b[38;5;241m.\u001b[39mon_chain_error(e)\n",


"File \u001b[0;32m~/langchain/langchain/agents/agent.py:953\u001b[0m, in \u001b[0;36mAgentExecutor._call\u001b[0;34m(self, inputs, run_manager)\u001b[0m\n\u001b[1;32m    951\u001b[0m \u001b[38;5;66;03m# We now enter the agent loop (until it returns something).\u001b[39;00m\n\u001b[1;32m    952\u001b[0m \u001b[38;5;28;01mwhile\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_should_continue(iterations, time_elapsed):\n\u001b[0;32m--> 953\u001b[0m     next_step_output \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_take_next_step\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    954\u001b[0m \u001b[43m        \u001b[49m\u001b[43mname_to_tool_map\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    955\u001b[0m \u001b[43m        \u001b[49m\u001b[43mcolor_mapping\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    956\u001b[0m \u001b[43m        \u001b[49m\u001b[43minputs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    957\u001b[0m \u001b[43m        \u001b[49m\u001b[43mintermediate_steps\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    958\u001b[0m \u001b[43m        \u001b[49m\u001b[43mrun_manager\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mrun_manager\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    959\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    960\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(next_step_output, AgentFinish):\n\u001b[1;32m    961\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_return(\n\u001b[1;32m    962\u001b[0m             next_step_output, intermediate_steps, run_manager\u001b[38;5;241m=\u001b[39mrun_manager\n\u001b[1;32m    963\u001b[0m         )\n",


"metadata": {},
   "outputs": [],
   "source": [
    "from langchain.agents import Tool\n",
    "from langchain_experimental.utilities import PythonREPL\n",
    "\n",
    "python_repl = PythonREPL()\n",
    "\n",
    "repl_tool = Tool(\n",
    "    name=\"python_repl\",\n",
    "    description=\"A Python shell. Use this to execute python commands. Input should be a valid python command. If you want to see the output of a value, you should print it out with `print(...)`.\",\n",
    "    func=python_repl.run,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain import hub\n",
    "\n",
    "llm = ChatOpenAI(temperature=0, model=\"gpt-4-1106-preview\")\n",
    "toolkit = CassandraDatabaseToolkit(db=db)\n",
    "\n",
    "# context = toolkit.get_context()\n",
    "# tools = toolkit.get_tools()\n",
    "tools = [schema_tool, select_data_tool, repl_tool]\n",
    "\n",
    "input = (\n",
    "    QUERY_PATH_PROMPT\n",
    "    + f\"\"\"\n",
    "\n",
    "Here is your task: In the {keyspace} keyspace, find the total number of times the temperature of each device has exceeded 23 degrees on July 14, 2020.\n",
    " Create a summary report including the name of the room. Use Pandas if helpful.\n",
    "\"\"\"\n",
    ")\n",
    "\n",
    "prompt = hub.pull(\"hwchase17/openai-tools-agent\")\n",
    "\n",
    "# messages = [\n",
    "#     HumanMessagePromptTemplate.from_template(input),\n",
    "#     AIMessage(content=QUERY_PATH_PROMPT),\n",
    "#     MessagesPlaceholder(variable_name=\"agent_scratchpad\"),\n",
    "# ]\n",
    "\n",
    "# prompt = ChatPromptTemplate.from_messages(messages)\n",
    "# print(prompt)\n",
    "\n",
    "# Choose the LLM that will drive the agent\n",
    "# Only certain models support this\n",
    "llm = ChatOpenAI(model=\"gpt-3.5-turbo-1106\", temperature=0)\n",
    "\n",
    "# Construct the OpenAI Tools agent\n",
    "agent = create_openai_tools_agent(llm, tools, prompt)\n",
    "\n",
    "print(\"Available tools:\")\n",
    "for tool in tools:\n",
    "    print(\"\\t\" + tool.name + \" - \" + tool.description + \" - \" + str(tool))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)\n",
    "\n",
    "response = agent_executor.invoke({\"input\": input})\n",
    "\n",
    "print(response[\"output\"])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}


{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ba5f8741",
   "metadata": {},
   "source": [
    "# Custom agent with tool retrieval\n",
    "\n",
    "The novel idea introduced in this notebook is the idea of using retrieval to select the set of tools to use to answer an agent query. This is useful when you have many many tools to select from. You cannot put the description of all the tools in the prompt (because of context length issues) so instead you dynamically select the N tools you do want to consider using at run time.\n",
    "\n",
    "In this notebook we will create a somewhat contrived example. We will have one legitimate tool (search) and then 99 fake tools which are just nonsense. We will then add a step in the prompt template that takes the user input and retrieves tool relevant to the query."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fea4812c",
   "metadata": {},
   "source": [
    "## Set up environment\n",
    "\n",
    "Do necessary imports, etc."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9af9734e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "from typing import Union\n",
    "\n",
    "from langchain.agents import (\n",
    "    AgentExecutor,\n",
    "    AgentOutputParser,\n",
    "    LLMSingleActionAgent,\n",
    "    Tool,\n",
    ")\n",
    "from langchain.chains import LLMChain\n",
    "from langchain.prompts import StringPromptTemplate\n",
    "from langchain_community.utilities import SerpAPIWrapper\n",
    "from langchain_core.agents import AgentAction, AgentFinish\n",
    "from langchain_openai import OpenAI"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6df0253f",
   "metadata": {},
   "source": [
    "## Set up tools\n",
    "\n",
    "We will create one legitimate tool (search) and then 99 fake tools."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "becda2a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define which tools the agent can use to answer user queries\n",
    "search = SerpAPIWrapper()\n",
    "search_tool = Tool(\n",
    "    name=\"Search\",\n",
    "    func=search.run,\n",
    "    description=\"useful for when you need to answer questions about current events\",\n",
    ")\n",
    "\n",
    "\n",
    "def fake_func(inp: str) -> str:\n",
    "    return \"foo\"\n",
    "\n",
    "\n",
    "fake_tools = [\n",
    "    Tool(\n",
    "        name=f\"foo-{i}\",\n",
    "        func=fake_func,\n",
    "        description=f\"a silly function that you can use to get more information about the number {i}\",\n",
    "    )\n",
    "    for i in range(99)\n",
    "]\n",
    "ALL_TOOLS = [search_tool] + fake_tools"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17362717",
   "metadata": {},
   "source": [
    "## Tool Retriever\n",
    "\n",
    "We will use a vector store to create embeddings for each tool description. Then, for an incoming query we can create embeddings for that query and do a similarity search for relevant tools."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "77c4be4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.vectorstores import FAISS\n",
    "from langchain_core.documents import Document\n",
    "from langchain_openai import OpenAIEmbeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9092a158",
   "metadata": {},
   "outputs": [],
   "source": [
    "docs = [\n",
    "    Document(page_content=t.description, metadata={\"index\": i})\n",
    "    for i, t in enumerate(ALL_TOOLS)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "affc4e56",
   "metadata": {},
   "outputs": [],
   "source": [
    "vector_store = FAISS.from_documents(docs, OpenAIEmbeddings())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "735a7566",
   "metadata": {},
   "outputs": [],
   "source": [
    "retriever = vector_store.as_retriever()\n",
    "\n",
    "\n",
    "def get_tools(query):\n",
    "    docs = retriever.invoke(query)\n",
    "    return [ALL_TOOLS[d.metadata[\"index\"]] for d in docs]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7699afd7",
   "metadata": {},
   "source": [
    "We can now test this retriever to see if it seems to work."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "425f2886",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Tool(name='Search', description='useful for when you need to answer questions about current events', return_direct=False, verbose=False, callback_manager=<langchain.callbacks.shared.SharedCallbackManager object at 0x114b28a90>, func=<bound method SerpAPIWrapper.run of SerpAPIWrapper(search_engine=<class 'serpapi.google_search.GoogleSearch'>, params={'engine': 'google', 'google_domain': 'google.com', 'gl': 'us', 'hl': 'en'}, serpapi_api_key='', aiosession=None)>, coroutine=None),\n",
       " Tool(name='foo-95', description='a silly function that you can use to get more information about the number 95', return_direct=False, verbose=False, callback_manager=<langchain.callbacks.shared.SharedCallbackManager object at 0x114b28a90>, func=<function fake_func at 0x15e5bd1f0>, coroutine=None),\n",
       " Tool(name='foo-12', description='a silly function that you can use to get more information about the number 12', return_direct=False, verbose=False, callback_manager=<langchain.callbacks.shared.SharedCallbackManager object at 0x114b28a90>, func=<function fake_func at 0x15e5bd1f0>, coroutine=None),\n",
       " Tool(name='foo-15', description='a silly function that you can use to get more information about the number 15', return_direct=False, verbose=False, callback_manager=<langchain.callbacks.shared.SharedCallbackManager object at 0x114b28a90>, func=<function fake_func at 0x15e5bd1f0>, coroutine=None)]"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "get_tools(\"whats the weather?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "4036dd19",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Tool(name='foo-13', description='a silly function that you can use to get more information about the number 13', return_direct=False, verbose=False, callback_manager=<langchain.callbacks.shared.SharedCallbackManager object at 0x114b28a90>, func=<function fake_func at 0x15e5bd1f0>, coroutine=None),\n",
       " Tool(name='foo-12', description='a silly function that you can use to get more information about the number 12', return_direct=False, verbose=False, callback_manager=<langchain.callbacks.shared.SharedCallbackManager object at 0x114b28a90>, func=<function fake_func at 0x15e5bd1f0>, coroutine=None),\n",
       " Tool(name='foo-14', description='a silly function that you can use to get more information about the number 14', return_direct=False, verbose=False, callback_manager=<langchain.callbacks.shared.SharedCallbackManager object at 0x114b28a90>, func=<function fake_func at 0x15e5bd1f0>, coroutine=None),\n",


"1                                       None                        None   \n",
       "2                                       None                        None   \n",
       "3                                       None                        None   \n",
       "4  Illinois J. Math. 52 (2008) no.2, 681-689                        None   \n",
       "\n",
       "          report-no       categories  \\\n",
       "0  ANL-HEP-PR-07-12           hep-ph   \n",
       "1              None    math.CO cs.CG   \n",
       "2              None   physics.gen-ph   \n",
       "3              None          math.CO   \n",
       "4              None  math.CA math.FA   \n",
       "\n",
       "                                             license  \\\n",
       "0                                               None   \n",
       "1  http://arxiv.org/licenses/nonexclusive-distrib...   \n",
       "2                                               None   \n",
       "3                                               None   \n",
       "4                                               None   \n",
       "\n",
       "                                            abstract  \\\n",
       "0    A fully differential calculation in perturba...   \n",
       "1    We describe a new algorithm, the $(k,\\ell)$-...   \n",
       "2    The evolution of Earth-Moon system is descri...   \n",
       "3    We show that a determinant of Stirling cycle...   \n",
       "4    In this paper we show how to compute the $\\L...   \n",
       "\n",
       "                                            versions update_date  \\\n",
       "0  [{'version': 'v1', 'created': 'Mon, 2 Apr 2007...  2008-11-26   \n",
       "1  [{'version': 'v1', 'created': 'Sat, 31 Mar 200...  2008-12-13   \n",
       "2  [{'version': 'v1', 'created': 'Sun, 1 Apr 2007...  2008-01-13   \n",
       "3  [{'version': 'v1', 'created': 'Sat, 31 Mar 200...  2007-05-23   \n",
       "4  [{'version': 'v1', 'created': 'Mon, 2 Apr 2007...  2013-10-15   \n",
       "\n",
       "                                      authors_parsed  \\\n",
       "0  [[Balázs, C., ], [Berger, E. L., ], [Nadolsky,...   \n",
       "1           [[Streinu, Ileana, ], [Theran, Louis, ]]   \n",
       "2                                 [[Pan, Hongjun, ]]   \n",
       "3                                [[Callan, David, ]]   \n",
       "4  [[Abu-Shammala, Wael, ], [Torchinsky, Alberto, ]]   \n",
       "\n",
       "                                           embedding  \n",
       "0  [0.0594153292, -0.0440569334, -0.0487333685, -...  \n",
       "1  [0.0247399714, -0.065658465, 0.0201423876, -0....  \n",
       "2  [0.0491479263, 0.0728017688, 0.0604138002, 0.0...  \n",
       "3  [0.0389556214, -0.0410280302, 0.0410280302, -0...  \n",
       "4  [0.118412666, -0.0127423415, 0.1185125113, 0.0...  "
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(len(dataset_df))\n",
    "dataset_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "o2gHwRjMfJlO"
   },
   "outputs": [],
   "source": [
    "from pymongo import MongoClient\n",
    "\n",
    "# Initialize MongoDB python client\n",
    "client = MongoClient(MONGO_URI, appname=\"devrel.content.ai_agent_firechain.python\")\n",
    "\n",
    "DB_NAME = \"agent_demo\"\n",
    "COLLECTION_NAME = \"knowledge\"\n",
    "ATLAS_VECTOR_SEARCH_INDEX_NAME = \"vector_index\"\n",
    "collection = client[DB_NAME][COLLECTION_NAME]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "zJkyy9UbffZT",
    "outputId": "c6f78ea3-fc93-4d57-95eb-98cea5bf15d3"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Data ingestion into MongoDB completed\n"
     ]
    }
   ],
   "source": [
    "# Delete any existing records in the collection\n",
    "collection.delete_many({})\n",
    "\n",
    "# Data Ingestion\n",
    "records = dataset_df.to_dict(\"records\")\n",
    "collection.insert_many(records)\n",
    "\n",
    "print(\"Data ingestion into MongoDB completed\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6S1Cz9dtGPwL"
   },
   "source": [
    "## Create Vector Search Index Defintion\n",
    "\n",
    "```\n",
    "{\n",
    "  \"fields\": [\n",
    "    {\n",
    "      \"type\": \"vector\",\n",
    "      \"path\": \"embedding\",\n",
    "      \"numDimensions\": 256,\n",
    "      \"similarity\": \"cosine\"\n",
    "    }\n",
    "  ]\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1a-0n9PpfqDj"
   },
   "source": [
    "## Create LangChain Retriever (MongoDB)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "HAxeTPimfxM-"
   },
   "outputs": [],
   "source": [
    "from langchain_mongodb import MongoDBAtlasVectorSearch\n",
    "from langchain_openai import OpenAIEmbeddings\n",
    "\n",
    "embedding_model = OpenAIEmbeddings(model=\"text-embedding-3-small\", dimensions=256)\n",
    "\n",
    "# Vector Store Creation\n",
    "vector_store = MongoDBAtlasVectorSearch.from_connection_string(\n",
    "    connection_string=MONGO_URI,\n",
    "    namespace=DB_NAME + \".\" + COLLECTION_NAME,\n",
    "    embedding=embedding_model,\n",
    "    index_name=ATLAS_VECTOR_SEARCH_INDEX_NAME,\n",
    "    text_key=\"abstract\",\n",
    ")\n",
    "\n",
    "retriever = vector_store.as_retriever(search_type=\"similarity\", search_kwargs={\"k\": 5})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Optional: Creating a retrevier with compression capabilities using LLMLingua\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install langchain_community llmlingua"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.retrievers import ContextualCompressionRetriever\n",
    "from langchain_community.document_compressors import LLMLinguaCompressor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/richmondalake/miniconda3/envs/langchain_workarea/lib/python3.12/site-packages/huggingface_hub/file_download.py:1132: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [


"compressor = LLMLinguaCompressor(model_name=\"openai-community/gpt2\", device_map=\"cpu\")\n",
    "compression_retriever = ContextualCompressionRetriever(\n",
    "    base_compressor=compressor, base_retriever=retriever\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Sm5QZdshwJLN"
   },
   "source": [
    "## Configure LLM Using Fireworks AI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {
    "id": "V4ztCMCtgme_"
   },
   "outputs": [],
   "source": [
    "from langchain_fireworks import ChatFireworks\n",
    "\n",
    "llm = ChatFireworks(model=\"accounts/fireworks/models/firefunction-v1\", max_tokens=256)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pZfheX5FiIhU"
   },
   "source": [
    "## Agent Tools Creation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {
    "id": "3eufR9H8gopU"
   },
   "outputs": [],
   "source": [
    "from langchain.agents import tool\n",
    "from langchain.tools.retriever import create_retriever_tool\n",
    "from langchain_community.document_loaders import ArxivLoader\n",
    "\n",
    "\n",
    "# Custom Tool Definiton\n",
    "@tool\n",
    "def get_metadata_information_from_arxiv(word: str) -> list:\n",
    "    \"\"\"\n",
    "    Fetches and returns metadata for a maximum of ten documents from arXiv matching the given query word.\n",
    "\n",
    "    Args:\n",
    "      word (str): The search query to find relevant documents on arXiv.\n",
    "\n",
    "    Returns:\n",
    "      list: Metadata about the documents matching the query.\n",
    "    \"\"\"\n",
    "    docs = ArxivLoader(query=word, load_max_docs=10).load()\n",
    "    # Extract just the metadata from each document\n",
    "    metadata_list = [doc.metadata for doc in docs]\n",
    "    return metadata_list\n",
    "\n",
    "\n",
    "@tool\n",
    "def get_information_from_arxiv(word: str) -> list:\n",
    "    \"\"\"\n",
    "    Fetches and returns metadata for a single research paper from arXiv matching the given query word, which is the ID of the paper, for example: 704.0001.\n",
    "\n",
    "    Args:\n",
    "      word (str): The search query to find the relevant paper on arXiv using the ID.\n",
    "\n",
    "    Returns:\n",
    "      list: Data about the paper matching the query.\n",
    "    \"\"\"\n",
    "    doc = ArxivLoader(query=word, load_max_docs=1).load()\n",
    "    return doc\n",
    "\n",
    "\n",
    "# If you created a retriever with compression capaitilies in the optional cell in an earlier cell, you can replace 'retriever' with 'compression_retriever'\n",
    "# Otherwise you can also create a compression procedure as a tool for the agent as shown in the `compress_prompt_using_llmlingua` tool definition function\n",
    "retriever_tool = create_retriever_tool(\n",
    "    retriever=retriever,\n",
    "    name=\"knowledge_base\",\n",
    "    description=\"This serves as the base knowledge source of the agent and contains some records of research papers from Arxiv. This tool is used as the first step for exploration and reseach efforts.\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.document_compressors import LLMLinguaCompressor\n",
    "\n",
    "compressor = LLMLinguaCompressor(model_name=\"openai-community/gpt2\", device_map=\"cpu\")\n",
    "\n",
    "\n",
    "@tool\n",
    "def compress_prompt_using_llmlingua(prompt: str, compression_rate: float = 0.5) -> str:\n",
    "    \"\"\"\n",
    "    Compresses a long data or prompt using the LLMLinguaCompressor.\n",
    "\n",
    "    Args:\n",
    "        data (str): The data or prompt to be compressed.\n",
    "        compression_rate (float): The rate at which to compress the data (default is 0.5).\n",
    "\n",
    "    Returns:\n",
    "        str: The compressed data or prompt.\n",
    "    \"\"\"\n",
    "    compressed_data = compressor.compress_prompt(\n",
    "        prompt,\n",
    "        rate=compression_rate,\n",
    "        force_tokens=[\"!\", \".\", \"?\", \"\\n\"],\n",
    "        drop_consecutive=True,\n",
    "    )\n",
    "    return compressed_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {
    "id": "AS8QmaKVjhbR"
   },
   "outputs": [],
   "source": [
    "tools = [\n",
    "    retriever_tool,\n",
    "    get_metadata_information_from_arxiv,\n",
    "    get_information_from_arxiv,\n",
    "    compress_prompt_using_llmlingua,\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ueEn73nlliNr"
   },
   "source": [
    "## Agent Prompt Creation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "metadata": {
    "id": "RY13DrVXFDrm"
   },
   "outputs": [],
   "source": [
    "from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\n",
    "\n",
    "agent_purpose = \"\"\"\n",
    "You are a helpful research assistant equipped with various tools to assist with your tasks efficiently. \n",
    "You have access to conversational history stored in your inpout as chat_history.\n",
    "You are cost-effective and utilize the compress_prompt_using_llmlingua tool whenever you determine that a prompt or conversational history is too long. \n",
    "Below are instructions on when and how to use each tool in your operations.\n",
    "\n",
    "1. get_metadata_information_from_arxiv\n",
    "\n",
    "Purpose: To fetch and return metadata for up to ten documents from arXiv that match a given query word.\n",
    "When to Use: Use this tool when you need to gather metadata about multiple research papers related to a specific topic.\n",
    "Example: If you are asked to provide an overview of recent papers on \"machine learning,\" use this tool to fetch metadata for relevant documents.\n",
    "\n",
    "2. get_information_from_arxiv\n",
    "\n",
    "Purpose: To fetch and return metadata for a single research paper from arXiv using the paper's ID.\n",
    "When to Use: Use this tool when you need detailed information about a specific research paper identified by its arXiv ID.\n",
    "Example: If you are asked to retrieve detailed information about the paper with the ID \"704.0001,\" use this tool.\n",
    "\n",
    "3. retriever_tool\n",
    "\n",
    "Purpose: To serve as your base knowledge, containing records of research papers from arXiv.\n",
    "When to Use: Use this tool as the first step for exploration and research efforts when dealing with topics covered by the documents in the knowledge base.\n",
    "Example: When beginning research on a new topic that is well-documented in the arXiv repository, use this tool to access the relevant papers.\n",
    "\n",
    "4. compress_prompt_using_llmlingua\n",
    "\n",
    "Purpose: To compress long prompts or conversational histories using the LLMLinguaCompressor.\n",
    "When to Use: Use this tool whenever you determine that a prompt or conversational history is too long to be efficiently processed.\n",
    "Example: If you receive a very lengthy query or conversation context that exceeds the typical token limits, compress it using this tool before proceeding with further processing.\n",
    "\n",
    "\"\"\"\n",
    "\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", agent_purpose),\n",


{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "707d13a7",
   "metadata": {},
   "source": [
    "# Databricks\n",
    "\n",
    "This notebook covers how to connect to the [Databricks runtimes](https://docs.databricks.com/runtime/index.html) and [Databricks SQL](https://www.databricks.com/product/databricks-sql) using the SQLDatabase wrapper of LangChain.\n",
    "It is broken into 3 parts: installation and setup, connecting to Databricks, and examples."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0076d072",
   "metadata": {},
   "source": [
    "## Installation and Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "739b489b",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install databricks-sql-connector"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73113163",
   "metadata": {},
   "source": [
    "## Connecting to Databricks\n",
    "\n",
    "You can connect to [Databricks runtimes](https://docs.databricks.com/runtime/index.html) and [Databricks SQL](https://www.databricks.com/product/databricks-sql) using the `SQLDatabase.from_databricks()` method.\n",
    "\n",
    "### Syntax\n",
    "```python\n",
    "SQLDatabase.from_databricks(\n",
    "    catalog: str,\n",
    "    schema: str,\n",
    "    host: Optional[str] = None,\n",
    "    api_token: Optional[str] = None,\n",
    "    warehouse_id: Optional[str] = None,\n",
    "    cluster_id: Optional[str] = None,\n",
    "    engine_args: Optional[dict] = None,\n",
    "    **kwargs: Any)\n",
    "```\n",
    "### Required Parameters\n",
    "* `catalog`: The catalog name in the Databricks database.\n",
    "* `schema`: The schema name in the catalog.\n",
    "\n",
    "### Optional Parameters\n",
    "There following parameters are optional. When executing the method in a Databricks notebook, you don't need to provide them in most of the cases.\n",
    "* `host`: The Databricks workspace hostname, excluding 'https://' part. Defaults to 'DATABRICKS_HOST' environment variable or current workspace if in a Databricks notebook.\n",
    "* `api_token`: The Databricks personal access token for accessing the Databricks SQL warehouse or the cluster. Defaults to 'DATABRICKS_TOKEN' environment variable or a temporary one is generated if in a Databricks notebook.\n",
    "* `warehouse_id`: The warehouse ID in the Databricks SQL.\n",
    "* `cluster_id`: The cluster ID in the Databricks Runtime. If running in a Databricks notebook and both 'warehouse_id' and 'cluster_id' are None, it uses the ID of the cluster the notebook is attached to.\n",
    "* `engine_args`: The arguments to be used when connecting Databricks.\n",
    "* `**kwargs`: Additional keyword arguments for the `SQLDatabase.from_uri` method."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b11c7e48",
   "metadata": {},
   "source": [
    "## Examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "8102bca0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Connecting to Databricks with SQLDatabase wrapper\n",
    "from langchain_community.utilities import SQLDatabase\n",
    "\n",
    "db = SQLDatabase.from_databricks(catalog=\"samples\", schema=\"nyctaxi\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9dd36f58",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Creating a OpenAI Chat LLM wrapper\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "llm = ChatOpenAI(temperature=0, model_name=\"gpt-4\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b5c5f1a",
   "metadata": {},
   "source": [
    "### SQL Chain example\n",
    "\n",
    "This example demonstrates the use of the [SQL Chain](https://python.langchain.com/en/latest/modules/chains/examples/sqlite.html) for answering a question over a Databricks database."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "36f2270b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.utilities import SQLDatabaseChain\n",
    "\n",
    "db_chain = SQLDatabaseChain.from_llm(llm, db, verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4e2b5f25",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new SQLDatabaseChain chain...\u001b[0m\n",
      "What is the average duration of taxi rides that start between midnight and 6am?\n",
      "SQLQuery:\u001b[32;1m\u001b[1;3mSELECT AVG(UNIX_TIMESTAMP(tpep_dropoff_datetime) - UNIX_TIMESTAMP(tpep_pickup_datetime)) as avg_duration\n",
      "FROM trips\n",
      "WHERE HOUR(tpep_pickup_datetime) >= 0 AND HOUR(tpep_pickup_datetime) < 6\u001b[0m\n",
      "SQLResult: \u001b[33;1m\u001b[1;3m[(987.8122786304605,)]\u001b[0m\n",
      "Answer:\u001b[32;1m\u001b[1;3mThe average duration of taxi rides that start between midnight and 6am is 987.81 seconds.\u001b[0m\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'The average duration of taxi rides that start between midnight and 6am is 987.81 seconds.'"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "db_chain.run(\n",
    "    \"What is the average duration of taxi rides that start between midnight and 6am?\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e496d5e5",
   "metadata": {},
   "source": [
    "### SQL Database Agent example\n",
    "\n",
    "This example demonstrates the use of the [SQL Database Agent](/docs/integrations/tools/sql_database) for answering questions over a Databricks database."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "9918e86a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.agents import create_sql_agent\n",
    "from langchain_community.agent_toolkits import SQLDatabaseToolkit\n",
    "\n",
    "toolkit = SQLDatabaseToolkit(db=db, llm=llm)\n",
    "agent = create_sql_agent(llm=llm, toolkit=toolkit, verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "c484a76e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3mAction: list_tables_sql_db\n",
      "Action Input: \u001b[0m\n",
      "Observation: \u001b[38;5;200m\u001b[1;3mtrips\u001b[0m\n",
      "Thought:\u001b[32;1m\u001b[1;3mI should check the schema of the trips table to see if it has the necessary columns for trip distance and duration.\n",
      "Action: schema_sql_db\n",


{
 "cells": [
  {
   "attachments": {
    "semantic-chunking-rag.png": {
     ""
    }
   },
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Retrieval Augmented Generation (RAG)\n",
    "\n",
    "This notebook demonstrates an example of using [LangChain](https://www.langchain.com/) to delvelop a Retrieval Augmented Generation (RAG) pattern. It uses Azure AI Document Intelligence as document loader, which can extracts tables, paragraphs, and layout information from pdf, image, office and html files. The output markdown can be used in LangChain's markdown header splitter, which enables semantic chunking of the documents. Then the chunked documents are indexed into Azure AI Search vectore store. Given a user query, it will use Azure AI Search to get the relevant chunks, then feed the context into the prompt with the query to generate the answer.\n",
    "\n",
    "![semantic-chunking-rag.png](attachment:semantic-chunking-rag.png)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "- An Azure AI Document Intelligence resource in one of the 3 preview regions: **East US**, **West US2**, **West Europe** - follow [this document](https://learn.microsoft.com/azure/ai-services/document-intelligence/create-document-intelligence-resource?view=doc-intel-4.0.0) to create one if you don't have.\n",
    "- An Azure AI Search resource - follow [this document](https://learn.microsoft.com/azure/search/search-create-service-portal) to create one if you don't have.\n",
    "- An Azure OpenAI resource and deployments for embeddings model and chat model - follow [this document](https://learn.microsoft.com/azure/ai-services/openai/how-to/create-resource?pivots=web-portal) to create one if you don't have.\n",
    "\n",
    "We’ll use an Azure OpenAI chat model and embeddings and Azure AI Search in this walkthrough, but everything shown here works with any ChatModel or LLM, Embeddings, and VectorStore or Retriever."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install python-dotenv langchain langchain-community langchain-openai langchainhub openai tiktoken azure-ai-documentintelligence azure-identity azure-search-documents==11.4.0b8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "This code loads environment variables using the `dotenv` library and sets the necessary environment variables for Azure services.\n",
    "The environment variables are loaded from the `.env` file in the same directory as this notebook.\n",
    "\"\"\"\n",
    "import os\n",
    "\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "os.environ[\"AZURE_OPENAI_ENDPOINT\"] = os.getenv(\"AZURE_OPENAI_ENDPOINT\")\n",
    "os.environ[\"AZURE_OPENAI_API_KEY\"] = os.getenv(\"AZURE_OPENAI_API_KEY\")\n",
    "doc_intelligence_endpoint = os.getenv(\"AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT\")\n",
    "doc_intelligence_key = os.getenv(\"AZURE_DOCUMENT_INTELLIGENCE_KEY\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain import hub\n",
    "from langchain.schema import StrOutputParser\n",
    "from langchain.schema.runnable import RunnablePassthrough\n",
    "from langchain.text_splitter import MarkdownHeaderTextSplitter\n",
    "from langchain.vectorstores.azuresearch import AzureSearch\n",
    "from langchain_community.document_loaders import AzureAIDocumentIntelligenceLoader\n",
    "from langchain_openai import AzureChatOpenAI, AzureOpenAIEmbeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load a document and split it into semantic chunks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initiate Azure AI Document Intelligence to load the document. You can either specify file_path or url_path to load the document.\n",
    "loader = AzureAIDocumentIntelligenceLoader(\n",
    "    file_path=\"<path to your file>\",\n",
    "    api_key=doc_intelligence_key,\n",
    "    api_endpoint=doc_intelligence_endpoint,\n",
    "    api_model=\"prebuilt-layout\",\n",
    ")\n",
    "docs = loader.load()\n",
    "\n",
    "# Split the document into chunks base on markdown headers.\n",
    "headers_to_split_on = [\n",
    "    (\"#\", \"Header 1\"),\n",
    "    (\"##\", \"Header 2\"),\n",
    "    (\"###\", \"Header 3\"),\n",
    "]\n",
    "text_splitter = MarkdownHeaderTextSplitter(headers_to_split_on=headers_to_split_on)\n",
    "\n",
    "docs_string = docs[0].page_content\n",
    "splits = text_splitter.split_text(docs_string)\n",
    "\n",
    "print(\"Length of splits: \" + str(len(splits)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Embed and index the chunks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Embed the splitted documents and insert into Azure Search vector store\n",
    "\n",
    "aoai_embeddings = AzureOpenAIEmbeddings(\n",
    "    azure_deployment=\"<Azure OpenAI embeddings model>\",\n",
    "    openai_api_version=\"<Azure OpenAI API version>\",  # e.g., \"2023-07-01-preview\"\n",
    ")\n",
    "\n",
    "vector_store_address: str = os.getenv(\"AZURE_SEARCH_ENDPOINT\")\n",
    "vector_store_password: str = os.getenv(\"AZURE_SEARCH_ADMIN_KEY\")\n",
    "\n",
    "index_name: str = \"<your index name>\"\n",
    "vector_store: AzureSearch = AzureSearch(\n",
    "    azure_search_endpoint=vector_store_address,\n",
    "    azure_search_key=vector_store_password,\n",
    "    index_name=index_name,\n",
    "    embedding_function=aoai_embeddings.embed_query,\n",
    ")\n",
    "\n",
    "vector_store.add_documents(documents=splits)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Retrive relevant chunks based on a question"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Retrieve relevant chunks based on the question\n",
    "\n",
    "retriever = vector_store.as_retriever(search_type=\"similarity\", search_kwargs={\"k\": 3})\n",
    "\n",
    "retrieved_docs = retriever.invoke(\"<your question>\")\n",
    "\n",
    "print(retrieved_docs[0].page_content)\n",
    "\n",
    "# Use a prompt for RAG that is checked into the LangChain prompt hub (https://smith.langchain.com/hub/rlm/rag-prompt?organizationId=989ad331-949f-4bac-9694-660074a208a7)\n",
    "prompt = hub.pull(\"rlm/rag-prompt\")\n",
    "llm = AzureChatOpenAI(\n",
    "    openai_api_version=\"<Azure OpenAI API version>\",  # e.g., \"2023-07-01-preview\"\n",
    "    azure_deployment=\"<your chat model deployment name>\",\n",
    "    temperature=0,\n",
    ")\n",
    "\n",
    "\n",
    "def format_docs(docs):\n",
    "    return \"\\n\\n\".join(doc.page_content for doc in docs)\n",
    "\n",
    "\n",
    "rag_chain = (\n",
    "    {\"context\": retriever | format_docs, \"question\": RunnablePassthrough()}\n",
    "    | prompt\n",
    "    | llm\n",
    "    | StrOutputParser()\n",
    ")"
   ]
  },
  {