File size: 4,019 Bytes
41e1fb9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 | import os
from openai import OpenAI
import gradio as gr
import anthropic
from dotenv import load_dotenv
load_dotenv(override=True)
openai_api_key = os.getenv('OPENAI_API_KEY')
anthropic_api_key = os.getenv('ANTHROPIC_API_KEY')
if openai_api_key:
print(f"OpenAI API Key exists and begins {openai_api_key[:8]}")
else:
print("OpenAI API Key not set")
if anthropic_api_key:
print(f"Anthropic API Key exists and begins {anthropic_api_key[:7]}")
else:
print("Anthropic API Key not set")
MODEL = "gpt-4o-mini"
openai = OpenAI()
CLAUDE_MODEL='claude-3-7-sonnet-latest'
claude= anthropic.Anthropic()
system_message = """
You are a Python expert assistant.
Your job is to add a clean, detailed, and professional Python docstring to a given function.
The docstring should:
- Start with a one-line summary of what the function does.
- Include an `Args:` section listing all parameters with type and description.
- Include a `Returns:` section describing the return value with type.
- Follow proper indentation and formatting standards (PEP 257 style).
- Do not modify the function body.
- Do not explain or comment — just return the updated function with the docstring inserted.
Example:
Input:
def multiply(a, b):
return a * b
Output:
def multiply(a, b):
\"\"\"
Multiply two numbers.
Args:
a (int): The first number.
b (int): The second number.
Returns:
int: The product of a and b.
\"\"\"
return a * b
"""
system_message_unittest = """
You are a Python testing assistant.
Your job is to generate clean, correct, and complete unit tests for the given Python function(s) using the built-in `unittest` framework.
Your output should:
- Import the function if needed (assume it's in the same file)
- Use a test class that inherits from `unittest.TestCase`
- Include 2 to 4 relevant test cases using `.assertEqual`, `.assertTrue`, or `.assertFalse`
- Cover edge cases (e.g. empty inputs, zero, negative, large numbers)
- Not explain the code — just return the full Python code block with the unit tests
Example:
Input:
def is_even(n):
return n % 2 == 0
Output:
import unittest
class TestIsEven(unittest.TestCase):
def test_even(self):
self.assertTrue(is_even(4))
self.assertTrue(is_even(0))
def test_odd(self):
self.assertFalse(is_even(3))
if __name__ == "__main__":
unittest.main()
"""
def userPrompt(code):
userPrompt=code
return userPrompt
def generator(code, mode):
response = ""
if mode == "docstring":
response = openai.chat.completions.create(
model=MODEL,
messages=[
{"role":"system","content":system_message},
{"role":"user","content":userPrompt(code)}
]
)
elif mode == "unit_test":
response = openai.chat.completions.create(
model=MODEL,
messages=[
{"role":"system","content":system_message_unittest},
{"role":"user","content":userPrompt(code)}
]
)
gpt_output = response.choices[0].message.content
# Remove triple backticks and 'python' label
if gpt_output.startswith("```python"):
gpt_output = gpt_output.replace("```python", "").strip()
if gpt_output.endswith("```"):
gpt_output = gpt_output[:-3].strip()
# Unescape \n and convert to real line breaks
cleaned = gpt_output.encode().decode('unicode_escape')
return cleaned.strip()
interface = gr.Interface(
fn=generator,
inputs=[
gr.Textbox(lines=12, label="Paste your Python function here"),
gr.Dropdown(choices=["docstring", "unit_test"], label="Select Mode", value="docstring")
],
outputs=gr.Code(label="Output", language="python"),
title="🧠 Python Docstring and Unit Test Generator",
description="Paste your Python function and select the mode to generate a docstring or a unit test automatically."
)
if __name__ == "__main__":
interface.launch() |