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
|
import openai
import pyjson5
class Env:
def __init__(self):
self.model = "gpt-4o"
self.output_tokens_min = 250
self.output_tokens_max = 500
self.temp: float = None
self.seed: int = None
self.max_tokens = 8000
class SourceFile:
def __init__(self):
self.name: str = None
self.category: str = None
self.contents: str = None
class LLMAPIException (Exception): ...
def __do_prompt (lang: str, extra_prompt: str, env: Env):
prompt = '''
Write me a program source code written in %s. Write anything you'd like.
Use more than %d words, but no more than %d words in the code.''' % (
lang, env.output_tokens_min, env.output_tokens_max)
if extra_prompt:
prompt = ' ' + extra_prompt
messages = [
{
"role": "system",
"content": "You're a helpful assistant that writes "
+ "computer program of any kind"
},
{ "role": "user", "content": prompt }
]
functions = [
{
"name": "get_result",
"description": "Output the source code",
"parameters": {
"type": "object",
"properties": {
"cat": {
"type": "string",
"description": "one word describing program functions"
},
"filename": {
"type": "string",
"description": "The source code file name" # TODO: strip unallowed characters in Unix fs
},
"code": {
"type": "string",
"description": "The source code"
}
}
}
}
]
return openai.chat.completions.create(
model = env.model,
messages = messages,
functions = functions,
function_call = "auto",
temperature = env.temp,
seed = env.seed,
max_tokens = env.max_tokens
)
def gen_rand_srccode (lang: str, extra_prompt: str, env: Env) -> SourceFile | str:
rsp = __do_prompt(lang, extra_prompt, env)
match rsp.choices[0].finish_reason:
case 'stop':
return rsp.choices[0].message.content
case 'function_call':
choice = pyjson5.loads(rsp.choices[0].message.function_call.arguments)
ret = SourceFile()
ret.name = choice["filename"]
ret.category = choice.get("cat")
ret.contents = choice["code"]
return ret
case _: raise LLMAPIException(rsp.choices[0].finish_reason)
|
