{
 "cells": [
  {
   "attachments": {
    "5fca0a3e-d13d-4bfa-95ea-58203640cc7a.png": {
     "image/png": "[encoded data, 222928 characters]"
    }
   },
   "cell_type": "markdown",
   "id": "848ba742-7443-4123-8115-061da9823309",
   "metadata": {},
   "source": [
    "# Self RAG using local LLMs\n",
    "\n",
    "Self-RAG is a strategy for RAG that incorporates self-reflection / self-grading on retrieved documents and generations. \n",
    "\n",
    "In the [paper](https://arxiv.org/abs/2310.11511), a few decisions are made:\n",
    "\n",
    "1. Should I retrieve from retriever, `R` -\n",
    "\n",
    "* Input: `x (question)` OR `x (question)`, `y (generation)`\n",
    "* Decides when to retrieve `D` chunks with `R`\n",
    "* Output: `yes, no, continue`\n",
    "\n",
    "2. Are the retrieved passages `D` relevant to the question `x` -\n",
    "\n",
    "* * Input: (`x (question)`, `d (chunk)`) for `d` in `D`\n",
    "* `d` provides useful information to solve `x`\n",
    "* Output: `relevant, irrelevant`\n",
    "\n",
    "3. Are the LLM generation from each chunk in `D` is relevant to the chunk (hallucinations, etc)  -\n",
    "\n",
    "* Input: `x (question)`, `d (chunk)`,  `y (generation)` for `d` in `D`\n",
    "* All of the verification-worthy statements in `y (generation)` are supported by `d`\n",
    "* Output: `{fully supported, partially supported, no support`\n",
    "\n",
    "4. The LLM generation from each chunk in `D` is a useful response to `x (question)` -\n",
    "\n",
    "* Input: `x (question)`, `y (generation)` for `d` in `D`\n",
    "* `y (generation)` is a useful response to `x (question)`.\n",
    "* Output: `{5, 4, 3, 2, 1}`\n",
    "\n",
    "We will implement some of these ideas from scratch using [LangGraph](https://langchain-ai.github.io/langgraph/).\n",
    "\n",
    "![Screenshot 2024-04-01 at 12.42.59 PM.png](attachment:5fca0a3e-d13d-4bfa-95ea-58203640cc7a.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ed0a85a-a33b-40a6-99fa-2444bf57a6cc",
   "metadata": {},
   "source": [
    "# Environment "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7f9cc6d-a70c-433a-b0ad-ea47c5a0717e",
   "metadata": {},
   "outputs": [],
   "source": ["%capture --no-stderr\n%pip install -U langchain-nomic langchain_community tiktoken langchainhub chromadb langchain langgraph nomic[local]"]
  },
  {
   "cell_type": "markdown",
   "id": "ccc6b6bd-a2fa-4a43-83b7-704b8b6fb855",
   "metadata": {},
   "source": [
    "### LLMs\n",
    "\n",
    "#### Local Embeddings\n",
    "\n",
    "You can use `GPT4AllEmbeddings()` from Nomic, which can access use Nomic's recently released [v1](https://blog.nomic.ai/posts/nomic-embed-text-v1) and [v1.5](https://blog.nomic.ai/posts/nomic-embed-matryoshka) embeddings.\n",
    "\n",
    "\n",
    "Follow the documentation [here](https://docs.gpt4all.io/gpt4all_python_embedding.html#supported-embedding-models).\n",
    "\n",
    "#### Local LLM\n",
    "\n",
    "(1) Download [Ollama app](https://ollama.ai/).\n",
    "\n",
    "(2) Download a `Mistral` model from various Mistral versions [here](https://ollama.ai/library/mistral) and Mixtral versions [here](https://ollama.ai/library/mixtral) available.\n",
    "```\n",
    "ollama pull mistral\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "bedffc73-6b10-42c8-8768-2085c8ed3398",
   "metadata": {},
   "outputs": [],
   "source": ["# Ollama model name\nlocal_llm = \"mistral\""]
  },
  {
   "cell_type": "markdown",
   "id": "6306aa61-2e6a-47b6-b306-4987d0636f97",
   "metadata": {},
   "source": [
    "### Tracing\n",
    "\n",
    "Optionally, use [LangSmith](https://docs.smith.langchain.com/) for tracing (shown at bottom)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2208f342-8163-4af3-8dc0-aa70f5e06143",
   "metadata": {},
   "outputs": [],
   "source": ["import os\n\nos.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": "ba68a46d-b617-4fdc-9113-fabdcf736feb",
   "metadata": {},
   "source": [
    "## Index\n",
    "\n",
    "Let's index 3 blog posts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3bb9060-ad74-4470-9991-2ba167b6b8d8",
   "metadata": {},
   "outputs": [],
   "source": ["from langchain.text_splitter import RecursiveCharacterTextSplitter\nfrom langchain_community.document_loaders import WebBaseLoader\nfrom langchain_community.vectorstores import Chroma\nfrom langchain_nomic.embeddings import NomicEmbeddings\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=NomicEmbeddings(model=\"nomic-embed-text-v1.5\", inference_mode=\"local\"),\n)\nretriever = vectorstore.as_retriever()"]
  },
  {
   "cell_type": "markdown",
   "id": "cc60ff95-7e12-4004-b18e-a067a2dcc201",
   "metadata": {},
   "source": [
    "## LLMs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3aad0c60-3208-48fb-af82-0024630b4da1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'score': 'yes'}\n"
     ]
    }
   ],
   "source": ["### Retrieval Grader\n\nfrom langchain.prompts import PromptTemplate\nfrom langchain_community.chat_models import ChatOllama\nfrom langchain_core.output_parsers import JsonOutputParser\n\n# LLM\nllm = ChatOllama(model=local_llm, format=\"json\", temperature=0)\n\nprompt = 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\nretrieval_grader = prompt | llm | JsonOutputParser()\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": 7,
   "id": "e5e45953-248d-492f-af28-d5e80c664c95",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " In an LLM-powered autonomous agent system, the Large Language Model (LLM) functions as the agent's brain. The agent has key components including memory, planning, and reflection mechanisms. The memory component is a long-term memory module that records a comprehensive list of agents’ experience in natural language. It includes a memory stream, which is an external database for storing past experiences. The reflection mechanism synthesizes memories into higher-level inferences over time and guides the agent's future behavior.\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 = ChatOllama(model=local_llm, 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": 8,
   "id": "56297862-df87-42a7-ba9d-310926dfb328",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'score': 'yes'}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": ["### Hallucination Grader\n\n# LLM\nllm = ChatOllama(model=local_llm, format=\"json\", temperature=0)\n\n# Prompt\nprompt = PromptTemplate(\n    template=\"\"\"You are a grader assessing whether an answer is grounded in / supported by a set of facts. \\n \n    Here are the facts:\n    \\n ------- \\n\n    {documents} \n    \\n ------- \\n\n    Here is the answer: {generation}\n    Give a binary score 'yes' or 'no' score to indicate whether the answer is grounded in / supported by a set of facts. \\n\n    Provide the binary score as a JSON with a single key 'score' and no preamble or explanation.\"\"\",\n    input_variables=[\"generation\", \"documents\"],\n)\n\nhallucination_grader = prompt | llm | JsonOutputParser()\nhallucination_grader.invoke({\"documents\": docs, \"generation\": generation})"]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "e1dd9174-2df6-45b1-8e69-13381f579c39",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'score': 'yes'}"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": ["### Answer Grader\n\n# LLM\nllm = ChatOllama(model=local_llm, format=\"json\", temperature=0)\n\n# Prompt\nprompt = PromptTemplate(\n    template=\"\"\"You are a grader assessing whether an answer is useful to resolve a question. \\n \n    Here is the answer:\n    \\n ------- \\n\n    {generation} \n    \\n ------- \\n\n    Here is the question: {question}\n    Give a binary score 'yes' or 'no' to indicate whether the answer is useful to resolve a question. \\n\n    Provide the binary score as a JSON with a single key 'score' and no preamble or explanation.\"\"\",\n    input_variables=[\"generation\", \"question\"],\n)\n\nanswer_grader = prompt | llm | JsonOutputParser()\nanswer_grader.invoke({\"question\": question, \"generation\": generation})"]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "5216d92b-1ca1-4bcf-a34f-c99ec2766b54",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "' What is agent memory and how can it be effectively utilized in vector database retrieval?'"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": ["### Question Re-writer\n\n# LLM\nllm = ChatOllama(model=local_llm, temperature=0)\n\n# Prompt\nre_write_prompt = PromptTemplate(\n    template=\"\"\"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 initial and formulate an improved question. \\n\n     Here is the initial question: \\n\\n {question}. Improved question with no preamble: \\n \"\"\",\n    input_variables=[\"generation\", \"question\"],\n)\n\nquestion_rewriter = re_write_prompt | llm | StrOutputParser()\nquestion_rewriter.invoke({\"question\": question})"]
  },
  {
   "cell_type": "markdown",
   "id": "3d3339b9-5f30-4d54-bfc9-9091c6035955",
   "metadata": {},
   "source": [
    "# Graph \n",
    "\n",
    "Capture the flow in as a graph.\n",
    "\n",
    "## Graph state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "90fb1dc6-c482-483a-8441-39965c401beb",
   "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        documents: list of documents\n    \"\"\"\n\n    question: str\n    generation: str\n    documents: List[str]"]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "5324ea49-5745-47b5-a0a5-bf58c8babe46",
   "metadata": {},
   "outputs": [],
   "source": ["### Nodes\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    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            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\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\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    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\ndef 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[\"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[\"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        print(\"---DECISION: GENERATION IS NOT GROUNDED IN DOCUMENTS, RE-TRY---\")\n        return \"not supported\""]
  },
  {
   "cell_type": "markdown",
   "id": "bdf3826c-b668-4f0b-bf83-81c40baaaf02",
   "metadata": {},
   "source": [
    "## Build Graph\n",
    "\n",
    "This just follows the flow we outlined in the figure above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "5605dee4-b2df-46ae-a640-cc2ed90c21a6",
   "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\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\", \"retrieve\")\nworkflow.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\napp = workflow.compile()"]
  },
  {
   "cell_type": "markdown",
   "id": "105ae1b5-6963-4186-bb83-6d6cb96d095f",
   "metadata": {},
   "source": [
    "## Run\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "26a64f7d-0c14-4e31-a67f-63021dee626e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---RETRIEVE---\n",
      "\"Node 'retrieve':\"\n",
      "'\\n---\\n'\n",
      "---CHECK DOCUMENT RELEVANCE TO QUESTION---\n",
      "---GRADE: DOCUMENT RELEVANT---\n",
      "---GRADE: DOCUMENT RELEVANT---\n",
      "---GRADE: DOCUMENT RELEVANT---\n",
      "---GRADE: DOCUMENT RELEVANT---\n",
      "\"Node 'grade_documents':\"\n",
      "'\\n---\\n'\n",
      "---ASSESS GRADED DOCUMENTS---\n",
      "---DECISION: GENERATE---\n",
      "---GENERATE---\n",
      "\"Node 'generate':\"\n",
      "'\\n---\\n'\n",
      "---CHECK HALLUCINATIONS---\n",
      "---DECISION: GENERATION IS GROUNDED IN DOCUMENTS---\n",
      "---GRADE GENERATION vs QUESTION---\n",
      "---DECISION: GENERATION ADDRESSES QUESTION---\n",
      "\"Node '__end__':\"\n",
      "'\\n---\\n'\n",
      "(' In a LLM-powered autonomous agent system, memory is a key component that '\n",
      " 'enables agents to store and retrieve information. There are different types '\n",
      " 'of memory in human brains, such as sensory memory which retains impressions '\n",
      " 'of sensory information for a few seconds, and long-term memory which records '\n",
      " \"experiences for extended periods (Lil'Log, 2023). In the context of LLM \"\n",
      " 'agents, memory is often implemented as an external database or memory stream '\n",
      " \"(Lil'Log, 2023). The agent can consult this memory to inform its behavior \"\n",
      " 'based on relevance, recency, and importance. Additionally, reflection '\n",
      " 'mechanisms synthesize memories into higher-level inferences over time and '\n",
      " \"guide the agent's future behavior (Lil'Log, 2023).\")\n"
     ]
    }
   ],
   "source": ["from pprint import pprint\n\n# Run\ninputs = {\"question\": \"Explain how the different types of agent memory 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": "f9a27907-2611-4791-910b-0d66c59f5cf5",
   "metadata": {},
   "source": [
    "Trace: \n",
    "\n",
    "https://smith.langchain.com/public/4163a342-5260-4852-8602-bda3f95177e7/r"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "953143c2-2f2a-4361-a36b-87db7cf21d63",
   "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
}
