From ca664dd23c3e427b5353da14f16f882dcb042b13 Mon Sep 17 00:00:00 2001 From: Ben Carter Date: Sun, 10 Mar 2024 23:33:52 -0400 Subject: [PATCH] added aidocs generator --- src/aidocs-generator-class.ts | 172 ++++++++++++++++++++++++++++++++++ 1 file changed, 172 insertions(+) create mode 100644 src/aidocs-generator-class.ts diff --git a/src/aidocs-generator-class.ts b/src/aidocs-generator-class.ts new file mode 100644 index 0000000..b72607c --- /dev/null +++ b/src/aidocs-generator-class.ts @@ -0,0 +1,172 @@ +import {OpenAI } from "openai"; +import Anthropic from "@anthropic-ai/sdk"; +import * as fs from "fs"; +import * as ts from "typescript"; +import * as dotenv from "dotenv"; +import { Command } from 'commander'; +const program = new Command(); + + +dotenv.config(); + +const openai = new OpenAI({ + apiKey: process.env.OPENAI_API_KEY, +}); + +const claude = new Anthropic({ + apiKey: process.env.CLAUDE_API_KEY +}); + +type ClassMethod = { + className: string; + methodName: string; + methodSignature: string; + comments: string; // Add this line to include comments +} + +type ParseSettings = { + className: string; + file: string; + exclude?: string[]; + include?: string[]; +} + +function parseDeclarations(settings: ParseSettings): ClassMethod[] { + const filePath = settings.file; + + const fileContent = fs.readFileSync(filePath, "utf8"); + const sourceFile = ts.createSourceFile( + filePath, + fileContent, + ts.ScriptTarget.Latest, + true + ); + let classMethods: ClassMethod[] = []; + + function getLeadingComments(node: ts.Node): string { + const fullText = node.getFullText(sourceFile); + const comments = ts.getLeadingCommentRanges(fullText, 0); + + if (!comments) return ''; + let commentText = ''; + comments.forEach((comment) => { + commentText += fullText.substring(comment.pos, comment.end) + '\n'; + }); + return commentText.trim(); + } + + function visit(node: ts.Node) { + if (ts.isClassDeclaration(node) && node.name) { + + if (settings.className && node.name.text !== settings.className) { + return; + } + + const className = node.name.text; + node.members.forEach((member) => { + if (ts.isMethodDeclaration(member) && member.name) { + + let methodName = member.name.getText(sourceFile); + + // exclude + if(settings.exclude && settings.exclude.includes(methodName)) { + return; + } + + if(settings.include && !settings.include.includes(methodName)) { + return; + } + + let methodSignature = member.getText(sourceFile); + let comments = getLeadingComments(member); // Extract comments + classMethods.push({ className, methodName, methodSignature, comments }); // Add comments to the object + } + }); + } + ts.forEachChild(node, visit); + } + visit(sourceFile); + return classMethods; +} + + +type AiModel = 'claude' | 'gpt3' | 'gpt4' + +async function createDocs(classMethods: ClassMethod[], aiModel: AiModel): Promise { + const fewShotExample = fs.readFileSync(`${__dirname}/player.ts`); + + let currentClass = ''; + for (const { className, methodName, methodSignature, comments } of classMethods) { + if (currentClass !== className) { + // When we encounter a new class, we reset currentClass to the new className + currentClass = className; + } + + let content: string; + let response: any; + const prompt = `${fewShotExample}\n\n Use the provided examples of documenting methods and knowneldge of mod-eluna, azerothcore, create markdown docs for this method: \n\n declare class ${className} {\n Inline Code Comment: ${comments} Method: ${methodName} MethodSignature ${methodSignature}\n} the examples should be not be too simple, and have 10-20 lines`; + switch(aiModel) { + case 'gpt3': + response = await openai.chat.completions.create({ + model: "gpt-3.5-turbo", + messages: [{ role: 'user', content: prompt}], + }); + + if(typeof response.choices[0].message.content == 'string') { + content = response.choices[0].message.content; + } + break; + case 'gpt4': + response = await openai.chat.completions.create({ + model: "gpt-4-turbo-preview", + messages: [{ role: 'user', content: prompt}], + }); + if(typeof response.choices[0].message.content == 'string') { + content = response.choices[0].message.content; + } + break; + case 'claude': + response = await claude.messages.create({ + model: 'claude-3-opus-20240229', + max_tokens: 2000, + messages: [{ role: "user", content: prompt}] + }); + if(typeof response.content[0].text == 'string') { + content = response.content[0].text; + } + break; + } + + + console.log(`writing documentation for ${className}.${methodName}`); + const documentation = content.replace(`filename: ${className.toLowerCase()}.md\n`, ''); + // Ensure directory exists + const dir = './docs/classes'; + if (!fs.existsSync(dir)){ + fs.mkdirSync(dir); + } + // Write to a separate markdown file for each class + fs.appendFileSync(`${dir}/${className}.md`, documentation + '\n\n'); + } + } + + ( async() => { + + program + .option('-m, --model ', 'The AI model to use for documentation generation. Options are gpt3, gpt4, claude', 'gpt3') + .requiredOption('-f, --file ', 'The file to parse for class methods') + .requiredOption('-c, --class ', 'Which classe to process from the file') + .option('-e, --exclude ', 'The file to parse for class methods') + .option('-i, --include ', 'The file to parse for class methods') + program.parse(); + + const options = program.opts(); + const classOutputs = parseDeclarations({ + file: options.file, + className: options.class, + exclude: options.exclude ? options.exclude.split(',') : undefined, + include: options.include ? options.include.split(',') : undefined + }); + console.log(classOutputs); + await createDocs(classOutputs, options.model); + })();