Chap3-RAG.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chapter 3. Retrieval: How to Chat with Your Data with RAG"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Basic Retrival-Augmented Generation pipeline\n",
    "\n",
    "## Indexing\n",
    "This stage involves preprocessing the external data source and storing embeddings that represent the data in a vector store where they can be easily retrieved.\n",
    "\n",
    "## Retrieval\n",
    "This stage involves retrieving the relevant embeddings and data stored in the Vector Store based on a user’s query.\n",
    "\n",
    "## Generation\n",
    "This stage involves synthesizing the original prompt with the retrieved relevant documents as one final prompt sent to the model for a prediction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "# Install langchain-chroma and lark, which is required for using the self-query retriever\n",
    "# See here for details: https://python.langchain.com/docs/how_to/self_query/\n",
    "%pip install --upgrade --quiet lark langchain-chroma"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import modules\n",
    "from langchain_community.document_loaders import TextLoader\n",
    "from langchain_community.vectorstores import FAISS\n",
    "from langchain_chroma import Chroma\n",
    "from langchain_openai import OpenAIEmbeddings\n",
    "from langchain_openai import OpenAIEmbeddings, ChatOpenAI, OpenAI\n",
    "from langchain_text_splitters import RecursiveCharacterTextSplitter\n",
    "from langchain_core.prompts import ChatPromptTemplate, PromptTemplate\n",
    "from langchain_core.runnables import chain\n",
    "from langchain_core.documents import Document\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_community.utils.math import cosine_similarity\n",
    "from langchain.chains.query_constructor.base import AttributeInfo\n",
    "from langchain.retrievers.self_query.base import SelfQueryRetriever\n",
    "from typing import Literal\n",
    "from pydantic import BaseModel, Field\n",
    "\n",
    "# Import modules for the SQL database query example\n",
    "import sqlite3\n",
    "import requests\n",
    "from langchain_community.utilities.sql_database import SQLDatabase\n",
    "from langchain_community.tools.sql_database.tool import QuerySQLDataBaseTool\n",
    "from langchain.chains import create_sql_query_chain\n",
    "from sqlalchemy import create_engine\n",
    "from sqlalchemy.pool import StaticPool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Load the document \n",
    "loader = TextLoader(\"TeachingwithGenerativeAI.txt\")\n",
    "doc = loader.load()\n",
    "\n",
    "## Split the document into chunks\n",
    "text_splitter = RecursiveCharacterTextSplitter(\n",
    "    chunk_size=1000,\n",
    "    chunk_overlap=20,\n",
    ")\n",
    "chunks = text_splitter.split_documents(doc)\n",
    "\n",
    "## Define the embedding model\n",
    "embed_model = OpenAIEmbeddings()\n",
    "\n",
    "# Create the vector store\n",
    "vector_db = FAISS.from_documents(\n",
    "    documents = chunks, \n",
    "    embedding = embed_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "number of docs found: 4\n",
      "\n",
      "* Students only learn from productive effort, and should understand how misuse or overuse of AI threatens that effort\n",
      "In order to help students understand these things, we recommend that instructors:\n",
      "* Explain your AI policy in your syllabus, and discuss the reasons you adopted it in class\n",
      "* Be specific about Dos and Don’ts—“Do acknowledge and describe any AI use”, or “Don’t use any AI for anything other than suggesting topics and sources”\n",
      "* Explain the limitations of generative AI. \n",
      "* Remember that students generally want to learn, and explain to them what they can learn from doing the work, not just the potential punishments for cheating\n",
      "Additional details can be found in the AI FAQ and in Adapting Assignments to AI.\n",
      "Academic Integrity and Generative AI\n"
     ]
    }
   ],
   "source": [
    "# Create a retriever\n",
    "retriever = vector_db.as_retriever()\n",
    "\n",
    "# Fetch relevant documents\n",
    "docs = retriever.invoke(\"Should a faculty member explain their AI policy in a class?\")\n",
    "print(f\"number of docs found: {len(docs)}\")\n",
    "print()\n",
    "print(docs[0].page_content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "number of docs found: 2\n"
     ]
    }
   ],
   "source": [
    "# Create a retriever with k=2\n",
    "retriever = vector_db.as_retriever(search_kwargs={\"k\": 2})\n",
    "\n",
    "# Fetch relevant documents\n",
    "docs = retriever.invoke(\"Should a faculty member explain their AI policy in a class?\")\n",
    "print(f\"number of docs found: {len(docs)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generating LLM Predictions Using Relevant Documents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "content='Yes, a faculty member should explain their AI policy in a class.' additional_kwargs={'refusal': None} response_metadata={'token_usage': {'completion_tokens': 14, 'prompt_tokens': 439, 'total_tokens': 453, 'completion_tokens_details': {'audio_tokens': 0, 'reasoning_tokens': 0, 'accepted_prediction_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-3.5-turbo-0125', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None} id='run-bf8fb9d3-e550-4231-bd63-5a6728943fca-0' usage_metadata={'input_tokens': 439, 'output_tokens': 14, 'total_tokens': 453, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}}\n"
     ]
    }
   ],
   "source": [
    "# Build a prompt template\n",
    "prompt_temp = ChatPromptTemplate.from_template(\"\"\"Answer the question based only on the following context:\n",
    "{context}\n",
    "Question: {question}\n",
    "\"\"\")\n",
    "\n",
    "# Choose a chatmodel\n",
    "chatmodel = ChatOpenAI(model_name=\"gpt-3.5-turbo\", temperature=0)\n",
    "\n",
    "# Create a chain\n",
    "chain = prompt_temp | chatmodel\n",
    "\n",
    "# Fetch relevant documents\n",
    "docs = retriever.invoke(\"Should a faculty member explain their AI policy in a class?\")\n",
    "\n",
    "# Invoke the chain to answer the question\n",
    "response = chain.invoke({\n",
    "    \"question\": \"Should a faculty member explain their AI policy in a class?\", \n",
    "    \"context\": docs})\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_1294/1939274990.py:5: LangChainDeprecationWarning: The method `BaseRetriever.get_relevant_documents` was deprecated in langchain-core 0.1.46 and will be removed in 1.0. Use :meth:`~invoke` instead.\n",
      "  docs = retriever.get_relevant_documents(question)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Yes, a faculty member should explain their AI policy in a class according to the provided context.\n"
     ]
    }
   ],
   "source": [
    "# Incoporating the above pipeline into a function\n",
    "@chain\n",
    "def qa(question):\n",
    "    # Fetch relevant documents \n",
    "    docs = retriever.get_relevant_documents(question)\n",
    "    # Invote the prompt template\n",
    "    prompt = prompt_temp.invoke({\"context\": docs, \"question\": question})\n",
    "    # Generate a response\n",
    "    response = chatmodel.invoke(prompt)\n",
    "    return response\n",
    "\n",
    "# Run the function\n",
    "# Note: Given the qa() function is a runnable chain, it should be invoked with qa.invoke()\n",
    "response = qa.invoke(\"Should a faculty member explain their AI policy in a class?\")\n",
    "print(response.content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Yes, a faculty member should explain their AI policy in a class.\n",
      "\n",
      "2\n",
      "\n",
      "* Students only learn from productive effort, and should understand how misuse or overuse of AI threatens that effort\n",
      "In order to help students understand these things, we recommend that instructors:\n",
      "* Explain your AI policy in your syllabus, and discuss the reasons you adopted it in class\n",
      "* Be specific about Dos and Don’ts—“Do acknowledge and describe any AI use”, or “Don’t use any AI for anything other than suggesting topics and sources”\n",
      "* Explain the limitations of generative AI. \n",
      "* Remember that students generally want to learn, and explain to them what they can learn from doing the work, not just the potential punishments for cheating\n",
      "Additional details can be found in the AI FAQ and in Adapting Assignments to AI.\n",
      "Academic Integrity and Generative AI\n",
      "\n",
      "Students and faculty report growing use of generative AI—tools that produce human-like writing (e.g ChatGPT), images (e.g. MidJourney), code (e.g. Microsoft Co-Pilot) and the like. The flexibility of these tools mean that there is no current default for acceptable vs. unacceptable use of these tools in coursework, and student adoption is moving faster than faculty adaptation. Many students are using AI without clear directions from their instructors about which uses are acceptable. \n",
      "Faculty should explain to students what is and is not allowed around AI use in their classes. (We know some faculty believe students are not using AI in their classes. If this is you, we assure you you are wrong.) The Provost’s Office recommends that faculty adopt the following three principles for student use of AI:\n",
      "* When a student uses these tools, they should acknowledge that use\n",
      "* The student is responsible for the content and accuracy of any work they submit, however created\n"
     ]
    }
   ],
   "source": [
    "# We can also return the retrieved documents for further inspection\n",
    "@chain\n",
    "def qa(question):\n",
    "    # Fetch relevant documents \n",
    "    docs = retriever.get_relevant_documents(question)\n",
    "    # Invote the prompt template\n",
    "    prompt = prompt_temp.invoke({\"context\": docs, \"question\": question})\n",
    "    # Generate a response\n",
    "    response = chatmodel.invoke(prompt)\n",
    "    return response, docs\n",
    "response, docs = qa.invoke(\"Should a faculty member explain their AI policy in a class?\")\n",
    "print(response.content)\n",
    "print()\n",
    "print(len(docs))\n",
    "print()\n",
    "print(docs[0].page_content)\n",
    "print()\n",
    "print(docs[1].page_content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Query Transformation\n",
    "In a production setting, a user is likely to construct their query in an incomplete, ambiguous, or poorly worded manner that leads to model hallucination. Query transformation is a subset of strategies designed to modify the user’s input to answer the first RAG problem question."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Rewrite-Retrieve-Read\n",
    "The Rewrite-Retrieve-Read strategy proposed by a Microsoft Research team simply prompts the LLM to rewrite the user’s query before performing retrieval."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Faculty should declare their AI policy to their students, as there is currently no default for acceptable vs. unacceptable use of generative AI tools in coursework. Students are using AI without clear directions from their instructors about which uses are acceptable, so it is important for faculty to explain to students what is and is not allowed around AI use in their classes.\n"
     ]
    }
   ],
   "source": [
    "# A demonstration of a poorly written prompt on RAG response\n",
    "@chain\n",
    "def qa(question):\n",
    "    # Fetch relevant documents \n",
    "    docs = retriever.get_relevant_documents(question)\n",
    "    # Invote the prompt template\n",
    "    prompt = prompt_temp.invoke({\"context\": docs, \"question\": question})\n",
    "    # Generate a response\n",
    "    response = chatmodel.invoke(prompt)\n",
    "    return response\n",
    "question = \"I don't know what to say. The weather is perfect for a walk. Should faculty declare their AI policy or not to their students?\"\n",
    "response = qa.invoke(question)\n",
    "print(response.content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Should faculty declare their AI policy to students?'"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create a rewrite prompt template\n",
    "rewrite_prompt_temp = ChatPromptTemplate.from_template(\"\"\"\n",
    "Provide a better search query for the vectore database search engine to answer the given question, end the queries with ’**’. Question: {x} Answer:\"\"\")\n",
    "\n",
    "# Define a function to parse the message from the chat model\n",
    "def parse_rewriter_output(message):\n",
    "    return message.content.strip('\"').strip(\"**\")\n",
    "\n",
    "# Define the rewriter chain\n",
    "rewriter = rewrite_prompt_temp | chatmodel | parse_rewriter_output\n",
    "\n",
    "# Test the rewriter chain\n",
    "rewriter.invoke(question)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Yes, faculty should declare their AI policy to students.\n"
     ]
    }
   ],
   "source": [
    "# Redefine the RAG pipeline function to incoporate the above rewrite chain\n",
    "@chain\n",
    "def qa_rrr(question):\n",
    "    # Invote the rewrite prompt template\n",
    "    rewritten_question = rewriter.invoke(question)\n",
    "    \n",
    "    # Fetch relevant documents\n",
    "    docs = retriever.get_relevant_documents(rewritten_question)\n",
    "    \n",
    "    # Assemble the prompt, which now includes the retrived documents and the rewritten question\n",
    "    prompt = prompt_temp.invoke({\"context\": docs, \"question\": rewritten_question})\n",
    "\n",
    "    # Generate a response\n",
    "    response = chatmodel.invoke(prompt)\n",
    "    return response\n",
    "question = \"I don't know what to say. The weather is perfect for a walk. Should faculty declare their AI policy or not to their students?\"\n",
    "response = qa_rrr.invoke(question)\n",
    "print(response.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Multi Query Retrieval\n",
    "A user’s single query can be insufficient to capture the full scope of information required to answer the query comprehensively. The Multi Query Retrieval strategy resolves this problem by instructing an LLM to generate multiple queries based on a user’s initial query, executing a parallel retrieval of each query from the data source, and then inserting the retrieved results as prompt context to generate a final model output.\n",
    "\n",
    "This strategy is particularly useful for use cases where a single question may rely on multiple perspectives to provide a comprehensive answer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5\n",
      "\n",
      "['1. Is it necessary for faculty to disclose their AI policy to students?', '2. Should faculty communicate their AI policy to students or keep it confidential?', '3. What are the implications of faculty sharing their AI policy with students?', '4. Is transparency about AI policies important for faculty-student relationships?', '5. How does faculty disclosure of their AI policy impact student trust and understanding?']\n"
     ]
    }
   ],
   "source": [
    "# Create a multi-perpspective prompt template\n",
    "perspectives_prompt_temp = ChatPromptTemplate.from_template(\n",
    "    \"\"\"You are an AI language model assistant. \n",
    "    Your task is to generate five different versions of the given user question to retrieve relevant documents from a vector database. \n",
    "    By generating multiple perspectives on the user question, your goal is to help the user overcome some of the limitations of the distance-based similarity search. \n",
    "    Provide these alternative questions separated by newlines. Original question: {question}\"\"\")\n",
    "def parse_queries_output(message):\n",
    "    return message.content.split('\\n')\n",
    "\n",
    "# Build and invoke the query generator chain\n",
    "query_gen = perspectives_prompt_temp | chatmodel | parse_queries_output\n",
    "lst_perspectives = query_gen.invoke(question)\n",
    "print(len(lst_perspectives))\n",
    "print()\n",
    "print(lst_perspectives)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Take the list of generated queries, retrieve the most relevant docs for each of them in parallel, \n",
    "# and then combine to get the unique union of all the retrieved relevant documents.\n",
    "def get_unique_union(document_lists):\n",
    "    # Flatten list of lists, and deduplicate them by including them in a dictionary.\n",
    "    # Note: a dictionoary can't have duplicate keys, so we can use it to deduplicate.\n",
    "    deduped_docs = {\n",
    "        doc.page_content: doc for sublist in document_lists for doc in sublist\n",
    "    }\n",
    "    # return a flat list of unique docs\n",
    "    return list(deduped_docs.values())\n",
    "\n",
    "# Build the retrieval chain\n",
    "# Note: Use the retriever.batch() method to retrieve documents in parallel.\n",
    "retrieval_chain = query_gen | retriever.batch | get_unique_union"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Faculty should declare their AI policy to their students.\n"
     ]
    }
   ],
   "source": [
    "# Redefine the RAG pipeline function to incoporate the above multi-query chain\n",
    "@chain\n",
    "def multi_query_qa(question):\n",
    "    # Invote the multi-query prompt template to obtain a list of queries\n",
    "    lst_perspectives = query_gen.invoke(question)\n",
    "    \n",
    "    # Fetch relevant documents\n",
    "    docs = retrieval_chain.invoke(lst_perspectives)\n",
    "    \n",
    "    # Assemble the prompt, which now includes the retrived documents and the rewritten question\n",
    "    prompt = prompt_temp.invoke({\"context\": docs, \"question\": question})\n",
    "\n",
    "    # Generate a response\n",
    "    response = chatmodel.invoke(prompt)\n",
    "    return response\n",
    "question = \"I don't know what to say. The weather is perfect for a walk. Should faculty declare their AI policy or not to their students?\"\n",
    "response = multi_query_qa.invoke(question)\n",
    "print(response.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RAG-Fusion\n",
    "The RAG-Fusion strategy shares similarities with the Multi Query Retrieval strategy, except that we apply a final reranking step to all the retrieved documents. This reranking step makes use of the Reciprocal Rank Fusion (RRF) algorithm, which involves combining the ranks of different search results to produce a single, unified ranking. By combining ranks from different queries, we pull the most relevant documents to the top of the final list."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4\n",
      "\n",
      "[\"1. How to overcome writer's block when you don't know what to say?\", '2. Benefits of walking in perfect weather for physical and mental health.', '3. Importance of transparency in faculty declaring their AI policy to students.', '4. Pros and cons of faculty disclosing their AI policy to students.']\n"
     ]
    }
   ],
   "source": [
    "# Create a prompt template for RAG-Fusion\n",
    "prompt_rag_fusion = ChatPromptTemplate.from_template(\"\"\"\n",
    "You are a helpful assistant that generates multiple search queries based on a single input query. \\n\n",
    "Generate multiple search queries related to: {question} \\n\n",
    "Output (4 queries):\"\"\")\n",
    "def parse_queries_output(message):\n",
    "    return message.content.split('\\n')\n",
    "chatmodel = ChatOpenAI(temperature=0)\n",
    "\n",
    "# Build and invoke the query generator chain\n",
    "query_gen = prompt_rag_fusion | chatmodel | parse_queries_output\n",
    "lst_perspectives = query_gen.invoke(question)\n",
    "print(len(lst_perspectives))\n",
    "print()\n",
    "print(lst_perspectives)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The function \"reciprocal_rank_fusion()\" takes a list of the search results of each query, so a list of lists of documents, where each inner list of documents is sorted by their relevance to that query. The RRF algorithm then calculates a new score for each document based on its ranks (or positions) in the different lists and sorts them to create a final reranked list. After calculating the fused scores, the function sorts the documents in descending order of these scores to get the final reranked list, which is then returned."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the reciprocal rank fusion function.\n",
    "# Note: Notice the function also takes a k parameter, \n",
    "# which determines how much influence documents in each query’s result sets have \n",
    "# over the final list of documents. \n",
    "# A higher value indicates that lower ranked documents have more influence.\n",
    "def reciprocal_rank_fusion(results: list[list], k=60):\n",
    "    \"\"\"reciprocal rank fusion on multiple lists of ranked documents \n",
    "       and an optional parameter k used in the RRF formula\n",
    "    \"\"\"\n",
    "    \n",
    "    # Initialize a dictionary to hold fused scores for each document\n",
    "    # Documents will be keyed by their contents to ensure uniqueness\n",
    "    fused_scores = {}\n",
    "    documents = {}\n",
    "    # Iterate through each list of ranked documents\n",
    "    for docs in results:\n",
    "        # Iterate through each document in the list, with its rank (position in the list)\n",
    "        for rank, doc in enumerate(docs):\n",
    "            # Use the document contents as the key for uniqueness\n",
    "            doc_str = doc.page_content\n",
    "            # If the document hasn't been seen yet,\n",
    "            # - initialize score to 0\n",
    "            # - save it for later\n",
    "            if doc_str not in fused_scores:\n",
    "                fused_scores[doc_str] = 0\n",
    "                documents[doc_str] = doc\n",
    "            # Update the score of the document using the RRF formula:\n",
    "            # 1 / (rank + k)\n",
    "            fused_scores[doc_str] += 1 / (rank + k)\n",
    "    # Sort the documents based on their fused scores in descending order to get the final reranked results\n",
    "    reranked_doc_strs = sorted(\n",
    "        fused_scores, key=lambda d: fused_scores[d], reverse=True\n",
    "    )\n",
    "    # retrieve the corresponding doc for each doc_str\n",
    "    return [\n",
    "        documents[doc_str]\n",
    "        for doc_str in reranked_doc_strs\n",
    "    ]\n",
    "\n",
    "# Build the retrieval chain\n",
    "retrieval_chain = query_gen | retriever.batch | reciprocal_rank_fusion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content=\"Based on the context provided, faculty should declare their AI policy to their students. It is recommended that faculty explain to students what is and is not allowed around AI use in their classes. This includes acknowledging the use of AI tools, being responsible for the content and accuracy of any work submitted, and discussing Dos and Don'ts related to AI use. It is important for faculty to help students understand the limitations of generative AI and the potential consequences of misuse or overuse of AI.\", additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 97, 'prompt_tokens': 688, 'total_tokens': 785, 'completion_tokens_details': {'audio_tokens': 0, 'reasoning_tokens': 0, 'accepted_prediction_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-3.5-turbo-0125', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None}, id='run-af0a6678-a46f-4638-8523-2fdfa1e45130-0', usage_metadata={'input_tokens': 688, 'output_tokens': 97, 'total_tokens': 785, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Build a prompt template that take the output from the retrieval chain and the question as input\n",
    "prompt = ChatPromptTemplate.from_template(\"\"\"\n",
    "Answer the following question based on this context:\n",
    "{context}\n",
    "Question: {question}\n",
    "\"\"\")\n",
    "\n",
    "# Build the multi-query chain with RAG-funtion\n",
    "@chain\n",
    "def multi_query_qa(input):\n",
    "    # fetch relevant documents \n",
    "    docs = retrieval_chain.invoke(input)\n",
    "    # format prompt\n",
    "    formatted = prompt.invoke({\"context\": docs, \"question\": input})\n",
    "    # generate answer\n",
    "    answer = chatmodel.invoke(formatted)\n",
    "    return answer\n",
    "\n",
    "# Invoke the chain\n",
    "multi_query_qa.invoke(question)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hypothetical Document Embeddings (HyDE)\n",
    "Hypothetical Document Embeddings (HyDE) is a strategy that involves creating a hypothetical document based on the user’s query, embedding the document, and retrieving relevant documents based on vector similarity. The intuition behind HyDE is that an LLM-generated hypothetical document will be more similar to the most relevant documents than the original query."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"In recent years, the integration of artificial intelligence (AI) technologies in educational settings has become increasingly prevalent. As such, the question of whether faculty should be required to declare their AI policy for their students has garnered significant attention. \\n\\nOne argument in favor of requiring faculty to declare their AI policy is the importance of transparency and accountability in the use of AI technologies in education. By clearly outlining their AI policy, faculty can ensure that students are aware of how AI is being utilized in their learning environment and can make informed decisions about their participation. This transparency can also help to build trust between faculty and students, as well as promote a culture of openness and communication.\\n\\nAdditionally, declaring an AI policy can help to mitigate potential ethical concerns surrounding the use of AI in education. Faculty can outline how AI technologies are being used, the data being collected, and how that data is being used to inform teaching and learning practices. This can help to ensure that students' privacy and data security are being protected, and that ethical considerations are being taken into account in the implementation of AI technologies.\\n\\nOn the other hand, some may argue that requiring faculty to declare their AI policy could be burdensome and unnecessary, particularly if the use of AI technologies is limited or if faculty are already following established guidelines and best practices. However, given the rapid pace of technological advancement and the potential impact of AI on education, it may be prudent for faculty to proactively declare their AI policy to ensure that they are staying current with emerging trends and addressing any potential concerns that may arise.\\n\\nIn conclusion, while there may be differing opinions on whether faculty should be required to declare their AI policy for their students, the benefits of transparency, accountability, and ethical considerations suggest that such a requirement could be beneficial in promoting responsible and effective use of AI technologies in educational settings. Further research and discussion on this topic are warranted to explore the potential implications and best practices for implementing AI policies in education.\""
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Build a prompt template for the HyDE\n",
    "prompt_hyde = ChatPromptTemplate.from_template(\"\"\"\n",
    "Please write a scientific paper passage to answer the question\n",
    "Question: {question}\n",
    "Passage:\"\"\")\n",
    "\n",
    "# Build the chain to generate document from the prompt\n",
    "generate_doc = prompt_hyde | ChatOpenAI(temperature=0) | StrOutputParser()\n",
    "\n",
    "# Invote the chain\n",
    "generate_doc.invoke(\"Should faculty be required to declare their AI policy for their students?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build the retrieval chain\n",
    "retrieval_chain = generate_doc | retriever"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content=\"Based on the context provided, it is recommended that faculty declare their AI policy to their students. This is important in order to help students understand the dos and don'ts of AI use, the limitations of generative AI, and to ensure academic integrity. By explaining the AI policy in the syllabus and discussing it in class, faculty can provide clear directions to students on what is and is not allowed in terms of AI use in coursework.\", additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 88, 'prompt_tokens': 455, 'total_tokens': 543, 'completion_tokens_details': {'audio_tokens': 0, 'reasoning_tokens': 0, 'accepted_prediction_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-3.5-turbo-0125', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None}, id='run-0d8fba3f-ed53-48af-86bb-9008b42d2134-0', usage_metadata={'input_tokens': 455, 'output_tokens': 88, 'total_tokens': 543, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Build a prompt template that take the output from the retrieval chain and the question as input\n",
    "prompt = ChatPromptTemplate.from_template(\"\"\"\n",
    "Answer the following question based on this context:\n",
    "{context}\n",
    "Question: {question}\n",
    "\"\"\")\n",
    "\n",
    "# Build the HyDe chain\n",
    "@chain\n",
    "def hyde_query_qa(input):\n",
    "    # fetch relevant documents \n",
    "    docs = retrieval_chain.invoke(input)\n",
    "    # format prompt\n",
    "    formatted = prompt.invoke({\"context\": docs, \"question\": input})\n",
    "    # generate answer\n",
    "    answer = chatmodel.invoke(formatted)\n",
    "    return answer\n",
    "\n",
    "# Invoke the chain\n",
    "hyde_query_qa.invoke(question)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Query Routing\n",
    "Although using a single Vector Store is useful, the required data may live in a variety of data sources, including relational databases or other Vector Stores.\n",
    "\n",
    "For example, you may have two Vector Stores: one for LangChain Python documentation and another for LangChain JS documentation. Given a user’s question, we would like to route the query to the appropriate inferred data source to retrieve relevant docs. Query routing is a strategy used to forward a user’s query to the relevant data source."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Logical Routing\n",
    "In logical routing, we give the LLM knowledge of the various data sources at our disposal and then let the LLM reason which data source to apply based on the user’s query."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Note: a RouteQuery model that validates input for a field called datasource, \n",
    "# ensuring it is either \"python_docs\" or \"js_docs\", \n",
    "# with the purpose of routing a user’s question to the appropriate documentation source.\n",
    "class RouteQuery(BaseModel):\n",
    "    \"\"\"Route a user query to the most relevant datasource.\"\"\"\n",
    "    # Note: (1) datasource: This attribute uses Literal, meaning it can only take specific values, \n",
    "    # \"python_docs\" or \"js_docs\". This restricts datasource to these two values, \n",
    "    # ensuring that any instance of RouteQuery must have datasource set to one of them.\n",
    "    # (2) Field: The Field function from pydantic adds metadata to the field. \n",
    "    # Here, ... (Ellipsis) is used to indicate that datasource is a required field (without default value), \n",
    "    # and description provides additional information about the field’s purpose.\n",
    "    datasource: Literal[\"python_docs\", \"js_docs\"] = Field(\n",
    "        ...,\n",
    "        description=\"Given a user question choose which datasource would be most relevant for answering their question\",\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "# LLM with function call\n",
    "chatmodel = ChatOpenAI(model=\"gpt-3.5-turbo\", temperature=0)\n",
    "structured_llm = chatmodel.with_structured_output(RouteQuery)\n",
    "\n",
    "# Prompt \n",
    "system = \"\"\"You are an expert at routing a user question to the appropriate data source.\n",
    "Based on the programming language the question is referring to, route it to the relevant data source.\"\"\"\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", system),\n",
    "        (\"human\", \"{question}\"),\n",
    "    ]\n",
    ")\n",
    "# Define router \n",
    "router = prompt | structured_llm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "python_docs\n",
      "\n",
      "datasource='python_docs'\n"
     ]
    }
   ],
   "source": [
    "question = '''\n",
    "Why doesn't the following code work:\n",
    "from PIL import Image\n",
    "import numpy as np\n",
    "\n",
    "def process_image(image_path, target_size=(224, 224)):\n",
    "    \"\"\"\n",
    "    Processes an image by resizing and normalizing it.\n",
    "    \n",
    "    Args:\n",
    "        image_path (str): Path to the input image file.\n",
    "        target_size (tuple): Desired output size (width, height).\n",
    "    \n",
    "    Returns:\n",
    "        np.ndarray: The processed image as a normalized numpy array.\n",
    "    \"\"\"\n",
    "    # Load the image\n",
    "    img = Image.open(image_path)\n",
    "    \n",
    "    # Resize the image to the target size\n",
    "    img = img.resize(target_size)\n",
    "    \n",
    "    # Convert the image to a numpy array\n",
    "    img_array = np.array(img)\n",
    "    \n",
    "    # Normalize the image data to a range of 0 to 1\n",
    "    img_array = img_array / 255.0\n",
    "    \n",
    "    # If the image is grayscale, expand dimensions to (height, width, 1)\n",
    "    if len(img_array.shape) == 2:\n",
    "        img_array = np.expand_dims(img_array, axis=-1)\n",
    "    \n",
    "    return img_array\n",
    "'''\n",
    "result = router.invoke({\"question\": question})\n",
    "print(result.datasource)\n",
    "print()\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a functioin to choose the route\n",
    "def choose_route(result):\n",
    "    if \"python_docs\" in result.datasource.lower():\n",
    "        ### Logic here \n",
    "        return \"chain for python_docs\"\n",
    "    else:\n",
    "        ### Logic here \n",
    "        return \"chain for js_docs\"\n",
    "\n",
    "# Then we can use the choose_route function to choose the route\n",
    "# full_chain = router | RunnableLambda(choose_route)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Semantic Routing\n",
    "Unlike logical routing, semantic routing involves embedding various prompts that represent various data sources alongside the user’s query, and then performing vector similarity search to retrieve the most similar prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Two prompts\n",
    "climate_template = \"\"\"You are a meteorologist. You are great at answering questions about weather and climate.\n",
    "Here is a question:\n",
    "{query}\"\"\"\n",
    "fire_template = \"\"\"You are a firefighter. You are great at answering questions about fire and emergency situations.\n",
    "Here is a question:\n",
    "{query}\"\"\"\n",
    "\n",
    "# Embed prompts\n",
    "embeddings = OpenAIEmbeddings()\n",
    "prompt_templates = [climate_template, fire_template]\n",
    "prompt_embeddings = embeddings.embed_documents(prompt_templates)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "text='You are a firefighter. You are great at answering questions about fire and emergency situations.\\nHere is a question:\\nWould calling a fire department be a good idea?'\n",
      "\n",
      "Yes, calling the fire department in an emergency situation is always a good idea. The fire department has trained professionals who can respond quickly and effectively to help mitigate the situation and ensure the safety of everyone involved. It is important to call 911 as soon as possible in the event of a fire or any other emergency.\n"
     ]
    }
   ],
   "source": [
    "# Define the function to route question to prompt that repesents the underling database\n",
    "@chain\n",
    "def prompt_router(query):\n",
    "    # Embed question\n",
    "    query_embedding = embeddings.embed_query(query)\n",
    "    # Compute similarity\n",
    "    similarity = cosine_similarity([query_embedding], prompt_embeddings)[0]\n",
    "    # Pick the prompt most similar to the input question\n",
    "    most_similar = prompt_templates[similarity.argmax()]\n",
    "    return PromptTemplate.from_template(most_similar)\n",
    "\n",
    "# Test the above function\n",
    "query = \"Would calling a fire department be a good idea?\"\n",
    "print(prompt_router.invoke(query))\n",
    "print()\n",
    "\n",
    "# Build the semantic router\n",
    "semantic_router = (\n",
    "    prompt_router\n",
    "    | ChatOpenAI()\n",
    "    | StrOutputParser()\n",
    ")\n",
    "\n",
    "# Invote the router\n",
    "response = semantic_router.invoke(query)\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Query Construction\n",
    "Retrieval-augmented generation is an effective strategy to embed and retrieve relevant unstructured data from a Vector Store based on a query. But most data available for use in production apps is structured and typically stored in relational databases. In addition, unstructured data embedded in a Vector Store also contains structured metadata that posesses important information.\n",
    "\n",
    "Query construction is the process of transforming a natural language query into the query language of the database or data source you are interacting with.\n",
    "\n",
    "For example, consider the query what are movies about aliens in the year 1980? This question contains an unstructured topic that can be retrieved via embeddings (aliens), but it also contains potential structured components (“year == 1980”)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Text-to-metadata-filter using self-querying retriever\n",
    "A self-querying retriever is one that has the ability to query itself. Specifically, given any natural language query, the retriever uses a query-constructing LLM chain to write a structured query and then applies that structured query to its underlying VectorStore. This allows the retriever to not only use the user-input query for semantic similarity comparison with the contents of stored documents but to also extract filters from the user query on the metadata of stored documents and to execute those filters.\n",
    "\n",
    "See here fore more details: https://python.langchain.com/docs/how_to/self_query/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a vectorstore from a list of Documents\n",
    "# Note: We have to use Chroma here because FAISS doesn't support self-querying retriever\n",
    "docs = [\n",
    "    Document(\n",
    "        page_content=\"A bunch of scientists bring back dinosaurs and mayhem breaks loose\",\n",
    "        metadata={\"year\": 1993, \"rating\": 7.7, \"genre\": \"science fiction\"},\n",
    "    ),\n",
    "    Document(\n",
    "        page_content=\"Leo DiCaprio gets lost in a dream within a dream within a dream within a ...\",\n",
    "        metadata={\"year\": 2010, \"director\": \"Christopher Nolan\", \"rating\": 8.2},\n",
    "    ),\n",
    "    Document(\n",
    "        page_content=\"A psychologist / detective gets lost in a series of dreams within dreams within dreams and Inception reused the idea\",\n",
    "        metadata={\"year\": 2006, \"director\": \"Satoshi Kon\", \"rating\": 8.6},\n",
    "    ),\n",
    "    Document(\n",
    "        page_content=\"A bunch of normal-sized women are supremely wholesome and some men pine after them\",\n",
    "        metadata={\"year\": 2019, \"director\": \"Greta Gerwig\", \"rating\": 8.3},\n",
    "    ),\n",
    "    Document(\n",
    "        page_content=\"Toys come alive and have a blast doing so\",\n",
    "        metadata={\"year\": 1995, \"genre\": \"animated\"},\n",
    "    ),\n",
    "    Document(\n",
    "        page_content=\"Three men walk into the Zone, three men walk out of the Zone\",\n",
    "        metadata={\n",
    "            \"year\": 1979,\n",
    "            \"director\": \"Andrei Tarkovsky\",\n",
    "            \"genre\": \"thriller\",\n",
    "            \"rating\": 9.9,\n",
    "        },\n",
    "    ),\n",
    "]\n",
    "\n",
    "# Create a vectorstore from the documents\n",
    "vectorstore = Chroma.from_documents(docs, OpenAIEmbeddings())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now we can instantiate our retriever. To do this we'll need to provide some information upfront about the metadata fields \n",
    "# # that our documents support and a short description of the document contents.\n",
    "metadata_field_info = [\n",
    "    AttributeInfo(\n",
    "        name=\"genre\",\n",
    "        description=\"The genre of the movie. One of ['science fiction', 'comedy', 'drama', 'thriller', 'romance', 'action', 'animated']\",\n",
    "        type=\"string\",\n",
    "    ),\n",
    "    AttributeInfo(\n",
    "        name=\"year\",\n",
    "        description=\"The year the movie was released\",\n",
    "        type=\"integer\",\n",
    "    ),\n",
    "    AttributeInfo(\n",
    "        name=\"director\",\n",
    "        description=\"The name of the movie director\",\n",
    "        type=\"string\",\n",
    "    ),\n",
    "    AttributeInfo(\n",
    "        name=\"rating\", description=\"A 1-10 rating for the movie\", type=\"float\"\n",
    "    ),\n",
    "]\n",
    "document_content_description = \"Brief summary of a movie\"\n",
    "\n",
    "# Define the llm\n",
    "llm = ChatOpenAI(temperature=0)\n",
    "\n",
    "# Build the self-query retriever\n",
    "retriever = SelfQueryRetriever.from_llm(\n",
    "    llm,\n",
    "    vectorstore,\n",
    "    document_content_description,\n",
    "    metadata_field_info,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Document(metadata={'director': 'Andrei Tarkovsky', 'genre': 'thriller', 'rating': 9.9, 'year': 1979}, page_content='Three men walk into the Zone, three men walk out of the Zone'),\n",
       " Document(metadata={'director': 'Satoshi Kon', 'rating': 8.6, 'year': 2006}, page_content='A psychologist / detective gets lost in a series of dreams within dreams within dreams and Inception reused the idea')]"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Test it out\n",
    "# Note: This example only specifies a filter.\n",
    "retriever.invoke(\"I want to watch a movie rated higher than 8.5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Document(metadata={'director': 'Greta Gerwig', 'rating': 8.3, 'year': 2019}, page_content='A bunch of normal-sized women are supremely wholesome and some men pine after them')]"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# This example specifies a query and a filter\n",
    "retriever.invoke(\"Has Greta Gerwig directed any movies about women\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Document(metadata={'director': 'Satoshi Kon', 'rating': 8.6, 'year': 2006}, page_content='A psychologist / detective gets lost in a series of dreams within dreams within dreams and Inception reused the idea'),\n",
       " Document(metadata={'director': 'Andrei Tarkovsky', 'genre': 'thriller', 'rating': 9.9, 'year': 1979}, page_content='Three men walk into the Zone, three men walk out of the Zone')]"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# This example specifies a composite filter\n",
    "retriever.invoke(\"What's a highly rated (above 8.5) science fiction film?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Document(metadata={'genre': 'animated', 'year': 1995}, page_content='Toys come alive and have a blast doing so')]"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# This example specifies a query and composite filter\n",
    "retriever.invoke(\n",
    "    \"What's a movie after 1990 but before 2005 that's all about toys, and preferably is animated\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SQL\n",
    "Create a chain to translate a question to an SQL query and then execute the query to get the answer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['Album', 'Artist', 'Customer', 'Employee', 'Genre', 'Invoice', 'InvoiceLine', 'MediaType', 'Playlist', 'PlaylistTrack', 'Track']\n"
     ]
    }
   ],
   "source": [
    "# Create a function to get the engine for the Chinook database\n",
    "# Code adopted from here: https://python.langchain.com/docs/integrations/tools/sql_database/\n",
    "def get_engine_for_chinook_db():\n",
    "    \"\"\"Pull sql file, populate in-memory database, and create engine.\"\"\"\n",
    "    url = \"https://raw.githubusercontent.com/lerocha/chinook-database/master/ChinookDatabase/DataSources/Chinook_Sqlite.sql\"\n",
    "    response = requests.get(url)\n",
    "    sql_script = response.text\n",
    "\n",
    "    connection = sqlite3.connect(\":memory:\", check_same_thread=False)\n",
    "    connection.executescript(sql_script)\n",
    "    return create_engine(\n",
    "        \"sqlite://\",\n",
    "        creator=lambda: connection,\n",
    "        poolclass=StaticPool,\n",
    "        connect_args={\"check_same_thread\": False},\n",
    "    )\n",
    "\n",
    "# Build the engine\n",
    "engine = get_engine_for_chinook_db()\n",
    "\n",
    "# Build the database\n",
    "db = SQLDatabase(engine)\n",
    "\n",
    "# Print the table names in the database\n",
    "print(db.get_usable_table_names())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SELECT COUNT(\"EmployeeId\") AS \"TotalEmployees\" FROM \"Employee\"\n"
     ]
    }
   ],
   "source": [
    "# Define the LLM\n",
    "llm = ChatOpenAI(model=\"gpt-3.5-turbo\", temperature=0)\n",
    "\n",
    "# Create the sql query chain to translate the question to SQL query\n",
    "chain = create_sql_query_chain(llm, db)\n",
    "\n",
    "# Invoke the chain with the question\n",
    "response = chain.invoke({\"question\": \"How many employees are there?\"})\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'[(8,)]'"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create the chain to translate the question to SQL query and execute the query\n",
    "execute_query = QuerySQLDataBaseTool(db=db)\n",
    "write_query = create_sql_query_chain(llm, db)\n",
    "chain = write_query | execute_query\n",
    "\n",
    "# Invoke the chain\n",
    "chain.invoke({\"question\": \"How many employees are there?\"})"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}