notebooks/react_repl_agent.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ReAct REPL Agent\n",
    "\n",
    "Large-language model based \"agent\" that can can execute small tasks by using a Python REPL and pre-implemented methods. The \"agent loop\" is based on [ReAct](https://arxiv.org/abs/2210.03629).\n",
    "\n",
    "The agent can look up methods using a `method_search()` method that uses semantic search on the indexed Python methods. Methods are indexed using OpenAI embedding of their signatures and descriptions.\n",
    "\n",
    "To run this notebook make sure all secrets are in place to be able to access all APIs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load_ext autoreload\n",
    "# %autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import logging\n",
    "from pathlib import Path\n",
    "from datetime import datetime, timedelta\n",
    "import time\n",
    "\n",
    "from openai import OpenAIError"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logger = logging.getLogger(\"root\")\n",
    "logger.setLevel(logging.INFO)\n",
    "logger.addHandler(logging.StreamHandler(sys.stderr))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from react_repl_agent.utils.output_format_logger import MdLogger\n",
    "from react_repl_agent.llm_api import set_model, get_nb_tokens\n",
    "from react_repl_agent.agent_steps.react_steps import (\n",
    "    react_step_1,\n",
    "    ReActFinished,\n",
    "    get_react_step,\n",
    ")\n",
    "from react_repl_agent.methods import get_method_index\n",
    "from react_repl_agent.apis.location import get_user_location\n",
    "from react_repl_agent.utils.timedelta import format_timedelta"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Notebook Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "llm_model = \"\"  # E.g. \"gpt-4\"\n",
    "task_name = \"\"\n",
    "task_description = \"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"{llm_model=}\")\n",
    "print(f\"{task_name=}\")\n",
    "print(f\"{task_description=!r}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Select OpenAI Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "set_model(llm_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create method index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "method_search, context_dict = get_method_index()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define Task"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ReAct Loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_time = time.time()\n",
    "md_log = MdLogger()\n",
    "# Get prompt to start with\n",
    "md_log.display(f\"# Task: {task_name}\")\n",
    "md_log.display(task_description)\n",
    "md_log.display(f\"Date: {datetime.now().strftime('%Y-%m-%d  %H:%M:%S %z')}\")\n",
    "md_log.display(f\"Model: {llm_model}\")\n",
    "\n",
    "md_log.display(\"\\n## Initial Prompt\")\n",
    "prompt_state, step_idx = react_step_1(\n",
    "    task=task_description,\n",
    "    method_search_fun=method_search,\n",
    "    current_loc_method=get_user_location,\n",
    "    context_dict=context_dict,\n",
    ")\n",
    "md_log.display_react(prompt_state)\n",
    "\n",
    "# Run the ReAct THOUGHT/ACTION/OBSERVATION loop till the task is finished\n",
    "max_steps = 20\n",
    "for i in range(step_idx, max_steps+1):\n",
    "    md_log.display(f\"\\n## STEP {i} - prompt size = {get_nb_tokens(prompt_state)}\")\n",
    "    try:\n",
    "        thought, action, observation, prompt_state = get_react_step(\n",
    "            prompt_state=prompt_state,\n",
    "            step_idx=i,\n",
    "            context_dict=context_dict,\n",
    "            method_search_fun=method_search,\n",
    "        )\n",
    "    except ReActFinished as exc:\n",
    "        md_log.display(f\"### THOUGHT {i}:\")\n",
    "        md_log.display(exc.thought)\n",
    "        md_log.display(f\"### ACTION {i}:\")\n",
    "        md_log.display(f\"```python\\n{exc.action}\\n```\")\n",
    "        prompt_state += f\"THOUGHT {i}:\\n{exc.thought}\\nACTION {i}:\\n```python\\n{exc.action}\\n```\\n\"\"\"\n",
    "        total_runtime = timedelta(seconds=time.time() - start_time)\n",
    "        md_log.display(f\"`stop()` called! Agent finished in {i} steps and {format_timedelta(total_runtime)}.\")\n",
    "        md_log.display(str(exc))\n",
    "        break\n",
    "    except OpenAIError as exc:\n",
    "        md_log.display(f\"## ERROR!\")\n",
    "        md_log.display(f\"{type(exc).__name__}: {exc!s}\")\n",
    "        prompt_state += f\"\\n...\\nERROR {type(exc).__name__}: {exc!s}.\"\"\"\n",
    "        logging.exception(exc)\n",
    "        break\n",
    "    md_log.display(f\"### THOUGHT {i}:\")\n",
    "    md_log.display(thought)\n",
    "    md_log.display(f\"### ACTION {i}:\")\n",
    "    md_log.display(f\"```python\\n{action}\\n```\")\n",
    "    md_log.display(f\"### OBSERVATION {i}:\")\n",
    "    md_log.display(f\"{observation}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Export full prompt history"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Export\n",
    "export_dir = Path(\"../react_logs\").resolve()\n",
    "now = datetime.now()\n",
    "export_name = f\"{task_name}__{llm_model}__{now.strftime('%Y_%m_%d__%H-%M-%S')}.txt\"\n",
    "export_path_plain = export_dir / llm_model / now.strftime('%Y') / now.strftime('%m') / now.strftime('%d') / export_name\n",
    "export_path_plain.parent.mkdir(parents=True, exist_ok=True)\n",
    "# Export plaintext\n",
    "with export_path_plain.open(\"w\") as f_out:\n",
    "    f_out.write(prompt_state.strip())\n",
    "# Export Markdown run log\n",
    "export_path_md = export_path_plain.with_suffix(\".md\")\n",
    "with export_path_md.open(\"w\") as f_out:\n",
    "    f_out.write(md_log.log.strip())\n",
    "\n",
    "print(f\"Exported run logs to '{export_path_plain}'\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "gpt3",
   "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.10"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "c9b6b824e19d80b215e51f65ef0704b185dd46a9386b77c5813a19deb9c87af4"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}