llamaindex_workflow.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/wandb/aihackercup/blob/main/llamaindex_workflow.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"http://wandb.me/logo-im-png\" width=\"400\" alt=\"Weights & Biases\" />\n",
    "\n",
    "# W&B Lighting Competition - AI Hacker Cup \n",
    "\n",
    "</a>\n",
    "\n",
    "[Weights & Biases](https://wandb.ai/site?utm_source=colab&utm_medium=code&utm_campaign=lightning-ai-hacker-cup) are running a 7-day Lightning Competition focussed on solving practice problems for the  [2024 NeurIPS AI Hacker Cup](https://hackercupai.github.io/) challenge.\n",
    "\n",
    "#### Goal\n",
    "The goal is to try and solve all 5 of the 2023 practice questions for the AI Hacker Cup using MistralAI's models. We’re offering free MistralAI api access via the code in this colab to get people started.\n",
    "\n",
    "#### Competition GitHub\n",
    "The competition [repo here](https://github.com/wandb/aihackercup) contains this colab, the code for the Code Generation Agent and the details on how to make a submission and the competition rules. Note that to run this notebook you'll need to be running it with a T4 GPU (15GB) or larger as the embedding model is run locally.\n",
    "\n",
    "#### Discord\n",
    "You can join the official NeurIPS AI Hacker Cup [discord here](discord.gg/wWeN9hTH32) to share ideas and discuss winning solutions.\n",
    "\n",
    "## Prizes\n",
    "\n",
    "Weights & Biases are giving away a pair of Meta Ray-Ban Smart Glasses for the first individual to submit code that solves:\n",
    "- 3 out of 5 correct solutions\n",
    "- 4 out of 5 correct solutions\n",
    "- 5 out of 5 correct solutions\n",
    "\n",
    "(i.e. in total 3 pairs of sunglasses to give away)\n",
    "\n",
    "## Entry Submissions, Rules & Deadline\n",
    "\n",
    "See the [competition README](https://github.com/wandb/aihackercup) for how to make a submissions the the competition rules."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## W&B Weave\n",
    "\n",
    "[W&B Weave](https://weave-docs.wandb.ai/tutorial-eval?utm_source=colab&utm_medium=code&utm_campaign=lightning-ai-hacker-cup) is used in this competition to run the evaluations. It is a lightweight toolkit for tracking and evaluating LLM applications, built by Weights & Biases. \n",
    "\n",
    "<img src=\"https://raw.githubusercontent.com/wandb/weave/master/docs/static/img/evals-hero.png\" width=\"800\" height=\"450\">\n",
    "\n",
    "If you want to learn more about Weave, you can [get started](https://weave-docs.wandb.ai/quickstart?utm_source=colab&utm_medium=code&utm_campaign=lightning-ai-hacker-cup) by decorating Python functions with `@weave.op`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LlamaIndex Workflow with retries\n",
    "\n",
    "A simple workflow to solve the problem using LlamaIndex. It is a good example on how to format your codebase to use llamaIndex!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note: You need to run this cell only once**\n",
    "We will clone the starter-kits repo\n",
    "Set the rag folder as our working directory\n",
    "and install the dependencies for the project.\n",
    "\n",
    "**You can comment out the cell after you have run it once.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Clone the starter-kits repo\n",
    "!git clone https://github.com/wandb/aihackercup\n",
    "# Change directory to the rag folder. Running the next line twice in the same session will raise an error.\n",
    "%cd aihackercup\n",
    "# Install dependencies\n",
    "!pip install -r requirements.txt -qq"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To run this colab, create a [free Weights & Biases (W&B) account here](https://wandb.ai/site?utm_source=colab&utm_medium=code&utm_campaign=lightning-ai-hacker-cup) and then copy your API key from https://wandb.ai/authorize into the input box below when requested."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import weave\n",
    "\n",
    "WEAVE_PROJECT = \"ai-hacker-cup\"\n",
    "weave_client = weave.init(WEAVE_PROJECT)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Select MistralAI models used depending if you want a fast or strong LLM\n",
    "# You can see the full range of MistralAI models here: https://docs.mistral.ai/getting-started/models/\n",
    "FAST_LLM = \"open-mistral-nemo-2407\"\n",
    "STRONG_LLM = \"mistral-large-latest\"\n",
    "\n",
    "os.environ[\"FAST_LLM\"] = STRONG_LLM  # We'll use stong model everywhere\n",
    "os.environ[\"STRONG_LLM\"] = STRONG_LLM\n",
    "\n",
    "# URL for the MistralAI api we'll be using\n",
    "os.environ[\"BASE_URL\"] = \"http://195.242.25.198:8000/v1\"\n",
    "os.environ[\"API_KEY\"] = \"dummy_key\"\n",
    "\n",
    "# Set the max tokens for the models and how many parallel requests to make in Weave Evaluations\n",
    "os.environ[\"MAX_TOKENS\"] = \"4096\"\n",
    "os.environ[\"WEAVE_PARALLELISM\"] = \"2\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import asyncio\n",
    "import logging\n",
    "\n",
    "# Start of workout\n",
    "from utils import Problem, async_client, STRONG_LLM, format_response, check_correctness"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Challenges Dataset\n",
    "We will use the **practice** dataset from the **2023** [HackerCup dataset](https://huggingface.co/datasets/hackercupai/hackercup).\n",
    "\n",
    "We have already processed the dataset and saved it as a [`weave.Dataset`](https://weave-docs.wandb.ai/guides/core-types/datasets/?utm_source=colab&utm_medium=code&utm_campaign=lightning-ai-hacker-cup). You can either use the Dataset by running the next cell or download the dataset using the instructions below.\n",
    "\n",
    "We will use this challenge dataset to load some practice problems and solutions from the HackerCup dataset and evaluate our agents on it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get dataset\n",
    "practice_dataset_uri = \"weave:///parambharat/hackercup/object/practice_dataset:R35fXf9N3FE2IOesg7bRPaPAxiE9YbpirhXO9HcHs8w\"\n",
    "problems_dataset = weave.ref(practice_dataset_uri).get().rows[:]\n",
    "problems = list(map(lambda x: Problem(**x), problems_dataset))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's define what we expect as a solution:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydantic import BaseModel, Field\n",
    "\n",
    "class Solution(BaseModel):\n",
    "    core_question: str = Field(..., description=\"Core question of the problem\")\n",
    "    problem_solving_info: str = Field(..., description=\"Problem-solving information related to the core question\")\n",
    "    plan: str = Field(..., description=\"Step by step plan to solve the problem\")\n",
    "    pseudocode: str = Field(..., description=\"Pseudocode to solve the problem\")\n",
    "    source_code: str = Field(..., description=\"Valid Python3 sourcecode to solve the problem.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## One Shot Solver with retires\n",
    "\n",
    "We will use a [llamaIndex workflow](https://docs.llamaindex.ai/en/stable/understanding/workflows/) based approach to solve the problem.\n",
    "\n",
    "> A workflow is an event-driven, step-based way to control the execution flow of an application.\n",
    "\n",
    "This structure is very flexible and works perfectly with `weave.op` decorators that we will use to define our steps.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "system_prompt = \"\"\"\n",
    "You are a world-class competitive programmer tasked with solving a programming problem. \n",
    "You will be provided with a problem statement, and you need to create a Python3 solution for it. \n",
    "Your task it to develop a winning solution to the problem in Python3 programming language.\n",
    "You will do this in a step-by-step manner.\n",
    "\n",
    "Step 1: Extract the core question and the problem-solving information from the problem statement.\n",
    "Step 2: Generate a step by step plan to solve the problem.\n",
    "Step 3: Generate the pseudocode to solve the problem.\n",
    "Step 4: Write the final solution in Python3 programming language to solve the problem.\n",
    "\n",
    "Competition Guidelines:\n",
    "    a. Do not use any external libraries; stick to Python 3 standard library\n",
    "    b. Handle input and output using standard input/output (stdin/stdout)\n",
    "    c. Use helper functions to improve readability of the code.\n",
    "    c. Use the `input()` function to take input from stdin and print the output to stdout.\n",
    "    d. Do not add extra print statements otherwise it will fail the test cases.\n",
    "    e. Make sure your code passes all potential test cases, including edge cases\n",
    "    f. Follow the input/output format specified in the problem statement and the sample test cases.\"\"\"\n",
    "\n",
    "prompt_template = \"\"\"\n",
    "Let's think step by step to solve the problem:\n",
    "\n",
    "Problem: \n",
    "{problem_description}\n",
    "\n",
    "Input: \n",
    "{sample_input}\n",
    "\n",
    "Output: \n",
    "{sample_output}\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's define the events that our workflow will emit.\n",
    "> The events attributes are the outputs of the steps in our workflow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.workflow import (\n",
    "    StartEvent,\n",
    "    StopEvent,\n",
    "    Workflow,\n",
    "    step,\n",
    "    Event,\n",
    "    Context,\n",
    ")\n",
    "\n",
    "class SetupEvent(Event):\n",
    "    problem: Problem\n",
    "    test_report: str = None\n",
    "\n",
    "class SolvedProblemEvent(Event):\n",
    "    problem: Problem\n",
    "    problem_solution: str\n",
    "\n",
    "class FormattedSolutionEvent(Event):\n",
    "    problem: Problem\n",
    "    solution: Solution\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, we can define our workflow.\n",
    "\n",
    "Each step in the workflow is a function decorated with `@weave.op` that takes the current context and the event as input and returns the next event.\n",
    "\n",
    "The `@step` decorator is used to mark the function as a step in the workflow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class OneShotSolverWorkflow(Workflow):\n",
    "\n",
    "    def __init__(self, retries: int = 2, temperature: float = 0.7, code_execution_timeout: int = 30, **kwargs):   \n",
    "        super().__init__(**kwargs)\n",
    "        self.retries = retries\n",
    "        self.temperature = temperature\n",
    "        self.code_execution_timeout = code_execution_timeout\n",
    "\n",
    "    @step\n",
    "    @weave.op\n",
    "    async def setup(self, ctx: Context, ev: StartEvent) -> SetupEvent:\n",
    "        problem = ev.problem\n",
    "        logging.info(f\"Solving problem: {problem.problem_name}\")\n",
    "        messages=[\n",
    "            {\"role\": \"system\", \"content\": ev.system_prompt},\n",
    "            {\"role\": \"user\", \"content\": ev.prompt_template.format(\n",
    "                    problem_description=problem.problem_description,\n",
    "                    sample_input=problem.sample_input,\n",
    "                    sample_output=problem.sample_output)}\n",
    "            ]\n",
    "        await ctx.set(\"messages\", messages)\n",
    "        return SetupEvent(problem=problem)\n",
    "    \n",
    "    @step\n",
    "    @weave.op\n",
    "    async def generate_code(self, ctx: Context, ev: SetupEvent) -> SolvedProblemEvent:\n",
    "        messages = await ctx.get(\"messages\")\n",
    "        if ev.test_report:\n",
    "            messages.append({\"role\": \"user\", \"content\": f\"Let's try again. The previous solution was incorrect:\\n {ev.test_report}\"})\n",
    "        logging.info(\"Calling model to solve the problem\")\n",
    "        model_output = await async_client.chat.completions.create(\n",
    "            model=STRONG_LLM,\n",
    "            messages=messages,\n",
    "            temperature=self.temperature,\n",
    "            response_model=None\n",
    "        )\n",
    "        problem_solution = model_output.choices[0].message.content\n",
    "        messages.append({\"role\": \"assistant\", \"content\": problem_solution})\n",
    "        await ctx.set(\"messages\", messages)\n",
    "        return SolvedProblemEvent(problem=ev.problem, problem_solution=problem_solution)\n",
    "\n",
    "    @step\n",
    "    @weave.op\n",
    "    async def format_solution(self, ev: SolvedProblemEvent) -> FormattedSolutionEvent:\n",
    "        logging.info(\"Formatting the response\")\n",
    "        solution = await format_response(ev.problem_solution, Solution)\n",
    "        return FormattedSolutionEvent(problem=ev.problem, solution=solution)\n",
    "\n",
    "    @step\n",
    "    @weave.op\n",
    "    async def check_solution(self, ev: FormattedSolutionEvent) -> StopEvent:\n",
    "        logging.info(\"Checking if the code is correct\")\n",
    "        test_report = await check_correctness(\n",
    "            ev.solution.source_code,\n",
    "            ev.problem.sample_input,\n",
    "            ev.problem.sample_output,\n",
    "            timeout=self.code_execution_timeout,\n",
    "        )\n",
    "        logging.info(f\"Test report: {test_report}\")\n",
    "        if (test_report.status != \"passed\") and self.retries > 0:\n",
    "            logging.info(f\"Retrying the solution. Retries left: {self.retries}\")\n",
    "            self.retries -= 1\n",
    "            return SetupEvent(problem=ev.problem, test_report=test_report.message)\n",
    "        else:\n",
    "            return StopEvent(result={\"solution\": ev.solution, \"test_report\": test_report})\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evaluation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we are ready to evaluate against the expected solutions.\n",
    "\n",
    "### Create a Weave Model\n",
    "First we create a Weave [\"Model\"](https://weave-docs.wandb.ai/guides/core-types/models?utm_source=colab&utm_medium=code&utm_campaign=lightning-ai-hacker-cup), which has a `predict` function that Weave Evaluations will call to generate a solution. It also has various attributes that we can set to adjust the behaviour of our pipeline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class OneShotSolverWithRetries(weave.Model):\n",
    "    llm_model: str = STRONG_LLM\n",
    "    system_prompt: str = system_prompt\n",
    "    prompt_template: str = prompt_template\n",
    "    temperature: float = 0.7\n",
    "    code_execution_timeout: int = 30\n",
    "    retries: int = 2\n",
    "\n",
    "    @weave.op\n",
    "    async def predict(self, problem: dict):\n",
    "        workflow = OneShotSolverWorkflow(\n",
    "            temperature=self.temperature,\n",
    "            retries=self.retries,\n",
    "            code_execution_timeout=self.code_execution_timeout,\n",
    "            timeout=600, # timeout for the entire workflow\n",
    "            \n",
    "        )\n",
    "        problem_obj = Problem(**problem)\n",
    "        \n",
    "        result = await workflow.run(\n",
    "            problem=problem_obj,\n",
    "            system_prompt=self.system_prompt,\n",
    "            prompt_template=self.prompt_template,\n",
    "        )\n",
    "        \n",
    "        return result\n",
    "\n",
    "model = OneShotSolverWithRetries()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create the Evals Dataset and a Scorer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We expect the output of the \"test_report\" from our agent above to be `\"passed\"` if the solution is correct. You can think of `expected_result` in the `evals_dataset` as the label that the `test_report` from our solver needs to return in order to ensure the generated solution is correct. In this case the scoring is actually happening in our agentic pipeline as the agent needs to know the result so it can decide whether or not to retry.\n",
    "\n",
    "Weave Evaluations expects data formatted as a list of dictionaries for the evaluation dataset. We dump `problem` as a dictionary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "evals_dataset = [{\"problem\": problem.model_dump(), \"expected_result\": \"passed\"} for problem in problems]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Weave Evaluations use a scorer function that returns a metric and its result in a dict. Here we define a metric that checks if the code generated by agent passed the test case"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@weave.op\n",
    "def scorer(expected_result: str, model_output: dict) -> dict:\n",
    "    if model_output is None or model_output[\"test_report\"].status is None:\n",
    "        return {\"solution_passed\": False}\n",
    "    return {\"solution_passed\": expected_result == model_output[\"test_report\"].status} # check if the test_report status == passed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluator = weave.Evaluation(dataset=evals_dataset, scorers=[scorer], trials=1)\n",
    "\n",
    "results = await evaluator.evaluate(model)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "include_colab_link": true,
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}