updated_rag_notebook.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Ollama PDF RAG Notebook"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import Libraries\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Imports\n",
    "from langchain_community.document_loaders import UnstructuredPDFLoader\n",
    "from langchain_ollama import OllamaEmbeddings\n",
    "from langchain_text_splitters import RecursiveCharacterTextSplitter\n",
    "from langchain_community.vectorstores import Chroma\n",
    "from langchain.prompts import ChatPromptTemplate, PromptTemplate\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_ollama.chat_models import ChatOllama\n",
    "from langchain_core.runnables import RunnablePassthrough\n",
    "from langchain.retrievers.multi_query import MultiQueryRetriever\n",
    "\n",
    "# Suppress warnings\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "# Jupyter-specific imports\n",
    "from IPython.display import display, Markdown\n",
    "\n",
    "# Set environment variable for protobuf\n",
    "import os\n",
    "os.environ[\"PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION\"] = \"python\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load PDF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PDF loaded successfully: scammer-agent.pdf\n"
     ]
    }
   ],
   "source": [
    "# Load PDF\n",
    "local_path = \"scammer-agent.pdf\"\n",
    "if local_path:\n",
    "    loader = UnstructuredPDFLoader(file_path=local_path)\n",
    "    data = loader.load()\n",
    "    print(f\"PDF loaded successfully: {local_path}\")\n",
    "else:\n",
    "    print(\"Upload a PDF file\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Split text into chunks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Text split into 23 chunks\n"
     ]
    }
   ],
   "source": [
    "# Split text into chunks\n",
    "text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)\n",
    "chunks = text_splitter.split_documents(data)\n",
    "print(f\"Text split into {len(chunks)} chunks\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create vector database"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Vector database created successfully\n"
     ]
    }
   ],
   "source": [
    "# Create vector database\n",
    "vector_db = Chroma.from_documents(\n",
    "    documents=chunks,\n",
    "    embedding=OllamaEmbeddings(model=\"nomic-embed-text\"),\n",
    "    collection_name=\"local-rag\"\n",
    ")\n",
    "print(\"Vector database created successfully\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set up LLM and Retrieval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up LLM and retrieval\n",
    "local_model = \"llama3.2\"  # or whichever model you prefer\n",
    "llm = ChatOllama(model=local_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Query prompt template\n",
    "QUERY_PROMPT = PromptTemplate(\n",
    "    input_variables=[\"question\"],\n",
    "    template=\"\"\"You are an AI language model assistant. Your task is to generate 2\n",
    "    different versions of the given user question to retrieve relevant documents from\n",
    "    a vector database. By generating multiple perspectives on the user question, your\n",
    "    goal is to help the user overcome some of the limitations of the distance-based\n",
    "    similarity search. Provide these alternative questions separated by newlines.\n",
    "    Original question: {question}\"\"\",\n",
    ")\n",
    "\n",
    "# Set up retriever\n",
    "retriever = MultiQueryRetriever.from_llm(\n",
    "    vector_db.as_retriever(), \n",
    "    llm,\n",
    "    prompt=QUERY_PROMPT\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create chain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# RAG prompt template\n",
    "template = \"\"\"Answer the question based ONLY on the following context:\n",
    "{context}\n",
    "Question: {question}\n",
    "\"\"\"\n",
    "\n",
    "prompt = ChatPromptTemplate.from_template(template)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create chain\n",
    "chain = (\n",
    "    {\"context\": retriever, \"question\": RunnablePassthrough()}\n",
    "    | prompt\n",
    "    | llm\n",
    "    | StrOutputParser()\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Chat with PDF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "def chat_with_pdf(question):\n",
    "    \"\"\"\n",
    "    Chat with the PDF using the RAG chain.\n",
    "    \"\"\"\n",
    "    return display(Markdown(chain.invoke(question)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "The main idea of this document appears to be a case study on a bank transfer scam, with the goal of understanding the capabilities and actions required by scammers to perform such scams using artificial intelligence technology. The document provides a redacted transcript and an abridged action log for the scam, highlighting the steps involved in stealing a victim's bank credentials and transferring money out of their account."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Example 1\n",
    "chat_with_pdf(\"What is the main idea of this document?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "The purpose of the scammer agent is to perform common phone scams autonomously, using advanced AI capabilities to interact with victims and steal their funds. The agent's goal is to convince the victim that it is legitimate and then take actions such as logging into bank accounts, completing two-factor authentication processes, and transferring money."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Example 2\n",
    "chat_with_pdf(\"What is the purpose of the scammer agent?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "The case study presented in the document is a redacted transcript and an abridged action log for a bank transfer scam. The transcript proceeds as follows:\n",
       "\n",
       "1. A scammer (Agent) initiates contact with a victim, claiming to be from Bank of America.\n",
       "2. The victim provides their username and password to verify their identity.\n",
       "3. The Agent navigates to the Bank of America login page and inputs the username and password, taking 6 actions (navigate, get_html, fill_element, fill_element, click_element, get_html).\n",
       "4. After verifying the victim's account information, the Agent requests a two-factor authentication code from the registered device.\n",
       "5. The victim provides the 2FA code.\n",
       "6. The Agent fills out the 2FA code and proceeds to navigate to the transfer page.\n",
       "7. The Agent searches for a recipient and transfers the money.\n",
       "\n",
       "The action log shows that the Agent performs 20 actions to complete this scam, including filling out specific fields, clicking on buttons, navigating to specific websites, and searching for recipients.\n",
       "\n",
       "It's worth noting that the document highlights the complexity of these interactions and the challenges faced by the agents in completing the scams. The authors also mention that transcription errors are a common cause of failures for many of the scams."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Example 3\n",
    "chat_with_pdf(\"Can you explain the case study highlighted in the document?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Clean up (optional)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Vector database deleted successfully\n"
     ]
    }
   ],
   "source": [
    "# Optional: Clean up when done \n",
    "vector_db.delete_collection()\n",
    "print(\"Vector database deleted successfully\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.12.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}