tutorials/math_problem_solving.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building MathSolving Problem with agentlite\n",
    "This tutorial is a simple case on how to design a math solving agent with agentlite and WolframAlpha. Designing the math solving agent contains three steps:\n",
    "\n",
    "- Defining the WolframAlphaSolver action with Wolframalpha api\n",
    "\n",
    "- Defining the MathSolveAgent with WolframAlphaSolver action\n",
    "\n",
    "- Test the MathSolveAgent"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define WolframAlphaSolver action\n",
    "For WolframAlphaSolver action, we need to provide three properties following the design of BaseAction.\n",
    "\n",
    "- action_name\n",
    "\n",
    "- action_desc\n",
    "\n",
    "- params_doc\n",
    "\n",
    "Agent will use these three property to understand how to use this action.\n",
    "And the action execution is implemented in __call__() method. The params_doc is a dictionary to explain the input parameters in __call__() method.\n",
    "\n",
    "You need to get your own api_keys from the corresponding website:\n",
    "\n",
    "- openai_api_key: https://openai.com/blog/openai-api\n",
    "\n",
    "- wolfram_alpha_api_key: https://developer.wolframalpha.com/access"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from agentlite.actions.BaseAction import BaseAction\n",
    "import wolframalpha\n",
    "from typing import Tuple\n",
    "from env import openai_api_key\n",
    "from env import wolfram_alpha_api_key\n",
    "from agentlite.commons import AgentAct, TaskPackage\n",
    "from agentlite.actions import ThinkAct, FinishAct\n",
    "from agentlite.actions.InnerActions import INNER_ACT_KEY\n",
    "\n",
    "\n",
    "\n",
    "# Define WolframAlphaSolver Action\n",
    "class WolframAlphaSolver(BaseAction):\n",
    "    def __init__(self, app_id) -> None:\n",
    "        action_name = \"WolframAlphaSolver\"\n",
    "        action_desc = \"Using this action to solve a math problem with WolframAlpha.\"\n",
    "        params_doc = {\"query\": \"the math equation to be solved. be simple\"}\n",
    "        self.client = wolframalpha.Client(app_id)\n",
    "        super().__init__(\n",
    "            action_name=action_name,\n",
    "            action_desc=action_desc,\n",
    "            params_doc=params_doc,\n",
    "        )\n",
    "        \n",
    "    def run(self, query: str) -> Tuple[str, bool]:\n",
    "        \"\"\"Run query through WolframAlpha and parse result.\"\"\"\n",
    "        from urllib.error import HTTPError\n",
    "\n",
    "        is_success = False  \n",
    "\n",
    "        res = self.client.query(query)\n",
    "\n",
    "        if not res[\"@success\"]:\n",
    "            return (\n",
    "                \"Your Wolfram query is invalid. Please try a new query for wolfram.\",\n",
    "                is_success,\n",
    "            )\n",
    "        assumption = next(res.pods).text\n",
    "        answer = \"\"\n",
    "        for result in res[\"pod\"]:\n",
    "            if result[\"@title\"] == \"Solution\":\n",
    "                answer = result[\"subpod\"][\"plaintext\"]\n",
    "            if result[\"@title\"] == \"Results\" or result[\"@title\"] == \"Solutions\":\n",
    "                for i, sub in enumerate(result[\"subpod\"]):\n",
    "                    answer += f\"ans {i}: \" + sub[\"plaintext\"] + \"\\n\"\n",
    "                break\n",
    "        if answer == \"\":\n",
    "            answer = next(res.results).text\n",
    "\n",
    "\n",
    "        if answer is None or answer == \"\":\n",
    "            # We don't want to return the assumption alone if answer is empty\n",
    "            return \"No good Wolfram Alpha Result was found\", is_success\n",
    "        is_success = True\n",
    "        return f\"Assumption: {assumption} \\nAnswer: {answer}\", is_success\n",
    "\n",
    "    def __call__(self, query):\n",
    "        res = self.run(query)\n",
    "        return res[0]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define the math solve agent\n",
    "For the MathSolveAgent, we need to define several properties:\n",
    "\n",
    "- actions: A list of available actions (WolframAlphaSolver) for MathSolveAgent\n",
    "\n",
    "- name: Name of the agent\n",
    "\n",
    "- role: The agent's role"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# defining the MathSolve agent as an agent to use the WolframAlphaSolver action\n",
    "from typing import List\n",
    "from agentlite.actions import BaseAction\n",
    "from agentlite.agents import ABCAgent, BaseAgent\n",
    "from agentlite.llm.agent_llms import BaseLLM, get_llm_backend\n",
    "from agentlite.llm.LLMConfig import LLMConfig\n",
    "from agentlite.logging.multi_agent_log import AgentLogger\n",
    "\n",
    "## using the agentrool logger to record the running log\n",
    "agent_logger = AgentLogger(PROMPT_DEBUG_FLAG=False)\n",
    "\n",
    "# Define the MathSolve agent\n",
    "class MathSolveAgent(BaseAgent):\n",
    "    def __init__(\n",
    "        self,\n",
    "        llm: BaseLLM,\n",
    "        actions: List[BaseAction] = [WolframAlphaSolver(wolfram_alpha_api_key)], # we add both actions into agents\n",
    "        manager: ABCAgent = None,\n",
    "        **kwargs\n",
    "    ):\n",
    "        name = \"MathSolve_agent\"\n",
    "        role = \"\"\"You can answer math questions by WolframAlphaSolver action. Finish it if you find the answer.\"\"\"\n",
    "        super().__init__(\n",
    "            name=name,\n",
    "            role=role,\n",
    "            llm=llm,\n",
    "            actions=actions,\n",
    "            manager=manager,\n",
    "            logger=agent_logger,\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize the MathSolveAgent need three steps\n",
    "\n",
    "- Initialize the backend LLM\n",
    "\n",
    "- Initialize the agent with LLM\n",
    "\n",
    "- Add an illustration example for the agent (Optional, buy highly recommended to stabilize the agent)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the backend LLM\n",
    "llm_config_dict = {\"llm_name\": \"gpt-4-32k\", \"temperature\": 0.0, \"openai_api_key\": openai_api_key}\n",
    "llm_config = LLMConfig(llm_config_dict)\n",
    "llm = get_llm_backend(llm_config)\n",
    "\n",
    "# Initialize the agent\n",
    "wolf_action = WolframAlphaSolver(wolfram_alpha_api_key) \n",
    "mathsolve_agent = MathSolveAgent(llm=llm, actions = [wolf_action])\n",
    "\n",
    "\n",
    "# Add an illustration example\n",
    "exp_task = \"What is the result of 1*98+23*45+87*45\"\n",
    "exp_task_pack = TaskPackage(instruction=exp_task)\n",
    "\n",
    "act_1 = AgentAct(\n",
    "    name=ThinkAct.action_name,\n",
    "    params={INNER_ACT_KEY: f\"\"\"The math equation is \"1*98+23*45+87*45\". And I can solve it with WolframAlphaSolver action.\"\"\"\n",
    "    },\n",
    ")\n",
    "obs_1 = \"ok\"\n",
    "\n",
    "act_2 = AgentAct(\n",
    "    name=wolf_action.action_name,\n",
    "    params={\"query\": \"1*98+23*45+87*45\",\n",
    "        },\n",
    ")\n",
    "obs_2 = \"\"\"5048\"\"\"\n",
    "\n",
    "act_3 = AgentAct(name=FinishAct.action_name, params={INNER_ACT_KEY: \"The result is 5048.\"})\n",
    "obs_3 = \"Task Completed.\"\n",
    "exp_act_obs = [(act_1, obs_1), (act_2, obs_2), (act_3, obs_3)]\n",
    "\n",
    "mathsolve_agent.prompt_gen.add_example(\n",
    "    task = exp_task_pack, action_chain = exp_act_obs\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test the agent on other math problems"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Agent \u001b[94mMathSolve_agent\u001b[0m receives the following \u001b[4mTaskPackage\u001b[0m:\n",
      "\u001b[96m[\n",
      "\tTask ID: 563daf80-6c85-403c-8d40-38d91a2395d9\n",
      "\tInstruction: What is the result of 76*23+87/2+45\n",
      "]\u001b[0m\n",
      "====\u001b[94mMathSolve_agent\u001b[0m starts execution on TaskPackage 563daf80-6c85-403c-8d40-38d91a2395d9====\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/liangwei.yang/miniconda3/lib/python3.11/site-packages/langchain_core/_api/deprecation.py:117: LangChainDeprecationWarning: The function `run` was deprecated in LangChain 0.1.0 and will be removed in 0.2.0. Use invoke instead.\n",
      "  warn_deprecated(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Agent \u001b[94mMathSolve_agent\u001b[0m takes 0-step \u001b[4mAction\u001b[0m:\n",
      "\u001b[94m{\n",
      "\tname: Think\n",
      "\tparams: {'response': 'The math equation is \"76*23+87/2+45\". And I can solve it with WolframAlphaSolver action.'}\n",
      "}\u001b[0m\n",
      "Observation: \u001b[92mOK\u001b[0m\n",
      "Agent \u001b[94mMathSolve_agent\u001b[0m takes 1-step \u001b[4mAction\u001b[0m:\n",
      "\u001b[94m{\n",
      "\tname: WolframAlphaSolver\n",
      "\tparams: {'query': '76*23+87/2+45'}\n",
      "}\u001b[0m\n",
      "Observation: \u001b[92mAssumption: 76×23 + 87/2 + 45 \n",
      "Answer: 3673/2\u001b[0m\n",
      "Agent \u001b[94mMathSolve_agent\u001b[0m takes 2-step \u001b[4mAction\u001b[0m:\n",
      "\u001b[94m{\n",
      "\tname: Finish\n",
      "\tparams: {'response': 'The result is 3673/2.'}\n",
      "}\u001b[0m\n",
      "Observation: \u001b[92mTask Completed.\u001b[0m\n",
      "=========\u001b[94mMathSolve_agent\u001b[0m finish execution. TaskPackage[ID:563daf80-6c85-403c-8d40-38d91a2395d9] status:\n",
      "\u001b[96m[\n",
      "\tcompletion: completed\n",
      "\tanswer: The result is 3673/2.\n",
      "]\u001b[0m\n",
      "==========\n",
      "response: The result is 3673/2.\n"
     ]
    }
   ],
   "source": [
    "# calling the agent with TaskPackage\n",
    "from agentlite.commons import TaskPackage\n",
    "test_task = \"What is the result of 76*23+87/2+45\"\n",
    "\n",
    "\n",
    "test_task_pack = TaskPackage(instruction=test_task)\n",
    "response = mathsolve_agent(test_task_pack)\n",
    "print(\"response:\", response)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}