{
 "cells": [
  {
   "attachments": {
    "683fae34-980f-43f0-a9c2-9894bebd9157.png": {
     "image/png": "[encoded data, 128984 characters]"
    }
   },
   "cell_type": "markdown",
   "id": "8889a307-fa3f-4d38-9127-d41e4686ae47",
   "metadata": {},
   "source": [
    "# Corrective RAG (CRAG)\n",
    "\n",
    "Corrective-RAG (CRAG) is a strategy for RAG that incorporates self-reflection / self-grading on retrieved documents. \n",
    "\n",
    "In the paper [here](https://arxiv.org/pdf/2401.15884.pdf), a few steps are taken:\n",
    "\n",
    "* If at least one document exceeds the threshold for relevance, then it proceeds to generation\n",
    "* Before generation, it performs knowledge refinement\n",
    "* This partitions the document into \"knowledge strips\"\n",
    "* It grades each strip, and filters our irrelevant ones\n",
    "* If all documents fall below the relevance threshold or if the grader is unsure, then the framework seeks an additional datasource\n",
    "* It will use web search to supplement retrieval\n",
    " \n",
    "We will implement some of these ideas from scratch using [LangGraph](https://langchain-ai.github.io/langgraph/):\n",
    "\n",
    "* Let's skip the knowledge refinement phase as a first pass. This can be added back as a node, if desired. \n",
    "* If *any* documents are irrelevant, let's opt to supplement retrieval with web search. \n",
    "* We'll use [Tavily Search](https://python.langchain.com/v0.2/docs/integrations/tools/tavily_search/) for web search.\n",
    "* Let's use query re-writing to optimize the query for web search.\n",
    "\n",
    "![Screenshot 2024-04-01 at 9.28.30 AM.png](attachment:683fae34-980f-43f0-a9c2-9894bebd9157.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4931ac25-99f9-4f04-b3d1-4683f7853667",
   "metadata": {},
   "source": [
    "# Environment "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "568c84d6-9df6-4b7b-b50d-476c0a64a04b",
   "metadata": {},
   "outputs": [],
   "source": ["! pip install langchain_community tiktoken langchain-openai langchainhub chromadb langchain langgraph tavily-python"]
  },
  {
   "cell_type": "markdown",
   "id": "c9ca92ae-c84f-423e-bab5-2a8bd06bd4cf",
   "metadata": {},
   "source": [
    "### LLMs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74710419-158d-4270-931c-de83db7b580d",
   "metadata": {},
   "outputs": [],
   "source": ["import os\n\nos.environ[\"OPENAI_API_KEY\"] = \"<your-api-key>\""]
  },
  {
   "cell_type": "markdown",
   "id": "db34bf9b-fa1e-45d3-a83e-4f5eaf897bad",
   "metadata": {},
   "source": [
    "### Search\n",
    " \n",
    "We'll use [Tavily Search](https://python.langchain.com/docs/integrations/tools/tavily_search) for web search."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3ac6e65-2d4e-48dd-9fff-40047373332d",
   "metadata": {},
   "outputs": [],
   "source": ["os.environ[\"TAVILY_API_KEY\"] = \"<your-api-key>\""]
  },
  {
   "cell_type": "markdown",
   "id": "20d2644a-043e-419c-9ff8-45118b8e9f17",
   "metadata": {},
   "source": [
    "### Tracing\n",
    "\n",
    "Optionally, use [LangSmith](https://docs.smith.langchain.com/) for tracing (shown at bottom) by setting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e205f57e-5218-478b-ad8e-1723bdb0d45e",
   "metadata": {},
   "outputs": [],
   "source": ["os.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\nos.environ[\"LANGCHAIN_ENDPOINT\"] = \"https://api.smith.langchain.com\"\nos.environ[\"LANGCHAIN_API_KEY\"] = \"<your-api-key>\""]
  },