scraipt
Version:
Scrape away inefficient code during compile-time using AI
123 lines (122 loc) • 4.52 kB
JavaScript
;
var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); }
return new (P || (P = Promise))(function (resolve, reject) {
function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); }
step((generator = generator.apply(thisArg, _arguments || [])).next());
});
};
var __importDefault = (this && this.__importDefault) || function (mod) {
return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.OpenAIAPI = void 0;
const openai_1 = __importDefault(require("openai"));
const gpt_tokenizer_1 = require("gpt-tokenizer");
const OPENAI_API_KEY = process.env.OPENAI_API_KEY;
const MAX_GPT_4_TOKENS = 4096;
// Messages to prepend to the prompt messages
const systemMessages = [
{
role: 'system',
content: 'Generate the most efficient, opimized version of the given code using the given code context. Do not include any other text in the response. The code should be valid JavaScript.',
},
];
/**
* Count the number of tokens in the messages.
* @param messages The messages to count.
* @returns The number of tokens in the messages.
*/
const countTokens = (messages) => {
return (0, gpt_tokenizer_1.encodeChat)(messages, 'gpt-4').length;
};
/**
* Limit the number of tokens in the input.
* @param input The input to limit.
* @returns The input with a limited number of tokens.
*/
const limitInputTokens = (input) => {
if (input.length > MAX_GPT_4_TOKENS) {
return input.substring(0, MAX_GPT_4_TOKENS);
}
return input;
};
/**
Add system messages to the prompt messages.
* @param context The context of the code.
* @param code The code to append.
* @returns The messages to use for the completion.
*/
const createMessages = (context, code) => {
return systemMessages.concat([
{
role: 'user',
content: context,
},
{
role: 'user',
content: code,
},
]);
};
/**
A class to interact with the OpenAI API.
*/
class OpenAIAPI {
/**
* Create a new instance of the OpenAIAPI class.
* @throws An error if the OpenAI API key is not found.
*/
constructor() {
if (!OPENAI_API_KEY) {
throw new Error('OpenAI API key not found');
}
this.openai = new openai_1.default({
apiKey: OPENAI_API_KEY,
});
this.totalTokensUsed = 0;
}
/**
* Create a text completion using the given messages.
* @param code The code to optimize.
* @param context The context of the code.
* @param model The model to use for the completion.
* @returns The completion text.
*/
createTextCompletion(code, context, model) {
return __awaiter(this, void 0, void 0, function* () {
// Add system messages to the prompt messages
let messages = createMessages(context, code);
const tokenCount = countTokens(messages);
this.totalTokensUsed += tokenCount;
if (this.maxTokenCount && this.totalTokensUsed > this.maxTokenCount) {
return;
}
// Don't surpass the maximum number of tokens for the model
if (tokenCount > MAX_GPT_4_TOKENS) {
// TODO: Don't hardcode the max tokens
const limitedCode = limitInputTokens(code);
const limitedMessages = createMessages(context, limitedCode);
messages = limitedMessages;
}
const result = yield this.openai.chat.completions.create({
messages,
model,
});
const choices = result.choices;
if (!choices) {
return;
}
const content = choices[0].message.content;
if (!content) {
return;
}
const completionTokenCount = countTokens([choices[0].message]);
this.totalTokensUsed += completionTokenCount;
return content;
});
}
}
exports.OpenAIAPI = OpenAIAPI;