Custom agent with tool retrieval.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "z0FJYcHSsHD8"
   },
   "outputs": [],
   "source": [
    "from langchain.agents import (\n",
    "    Tool,\n",
    "    AgentExecutor,\n",
    "    LLMSingleActionAgent,\n",
    "    AgentOutputParser,\n",
    ")\n",
    "from langchain.prompts import StringPromptTemplate\n",
    "from langchain import OpenAI, SerpAPIWrapper, LLMChain\n",
    "from typing import List, Union\n",
    "from langchain.schema import AgentAction, AgentFinish\n",
    "import re"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "L_--sBxusVFY"
   },
   "outputs": [],
   "source": [
    "# Define which tools the agent can use to answer user queries\n",
    "search = SerpAPIWrapper()\n",
    "search_tool = Tool(\n",
    "    name=\"Search\",\n",
    "    func=search.run,\n",
    "    description=\"useful for when you need to answer questions about current events\",\n",
    ")\n",
    "\n",
    "\n",
    "def fake_func(inp: str) -> str:\n",
    "    return \"foo\"\n",
    "\n",
    "\n",
    "fake_tools = [\n",
    "    Tool(\n",
    "        name=f\"foo-{i}\",\n",
    "        func=fake_func,\n",
    "        description=f\"a silly function that you can use to get more information about the number {i}\",\n",
    "    )\n",
    "    for i in range(99)\n",
    "]\n",
    "ALL_TOOLS = [search_tool] + fake_tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "Ucv8zcLvsWoN"
   },
   "outputs": [],
   "source": [
    "from langchain.vectorstores import FAISS\n",
    "from langchain.embeddings import OpenAIEmbeddings\n",
    "from langchain.schema import Document"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "d3neMbYnsYBk"
   },
   "outputs": [],
   "source": [
    "docs = [\n",
    "    Document(page_content=t.description, metadata={\"index\": i})\n",
    "    for i, t in enumerate(ALL_TOOLS)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "K9qj75q0sZZv"
   },
   "outputs": [],
   "source": [
    "vector_store = FAISS.from_documents(docs, OpenAIEmbeddings())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "eq7r_FtKsb1d"
   },
   "outputs": [],
   "source": [
    "retriever = vector_store.as_retriever()\n",
    "\n",
    "\n",
    "def get_tools(query):\n",
    "    docs = retriever.get_relevant_documents(query)\n",
    "    return [ALL_TOOLS[d.metadata[\"index\"]] for d in docs]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "aXIPDKC3sdT3"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Tool(name='Search', description='useful for when you need to answer questions about current events', args_schema=None, return_direct=False, verbose=False, callbacks=None, callback_manager=None, tags=None, metadata=None, handle_tool_error=False, func=<bound method SerpAPIWrapper.run of SerpAPIWrapper(search_engine=<class 'serpapi.google_search.GoogleSearch'>, params={'engine': 'google', 'google_domain': 'google.com', 'gl': 'us', 'hl': 'en'}, serpapi_api_key='f70b083fab9292ba02681ffb67fcb93f30575ff9d0de93478e5c8a0a9d34a67a', aiosession=None)>, coroutine=None),\n",
       " Tool(name='foo-95', description='a silly function that you can use to get more information about the number 95', args_schema=None, return_direct=False, verbose=False, callbacks=None, callback_manager=None, tags=None, metadata=None, handle_tool_error=False, func=<function fake_func at 0x7fda9a336e60>, coroutine=None),\n",
       " Tool(name='foo-12', description='a silly function that you can use to get more information about the number 12', args_schema=None, return_direct=False, verbose=False, callbacks=None, callback_manager=None, tags=None, metadata=None, handle_tool_error=False, func=<function fake_func at 0x7fda9a336e60>, coroutine=None),\n",
       " Tool(name='foo-15', description='a silly function that you can use to get more information about the number 15', args_schema=None, return_direct=False, verbose=False, callbacks=None, callback_manager=None, tags=None, metadata=None, handle_tool_error=False, func=<function fake_func at 0x7fda9a336e60>, coroutine=None)]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "get_tools(\"whats the weather?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "IaBUDwYTsi6F"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Tool(name='foo-13', description='a silly function that you can use to get more information about the number 13', args_schema=None, return_direct=False, verbose=False, callbacks=None, callback_manager=None, tags=None, metadata=None, handle_tool_error=False, func=<function fake_func at 0x7fda9a336e60>, coroutine=None),\n",
       " Tool(name='foo-12', description='a silly function that you can use to get more information about the number 12', args_schema=None, return_direct=False, verbose=False, callbacks=None, callback_manager=None, tags=None, metadata=None, handle_tool_error=False, func=<function fake_func at 0x7fda9a336e60>, coroutine=None),\n",
       " Tool(name='foo-14', description='a silly function that you can use to get more information about the number 14', args_schema=None, return_direct=False, verbose=False, callbacks=None, callback_manager=None, tags=None, metadata=None, handle_tool_error=False, func=<function fake_func at 0x7fda9a336e60>, coroutine=None),\n",
       " Tool(name='foo-11', description='a silly function that you can use to get more information about the number 11', args_schema=None, return_direct=False, verbose=False, callbacks=None, callback_manager=None, tags=None, metadata=None, handle_tool_error=False, func=<function fake_func at 0x7fda9a336e60>, coroutine=None)]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "get_tools(\"whats the number 13?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "id": "cb4akruQsnG_"
   },
   "outputs": [],
   "source": [
    "# Set up the base template\n",
    "template = \"\"\"Answer the following questions as best you can, but speaking as a pirate might speak. You have access to the following tools:\n",
    "\n",
    "{tools}\n",
    "\n",
    "Use the following format:\n",
    "\n",
    "Question: the input question you must answer\n",
    "Thought: you should always think about what to do\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",
    "... (this Thought/Action/Action Input/Observation can repeat N times)\n",
    "Thought: I now know the final answer\n",
    "Final Answer: the final answer to the original input question\n",
    "\n",
    "Begin! Remember to speak as a pirate when giving your final answer. Use lots of \"Arg\"s\n",
    "\n",
    "Question: {input}\n",
    "{agent_scratchpad}\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "id": "tnLeG4EBsowP"
   },
   "outputs": [],
   "source": [
    "from typing import Callable\n",
    "\n",
    "\n",
    "# Set up a prompt template\n",
    "class CustomPromptTemplate(StringPromptTemplate):\n",
    "    # The template to use\n",
    "    template: str\n",
    "    ############## NEW ######################\n",
    "    # The list of tools available\n",
    "    tools_getter: Callable\n",
    "\n",
    "    def format(self, **kwargs) -> str:\n",
    "        # Get the intermediate steps (AgentAction, Observation tuples)\n",
    "        # Format them in a particular way\n",
    "        intermediate_steps = kwargs.pop(\"intermediate_steps\")\n",
    "        thoughts = \"\"\n",
    "        for action, observation in intermediate_steps:\n",
    "            thoughts += action.log\n",
    "            thoughts += f\"\\nObservation: {observation}\\nThought: \"\n",
    "        # Set the agent_scratchpad variable to that value\n",
    "        kwargs[\"agent_scratchpad\"] = thoughts\n",
    "        ############## NEW ######################\n",
    "        tools = self.tools_getter(kwargs[\"input\"])\n",
    "        # Create a tools variable from the list of tools provided\n",
    "        kwargs[\"tools\"] = \"\\n\".join(\n",
    "            [f\"{tool.name}: {tool.description}\" for tool in tools]\n",
    "        )\n",
    "        # Create a list of tool names for the tools provided\n",
    "        kwargs[\"tool_names\"] = \", \".join([tool.name for tool in tools])\n",
    "        return self.template.format(**kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "id": "Z3l7fOF_svBf"
   },
   "outputs": [],
   "source": [
    "prompt = CustomPromptTemplate(\n",
    "    template=template,\n",
    "    tools_getter=get_tools,\n",
    "    # This omits the `agent_scratchpad`, `tools`, and `tool_names` variables because those are generated dynamically\n",
    "    # This includes the `intermediate_steps` variable because that is needed\n",
    "    input_variables=[\"input\", \"intermediate_steps\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "id": "nO6RZwa4swI9"
   },
   "outputs": [],
   "source": [
    "class CustomOutputParser(AgentOutputParser):\n",
    "    def parse(self, llm_output: str) -> Union[AgentAction, AgentFinish]:\n",
    "        # Check if agent should finish\n",
    "        if \"Final Answer:\" in llm_output:\n",
    "            return AgentFinish(\n",
    "                # Return values is generally always a dictionary with a single `output` key\n",
    "                # It is not recommended to try anything else at the moment :)\n",
    "                return_values={\"output\": llm_output.split(\"Final Answer:\")[-1].strip()},\n",
    "                log=llm_output,\n",
    "            )\n",
    "        # Parse out the action and action input\n",
    "        regex = r\"Action\\s*\\d*\\s*:(.*?)\\nAction\\s*\\d*\\s*Input\\s*\\d*\\s*:[\\s]*(.*)\"\n",
    "        match = re.search(regex, llm_output, re.DOTALL)\n",
    "        if not match:\n",
    "            raise ValueError(f\"Could not parse LLM output: `{llm_output}`\")\n",
    "        action = match.group(1).strip()\n",
    "        action_input = match.group(2)\n",
    "        # Return the action and action input\n",
    "        return AgentAction(\n",
    "            tool=action, tool_input=action_input.strip(\" \").strip('\"'), log=llm_output\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "id": "lPe9hQdusxje"
   },
   "outputs": [],
   "source": [
    "output_parser = CustomOutputParser()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "id": "xa9sTpQIsyuC"
   },
   "outputs": [],
   "source": [
    "llm = OpenAI(temperature=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "id": "jOgcI8Lbszw9"
   },
   "outputs": [],
   "source": [
    "# LLM chain consisting of the LLM and a prompt\n",
    "llm_chain = LLMChain(llm=llm, prompt=prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "id": "PDcdJBVcs08l"
   },
   "outputs": [],
   "source": [
    "tools = get_tools(\"whats the weather?\")\n",
    "tool_names = [tool.name for tool in tools]\n",
    "agent = LLMSingleActionAgent(\n",
    "    llm_chain=llm_chain,\n",
    "    output_parser=output_parser,\n",
    "    stop=[\"\\nObservation:\"],\n",
    "    allowed_tools=tool_names,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "id": "YyambDSys2UA"
   },
   "outputs": [],
   "source": [
    "agent_executor = AgentExecutor.from_agent_and_tools(\n",
    "    agent=agent, tools=tools, verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "h8gl78URs356"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new  chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3m\n",
      "Thought: I need to use the search tool to find out the odds of New York winning the world series.\n",
      "Action: foo-19\n",
      "Action Input: New York winning the world series\u001b[0m\n",
      "\n",
      "Observation:foo-19 is not a valid tool, try another one.\n",
      "\u001b[32;1m\u001b[1;3m I need to find a tool that can give me the odds of New York winning the world series.\n",
      "Action: foo-10\n",
      "Action Input: New York winning the world series\u001b[0m\n",
      "\n",
      "Observation:foo-10 is not a valid tool, try another one.\n",
      "\u001b[32;1m\u001b[1;3m I need to find a tool that can give me the odds of New York winning the world series.\n",
      "Action: foo-1\n",
      "Action Input: New York winning the world series\u001b[0m\n",
      "\n",
      "Observation:foo-1 is not a valid tool, try another one.\n",
      "\u001b[32;1m\u001b[1;3m I need to find a tool that can give me the odds of New York winning the world series.\n",
      "Action: foo-9\n",
      "Action Input: New York winning the world series\u001b[0m\n",
      "\n",
      "Observation:foo-9 is not a valid tool, try another one.\n"
     ]
    }
   ],
   "source": [
    "agent_executor.run(\"Using the search tool, tell me what are the odds that New York wins the world series this year?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "authorship_tag": "ABX9TyOVHB1O9m2QviDk2uMu8GHW",
   "provenance": []
  },
  "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}