-
Notifications
You must be signed in to change notification settings - Fork 0
/
prompt.py
73 lines (56 loc) · 1.87 KB
/
prompt.py
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
from sys import argv, exit as abort
from os import getenv as ENV
import openai
openai.api_key = ENV("OPENAI_KEY")
def get_file_content(file_name: str) -> str:
with open(f"{file_name}", "r") as file: return file.read()
def prompt_with(tests:str) -> list[dict]:
return [
{"role": "system", "content": "You are a programmer using TDD"},
{"role": "user", "content": "Which language do you use?"},
{"role": "assistant", "content": "I use python"},
{"role": "user", "content": "Which test framework do you use?"},
{"role": "assistant", "content": "I use pytest"},
{
"role": "user",
"content": f"""
Implement the code tested in the tests provided so that all pass.
Include nothing else in your reply, as it will be written
directly to a single source code file.
The tests are as follows:
```
{tests}
```
"""
}
]
def generate_code(
tests: str,
model: str = "gpt-3.5-turbo",
) -> str:
messages = prompt_with(tests)
response = openai.ChatCompletion.create(
model=model,
messages=messages
)['choices'][0]['message']['content']
return response.replace('```python', '').replace('```', '')
def write_file_content(file_name: str, content: str = "") -> None:
try:
with open(f"{file_name}", "w") as file:
file.write(content)
except:
print("Something has gone wrong... have a look anyways:")
print(content)
def main():
dir_name: str
try:
dir_name = argv[1]
except:
print("You should specify a directory...")
abort()
tests = get_file_content(f"{dir_name}/test.py")
code = generate_code(tests)
print(code)
write_file_content(f"{dir_name}/src.py", code)
print("...And done!")
__name__ == "__main__" and main()