main.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# CodeQuest Tutorial\n",
    "***"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "This notebook is a tutorial for CodeQUEST, a novel framework leveraging Large Language Models (LLMs) to evaluate and enhance code quality across multiple dimensions including readability, maintainability, efficiency, and security.  \n",
    "\n",
    "The framework is divided into two main components: an Evaluator, which assesses code quality across ten dimensions providing both quantitative scores and qualitative summaries, and an Optimizer, which iteratively improves the code based on feedback from the Evaluator.  \n",
    "\n",
    "Our study demonstrates that CodeQUEST can effectively evaluate code quality, with its assessments aligning closely with established code quality metrics.  \n",
    "\n",
    "This highlights the potential of LLMs in automating code quality evaluation and improvement processes, presenting a significant advancement toward enhancing software development practices.  \n",
    "\n",
    "Note that evaluation scores may vary slightly across runs, due to the inherent stochasticity of the LLM. Increased output stability can be achieved by increasing the number of retries per evaluation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##  Loading \n",
    "***"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "class Node: \n",
      "\tdef __init__(self, data): \n",
      "\t\tself.data = data \n",
      "\t\tself.left = None\n",
      "\t\tself.right = None\n",
      "def max_height(node): \n",
      "\tif node is None: \n",
      "\t\treturn 0 ; \n",
      "\telse : \n",
      "\t\tleft_height = max_height(node.left) \n",
      "\t\tright_height = max_height(node.right) \n",
      "\t\tif (left_height > right_height): \n",
      "\t\t\treturn left_height+1\n",
      "\t\telse: \n",
      "\t\t\treturn right_height+1\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "# Load the environment variables if using a .env file at root\n",
    "# from dotenv import load_dotenv \n",
    "# load_dotenv()\n",
    "\n",
    "import os \n",
    "from codequest.evaluator import Evaluator, CodeQUESTEvaluator \n",
    "from codequest.optimizer import Optimizer, CodeQUESTOptimizer \n",
    "from codequest.codequest import format_dimwise_feedback_for_improver, code_tester, check_syntax, QUESTer \n",
    "\n",
    "script_path = \"examples/python/code.py\"\n",
    "code = open(script_path, 'r').read().strip() \n",
    "testcases_path = \"tests/test_cases/python/code.py\"\n",
    "print(code)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluation "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section, we create a *baseline evaluator* based on Chain-of-Thought prompting technique and we also create a *codeQUEST evaluator* based on our novel approach. \n",
    "We compare the performance of both evaluations on an example script in the zero-shot setting. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "baseline_evaluator = Evaluator(num_retries=1)\n",
    "codequest_evaluator = CodeQUESTEvaluator(num_retries=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "baseline_eval = baseline_evaluator(code)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>code_score</th>\n",
       "      <th>code_insight</th>\n",
       "      <th>code_runtime</th>\n",
       "      <th>code_runcost</th>\n",
       "      <th>code_scores</th>\n",
       "      <th>code_insights</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>4.0</td>\n",
       "      <td>The code is well-structured and correctly impl...</td>\n",
       "      <td>2</td>\n",
       "      <td>0.002</td>\n",
       "      <td>[4]</td>\n",
       "      <td>[The code is well-structured and correctly imp...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   code_score                                       code_insight  \\\n",
       "0         4.0  The code is well-structured and correctly impl...   \n",
       "\n",
       "   code_runtime  code_runcost code_scores  \\\n",
       "0             2         0.002         [4]   \n",
       "\n",
       "                                       code_insights  \n",
       "0  [The code is well-structured and correctly imp...  "
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "baseline_eval['report']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "codequest_eval = codequest_evaluator(code)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>quality_dimension</th>\n",
       "      <th>dimension_score</th>\n",
       "      <th>dimension_insights</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Readability</td>\n",
       "      <td>1</td>\n",
       "      <td>[The code is generally readable with descripti...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Maintainability</td>\n",
       "      <td>5</td>\n",
       "      <td>[The code is logically organized and easy to u...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Testability</td>\n",
       "      <td>4</td>\n",
       "      <td>[The code provided is well-structured for test...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Efficiency</td>\n",
       "      <td>2</td>\n",
       "      <td>[The code provided is a basic implementation o...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Robustness</td>\n",
       "      <td>-4</td>\n",
       "      <td>[The code provided is a basic implementation o...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>Security</td>\n",
       "      <td>1</td>\n",
       "      <td>[The provided code is a simple implementation ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>Documentation</td>\n",
       "      <td>-5</td>\n",
       "      <td>[The code lacks documentation. There are no co...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>Modularity</td>\n",
       "      <td>-2</td>\n",
       "      <td>[The code provided is simple and functional bu...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>Scalability</td>\n",
       "      <td>-5</td>\n",
       "      <td>[The code provided is a basic implementation o...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>Portability</td>\n",
       "      <td>4</td>\n",
       "      <td>[The code provided is highly portable as it do...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  quality_dimension  dimension_score  \\\n",
       "0       Readability                1   \n",
       "1   Maintainability                5   \n",
       "2       Testability                4   \n",
       "3        Efficiency                2   \n",
       "4        Robustness               -4   \n",
       "5          Security                1   \n",
       "6     Documentation               -5   \n",
       "7        Modularity               -2   \n",
       "8       Scalability               -5   \n",
       "9       Portability                4   \n",
       "\n",
       "                                  dimension_insights  \n",
       "0  [The code is generally readable with descripti...  \n",
       "1  [The code is logically organized and easy to u...  \n",
       "2  [The code provided is well-structured for test...  \n",
       "3  [The code provided is a basic implementation o...  \n",
       "4  [The code provided is a basic implementation o...  \n",
       "5  [The provided code is a simple implementation ...  \n",
       "6  [The code lacks documentation. There are no co...  \n",
       "7  [The code provided is simple and functional bu...  \n",
       "8  [The code provided is a basic implementation o...  \n",
       "9  [The code provided is highly portable as it do...  "
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "codequest_eval['report']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.1"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "codequest_eval['score']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Noticed how the baseline eval based on Chain-of-Thought over-estimate the code quality*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optimization "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section, we create a *baseline optimizer* based on Chain-of-Thought prompting technique and we also create a *codeQUEST optimizer* based on our novel approach. \n",
    "We compare the optimized code on an example script in the zero-shot setting. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "baseline_optimizer = Optimizer()\n",
    "codequest_optimizer = CodeQUESTOptimizer() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "class Node:\n",
      "    def __init__(self, data):\n",
      "        \"\"\"\n",
      "        Initialize a new Node with given data.\n",
      "        \"\"\"\n",
      "        self.data = data\n",
      "        self.left = None\n",
      "        self.right = None\n",
      "\n",
      "def max_height(node):\n",
      "    \"\"\"\n",
      "    Calculate the maximum height of a binary tree.\n",
      "    \n",
      "    Args:\n",
      "    node (Node): The root node of the binary tree.\n",
      "    \n",
      "    Returns:\n",
      "    int: The maximum height of the tree.\n",
      "    \"\"\"\n",
      "    if not isinstance(node, Node) and node is not None:\n",
      "        raise ValueError(\"Input must be a Node object or None\")\n",
      "    \n",
      "    if node is None:\n",
      "        return 0\n",
      "    else:\n",
      "        # Recursively calculate the height of the left and right subtrees\n",
      "        left_subtree_height = max_height(node.left)\n",
      "        right_subtree_height = max_height(node.right)\n",
      "        \n",
      "        # Return the greater height between the two subtrees, plus one for the current node\n",
      "        if left_subtree_height > right_subtree_height:\n",
      "            return left_subtree_height + 1\n",
      "        else:\n",
      "            return right_subtree_height + 1\n",
      "\n"
     ]
    }
   ],
   "source": [
    "baseline_improved_code = baseline_optimizer(code, baseline_eval['insight'])['code']\n",
    "print(baseline_improved_code)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "class Node:\n",
      "    \"\"\"\n",
      "    A class to represent a node in a binary tree.\n",
      "    \n",
      "    Attributes:\n",
      "    data : any type\n",
      "        The data stored in the node.\n",
      "    left : Node\n",
      "        The left child node.\n",
      "    right : Node\n",
      "        The right child node.\n",
      "    \"\"\"\n",
      "    def __init__(self, data):\n",
      "        self.data = data\n",
      "        self.left = None\n",
      "        self.right = None\n",
      "\n",
      "def max_height(node, memo=None):\n",
      "    \"\"\"\n",
      "    Calculate the maximum height of a binary tree.\n",
      "    \n",
      "    Parameters:\n",
      "    node : Node\n",
      "        The root node of the binary tree.\n",
      "    memo : dict, optional\n",
      "        A dictionary to store the heights of subtrees to avoid repeated calculations.\n",
      "    \n",
      "    Returns:\n",
      "    int\n",
      "        The maximum height of the binary tree.\n",
      "    \"\"\"\n",
      "    if memo is None:\n",
      "        memo = {}\n",
      "    \n",
      "    if node is None:\n",
      "        return 0\n",
      "    \n",
      "    if node in memo:\n",
      "        return memo[node]\n",
      "    \n",
      "    try:\n",
      "        left_height = max_height(node.left, memo)\n",
      "        right_height = max_height(node.right, memo)\n",
      "        height = max(left_height, right_height) + 1\n",
      "        memo[node] = height\n",
      "        return height\n",
      "    except Exception as e:\n",
      "        print(f\"An error occurred: {e}\")\n",
      "        return -1\n",
      "\n",
      "# Example usage:\n",
      "# root = Node(1)\n",
      "# root.left = Node(2)\n",
      "# root.right = Node(3)\n",
      "# print(max_height(root))  # Output should be 2\n",
      "\n"
     ]
    }
   ],
   "source": [
    "codequest_improved_code = codequest_optimizer(\n",
    "    code, \n",
    "    format_dimwise_feedback_for_improver(codequest_eval['report'])\n",
    ")['code']\n",
    "print(codequest_improved_code)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Noticed how the codeQUEST Optimizer achieved much well-rounded improvement*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## CodeQUEST Cycle \n",
    "\n",
    "Putting Evaluator and Optimizer in an Actor-Critic Loop "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "quester = QUESTer(codequest_evaluator, codequest_optimizer) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "quest_result = quester(script_path, testcases_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Final Version "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "class Node:\n",
      "    \"\"\"\n",
      "    A class to represent a node in a binary tree.\n",
      "\n",
      "    Attributes:\n",
      "    data : any type\n",
      "        The data stored in the node.\n",
      "    left : Node\n",
      "        The left child node.\n",
      "    right : Node\n",
      "        The right child node.\n",
      "    \"\"\"\n",
      "    def __init__(self, data):\n",
      "        \"\"\"\n",
      "        Constructs all the necessary attributes for the node object.\n",
      "\n",
      "        Parameters:\n",
      "        data : any type\n",
      "            The data stored in the node.\n",
      "        \"\"\"\n",
      "        self.data = data\n",
      "        self.left = None\n",
      "        self.right = None\n",
      "\n",
      "def max_height(node):\n",
      "    \"\"\"\n",
      "    Calculate the maximum height of a binary tree.\n",
      "\n",
      "    Parameters:\n",
      "    node : Node\n",
      "        The root node of the binary tree.\n",
      "\n",
      "    Returns:\n",
      "    int\n",
      "        The maximum height of the binary tree.\n",
      "    \"\"\"\n",
      "    if not isinstance(node, Node) and node is not None:\n",
      "        raise ValueError(\"Input must be a Node object or None\")\n",
      "\n",
      "    if node is None:\n",
      "        return 0\n",
      "\n",
      "    # Use an iterative approach with a stack to avoid deep recursion\n",
      "    stack = [(node, 1)]\n",
      "    max_height = 0\n",
      "\n",
      "    while stack:\n",
      "        current_node, current_height = stack.pop()\n",
      "        if current_node:\n",
      "            max_height = max(max_height, current_height)\n",
      "            stack.append((current_node.left, current_height + 1))\n",
      "            stack.append((current_node.right, current_height + 1))\n",
      "\n",
      "    return max_height\n",
      "\n"
     ]
    }
   ],
   "source": [
    "trajecs = quest_result['trajectories']\n",
    "final_version = [trajec['code'] for trajec in trajecs if trajec['accepted']][-1]\n",
    "print(final_version)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.4"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}