Bootcamp/Euclid (Ephraim): Week 4 - Added an AI coding assistant for code comments, unittest and code fixing

week4/community-contributions/edfolmi/week4_exercise_ai_coding_assistant.ipynb

@@ -0,0 +1,243 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "1f4feb01",
+   "metadata": {},
+   "source": [
+    "# AI coding Assistant\n",
+    "\n",
+    "Help you with writing code comments, unittest, and code fixing.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "c1b490f5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# packages\n",
+    "import os\n",
+    "import gradio as gr\n",
+    "from enum import Enum\n",
+    "from dotenv import load_dotenv\n",
+    "from openai import OpenAI"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "3af015df",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Providers\n",
+    "\n",
+    "# Load environment variables\n",
+    "load_dotenv()\n",
+    "\n",
+    "\n",
+    "class Provider(str, Enum):\n",
+    "    OPENROUTER = \"OpenRouter\"\n",
+    "    OPENAI = \"OpenAI\"\n",
+    "    OLLAMA = \"Ollama\"\n",
+    "\n",
+    "\n",
+    "clients = {\n",
+    "    Provider.OPENROUTER: OpenAI(\n",
+    "        api_key=os.getenv(\"OPENROUTER_API_KEY\"),\n",
+    "        base_url=\"https://openrouter.ai/api/v1\"\n",
+    "    ),\n",
+    "    Provider.OPENAI: OpenAI(api_key=os.getenv(\"OPENAI_API_KEY\")),\n",
+    "    Provider.OLLAMA: OpenAI(\n",
+    "        api_key=\"ollama\",\n",
+    "        base_url=\"http://localhost:11434/v1\"\n",
+    "    ),\n",
+    "}\n",
+    "\n",
+    "models = {\n",
+    "    Provider.OPENROUTER: [\"openai/gpt-4o-mini\", \"mistralai/mistral-7b-instruct\"],\n",
+    "    Provider.OPENAI: [\"gpt-4o-mini\", \"gpt-4.1-mini\"],\n",
+    "    Provider.OLLAMA: [\"llama3\", \"deepseek-coder\"],\n",
+    "}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "21dc43fb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Tasks\n",
+    "TASK_OPTIONS = [\n",
+    "    \"Add Comments\",\n",
+    "    \"Explain Code\",\n",
+    "    \"Write Unit Tests\",\n",
+    "    \"Fix Code\"\n",
+    "]\n",
+    "\n",
+    "\n",
+    "def build_prompt(tasks, code):\n",
+    "    instructions = []\n",
+    "\n",
+    "    if \"Add Comments\" in tasks:\n",
+    "        instructions.append(\"Add helpful comments to the code.\")\n",
+    "    if \"Explain Code\" in tasks:\n",
+    "        instructions.append(\"Explain what the code does.\")\n",
+    "    if \"Write Unit Tests\" in tasks:\n",
+    "        instructions.append(\"Write comprehensive unit tests.\")\n",
+    "    if \"Fix Code\" in tasks:\n",
+    "        instructions.append(\"Fix any issues in the code.\")\n",
+    "\n",
+    "    system_prompt = \"You are a senior software engineer.\"\n",
+    "\n",
+    "    user_prompt = f\"\"\"\n",
+    "Tasks:\n",
+    "- {\"\\n- \".join(instructions)}\n",
+    "\n",
+    "Code:\n",
+    "{code}\n",
+    "\"\"\"\n",
+    "\n",
+    "    return [\n",
+    "        {\"role\": \"system\", \"content\": system_prompt},\n",
+    "        {\"role\": \"user\", \"content\": user_prompt}\n",
+    "    ]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "02245980",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Generation\n",
+    "def run_code_assistant(provider, model, tasks, code):\n",
+    "    if not code.strip():\n",
+    "        return \"Please provide code.\"\n",
+    "\n",
+    "    client = clients[Provider(provider)]\n",
+    "\n",
+    "    messages = build_prompt(tasks, code)\n",
+    "\n",
+    "    response = client.chat.completions.create(\n",
+    "        model=model,\n",
+    "        messages=messages,\n",
+    "        temperature=0.3,\n",
+    "    )\n",
+    "\n",
+    "    return response.choices[0].message.content"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "7ea8e54b",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "* Running on local URL:  http://127.0.0.1:7863\n",
+      "* To create a public link, set `share=True` in `launch()`.\n"
+     ]
+    },
+    {
+     "data": {
+      "text/html": [
+       "<div><iframe src=\"http://127.0.0.1:7863/\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/plain": []
+     },
+     "execution_count": 14,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# UI\n",
+    "with gr.Blocks(title=\"AI Coding Assistant\") as ui:\n",
+    "    gr.Markdown(\"# Help you with writing code comments, unittest, and code fixing.\")\n",
+    "    gr.Markdown(\"Perform tasks on your code using different LLM providers.\")\n",
+    "\n",
+    "    with gr.Row():\n",
+    "        provider = gr.Dropdown(\n",
+    "            choices=[p.value for p in Provider],\n",
+    "            value=Provider.OPENROUTER.value,\n",
+    "            label=\"Provider\"\n",
+    "        )\n",
+    "\n",
+    "        model = gr.Dropdown(\n",
+    "            choices=models[Provider.OPENROUTER],\n",
+    "            value=models[Provider.OPENROUTER][0],\n",
+    "            label=\"Model\"\n",
+    "        )\n",
+    "\n",
+    "    def update_models(selected_provider):\n",
+    "        return gr.Dropdown(\n",
+    "            choices=models[Provider(selected_provider)],\n",
+    "            value=models[Provider(selected_provider)][0]\n",
+    "        )\n",
+    "\n",
+    "    provider.change(update_models, provider, model)\n",
+    "\n",
+    "    tasks = gr.CheckboxGroup(\n",
+    "        choices=TASK_OPTIONS,\n",
+    "        value=[\"Add Comments\"],\n",
+    "        label=\"Tasks\"\n",
+    "    )\n",
+    "\n",
+    "    code_input = gr.Code(\n",
+    "        label=\"Code\",\n",
+    "        lines=20\n",
+    "    )\n",
+    "\n",
+    "    run_button = gr.Button(\"Run\")\n",
+    "\n",
+    "    output = gr.Markdown()\n",
+    "\n",
+    "    run_button.click(\n",
+    "        run_code_assistant,\n",
+    "        inputs=[provider, model, tasks, code_input],\n",
+    "        outputs=output\n",
+    "    )\n",
+    "\n",
+    "ui.launch()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv (3.12.12)",
+   "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.12.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}