This repository was archived by the owner on Feb 16, 2025. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 2
Expand file tree
/
Copy pathcodex.py
More file actions
226 lines (180 loc) · 5.81 KB
/
codex.py
File metadata and controls
226 lines (180 loc) · 5.81 KB
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
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
#!/usr/bin/env python
import os
import argparse
from typing import Optional, Callable
import openai
openai.api_key = os.getenv("OPENAI_API_KEY")
import languages
g_source, g_prompt = [], ""
g_last_cmd = None
g_last_cmd_args = list()
g_last_cmd_kwargs = dict()
def get_source() -> Optional[list[str]]:
""" Returns a multiline string representation of the source code.
"""
global g_source
return "\n".join(g_source)
def set_source(source: list[str]):
""" Sets value of the source code.
"""
global g_source
g_source = source
def get_prompt() -> Optional[str]:
""" Returns string representation of the current prompt value.
"""
global g_prompt
return g_prompt
def set_prompt(prompt: str):
""" Sets value of the current prompt.
"""
global g_prompt
g_prompt = prompt
def save(source : str, path : str):
""" Save input string in a file located in given path.
"""
with open(path, "w") as file_:
file_.write(source)
def load(path : str) -> list[str]:
""" Returns source loaded from the file located under given path.
"""
with open(path, "r") as file_:
lines = file_.read()
return lines.splitlines()
def on_save(path: str, *args, **kwargs):
""" Handler for the .save <path> command.
"""
print(f"Saving code to {path} ...")
save(get_source(), path)
def on_load(path: str, *args, **kwargs):
""" Handler for the .load <path> command.
"""
print(f"Loading context from {path} ...")
set_source(load(path))
def on_erase(*args, **kwargs):
""" Handler for the .erase command.
"""
set_source(get_source().splitlines()[:-1])
def on_exit(*args, **kwargs):
""" Handler for the .exit command.
"""
print("Goodbye...")
exit()
def on_help(*args, **kwargs):
""" Handler for the .help command.
Will print available dot commands.
"""
print("""
.save <path> : save current context to the file
.load <path> : replace current context with the content loaded from a file
.erase : remove last line of the current context
.exit : exit from the script
.help : print this help message
. : repeat last command that was executed
""")
input("Press Enter to continue...")
def on_repeat_last_cmd(*args, **kwargs):
""" Handler for the . command.
"""
global g_last_cmd
global g_last_cmd_args
global g_last_cmd_kwargs
if g_last_cmd is None:
print("Command history is empty...")
else:
g_last_cmd(*g_last_cmd_args, **g_last_cmd_kwargs)
COMMANDS = {
".save" : on_save,
".exit" : on_exit,
".load" : on_load,
".erase": on_erase,
".help": on_help,
".": on_repeat_last_cmd
}
def format_prompt(source: str, prompt: Optional[str], language: str) -> str:
""" Format source and prompt into a soure code that will be fed into Codex API for completion.
"""
if not languages.is_language_supported(language):
raise KeyError("Unknown language: {language}...")
if prompt is None:
prompt = ""
context = source
if len(prompt) > 0:
comment = languages.CONFIG[language]["prompt"].format(prompt)
context = f"{source}\n\n{comment}"
return context
def codex(source: str, prompt: Optional[str], max_tokens: int = 64, language: str = 'javascript', temperature=0.1, concat=True, stop=None) -> str:
""" Write code with the help from openAI Codex API.
* context : source code which AI is supposed to complete
* prompt : natural language command that will be appended as a comment to the source.
"""
if stop is None:
stop = languages.CONFIG[language]["stop"]
context = format_prompt(source, prompt, language)
response = openai.Completion.create(
engine="davinci-codex",
prompt=context,
max_tokens=max_tokens,
temperature=temperature,
stop=stop
)
completion = response["choices"][0]["text"]
if concat:
return "{}{}".format(context, completion).splitlines()
return completion.splitlines()
def update_last_command(command: Callable, args=list(), kwargs=dict()):
""" Save last executed command with arguments to be able to repeat it.
"""
global g_last_cmd
global g_last_cmd_args
global g_last_cmd_kwargs
g_last_cmd = command
g_last_cmd_args = args
g_last_cmd_kwargs = kwargs
def handle_command_prompt(prompt: str):
""" Handle .<command> [<args> ...] prompt.
"""
cmd, *args = prompt.split()
if not cmd in COMMANDS.keys():
print("Unknown command...")
else:
command = COMMANDS[cmd]
command(*args)
if cmd != ".":
update_last_command(command, args)
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument(
"--language",
choices=["cpp", "javascript", "python", "html"],
default="javascript",
help="Hint which programming language should be used for generated code."
)
parser.add_argument(
"--context",
nargs="?",
const=None,
default=None,
help="Path to the context file."
)
parser.add_argument("--temperature", type=float, default=0.1)
args = parser.parse_args()
if args.context is not None:
set_source(load(args.context))
else:
set_source([languages.CONFIG[args.language]["hint"]])
while True:
os.system("clear")
print(get_source())
prompt = input("# ")
if prompt.startswith("."):
handle_command_prompt(prompt)
else:
set_source(
codex(
get_source(),
prompt,
max_tokens=64,
language=args.language,
temperature=args.temperature
)
)