/
openai.ts
90 lines (82 loc) · 2.94 KB
/
openai.ts
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
import { Configuration, OpenAIApi } from "openai";
import { splitCode } from "./split-file.js";
import {
Rename,
renameVariablesAndFunctions,
} from "./rename-variables-and-functions.js";
import { mapPromisesParallel } from "./run-promises-in-parallel.js";
type Options = {
apiKey: string;
use4k: boolean;
};
export default ({ apiKey, use4k }: Options) => {
const client = new OpenAIApi(new Configuration({ apiKey }));
return async (code: string): Promise<string> => {
const codeBlocks = await splitCode(code, use4k);
let variablesAndFunctionsToRename: Rename[] = [];
await mapPromisesParallel(10, codeBlocks, async (codeBlock) => {
const renames = await codeToVariableRenames(codeBlock);
variablesAndFunctionsToRename =
variablesAndFunctionsToRename.concat(renames);
});
return renameVariablesAndFunctions(code, variablesAndFunctionsToRename);
};
async function codeToVariableRenames(code: string) {
const chatCompletion = await client.createChatCompletion({
model: use4k ? "gpt-3.5-turbo" : "gpt-3.5-turbo-16k",
functions: [
{
name: "rename_variables_and_functions",
description: "Rename variables and function names in Javascript code",
parameters: {
type: "object",
properties: {
variablesAndFunctionsToRename: {
type: "array",
items: {
type: "object",
properties: {
name: {
type: "string",
description:
"The name of the variable or function name to rename",
},
newName: {
type: "string",
description:
"The new name of the variable or function name",
},
},
required: ["name", "newName"],
},
},
},
required: ["variablesToRename"],
},
},
],
messages: [
{
role: "assistant",
content:
"Rename all Javascript variables and functions to have descriptive names based on their usage in the code.",
},
{ role: "user", content: code },
],
});
const data = chatCompletion.data.choices[0];
if (data.finish_reason !== "function_call") return [];
const {
variablesAndFunctionsToRename,
}: { variablesAndFunctionsToRename: Rename[] } = JSON.parse(
fixPerhapsBrokenResponse(data.message?.function_call?.arguments!)
);
return variablesAndFunctionsToRename;
}
};
function fixPerhapsBrokenResponse(jsonResponse: string) {
// Sometimes the response has an extra comma at the end of the array, like:
// {"result": [{"foo": "bar"}, { "foo": "baz" },\n ]}
// This is invalid JSON, so we need to remove the comma.
return jsonResponse.replace(/},\s*]/im, "}]");
}