
  {
   "cell_type": "markdown",
   "id": "a21f32d2-92ce-4995-b309-99347bafe3be",
   "metadata": {},
   "source": [
    "## Index\n",
    " \n",
    "Let's index 3 blog posts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3a566a30-cf0e-4330-ad4d-9bf994bdfa86",
   "metadata": {},
   "outputs": [],
   "source": ["from langchain.text_splitter import RecursiveCharacterTextSplitter\nfrom langchain_community.document_loaders import WebBaseLoader\nfrom langchain_community.vectorstores import Chroma\nfrom langchain_openai import OpenAIEmbeddings\n\nurls = [\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\ndocs = [WebBaseLoader(url).load() for url in urls]\ndocs_list = [item for sublist in docs for item in sublist]\n\ntext_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(\n    chunk_size=250, chunk_overlap=0\n)\ndoc_splits = text_splitter.split_documents(docs_list)\n\n# Add to vectorDB\nvectorstore = Chroma.from_documents(\n    documents=doc_splits,\n    collection_name=\"rag-chroma\",\n    embedding=OpenAIEmbeddings(),\n)\nretriever = vectorstore.as_retriever()"]
  },
  {
   "cell_type": "markdown",
   "id": "6fca2db8-8d68-42b0-981d-4be5ccdbe293",
   "metadata": {},
   "source": [
    "## LLMs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7ece414c-2df5-4ffd-aa82-550a65775261",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "binary_score='yes'\n"
     ]
    }
   ],
   "source": ["### Retrieval Grader\n\nfrom langchain_core.prompts import ChatPromptTemplate\nfrom langchain_core.pydantic_v1 import BaseModel, Field\nfrom langchain_openai import ChatOpenAI\n\n\n# Data model\nclass 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\nllm = ChatOpenAI(model=\"gpt-3.5-turbo-0125\", temperature=0)\nstructured_llm_grader = llm.with_structured_output(GradeDocuments)\n\n# Prompt\nsystem = \"\"\"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 question, grade it as relevant. \\n\n    Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question.\"\"\"\ngrade_prompt = ChatPromptTemplate.from_messages(\n    [\n        (\"system\", system),\n        (\"human\", \"Retrieved document: \\n\\n {document} \\n\\n User question: {question}\"),\n    ]\n)\n\nretrieval_grader = grade_prompt | structured_llm_grader\nquestion = \"agent memory\"\ndocs = retriever.get_relevant_documents(question)\ndoc_txt = docs[1].page_content\nprint(retrieval_grader.invoke({\"question\": question, \"document\": doc_txt}))"]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a207c85f-e414-46b7-8999-4c0ead1493da",
   "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 conditioned on past experience. Memory stream is a long-term memory module that records a comprehensive list of agents' experience in natural language. Short-term memory is utilized for in-context learning, while long-term memory allows agents to retain and recall information over extended periods.\n"
     ]
    }
   ],
   "source": ["### Generate\n\nfrom langchain import hub\nfrom langchain_core.output_parsers import StrOutputParser\n\n# Prompt\nprompt = hub.pull(\"rlm/rag-prompt\")\n\n# LLM\nllm = ChatOpenAI(model_name=\"gpt-3.5-turbo\", temperature=0)\n\n\n# Post-processing\ndef format_docs(docs):\n    return \"\\n\\n\".join(doc.page_content for doc in docs)\n\n\n# Chain\nrag_chain = prompt | llm | StrOutputParser()\n\n# Run\ngeneration = rag_chain.invoke({\"context\": docs, \"question\": question})\nprint(generation)"]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "30d0a69b-9087-4f85-af26-cab55b567872",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'What is the role of memory in artificial intelligence agents?'"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": ["### Question Re-writer\n\n# LLM\nllm = ChatOpenAI(model=\"gpt-3.5-turbo-0125\", temperature=0)\n\n# Prompt\nsystem = \"\"\"You a question re-writer that converts an input question to a better version that is optimized \\n \n     for web search. Look at the input and try to reason about the underlying semantic intent / meaning.\"\"\"\nre_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\nquestion_rewriter = re_write_prompt | llm | StrOutputParser()\nquestion_rewriter.invoke({\"question\": question})"]
  },
  {
   "cell_type": "markdown",
   "id": "e4538467-4a15-4733-b93c-2b79d4d6bf25",
   "metadata": {},
   "source": [
    "## Web Search Tool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "46d51b53-54a9-4e0a-9f14-e39998f5b340",
   "metadata": {},
   "outputs": [],
   "source": ["### Search\n\nfrom langchain_community.tools.tavily_search import TavilySearchResults\n\nweb_search_tool = TavilySearchResults(k=3)"]
  },
  {
   "cell_type": "markdown",
   "id": "87194a1b-535a-4593-ab95-5736fae176d1",
   "metadata": {},
   "source": [
    "# Graph \n",
    "\n",
    "Capture the flow in as a graph.\n",
    "\n",
    "## Graph state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "94b3945f-ef0f-458d-a443-f763903550b0",
   "metadata": {},
   "outputs": [],
   "source": ["from typing import List\n\nfrom typing_extensions import TypedDict\n\n\nclass GraphState(TypedDict):\n    \"\"\"\n    Represents the state of our graph.\n\n    Attributes:\n        question: question\n        generation: LLM generation\n        web_search: whether to add search\n        documents: list of documents\n    \"\"\"\n\n    question: str\n    generation: str\n    web_search: str\n    documents: List[str]"]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "efd639c5-82e2-45e6-a94a-6a4039646ef5",
   "metadata": {},
   "outputs": [],
   "source": ["from langchain.schema import Document\n\n\ndef 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.get_relevant_documents(question)\n    return {\"documents\": documents, \"question\": question}\n\n\ndef 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\ndef 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    web_search = \"No\"\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            web_search = \"Yes\"\n            continue\n    return {\"documents\": filtered_docs, \"question\": question, \"web_search\": web_search}\n\n\ndef 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\ndef 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    documents = state[\"documents\"]\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    documents.append(web_results)\n\n    return {\"documents\": documents, \"question\": question}\n\n\n### Edges\n\n\ndef 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    web_search = state[\"web_search\"]\n    state[\"documents\"]\n\n    if web_search == \"Yes\":\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\""]
  },
  {
   "cell_type": "markdown",
   "id": "fa076e90-7132-4fcf-8507-db5990314c4f",
   "metadata": {},
   "source": [
    "## Build Graph\n",
    "\n",
    "The just follows the flow we outlined in the figure above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "dedae17a-98c6-474d-90a7-9234b7c8cea0",
   "metadata": {},
   "outputs": [],
   "source": ["from langgraph.graph import END, StateGraph, START\n\nworkflow = StateGraph(GraphState)\n\n# Define the nodes\nworkflow.add_node(\"retrieve\", retrieve)  # retrieve\nworkflow.add_node(\"grade_documents\", grade_documents)  # grade documents\nworkflow.add_node(\"generate\", generate)  # generatae\nworkflow.add_node(\"transform_query\", transform_query)  # transform_query\nworkflow.add_node(\"web_search_node\", web_search)  # web search\n\n# Build graph\nworkflow.add_edge(START, \"retrieve\")\nworkflow.add_edge(\"retrieve\", \"grade_documents\")\nworkflow.add_conditional_edges(\n    \"grade_documents\",\n    decide_to_generate,\n    {\n        \"transform_query\": \"transform_query\",\n        \"generate\": \"generate\",\n    },\n)\nworkflow.add_edge(\"transform_query\", \"web_search_node\")\nworkflow.add_edge(\"web_search_node\", \"generate\")\nworkflow.add_edge(\"generate\", END)\n\n# Compile\napp = workflow.compile()"]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "f5b7c2fe-1fc7-4b76-bf93-ba701a40aa6b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---RETRIEVE---\n",
      "\"Node 'retrieve':\"\n",
      "'\\n---\\n'\n",
      "---CHECK DOCUMENT RELEVANCE TO QUESTION---\n",
      "---GRADE: DOCUMENT NOT RELEVANT---\n",
      "---GRADE: DOCUMENT NOT RELEVANT---\n",
      "---GRADE: DOCUMENT RELEVANT---\n",
      "---GRADE: DOCUMENT RELEVANT---\n",
      "\"Node 'grade_documents':\"\n",
      "'\\n---\\n'\n",
      "---ASSESS GRADED DOCUMENTS---\n",
      "---DECISION: ALL DOCUMENTS ARE NOT RELEVANT TO QUESTION, TRANSFORM QUERY---\n",
      "---TRANSFORM QUERY---\n",
      "\"Node 'transform_query':\"\n",
      "'\\n---\\n'\n",
      "---WEB SEARCH---\n",
      "\"Node 'web_search_node':\"\n",
      "'\\n---\\n'\n",
      "---GENERATE---\n",
      "\"Node 'generate':\"\n",
      "'\\n---\\n'\n",
      "\"Node '__end__':\"\n",
      "'\\n---\\n'\n",
      "('Agents possess short-term memory, which is utilized for in-context learning, '\n",
      " 'and long-term memory, allowing them to retain and recall vast amounts of '\n",
      " 'information over extended periods. Some experts also classify working memory '\n",
      " 'as a distinct type, although it can be considered a part of short-term '\n",
      " 'memory in many cases.')\n"
     ]
    }
   ],
   "source": ["from pprint import pprint\n\n# Run\ninputs = {\"question\": \"What are the types of agent memory?\"}\nfor 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\npprint(value[\"generation\"])"]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "41ea1108-f385-4774-962d-db157922e231",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---RETRIEVE---\n",
      "\"Node 'retrieve':\"\n",
      "'\\n---\\n'\n",
      "---CHECK DOCUMENT RELEVANCE TO QUESTION---\n",
      "---GRADE: DOCUMENT NOT RELEVANT---\n",
      "---GRADE: DOCUMENT NOT RELEVANT---\n",
      "---GRADE: DOCUMENT NOT RELEVANT---\n",
      "---GRADE: DOCUMENT RELEVANT---\n",
      "\"Node 'grade_documents':\"\n",
      "'\\n---\\n'\n",
      "---ASSESS GRADED DOCUMENTS---\n",
      "---DECISION: ALL DOCUMENTS ARE NOT RELEVANT TO QUESTION, TRANSFORM QUERY---\n",
      "---TRANSFORM QUERY---\n",
      "\"Node 'transform_query':\"\n",
      "'\\n---\\n'\n",
      "---WEB SEARCH---\n",
      "\"Node 'web_search_node':\"\n",
      "'\\n---\\n'\n",
      "---GENERATE---\n",
      "\"Node 'generate':\"\n",
      "'\\n---\\n'\n",
      "\"Node '__end__':\"\n",
      "'\\n---\\n'\n",
      "('The AlphaCodium paper functions by proposing a code-oriented iterative flow '\n",
      " 'that involves repeatedly running and fixing generated code against '\n",
      " 'input-output tests. Its key mechanisms include generating additional data '\n",
      " 'like problem reflection and test reasoning to aid the iterative process, as '\n",
      " 'well as enriching the code generation process. AlphaCodium aims to improve '\n",
      " 'the performance of Large Language Models on code problems by following a '\n",
      " 'test-based, multi-stage approach.')\n"
     ]
    }
   ],
   "source": ["from pprint import pprint\n\n# Run\ninputs = {\"question\": \"How does the AlphaCodium paper work?\"}\nfor 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\npprint(value[\"generation\"])"]
  },
  {
   "cell_type": "markdown",
   "id": "a7e44593-1959-4abf-8405-5e23aa9398f5",
   "metadata": {},
   "source": [
    "LangSmith Traces - \n",
    " \n",
    "* https://smith.langchain.com/public/f6b1716c-e842-4282-9112-1026b93e246b/r\n",
    "\n",
    "* https://smith.langchain.com/public/497c8ed9-d9e2-429e-8ada-e64de3ec26c9/r"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ce65be5-fd12-4ffc-984c-34c132693e69",
   "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
}
