Untitled

import openai
import os

openai.api_key = os.environ["OPENAI_API_KEY"]
model = "gpt-3.5-turbo"
messages = [
    {
        "role": "system",
        "content": """
Infer the types of variables, parameters and return values in the following python code.
""",
    },
    {
        "role": "user",
        "content": """
def min(xs):
    m = None
    for x in xs:
        if m is None or x < m:
        m = x
    return m

a = min([1, 4, 2, 8, 5, 7])
""",
    },
    {
        "role": "assistant",
        "content": """
min.xs\tIterable[int|float]
min.return\tint|float
min.s\tint|float
a\tint
""",
    },
    {
        "role": "user",
        "content": """
def squared(xs):
    for x in xs:
        yield x * x

xs = squared([1,2,3])
""",
    },
]

response = openai.ChatCompletion.create(model=model, messages=messages)
generated_text = getattr(response, "choices")[0]["message"]["content"]
print(generated_text)