{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "437161e4",
   "metadata": {},
   "source": [
    "# LangGraph Handling LangChain Agent Intermediate_Steps \n",
    "\n",
    "\n",
    "In this notebook we will learn how to build a basic [agent executor](https://api.python.langchain.com/en/latest/agents/langchain.agents.agent.AgentExecutor.html) leveraging [langGraph](https://github.com/langchain-ai/langgraph).\n",
    "\n",
    "We demonstrate how to handle the logic of the intermediate steps from the agent leveraging different provided tools within langGraph.\n",
    "\n",
    "\n",
    "- We will be leveraging LLM [mixtral-8x7b-instruct-v0.1 from NVIDIA API Catalog](https://build.nvidia.com/mistralai/mixtral-8x7b-instruct).\n",
    "\n",
    "- Simple Faiss Retriever as one of the tools with the [NV-Embed-QA from NVIDIA API Catalog](https://build.nvidia.com/nvidia/embed-qa-4).\n",
    "\n",
    "- Wikipedia (the pip installable package) as one of the tools.\n",
    "\n",
    "Then we will utilize with LangGraph to control and intervene intermediate steps as well as the outputs from the agent.\n",
    "\n",
    "\n",
    "## Prerequisites \n",
    "\n",
    "To run this notebook, you need to complete the [setup](https://python.langchain.com/docs/integrations/text_embedding/nvidia_ai_endpoints#setup) and generate an API key."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d62742d",
   "metadata": {},
   "source": [
    "## Install additional Python packages \n",
    "\n",
    "Install the additional packages that required for this example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20cb957f",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install --upgrade pip\n",
    "!pip install langchain==0.2.5\n",
    "!pip install langchain-nvidia-ai-endpoints==0.1.2\n",
    "!pip install langchain-community==0.2.2\n",
    "!pip install langgraph==0.0.62\n",
    "!pip install faiss-gpu==1.7.2\n",
    "!pip install wikipedia==1.4.0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd39cd00",
   "metadata": {},
   "source": [
    "## Step 1  - Export the NVIDIA_API_KEY "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b52cefe2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "## API Key can be found by going to NVIDIA NGC -> AI Foundation Models -> (some model) -> Get API Code or similar.\n",
    "## 10K free queries to any endpoint (which is a lot actually).\n",
    "\n",
    "# del os.environ['NVIDIA_API_KEY']  ## delete key and reset\n",
    "if os.environ.get(\"NVIDIA_API_KEY\", \"\").startswith(\"nvapi-\"):\n",
    "    print(\"Valid NVIDIA_API_KEY already in environment. Delete to reset\")\n",
    "else:\n",
    "    nvapi_key = getpass.getpass(\"NVAPI Key (starts with nvapi-): \")\n",
    "    assert nvapi_key.startswith(\"nvapi-\"), f\"{nvapi_key[:5]}... is not a valid key\"\n",
    "    os.environ[\"NVIDIA_API_KEY\"] = nvapi_key"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3bdf4d3",
   "metadata": {},
   "source": [
    "Optionally, we can set API key for [LangSmith tracing](https://smith.langchain.com/), which will give us best-in-class observability."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db53ef9e",
   "metadata": {},
   "source": [
    "## Step 2 - Initialize the LLM and embedding models\n",
    "\n",
    "The following code sets ai-mixtral-8x7b-instruct as the main LLM and ai-embed-qa-4 as the embedding model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bdd1cd0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.vectorstores import FAISS\n",
    "\n",
    "from langchain_nvidia_ai_endpoints import ChatNVIDIA\n",
    "from langchain_nvidia_ai_endpoints import NVIDIAEmbeddings\n",
    "\n",
    "llm = ChatNVIDIA(model=\"mistralai/mixtral-8x7b-instruct-v0.1\", nvidia_api_key=nvapi_key, max_tokens=2048)\n",
    "embedder = NVIDIAEmbeddings(model=\"NV-Embed-QA\", truncate=\"END\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07e5fb59",
   "metadata": {},
   "source": [
    "## Step 3 - Retriever from FAISS vector store\n",
    "\n",
    "We need to process a toy example, here we use `Sweden.txt` from the `data` folder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2192ca55",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.vectorstores import FAISS\n",
    "from langchain.text_splitter import CharacterTextSplitter\n",
    "\n",
    "# We need to process the text data and prepare them.\n",
    "p = \"Sweden.txt\"\n",
    "data = []\n",
    "sources = []\n",
    "path2file = \"./data/\" + p\n",
    "with open(path2file, encoding=\"utf-8\") as f:\n",
    "    lines = f.readlines()\n",
    "    for line in lines:\n",
    "        if len(line) >= 1:\n",
    "            data.append(line)\n",
    "            sources.append(path2file)\n",
    "documents = [d for d in data if d != '\\n']\n",
    "\n",
    "# create docs and metadatas\n",
    "text_splitter = CharacterTextSplitter(chunk_size=400, separator=\" \")\n",
    "docs = []\n",
    "metadatas = []\n",
    "\n",
    "for i, d in enumerate(documents):\n",
    "    splits = text_splitter.split_text(d)\n",
    "    docs.extend(splits)\n",
    "    metadatas.extend([{\"source\": sources[i]}] * len(splits))\n",
    "\n",
    "# you only need to do this once, in the future, when re-run this notebook, skip to below and load the vector store from disk\n",
    "store = FAISS.from_texts(docs, embedder , metadatas=metadatas)\n",
    "store.save_local('./data/save_embedding/sv')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8dd36e06",
   "metadata": {},
   "outputs": [],
   "source": [
    "## If you previously preprocessed and saved the vector store to disk, then reload it here\n",
    "faissDB = FAISS.load_local(\"./data/save_embedding/sv\", embedder, allow_dangerous_deserialization=True)\n",
    "retriever = faissDB.as_retriever()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8ecdf18",
   "metadata": {},
   "source": [
    "## Step 4 - Construct a Retriever for Sweden data\n",
    "\n",
    "The following code creates a `SwedenRetriever` class that inherits from LangChain's `BaseTool` class.\n",
    "\n",
    "We'll use the class as a tool for retrieving data about Sweden to augment responses."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75e46892",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.tools import BaseTool\n",
    "\n",
    "class SwedenRetriever(BaseTool):\n",
    "    name = \"AboutSweden\"\n",
    "    description = \"Useful for when you need to answer questions about Sweden's population, history, and so on.\"\n",
    "\n",
    "    def _run(self, query):\n",
    "        out = retriever.invoke(query)\n",
    "        o = out[0]\n",
    "        item=o.page_content.split('|')\n",
    "        output = '\\n'.join(item)\n",
    "        return output\n",
    "\n",
    "    def _arun(self, query: str):\n",
    "        raise NotImplementedError(\"This tool does not support async\")\n",
    "sv=SwedenRetriever()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b90845b7",
   "metadata": {},
   "source": [
    "## Step 5 - Construct wikipedia as the second tool\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fc872c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.tools.wikipedia.tool import WikipediaQueryRun\n",
    "from langchain_community.utilities import WikipediaAPIWrapper\n",
    "\n",
    "wikipedia = WikipediaQueryRun(api_wrapper=WikipediaAPIWrapper())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20c4c3f9",
   "metadata": {},
   "source": [
    "## Step 6 - Give your tools a good name and populate the description "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab96fe6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.tools import Tool\n",
    "\n",
    "## Make sure you give it a proper name and a good description on how to use the tools\n",
    "wiki_tool = Tool.from_function(\n",
    "    func=wikipedia.run,\n",
    "    name=\"Wiki\",\n",
    "    description=\"useful for when you need to search certain topic on Wikipedia, aka wiki\")\n",
    "retriever_tool=Tool.from_function(\n",
    "    func=sv.invoke,\n",
    "    name=\"AboutSweden\",\n",
    "    description=\"useful for when you need to find information about Sweden\")\n",
    "\n",
    "tools = [wiki_tool, retriever_tool]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11a66119",
   "metadata": {},
   "source": [
    "## Step 7 - Wrap tools into ToolExecutor \n",
    "\n",
    "We will use these ToolExecutor to invoke tool in LangGraph nodes later on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6426610c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.agents import AgentFinish\n",
    "from langgraph.prebuilt.tool_executor import ToolExecutor\n",
    "\n",
    "# This a helper class we have that is useful for running tools\n",
    "# It takes in an agent action and calls that tool and returns the result\n",
    "tool_executor = ToolExecutor(tools)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d97af339",
   "metadata": {},
   "source": [
    "## Step 8 - Create the prompt template and conversation memory\n",
    "\n",
    "The following code creates a memory buffer for storing queries and responses.\n",
    "It also demonstrates how to write a prompt template for a Mistral mode that uses conversation memory and the Wiki and retriever tools."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cfac120",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.memory import ConversationBufferMemory\n",
    "from langchain_core.prompts import PromptTemplate\n",
    "\n",
    "## set up memory\n",
    "memory = ConversationBufferMemory(memory_key=\"chat_history\", input_key='input', output_key=\"output\")\n",
    "\n",
    "\n",
    "prompt_template = \"\"\"\n",
    "### [INST]\n",
    "\n",
    "Assistant is a large language model trained by Mistral.\n",
    "\n",
    "Assistant is designed to be able to assist with a wide range of tasks, from answering simple questions to providing in-depth explanations and discussions on a wide range of topics. As a language model, Assistant is able to generate human-like text based on the input it receives, allowing it to engage in natural-sounding conversations and provide responses that are coherent and relevant to the topic at hand.\n",
    "\n",
    "Assistant is constantly learning and improving, and its capabilities are constantly evolving. It is able to process and understand large amounts of text, and can use this knowledge to provide accurate and informative responses to a wide range of questions. Additionally, Assistant is able to generate its own text based on the input it receives, allowing it to engage in discussions and provide explanations and descriptions on a wide range of topics.\n",
    "\n",
    "Overall, Assistant is a powerful tool that can help with a wide range of tasks and provide valuable insights and information on a wide range of topics. Whether you need help with a specific question or just want to have a conversation about a particular topic, Assistant is here to assist.\n",
    "\n",
    "Context:\n",
    "------\n",
    "\n",
    "Assistant has access to the following tools:\n",
    "\n",
    "{tools}\n",
    "\n",
    "To use a tool, please use the following format:\n",
    "\n",
    "'''\n",
    "Thought: Do I need to use a tool? Yes\n",
    "Action: the action to take, should be one of [{tool_names}]\n",
    "Action Input: the input to the action\n",
    "Observation: the result of the action\n",
    "'''\n",
    "\n",
    "When you have a response to say to the Human, or if you do not need to use a tool, you MUST use the format:\n",
    "\n",
    "'''\n",
    "Thought: Do I need to use a tool? No\n",
    "Final Answer: [your response here]\n",
    "'''\n",
    "\n",
    "Begin!\n",
    "\n",
    "Previous conversation history:\n",
    "{chat_history}\n",
    "\n",
    "New input: {input}\n",
    "\n",
    "Current Scratchpad:\n",
    "{agent_scratchpad}\n",
    "\n",
    "[/INST]\n",
    " \"\"\"\n",
    "\n",
    "# Create prompt from prompt template\n",
    "prompt = PromptTemplate(\n",
    "    input_variables=['agent_scratchpad', 'chat_history', 'input', 'tool_names', 'tools'],\n",
    "    template=prompt_template,\n",
    ")\n",
    "\n",
    "prompt = prompt.partial(\n",
    "    tools=[t.name for t in tools],\n",
    "    tool_names=\", \".join([t.name for t in tools]),\n",
    ")\n",
    "print(\"prompt ---> \\n\", prompt)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a217e64",
   "metadata": {},
   "source": [
    "## Step 9 - Establish agent executor using LangChain  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33c9f5f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.agents.agent import AgentExecutor\n",
    "from langchain.agents.agent_types import AgentType\n",
    "from langchain.agents.loading import AGENT_TO_CLASS\n",
    "\n",
    "agent_cls = AGENT_TO_CLASS[AgentType.CONVERSATIONAL_REACT_DESCRIPTION]\n",
    "agent_kwargs = {}\n",
    "agent_obj = agent_cls.from_llm_and_tools(\n",
    "    llm, tools, callback_manager=None, **agent_kwargs)\n",
    "\n",
    "agent_execute=AgentExecutor.from_agent_and_tools(\n",
    "        agent=agent_obj,\n",
    "        tools=tools,\n",
    "        callback_manager=None,\n",
    "        handle_parsing_errors=True,\n",
    "        verbose=True,\n",
    "        output_key = \"output\",\n",
    "        max_iterations=3,\n",
    "        return_intermediate_steps=True,\n",
    "        early_stopping_method=\"generate\", # or use **force**\n",
    "        memory=ConversationBufferMemory(memory_key=\"chat_history\", input_key='input', output_key=\"output\")\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2448f7b9",
   "metadata": {},
   "source": [
    "## Step 10 - Define the graph state\n",
    "\n",
    "We now define the graph state. The state for the traditional LangChain agent has a few attributes:\n",
    "\n",
    "1. `input`: This is the input string representing the main ask from the user, passed in as input.\n",
    "2. `chat_history`: This is any previous conversation messages, also passed in as input.\n",
    "3. `intermediate_steps`: This is list of actions and corresponding observations that the agent takes over time. This is updated each iteration of the agent.\n",
    "4. `agent_outcome`: This is the response from the agent, either an AgentAction or AgentFinish. The AgentExecutor should finish when this is an AgentFinish, otherwise it should call the requested tools.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "788076e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import TypedDict, Annotated, Union\n",
    "from langchain_core.agents import AgentAction, AgentFinish\n",
    "from langchain_core.messages import BaseMessage\n",
    "import operator\n",
    "\n",
    "\n",
    "class AgentState(TypedDict):\n",
    "    # The input string\n",
    "    input: str\n",
    "    # The list of previous messages in the conversation\n",
    "    chat_history: list[BaseMessage]\n",
    "    # The outcome of a given call to the agent\n",
    "    # Needs `None` as a valid type, since this is what this will start as\n",
    "    agent_outcome: Union[AgentAction, AgentFinish, None]\n",
    "    # List of actions and corresponding observations\n",
    "    # Here we annotate this with `operator.add` to indicate that operations to\n",
    "    # this state should be ADDED to the existing values (not overwrite it)\n",
    "    intermediate_steps: Annotated[list[tuple[AgentAction, str]], operator.add]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "191a53ed",
   "metadata": {},
   "source": [
    "## Step 11 - Define the nodes\n",
    "\n",
    "We now need to define a few different nodes in our graph.\n",
    "In LangGraph, a node can be either a function or a [runnable](https://python.langchain.com/docs/expression_language/).\n",
    "There are two main nodes we need for this:\n",
    "\n",
    "1. The agent (`run_agent`): responsible for deciding what (if any) actions to take.\n",
    "2. A function to invoke tools (`execute_tools`): if the agent decides to take an action, this node will then execute that action.\n",
    "\n",
    "We will also need to define some edges.\n",
    "Some of these edges may be conditional.\n",
    "The reason they are conditional is that based on the output of a node, one of several paths may be taken.\n",
    "The path that is taken is not known until that node is run (the LLM decides).\n",
    "\n",
    "1. Conditional Edge (`should_continue`): after the agent is called, we should either:\n",
    "   - If the agent said to take an action, then the function to invoke tools is called.\n",
    "   - If the agent said that it was finished, then it finishes.\n",
    "   \n",
    "2. Normal Edge: after the tools are invoked, it should always go back to the agent to decide what to do next.\n",
    "\n",
    "Let's define the nodes, as well as a function to decide how what conditional edge to take."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77747748",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the agent\n",
    "from langchain_core.agents import AgentActionMessageLog\n",
    "\n",
    "def run_agent(data):\n",
    "    inputs = data.copy()\n",
    "    text = inputs['input']\n",
    "    agent_outcome = agent_execute.invoke({\"input\":text})\n",
    "    return {\"agent_outcome\": agent_outcome}\n",
    "\n",
    "# Define the function to execute tools\n",
    "def execute_tools(data):\n",
    "    # Get the most recent agent_outcome - this is the key added in the `agent` above\n",
    "    agent_output = data[\"agent_outcome\"]\n",
    "    if len(agent_output['intermediate_steps'])>=1 :\n",
    "        agent_action = agent_output['intermediate_steps'][0][0]\n",
    "        output = tool_executor.invoke(agent_action)\n",
    "        return {\"intermediate_steps\": [(agent_action, str(output))]}\n",
    "    else:\n",
    "        return {\"intermediate_steps\":[]}\n",
    "\n",
    "# Define logic that is used to determine which conditional edge to go down\n",
    "def should_continue(data):\n",
    "    # If the agent outcome is an AgentFinish, then we return `exit` string\n",
    "    # This will be used when setting up the graph to define the flow\n",
    "    if data[\"agent_outcome\"][\"output\"] is not None:\n",
    "        print(\" **AgentFinish** \" )\n",
    "        return \"end\"\n",
    "    # Otherwise, an AgentAction is returned\n",
    "    # Here we return `continue` string\n",
    "    # This will be used when setting up the graph to define the flow\n",
    "    else:\n",
    "        print(\" **continue** \" )\n",
    "        return \"continue\"\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6fdb848d",
   "metadata": {},
   "source": [
    "## Step 12 - Connect the nodes with edges to form the graph, let's call it **app**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3962cf7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.graph import END, StateGraph\n",
    "\n",
    "# Define a new graph\n",
    "workflow = StateGraph(AgentState)\n",
    "\n",
    "# Define the two nodes we will cycle between\n",
    "workflow.add_node(\"agent\", run_agent)\n",
    "workflow.add_node(\"action\", execute_tools)\n",
    "\n",
    "# Set the entrypoint as `agent`\n",
    "# This means that this node is the first one called\n",
    "workflow.set_entry_point(\"agent\")\n",
    "\n",
    "# We now add a conditional edge\n",
    "workflow.add_conditional_edges(\n",
    "    # First, we define the start node. We use `agent`.\n",
    "    # This means these are the edges taken after the `agent` node is called.\n",
    "    \"agent\",\n",
    "    # Next, we pass in the function that will determine which node is called next.\n",
    "    should_continue,\n",
    "    # Finally we pass in a mapping.\n",
    "    # The keys are strings, and the values are other nodes.\n",
    "    # END is a special node marking that the graph should finish.\n",
    "    # What will happen is we will call `should_continue`, and then the output of that\n",
    "    # will be matched against the keys in this mapping.\n",
    "    # Based on which one it matches, that node will then be called.\n",
    "    {\n",
    "        # If `tools`, then we call the tool node.\n",
    "        \"continue\": \"action\",\n",
    "        # Otherwise we finish.\n",
    "        \"end\": END,\n",
    "    },\n",
    ")\n",
    "\n",
    "# We now add a normal edge from `tools` to `agent`.\n",
    "# This means that after `tools` is called, `agent` node is called next.\n",
    "workflow.add_edge(\"action\", \"agent\")\n",
    "\n",
    "# Finally, we compile it!\n",
    "# This compiles it into a LangChain Runnable,\n",
    "# meaning you can use it as you would any other runnable\n",
    "app = workflow.compile()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc5420b9",
   "metadata": {},
   "source": [
    "## Step 13 - Time to test it out\n",
    "\n",
    "Let's start by seeing if we can trigger the retriever tool (tool name: AboutSweden).\n",
    "\n",
    "\n",
    "Then, we will try to call the Wikipedia tool (tool name : Wiki)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34d7b094",
   "metadata": {},
   "outputs": [],
   "source": [
    "## first let's see if we can trigger our custom retriever tool named : AboutSweden\n",
    "\n",
    "inputs = {\"input\": \"What is Sweden's population?\"}\n",
    "outputs = app.invoke(inputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d94eb98",
   "metadata": {},
   "outputs": [],
   "source": [
    "## let's see if we can trigger our Wikipedia tool named : Wiki\n",
    "\n",
    "inputs = {\"input\": \"Find me Taylor Swift information on wiki?\"}\n",
    "outputs=app.invoke(inputs)"
   ]
  }
 ],
 "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.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}