gptcode.py

import inspect
from functools import wraps

import openai


def generate_code(signature, docstring):
    init_prompt = "You are a Python expert who can implement the given function."
    definition = f"def {signature}"
    prompt = f"Read this incomplete Python code:\n```python\n{definition}\n```"
    prompt += "\n"
    prompt += f"Complete the Python code that follows this instruction: '{docstring}'. Your response must start with code block '```python'."

    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        temperature=0,
        max_tokens=1024,
        top_p=1,
        messages=[
            {
                "role": "system",
                "content": init_prompt,
            },
            {
                "role": "user",
                "content": prompt,
            },
        ],
    )

    codeblock = response.choices[0].message.content
    code = next(filter(None, codeblock.split("```python"))).rsplit("```", 1)[0]
    code = code.strip()

    return code


def gpt_code(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        signature = f'{func.__name__}({", ".join(inspect.signature(func).parameters)}):'
        docstring = func.__doc__.strip()
        code = generate_code(signature, docstring)
        print(f"generated code:\n```python\n{code}\n```")
        exec(code)
        return locals()[func.__name__](*args, **kwargs)

    return wrapper


if __name__ == "__main__":

    @gpt_code
    def sum_two(arg1, arg2):
        """
        Sum two numbers.
        """

    print(sum_two(2, 3))

    @gpt_code
    def fizzbuzz(n):
        """
        Given an integer n, return a string array answer (1-indexed) where:

        answer[i] == "FizzBuzz" if i is divisible by 3 and 5.
        answer[i] == "Fizz" if i is divisible by 3.
        answer[i] == "Buzz" if i is divisible by 5.
        answer[i] == i (as a string) if none of the above conditions are true.
        """

    print(fizzbuzz(15))