chapter_04_ai_agents.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chapter 04: AI Agents\n",
    "\n",
    "We've explored the incredible capabilities of LLMs, the art of prompting, the integration of tools, and the generation of structured outputs. With this strong foundation, we're now ready to dive into something even more powerful - creating our very own LLM-powered AI agent. Here are just a few examples of what they can be used for: \n",
    "- **Customer support**: Provide instant, accurate, and empathetic responses to customer inquiries, improving satisfaction and efficiency. \n",
    "- **Content creation**: Generate high-quality articles, social media posts, and marketing copy with ease. \n",
    "- **Research assistance**: Summarize complex documents, find relevant studies, and provide insightful analysis. \n",
    "- **Personal assistants**: Manage schedules, set reminders, and help you stay organized. \n",
    "- **Creative expression**: Write stories, poems, and even act as a brainstorming partner for new ideas. \n",
    "- **Data analysis**: Process and interpret large datasets to extract meaningful patterns and insights.  \n",
    "\n",
    "With the ability to configure your LLM-powered AI agent for different scenarios and outputs, you unlock an incredible tool that can adapt to various roles and industries. The potential applications are vast, limited only by your imagination and creativity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import annotations\n",
    "import sys\n",
    "import tiktoken\n",
    "import requests\n",
    "from enum import Enum\n",
    "from typing import Any\n",
    "from pydantic import BaseModel, Field, ValidationError, create_model\n",
    "from datetime import datetime, timedelta\n",
    "from loguru import logger\n",
    "from language_models.agent.output_parser import (\n",
    "    FINAL_ANSWER_INSTRUCTIONS,\n",
    "    AgentOutputParser,\n",
    "    LLMFinalAnswer,\n",
    "    LLMToolUse,\n",
    "    OutputType,\n",
    "    get_schema_from_args\n",
    ")\n",
    "from language_models.models.llm import ChatMessage, ChatMessageRole, OpenAILanguageModel\n",
    "from language_models.tools import Tool, current_date\n",
    "from language_models.models.llm import ChatMessage\n",
    "from language_models.proxy_client import ProxyClient\n",
    "from language_models.settings import settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "logger.remove()\n",
    "logger.add(sys.stderr, format=\"{message}\", level=\"INFO\")\n",
    "\n",
    "proxy_client = ProxyClient(\n",
    "    client_id=settings.CLIENT_ID,\n",
    "    client_secret=settings.CLIENT_SECRET,\n",
    "    auth_url=settings.AUTH_URL,\n",
    "    api_base=settings.API_BASE,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = OpenAILanguageModel(\n",
    "    proxy_client=proxy_client,\n",
    "    model='gpt-4',\n",
    "    max_tokens=1000,\n",
    "    temperature=0.2,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we define instructions for an AI agent to follow when responding to different types of prompts:\n",
    "- **Single completion:** Provide a straightforward answer without using any tools. For example, if the prompt is \"What's the capital of France?\", respond with \"The capital of France is Paris.\"  \n",
    "- **Chain-of-Thought with tools:** Multi-step tasks that require tools. If the prompt is \"Find the current weather in NYC and suggest an outfit,\" first, query a weather API to get the current weather (for instance, \"22°C, partly cloudy\"). Then suggest an appropriate outfit based on the weather, such as \"Wear a t-shirt and a light jacket.\"  \n",
    "- **Chain-of-Thought without tools:** Tasks that require (complex) structured outputs. If you want error handling logic due to the possibility of mistakes by the LLM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "SINGLE_COMPLETION_INSTRUCTIONS = \"\"\"### Instructions ###\n",
    "\n",
    "Your goal is to solve the problem you will be provided with\n",
    "\n",
    "You should respond with:\n",
    "```\n",
    "<response to the prompt>\n",
    "```\"\"\"\n",
    "\n",
    "CHAIN_OF_THOUGHT_INSTRUCTIONS_WITH_TOOLS = \"\"\"### Tools ###\n",
    "\n",
    "You have access to the following tools:\n",
    "{tools}\n",
    "\n",
    "### Instructions ###\n",
    "\n",
    "Your goal is to solve the problem you will be provided with\n",
    "\n",
    "You should respond with:\n",
    "```\n",
    "Thought: <thought process on how to respond to the prompt>\n",
    "\n",
    "Tool: <name of the tool to use>\n",
    "\n",
    "Tool Input: <input of the tool to use>\n",
    "```\n",
    "\n",
    "Your <input of the tool to use> should be a JSON format and must be keyword arguments of the properties of <name of the tool to use>\n",
    "\n",
    "When you know the final answer to the user's query you should respond with:\n",
    "```\n",
    "Thought: <thought process on how to respond to the prompt>\n",
    "\n",
    "Final Answer: <response to the prompt>\n",
    "```\"\"\"\n",
    "\n",
    "CHAIN_OF_THOUGHT_INSTRUCTIONS_WITHOUT_TOOLS = \"\"\"### Instructions ###\n",
    "\n",
    "Your goal is to solve the problem you will be provided with\n",
    "\n",
    "You should respond with:\n",
    "```\n",
    "Thought: <thought process on how to respond to the prompt>\n",
    "\n",
    "Final Answer: <response to the prompt>\n",
    "```\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we define a `Chat` class that manages the chat history and a list of intermediate steps. Given the model's inability to recall previous conversations, we introduce a mechanism to track and provide previous steps to the model. This ensures that the model can make informed decisions on what to do next or when to provide the user with the final answer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Chat(BaseModel):\n",
    "    \"\"\"Class that implements the chat history.\n",
    "\n",
    "    Attributes:\n",
    "        messages: The conversation history.\n",
    "        steps: The intermediate steps.\n",
    "    \"\"\"\n",
    "\n",
    "    messages: list[ChatMessage]\n",
    "    previous_steps: list[str] = []\n",
    "\n",
    "    def update(self, prompt: str) -> None:\n",
    "        \"\"\"Modifies the user prompt to include intermediate steps.\"\"\"\n",
    "        self.messages[-1].content = \"\\n\\n\".join(\n",
    "            [prompt, f\"These were your previous steps:\\n{'\\n\\n'.join(self.previous_steps)}\"]\n",
    "        )\n",
    "\n",
    "    def reset(self) -> None:\n",
    "        \"\"\"Resets the chat.\"\"\"\n",
    "        self.messages = [self.messages[0]]\n",
    "        self.previous_steps = []"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We count the number of tokens in a conversation history to ensure it stays within the predefined limits for different models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_TOKEN_LIMIT = {\n",
    "    \"gpt-4\": 8192,\n",
    "    \"gpt-4-32k\": 32768,\n",
    "    \"gpt-35-turbo\": 4096,\n",
    "    \"gpt-35-turbo-16k\": 16385,\n",
    "}\n",
    "\n",
    "def num_tokens_from_messages(messages: list[ChatMessage]) -> int:\n",
    "    \"\"\"Counts the number of tokens in the conversation history.\"\"\"\n",
    "    encoding = tiktoken.get_encoding(\"cl100k_base\")\n",
    "    tokens_per_message = 3\n",
    "    tokens_per_name = 1\n",
    "    num_tokens = 0\n",
    "    for message in messages:\n",
    "        num_tokens += tokens_per_message\n",
    "        for key, value in message.model_dump().items():\n",
    "            num_tokens += len(encoding.encode(value))\n",
    "            if key == \"name\":\n",
    "                num_tokens += tokens_per_name\n",
    "    num_tokens += 3\n",
    "    return num_tokens"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we create an `Agent` to encapsulate an LLM, providing methods for managing conversation context, parsing outputs, and executing the agent logic. \n",
    "\n",
    "It's important to recognize that LLMs can sometimes produce errors in their output, such as generating incorrect output structure or invalid tool input arguments. To address this in Chain-of-Thought prompting, it is crucial to highlight these errors to the LLM in hopes that it can correct them. However, the LLM may not always be able to fix its own mistakes. In such instances, after several iterations of feedback and corrections, we may need to provide a fallback response. This iterative process, combined with structured output formats and tool usage, forms the basis of an LLM-powered AI agent. While it may seem like magic in demonstrations, the underlying mechanism is actually quite straightforward.\n",
    "\n",
    "For single completion, we perform only one iteration and present the user with the final answer. If parsing fails when requesting structured output, we also provide a fallback response to ensure clarity.\n",
    "\n",
    "![ReAct prompting](./assets/images/react.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PromptingStrategy(str, Enum):\n",
    "    SINGLE_COMPLETION = \"single_completion\"\n",
    "    CHAIN_OF_THOUGHT = \"chain_of_thought\"\n",
    "\n",
    "class AgentOutput(BaseModel):\n",
    "    \"\"\"Class that represents the agent output.\"\"\"\n",
    "\n",
    "    prompt: str\n",
    "    final_answer: (\n",
    "        str\n",
    "        | int\n",
    "        | float\n",
    "        | dict[str, Any]\n",
    "        | BaseModel\n",
    "        | list[str]\n",
    "        | list[int]\n",
    "        | list[float]\n",
    "        | list[dict[str, Any]]\n",
    "        | list[BaseModel]\n",
    "        | None\n",
    "    )\n",
    "\n",
    "class Agent(BaseModel):\n",
    "    \"\"\"Class that implements an LLM-powered AI agent.\n",
    "\n",
    "    Attributes:\n",
    "        llm: The OpenAI LLM to use (in practice any model you choose).\n",
    "        tools: The tools the LLM can use.\n",
    "        prompt: The task prompt the LLM should solve.\n",
    "        prompt_variables: The task prompt variables that are needed to run the agent.\n",
    "        output_parser: The parser that handles LLM responses.\n",
    "        chat: The chat history.\n",
    "        prompt_strategy: The prompting strategy to use\n",
    "            (single completion for input/output type of queries, chain-of-thought for multi-step queries that involve tools).\n",
    "        iterations: The number of steps the LLM can take to solve the user query.\n",
    "        verbose: Enable logging.\n",
    "    \"\"\"\n",
    "\n",
    "    llm: OpenAILanguageModel\n",
    "    tools: dict[str, Tool] | None\n",
    "    prompt: str\n",
    "    prompt_variables: list[str]\n",
    "    output_parser: AgentOutputParser\n",
    "    chat: Chat\n",
    "    prompting_strategy: PromptingStrategy = PromptingStrategy.CHAIN_OF_THOUGHT\n",
    "    iterations: int = 5\n",
    "    verbose: bool = True\n",
    "\n",
    "    def trim_conversation(self) -> None:\n",
    "        \"\"\"Trims the chat messages to fit the LLM context length.\"\"\"\n",
    "        num_tokens = num_tokens_from_messages(self.chat.messages)\n",
    "        while num_tokens + self.llm.max_tokens >= MODEL_TOKEN_LIMIT[self.llm.model]:\n",
    "            del self.chat.messages[1]\n",
    "            num_tokens = num_tokens_from_messages(self.chat.messages)\n",
    "\n",
    "    def parse_output(self, output: str) -> LLMToolUse | LLMFinalAnswer:\n",
    "        \"\"\"Parses the LLM output.\"\"\"\n",
    "        try:\n",
    "            output = self.output_parser.parse(output)\n",
    "            observation = None\n",
    "        except (ValueError, ValidationError) as error:\n",
    "            output = None\n",
    "            observation = error\n",
    "        return output, observation\n",
    "\n",
    "    def invoke(self, prompt: dict[str, Any]) -> AgentOutput:\n",
    "        \"\"\"Runs the agent given a prompt.\"\"\"\n",
    "        prompt = self.prompt.format(**{variable: prompt.get(variable) for variable in self.prompt_variables})\n",
    "\n",
    "        self.chat.messages.append(ChatMessage(role=ChatMessageRole.USER, content=prompt))\n",
    "        self.chat.previous_steps = []\n",
    "\n",
    "        iteration = 0\n",
    "        while iteration <= self.iterations:\n",
    "            self.trim_conversation()\n",
    "            output = self.llm.get_completion(self.chat.messages)\n",
    "            output, observation = self.parse_output(output)\n",
    "            if self.prompting_strategy == PromptingStrategy.CHAIN_OF_THOUGHT:\n",
    "                if output is not None:\n",
    "                    if self.verbose:\n",
    "                        logger.opt(colors=True).info(f\"<b><fg #2D72D2>Thought</fg #2D72D2></b>: {output.thought}\")\n",
    "\n",
    "                    self.chat.previous_steps.append(f\"Thought: {output.thought}\")\n",
    "                    if isinstance(output, LLMFinalAnswer):\n",
    "                        if self.verbose:\n",
    "                            logger.opt(colors=True).success(\n",
    "                                f\"<b><fg #32A467>Final Answer</fg #32A467></b>: {output.final_answer}\"\n",
    "                            )\n",
    "\n",
    "                        self.chat.messages.append(\n",
    "                            ChatMessage(role=ChatMessageRole.ASSISTANT, content=str(output.final_answer))\n",
    "                        )\n",
    "                        return AgentOutput(prompt=prompt, final_answer=output.final_answer)\n",
    "                    else:\n",
    "                        if self.tools is not None:\n",
    "                            if self.verbose:\n",
    "                                logger.opt(colors=True).info(f\"<b><fg #EC9A3C>Tool</fg #EC9A3C></b>: {output.tool}\")\n",
    "                                logger.opt(colors=True).info(\n",
    "                                    f\"<b><fg #EC9A3C>Tool Input</fg #EC9A3C></b>: {output.tool_input}\"\n",
    "                                )\n",
    "\n",
    "                            tool = self.tools.get(output.tool)\n",
    "                            if tool is not None:\n",
    "                                self.chat.previous_steps.append(f\"Tool: {output.tool}\")\n",
    "                                self.chat.previous_steps.append(f\"Tool Input: {output.tool_input}\")\n",
    "                                tool_output = tool.invoke(output.tool_input)\n",
    "                                observation = f\"Tool Output: {tool_output}\"\n",
    "                                if self.verbose:\n",
    "                                    logger.opt(colors=True).info(\n",
    "                                        f\"<b><fg #EC9A3C>Tool Output</fg #EC9A3C></b>: {tool_output}\"\n",
    "                                    )\n",
    "                            else:\n",
    "                                tool_names = \", \".join(list(self.tools.keys()))\n",
    "                                observation = f\"{output.tool} tool doesn't exist. Try one of these tools: {tool_names}\"\n",
    "\n",
    "                self.chat.previous_steps.append(f\"Observation: {observation}\")\n",
    "                self.chat.update(prompt)\n",
    "            else:\n",
    "                if isinstance(output, LLMFinalAnswer):\n",
    "                    if self.verbose:\n",
    "                        logger.opt(colors=True).success(\n",
    "                            f\"<b><fg #32A467>Final Answer</fg #32A467></b>: {output.final_answer}\"\n",
    "                        )\n",
    "\n",
    "                    self.chat.messages.append(\n",
    "                        ChatMessage(role=ChatMessageRole.ASSISTANT, content=str(output.final_answer))\n",
    "                    )\n",
    "                    return AgentOutput(prompt=prompt, final_answer=output.final_answer)\n",
    "\n",
    "            iteration += 1\n",
    "\n",
    "        if self.output_parser.output_type == OutputType.STRUCT:\n",
    "            final_answer = {key: None for key in self.output_parser.output_schema.model_json_schema()[\"properties\"]}\n",
    "        elif self.output_parser.output_type == OutputType.ARRAY_STRUCT:\n",
    "            final_answer = [{key: None for key in self.output_parser.output_schema.model_json_schema()[\"properties\"]}]\n",
    "        elif self.output_parser.output_type in (OutputType.OBJECT, OutputType.ARRAY_OBJECT):\n",
    "            fields = self.output_parser.output_schema.__annotations__\n",
    "            optional_fields = {field: (data_type | None, None) for field, data_type in fields.items()}\n",
    "            model = create_model(self.output_parser.output_schema.__name__, **optional_fields)\n",
    "            final_answer = model() if self.output_parser.output_type == OutputType.OBJECT else [model()]\n",
    "        else:\n",
    "            final_answer = None\n",
    "\n",
    "        if self.verbose:\n",
    "            logger.opt(colors=True).warning(f\"<b><fg #CD4246>Final Answer</fg #CD4246></b>: {final_answer}\")\n",
    "\n",
    "        return AgentOutput(prompt=prompt, final_answer=final_answer)\n",
    "\n",
    "    @classmethod\n",
    "    def create(\n",
    "        cls,\n",
    "        llm: OpenAILanguageModel,\n",
    "        system_prompt: str,\n",
    "        prompt: str,\n",
    "        prompt_variables: list[str],\n",
    "        output_type: OutputType,\n",
    "        output_schema: type[BaseModel] | str | None = None,\n",
    "        tools: list[Tool] | None = None,\n",
    "        prompting_strategy: PromptingStrategy = PromptingStrategy.CHAIN_OF_THOUGHT,\n",
    "        verbose: bool = True,\n",
    "    ) -> Agent:\n",
    "        \"\"\"Creates an instance of the Agent.\"\"\"\n",
    "        if prompting_strategy == PromptingStrategy.CHAIN_OF_THOUGHT:\n",
    "            if tools is None:\n",
    "                instructions = CHAIN_OF_THOUGHT_INSTRUCTIONS_WITHOUT_TOOLS\n",
    "                tool_use = False\n",
    "                tools = None\n",
    "                iterations = 5\n",
    "            else:\n",
    "                instructions = CHAIN_OF_THOUGHT_INSTRUCTIONS_WITH_TOOLS.format(\n",
    "                    tools=\"\\n\\n\".join([str(tool) for tool in tools])\n",
    "                )\n",
    "                tool_use = True\n",
    "                tools = {tool.name: tool for tool in tools}\n",
    "                iterations = max(5, len(tools) * 2)\n",
    "        else:\n",
    "            instructions = SINGLE_COMPLETION_INSTRUCTIONS\n",
    "            tool_use = False\n",
    "            iterations = 1\n",
    "\n",
    "        if output_type in (OutputType.OBJECT, OutputType.ARRAY_OBJECT):\n",
    "            if output_schema is None:\n",
    "                raise ValueError(f\"When using {output_type} as the output type a schema must be provided.\")\n",
    "\n",
    "            schema = output_schema.model_json_schema()\n",
    "            if \"$defs\" not in schema:\n",
    "                schema = get_schema_from_args(schema[\"properties\"])\n",
    "            final_answer_instructions = FINAL_ANSWER_INSTRUCTIONS[output_type].format(output_schema=schema)\n",
    "        elif output_type in (OutputType.DATE, OutputType.TIMESTAMP):\n",
    "            final_answer_instructions = FINAL_ANSWER_INSTRUCTIONS[output_type].format(output_schema=output_schema)\n",
    "        else:\n",
    "            final_answer_instructions = FINAL_ANSWER_INSTRUCTIONS[output_type]\n",
    "\n",
    "        chat = Chat(\n",
    "            messages=[\n",
    "                ChatMessage(\n",
    "                    role=ChatMessageRole.SYSTEM,\n",
    "                    content=\"\\n\\n\".join([system_prompt, instructions, final_answer_instructions]),\n",
    "                )\n",
    "            ]\n",
    "        )\n",
    "\n",
    "        output_parser = AgentOutputParser(\n",
    "            output_type=output_type,\n",
    "            output_schema=output_schema,\n",
    "            prompting_strategy=prompting_strategy,\n",
    "            tool_use=tool_use,\n",
    "        )\n",
    "\n",
    "        return Agent(\n",
    "            llm=llm,\n",
    "            tools=tools,\n",
    "            prompt=prompt,\n",
    "            prompt_variables=prompt_variables,\n",
    "            output_parser=output_parser,\n",
    "            chat=chat,\n",
    "            prompting_strategy=prompting_strategy,\n",
    "            iterations=iterations,\n",
    "            verbose=verbose,\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Single Completion\n",
    "\n",
    "The following code creates an AI agent that can answer user questions based on the internal knowledge of the LLM. It's designed to provide precise and detailed responses by using a single completion prompting strategy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "system_prompt = \"You are an AI assistant designed to help users with a variety of tasks.\"\n",
    "\n",
    "agent = Agent.create(\n",
    "    llm=llm,\n",
    "    system_prompt=system_prompt,\n",
    "    prompt=\"{question}\",\n",
    "    prompt_variables=[\"question\"],\n",
    "    output_type=OutputType.STRING,\n",
    "    prompting_strategy=PromptingStrategy.SINGLE_COMPLETION,\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[38;2;50;164;103mFinal Answer\u001b[0m\u001b[1m\u001b[0m: Sure, I'd be happy to help explain the basics of quantum computing.\n",
      "\n",
      "Quantum computing is a type of computation that uses quantum bits, or qubits, to encode information as 0s, 1s, or both at the same time. This is different from classical computing, which uses bits to encode information as either 0 or 1.\n",
      "\n",
      "The key principles of quantum computing include superposition, entanglement, and quantum interference.\n",
      "\n",
      "1. Superposition: In classical computing, a bit is either a 0 or a 1. In quantum computing, thanks to superposition, a qubit can be in a state of 0, 1, or both at the same time. This allows quantum computers to process a higher volume of information compared to classical computers.\n",
      "\n",
      "2. Entanglement: This is a phenomenon where two qubits become linked, such that the state of one qubit is directly related to the state of another, no matter how far apart they are. This interconnectedness helps in faster processing of information.\n",
      "\n",
      "3. Quantum Interference: Quantum computers can manipulate the state of qubits to perform calculations. They can adjust the probability of a qubit's state, steering it towards the correct answer. This is done through quantum gates, which are similar to logic gates in classical computing but operate under the rules of quantum physics.\n",
      "\n",
      "Quantum computing has the potential to solve complex problems much more efficiently than classical computers. However, it's still in its early stages and there are many technical challenges to overcome before it becomes mainstream.\n"
     ]
    }
   ],
   "source": [
    "output = agent.invoke({\"question\": \"Hi, can you help me understand the basics of quantum computing?\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sure, I'd be happy to help explain the basics of quantum computing.\n",
      "\n",
      "Quantum computing is a type of computation that uses quantum bits, or qubits, to encode information as 0s, 1s, or both at the same time. This is different from classical computing, which uses bits to encode information as either 0 or 1.\n",
      "\n",
      "The key principles of quantum computing include superposition, entanglement, and quantum interference.\n",
      "\n",
      "1. Superposition: In classical computing, a bit is either a 0 or a 1. In quantum computing, thanks to superposition, a qubit can be in a state of 0, 1, or both at the same time. This allows quantum computers to process a higher volume of information compared to classical computers.\n",
      "\n",
      "2. Entanglement: This is a phenomenon where two qubits become linked, such that the state of one qubit is directly related to the state of another, no matter how far apart they are. This interconnectedness helps in faster processing of information.\n",
      "\n",
      "3. Quantum Interference: Quantum computers can manipulate the state of qubits to perform calculations. They can adjust the probability of a qubit's state, steering it towards the correct answer. This is done through quantum gates, which are similar to logic gates in classical computing but operate under the rules of quantum physics.\n",
      "\n",
      "Quantum computing has the potential to solve complex problems much more efficiently than classical computers. However, it's still in its early stages and there are many technical challenges to overcome before it becomes mainstream.\n"
     ]
    }
   ],
   "source": [
    "print(output.final_answer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[38;2;50;164;103mFinal Answer\u001b[0m\u001b[1m\u001b[0m: Let's denote the number of apples Ben has as x. According to the problem, Alex has 3 times as many apples as Ben, so Alex has 3x apples. Together, they have 24 apples. So, we can set up the following equation:\n",
      "\n",
      "x (Ben's apples) + 3x (Alex's apples) = 24\n",
      "\n",
      "Solving this equation gives:\n",
      "\n",
      "4x = 24\n",
      "x = 24 / 4\n",
      "x = 6\n",
      "\n",
      "So, Ben has 6 apples and Alex, having three times as many, has 18 apples.\n"
     ]
    }
   ],
   "source": [
    "output = agent.invoke({\"question\": \"Alex has three times as many apples as Ben. Together, they have 24 apples. How many apples does each person have?\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Let's denote the number of apples Ben has as x. According to the problem, Alex has 3 times as many apples as Ben, so Alex has 3x apples. Together, they have 24 apples. So, we can set up the following equation:\n",
      "\n",
      "x (Ben's apples) + 3x (Alex's apples) = 24\n",
      "\n",
      "Solving this equation gives:\n",
      "\n",
      "4x = 24\n",
      "x = 24 / 4\n",
      "x = 6\n",
      "\n",
      "So, Ben has 6 apples and Alex, having three times as many, has 18 apples.\n"
     ]
    }
   ],
   "source": [
    "print(output.final_answer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Chain-of-Thought\n",
    "\n",
    "The following code configures an LLM to specialize in answering earthquake-related questions by simulating the expertise of a United States Geological Survey (USGS) expert using Chain-of-Thought prompting. It integrates various tools, including those for earthquake information and current date retrieval, to enhance its responses."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "class USGeopoliticalSurveyEarthquakeAPI(BaseModel):\n",
    "    \"\"\"Class that implements the API interface.\"\"\"\n",
    "\n",
    "    start_time: str = Field(\n",
    "        None,\n",
    "        description=(\n",
    "            \"Limit to events on or after the specified start time. NOTE: All times use ISO8601 Date/Time format.\"\n",
    "            + \" Unless a timezone is specified, UTC is assumed.\"\n",
    "        ),\n",
    "    )\n",
    "    end_time: str = Field(\n",
    "        None,\n",
    "        description=(\n",
    "            \"Limit to events on or before the specified end time. NOTE: All times use ISO8601 Date/Time format.\"\n",
    "            + \" Unless a timezone is specified, UTC is assumed.\"\n",
    "        ),\n",
    "    )\n",
    "    limit: int = Field(\n",
    "        20000,\n",
    "        description=(\n",
    "            \"Limit the results to the specified number of events. NOTE: The service limits queries to 20000,\"\n",
    "            + \" and any that exceed this limit will generate a HTTP response code 400 Bad Request.\"\n",
    "        ),\n",
    "    )\n",
    "    min_depth: int = Field(\n",
    "        -100,\n",
    "        description=\"Limit to events with depth more than the specified minimum.\",\n",
    "    )\n",
    "    max_depth: int = Field(\n",
    "        1000,\n",
    "        description=\"Limit to events with depth less than the specified maximum.\",\n",
    "    )\n",
    "    min_magnitude: int = Field(\n",
    "        None,\n",
    "        description=\"Limit to events with a magnitude larger than the specified minimum.\",\n",
    "    )\n",
    "    max_magnitude: int = Field(\n",
    "        None,\n",
    "        description=\"Limit to events with a magnitude smaller than the specified maximum.\",\n",
    "    )\n",
    "    alert_level: str = Field(\n",
    "        None,\n",
    "        description=(\n",
    "            \"Limit to events with a specific PAGER alert level.\"\n",
    "            + \" The allowed values are: alert_level=green Limit to events with PAGER\"\n",
    "            + ' alert level \"green\". alert_level=yellow Limit to events with PAGER alert level \"yellow\".'\n",
    "            + ' alert_level=orange Limit to events with PAGER alert level \"orange\".'\n",
    "            + ' alert_level=red Limit to events with PAGER alert level \"red\".'\n",
    "        ),\n",
    "    )\n",
    "\n",
    "def get_earthquakes(\n",
    "    endpoint: str,\n",
    "    start_time: datetime = (datetime.now() - timedelta(days=30)).date(),\n",
    "    end_time: datetime = datetime.now().date(),\n",
    "    limit: int = 20000,\n",
    "    min_depth: int = -100,\n",
    "    max_depth: int = 1000,\n",
    "    min_magnitude: int | None = None,\n",
    "    max_magnitude: int | None = None,\n",
    "    alert_level: str | None = None,\n",
    ") -> Any:\n",
    "    params = {\n",
    "        \"format\": \"geojson\",\n",
    "        \"starttime\": start_time,\n",
    "        \"endtime\": end_time,\n",
    "        \"limit\": limit,\n",
    "        \"mindepth\": min_depth,\n",
    "        \"maxdepth\": max_depth,\n",
    "        \"minmagnitude\": min_magnitude,\n",
    "        \"maxmagnitude\": max_magnitude,\n",
    "        \"alertlevel\": alert_level,\n",
    "        \"eventtype\": \"earthquake\",\n",
    "    }\n",
    "    response = requests.get(\n",
    "        f\"https://earthquake.usgs.gov/fdsnws/event/1/{endpoint}\",\n",
    "        params=params,\n",
    "        timeout=None,\n",
    "    )\n",
    "    return response.json()\n",
    "\n",
    "def query_earthquakes(\n",
    "    start_time: datetime = (datetime.now() - timedelta(days=30)).date(),\n",
    "    end_time: datetime = datetime.now().date(),\n",
    "    limit: int = 20000,\n",
    "    min_depth: int = -100,\n",
    "    max_depth: int = 1000,\n",
    "    min_magnitude: int | None = None,\n",
    "    max_magnitude: int | None = None,\n",
    "    alert_level: str | None = None,\n",
    ") -> Any:\n",
    "    return get_earthquakes(\n",
    "        endpoint=\"query\",\n",
    "        start_time=start_time,\n",
    "        end_time=end_time,\n",
    "        limit=limit,\n",
    "        min_depth=min_depth,\n",
    "        max_depth=max_depth,\n",
    "        min_magnitude=min_magnitude,\n",
    "        max_magnitude=max_magnitude,\n",
    "        alert_level=alert_level,\n",
    "    )\n",
    "\n",
    "def count_earthquakes(\n",
    "    start_time: datetime = (datetime.now() - timedelta(days=30)).date(),\n",
    "    end_time: datetime = datetime.now().date(),\n",
    "    limit: int = 20000,\n",
    "    min_depth: int = -100,\n",
    "    max_depth: int = 1000,\n",
    "    min_magnitude: int | None = None,\n",
    "    max_magnitude: int | None = None,\n",
    "    alert_level: str | None = None,\n",
    ") -> Any:\n",
    "    return get_earthquakes(\n",
    "        endpoint=\"count\",\n",
    "        start_time=start_time,\n",
    "        end_time=end_time,\n",
    "        limit=limit,\n",
    "        min_depth=min_depth,\n",
    "        max_depth=max_depth,\n",
    "        min_magnitude=min_magnitude,\n",
    "        max_magnitude=max_magnitude,\n",
    "        alert_level=alert_level,\n",
    "    )\n",
    "\n",
    "query_earthquakes_tool = Tool(\n",
    "    function=query_earthquakes,\n",
    "    name=\"Query Earthquakes\",\n",
    "    description=\"Use this tool to search recent earthquakes\",\n",
    "    args_schema=USGeopoliticalSurveyEarthquakeAPI,\n",
    ")\n",
    "\n",
    "count_earthquakes_tool = Tool(\n",
    "    function=count_earthquakes,\n",
    "    name=\"Count Earthquakes\",\n",
    "    description=\"Use this tool to count and aggregate recent earthquakes\",\n",
    "    args_schema=USGeopoliticalSurveyEarthquakeAPI,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "system_prompt = \"You are an United States Geological Survey expert who can answer questions regarding earthquakes.\"\n",
    "\n",
    "agent = Agent.create(\n",
    "    llm=llm,\n",
    "    system_prompt=system_prompt,\n",
    "    prompt=\"{question}\",\n",
    "    prompt_variables=[\"question\"],\n",
    "    output_type=OutputType.STRING,\n",
    "    tools=[current_date, count_earthquakes_tool, query_earthquakes_tool],\n",
    "    prompting_strategy=PromptingStrategy.CHAIN_OF_THOUGHT,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following question demonstrates how the LLM processes a user's query about the number of earthquakes that occurred on the current day. Using the provided tools, the LLM first retrieves the current date and time. This initial step is crucial for identifying and counting the seismic activities for the day. This process shows the power of Chain-of-Thought, particularly ReAct, which enables the LLM to tackle multi-step problems. By breaking down the task into smaller, manageable steps - first obtaining the current date, then querying the relevant seismic data - the LLM can deliver correct answers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[38;2;45;114;210mThought\u001b[0m\u001b[1m\u001b[0m: To answer this question, I need to use the \"Count Earthquakes\" tool. I need to set the start_time to the beginning of today and the end_time to the current time.\n",
      "\u001b[1m\u001b[38;2;236;154;60mTool\u001b[0m\u001b[1m\u001b[0m: Current Date\n",
      "\u001b[1m\u001b[38;2;236;154;60mTool Input\u001b[0m\u001b[1m\u001b[0m: {}\n",
      "\u001b[1m\u001b[38;2;236;154;60mTool Output\u001b[0m\u001b[1m\u001b[0m: 2024-08-22 07:40:14.698953\n",
      "\u001b[1m\u001b[38;2;45;114;210mThought\u001b[0m\u001b[1m\u001b[0m: Now that I have the current date and time, I can use the \"Count Earthquakes\" tool to find out how many earthquakes occurred today. The start_time should be the start of today (2024-08-22 00:00:00) and the end_time should be the current time (2024-08-22 07:40:14).\n",
      "\u001b[1m\u001b[38;2;236;154;60mTool\u001b[0m\u001b[1m\u001b[0m: Count Earthquakes\n",
      "\u001b[1m\u001b[38;2;236;154;60mTool Input\u001b[0m\u001b[1m\u001b[0m: {'start_time': '2024-08-22T00:00:00', 'end_time': '2024-08-22T07:40:14'}\n",
      "\u001b[1m\u001b[38;2;236;154;60mTool Output\u001b[0m\u001b[1m\u001b[0m: {'count': 48, 'maxAllowed': 20000}\n",
      "\u001b[1m\u001b[38;2;45;114;210mThought\u001b[0m\u001b[1m\u001b[0m: The tool output shows that there have been 48 earthquakes today.\n",
      "\u001b[1m\u001b[38;2;50;164;103mFinal Answer\u001b[0m\u001b[1m\u001b[0m: There have been 48 earthquakes today.\n"
     ]
    }
   ],
   "source": [
    "output = agent.invoke({\"question\": \"How many earthquakes occurred today?\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "There have been 48 earthquakes today.\n"
     ]
    }
   ],
   "source": [
    "print(output.final_answer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When we provide the LLM with a follow-up question where ```Show me 3``` refers to the earthquakes that occurred today, it uses the context and continuity from the chat history to understand that \"today\" pertains to the current day. This capability allows the LLM to skip re-fetching the current date and directly query the earthquakes that happened on the same day, based on the date from the previous interaction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[38;2;45;114;210mThought\u001b[0m\u001b[1m\u001b[0m: To show the details of 3 earthquakes that occurred today, I need to use the \"Query Earthquakes\" tool. I will set the start_time to the beginning of today, the end_time to the current time, and the limit to 3.\n",
      "\u001b[1m\u001b[38;2;236;154;60mTool\u001b[0m\u001b[1m\u001b[0m: Query Earthquakes\n",
      "\u001b[1m\u001b[38;2;236;154;60mTool Input\u001b[0m\u001b[1m\u001b[0m: {'start_time': '2024-08-22T00:00:00', 'end_time': '2024-08-22T07:40:14', 'limit': 3}\n",
      "\u001b[1m\u001b[38;2;236;154;60mTool Output\u001b[0m\u001b[1m\u001b[0m: {'type': 'FeatureCollection', 'metadata': {'generated': 1724305233000, 'url': 'https://earthquake.usgs.gov/fdsnws/event/1/query?format=geojson&starttime=2024-08-22T00%3A00%3A00&endtime=2024-08-22T07%3A40%3A14&limit=3&mindepth=-100&maxdepth=1000&eventtype=earthquake', 'title': 'USGS Earthquakes', 'status': 200, 'api': '1.14.1', 'limit': 3, 'offset': 1, 'count': 3}, 'features': [{'type': 'Feature', 'properties': {'mag': 1.31, 'place': '8 km NNW of The Geysers, CA', 'time': 1724304849300, 'updated': 1724304947189, 'tz': None, 'url': 'https://earthquake.usgs.gov/earthquakes/eventpage/nc75051621', 'detail': 'https://earthquake.usgs.gov/fdsnws/event/1/query?eventid=nc75051621&format=geojson', 'felt': None, 'cdi': None, 'mmi': None, 'alert': None, 'status': 'automatic', 'tsunami': 0, 'sig': 26, 'net': 'nc', 'code': '75051621', 'ids': ',nc75051621,', 'sources': ',nc,', 'types': ',nearby-cities,origin,phase-data,', 'nst': 24, 'dmin': 0.008473, 'rms': 0.02, 'gap': 45, 'magType': 'md', 'type': 'earthquake', 'title': 'M 1.3 - 8 km NNW of The Geysers, CA'}, 'geometry': {'type': 'Point', 'coordinates': [-122.81600189209, 38.8305015563965, 1.78999996185303]}, 'id': 'nc75051621'}, {'type': 'Feature', 'properties': {'mag': 2, 'place': '54 km SE of Denali National Park, Alaska', 'time': 1724304129138, 'updated': 1724304243284, 'tz': None, 'url': 'https://earthquake.usgs.gov/earthquakes/eventpage/ak024asqegj5', 'detail': 'https://earthquake.usgs.gov/fdsnws/event/1/query?eventid=ak024asqegj5&format=geojson', 'felt': None, 'cdi': None, 'mmi': None, 'alert': None, 'status': 'automatic', 'tsunami': 0, 'sig': 62, 'net': 'ak', 'code': '024asqegj5', 'ids': ',ak024asqegj5,', 'sources': ',ak,', 'types': ',origin,phase-data,', 'nst': None, 'dmin': None, 'rms': 0.91, 'gap': None, 'magType': 'ml', 'type': 'earthquake', 'title': 'M 2.0 - 54 km SE of Denali National Park, Alaska'}, 'geometry': {'type': 'Point', 'coordinates': [-150.943, 63.1935, 2.4]}, 'id': 'ak024asqegj5'}, {'type': 'Feature', 'properties': {'mag': 2.01, 'place': '5 km SSW of Volcano, Hawaii', 'time': 1724304023760, 'updated': 1724304183530, 'tz': None, 'url': 'https://earthquake.usgs.gov/earthquakes/eventpage/hv74418861', 'detail': 'https://earthquake.usgs.gov/fdsnws/event/1/query?eventid=hv74418861&format=geojson', 'felt': None, 'cdi': None, 'mmi': None, 'alert': None, 'status': 'automatic', 'tsunami': 0, 'sig': 62, 'net': 'hv', 'code': '74418861', 'ids': ',hv74418861,', 'sources': ',hv,', 'types': ',origin,phase-data,', 'nst': 16, 'dmin': 0.008996, 'rms': 0.109999999, 'gap': 53, 'magType': 'ml', 'type': 'earthquake', 'title': 'M 2.0 - 5 km SSW of Volcano, Hawaii'}, 'geometry': {'type': 'Point', 'coordinates': [-155.252334594727, 19.3945007324219, 0.610000014305115]}, 'id': 'hv74418861'}], 'bbox': [-155.25233459473, 19.394500732422, 0.61000001430511, -122.81600189209, 63.1935, 2.4]}\n",
      "\u001b[1m\u001b[38;2;45;114;210mThought\u001b[0m\u001b[1m\u001b[0m: Based on the output from the Query Earthquakes tool, I can provide the details of the three earthquakes that occurred today.\n",
      "\u001b[1m\u001b[38;2;50;164;103mFinal Answer\u001b[0m\u001b[1m\u001b[0m: Here are the details of 3 earthquakes that occurred today:\n",
      "\n",
      "1. A magnitude 1.31 earthquake occurred 8 km NNW of The Geysers, CA. More details can be found [here](https://earthquake.usgs.gov/earthquakes/eventpage/nc75051621).\n",
      "\n",
      "2. A magnitude 2.0 earthquake occurred 54 km SE of Denali National Park, Alaska. More details can be found [here](https://earthquake.usgs.gov/earthquakes/eventpage/ak024asqegj5).\n",
      "\n",
      "3. A magnitude 2.01 earthquake occurred 5 km SSW of Volcano, Hawaii. More details can be found [here](https://earthquake.usgs.gov/earthquakes/eventpage/hv74418861).\n"
     ]
    }
   ],
   "source": [
    "output = agent.invoke({\"question\": \"Show me 3.\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Here are the details of 3 earthquakes that occurred today:\n",
      "\n",
      "1. A magnitude 1.31 earthquake occurred 8 km NNW of The Geysers, CA. More details can be found [here](https://earthquake.usgs.gov/earthquakes/eventpage/nc75051621).\n",
      "\n",
      "2. A magnitude 2.0 earthquake occurred 54 km SE of Denali National Park, Alaska. More details can be found [here](https://earthquake.usgs.gov/earthquakes/eventpage/ak024asqegj5).\n",
      "\n",
      "3. A magnitude 2.01 earthquake occurred 5 km SSW of Volcano, Hawaii. More details can be found [here](https://earthquake.usgs.gov/earthquakes/eventpage/hv74418861).\n"
     ]
    }
   ],
   "source": [
    "print(output.final_answer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The LLM can also handle straightforward informational questions that don't require the use of additional tools. In this case, the question about the possibility of MegaQuakes (magnitude 10 or larger) can be answered directly by the model without invoking any tools. This shows the LLM's ability to recognize when it can rely on its internal knowledge base to provide an accurate response."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[38;2;45;114;210mThought\u001b[0m\u001b[1m\u001b[0m: The user is asking about the possibility of a magnitude 10 or larger earthquake, also known as a \"MegaQuake\". This is a question related to the science of earthquakes and does not require the use of any tools.\n",
      "\u001b[1m\u001b[38;2;50;164;103mFinal Answer\u001b[0m\u001b[1m\u001b[0m: Theoretically, a magnitude 10 earthquake is possible, but it's extremely unlikely. The magnitude of an earthquake is related to the length of the fault on which it occurs - the larger the fault, the larger the earthquake. The largest earthquake ever recorded was a magnitude 9.5 in Chile in 1960. A magnitude 10 quake would require a fault rupture along the entire length of the earth's largest fault, the San Andreas, which is highly unlikely. Furthermore, the energy release of a magnitude 10 quake would be equivalent to the explosion of 1 billion nuclear bombs, which would cause catastrophic damage.\n"
     ]
    }
   ],
   "source": [
    "output = agent.invoke({\"question\": \"Can MegaQuakes really happen? Like a magnitude 10 or larger?\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Theoretically, a magnitude 10 earthquake is possible, but it's extremely unlikely. The magnitude of an earthquake is related to the length of the fault on which it occurs - the larger the fault, the larger the earthquake. The largest earthquake ever recorded was a magnitude 9.5 in Chile in 1960. A magnitude 10 quake would require a fault rupture along the entire length of the earth's largest fault, the San Andreas, which is highly unlikely. Furthermore, the energy release of a magnitude 10 quake would be equivalent to the explosion of 1 billion nuclear bombs, which would cause catastrophic damage.\n"
     ]
    }
   ],
   "source": [
    "print(output.final_answer)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llm-powered-ai-agents",
   "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}