{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "feb31cc6",
   "metadata": {},
   "source": [
    "# Running Agent as an Iterator\n",
    "\n",
    "To demonstrate the `AgentExecutorIterator` functionality, we will set up a problem where an Agent must:\n",
    "\n",
    "- Retrieve three prime numbers from a Tool\n",
    "- Multiply these together. \n",
    "\n",
    "In this simple problem we can demonstrate adding some logic to verify intermediate steps by checking whether their outputs are prime."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8167db11",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import dotenv\n",
    "import pydantic\n",
    "from langchain.agents import AgentExecutor, initialize_agent, AgentType\n",
    "from langchain.schema import AgentFinish\n",
    "from langchain.agents.tools import Tool\n",
    "from langchain.chains import LLMMathChain\n",
    "from langchain.chat_models import ChatOpenAI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "7e41b9e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uncomment if you have a .env in root of repo contains OPENAI_API_KEY\n",
    "# dotenv.load_dotenv(\"../../../../../.env\")\n",
    "\n",
    "# need to use GPT-4 here as GPT-3.5 does not understand, however hard you insist, that\n",
    "# it should use the calculator to perform the final calculation\n",
    "llm = ChatOpenAI(temperature=0, model=\"gpt-4\")\n",
    "llm_math_chain = LLMMathChain.from_llm(llm=llm, verbose=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81e88aa5",
   "metadata": {},
   "source": [
    "Define tools which provide:\n",
    "- The `n`th prime number (using a small subset for this example) \n",
    "- The `LLMMathChain` to act as a calculator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "86f04b55",
   "metadata": {},
   "outputs": [],
   "source": [
    "primes = {998: 7901, 999: 7907, 1000: 7919}\n",
    "\n",
    "\n",
    "class CalculatorInput(pydantic.BaseModel):\n",
    "    question: str = pydantic.Field()\n",
    "\n",
    "\n",
    "class PrimeInput(pydantic.BaseModel):\n",
    "    n: int = pydantic.Field()\n",
    "\n",
    "\n",
    "def is_prime(n: int) -> bool:\n",
    "    if n <= 1 or (n % 2 == 0 and n > 2):\n",
    "        return False\n",
    "    for i in range(3, int(n**0.5) + 1, 2):\n",
    "        if n % i == 0:\n",
    "            return False\n",
    "    return True\n",
    "\n",
    "\n",
    "def get_prime(n: int, primes: dict = primes) -> str:\n",
    "    return str(primes.get(int(n)))\n",
    "\n",
    "\n",
    "async def aget_prime(n: int, primes: dict = primes) -> str:\n",
    "    return str(primes.get(int(n)))\n",
    "\n",
    "\n",
    "tools = [\n",
    "    Tool(\n",
    "        name=\"GetPrime\",\n",
    "        func=get_prime,\n",
    "        description=\"A tool that returns the `n`th prime number\",\n",
    "        args_schema=PrimeInput,\n",
    "        coroutine=aget_prime,\n",
    "    ),\n",
    "    Tool.from_function(\n",
    "        func=llm_math_chain.run,\n",
    "        name=\"Calculator\",\n",
    "        description=\"Useful for when you need to compute mathematical expressions\",\n",
    "        args_schema=CalculatorInput,\n",
    "        coroutine=llm_math_chain.arun,\n",
    "    ),\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e660ee6",
   "metadata": {},
   "source": [
    "Construct the agent. We will use the default agent type here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "21c775b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "agent = initialize_agent(\n",
    "    tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a233fe4e",
   "metadata": {},
   "source": [
    "Run the iteration and perform a custom check on certain steps:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "582d61f4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new  chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3mI need to find the 998th, 999th and 1000th prime numbers first.\n",
      "Action: GetPrime\n",
      "Action Input: 998\u001b[0m\n",
      "Observation: \u001b[36;1m\u001b[1;3m7901\u001b[0m\n",
      "Thought:Checking whether 7901 is prime...\n",
      "Should the agent continue (Y/n)?:\n",
      "Y\n",
      "\u001b[32;1m\u001b[1;3mI have the 998th prime number. Now I need to find the 999th prime number.\n",
      "Action: GetPrime\n",
      "Action Input: 999\u001b[0m\n",
      "Observation: \u001b[36;1m\u001b[1;3m7907\u001b[0m\n",
      "Thought:Checking whether 7907 is prime...\n",
      "Should the agent continue (Y/n)?:\n",
      "Y\n",
      "\u001b[32;1m\u001b[1;3mI have the 999th prime number. Now I need to find the 1000th prime number.\n",
      "Action: GetPrime\n",
      "Action Input: 1000\u001b[0m\n",
      "Observation: \u001b[36;1m\u001b[1;3m7919\u001b[0m\n",
      "Thought:Checking whether 7919 is prime...\n",
      "Should the agent continue (Y/n)?:\n",
      "Y\n",
      "\u001b[32;1m\u001b[1;3mI have all three prime numbers. Now I need to calculate the product of these numbers.\n",
      "Action: Calculator\n",
      "Action Input: 7901 * 7907 * 7919\u001b[0m\n",
      "\n",
      "\u001b[1m> Entering new  chain...\u001b[0m\n",
      "7901 * 7907 * 7919\u001b[32;1m\u001b[1;3m```text\n",
      "7901 * 7907 * 7919\n",
      "```\n",
      "...numexpr.evaluate(\"7901 * 7907 * 7919\")...\n",
      "\u001b[0m\n",
      "Answer: \u001b[33;1m\u001b[1;3m494725326233\u001b[0m\n",
      "\u001b[1m> Finished chain.\u001b[0m\n",
      "\n",
      "Observation: \u001b[33;1m\u001b[1;3mAnswer: 494725326233\u001b[0m\n",
      "Thought:Should the agent continue (Y/n)?:\n",
      "Y\n",
      "\u001b[32;1m\u001b[1;3mI now know the final answer\n",
      "Final Answer: 494725326233\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "question = \"What is the product of the 998th, 999th and 1000th prime numbers?\"\n",
    "\n",
    "for step in agent.iter(question):\n",
    "    if output := step.get(\"intermediate_step\"):\n",
    "        action, value = output[0]\n",
    "        if action.tool == \"GetPrime\":\n",
    "            print(f\"Checking whether {value} is prime...\")\n",
    "            assert is_prime(int(value))\n",
    "        # Ask user if they want to continue\n",
    "        _continue = input(\"Should the agent continue (Y/n)?:\\n\")\n",
    "        if _continue != \"Y\":\n",
    "            break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6934ff8e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "language": "python",
   "name": "venv"
  },
  "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}