Adding a new parameter for model, and updating the default to 2.5 Flash. (#18)

7 files changed, 81 insertions, 7 deletions

diff --git a/packages/cli/src/config/args.ts b/packages/cli/src/config/args.ts
index 7c4ebbc0..f36e7e58 100644
--- a/packages/cli/src/config/args.ts
+++ b/packages/cli/src/config/args.ts
@@ -1,8 +1,11 @@
 import yargs from 'yargs/yargs';
 import { hideBin } from 'yargs/helpers';
 
+const DEFAULT_GEMINI_MODEL = 'gemini-2.5-flash-preview-04-17';
+
 export interface CliArgs {
   target_dir: string | undefined;
+  model: string | undefined;
   _: (string | number)[]; // Captures positional arguments
   // Add other expected args here if needed
   // e.g., verbose?: boolean;
@@ -16,6 +19,12 @@ export async function parseArguments(): Promise<CliArgs> {
       description:
         'The target directory for Gemini operations. Defaults to the current working directory.',
     })
+    .option('model', {
+      alias: 'm',
+      type: 'string',
+      description: `The Gemini model to use. Defaults to ${DEFAULT_GEMINI_MODEL}.`,
+      default: DEFAULT_GEMINI_MODEL,
+    })
     .help()
     .alias('h', 'help')
     .strict() // Keep strict mode to error on unknown options
diff --git a/packages/cli/src/config/globalConfig.ts b/packages/cli/src/config/globalConfig.ts
new file mode 100644
index 00000000..2b6ad518
--- /dev/null
+++ b/packages/cli/src/config/globalConfig.ts
@@ -0,0 +1,50 @@
+import { CliArgs } from './args.js'; // Assuming CliArgs contains the needed fields
+
+interface GlobalConfig {
+  model: string;
+  // Add other global config values here if needed
+  // e.g., targetDir?: string;
+}
+
+let config: GlobalConfig | null = null;
+
+/**
+ * Initializes the global configuration. Should only be called once at application startup.
+ * @param args The parsed command-line arguments.
+ */
+export function initializeConfig(args: Pick<CliArgs, 'model'>): void {
+  if (config) {
+    console.warn('Global configuration already initialized.');
+    return;
+  }
+  if (!args.model) {
+    // This shouldn't happen if default is set correctly in args.ts
+    throw new Error('Model not provided during config initialization.');
+  }
+  config = {
+    model: args.model,
+    // Initialize other config values from args here
+  };
+}
+
+/**
+ * Retrieves the globally stored configuration.
+ * Throws an error if the configuration has not been initialized.
+ * @returns The global configuration object.
+ */
+export function getConfig(): GlobalConfig {
+  if (!config) {
+    throw new Error(
+      'Global configuration accessed before initialization. Call initializeConfig() first.',
+    );
+  }
+  return config;
+}
+
+/**
+ * Helper function to get the configured Gemini model name.
+ * @returns The model name string.
+ */
+export function getModel(): string {
+  return getConfig().model;
+} 
\ No newline at end of file
diff --git a/packages/cli/src/core/gemini-client.ts b/packages/cli/src/core/gemini-client.ts
index 41cabdb7..c7c7b5f6 100644
--- a/packages/cli/src/core/gemini-client.ts
+++ b/packages/cli/src/core/gemini-client.ts
@@ -9,6 +9,7 @@ import {
   Content,
 } from '@google/genai';
 import { getApiKey } from '../config/env.js';
+import { getModel } from '../config/globalConfig.js';
 import { CoreSystemPrompt } from './prompts.js';
 import {
   type ToolCallEvent,
@@ -45,6 +46,7 @@ export class GeminiClient {
 
   public async startChat(): Promise<Chat> {
     const tools = toolRegistry.getToolSchemas();
+    const model = getModel();
 
     // --- Get environmental information ---
     const cwd = process.cwd();
@@ -73,7 +75,7 @@ ${folderStructure}
 
     try {
       const chat = this.ai.chats.create({
-        model: 'gemini-2.0-flash', //'gemini-2.0-flash',
+        model: model,
         config: {
           systemInstruction: CoreSystemPrompt,
           ...this.defaultHyperParameters,
@@ -446,9 +448,10 @@ Respond *only* in JSON format according to the following schema. Do not include
     contents: Content[],
     schema: SchemaUnion,
   ): Promise<any> {
+    const model = getModel();
     try {
       const result = await this.ai.models.generateContent({
-        model: 'gemini-2.0-flash', // Using flash for potentially faster structured output
+        model: model,
         config: {
           ...this.defaultHyperParameters,
           systemInstruction: CoreSystemPrompt,
diff --git a/packages/cli/src/gemini.ts b/packages/cli/src/gemini.ts
index 175dd9c6..a2797dc8 100644
--- a/packages/cli/src/gemini.ts
+++ b/packages/cli/src/gemini.ts
@@ -3,6 +3,7 @@ import { render } from 'ink';
 import App from './ui/App.js';
 import { parseArguments } from './config/args.js';
 import { loadEnvironment } from './config/env.js';
+import { initializeConfig } from './config/globalConfig.js';
 import { getTargetDirectory } from './utils/paths.js';
 import { toolRegistry } from './tools/tool-registry.js';
 import { LSTool } from './tools/ls.tool.js';
@@ -16,14 +17,19 @@ import { WriteFileTool } from './tools/write-file.tool.js';
 async function main() {
   // 1. Configuration
   loadEnvironment();
-  const argv = await parseArguments(); // Ensure args.ts imports printWarning from ui/display
+  const argv = await parseArguments();
+  initializeConfig({ model: argv.model as string });
   const targetDir = getTargetDirectory(argv.target_dir);
 
   // 2. Configure tools
   registerTools(targetDir);
 
   // 3. Render UI
-  render(React.createElement(App, { directory: targetDir }));
+  render(
+    React.createElement(App, {
+      directory: targetDir,
+    }),
+  );
 }
 
 // --- Global Unhandled Rejection Handler ---
diff --git a/packages/cli/src/tools/terminal.tool.ts b/packages/cli/src/tools/terminal.tool.ts
index a51f8e6b..3e2d7bf8 100644
--- a/packages/cli/src/tools/terminal.tool.ts
+++ b/packages/cli/src/tools/terminal.tool.ts
@@ -142,7 +142,10 @@ export class TerminalTool extends BaseTool<
   private rejectShellReady: ((reason?: any) => void) | undefined; // Definite assignment assertion
   private readonly backgroundTerminalAnalyzer: BackgroundTerminalAnalyzer;
 
-  constructor(rootDirectory: string, outputLimit: number = MAX_OUTPUT_LENGTH) {
+  constructor(
+    rootDirectory: string,
+    outputLimit: number = MAX_OUTPUT_LENGTH,
+  ) {
     const toolDisplayName = 'Terminal';
     // --- LLM-Facing Description ---
     // Updated description for background tasks to mention polling and LLM analysis
diff --git a/packages/cli/src/ui/components/InputPrompt.tsx b/packages/cli/src/ui/components/InputPrompt.tsx
index b1832c04..cf28960e 100644
--- a/packages/cli/src/ui/components/InputPrompt.tsx
+++ b/packages/cli/src/ui/components/InputPrompt.tsx
@@ -1,6 +1,7 @@
 import React from 'react';
 import { Box, Text } from 'ink';
 import TextInput from 'ink-text-input';
+import { getModel } from '../../config/globalConfig.js';
 
 interface InputPromptProps {
   query: string;
@@ -14,6 +15,8 @@ const InputPrompt: React.FC<InputPromptProps> = ({
   setQuery,
   onSubmit,
 }) => {
+  const model = getModel();
+
   return (
     <Box marginTop={1} borderStyle="round" borderColor={'white'} paddingX={1}>
       <Text color={'white'}>&gt; </Text>
@@ -24,7 +27,7 @@ const InputPrompt: React.FC<InputPromptProps> = ({
           onSubmit={onSubmit}
           showCursor={true}
           focus={true}
-          placeholder={'Ask Gemini... (try "/init" or "/help")'}
+          placeholder={`Ask Gemini (${model})... (try "/init" or "/help")`}
         />
       </Box>
     </Box>
diff --git a/packages/cli/src/utils/BackgroundTerminalAnalyzer.ts b/packages/cli/src/utils/BackgroundTerminalAnalyzer.ts
index b6c96290..3e958d09 100644
--- a/packages/cli/src/utils/BackgroundTerminalAnalyzer.ts
+++ b/packages/cli/src/utils/BackgroundTerminalAnalyzer.ts
@@ -60,7 +60,7 @@ export class BackgroundTerminalAnalyzer {
       initialDelayMs?: number;
     } = {}, // Provide default options
   ) {
-    this.ai = aiClient || new GeminiClient(); // Use injected client or default
+    this.ai = aiClient || new GeminiClient(); // Call constructor without model
     this.pollIntervalMs = options.pollIntervalMs ?? 5000; // Default 5 seconds
     this.maxAttempts = options.maxAttempts ?? 6; // Default 6 attempts (approx 30s total)
     this.initialDelayMs = options.initialDelayMs ?? 500; // Default 0.5s initial delay