|
| 1 | +/** |
| 2 | + * @license |
| 3 | + * Copyright Google LLC |
| 4 | + * |
| 5 | + * Use of this source code is governed by an MIT-style license that can be |
| 6 | + * found in the LICENSE file at https://angular.io/license |
| 7 | + */ |
| 8 | + |
| 9 | +import {createPartFromUri, FileState, GoogleGenAI, Part} from '@google/genai'; |
| 10 | +import {setTimeout} from 'node:timers/promises'; |
| 11 | +import {readFile, writeFile} from 'node:fs/promises'; |
| 12 | +import {basename} from 'node:path'; |
| 13 | +import glob from 'fast-glob'; |
| 14 | +import assert from 'node:assert'; |
| 15 | +import {Bar} from 'cli-progress'; |
| 16 | +import {Argv, Arguments, CommandModule} from 'yargs'; |
| 17 | +import {randomUUID} from 'node:crypto'; |
| 18 | +import {DEFAULT_MODEL, DEFAULT_TEMPERATURE, DEFAULT_API_KEY} from './consts.js'; |
| 19 | +import {Spinner} from '../utils/spinner.js'; |
| 20 | +import {Log} from '../utils/logging.js'; |
| 21 | + |
| 22 | +/** Command line options. */ |
| 23 | +export interface Options { |
| 24 | + /** Files that the fix should apply to. */ |
| 25 | + files: string[]; |
| 26 | + |
| 27 | + /** Error message(s) to be resolved. */ |
| 28 | + error: string; |
| 29 | + |
| 30 | + /** Model that should be used to apply the prompt. */ |
| 31 | + model: string; |
| 32 | + |
| 33 | + /** Temperature for the model. */ |
| 34 | + temperature: number; |
| 35 | + |
| 36 | + /** API key to use when making requests. */ |
| 37 | + apiKey?: string; |
| 38 | +} |
| 39 | + |
| 40 | +interface FixedFileContent { |
| 41 | + filePath: string; |
| 42 | + content: string; |
| 43 | +} |
| 44 | + |
| 45 | +/** Yargs command builder for the command. */ |
| 46 | +function builder(argv: Argv): Argv<Options> { |
| 47 | + return argv |
| 48 | + .positional('files', { |
| 49 | + description: `One or more glob patterns to find target files (e.g., 'src/**/*.ts' 'test/**/*.ts').`, |
| 50 | + type: 'string', |
| 51 | + array: true, |
| 52 | + demandOption: true, |
| 53 | + }) |
| 54 | + .option('error', { |
| 55 | + alias: 'e', |
| 56 | + description: 'Full error description from the build process', |
| 57 | + type: 'string', |
| 58 | + demandOption: true, |
| 59 | + }) |
| 60 | + .option('model', { |
| 61 | + type: 'string', |
| 62 | + alias: 'm', |
| 63 | + description: 'Model to use for the migration', |
| 64 | + default: DEFAULT_MODEL, |
| 65 | + }) |
| 66 | + .option('temperature', { |
| 67 | + type: 'number', |
| 68 | + alias: 't', |
| 69 | + default: DEFAULT_TEMPERATURE, |
| 70 | + description: 'Temperature for the model. Lower temperature reduces randomness/creativity', |
| 71 | + }) |
| 72 | + .option('apiKey', { |
| 73 | + type: 'string', |
| 74 | + alias: 'a', |
| 75 | + default: DEFAULT_API_KEY, |
| 76 | + description: 'API key used when making calls to the Gemini API', |
| 77 | + }); |
| 78 | +} |
| 79 | + |
| 80 | +/** Yargs command handler for the command. */ |
| 81 | +async function handler(options: Arguments<Options>) { |
| 82 | + assert( |
| 83 | + options.apiKey, |
| 84 | + [ |
| 85 | + 'No API key configured. A Gemini API key must be set as the `GEMINI_API_KEY` environment ' + |
| 86 | + 'variable, or passed in using the `--api-key` flag.', |
| 87 | + 'For internal users, see go/aistudio-apikey', |
| 88 | + ].join('\n'), |
| 89 | + ); |
| 90 | + |
| 91 | + const fixedContents = await fixFilesWithAI( |
| 92 | + options.apiKey, |
| 93 | + options.files, |
| 94 | + options.error, |
| 95 | + options.model, |
| 96 | + options.temperature, |
| 97 | + ); |
| 98 | + Log.info('\n--- AI Suggested Fixes Summary ---'); |
| 99 | + if (fixedContents.length === 0) { |
| 100 | + Log.info( |
| 101 | + 'No files were fixed or found matching the pattern. Check your glob pattern and check whether the files exist.', |
| 102 | + ); |
| 103 | + |
| 104 | + return; |
| 105 | + } |
| 106 | + |
| 107 | + Log.info('Updated files:'); |
| 108 | + const writeTasks = fixedContents.map(({filePath, content}) => |
| 109 | + writeFile(filePath, content).then(() => Log.info(` - ${filePath}`)), |
| 110 | + ); |
| 111 | + await Promise.all(writeTasks); |
| 112 | +} |
| 113 | + |
| 114 | +async function fixFilesWithAI( |
| 115 | + apiKey: string, |
| 116 | + globPatterns: string[], |
| 117 | + errorDescription: string, |
| 118 | + model: string, |
| 119 | + temperature: number, |
| 120 | +): Promise<FixedFileContent[]> { |
| 121 | + const filePaths = await glob(globPatterns, { |
| 122 | + onlyFiles: true, |
| 123 | + absolute: false, |
| 124 | + }); |
| 125 | + |
| 126 | + if (filePaths.length === 0) { |
| 127 | + Log.error(`No files found matching the patterns: ${JSON.stringify(globPatterns, null, 2)}.`); |
| 128 | + return []; |
| 129 | + } |
| 130 | + |
| 131 | + const ai = new GoogleGenAI({vertexai: false, apiKey}); |
| 132 | + let uploadedFileNames: string[] = []; |
| 133 | + |
| 134 | + const progressBar = new Bar({ |
| 135 | + format: `{step} [{bar}] ETA: {eta}s | {value}/{total} files`, |
| 136 | + clearOnComplete: true, |
| 137 | + }); |
| 138 | + |
| 139 | + try { |
| 140 | + const { |
| 141 | + fileNameMap, |
| 142 | + partsForGeneration, |
| 143 | + uploadedFileNames: uploadedFiles, |
| 144 | + } = await uploadFiles(ai, filePaths, progressBar); |
| 145 | + |
| 146 | + uploadedFileNames = uploadedFiles; |
| 147 | + |
| 148 | + const spinner = new Spinner('AI is analyzing the files and generating potential fixes...'); |
| 149 | + const response = await ai.models.generateContent({ |
| 150 | + model, |
| 151 | + contents: [{text: generatePrompt(errorDescription, fileNameMap)}, ...partsForGeneration], |
| 152 | + config: { |
| 153 | + responseMimeType: 'application/json', |
| 154 | + candidateCount: 1, |
| 155 | + maxOutputTokens: Infinity, |
| 156 | + temperature, |
| 157 | + }, |
| 158 | + }); |
| 159 | + |
| 160 | + const responseText = response.text; |
| 161 | + if (!responseText) { |
| 162 | + spinner.failure(`AI returned an empty response.`); |
| 163 | + return []; |
| 164 | + } |
| 165 | + |
| 166 | + const fixes = JSON.parse(responseText) as FixedFileContent[]; |
| 167 | + |
| 168 | + if (!Array.isArray(fixes)) { |
| 169 | + throw new Error('AI response is not a JSON array.'); |
| 170 | + } |
| 171 | + |
| 172 | + spinner.complete(); |
| 173 | + return fixes; |
| 174 | + } finally { |
| 175 | + if (uploadedFileNames.length) { |
| 176 | + progressBar.start(uploadedFileNames.length, 0, { |
| 177 | + step: 'Deleting temporary uploaded files', |
| 178 | + }); |
| 179 | + const deleteTasks = uploadedFileNames.map((name) => { |
| 180 | + return ai.files |
| 181 | + .delete({name}) |
| 182 | + .catch((error) => Log.warn(`WARNING: Failed to delete temporary file ${name}:`, error)) |
| 183 | + .finally(() => progressBar.increment()); |
| 184 | + }); |
| 185 | + |
| 186 | + await Promise.allSettled(deleteTasks).finally(() => progressBar.stop()); |
| 187 | + } |
| 188 | + } |
| 189 | +} |
| 190 | + |
| 191 | +async function uploadFiles( |
| 192 | + ai: GoogleGenAI, |
| 193 | + filePaths: string[], |
| 194 | + progressBar: Bar, |
| 195 | +): Promise<{ |
| 196 | + uploadedFileNames: string[]; |
| 197 | + partsForGeneration: Part[]; |
| 198 | + fileNameMap: Map<string, string>; |
| 199 | +}> { |
| 200 | + const uploadedFileNames: string[] = []; |
| 201 | + const partsForGeneration: Part[] = []; |
| 202 | + const fileNameMap = new Map<string, string>(); |
| 203 | + |
| 204 | + progressBar.start(filePaths.length, 0, {step: 'Uploading files'}); |
| 205 | + |
| 206 | + const uploadPromises = filePaths.map(async (filePath) => { |
| 207 | + try { |
| 208 | + const uploadedFile = await ai.files.upload({ |
| 209 | + file: new Blob([await readFile(filePath, {encoding: 'utf8'})], { |
| 210 | + type: 'text/plain', |
| 211 | + }), |
| 212 | + config: { |
| 213 | + displayName: `fix_request_${basename(filePath)}_${randomUUID()}`, |
| 214 | + }, |
| 215 | + }); |
| 216 | + |
| 217 | + assert(uploadedFile.name, 'File name cannot be undefined after upload.'); |
| 218 | + |
| 219 | + let getFile = await ai.files.get({name: uploadedFile.name}); |
| 220 | + while (getFile.state === FileState.PROCESSING) { |
| 221 | + await setTimeout(500); // Wait for 500ms before re-checking |
| 222 | + getFile = await ai.files.get({name: uploadedFile.name}); |
| 223 | + } |
| 224 | + |
| 225 | + if (getFile.state === FileState.FAILED) { |
| 226 | + throw new Error(`File processing failed on API for ${filePath}. Skipping this file.`); |
| 227 | + } |
| 228 | + |
| 229 | + if (getFile.uri && getFile.mimeType) { |
| 230 | + const filePart = createPartFromUri(getFile.uri, getFile.mimeType); |
| 231 | + partsForGeneration.push(filePart); |
| 232 | + fileNameMap.set(filePath, uploadedFile.name); |
| 233 | + progressBar.increment(); |
| 234 | + return uploadedFile.name; // Return the name on success |
| 235 | + } else { |
| 236 | + throw new Error( |
| 237 | + `Uploaded file for ${filePath} is missing URI or MIME type after processing. Skipping.`, |
| 238 | + ); |
| 239 | + } |
| 240 | + } catch (error: any) { |
| 241 | + Log.error(`Error uploading or processing file ${filePath}: ${error.message}`); |
| 242 | + return null; // Indicate failure for this specific file |
| 243 | + } |
| 244 | + }); |
| 245 | + |
| 246 | + const results = await Promise.allSettled(uploadPromises).finally(() => progressBar.stop()); |
| 247 | + |
| 248 | + for (const result of results) { |
| 249 | + if (result.status === 'fulfilled' && result.value !== null) { |
| 250 | + uploadedFileNames.push(result.value); |
| 251 | + } |
| 252 | + } |
| 253 | + |
| 254 | + return {uploadedFileNames, fileNameMap, partsForGeneration}; |
| 255 | +} |
| 256 | + |
| 257 | +function generatePrompt(errorDescription: string, fileNameMap: Map<string, string>): string { |
| 258 | + return ` |
| 259 | + You are a highly skilled software engineer, specializing in Bazel, Starlark, Python, Angular, JavaScript, |
| 260 | + TypeScript, and everything related. |
| 261 | + The following files are part of a build process that failed with the error: |
| 262 | + \`\`\` |
| 263 | + ${errorDescription} |
| 264 | + \`\`\` |
| 265 | + Please analyze the content of EACH provided file and suggest modifications to resolve the issue. |
| 266 | +
|
| 267 | + Your response MUST be a JSON array of objects. Each object in the array MUST have two properties: |
| 268 | + 'filePath' (the full path from the mappings provided.) and 'content' (the complete corrected content of that file). |
| 269 | + DO NOT include any additional text, non modified files, commentary, or markdown outside the JSON array. |
| 270 | + For example: |
| 271 | + [ |
| 272 | + {"filePath": "/full-path-from-mappings/file1.txt", "content": "Corrected content for file1."}, |
| 273 | + {"filePath": "/full-path-from-mappings/file2.js", "content": "console.log('Fixed JS');"} |
| 274 | + ] |
| 275 | +
|
| 276 | + IMPORTANT: The input files are mapped as follows: ${Array.from(fileNameMap.entries())} |
| 277 | +`; |
| 278 | +} |
| 279 | + |
| 280 | +/** CLI command module. */ |
| 281 | +export const FixModule: CommandModule<{}, Options> = { |
| 282 | + builder, |
| 283 | + handler, |
| 284 | + command: 'fix <files..>', |
| 285 | + describe: 'Fixes errors from the specified error output', |
| 286 | +}; |
0 commit comments