From cb303514035440b1631964cad5093a4c80cd3e43 Mon Sep 17 00:00:00 2001 From: Evan Senter Date: Fri, 18 Apr 2025 17:06:16 +0100 Subject: [PATCH] Adding a new parameter for model, and updating the default to 2.5 Flash. (#18) --- packages/cli/src/config/args.ts | 9 ++++ packages/cli/src/config/globalConfig.ts | 50 +++++++++++++++++++ packages/cli/src/core/gemini-client.ts | 7 ++- packages/cli/src/gemini.ts | 10 +++- packages/cli/src/tools/terminal.tool.ts | 5 +- .../cli/src/ui/components/InputPrompt.tsx | 5 +- .../src/utils/BackgroundTerminalAnalyzer.ts | 2 +- 7 files changed, 81 insertions(+), 7 deletions(-) create mode 100644 packages/cli/src/config/globalConfig.ts 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 { 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): 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 { 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 { + 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 = ({ setQuery, onSubmit, }) => { + const model = getModel(); + return ( > @@ -24,7 +27,7 @@ const InputPrompt: React.FC = ({ onSubmit={onSubmit} showCursor={true} focus={true} - placeholder={'Ask Gemini... (try "/init" or "/help")'} + placeholder={`Ask Gemini (${model})... (try "/init" or "/help")`} /> 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