talhafewfefefew/Python_Docstring_and_Unit_Test_Generator

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