{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ff022957-2e81-4ea9-84d3-e52d5753e133",
   "metadata": {},
   "source": [
    "### Comment and Unit Test Generater \n",
    "\n",
    "The requirement: \n",
    "* use an LLM to generate docstring and comments for Python code\n",
    "* use an LLM to generate unit test\n",
    "\n",
    "This is my week 4 day 5 project."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea1841f6-4afc-4d29-ace8-5ca5a3915c8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# imports\n",
    "\n",
    "import os\n",
    "import io\n",
    "import sys\n",
    "import json\n",
    "import requests\n",
    "from dotenv import load_dotenv\n",
    "from openai import OpenAI\n",
    "import google.generativeai\n",
    "import anthropic\n",
    "from IPython.display import Markdown, display, update_display\n",
    "import gradio as gr\n",
    "import subprocess\n",
    "from huggingface_hub import login, InferenceClient\n",
    "from transformers import AutoTokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11957fd3-6c61-4496-aef1-8223cb9ec4ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "# environment\n",
    "\n",
    "load_dotenv()\n",
    "os.environ['OPENAI_API_KEY'] = os.getenv('OPENAI_API_KEY', 'your-key-if-not-using-env')\n",
    "os.environ['ANTHROPIC_API_KEY'] = os.getenv('ANTHROPIC_API_KEY', 'your-key-if-not-using-env')\n",
    "os.environ['HF_TOKEN'] = os.getenv('HF_TOKEN', 'your-key-if-not-using-env')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee7b08fd-e678-4234-895e-4e3a925e60f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize\n",
    "\n",
    "openai = OpenAI()\n",
    "claude = anthropic.Anthropic()\n",
    "OPENAI_MODEL = \"gpt-4o\"\n",
    "CLAUDE_MODEL = \"claude-3-5-sonnet-20240620\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8023255-9c98-4fbc-92e4-c553bed3b605",
   "metadata": {},
   "outputs": [],
   "source": [
    "hf_token = os.environ['HF_TOKEN']\n",
    "login(hf_token, add_to_git_credential=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8ce3f5e-74c4-4d35-bfbc-91c5be85e094",
   "metadata": {},
   "outputs": [],
   "source": [
    "code_qwen = \"Qwen/CodeQwen1.5-7B-Chat\"\n",
    "CODE_QWEN_URL = \"https://g39mbjooiiwkbgyz.us-east-1.aws.endpoints.huggingface.cloud\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1bbc66b6-52ae-465e-a368-edc8f097fe9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def system_prompt_for_comment():\n",
    "    system=\"\"\"\n",
    "    You are a Python documentation expert. When writing documentation:\n",
    "    - Follow PEP 257 and Google docstring style guidelines\n",
    "    - Write clear, concise explanations\n",
    "    - Include practical examples\n",
    "    - Highlight edge cases and limitations\n",
    "    - Use type hints in docstrings\n",
    "    - Add inline comments only for complex logic\n",
    "    - Never skip documenting parameters or return values\n",
    "    - Validate that all documentation is accurate and complete\n",
    "    \"\"\"\n",
    "    return system"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b089f87b-53ae-40ad-8d06-b9924bb998a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def system_prompt_for_unit_test():\n",
    "    system=\"\"\"\n",
    "    You are an expert Python testing engineer who specializes in creating comprehensive unit tests. Follow these principles:\n",
    "    - Use pytest as the testing framework\n",
    "    - Follow the Arrange-Act-Assert pattern\n",
    "    - Test both valid and invalid inputs\n",
    "    - Include edge cases and boundary conditions\n",
    "    - Write descriptive test names that explain the scenario being tested\n",
    "    - Create independent tests that don't rely on each other\n",
    "    - Use appropriate fixtures and parametrize when needed\n",
    "    - Add clear comments explaining complex test logic\n",
    "    - Cover error cases and exceptions\n",
    "    - Achieve high code coverage while maintaining meaningful tests\n",
    "    \"\"\"\n",
    "    return system"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22193622-f3a0-4894-a6c4-eb6d88097861",
   "metadata": {},
   "outputs": [],
   "source": [
    "def user_prompt_for_comment(code):\n",
    "    user = f\"\"\"\n",
    "        Please document this Python code with:\n",
    "        \n",
    "        1. A docstring containing:\n",
    "        - A clear description of purpose and functionality\n",
    "        - All parameters with types and descriptions\n",
    "        - Return values with types\n",
    "        - Exceptions that may be raised\n",
    "        - Any important notes or limitations\n",
    "        \n",
    "        2. Strategic inline comments for:\n",
    "        - Complex algorithms or business logic\n",
    "        - Non-obvious implementation choices\n",
    "        - Performance considerations\n",
    "        - Edge cases\n",
    "        \n",
    "        Here's the code to document:\n",
    "        \\n{code}\n",
    "        \"\"\"\n",
    "    return user;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81e61752-ec2f-44c1-86a2-ff3234a0358c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def user_prompt_for_unit_test(code):\n",
    "    user = f\"\"\"\n",
    "        Please generate unit tests for the following Python code. Include:\n",
    "        \n",
    "        1. Test cases for:\n",
    "        - Normal/expected inputs\n",
    "        - Edge cases and boundary values\n",
    "        - Invalid inputs and error conditions\n",
    "        - Different combinations of parameters\n",
    "        - All public methods and functions\n",
    "        \n",
    "        2. For each test:\n",
    "        - Clear test function names describing the scenario\n",
    "        - Setup code (fixtures if needed)\n",
    "        - Test data preparation\n",
    "        - Expected outcomes\n",
    "        - Assertions checking results\n",
    "        - Comments explaining complex test logic\n",
    "        \n",
    "        3. Include any necessary:\n",
    "        - Imports\n",
    "        - Fixtures\n",
    "        - Mock objects\n",
    "        - Helper functions\n",
    "        - Test data generators\n",
    "        \n",
    "        Here's the code to test:\n",
    "        \\n{code}\n",
    "    \"\"\"\n",
    "    return user"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f31ceed3-0eb2-4962-ab86-2d0302185560",
   "metadata": {},
   "outputs": [],
   "source": [
    "pi = \"\"\"\n",
    "import time\n",
    "\n",
    "def calculate(iterations, param1, param2):\n",
    "    result = 1.0\n",
    "    for i in range(1, iterations+1):\n",
    "        j = i * param1 - param2\n",
    "        result -= (1/j)\n",
    "        j = i * param1 + param2\n",
    "        result += (1/j)\n",
    "    return result\n",
    "\n",
    "start_time = time.time()\n",
    "result = calculate(100_000_000, 4, 1) * 4\n",
    "end_time = time.time()\n",
    "\n",
    "print(f\"Result: {result:.12f}\")\n",
    "print(f\"Execution Time: {(end_time - start_time):.6f} seconds\")\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "192c30f5-4be6-49b7-a054-11bfcffa91e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "exec(pi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4e920dc-4094-42d8-9255-18f2919df2d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def messages_for_comment(python):\n",
    "    return [\n",
    "        {\"role\": \"system\", \"content\": system_prompt_for_comment()},\n",
    "        {\"role\": \"user\", \"content\": user_prompt_for_comment(python)}\n",
    "    ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77500cae-bf84-405c-8b03-2f984108951b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def messages_for_unit_test(python):\n",
    "    return [\n",
    "        {\"role\": \"system\", \"content\": system_prompt_for_unit_test()},\n",
    "        {\"role\": \"user\", \"content\": user_prompt_for_unit_test(python)}\n",
    "    ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ec58bf1-4a44-4c21-a71a-2cac359884e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def stream_comment_gpt(code):\n",
    "    stream = openai.chat.completions.create(model=OPENAI_MODEL, messages=messages_for_comment(code), stream=True)\n",
    "    reply = \"\"\n",
    "    for chunk in stream:\n",
    "        fragment = chunk.choices[0].delta.content or \"\"\n",
    "        reply += fragment\n",
    "        #print(fragment, end='', flush=True)\n",
    "        yield reply.replace('```','') \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47c615e2-4eb6-4ce1-ad09-7f2e6dbc3934",
   "metadata": {},
   "outputs": [],
   "source": [
    "stream_comment_gpt(pi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b990875-31fd-40e5-bc8c-f6099d362249",
   "metadata": {},
   "outputs": [],
   "source": [
    "def stream_unit_test_gpt(code):\n",
    "    stream = openai.chat.completions.create(model=OPENAI_MODEL, messages=messages_for_unit_test(code), stream=True)\n",
    "    reply = \"\"\n",
    "    for chunk in stream:\n",
    "        fragment = chunk.choices[0].delta.content or \"\"\n",
    "        reply += fragment\n",
    "        #print(fragment, end='', flush=True)\n",
    "        yield reply.replace('```','')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3dc90578-4f5e-47f1-b30f-c21b5795e82f",
   "metadata": {},
   "outputs": [],
   "source": [
    "stream_unit_test_gpt(pi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17380c0f-b851-472b-a234-d86f5c219e50",
   "metadata": {},
   "outputs": [],
   "source": [
    "def stream_comment_claude(code):\n",
    "    result = claude.messages.stream(\n",
    "        model=CLAUDE_MODEL,\n",
    "        max_tokens=2000,\n",
    "        system=system_prompt_for_comment(),\n",
    "        messages=[{\"role\": \"user\", \"content\": user_prompt_for_comment(code)}],\n",
    "    )\n",
    "    reply = \"\"\n",
    "    with result as stream:\n",
    "        for text in stream.text_stream:\n",
    "            reply += text\n",
    "            #print(text, end=\"\", flush=True)\n",
    "            yield reply.replace('```','')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a2d016d-76a2-4752-bd4d-6f93ddec46be",
   "metadata": {},
   "outputs": [],
   "source": [
    "def stream_unit_test_claude(code):\n",
    "    result = claude.messages.stream(\n",
    "        model=CLAUDE_MODEL,\n",
    "        max_tokens=2000,\n",
    "        system=system_prompt_for_unit_test(),\n",
    "        messages=[{\"role\": \"user\", \"content\": user_prompt_for_unit_test(code)}],\n",
    "    )\n",
    "    reply = \"\"\n",
    "    with result as stream:\n",
    "        for text in stream.text_stream:\n",
    "            reply += text\n",
    "            #print(text, end=\"\", flush=True)\n",
    "            yield reply.replace('```','')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee43428e-b577-4e95-944d-399f2f3b89ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "stream_comment_claude(pi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0565e33b-9f14-48b7-ae8d-d22dc03b93c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "stream_unit_test_claude(pi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f13b3a5b-366d-4b28-adda-977a313e6b4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def stream_comment_model(model, model_url, code):\n",
    "    tokenizer = AutoTokenizer.from_pretrained(model)\n",
    "    messages = messages_for_comment(code)\n",
    "    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n",
    "    client = InferenceClient(model_url, token=hf_token)\n",
    "    stream = client.text_generation(text, stream=True, details=True, max_new_tokens=5000)\n",
    "    result = \"\"\n",
    "    for r in stream:\n",
    "        #print(r.token.text, end = \"\")\n",
    "        result += r.token.text\n",
    "        yield result    \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2efdb92-fc7a-4952-ab46-ae942cb996bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "def stream_unit_test_model(model, model_url, code):\n",
    "    tokenizer = AutoTokenizer.from_pretrained(model)\n",
    "    messages = messages_for_unit_test(code)\n",
    "    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n",
    "    client = InferenceClient(model_url, token=hf_token)\n",
    "    stream = client.text_generation(text, stream=True, details=True, max_new_tokens=3000)\n",
    "    result = \"\"\n",
    "    for r in stream:\n",
    "        #print(r.token.text, end = \"\")\n",
    "        result += r.token.text\n",
    "        yield result    \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a756193-fcba-43da-a981-203c10d36488",
   "metadata": {},
   "outputs": [],
   "source": [
    "stream_comment_model(code_qwen, CODE_QWEN_URL, pi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12ddcbf4-6286-47a8-847b-5be78e7aa995",
   "metadata": {},
   "outputs": [],
   "source": [
    "stream_unit_test_model(code_qwen, CODE_QWEN_URL, pi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "321609ee-b64a-44fc-9090-39f87e1f8e0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def comment_code(python, model):\n",
    "    if model==\"GPT\":\n",
    "        result = stream_comment_gpt(python)\n",
    "    elif model==\"Claude\":\n",
    "        result = stream_comment_claude(python)\n",
    "    elif model==\"CodeQwen\":\n",
    "        result = stream_comment_model(code_qwen, CODE_QWEN_URL, python)\n",
    "    else:\n",
    "        raise ValueError(\"Unknown model\")\n",
    "    for stream_so_far in result:\n",
    "        yield stream_so_far     "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4c560c9-922d-4893-941f-42893373b1be",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_unit_test(python, model):\n",
    "    if model==\"GPT\":\n",
    "        result = stream_unit_test_gpt(python)\n",
    "    elif model==\"Claude\":\n",
    "        result = stream_unit_test_claude(python)\n",
    "    elif model==\"CodeQwen\":\n",
    "        result = stream_unit_test_model(code_qwen, CODE_QWEN_URL, python)\n",
    "    else:\n",
    "        raise ValueError(\"Unknown model\")\n",
    "    for stream_so_far in result:\n",
    "        yield stream_so_far     "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f85bc777-bebe-436b-88cc-b9ecdb6306c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "css = \"\"\"\n",
    ".python {background-color: #306998;}\n",
    ".cpp {background-color: #050;}\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee27cc91-81e6-42c8-ae3c-c04161229d8c",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "with gr.Blocks(css=css) as ui:\n",
    "    gr.Markdown(\"## Convert code from Python to C++\")\n",
    "    with gr.Row():\n",
    "        python = gr.Textbox(label=\"Python code:\", value=pi, lines=10)\n",
    "        result = gr.Textbox(label=\"Result code:\", lines=10)\n",
    "    with gr.Row():\n",
    "        model = gr.Dropdown([\"GPT\", \"Claude\",\"CodeQwen\"], label=\"Select model\", value=\"GPT\")\n",
    "    with gr.Row():\n",
    "        comment_button = gr.Button(\"Comment code\")\n",
    "    with gr.Row():\n",
    "        unit_test_button = gr.Button(\"Unit Test code\")\n",
    "        \n",
    "    comment_button.click(comment_code, inputs=[python, model], outputs=[result])\n",
    "    unit_test_button.click(get_unit_test, inputs=[python, model], outputs=[result])\n",
    "ui.launch(inbrowser=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06e8279c-b488-4807-9bed-9d26be11c057",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}