{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fedd6d23",
   "metadata": {},
   "source": [
    "This directory is retained purely for archival purposes and is no longer updated. Please see the newly [consolidated LangChain documentation](https://docs.langchain.com/oss/python/langgraph/overview) for the most current information and resources."
   ]
  },
  {
   "attachments": {
    "36fa621a-9d3d-4860-a17c-5d20e6987481.png": {
     "image/png": "[encoded data, 186088 characters]"
    }
   },
   "cell_type": "markdown",
   "id": "5afcaed0-3d55-4e1f-95d3-c32c751c29d8",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "# Adaptive RAG\n",
    "\n",
    "Adaptive RAG is a strategy for RAG that unites (1) [query analysis](https://blog.langchain.dev/query-construction/) with (2) [active / self-corrective RAG](https://blog.langchain.dev/agentic-rag-with-langgraph/).\n",
    "\n",
    "In the [paper](https://arxiv.org/abs/2403.14403), they report query analysis to route across:\n",
    "\n",
    "* No Retrieval\n",
    "* Single-shot RAG\n",
    "* Iterative RAG\n",
    "\n",
    "Let's build on this using LangGraph. \n",
    "\n",
    "In our implementation, we will route between:\n",
    "\n",
    "* Web search: for questions related to recent events\n",
    "* Self-corrective RAG: for questions related to our index\n",
    "\n",
    "![Screenshot 2024-03-26 at 1.36.03 PM.png](attachment:36fa621a-9d3d-4860-a17c-5d20e6987481.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a85501ca-eb89-4795-aeab-cdab050ead6b",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "First, let's install our required packages and set our API keys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53d1a740-9fea-4a6e-8f95-fb9dbf1c80a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture --no-stderr\n",
    "! pip install -U langchain_community tiktoken langchain-openai langchain-cohere langchainhub chromadb langchain langgraph  tavily-python"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "222f204d-956f-4128-b597-2c698120edda",
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "\n",
    "def _set_env(var: str):\n",
    "    if not os.environ.get(var):\n",
    "        os.environ[var] = getpass.getpass(f\"{var}: \")\n",
    "\n",
    "\n",
    "_set_env(\"OPENAI_API_KEY\")\n",
    "_set_env(\"COHERE_API_KEY\")\n",
    "_set_env(\"TAVILY_API_KEY\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47e04b18",
   "metadata": {},
   "source": [
    "<div class=\"admonition tip\">\n",
    "    <p class=\"admonition-title\">Set up <a href=\"https://smith.langchain.com\">LangSmith</a> for LangGraph development</p>\n",
    "    <p style=\"padding-top: 5px;\">\n",
    "        Sign up for LangSmith to quickly spot issues and improve the performance of your LangGraph projects. LangSmith lets you use trace data to debug, test, and monitor your LLM apps built with LangGraph — read more about how to get started <a href=\"https://docs.smith.langchain.com\">here</a>. \n",
    "    </p>\n",
    "</div>    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ac1c2cd-81fb-40eb-8ba1-e9197800cba6",
   "metadata": {},
   "source": [
    "## Create Index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b224e5ba-50ca-495a-a7fa-0f75a080e03c",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Build Index\n",
    "\n",
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "from langchain_community.document_loaders import WebBaseLoader\n",
    "from langchain_community.vectorstores import Chroma\n",
    "from langchain_openai import OpenAIEmbeddings\n",
    "\n",
    "### from langchain_cohere import CohereEmbeddings\n",
    "\n",
    "# Set embeddings\n",
    "embd = OpenAIEmbeddings()\n",
    "\n",
    "# Docs to index\n",
    "urls = [\n",
    "    \"https://lilianweng.github.io/posts/2023-06-23-agent/\",\n",
    "    \"https://lilianweng.github.io/posts/2023-03-15-prompt-engineering/\",\n",
    "    \"https://lilianweng.github.io/posts/2023-10-25-adv-attack-llm/\",\n",
    "]\n",
    "\n",
    "# Load\n",
    "docs = [WebBaseLoader(url).load() for url in urls]\n",
    "docs_list = [item for sublist in docs for item in sublist]\n",
    "\n",
    "# Split\n",
    "text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(\n",
    "    chunk_size=500, chunk_overlap=0\n",
    ")\n",
    "doc_splits = text_splitter.split_documents(docs_list)\n",
    "\n",
    "# Add to vectorstore\n",
    "vectorstore = Chroma.from_documents(\n",
    "    documents=doc_splits,\n",
    "    collection_name=\"rag-chroma\",\n",
    "    embedding=embd,\n",
    ")\n",
    "retriever = vectorstore.as_retriever()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f52b427-750c-40f8-8893-e9caab3afd8d",
   "metadata": {},
   "source": [
    "## LLMs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4dec9d98-f3dc-4b7f-abc0-9d01c754f2be",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "datasource='web_search'\n",
      "datasource='vectorstore'\n"
     ]
    }
   ],
   "source": [
    "### Router\n",
    "\n",
    "from typing import Literal\n",
    "\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_core.pydantic_v1 import BaseModel, Field\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "\n",
    "# Data model\n",
    "class RouteQuery(BaseModel):\n",
    "    \"\"\"Route a user query to the most relevant datasource.\"\"\"\n",
    "\n",
    "    datasource: Literal[\"vectorstore\", \"web_search\"] = Field(\n",
    "        ...,\n",
    "        description=\"Given a user question choose to route it to web search or a vectorstore.\",\n",
    "    )\n",
    "\n",
    "\n",
    "# LLM with function call\n",
    "llm = ChatOpenAI(model=\"gpt-4o-mini\", temperature=0)\n",
    "structured_llm_router = llm.with_structured_output(RouteQuery)\n",
    "\n",
    "# Prompt\n",
    "system = \"\"\"You are an expert at routing a user question to a vectorstore or web search.\n",
    "The vectorstore contains documents related to agents, prompt engineering, and adversarial attacks.\n",
    "Use the vectorstore for questions on these topics. Otherwise, use web-search.\"\"\"\n",
    "route_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", system),\n",
    "        (\"human\", \"{question}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "question_router = route_prompt | structured_llm_router\n",
    "print(\n",
    "    question_router.invoke(\n",
    "        {\"question\": \"Who will the Bears draft first in the NFL draft?\"}\n",
    "    )\n",
    ")\n",
    "print(question_router.invoke({\"question\": \"What are the types of agent memory?\"}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "856801cb-f42a-44e7-956f-47845e3664ca",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "binary_score='no'\n"
     ]
    }
   ],
   "source": [
    "### Retrieval Grader\n",
    "\n",
    "\n",
    "# Data model\n",
    "class GradeDocuments(BaseModel):\n",
    "    \"\"\"Binary score for relevance check on retrieved documents.\"\"\"\n",
    "\n",
    "    binary_score: str = Field(\n",
    "        description=\"Documents are relevant to the question, 'yes' or 'no'\"\n",
    "    )\n",
    "\n",
    "\n",
    "# LLM with function call\n",
    "llm = ChatOpenAI(model=\"gpt-4o-mini\", temperature=0)\n",
    "structured_llm_grader = llm.with_structured_output(GradeDocuments)\n",
    "\n",
    "# Prompt\n",
    "system = \"\"\"You are a grader assessing relevance of a retrieved document to a user question. \\n \n",
    "    If the document contains keyword(s) or semantic meaning related to the user question, grade it as relevant. \\n\n",
    "    It does not need to be a stringent test. The goal is to filter out erroneous retrievals. \\n\n",
    "    Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question.\"\"\"\n",
    "grade_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", system),\n",
    "        (\"human\", \"Retrieved document: \\n\\n {document} \\n\\n User question: {question}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "retrieval_grader = grade_prompt | structured_llm_grader\n",
    "question = \"agent memory\"\n",
    "docs = retriever.get_relevant_documents(question)\n",
    "doc_txt = docs[1].page_content\n",
    "print(retrieval_grader.invoke({\"question\": question, \"document\": doc_txt}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2272333e-50b2-42ab-b472-e1055a3b94a8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The design of generative agents combines LLM with memory, planning, and reflection mechanisms to enable agents to behave based on past experience and interact with other agents. Memory stream is a long-term memory module that records agents' experiences in natural language. The retrieval model surfaces context to inform the agent's behavior based on relevance, recency, and importance.\n"
     ]
    }
   ],
   "source": [
    "### Generate\n",
    "\n",
    "from langchain import hub\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "\n",
    "# Prompt\n",
    "prompt = hub.pull(\"rlm/rag-prompt\")\n",
    "\n",
    "# LLM\n",
    "llm = ChatOpenAI(model_name=\"gpt-3.5-turbo\", temperature=0)\n",
    "\n",
    "\n",
    "# Post-processing\n",
    "def format_docs(docs):\n",
    "    return \"\\n\\n\".join(doc.page_content for doc in docs)\n",
    "\n",
    "\n",
    "# Chain\n",
    "rag_chain = prompt | llm | StrOutputParser()\n",
    "\n",
    "# Run\n",
    "generation = rag_chain.invoke({\"context\": docs, \"question\": question})\n",
    "print(generation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f0c08d14-77a0-4eed-b882-2d636abb22a3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "GradeHallucinations(binary_score='yes')"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "### Hallucination Grader\n",
    "\n",
    "\n",
    "# Data model\n",
    "class GradeHallucinations(BaseModel):\n",
    "    \"\"\"Binary score for hallucination present in generation answer.\"\"\"\n",
    "\n",
    "    binary_score: str = Field(\n",
    "        description=\"Answer is grounded in the facts, 'yes' or 'no'\"\n",
    "    )\n",
    "\n",
    "\n",
    "# LLM with function call\n",
    "llm = ChatOpenAI(model=\"gpt-4o-mini\", temperature=0)\n",
    "structured_llm_grader = llm.with_structured_output(GradeHallucinations)\n",
    "\n",
    "# Prompt\n",
    "system = \"\"\"You are a grader assessing whether an LLM generation is grounded in / supported by a set of retrieved facts. \\n \n",
    "     Give a binary score 'yes' or 'no'. 'Yes' means that the answer is grounded in / supported by the set of facts.\"\"\"\n",
    "hallucination_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", system),\n",
    "        (\"human\", \"Set of facts: \\n\\n {documents} \\n\\n LLM generation: {generation}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "hallucination_grader = hallucination_prompt | structured_llm_grader\n",
    "hallucination_grader.invoke({\"documents\": docs, \"generation\": generation})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ded99680-437a-4c9d-b860-619c88949d84",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "GradeAnswer(binary_score='yes')"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "### Answer Grader\n",
    "\n",
    "\n",
    "# Data model\n",
    "class GradeAnswer(BaseModel):\n",
    "    \"\"\"Binary score to assess answer addresses question.\"\"\"\n",
    "\n",
    "    binary_score: str = Field(\n",
    "        description=\"Answer addresses the question, 'yes' or 'no'\"\n",
    "    )\n",
    "\n",
    "\n",
    "# LLM with function call\n",
    "llm = ChatOpenAI(model=\"gpt-4o-mini\", temperature=0)\n",
    "structured_llm_grader = llm.with_structured_output(GradeAnswer)\n",
    "\n",
    "# Prompt\n",
    "system = \"\"\"You are a grader assessing whether an answer addresses / resolves a question \\n \n",
    "     Give a binary score 'yes' or 'no'. Yes' means that the answer resolves the question.\"\"\"\n",
    "answer_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", system),\n",
    "        (\"human\", \"User question: \\n\\n {question} \\n\\n LLM generation: {generation}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "answer_grader = answer_prompt | structured_llm_grader\n",
    "answer_grader.invoke({\"question\": question, \"generation\": generation})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "9d75f1d7-a47a-4577-bb0d-84b504b0867e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"What is the role of memory in an agent's functioning?\""
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "### Question Re-writer\n",
    "\n",
    "# LLM\n",
    "llm = ChatOpenAI(model=\"gpt-3.5-turbo-0125\", temperature=0)\n",
    "\n",
    "# Prompt\n",
    "system = \"\"\"You a question re-writer that converts an input question to a better version that is optimized \\n \n",
    "     for vectorstore retrieval. Look at the input and try to reason about the underlying semantic intent / meaning.\"\"\"\n",
    "re_write_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", system),\n",
    "        (\n",
    "            \"human\",\n",
    "            \"Here is the initial question: \\n\\n {question} \\n Formulate an improved question.\",\n",
    "        ),\n",
    "    ]\n",
    ")\n",
    "\n",
    "question_rewriter = re_write_prompt | llm | StrOutputParser()\n",
    "question_rewriter.invoke({\"question\": question})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d07c0b31-b919-4498-869f-9673125c2473",
   "metadata": {},
   "source": [
    "## Web Search Tool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "01d829bb-1074-4976-b650-ead41dcb9788",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Search\n",
    "\n",
    "from langchain_community.tools.tavily_search import TavilySearchResults\n",
    "\n",
    "web_search_tool = TavilySearchResults(k=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efbbff0e-8843-45bb-b2ff-137bef707ef4",
   "metadata": {},
   "source": [
    "## Construct the Graph \n",
    "\n",
    "Capture the flow in as a graph.\n",
    "\n",
    "### Define Graph State"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "e723fcdb-06e6-402d-912e-899795b78408",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "\n",
    "from typing_extensions import TypedDict\n",
    "\n",
    "\n",
    "class GraphState(TypedDict):\n",
    "    \"\"\"\n",
    "    Represents the state of our graph.\n",
    "\n",
    "    Attributes:\n",
    "        question: question\n",
    "        generation: LLM generation\n",
    "        documents: list of documents\n",
    "    \"\"\"\n",
    "\n",
    "    question: str\n",
    "    generation: str\n",
    "    documents: List[str]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e2d6c0d-42e8-4399-9751-e315be16607a",
   "metadata": {},
   "source": [
    "### Define Graph Flow "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "b76b5ec3-0720-443d-85b1-c0e79659ca0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.schema import Document\n",
    "\n",
    "\n",
    "def retrieve(state):\n",
    "    \"\"\"\n",
    "    Retrieve documents\n",
    "\n",
    "    Args:\n",
    "        state (dict): The current graph state\n",
    "\n",
    "    Returns:\n",
    "        state (dict): New key added to state, documents, that contains retrieved documents\n",
    "    \"\"\"\n",
    "    print(\"---RETRIEVE---\")\n",
    "    question = state[\"question\"]\n",
    "\n",
    "    # Retrieval\n",
    "    documents = retriever.invoke(question)\n",
    "    return {\"documents\": documents, \"question\": question}\n",
    "\n",
    "\n",
    "def generate(state):\n",
    "    \"\"\"\n",
    "    Generate answer\n",
    "\n",
    "    Args:\n",
    "        state (dict): The current graph state\n",
    "\n",
    "    Returns:\n",
    "        state (dict): New key added to state, generation, that contains LLM generation\n",
    "    \"\"\"\n",
    "    print(\"---GENERATE---\")\n",
    "    question = state[\"question\"]\n",
    "    documents = state[\"documents\"]\n",
    "\n",
    "    # RAG generation\n",
    "    generation = rag_chain.invoke({\"context\": documents, \"question\": question})\n",
    "    return {\"documents\": documents, \"question\": question, \"generation\": generation}\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 graph state\n",
    "\n",
    "    Returns:\n",
    "        state (dict): Updates documents key with only filtered relevant documents\n",
    "    \"\"\"\n",
    "\n",
    "    print(\"---CHECK DOCUMENT RELEVANCE TO QUESTION---\")\n",
    "    question = state[\"question\"]\n",
    "    documents = state[\"documents\"]\n",
    "\n",
    "    # Score each doc\n",
    "    filtered_docs = []\n",
    "    for d in documents:\n",
    "        score = retrieval_grader.invoke(\n",
    "            {\"question\": question, \"document\": d.page_content}\n",
    "        )\n",
    "        grade = score.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",
    "    return {\"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 graph state\n",
    "\n",
    "    Returns:\n",
    "        state (dict): Updates question key with a re-phrased question\n",
    "    \"\"\"\n",
    "\n",
    "    print(\"---TRANSFORM QUERY---\")\n",
    "    question = state[\"question\"]\n",
    "    documents = state[\"documents\"]\n",
    "\n",
    "    # Re-write question\n",
    "    better_question = question_rewriter.invoke({\"question\": question})\n",
    "    return {\"documents\": documents, \"question\": better_question}\n",
    "\n",
    "\n",
    "def web_search(state):\n",
    "    \"\"\"\n",
    "    Web search based on the re-phrased question.\n",
    "\n",
    "    Args:\n",
    "        state (dict): The current graph state\n",
    "\n",
    "    Returns:\n",
    "        state (dict): Updates documents key with appended web results\n",
    "    \"\"\"\n",
    "\n",
    "    print(\"---WEB SEARCH---\")\n",
    "    question = state[\"question\"]\n",
    "\n",
    "    # Web search\n",
    "    docs = web_search_tool.invoke({\"query\": question})\n",
    "    web_results = \"\\n\".join([d[\"content\"] for d in docs])\n",
    "    web_results = Document(page_content=web_results)\n",
    "\n",
    "    return {\"documents\": web_results, \"question\": question}\n",
    "\n",
    "\n",
    "### Edges ###\n",
    "\n",
    "\n",
    "def route_question(state):\n",
    "    \"\"\"\n",
    "    Route question to web search or RAG.\n",
    "\n",
    "    Args:\n",
    "        state (dict): The current graph state\n",
    "\n",
    "    Returns:\n",
    "        str: Next node to call\n",
    "    \"\"\"\n",
    "\n",
    "    print(\"---ROUTE QUESTION---\")\n",
    "    question = state[\"question\"]\n",
    "    source = question_router.invoke({\"question\": question})\n",
    "    if source.datasource == \"web_search\":\n",
    "        print(\"---ROUTE QUESTION TO WEB SEARCH---\")\n",
    "        return \"web_search\"\n",
    "    elif source.datasource == \"vectorstore\":\n",
    "        print(\"---ROUTE QUESTION TO RAG---\")\n",
    "        return \"vectorstore\"\n",
    "\n",
    "\n",
    "def decide_to_generate(state):\n",
    "    \"\"\"\n",
    "    Determines whether to generate an answer, or re-generate a question.\n",
    "\n",
    "    Args:\n",
    "        state (dict): The current graph state\n",
    "\n",
    "    Returns:\n",
    "        str: Binary decision for next node to call\n",
    "    \"\"\"\n",
    "\n",
    "    print(\"---ASSESS GRADED DOCUMENTS---\")\n",
    "    state[\"question\"]\n",
    "    filtered_documents = state[\"documents\"]\n",
    "\n",
    "    if not filtered_documents:\n",
    "        # All documents have been filtered check_relevance\n",
    "        # We will re-generate a new query\n",
    "        print(\n",
    "            \"---DECISION: ALL DOCUMENTS ARE NOT RELEVANT TO QUESTION, TRANSFORM QUERY---\"\n",
    "        )\n",
    "        return \"transform_query\"\n",
    "    else:\n",
    "        # We have relevant documents, so generate answer\n",
    "        print(\"---DECISION: GENERATE---\")\n",
    "        return \"generate\"\n",
    "\n",
    "\n",
    "def grade_generation_v_documents_and_question(state):\n",
    "    \"\"\"\n",
    "    Determines whether the generation is grounded in the document and answers question.\n",
    "\n",
    "    Args:\n",
    "        state (dict): The current graph state\n",
    "\n",
    "    Returns:\n",
    "        str: Decision for next node to call\n",
    "    \"\"\"\n",
    "\n",
    "    print(\"---CHECK HALLUCINATIONS---\")\n",
    "    question = state[\"question\"]\n",
    "    documents = state[\"documents\"]\n",
    "    generation = state[\"generation\"]\n",
    "\n",
    "    score = hallucination_grader.invoke(\n",
    "        {\"documents\": documents, \"generation\": generation}\n",
    "    )\n",
    "    grade = score.binary_score\n",
    "\n",
    "    # Check hallucination\n",
    "    if grade == \"yes\":\n",
    "        print(\"---DECISION: GENERATION IS GROUNDED IN DOCUMENTS---\")\n",
    "        # Check question-answering\n",
    "        print(\"---GRADE GENERATION vs QUESTION---\")\n",
    "        score = answer_grader.invoke({\"question\": question, \"generation\": generation})\n",
    "        grade = score.binary_score\n",
    "        if grade == \"yes\":\n",
    "            print(\"---DECISION: GENERATION ADDRESSES QUESTION---\")\n",
    "            return \"useful\"\n",
    "        else:\n",
    "            print(\"---DECISION: GENERATION DOES NOT ADDRESS QUESTION---\")\n",
    "            return \"not useful\"\n",
    "    else:\n",
    "        pprint(\"---DECISION: GENERATION IS NOT GROUNDED IN DOCUMENTS, RE-TRY---\")\n",
    "        return \"not supported\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ab01f36-5628-49ab-bfd3-84bb6f1a1b0f",
   "metadata": {},
   "source": [
    "### Compile Graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "67854e07-9293-4c3c-bf9a-bc9a605570ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.graph import END, StateGraph, START\n",
    "\n",
    "workflow = StateGraph(GraphState)\n",
    "\n",
    "# Define the nodes\n",
    "workflow.add_node(\"web_search\", web_search)  # web search\n",
    "workflow.add_node(\"retrieve\", retrieve)  # retrieve\n",
    "workflow.add_node(\"grade_documents\", grade_documents)  # grade documents\n",
    "workflow.add_node(\"generate\", generate)  # generate\n",
    "workflow.add_node(\"transform_query\", transform_query)  # transform_query\n",
    "\n",
    "# Build graph\n",
    "workflow.add_conditional_edges(\n",
    "    START,\n",
    "    route_question,\n",
    "    {\n",
    "        \"web_search\": \"web_search\",\n",
    "        \"vectorstore\": \"retrieve\",\n",
    "    },\n",
    ")\n",
    "workflow.add_edge(\"web_search\", \"generate\")\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_and_question,\n",
    "    {\n",
    "        \"not supported\": \"generate\",\n",
    "        \"useful\": END,\n",
    "        \"not useful\": \"transform_query\",\n",
    "    },\n",
    ")\n",
    "\n",
    "# Compile\n",
    "app = workflow.compile()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85bce541",
   "metadata": {},
   "source": [
    "## Use Graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "29acc541-d726-4b75-84d1-a215845fe88a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---ROUTE QUESTION---\n",
      "---ROUTE QUESTION TO WEB SEARCH---\n",
      "---WEB SEARCH---\n",
      "\"Node 'web_search':\"\n",
      "'\\n---\\n'\n",
      "---GENERATE---\n",
      "---CHECK HALLUCINATIONS---\n",
      "---DECISION: GENERATION IS GROUNDED IN DOCUMENTS---\n",
      "---GRADE GENERATION vs QUESTION---\n",
      "---DECISION: GENERATION ADDRESSES QUESTION---\n",
      "\"Node 'generate':\"\n",
      "'\\n---\\n'\n",
      "('It is expected that the Chicago Bears could have the opportunity to draft '\n",
      " 'the first defensive player in the 2024 NFL draft. The Bears have the first '\n",
      " 'overall pick in the draft, giving them a prime position to select top '\n",
      " 'talent. The top wide receiver Marvin Harrison Jr. from Ohio State is also '\n",
      " 'mentioned as a potential pick for the Cardinals.')\n"
     ]
    }
   ],
   "source": [
    "from pprint import pprint\n",
    "\n",
    "# Run\n",
    "inputs = {\n",
    "    \"question\": \"What player at the Bears expected to draft first in the 2024 NFL draft?\"\n",
    "}\n",
    "for output in app.stream(inputs):\n",
    "    for key, value in output.items():\n",
    "        # Node\n",
    "        pprint(f\"Node '{key}':\")\n",
    "        # Optional: print full state at each node\n",
    "        # pprint.pprint(value[\"keys\"], indent=2, width=80, depth=None)\n",
    "    pprint(\"\\n---\\n\")\n",
    "\n",
    "# Final generation\n",
    "pprint(value[\"generation\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11fddd00-58bf-4910-bf36-be9e5bfba778",
   "metadata": {},
   "source": [
    "Trace: \n",
    "\n",
    "https://smith.langchain.com/public/7e3aa7e5-c51f-45c2-bc66-b34f17ff2263/r"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "69a985dd-03c6-45af-a67b-b15746a2cb5f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---ROUTE QUESTION---\n",
      "---ROUTE QUESTION TO RAG---\n",
      "---RETRIEVE---\n",
      "\"Node 'retrieve':\"\n",
      "'\\n---\\n'\n",
      "---CHECK DOCUMENT RELEVANCE TO QUESTION---\n",
      "---GRADE: DOCUMENT RELEVANT---\n",
      "---GRADE: DOCUMENT RELEVANT---\n",
      "---GRADE: DOCUMENT NOT RELEVANT---\n",
      "---GRADE: DOCUMENT RELEVANT---\n",
      "---ASSESS GRADED DOCUMENTS---\n",
      "---DECISION: GENERATE---\n",
      "\"Node 'grade_documents':\"\n",
      "'\\n---\\n'\n",
      "---GENERATE---\n",
      "---CHECK HALLUCINATIONS---\n",
      "---DECISION: GENERATION IS GROUNDED IN DOCUMENTS---\n",
      "---GRADE GENERATION vs QUESTION---\n",
      "---DECISION: GENERATION ADDRESSES QUESTION---\n",
      "\"Node 'generate':\"\n",
      "'\\n---\\n'\n",
      "('The types of agent memory include Sensory Memory, Short-Term Memory (STM) or '\n",
      " 'Working Memory, and Long-Term Memory (LTM) with subtypes of Explicit / '\n",
      " 'declarative memory and Implicit / procedural memory. Sensory memory retains '\n",
      " 'sensory information briefly, STM stores information for cognitive tasks, and '\n",
      " 'LTM stores information for a long time with different types of memories.')\n"
     ]
    }
   ],
   "source": [
    "# Run\n",
    "inputs = {\"question\": \"What are the types of agent memory?\"}\n",
    "for output in app.stream(inputs):\n",
    "    for key, value in output.items():\n",
    "        # Node\n",
    "        pprint(f\"Node '{key}':\")\n",
    "        # Optional: print full state at each node\n",
    "        # pprint.pprint(value[\"keys\"], indent=2, width=80, depth=None)\n",
    "    pprint(\"\\n---\\n\")\n",
    "\n",
    "# Final generation\n",
    "pprint(value[\"generation\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ebf41097-fc4c-4072-95b3-e7e07731ada1",
   "metadata": {},
   "source": [
    "Trace: \n",
    "\n",
    "https://smith.langchain.com/public/fdf0a180-6d15-4d09-bb92-f84f2105ca51/r"
   ]
  }
 ],
 "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.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
