genaicode 0.0.31 → 0.0.33

package/README.md

@@ -1,13 +1,3 @@
-<p align="center">
-  <picture>
-    <source media="(prefers-color-scheme: dark)" srcset="media/logo-dark.png">
-    <source media="(prefers-color-scheme: light)" srcset="media/logo.png">
-    <img alt="GenAIcode Logo." src="media/logo.png" width="100%" height="auto">
-  </picture>
-</p>
-
-<div align="center">
-
 # Programming on steroids
 
   <a href="https://www.npmjs.com/package/genaicode">
@@ -77,6 +67,8 @@ GenAIcode supports various command-line options to customize its behavior:
 - `--temperature=<value>`: Sets the temperature parameter for the AI model (default: 0.7).
 - `--vision`: Enables vision capabilities for processing image inputs.
 - `--imagen`: Enables image generation capabilities using AI models.
+- `--cheap`: Uses a cheaper, faster model for code generation, which may provide lower quality results but is more cost-effective for simpler tasks.
+- `--content-mask=<path>`: Applies a content mask to limit the initial source code files included in the request. The value should be a prefix of the path relative to rootDir.
 - `--help`: Displays the help message with all available options.
 
 ## Configuration (.genaicoderc)

package/bin/vertex-monkey-patch.cjs

