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)