examples/ex_babyagi.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IY8_JMKsgTdE"
      },
      "source": [
        "# BabyAGI Implementation with LLFn\n",
        "\n",
        "<img src=\"https://user-images.githubusercontent.com/21254008/235015461-543a897f-70cc-4b63-941a-2ae3c9172b11.png\" width=\"600\"/>\n",
        "\n",
        "In this example, we're going to rebuild [**BabyAGI**](https://github.com/yoheinakajima/babyagi#how-to-use) (originally created by\n",
        "Yohei Nakajima) using [**LLFn**](https://github.com/orgexyz/LLFn).\n",
        "\n",
        "### Modifications vs Original Implementation\n",
        "\n",
        "- We will use almost exactly the same prompts as the original implementation (with minor modification for readability).\n",
        "- We will NOT be using vectorstore to enrich the data, but instead including the past 5 tasks executed and their results for simplicity."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FuzU84TPfzFU"
      },
      "source": [
        "## **Before we start**: let's setup objective and keys for our BabyAGI"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Set objective here\n",
        "OBJECTIVE = \"Solve world hunger\"\n",
        "\n",
        "# Set API key here\n",
        "OPENAI_MODEL = \"gpt-3.5-turbo\"\n",
        "OPENAI_API_KEY = \"<OPENAI_API_KEY>\""
      ],
      "metadata": {
        "id": "Zwd7K4LNV7VY"
      },
      "execution_count": 43,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YP50Nf87Vwuo"
      },
      "source": [
        "## **Step 1**: Installing dependencies and setup `function_prompt`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 44,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "GrhBvuGQcdhb",
        "outputId": "453f88a4-6164-49d9-d79a-8e418cc1a8a2"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: langchain in /usr/local/lib/python3.10/dist-packages (0.0.238)\n",
            "Requirement already satisfied: llfn in /usr/local/lib/python3.10/dist-packages (0.1.4)\n",
            "Requirement already satisfied: PyYAML>=5.4.1 in /usr/local/lib/python3.10/dist-packages (from langchain) (6.0)\n",
            "Requirement already satisfied: SQLAlchemy<3,>=1.4 in /usr/local/lib/python3.10/dist-packages (from langchain) (2.0.18)\n",
            "Requirement already satisfied: aiohttp<4.0.0,>=3.8.3 in /usr/local/lib/python3.10/dist-packages (from langchain) (3.8.4)\n",
            "Requirement already satisfied: async-timeout<5.0.0,>=4.0.0 in /usr/local/lib/python3.10/dist-packages (from langchain) (4.0.2)\n",
            "Requirement already satisfied: dataclasses-json<0.6.0,>=0.5.7 in /usr/local/lib/python3.10/dist-packages (from langchain) (0.5.13)\n",
            "Requirement already satisfied: langsmith<0.1.0,>=0.0.11 in /usr/local/lib/python3.10/dist-packages (from langchain) (0.0.12)\n",
            "Requirement already satisfied: numexpr<3.0.0,>=2.8.4 in /usr/local/lib/python3.10/dist-packages (from langchain) (2.8.4)\n",
            "Requirement already satisfied: numpy<2,>=1 in /usr/local/lib/python3.10/dist-packages (from langchain) (1.22.4)\n",
            "Requirement already satisfied: openapi-schema-pydantic<2.0,>=1.2 in /usr/local/lib/python3.10/dist-packages (from langchain) (1.2.4)\n",
            "Requirement already satisfied: pydantic<2,>=1 in /usr/local/lib/python3.10/dist-packages (from langchain) (1.10.11)\n",
            "Requirement already satisfied: requests<3,>=2 in /usr/local/lib/python3.10/dist-packages (from langchain) (2.27.1)\n",
            "Requirement already satisfied: tenacity<9.0.0,>=8.1.0 in /usr/local/lib/python3.10/dist-packages (from langchain) (8.2.2)\n",
            "Requirement already satisfied: openai>=0.27.0 in /usr/local/lib/python3.10/dist-packages (from llfn) (0.27.8)\n",
            "Requirement already satisfied: python-dotenv>=1.0.0 in /usr/local/lib/python3.10/dist-packages (from llfn) (1.0.0)\n",
            "Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.8.3->langchain) (23.1.0)\n",
            "Requirement already satisfied: charset-normalizer<4.0,>=2.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.8.3->langchain) (2.0.12)\n",
            "Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.8.3->langchain) (6.0.4)\n",
            "Requirement already satisfied: yarl<2.0,>=1.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.8.3->langchain) (1.9.2)\n",
            "Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.8.3->langchain) (1.4.0)\n",
            "Requirement already satisfied: aiosignal>=1.1.2 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.8.3->langchain) (1.3.1)\n",
            "Requirement already satisfied: marshmallow<4.0.0,>=3.18.0 in /usr/local/lib/python3.10/dist-packages (from dataclasses-json<0.6.0,>=0.5.7->langchain) (3.20.1)\n",
            "Requirement already satisfied: typing-inspect<1,>=0.4.0 in /usr/local/lib/python3.10/dist-packages (from dataclasses-json<0.6.0,>=0.5.7->langchain) (0.9.0)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.10/dist-packages (from openai>=0.27.0->llfn) (4.65.0)\n",
            "Requirement already satisfied: typing-extensions>=4.2.0 in /usr/local/lib/python3.10/dist-packages (from pydantic<2,>=1->langchain) (4.7.1)\n",
            "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests<3,>=2->langchain) (1.26.16)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests<3,>=2->langchain) (2023.5.7)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests<3,>=2->langchain) (3.4)\n",
            "Requirement already satisfied: greenlet!=0.4.17 in /usr/local/lib/python3.10/dist-packages (from SQLAlchemy<3,>=1.4->langchain) (2.0.2)\n",
            "Requirement already satisfied: packaging>=17.0 in /usr/local/lib/python3.10/dist-packages (from marshmallow<4.0.0,>=3.18.0->dataclasses-json<0.6.0,>=0.5.7->langchain) (23.1)\n",
            "Requirement already satisfied: mypy-extensions>=0.3.0 in /usr/local/lib/python3.10/dist-packages (from typing-inspect<1,>=0.4.0->dataclasses-json<0.6.0,>=0.5.7->langchain) (1.0.0)\n"
          ]
        }
      ],
      "source": [
        "!pip install langchain llfn"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 45,
      "metadata": {
        "id": "BybSe4f4dCs3"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "from llfn import LLFn\n",
        "from pydantic import BaseModel, Field\n",
        "from typing import List\n",
        "from langchain.chat_models import ChatOpenAI\n",
        "\n",
        "llm = ChatOpenAI(\n",
        "    temperature=0.7,\n",
        "    model=OPENAI_MODEL,\n",
        "    openai_api_key=OPENAI_API_KEY,\n",
        ")\n",
        "\n",
        "function_prompt = LLFn()\n",
        "function_prompt.bind(llm)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WV7UWctIgP3j"
      },
      "source": [
        "## **Step 2**: Defining `task_creation_agent`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 46,
      "metadata": {
        "id": "xdizw0_VhwPy"
      },
      "outputs": [],
      "source": [
        "class TaskCreationOutput(BaseModel):\n",
        "    new_tasks: List[str] = Field(..., description=\"a list of new tasks (may be empty)\")\n",
        "    reasoning: str = Field(..., description=\"why you think that way\")\n",
        "\n",
        "@function_prompt\n",
        "def task_creation_agent(\n",
        "    objective: str,\n",
        "    last_task: str,\n",
        "    last_task_result: str,\n",
        "    task_list: List[str]\n",
        ") -> TaskCreationOutput:\n",
        "    return f\"\"\"\n",
        "You are to use the result from an execution agent to create new tasks with the following objective: {objective}.\n",
        "The last completed task has the result:\n",
        "{last_task_result}\n",
        "This result was based on this task description: {last_task}.\n",
        "\n",
        "These are {len(task_list)} incomplete tasks: {', '.join(task_list)}\n",
        "\n",
        "Based on the result, return a list of tasks to be completed in order to meet the objective.\n",
        "These new tasks must not overlap with incomplete tasks.\n",
        "\"\"\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lvFee7Xl6Ig9"
      },
      "source": [
        "## **Step 3**: Defining `task_prioritization_agent`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 47,
      "metadata": {
        "id": "EqZb4qNQ6MkL"
      },
      "outputs": [],
      "source": [
        "class TaskPrioritizationOutput(BaseModel):\n",
        "    tasks: List[str] = Field(..., description=\"a list of new tasks (may be empty)\")\n",
        "    reasoning: str = Field(..., description=\"why you prioritize tasks that way\")\n",
        "\n",
        "@function_prompt\n",
        "def task_prioritization_agent(\n",
        "    objective: str,\n",
        "    task_list: List[str]\n",
        ") -> TaskPrioritizationOutput:\n",
        "    task_bullets = \"\\n\".join([f\"  - {task}\" for task in task_list])\n",
        "    return f\"\"\"\n",
        "You are tasked with prioritizing the following tasks:\n",
        "{task_bullets}\n",
        "\n",
        "Consider the ultimate objective of your team: {objective}.\n",
        "\"\"\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PSM4SfgZ8vyb"
      },
      "source": [
        "## **Step 4**: Defining `task_execution_agent`"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "class TaskExecutionOutput(BaseModel):\n",
        "    task: str = Field(..., description=\"task name\")\n",
        "    result: str = Field(..., description=\"result and effect of your execution and why\")\n",
        "\n",
        "@function_prompt\n",
        "def task_execution_agent(\n",
        "    objective: str,\n",
        "    task: str,\n",
        "    completed_tasks: List[TaskExecutionOutput]\n",
        ") -> TaskExecutionOutput:\n",
        "    completed_task_bullets = \"\\n\".join([f\"  - {execution.task}: {execution.result}\" for execution in completed_tasks])\n",
        "    return  f\"\"\"\n",
        "Perform one task based on the following objective: {objective}.\n",
        "\n",
        "Take into account these previously completed tasks:\n",
        "{completed_task_bullets}\n",
        "\n",
        "Your task: {task}\n",
        "Response:\n",
        "\"\"\""
      ],
      "metadata": {
        "id": "Eo5IW4Sd86A3"
      },
      "execution_count": 48,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OOtp40L2DIii"
      },
      "source": [
        "## **Step 5**: Wire up everything and let BabyAGI run wild!"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "tasks = [\"Develop a task list\"]\n",
        "completed_tasks = []\n",
        "\n",
        "max_rounds = 5\n",
        "\n",
        "for r in range(1, max_rounds+1):\n",
        "    print(f\"\\n=====ROUND#{r}=====\")\n",
        "\n",
        "    # Print the task list\n",
        "    print(\"\\n*****TASK LIST*****\")\n",
        "    for t in tasks:\n",
        "        print(f\"- {t}\")\n",
        "\n",
        "    # Step 1: Pull the first incomplete task\n",
        "    task = tasks.pop(0)\n",
        "    print(\"\\n*****NEXT TASK*****\")\n",
        "    print(task)\n",
        "\n",
        "    # Send to execution function to complete the task based on the context\n",
        "    execution = task_execution_agent(OBJECTIVE, task, completed_tasks[-5:])\n",
        "    print(\"\\n*****TASK RESULT*****\")\n",
        "    print(execution.result)\n",
        "\n",
        "    # Step 2: Store the result in completed task\n",
        "    completed_tasks.append(execution)\n",
        "\n",
        "    # Step 3: Create new tasks re-prioritize task list\n",
        "    new_tasks = task_creation_agent(OBJECTIVE, execution.task, execution.result, tasks).new_tasks\n",
        "\n",
        "    print('\\nAdding new tasks to task_storage')\n",
        "    for t in new_tasks:\n",
        "        print(f\"- {t}\")\n",
        "\n",
        "    # Re-prioritize tasks\n",
        "    tasks = task_prioritization_agent(OBJECTIVE, tasks + new_tasks).tasks\n",
        "\n",
        "    # If no more tasks, terminate\n",
        "    if not len(tasks):\n",
        "        print(\"\\nDone.\")\n",
        "        break\n",
        ""
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8vmL6GD3DbOn",
        "outputId": "bb394e64-e456-4169-dab6-9ba3f80504e5"
      },
      "execution_count": 49,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "=====ROUND#1=====\n",
            "\n",
            "*****TASK LIST*****\n",
            "- Develop a task list\n",
            "\n",
            "*****NEXT TASK*****\n",
            "Develop a task list\n",
            "\n",
            "*****TASK RESULT*****\n",
            "To solve world hunger, we need to address various factors such as poverty, access to resources, agricultural practices, and distribution systems. Some possible tasks to work towards this objective include: \n",
            "\n",
            "1. Promote sustainable agriculture practices that increase food production and reduce environmental impact.\n",
            "2. Implement social safety nets and poverty alleviation programs to ensure access to food for vulnerable populations.\n",
            "3. Improve infrastructure and transportation systems to facilitate the distribution of food to remote areas.\n",
            "4. Increase investment in research and development of innovative solutions for food production and preservation.\n",
            "5. Collaborate with international organizations and governments to develop and implement comprehensive strategies to address world hunger.\n",
            "\n",
            "These are just a few examples, and the specific tasks needed to solve world hunger will vary depending on the context and resources available.\n",
            "\n",
            "Adding new tasks to task_storage\n",
            "- Promote sustainable agriculture practices that increase food production and reduce environmental impact.\n",
            "- Implement social safety nets and poverty alleviation programs to ensure access to food for vulnerable populations.\n",
            "- Improve infrastructure and transportation systems to facilitate the distribution of food to remote areas.\n",
            "- Increase investment in research and development of innovative solutions for food production and preservation.\n",
            "- Collaborate with international organizations and governments to develop and implement comprehensive strategies to address world hunger.\n",
            "\n",
            "=====ROUND#2=====\n",
            "\n",
            "*****TASK LIST*****\n",
            "- Promote sustainable agriculture practices that increase food production and reduce environmental impact.\n",
            "- Increase investment in research and development of innovative solutions for food production and preservation.\n",
            "- Implement social safety nets and poverty alleviation programs to ensure access to food for vulnerable populations.\n",
            "- Improve infrastructure and transportation systems to facilitate the distribution of food to remote areas.\n",
            "- Collaborate with international organizations and governments to develop and implement comprehensive strategies to address world hunger.\n",
            "\n",
            "*****NEXT TASK*****\n",
            "Promote sustainable agriculture practices that increase food production and reduce environmental impact.\n",
            "\n",
            "*****TASK RESULT*****\n",
            "Increase awareness about sustainable agriculture practices and their benefits, provide training and support for farmers to adopt these practices, and incentivize sustainable agriculture through government policies and programs.\n",
            "\n",
            "Adding new tasks to task_storage\n",
            "- Increase investment in research and development of innovative solutions for food production and preservation.\n",
            "- Implement social safety nets and poverty alleviation programs to ensure access to food for vulnerable populations.\n",
            "- Improve infrastructure and transportation systems to facilitate the distribution of food to remote areas.\n",
            "- Collaborate with international organizations and governments to develop and implement comprehensive strategies to address world hunger.\n",
            "\n",
            "=====ROUND#3=====\n",
            "\n",
            "*****TASK LIST*****\n",
            "- Increase investment in research and development of innovative solutions for food production and preservation.\n",
            "- Implement social safety nets and poverty alleviation programs to ensure access to food for vulnerable populations.\n",
            "- Improve infrastructure and transportation systems to facilitate the distribution of food to remote areas.\n",
            "- Collaborate with international organizations and governments to develop and implement comprehensive strategies to address world hunger.\n",
            "\n",
            "*****NEXT TASK*****\n",
            "Increase investment in research and development of innovative solutions for food production and preservation.\n",
            "\n",
            "*****TASK RESULT*****\n",
            "Increasing investment in research and development of innovative solutions for food production and preservation will help address the challenges of world hunger by finding new ways to increase food production, improve crop yields, and enhance food preservation techniques. This can be done by allocating more funds towards research institutions, supporting initiatives that promote innovation in agriculture, and collaborating with scientists and experts in the field to develop and implement effective solutions.\n",
            "\n",
            "Adding new tasks to task_storage\n",
            "- Implement social safety nets and poverty alleviation programs to ensure access to food for vulnerable populations.\n",
            "- Improve infrastructure and transportation systems to facilitate the distribution of food to remote areas.\n",
            "- Collaborate with international organizations and governments to develop and implement comprehensive strategies to address world hunger.\n",
            "\n",
            "=====ROUND#4=====\n",
            "\n",
            "*****TASK LIST*****\n",
            "- Implement social safety nets and poverty alleviation programs to ensure access to food for vulnerable populations.\n",
            "- Improve infrastructure and transportation systems to facilitate the distribution of food to remote areas.\n",
            "- Collaborate with international organizations and governments to develop and implement comprehensive strategies to address world hunger.\n",
            "\n",
            "*****NEXT TASK*****\n",
            "Implement social safety nets and poverty alleviation programs to ensure access to food for vulnerable populations.\n",
            "\n",
            "*****TASK RESULT*****\n",
            "Implementing social safety nets and poverty alleviation programs will help ensure access to food for vulnerable populations. These programs can include initiatives such as cash transfers, food subsidies, and employment opportunities that provide income and resources for individuals and families in need. By addressing the root causes of poverty and providing support to those in need, these programs can help reduce hunger and improve food security for vulnerable populations.\n",
            "\n",
            "Adding new tasks to task_storage\n",
            "- Improve infrastructure and transportation systems to facilitate the distribution of food to remote areas.\n",
            "- Collaborate with international organizations and governments to develop and implement comprehensive strategies to address world hunger.\n",
            "\n",
            "=====ROUND#5=====\n",
            "\n",
            "*****TASK LIST*****\n",
            "- Collaborate with international organizations and governments to develop and implement comprehensive strategies to address world hunger.\n",
            "- Improve infrastructure and transportation systems to facilitate the distribution of food to remote areas.\n",
            "\n",
            "*****NEXT TASK*****\n",
            "Collaborate with international organizations and governments to develop and implement comprehensive strategies to address world hunger.\n",
            "\n",
            "*****TASK RESULT*****\n",
            "Collaborating with international organizations and governments is crucial in developing and implementing comprehensive strategies to address world hunger. By working together, we can leverage resources, expertise, and global networks to tackle the complex challenges of food insecurity and ensure sustainable solutions for all.\n",
            "\n",
            "Adding new tasks to task_storage\n",
            "- Improve agricultural practices to increase food production\n",
            "- Promote sustainable farming methods to protect the environment\n",
            "- Invest in research and development for innovative food solutions\n",
            "- Enhance access to education and training on nutrition and food security\n",
            "- Support local farmers and food producers to strengthen local food systems\n"
          ]
        }
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}