sourcesailor
Version:
A CLI tool for analyzing and documenting codebases
161 lines (160 loc) • 8.61 kB
JavaScript
import OpenAI from 'openai';
import { prompts } from "./prompts.mjs";
import { readConfig } from "./utils.mjs";
// interface Tool {
// type: string
// function: Function
// }
export class OpenAIInferrence {
constructor() {
this.modelLimits = [
{ name: 'gpt-4', limit: 8000 },
{ name: 'gpt-4-32k', limit: 32000 },
{ name: 'gpt-4-0125-preview', limit: 128000 },
{ name: 'gpt-4-turbo', limit: 128000 },
{ name: 'gpt-4o', limit: 128000 },
{ name: 'gpt-4o-mini', limit: 128000 },
{ name: 'chatgpt-4o-latest', limit: 128000 },
{ name: 'gpt-4o-2024-08-06', limit: 128000 },
{ name: 'gpt-4-1106-preview', limit: 128000 },
{ name: 'gpt-4-turbo-preview', limit: 128000 },
{ name: 'gpt-3.5-turbo', limit: 4000 },
{ name: 'gpt-3.5-turbo-16k', limit: 16000 }
];
}
getName() {
return 'OpenAI';
}
createPrompt(systemPrompt, userPrompt, isVerbose, userExpertise) {
let finalSystemPrompt = systemPrompt;
if (userExpertise) {
finalSystemPrompt += `\n<Expertise>${JSON.stringify(userExpertise)}</Expertise>`;
}
const compatibilityMessage = [{
role: "system",
content: finalSystemPrompt
}, {
role: "user",
content: userPrompt
}];
if (isVerbose) {
console.log(`System Prompt: ${finalSystemPrompt}`);
console.log(`User Prompt: ${userPrompt}`);
}
return compatibilityMessage;
}
async callApiAndReturnResult(openai, model, compatibilityMessage, isStreaming, isVerbose, tools) {
const apiParams = {
model,
messages: compatibilityMessage,
temperature: 0,
stream: isStreaming
};
if (tools) {
apiParams.tools = tools;
apiParams.tool_choice = { type: "function", function: { name: tools[0].function.name } };
delete apiParams.stream;
isStreaming = false;
}
if (isVerbose) {
console.log(JSON.stringify(apiParams, null, 2));
}
const matchJson = await openai.chat.completions.create(apiParams);
if (isVerbose && !isStreaming) {
console.log(JSON.stringify(matchJson.choices[0], null, 2));
}
if (isStreaming) {
// @ts-expect-erro Exclude streaming from coverage
const matchJsonStream = matchJson;
return this.convertStreamToStringStream(matchJsonStream);
}
else {
const completionData = matchJson;
if (completionData.choices.length === 0) {
throw new Error('Invalid response from OpenAI');
}
if (completionData.choices[0].finish_reason === 'tool_calls' || (completionData.choices[0].message?.tool_calls?.length ?? 0 > 0)) {
const response = completionData.choices[0].message?.tool_calls?.flatMap(toolCall => toolCall?.function?.arguments);
return response?.join('');
}
else {
return completionData.choices[0].message.content || undefined;
}
}
}
async getModel(modelName) {
const config = readConfig();
const defaultModel = config.DEFAULT_OPENAI_MODEL || process.env.DEFAULT_OPENAI_MODEL || 'gpt-4-preview';
return modelName || defaultModel;
}
getOpenAiClient() {
const config = readConfig();
return new OpenAI({ apiKey: config.OPENAI_API_KEY || process.env.OPENAI_API_KEY, timeout: 60000 });
}
async inferProjectDirectory(projectDirectory, isStreaming = false, isVerbose = false, userExpertise, modelName) {
const openai = this.getOpenAiClient();
const model = await this.getModel(modelName);
const compatibilityMessage = this.createPrompt(`${prompts.commonSystemPrompt.prompt}\n${prompts.rootUnderstanding.prompt}`, `<FileStructure>${JSON.stringify(projectDirectory)}</FileStructure>`, isVerbose, userExpertise);
const tools = [];
if (prompts.rootUnderstanding.params) {
tools.push({
type: "function",
function: {
name: prompts.rootUnderstanding.params.name,
parameters: prompts.rootUnderstanding.params.parameters,
description: prompts.rootUnderstanding.params.description
}
});
}
return this.callApiAndReturnResult(openai, model, compatibilityMessage, isStreaming, isVerbose, tools);
}
async inferDependency(dependencyFile, workflow, isStreaming = false, isVerbose = false, userExpertise, modelName) {
// @ts-expect-error Exclude streaming from coverage
const openai = this.getOpenAiClient(isVerbose);
const model = await this.getModel(modelName);
const compatibilityMessage = this.createPrompt(`${prompts.commonSystemPrompt.prompt}\n${prompts.dependencyUnderstanding.prompt}`, `<DependencyFile>${JSON.stringify(dependencyFile)}</DependencyFile>\n<Workflow>${workflow}</Workflow> ${prompts.commonMarkdownPrompt.prompt}`, isVerbose, userExpertise);
return this.callApiAndReturnResult(openai, model, compatibilityMessage, isStreaming, isVerbose);
}
async inferCode(code, isStreaming = false, isVerbose = false, userExpertise, modelName) {
// @ts-expect-error Exclude streaming from coverage
const openai = this.getOpenAiClient(isVerbose);
const model = await this.getModel(modelName);
const compatibilityMessage = this.createPrompt(`${prompts.commonSystemPrompt.prompt}\n${prompts.codeUnderstanding.prompt}`, `<Code>${JSON.stringify(code)}</Code> ${prompts.commonMarkdownPrompt.prompt}`, isVerbose, userExpertise);
return this.callApiAndReturnResult(openai, model, compatibilityMessage, isStreaming, isVerbose);
}
async inferInterestingCode(code, isStreaming = false, isVerbose = false, userExpertise, modelName) {
// @ts-expect-error Exclude streaming from coverage
const openai = this.getOpenAiClient(isVerbose);
const model = await this.getModel(modelName);
const compatibilityMessage = this.createPrompt(prompts.interestingCodeParts.prompt, `<Code>${JSON.stringify(code)}</Code> ${prompts.commonMarkdownPrompt.prompt}`, isVerbose, userExpertise);
return this.callApiAndReturnResult(openai, model, compatibilityMessage, isStreaming, isVerbose);
}
async generateReadme(directoryStructure, dependencyInference, codeInference, isStreaming = false, isVerbose = false, userExpertise, modelName) {
// @ts-expect-error Exclude streaming from coverage
const openai = this.getOpenAiClient(isVerbose);
const model = await this.getModel(modelName);
const compatibilityMessage = this.createPrompt(prompts.readmePrompt.prompt, `<DirectoryStructure>${JSON.stringify(directoryStructure)}</DirectoryStructure>\n<DependencyInferrence>${JSON.stringify(dependencyInference)}</DependencyInferrence>\n<CodeInferrence>${JSON.stringify(codeInference)}</CodeInferrence> ${prompts.commonMarkdownPrompt.prompt}`, isVerbose, userExpertise);
return this.callApiAndReturnResult(openai, model, compatibilityMessage, isStreaming, isVerbose);
}
async generateMonorepoReadme(monorepoInferrenceInfo, isStreaming = false, isVerbose = false, userExpertise, modelName) {
// @ts-expect-error Exclude streaming from coverage
const openai = this.getOpenAiClient(isVerbose);
const model = await this.getModel(modelName);
const compatibilityMessage = this.createPrompt(prompts.consolidatedInferrenceForMonoRepo.prompt, `<MonoRepoInferrence>${JSON.stringify(monorepoInferrenceInfo)}</MonoRepoInferrence> ${prompts.commonMarkdownPrompt.prompt}`, isVerbose, userExpertise);
return this.callApiAndReturnResult(openai, model, compatibilityMessage, isStreaming, isVerbose);
}
async listModels(isVerbose = false) {
const openai = this.getOpenAiClient();
const models = await openai.models.list();
if (isVerbose) {
console.log(models.data);
}
return models.data.sort((a, b) => b.created - a.created).flatMap(model => model.id);
}
async *convertStreamToStringStream(response) {
for await (const chunk of response) {
yield chunk.choices[0]?.delta.content || "";
}
}
}
export default OpenAIInferrence;