{
 "cells": [
  {
   "attachments": {
    "8bd22c04-50d3-441f-8cbf-f4f225b8afda.png": {
     "image/png": "[encoded data, 256468 characters]"
    }
   },
   "cell_type": "markdown",
   "id": "02ea5aaf-fcff-4fff-aa18-c637a4bfe083",
   "metadata": {},
   "source": [
    "# Building and Testing Local Corrective RAG\n",
    "\n",
    "[Corrective-RAG (CRAG)](https://arxiv.org/abs/2401.15884) is a strategy for RAG that incorporates self-reflection / self-grading on retrieved documents. \n",
    "\n",
    "We will implement this agent as a custom / user-defined control flow in LangGraph. \n",
    "\n",
    "Then, we will show how to test these agents. \n",
    "\n",
    "(1) We'll use an [evaluation set of question-answer pairs for RAG in LangSmith](https://docs.smith.langchain.com/tutorials/Developers/agents#eval). \n",
    "\n",
    "(2) We'll evaluate end-to-end performance of our agents along with the specific reasoning trace of each one.\n",
    "\n",
    "This notebook will reference ideas from [these slides](https://docs.google.com/presentation/d/1QWkXi4DYjfw94eHcy9RMLqpQdJtS2C_kx_u7wAUvlZE/edit?usp=sharing).\n",
    "\n",
    "\n",
    "\n",
    "![Screenshot 2024-07-23 at 11.38.32 AM.png](attachment:8bd22c04-50d3-441f-8cbf-f4f225b8afda.png)\n",
    "\n",
    "## Environment\n",
    "\n",
    "We'll use [llama3.1 with Ollama](https://ai.meta.com/blog/meta-llama-3-1/) as our LLM, but this can be easily swapped. \n",
    "\n",
    "We'll use [Tavily](https://python.langchain.com/v0.2/docs/integrations/tools/tavily_search/) for web search.\n",
    "\n",
    "We'll use a vectorstore with [OpenAI embeddings](https://python.langchain.com/v0.2/docs/integrations/text_embedding/openai/#embed-documents), but this can be swapped for Nomic local embeddings.\n",
    "\n",
    "We'll use [LangSmith](https://docs.smith.langchain.com/) for tracing and evaluation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df96d419-e12a-4640-8131-8460616d45f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture --no-stderr\n",
    "%pip install -U langchain tavily-python langgraph matplotlib langchain_community langchain-openai scikit-learn langchainhub langchain-ollama nomic[local]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26b0bea7-a99c-4bf8-a375-08d2d610f7bf",
   "metadata": {},
   "source": [
    "## Custom LangGraph Agent\n",
    "\n",
    "Here we'll explicitly define the majority of the control flow, only using an LLM to define a single branch point following grading."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d5925281-6762-43f0-8efb-8e05d4391b8d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "USER_AGENT environment variable not set, consider setting it to identify your requests.\n"
     ]
    }
   ],
   "source": [
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "from langchain_community.document_loaders import WebBaseLoader\n",
    "from langchain_community.vectorstores import SKLearnVectorStore\n",
    "from langchain_nomic.embeddings import NomicEmbeddings\n",
    "from langchain_openai import OpenAIEmbeddings\n",
    "from langchain_core.tools import tool\n",
    "\n",
    "# List of URLs to load documents from\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 documents from the URLs\n",
    "docs = [WebBaseLoader(url).load() for url in urls]\n",
    "docs_list = [item for sublist in docs for item in sublist]\n",
    "\n",
    "# Initialize a text splitter with specified chunk size and overlap\n",
    "text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(\n",
    "    chunk_size=250, chunk_overlap=0\n",
    ")\n",
    "\n",
    "# Split the documents into chunks\n",
    "doc_splits = text_splitter.split_documents(docs_list)\n",
    "\n",
    "# Add the document chunks to the \"vector store\" using NomicEmbeddings\n",
    "vectorstore = SKLearnVectorStore.from_documents(\n",
    "    documents=doc_splits,\n",
    "    # embedding=NomicEmbeddings(model=\"nomic-embed-text-v1.5\", inference_mode=\"local\"),\n",
    "    embedding=OpenAIEmbeddings(),\n",
    ")\n",
    "retriever = vectorstore.as_retriever(k=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "cb3dd4c4-8319-436a-b612-9e2e99a06ee4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.schema import Document\n",
    "from langchain_community.tools.tavily_search import TavilySearchResults\n",
    "\n",
    "web_search_tool = TavilySearchResults()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "93db9b3d-e93d-4837-baf5-dc1eb3126ede",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_ollama import ChatOllama\n",
    "from langchain.prompts import PromptTemplate\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "\n",
    "prompt = PromptTemplate(\n",
    "    template=\"\"\"You are an assistant for question-answering tasks. \n",
    "    \n",
    "    Use the following documents to answer the question. \n",
    "    \n",
    "    If you don't know the answer, just say that you don't know. \n",
    "    \n",
    "    Use three sentences maximum and keep the answer concise:\n",
    "    Question: {question} \n",
    "    Documents: {documents} \n",
    "    Answer: \n",
    "    \"\"\",\n",
    "    input_variables=[\"question\", \"documents\"],\n",
    ")\n",
    "\n",
    "llm = ChatOllama(\n",
    "    model=\"llama3.1\",\n",
    "    temperature=0,\n",
    ")\n",
    "\n",
    "rag_chain = prompt | llm | StrOutputParser()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "44a0f1e8-be63-4e64-98b8-582c390f263b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_core.pydantic_v1 import BaseModel, Field\n",
    "from langchain_core.output_parsers import JsonOutputParser\n",
    "\n",
    "# JSON\n",
    "llm = ChatOllama(model=\"llama3.1\", \n",
    "                 format=\"json\", \n",
    "                 temperature=0)\n",
    "\n",
    "\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 {document} \\n\\n\n",
    "    Here is the user question: {question} \\n\n",
    "    If the document contains keywords 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\n",
    "    Provide the binary score as a JSON with a single key 'score' and no premable or explanation.\"\"\",\n",
    "    input_variables=[\"question\", \"document\"],\n",
    ")\n",
    "\n",
    "retrieval_grader = prompt | llm | JsonOutputParser()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "fc6dbd23-ad44-4bf1-9216-86fa10420e5a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 19676 characters]",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from typing_extensions import TypedDict, List\n",
    "from IPython.display import Image, display\n",
    "from langgraph.graph import START, END, StateGraph\n",
    "\n",
    "class GraphState(TypedDict):\n",
    "    \"\"\"\n",
    "    Represents the state of our graph.\n",
    "\n",
    "    Attributes:\n",
    "        question: question\n",
    "        generation: LLM generation\n",
    "        search: whether to add search\n",
    "        documents: list of documents\n",
    "    \"\"\"\n",
    "\n",
    "    question: str\n",
    "    generation: str\n",
    "    search: str\n",
    "    documents: List[str]\n",
    "    steps: List[str]\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",
    "    question = state[\"question\"]\n",
    "    documents = retriever.invoke(question)\n",
    "    steps = state[\"steps\"]\n",
    "    steps.append(\"retrieve_documents\")\n",
    "    return {\"documents\": documents, \"question\": question, \"steps\": steps}\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",
    "\n",
    "    question = state[\"question\"]\n",
    "    documents = state[\"documents\"]\n",
    "    generation = rag_chain.invoke({\"documents\": documents, \"question\": question})\n",
    "    steps = state[\"steps\"]\n",
    "    steps.append(\"generate_answer\")\n",
    "    return {\n",
    "        \"documents\": documents,\n",
    "        \"question\": question,\n",
    "        \"generation\": generation,\n",
    "        \"steps\": steps,\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 graph state\n",
    "\n",
    "    Returns:\n",
    "        state (dict): Updates documents key with only filtered relevant documents\n",
    "    \"\"\"\n",
    "\n",
    "    question = state[\"question\"]\n",
    "    documents = state[\"documents\"]\n",
    "    steps = state[\"steps\"]\n",
    "    steps.append(\"grade_document_retrieval\")\n",
    "    filtered_docs = []\n",
    "    search = \"No\"\n",
    "    for d in documents:\n",
    "        score = retrieval_grader.invoke(\n",
    "            {\"question\": question, \"document\": d.page_content}\n",
    "        )\n",
    "        grade = score[\"score\"]\n",
    "        if grade == \"yes\":\n",
    "            filtered_docs.append(d)\n",
    "        else:\n",
    "            search = \"Yes\"\n",
    "            continue\n",
    "    return {\n",
    "        \"documents\": filtered_docs,\n",
    "        \"question\": question,\n",
    "        \"search\": search,\n",
    "        \"steps\": steps,\n",
    "    }\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",
    "    question = state[\"question\"]\n",
    "    documents = state.get(\"documents\", [])\n",
    "    steps = state[\"steps\"]\n",
    "    steps.append(\"web_search\")\n",
    "    web_results = web_search_tool.invoke({\"query\": question})\n",
    "    documents.extend(\n",
    "        [\n",
    "            Document(page_content=d[\"content\"], metadata={\"url\": d[\"url\"]})\n",
    "            for d in web_results\n",
    "        ]\n",
    "    )\n",
    "    return {\"documents\": documents, \"question\": question, \"steps\": steps}\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",
    "    search = state[\"search\"]\n",
    "    if search == \"Yes\":\n",
    "        return \"search\"\n",
    "    else:\n",
    "        return \"generate\"\n",
    "\n",
    "\n",
    "# Graph\n",
    "workflow = StateGraph(GraphState)\n",
    "\n",
    "# Define the nodes\n",
    "workflow.add_node(\"retrieve\", retrieve)  # retrieve\n",
    "workflow.add_node(\"grade_documents\", grade_documents)  # grade documents\n",
    "workflow.add_node(\"generate\", generate)  # generatae\n",
    "workflow.add_node(\"web_search\", web_search)  # web search\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",
    "        \"search\": \"web_search\",\n",
    "        \"generate\": \"generate\",\n",
    "    },\n",
    ")\n",
    "workflow.add_edge(\"web_search\", \"generate\")\n",
    "workflow.add_edge(\"generate\", END)\n",
    "\n",
    "custom_graph = workflow.compile()\n",
    "\n",
    "display(Image(custom_graph.get_graph(xray=True).draw_mermaid_png()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3689edb1-a852-40a8-8a60-0fb2de856ba3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import uuid \n",
    "\n",
    "def predict_custom_agent_answer(example: dict):\n",
    "    \n",
    "    config = {\"configurable\": {\"thread_id\": str(uuid.uuid4())}}\n",
    "    \n",
    "    state_dict = custom_graph.invoke(\n",
    "        {\"question\": example[\"input\"], \"steps\": []}, config\n",
    "    )\n",
    "    \n",
    "    return {\"response\": state_dict[\"generation\"], \"steps\": state_dict[\"steps\"]}\n",
    "\n",
    "example = {\"input\": \"What are the types of agent memory?\"}\n",
    "#response = predict_custom_agent_answer(example)\n",
    "#response"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a7a8fe4-27c4-48a8-a2da-dc369c1cc09c",
   "metadata": {},
   "source": [
    "## Evaluation\n",
    "\n",
    "Now we've defined two different agent architectures that do roughly the same thing!\n",
    "\n",
    "We can evaluate them. See our [conceptual guide](https://docs.smith.langchain.com/concepts/evaluation#agents) for context on agent evaluation.\n",
    "\n",
    "### Response\n",
    "\n",
    "First, we can assess how well [our agent performs on a set of question-answer pairs](https://docs.smith.langchain.com/tutorials/Developers/agents#response-evaluation).\n",
    "\n",
    "We'll use a dataset, `\"Corrective RAG Agent Testing`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5e43252-e5e8-4929-ad7d-3eff3052c1cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langsmith import Client\n",
    "\n",
    "client = Client()\n",
    "\n",
    "# Create a dataset\n",
    "examples = [\n",
    "    (\n",
    "        \"How does the ReAct agent use self-reflection? \",\n",
    "        \"ReAct integrates reasoning and acting, performing actions - such tools like Wikipedia search API - and then observing / reasoning about the tool outputs.\",\n",
    "    ),\n",
    "    (\n",
    "        \"What are the types of biases that can arise with few-shot prompting?\",\n",
    "        \"The biases that can arise with few-shot prompting include (1) Majority label bias, (2) Recency bias, and (3) Common token bias.\",\n",
    "    ),\n",
    "    (\n",
    "        \"What are five types of adversarial attacks?\",\n",
    "        \"Five types of adversarial attacks are (1) Token manipulation, (2) Gradient based attack, (3) Jailbreak prompting, (4) Human red-teaming, (5) Model red-teaming.\",\n",
    "    ),\n",
    "    (\n",
    "        \"Who did the Chicago Bears draft first in the 2024 NFL draft”?\",\n",
    "        \"The Chicago Bears drafted Caleb Williams first in the 2024 NFL draft.\",\n",
    "    ),\n",
    "    (\"Who won the 2024 NBA finals?\", \"The Boston Celtics on the 2024 NBA finals\"),\n",
    "]\n",
    "\n",
    "# Save it\n",
    "dataset_name = \"Corrective RAG Agent Testing\"\n",
    "if not client.has_dataset(dataset_name=dataset_name):\n",
    "    dataset = client.create_dataset(dataset_name=dataset_name)\n",
    "    inputs, outputs = zip(\n",
    "        *[({\"input\": text}, {\"output\": label}) for text, label in examples]\n",
    "    )\n",
    "    client.create_examples(inputs=inputs, outputs=outputs, dataset_id=dataset.id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5991e863-7500-4ede-b76a-e7e59c44794b",
   "metadata": {},
   "source": [
    "Now, we'll use an `LLM as a grader` to compare both agent responses to our ground truth reference answer.\n",
    "\n",
    "[Here](https://smith.langchain.com/hub/rlm/rag-answer-vs-reference) is the default prompt that we can use.\n",
    "\n",
    "We'll use `gpt-4o` as our LLM grader."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "1263f1a0-078c-44c6-98e2-b3efd0371072",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain import hub\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "# Grade prompt\n",
    "grade_prompt_answer_accuracy = hub.pull(\"langchain-ai/rag-answer-vs-reference\")\n",
    "\n",
    "\n",
    "def answer_evaluator(run, example) -> dict:\n",
    "    \"\"\"\n",
    "    A simple evaluator for RAG answer accuracy\n",
    "    \"\"\"\n",
    "\n",
    "    # Get the question, the ground truth reference answer, RAG chain answer prediction\n",
    "    input_question = example.inputs[\"input\"]\n",
    "    reference = example.outputs[\"output\"]\n",
    "    prediction = run.outputs[\"response\"]\n",
    "\n",
    "    # Define an LLM grader\n",
    "    llm = ChatOpenAI(model=\"gpt-4o\", temperature=0)\n",
    "    answer_grader = grade_prompt_answer_accuracy | llm\n",
    "\n",
    "    # Run evaluator\n",
    "    score = answer_grader.invoke(\n",
    "        {\n",
    "            \"question\": input_question,\n",
    "            \"correct_answer\": reference,\n",
    "            \"student_answer\": prediction,\n",
    "        }\n",
    "    )\n",
    "    score = score[\"Score\"]\n",
    "    return {\"key\": \"answer_v_reference_score\", \"score\": score}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08108a87-5e52-4c6b-a376-4fe81c594603",
   "metadata": {},
   "source": [
    "### Trajectory\n",
    "\n",
    "Second, [we can assess the list of tool calls](https://docs.smith.langchain.com/tutorials/Developers/agents#trajectory) that each agent makes relative to expected trajectories.\n",
    "\n",
    "This evaluates the specific reasoning traces taken by our agents!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d6cbc8c3-c312-4ff4-876f-59a8928f18d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langsmith.schemas import Example, Run\n",
    "\n",
    "# Reasoning traces that we expect the agents to take\n",
    "expected_trajectory_1 = [\n",
    "    \"retrieve_documents\",\n",
    "    \"grade_document_retrieval\",\n",
    "    \"web_search\",\n",
    "    \"generate_answer\",\n",
    "]\n",
    "expected_trajectory_2 = [\n",
    "    \"retrieve_documents\",\n",
    "    \"grade_document_retrieval\",\n",
    "    \"generate_answer\",\n",
    "]\n",
    "\n",
    "def check_trajectory_custom(root_run: Run, example: Example) -> dict:\n",
    "    \"\"\"\n",
    "    Check if all expected tools are called in exact order and without any additional tool calls.\n",
    "    \"\"\"\n",
    "    tool_calls = root_run.outputs[\"steps\"]\n",
    "    print(f\"Tool calls custom agent: {tool_calls}\")\n",
    "    if tool_calls == expected_trajectory_1 or tool_calls == expected_trajectory_2:\n",
    "        score = 1\n",
    "    else:\n",
    "        score = 0\n",
    "\n",
    "    return {\"score\": int(score), \"key\": \"tool_calls_in_exact_order\"}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17737eb1-d25d-4c52-a2f5-34f90a68d478",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langsmith.evaluation import evaluate\n",
    "\n",
    "dataset_name = \"Corrective RAG Agent Testing\"\n",
    "model_tested = \"llama3.1\"\n",
    "metadata = \"CRAG, llama3.1\"\n",
    "experiment_prefix = f\"custom-agent-{model_tested}\"\n",
    "experiment_results = evaluate(\n",
    "    predict_custom_agent_answer,\n",
    "    data=dataset_name,\n",
    "    evaluators=[answer_evaluator, check_trajectory_custom],\n",
    "    experiment_prefix=experiment_prefix + \"-answer-and-tool-use\",\n",
    "    num_repetitions=3,\n",
    "    max_concurrency=1,\n",
    "    metadata={\"version\": metadata},\n",
    ")"
   ]
  },
  {
   "attachments": {
    "5e0fb308-080e-4621-ba20-55a655bb981e.png": {
     "image/png": "[encoded data, 154056 characters]"
    }
   },
   "cell_type": "markdown",
   "id": "4be2b4ec-7489-4475-ac91-ce992e9503a3",
   "metadata": {},
   "source": [
    "Dataset - \n",
    "\n",
    "https://smith.langchain.com/public/6c6c4543-1e56-4126-ab18-a84a31ff4946/d\n",
    "\n",
    "Results -\n",
    "\n",
    "* Latency for llama3.1 appears higher than llama3.\n",
    "* Compares well to other, larger llms and beat llama3.\n",
    "\n",
    "![Screenshot 2024-07-23 at 12.39.27 PM.png](attachment:5e0fb308-080e-4621-ba20-55a655bb981e.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8eb46884-08aa-4ab4-8bb9-f72277c3b35b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}
