{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Install dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture --no-stderr\n",
    "%pip install -U --quiet langchain-community tiktoken langchain-openai langchainhub chromadb langchain langgraph langchain-text-splitters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Set environment variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "import getpass\n",
    "import time\n",
    "initial_time = time.time()\n",
    "\n",
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = getpass.getpass(\"Enter your OpenAI API key: \")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup Retriever"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.document_loaders import WebBaseLoader\n",
    "from langchain_community.vectorstores import Chroma\n",
    "from langchain_openai import OpenAIEmbeddings\n",
    "from langchain_text_splitters import RecursiveCharacterTextSplitter\n",
    "\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",
    "docs = [WebBaseLoader(url).load() for url in urls]\n",
    "docs_list = [item for sublist in docs for item in sublist]\n",
    "\n",
    "text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(\n",
    "    chunk_size=100, chunk_overlap=50\n",
    ")\n",
    "doc_splits = text_splitter.split_documents(docs_list)\n",
    "\n",
    "# Add to vectorDB\n",
    "vectorstore = Chroma.from_documents(\n",
    "    documents=doc_splits,\n",
    "    collection_name=\"rag-chroma\",\n",
    "    embedding=OpenAIEmbeddings(),\n",
    ")\n",
    "retriever = vectorstore.as_retriever()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create Retriever Tool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.tools.retriever import create_retriever_tool\n",
    "\n",
    "retriever_tool = create_retriever_tool(\n",
    "    retriever,\n",
    "    \"retrieve_blog_posts\",\n",
    "    \"Search and return information about Lilian Weng blog posts on LLM agents, prompt engineering, and adversarial attacks on LLMs.\",\n",
    ")\n",
    "\n",
    "tools = [retriever_tool]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Agent State"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Annotated, Sequence\n",
    "from typing_extensions import TypedDict\n",
    "\n",
    "from langchain_core.messages import BaseMessage\n",
    "\n",
    "from langgraph.graph.message import add_messages\n",
    "\n",
    "\n",
    "class AgentState(TypedDict):\n",
    "    # The add_messages function defines how an update should be processed\n",
    "    # Default is to replace. add_messages says \"append\"\n",
    "    messages: Annotated[Sequence[BaseMessage], add_messages]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Nodes and Edges"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "********************Prompt[rlm/rag-prompt]********************\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/joaomoura/.pyenv/versions/3.11.7/lib/python3.11/site-packages/langsmith/client.py:241: LangSmithMissingAPIKeyWarning: API key must be provided when using hosted LangSmith API\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "You are an assistant for question-answering tasks. Use the following pieces of retrieved context to answer the question. If you don't know the answer, just say that you don't know. Use three sentences maximum and keep the answer concise.\n",
      "Question: \u001b[33;1m\u001b[1;3m{question}\u001b[0m \n",
      "Context: \u001b[33;1m\u001b[1;3m{context}\u001b[0m \n",
      "Answer:\n"
     ]
    }
   ],
   "source": [
    "from typing import Annotated, Literal, Sequence\n",
    "from typing_extensions import TypedDict\n",
    "\n",
    "from langchain import hub\n",
    "from langchain_core.messages import BaseMessage, HumanMessage\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.prompts import PromptTemplate\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "from pydantic import BaseModel, Field\n",
    "\n",
    "\n",
    "from langgraph.prebuilt import tools_condition\n",
    "\n",
    "### Edges\n",
    "\n",
    "\n",
    "def grade_documents(state) -> Literal[\"generate\", \"rewrite\"]:\n",
    "    \"\"\"\n",
    "    Determines whether the retrieved documents are relevant to the question.\n",
    "\n",
    "    Args:\n",
    "        state (messages): The current state\n",
    "\n",
    "    Returns:\n",
    "        str: A decision for whether the documents are relevant or not\n",
    "    \"\"\"\n",
    "\n",
    "    print(\"---CHECK RELEVANCE---\")\n",
    "\n",
    "    # Data model\n",
    "    class grade(BaseModel):\n",
    "        \"\"\"Binary score for relevance check.\"\"\"\n",
    "\n",
    "        binary_score: str = Field(description=\"Relevance score 'yes' or 'no'\")\n",
    "\n",
    "    # LLM\n",
    "    model = ChatOpenAI(temperature=0, model=\"gpt-4o-mini\", streaming=True)\n",
    "\n",
    "    # LLM with tool and validation\n",
    "    llm_with_tool = model.with_structured_output(grade)\n",
    "\n",
    "    # Prompt\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 {context} \\n\\n\n",
    "        Here is the user question: {question} \\n\n",
    "        If the document contains keyword(s) or semantic meaning related to the user 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.\"\"\",\n",
    "        input_variables=[\"context\", \"question\"],\n",
    "    )\n",
    "\n",
    "    # Chain\n",
    "    chain = prompt | llm_with_tool\n",
    "\n",
    "    messages = state[\"messages\"]\n",
    "    last_message = messages[-1]\n",
    "\n",
    "    question = messages[0].content\n",
    "    docs = last_message.content\n",
    "\n",
    "    scored_result = chain.invoke({\"question\": question, \"context\": docs})\n",
    "\n",
    "    score = scored_result.binary_score\n",
    "\n",
    "    if score == \"yes\":\n",
    "        print(\"---DECISION: DOCS RELEVANT---\")\n",
    "        return \"generate\"\n",
    "\n",
    "    else:\n",
    "        print(\"---DECISION: DOCS NOT RELEVANT---\")\n",
    "        print(score)\n",
    "        return \"rewrite\"\n",
    "\n",
    "\n",
    "### Nodes\n",
    "\n",
    "\n",
    "def agent(state):\n",
    "    \"\"\"\n",
    "    Invokes the agent model to generate a response based on the current state. Given\n",
    "    the question, it will decide to retrieve using the retriever tool, or simply end.\n",
    "\n",
    "    Args:\n",
    "        state (messages): The current state\n",
    "\n",
    "    Returns:\n",
    "        dict: The updated state with the agent response appended to messages\n",
    "    \"\"\"\n",
    "    print(\"---CALL AGENT---\")\n",
    "    messages = state[\"messages\"]\n",
    "    model = ChatOpenAI(temperature=0, streaming=True, model=\"gpt-4-turbo\")\n",
    "    model = model.bind_tools(tools)\n",
    "    response = model.invoke(messages)\n",
    "    # We return a list, because this will get added to the existing list\n",
    "    return {\"messages\": [response]}\n",
    "\n",
    "\n",
    "def rewrite(state):\n",
    "    \"\"\"\n",
    "    Transform the query to produce a better question.\n",
    "\n",
    "    Args:\n",
    "        state (messages): The current state\n",
    "\n",
    "    Returns:\n",
    "        dict: The updated state with re-phrased question\n",
    "    \"\"\"\n",
    "\n",
    "    print(\"---TRANSFORM QUERY---\")\n",
    "    messages = state[\"messages\"]\n",
    "    question = messages[0].content\n",
    "\n",
    "    msg = [\n",
    "        HumanMessage(\n",
    "            content=f\"\"\" \\n\n",
    "    Look at the input and try to reason about the underlying semantic intent / meaning. \\n\n",
    "    Here is the initial question:\n",
    "    \\n ------- \\n\n",
    "    {question}\n",
    "    \\n ------- \\n\n",
    "    Formulate an improved question: \"\"\",\n",
    "        )\n",
    "    ]\n",
    "\n",
    "    # Grader\n",
    "    model = ChatOpenAI(temperature=0, model=\"gpt-4o-mini\", streaming=True)\n",
    "    response = model.invoke(msg)\n",
    "    return {\"messages\": [response]}\n",
    "\n",
    "\n",
    "def generate(state):\n",
    "    \"\"\"\n",
    "    Generate answer\n",
    "\n",
    "    Args:\n",
    "        state (messages): The current state\n",
    "\n",
    "    Returns:\n",
    "         dict: The updated state with re-phrased question\n",
    "    \"\"\"\n",
    "    print(\"---GENERATE---\")\n",
    "    messages = state[\"messages\"]\n",
    "    question = messages[0].content\n",
    "    last_message = messages[-1]\n",
    "\n",
    "    docs = last_message.content\n",
    "\n",
    "    # Prompt\n",
    "    prompt = hub.pull(\"rlm/rag-prompt\")\n",
    "\n",
    "    # LLM\n",
    "    llm = ChatOpenAI(model_name=\"gpt-4o-mini\", temperature=0, streaming=True)\n",
    "\n",
    "    # Post-processing\n",
    "    def format_docs(docs):\n",
    "        return \"\\n\\n\".join(doc.page_content for doc in docs)\n",
    "\n",
    "    # Chain\n",
    "    rag_chain = prompt | llm | StrOutputParser()\n",
    "\n",
    "    # Run\n",
    "    response = rag_chain.invoke({\"context\": docs, \"question\": question})\n",
    "    return {\"messages\": [response]}\n",
    "\n",
    "\n",
    "print(\"*\" * 20 + \"Prompt[rlm/rag-prompt]\" + \"*\" * 20)\n",
    "prompt = hub.pull(\"rlm/rag-prompt\").pretty_print()  # Show what the prompt looks like"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create Graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.graph import END, StateGraph, START\n",
    "from langgraph.prebuilt import ToolNode\n",
    "\n",
    "# Define a new graph\n",
    "workflow = StateGraph(AgentState)\n",
    "\n",
    "# Define the nodes we will cycle between\n",
    "workflow.add_node(\"agent\", agent)  # agent\n",
    "retrieve = ToolNode([retriever_tool])\n",
    "workflow.add_node(\"retrieve\", retrieve)  # retrieval\n",
    "workflow.add_node(\"rewrite\", rewrite)  # Re-writing the question\n",
    "workflow.add_node(\n",
    "    \"generate\", generate\n",
    ")  # Generating a response after we know the documents are relevant\n",
    "# Call agent node to decide to retrieve or not\n",
    "workflow.add_edge(START, \"agent\")\n",
    "\n",
    "# Decide whether to retrieve\n",
    "workflow.add_conditional_edges(\n",
    "    \"agent\",\n",
    "    # Assess agent decision\n",
    "    tools_condition,\n",
    "    {\n",
    "        # Translate the condition outputs to nodes in our graph\n",
    "        \"tools\": \"retrieve\",\n",
    "        END: END,\n",
    "    },\n",
    ")\n",
    "\n",
    "# Edges taken after the `action` node is called.\n",
    "workflow.add_conditional_edges(\n",
    "    \"retrieve\",\n",
    "    # Assess agent decision\n",
    "    grade_documents,\n",
    ")\n",
    "workflow.add_edge(\"generate\", END)\n",
    "workflow.add_edge(\"rewrite\", \"agent\")\n",
    "\n",
    "# Compile\n",
    "graph = workflow.compile()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plot Graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 25060 characters]",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.display import Image, display\n",
    "\n",
    "try:\n",
    "    display(Image(graph.get_graph(xray=True).draw_mermaid_png()))\n",
    "except Exception:\n",
    "    # This requires some extra dependencies and is optional\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Run Graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---CALL AGENT---\n",
      "\"Output from node 'agent':\"\n",
      "'---'\n",
      "{ 'messages': [ AIMessage(content='', additional_kwargs={'tool_calls': [{'index': 0, 'id': 'call_R8JiVuIfmbTxeMNldY63YMgi', 'function': {'arguments': '{\"query\":\"types of agent memory\"}', 'name': 'retrieve_blog_posts'}, 'type': 'function'}]}, response_metadata={'finish_reason': 'tool_calls', 'model_name': 'gpt-4-turbo-2024-04-09', 'system_fingerprint': 'fp_1e9a2d1e89'}, id='run-d81a8fad-7d66-4690-bc8c-9cf2a8dc02ec-0', tool_calls=[{'name': 'retrieve_blog_posts', 'args': {'query': 'types of agent memory'}, 'id': 'call_R8JiVuIfmbTxeMNldY63YMgi', 'type': 'tool_call'}])]}\n",
      "'\\n---\\n'\n",
      "---CHECK RELEVANCE---\n",
      "---DECISION: DOCS NOT RELEVANT---\n",
      "no\n",
      "\"Output from node 'retrieve':\"\n",
      "'---'\n",
      "{ 'messages': [ ToolMessage(content='Table of Contents\\n\\n\\n\\nAgent System Overview\\n\\nComponent One: Planning\\n\\nTask Decomposition\\n\\nSelf-Reflection\\n\\n\\nComponent Two: Memory\\n\\nTypes of Memory\\n\\nMaximum Inner Product Search (MIPS)\\n\\n\\nComponent Three: Tool Use\\n\\nCase Studies\\n\\nScientific Discovery Agent\\n\\nGenerative Agents Simulation\\n\\nProof-of-Concept Examples\\n\\n\\nChallenges\\n\\nCitation\\n\\nReferences\\n\\nTable of Contents\\n\\n\\n\\nAgent System Overview\\n\\nComponent One: Planning\\n\\nTask Decomposition\\n\\nSelf-Reflection\\n\\n\\nComponent Two: Memory\\n\\nTypes of Memory\\n\\nMaximum Inner Product Search (MIPS)\\n\\n\\nComponent Three: Tool Use\\n\\nCase Studies\\n\\nScientific Discovery Agent\\n\\nGenerative Agents Simulation\\n\\nProof-of-Concept Examples\\n\\n\\nChallenges\\n\\nCitation\\n\\nReferences\\n\\nTable of Contents\\n\\n\\n\\nAgent System Overview\\n\\nComponent One: Planning\\n\\nTask Decomposition\\n\\nSelf-Reflection\\n\\n\\nComponent Two: Memory\\n\\nTypes of Memory\\n\\nMaximum Inner Product Search (MIPS)\\n\\n\\nComponent Three: Tool Use\\n\\nCase Studies\\n\\nScientific Discovery Agent\\n\\nGenerative Agents Simulation\\n\\nProof-of-Concept Examples\\n\\n\\nChallenges\\n\\nCitation\\n\\nReferences\\n\\nPlanning\\n\\nSubgoal and decomposition: The agent breaks down large tasks into smaller, manageable subgoals, enabling efficient handling of complex tasks.\\nReflection and refinement: The agent can do self-criticism and self-reflection over past actions, learn from mistakes and refine them for future steps, thereby improving the quality of final results.\\n\\n\\nMemory', name='retrieve_blog_posts', id='936c1034-d281-46d5-9225-be750a61b60c', tool_call_id='call_R8JiVuIfmbTxeMNldY63YMgi')]}\n",
      "'\\n---\\n'\n",
      "---TRANSFORM QUERY---\n",
      "\"Output from node 'rewrite':\"\n",
      "'---'\n",
      "{ 'messages': [ AIMessage(content='What insights does Lilian Weng provide regarding the different types of memory in agents?', additional_kwargs={}, response_metadata={'finish_reason': 'stop', 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_6fc10e10eb'}, id='run-fdb07ad3-6d05-44c3-b6f8-e6a338614f41-0')]}\n",
      "'\\n---\\n'\n",
      "---CALL AGENT---\n",
      "\"Output from node 'agent':\"\n",
      "'---'\n",
      "{ 'messages': [ AIMessage(content='Lilian Weng discusses the concept of memory in agents under the \"Component Two: Memory\" section of her blog. She categorizes memory into different types, although the specific details or categories of these types are not explicitly listed in the brief content overview provided. However, she does mention the use of techniques like Maximum Inner Product Search (MIPS) which is often related to memory systems in agents, suggesting a focus on efficient retrieval mechanisms which could be part of how different types of memory are utilized or implemented in agent systems. This indicates a technical approach to understanding and optimizing how memory functions within agents, possibly differentiating between types based on functionality, such as short-term, long-term, or task-specific memories.', additional_kwargs={}, response_metadata={'finish_reason': 'stop', 'model_name': 'gpt-4-turbo-2024-04-09', 'system_fingerprint': 'fp_1e9a2d1e89'}, id='run-5a1f3733-e1ca-4a14-840a-b65c672be530-0')]}\n",
      "'\\n---\\n'\n"
     ]
    }
   ],
   "source": [
    "import pprint\n",
    "\n",
    "inputs = {\n",
    "    \"messages\": [\n",
    "        (\"user\", \"What does Lilian Weng say about the types of agent memory?\"),\n",
    "    ]\n",
    "}\n",
    "for output in graph.stream(inputs):\n",
    "    for key, value in output.items():\n",
    "        pprint.pprint(f\"Output from node '{key}':\")\n",
    "        pprint.pprint(\"---\")\n",
    "        pprint.pprint(value, indent=2, width=80, depth=None)\n",
    "    pprint.pprint(\"\\n---\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total execution time: 914.08 seconds\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "final_time = time.time()\n",
    "print(f\"Total execution time: {final_time - initial_time:.2f} seconds\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
