181 lines
7.1 KiB
Plaintext
181 lines
7.1 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "ed8c52b6",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import os\n",
|
|
"from dotenv import load_dotenv\n",
|
|
"from openai import OpenAI\n",
|
|
"import gradio as gr\n",
|
|
"\n",
|
|
"load_dotenv(override=True)\n",
|
|
"\n",
|
|
"openai_api_key = os.getenv('OPENAI_API_KEY')\n",
|
|
"ollama_api_key = os.getenv('OLLAMA_API_KEY')\n",
|
|
"\n",
|
|
"if openai_api_key:\n",
|
|
" print(f\"OpenAI API Key exists and begins {openai_api_key[:8]}\")\n",
|
|
"else:\n",
|
|
" print(\"OpenAI API Key not set\")\n",
|
|
"\n",
|
|
"if ollama_api_key:\n",
|
|
" print(f\"OLLAMA API Key exists and begins {ollama_api_key[:2]}\")\n",
|
|
"else:\n",
|
|
" print(\"OLLAMA API Key not set (and this is optional)\")\n",
|
|
"\n",
|
|
"ollama_url = \"http://localhost:11434/v1\"\n",
|
|
"\n",
|
|
"openai = OpenAI()\n",
|
|
"ollama = OpenAI(api_key=ollama_api_key, base_url=ollama_url)\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 38,
|
|
"id": "c628f95e",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"system_prompt_doc = \"\"\"You are an expert Python developer and code reviewer.\n",
|
|
"Your job is to read the user's provided function, and return:\n",
|
|
"1. A concise, PEP-257-compliant docstring summarizing what the function does, clarifying types, parameters, return values, and side effects.\n",
|
|
"2. Helpful inline comments that improve both readability and maintainability, without restating what the code obviously does.\n",
|
|
"\n",
|
|
"Only output the function, not explanations or additional text. \n",
|
|
"Do not modify variable names or refactor the function logic.\n",
|
|
"Your response should improve the code's clarity and documentation, making it easier for others to understand and maintain.\n",
|
|
"Don't be extremely verbose.\n",
|
|
"Your answer should be at a senior level of expertise.\n",
|
|
"\"\"\"\n",
|
|
"\n",
|
|
"system_prompt_tests = \"\"\"You are a seasoned Python developer and testing expert.\n",
|
|
"Your task is to read the user's provided function, and generate:\n",
|
|
"1. A concise set of meaningful unit tests that thoroughly validate the function's correctness, including typical, edge, and error cases.\n",
|
|
"2. The tests should be written for pytest (or unittest if pytest is not appropriate), use clear, descriptive names, and avoid unnecessary complexity.\n",
|
|
"3. If dependencies or mocking are needed, include minimal necessary setup code (but avoid over-mocking).\n",
|
|
"\n",
|
|
"Only output the relevant test code, not explanations or extra text.\n",
|
|
"Do not change the original function; focus solely on comprehensive, maintainable test coverage that other developers can easily understand and extend.\n",
|
|
"\"\"\""
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 39,
|
|
"id": "4bb84e6c",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"models = [\"gpt-4.1-mini\", \"llama3.1\"]\n",
|
|
"clients = {\"gpt-4.1-mini\": openai, \"llama3.1\": ollama}\n",
|
|
"\n",
|
|
"def generate_documentation(code, model):\n",
|
|
" response = clients[model].chat.completions.create(\n",
|
|
" model=model,\n",
|
|
" messages=[\n",
|
|
" {\"role\": \"system\", \"content\": system_prompt_doc},\n",
|
|
" {\"role\": \"user\", \"content\": code}\n",
|
|
" ],\n",
|
|
" stream=True\n",
|
|
" )\n",
|
|
" output = \"\"\n",
|
|
" for chunk in response:\n",
|
|
" output += chunk.choices[0].delta.content or \"\"\n",
|
|
" yield output.replace(\"```python\", \"\").replace(\"```\", \"\")\n",
|
|
"\n",
|
|
"def generate_tests(code, model):\n",
|
|
" response = clients[model].chat.completions.create(\n",
|
|
" model=model,\n",
|
|
" messages=[\n",
|
|
" {\"role\": \"system\", \"content\": system_prompt_tests},\n",
|
|
" {\"role\": \"user\", \"content\": code}\n",
|
|
" ],\n",
|
|
" stream=True\n",
|
|
" )\n",
|
|
" output = \"\"\n",
|
|
" for chunk in response:\n",
|
|
" output += chunk.choices[0].delta.content or \"\"\n",
|
|
" yield output.replace(\"```python\", \"\").replace(\"```\", \"\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "a4e65b26",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"with gr.Blocks(theme=gr.themes.Soft(spacing_size=gr.themes.sizes.spacing_sm, radius_size=gr.themes.sizes.radius_none)) as ui:\n",
|
|
" gr.Markdown(\"# Python Toolbox\", elem_id=\"app-title\")\n",
|
|
" \n",
|
|
" with gr.Tab(\"Docstring Generator\") as tab1:\n",
|
|
" gr.Markdown(\"## Docstring & Comment Generator\")\n",
|
|
" gr.Markdown(\"Paste your function below to generate helpful docstrings and inline comments!\")\n",
|
|
"\n",
|
|
" with gr.Row():\n",
|
|
" with gr.Column():\n",
|
|
" code_input = gr.Code(label=\"Your Python function here\", lines=20, language=\"python\")\n",
|
|
" model_dropdown = gr.Dropdown(choices=models, value=models[0], label=\"Select model\")\n",
|
|
" submit_doc_btn = gr.Button(\"Generate docstring & comments\")\n",
|
|
" with gr.Column():\n",
|
|
" code_output = gr.Code(label=\"New function with docstring and comments\", language=\"python\")\n",
|
|
"\n",
|
|
" submit_doc_btn.click(\n",
|
|
" generate_documentation, \n",
|
|
" inputs=[code_input, model_dropdown], \n",
|
|
" outputs=code_output\n",
|
|
" )\n",
|
|
"\n",
|
|
" with gr.Tab(\"Unit Tests Generator\") as tab2:\n",
|
|
" gr.Markdown(\"## Unit Test Generator\")\n",
|
|
" gr.Markdown(\"Paste your function below to generate helpful unit tests!\")\n",
|
|
"\n",
|
|
" with gr.Row():\n",
|
|
" with gr.Column():\n",
|
|
" code_input_2 = gr.Code(label=\"Your Python function here\", lines=20, language=\"python\")\n",
|
|
" model_dropdown_2 = gr.Dropdown(choices=models, value=models[0], label=\"Select model\")\n",
|
|
" submit_test_btn = gr.Button(\"Generate unit tests\")\n",
|
|
" with gr.Column():\n",
|
|
" code_output_2 = gr.Code(label=\"Generated unit tests\", language=\"python\")\n",
|
|
"\n",
|
|
" submit_test_btn.click(\n",
|
|
" generate_tests, \n",
|
|
" inputs=[code_input_2, model_dropdown_2], \n",
|
|
" outputs=code_output_2\n",
|
|
" )\n",
|
|
" \n",
|
|
" \n",
|
|
" tab1.select(lambda x: x, inputs=code_input_2, outputs=code_input)\n",
|
|
" tab2.select(lambda x: x, inputs=code_input, outputs=code_input_2)\n",
|
|
"\n",
|
|
"ui.launch(share=False, inbrowser=True)"
|
|
]
|
|
}
|
|
],
|
|
"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.12.12"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 5
|
|
}
|