@@ -27,7 +27,7 @@ if (content.includes(MONKEY_PATCH_TOOL_CONFIG)) {
 const newContent = content.replaceAll(
   'data: generateContentRequest,',
   `// MONKEY PATCH TOOL_CONFIG
-        data: {...generateContentRequest,tool_config: {function_calling_config: { mode: "ANY", allowed_function_names: request.toolConfig.functionCallingConfig.allowedFunctionNames }}},`,
+        data: {...generateContentRequest,tool_config: {function_calling_config: { mode: request.toolConfig?.functionCallingConfig?.mode, allowed_function_names: request.toolConfig?.functionCallingConfig.allowedFunctionNames }}},`,
 );
 fs.writeFileSync(path, newContent, 'utf-8');
 console.log('Vertex monkey patch applied');

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "genaicode",
-  "version": "0.0.31",
+  "version": "0.0.33",
   "author": "Grzegorz Tańczyk",
   "repository": {
     "type": "git",
@@ -16,7 +16,7 @@
   "engines": {
     "node": ">=18"
   },
-  "license": "MIT",
+  "license": "Apache-2.0",
   "files": [
     "bin",
     "src",
@@ -46,11 +46,14 @@
     "@anthropic-ai/sdk": "^0.24.3",
     "@anthropic-ai/vertex-sdk": "^0.4.1",
     "@google-cloud/vertexai": "^1.3.0",
+    "@google-cloud/aiplatform": "^3.25.0",
+    "@imgly/background-removal-node": "^1.4.5",
     "diff": "^5.2.0",
     "image-size": "^1.1.1",
     "mime-types": "^2.1.35",
     "openai": "^4.52.7",
-    "please-upgrade-node": "^3.2.0"
+    "please-upgrade-node": "^3.2.0",
+    "sharp": "^0.33.4"
   },
   "lint-staged": {
     "*.{js,ts,css,md}": "prettier --write"

package/src/ai-service/anthropic.js

@@ -4,7 +4,7 @@ import { printTokenUsageAndCost, processFunctionCalls } from './common.js';
 /**
  * This function generates content using the Anthropic Claude model.
  */
-export async function generateContent(prompt, functionDefs, requiredFunctionName, temperature) {
+export async function generateContent(prompt, functionDefs, requiredFunctionName, temperature, cheap = false) {
   const anthropic = new Anthropic({
     defaultHeaders: {
       'anthropic-beta': 'max-tokens-3-5-sonnet-2024-07-15',
@@ -53,8 +53,11 @@ export async function generateContent(prompt, functionDefs, requiredFunctionName
       }
     });
 
+  const model = cheap ? 'claude-3-haiku-20240307' : 'claude-3-5-sonnet-20240620';
+  console.log(`Using Anthropic model: ${model}`);
+
   const response = await anthropic.messages.create({
-    model: 'claude-3-5-sonnet-20240620',
+    model: model,
     system: prompt.find((item) => item.type === 'systemPrompt').systemPrompt,
     messages,
     tools: functionDefs.map((fd) => ({
@@ -63,7 +66,7 @@ export async function generateContent(prompt, functionDefs, requiredFunctionName
       input_schema: fd.parameters,
     })),
     tool_choice: requiredFunctionName ? { type: 'tool', name: requiredFunctionName } : { type: 'any' },
-    max_tokens: 8192,
+    max_tokens: cheap ? 4096 : 8192,
     temperature: temperature,
   });
 

package/src/ai-service/chat-gpt.js

@@ -4,7 +4,7 @@ import { printTokenUsageAndCost, processFunctionCalls } from './common.js';
 /**
  * This function generates content using the OpenAI chat model.
  */
-export async function generateContent(prompt, functionDefs, requiredFunctionName, temperature) {
+export async function generateContent(prompt, functionDefs, requiredFunctionName, temperature, cheap = false) {
   const openai = new OpenAI();
 
   const messages = prompt
@@ -55,8 +55,11 @@ export async function generateContent(prompt, functionDefs, requiredFunctionName
     })
     .flat();
 
+  const model = cheap ? 'gpt-4o-mini' : 'gpt-4o-2024-08-06';
+  console.log(`Using OpenAI model: ${model}`);
+
   const response = await openai.chat.completions.create({
-    model: 'gpt-4o',
+    model: model,
     messages,
     tools: functionDefs.map((funDef) => ({ type: 'function', function: funDef })),
     tool_choice: requiredFunctionName ? { type: 'function', function: { name: requiredFunctionName } } : 'required',

package/src/ai-service/common.test.js

@@ -0,0 +1,106 @@
+import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
+import { printTokenUsageAndCost, processFunctionCalls } from './common.js';
+import { functionDefs } from './function-calling.js';
+
+// Mock cli-params.js
+vi.mock('../cli/cli-params.js', () => ({
+  chatGpt: true,
+  anthropic: false,
+  vertexAi: false,
+  vertexAiClaude: false,
+  temperature: 0.7,
+  cheap: false,
+  geminiBlockNone: false,
+  requireExplanations: false,
+}));
+
+describe('printTokenUsageAndCost', () => {
+  let consoleLogSpy;
+
+  beforeEach(() => {
+    consoleLogSpy = vi.spyOn(console, 'log').mockImplementation(() => {});
+  });
+
+  afterEach(() => {
+    consoleLogSpy.mockRestore();
+  });
+
+  it('should correctly log token usage and estimated cost', () => {
+    const usage = {
+      inputTokens: 100,
+      outputTokens: 50,
+      totalTokens: 150,
+    };
+    const inputCostPerToken = 0.0001;
+    const outputCostPerToken = 0.0002;
+
+    printTokenUsageAndCost(usage, inputCostPerToken, outputCostPerToken);
+
+    expect(consoleLogSpy).toHaveBeenCalledTimes(5);
+    expect(consoleLogSpy).toHaveBeenCalledWith('Token Usage:');
+    expect(consoleLogSpy).toHaveBeenCalledWith('  - Input tokens: ', 100);
+    expect(consoleLogSpy).toHaveBeenCalledWith('  - Output tokens: ', 50);
+    expect(consoleLogSpy).toHaveBeenCalledWith('  - Total tokens: ', 150);
+    expect(consoleLogSpy).toHaveBeenCalledWith('  - Estimated cost: ', '0.020000', ' USD');
+  });
+
+  it('should handle zero tokens correctly', () => {
+    const usage = {
+      inputTokens: 0,
+      outputTokens: 0,
+      totalTokens: 0,
+    };
+    const inputCostPerToken = 0.0001;
+    const outputCostPerToken = 0.0002;
+
+    printTokenUsageAndCost(usage, inputCostPerToken, outputCostPerToken);
+
+    expect(consoleLogSpy).toHaveBeenCalledWith('  - Estimated cost: ', '0.000000', ' USD');
+  });
+});
+
+describe('processFunctionCalls', () => {
+  it('should process valid function calls correctly', () => {
+    const validFunctionCalls = [
+      { name: 'explanation', args: { text: 'This is an explanation' } },
+      { name: functionDefs[0].name, args: {} },
+    ];
+
+    const result = processFunctionCalls(validFunctionCalls);
+
+    expect(result).toEqual(validFunctionCalls);
+  });
+
+  it('should throw an error for unknown function names', () => {
+    const invalidFunctionCalls = [{ name: 'unknownFunction', args: {} }];
+
+    expect(() => processFunctionCalls(invalidFunctionCalls)).toThrow('Unknown function name: unknownFunction');
+  });
+
+  it('should correctly handle explanations', () => {
+    const consoleLogSpy = vi.spyOn(console, 'log').mockImplementation(() => {});
+
+    const functionCallsWithExplanation = [
+      { name: 'explanation', args: { text: 'This is an explanation' } },
+      { name: functionDefs[0].name, args: {} },
+    ];
+
+    processFunctionCalls(functionCallsWithExplanation);
+
+    expect(consoleLogSpy).toHaveBeenCalledWith('Explanations:', ['This is an explanation']);
+
+    consoleLogSpy.mockRestore();
+  });
+
+  it('should return all function calls including explanations', () => {
+    const functionCalls = [
+      { name: 'explanation', args: { text: 'This is an explanation' } },
+      { name: functionDefs[0].name, args: {} },
+    ];
+
+    const result = processFunctionCalls(functionCalls);
+
+    expect(result).toEqual(functionCalls);
+    expect(result.length).toBe(2);
+  });
+});

package/src/ai-service/dall-e.js

@@ -1,24 +1,45 @@
-import OpenAI from 'openai';
+import OpenAI, { toFile } from 'openai';
+
+import { setTempBuffer } from '../files/temp-buffer.js';
+import { resizeImageBuffer } from '../images/resize-image.js';
+import { ensureAlpha } from '../images/ensure-alpha.js';
 
 /**
  * Generate an image using OpenAI's DALL-E model and save it to a file
  * @param {string} prompt - The description of the image to generate
- * @param {string} size - The size of the image to generate ('256x256', '512x512', or '1024x1024')
+ * @param {string|undefined} contextImagePath - The image to be used as a context
+ * @param {{width: number, height: number}} size - The size of the image to generate
+ * @param {boolean} cheap - Whether to use a cheaper model
  * @returns {Promise<string>} - The url of the image
  */
-export async function generateImage(prompt, size = '1024x1024') {
+export async function generateImage(prompt, contextImagePath, size, cheap = false) {
   const openai = new OpenAI();
 
   try {
-    const response = await openai.images.generate({
-      model: 'dall-e-3',
+    const model = contextImagePath ? 'dall-e-2' : cheap ? 'dall-e-2' : 'dall-e-3';
+    console.log(`Using DALL-E model: ${model}`);
+
+    const options = {
+      model: model,
       prompt: prompt,
       n: 1,
-      size: size,
+      size: '1024x1024',
       response_format: 'url',
-    });
+    };
+    const response = contextImagePath
+      ? await openai.images.edit({ ...options, image: await toFile(await ensureAlpha(contextImagePath)) })
+      : await openai.images.generate(options);
+
+    let imageUrl = response.data[0].url;
+
+    if ((size.width !== 1024) | (size.height !== 1024)) {
+      console.log('Resizing image to desired size', size);
+      const imageResponse = await fetch(imageUrl);
+      const arrayBuffer = await imageResponse.arrayBuffer();
+      const buffer = Buffer.from(arrayBuffer);
 
-    const imageUrl = response.data[0].url;
+      imageUrl = setTempBuffer(await resizeImageBuffer(buffer, size));
+    }
 
     console.log(`Image generated, url: ${imageUrl}`);
     return imageUrl;

package/src/ai-service/function-calling.js

@@ -1,4 +1,4 @@
-import { requireExplanations, temperature } from '../cli/cli-params.js';
+import { requireExplanations, temperature, cheap } from '../cli/cli-params.js';
 
 /**
  * Function definitions for function calling feature
@@ -7,7 +7,7 @@ export const functionDefs = [
   {
     name: 'getSourceCode',
     description:
-      'This function returns source code of the application in Map format, where absolute file path is the key, and file content is the value. This function can be called only once during the conversation, and only if suggested by the user.',
+      'This function returns source code of the application in Map format, where absolute file path is the key, and the value is an object, where one of the properties may be the content of the file. Some keys may not provide content. This function can be called only once during the conversation, and only if suggested by the user.',
     parameters: {
       type: 'object',
       properties: {
@@ -66,19 +66,28 @@ export const functionDefs = [
                   'moveFile',
                   'generateImage',
                   'downloadFile',
+                  'splitImage',
+                  'resizeImage',
+                  'imglyRemoveBackground',
                 ],
                 description: 'A name of the tool that will be used to perform the update.',
               },
               temperature: {
                 type: 'number',
                 description:
-                  'Temperature parameter that will be used for LLM request. The value is adjusted to the characteristic of the update. If there is a need for more creative solution, the value should be lower, but stil within [0.0, 2.0] range. The default value is: ' +
+                  'Temperature parameter that will be used for LLM request. The value is adjusted to the characteristic of the update. If there is a need for a more creative solution, the value should be lower, but stil within [0.0, 2.0] range. The default value is: ' +
                   temperature,
               },
+              cheap: {
+                type: 'boolean',
+                description:
+                  'true value means that the prompt will be executed with cheaper model, which work faster, but provides lower quality results, so please use it only in situation when lower quality results are acceptable for the prompt. The default value is: ' +
+                  !!cheap,
+              },
               prompt: {
                 type: 'string',
                 description:
-                  'Prompt that will be added to the LLM request together with the tool request. It summarizes the planned changes for this particular file.',
+                  'Prompt that will be passed to the model request together with the tool request. It summarizes the planned changes for this particular file, so it should be detailed enough for the model to generate necessary changes.',
               },
               contextImageAssets: {
                 type: 'array',
@@ -87,13 +96,13 @@ export const functionDefs = [
                 items: { type: 'string' },
               },
             },
-            required: ['path', 'updateToolName', 'temperature', 'prompt', 'contextImageAssets'],
+            required: ['path', 'updateToolName', 'temperature', 'prompt', 'contextImageAssets', 'cheap'],
           },
         },
         contextPaths: {
           type: 'array',
           description:
-            'An array of absolute paths of files that should be used to provided context. Context files could be for example the dependencies, or files that depend on one of the files that we want to update in the next step.',
+            'An array of absolute paths of files that should be used to provide context for the following updates. Context files could be for example the dependencies, or files that depend on one of the files that we want to update in the next step.',
           items: {
             type: 'string',
           },
@@ -268,23 +277,44 @@ Index: filename.js
       properties: {
         prompt: {
           type: 'string',
-          description: 'The prompt to generate the image.',
+          description:
+            'The prompt that will be used to generate the image. This prompt must be detailed, it will be used by image generation model.',
         },
         filePath: {
           type: 'string',
           description: 'The file path to save the generated image.',
         },
-        size: {
+        contextImagePath: {
           type: 'string',
-          enum: ['256x256', '512x512', '1024x1024'],
+          description:
+            'Path to a image file that will be used as a context for image generation. It is useful if there is a need to edit an image with genAI.',
+        },
+        size: {
+          type: 'object',
+          properties: {
+            width: {
+              type: 'number',
+              description: 'width of the image',
+            },
+            height: {
+              type: 'number',
+              description: 'height of the image',
+            },
+          },
+          required: ['width', 'height'],
           description: 'The size of the image to generate.',
         },
+        cheap: {
+          type: 'boolean',
+          description:
+            'true value means that the prompt will be executed with cheaper model, which work faster, but provides lower quality results, so please use it only in situation when lower quality results are acceptable for the prompt.',
+        },
         explanation: {
           type: 'string',
           description: 'The explanation of the reasoning behind generating this image',
         },
       },
-      required: ['prompt', 'filePath', 'size'],
+      required: ['prompt', 'filePath', 'size', 'cheap'],
     },
   },
   {
@@ -309,6 +339,103 @@ Index: filename.js
       required: ['filePath', 'downloadUrl'],
     },
   },
+  {
+    name: 'imglyRemoveBackground',
+    description: 'Removes background from an image using @imgly/background-removal-node',
+    parameters: {
+      type: 'object',
+      properties: {
+        inputFilePath: {
+          type: 'string',
+          description: 'The file path of the input image.',
+        },
+        outputFilePath: {
+          type: 'string',
+          description: 'The file path to save the output image with removed background.',
+        },
+        explanation: {
+          type: 'string',
+          description: 'The explanation of the reasoning behind removing the background from this image',
+        },
+      },
+      required: ['inputFilePath', 'outputFilePath'],
+    },
+  },
+  {
+    name: 'resizeImage',
+    description: 'Resize image to the desired size',
+    parameters: {
+      type: 'object',
+      properties: {
+        filePath: {
+          type: 'string',
+          description: 'The file path of the image.',
+        },
+        size: {
+          type: 'object',
+          properties: {
+            width: {
+              type: 'number',
+              description: 'width of the image',
+            },
+            height: {
+              type: 'number',
+              description: 'height of the image',
+            },
+          },
+          required: ['width', 'height'],
+          description: 'The size of the image to generate.',
+        },
+        explanation: {
+          type: 'string',
+          description: 'The explanation of the reasoning behind removing the background from this image',
+        },
+      },
+      required: ['filePath', 'size'],
+    },
+  },
+  {
+    name: 'splitImage',
+    description: 'Split an image into multiple parts and save them as separate files.',
+    parameters: {
+      type: 'object',
+      properties: {
+        inputFilePath: {
+          type: 'string',
+          description: 'The file path of the input image to be split.',
+        },
+        parts: {
+          type: 'array',
+          items: {
+            type: 'object',
+            properties: {
+              rect: {
+                type: 'object',
+                properties: {
+                  x: { type: 'number', description: 'The x-coordinate of the top-left corner of the rectangle.' },
+                  y: { type: 'number', description: 'The y-coordinate of the top-left corner of the rectangle.' },
+                  width: { type: 'number', description: 'The width of the rectangle.' },
+                  height: { type: 'number', description: 'The height of the rectangle.' },
+                },
+                required: ['x', 'y', 'width', 'height'],
+              },
+              outputFilePath: {
+                type: 'string',
+                description: 'The file path to save the extracted part of the image.',
+              },
+            },
+            required: ['rect', 'outputFilePath'],
+          },
+          description: 'An array of parts to extract from the image, each with a rectangle and output file path.',
+        },
+        explanation: {
+          type: 'string',
+          description: 'The explanation of the reasoning behind splitting this image',
+        },
+      },
+      required: ['inputFilePath', 'parts'],
+    },
+  },
 ].map((fd) => {
   if (requireExplanations && fd.parameters.properties.explanation && !fd.parameters.required.includes('explanation')) {
     fd.parameters.required.push('explanation');

package/src/ai-service/vertex-ai-claude.js

@@ -4,7 +4,7 @@ import { printTokenUsageAndCost, processFunctionCalls } from './common.js';
 /**
  * This function generates content using the Anthropic Claude model via Vertex AI.
  */
-export async function generateContent(prompt, functionDefs, requiredFunctionName, temperature) {
+export async function generateContent(prompt, functionDefs, requiredFunctionName, temperature, cheap = false) {
   const projectId = process.env.GOOGLE_CLOUD_PROJECT;
   const region = process.env.GOOGLE_CLOUD_REGION;
 
@@ -56,8 +56,11 @@ export async function generateContent(prompt, functionDefs, requiredFunctionName
       }
     });
 
+  const model = cheap ? 'claude-3-haiku@20240307' : 'claude-3-5-sonnet@20240620';
+  console.log(`Using Vertex AI Claude model: ${model}`);
+
   const response = await client.messages.create({
-    model: 'claude-3-5-sonnet@20240620',
+    model: model,
     max_tokens: 4096,
     temperature: temperature,
     system: prompt.find((item) => item.type === 'systemPrompt').systemPrompt,

package/src/ai-service/vertex-ai-imagen.js

@@ -0,0 +1,78 @@
+import fs from 'fs';
+import mime from 'mime-types';
+import { PredictionServiceClient, helpers } from '@google-cloud/aiplatform';
+import { setTempBuffer } from '../files/temp-buffer.js';
+import { resizeImageBuffer } from '../images/resize-image.js';
+
+/**
+ * Generate an image using Vertex AI's Imagen model and return the image URL
+ * @param {string} prompt - The description of the image to generate
+ * @param {string|undefined} contextImagePath - The image to be used as a context
+ * @param {{width: number, height: number}} size - The size of the image to generate
+ * @param {boolean} cheap - Whether to use a cheaper model
+ * @returns {Promise<string>} - The url of the generated image
+ */
+export async function generateImage(prompt, contextImagePath, size, cheap = false) {
+  // Initialize the PredictionServiceClient
+  const client = new PredictionServiceClient({
+    apiEndpoint: `${process.env.GOOGLE_CLOUD_REGION}-aiplatform.googleapis.com`,
+  });
+
+  // Set the project and location
+  const projectId = process.env.GOOGLE_CLOUD_PROJECT;
+  const location = process.env.GOOGLE_CLOUD_REGION;
+
+  // Set the model name based on the cheap parameter
+  const modelName = contextImagePath
+    ? 'imagegeneration@002'
+    : cheap
+      ? 'imagen-3.0-fast-generate-001'
+      : 'imagen-3.0-generate-001';
+  console.log(`Using Vertex AI Imagen model: ${modelName}`);
+
+  try {
+    // Prepare the request
+    const request = {
+      endpoint: `projects/${projectId}/locations/${location}/publishers/google/models/${modelName}`,
+      instances: [
+        helpers.toValue({
+          prompt: prompt,
+          ...(contextImagePath
+            ? {
+                image: {
+                  bytesBase64Encoded: fs.readFileSync(contextImagePath, 'base64'),
+                  mediaType: mime.lookup(contextImagePath),
+                },
+              }
+            : {}),
+        }),
+      ],
+      parameters: helpers.toValue({
+        sampleCount: 1,
+        safetySetting: 'block_most',
+        personGeneration: 'allow_adult',
+        includeRaiReason: true,
+        language: 'auto',
+        aspectRatio: '1:1',
+        addWatermark: false,
+      }),
+    };
+
+    // Make the prediction request
+    const [response] = await client.predict(request);
+
+    if (response.predictions && response.predictions.length > 0) {
+      const prediction = helpers.fromValue(response.predictions[0]);
+      if (prediction.bytesBase64Encoded) {
+        console.log(`Image generated successfully, resizing to desired dimension`, size);
+        const buffer = Buffer.from(prediction.bytesBase64Encoded, 'base64');
+        return setTempBuffer(await resizeImageBuffer(buffer, size));
+      }
+    }
+
+    throw new Error('No image generated in the response');
+  } catch (error) {
+    console.error('Error generating image:', error);
+    throw error;
+  }
+}

package/src/ai-service/vertex-ai.js

@@ -7,7 +7,7 @@ import { geminiBlockNone } from '../cli/cli-params.js';
  * This function generates content using the Gemini Pro model.
  */
 
-export async function generateContent(prompt, functionDefs, requiredFunctionName, temperature) {
+export async function generateContent(prompt, functionDefs, requiredFunctionName, temperature, cheap = false) {
   const messages = prompt
     .filter((item) => item.type !== 'systemPrompt')
     .map((item) => {
@@ -53,16 +53,15 @@ export async function generateContent(prompt, functionDefs, requiredFunctionName
         functionDeclarations: functionDefs,
       },
     ],
-    // TODO: add tool_config once [it is supported](https://github.com/googleapis/nodejs-vertexai/issues/331)
     toolConfig: {
       functionCallingConfig: {
-        mode: 'ANY',
-        ...(requiredFunctionName ? { allowedFunctionNames: [requiredFunctionName] } : {}),
+        mode: cheap ? undefined : 'ANY',
+        ...(!cheap && requiredFunctionName ? { allowedFunctionNames: [requiredFunctionName] } : {}),
       },
     },
   };
 
-  const model = await getGenModel(prompt.find((item) => item.type === 'systemPrompt').systemPrompt, temperature);
+  const model = await getGenModel(prompt.find((item) => item.type === 'systemPrompt').systemPrompt, temperature, cheap);
 
   assert(await verifyVertexMonkeyPatch(), 'Vertex AI Tool Config was not monkey patched');
 
@@ -105,11 +104,13 @@ export async function generateContent(prompt, functionDefs, requiredFunctionName
 }
 
 // A function to get the generative model
-// Modified to accept temperature parameter
-export function getGenModel(systemPrompt, temperature) {
+// Modified to accept temperature parameter and cheap flag
+export function getGenModel(systemPrompt, temperature, cheap = false) {
   // Initialize Vertex with your Cloud project and location
   const vertex_ai = new VertexAI({});
-  const model = 'gemini-1.5-pro-001';
+  const model = cheap ? 'gemini-1.5-flash-001' : 'gemini-1.5-pro-001';
+
+  console.log(`Using Vertex AI model: ${model}`);
 
   // Instantiate the models
   return vertex_ai.preview.getGenerativeModel({

package/src/cli/cli-options.js

@@ -86,6 +86,20 @@ const cliOptions = [
     description:
       'Enable vision capabilities for processing image inputs. This option allows the tool to analyze and generate code based on image content when used with compatible AI models.',
   },
+  {
+    name: '--imagen=<service>',
+    description:
+      'Enable image generation functionality and specify the service to use (either "vertex-ai" or "dall-e").',
+  },
+  {
+    name: '--cheap',
+    description: 'Switch to cheaper models in AI services for content and image generation.',
+  },
+  {
+    name: '--content-mask=<path>',
+    description:
+      'Apply a content mask to limit the initial source code files included in the request. The value should be a prefix of the path relative to rootDir.',
+  },
 ];
 
 /**