From 8e0fb9ee2f9fa1a6c97fc2e8cebd67da0386c5e9 Mon Sep 17 00:00:00 2001 From: Jaana Dogan Date: Tue, 22 Apr 2025 11:01:09 -0700 Subject: [PATCH] Initiate the GeminiClient with a config Also address the open readability improvement comments from #104. --- packages/cli/src/gemini.ts | 72 +++++++------------ packages/cli/src/ui/App.tsx | 14 +--- packages/cli/src/ui/hooks/useGeminiStream.ts | 38 ++-------- packages/server/src/core/client.ts | 19 +++-- .../src/utils/BackgroundTerminalAnalyzer.ts | 5 +- 5 files changed, 43 insertions(+), 105 deletions(-) diff --git a/packages/cli/src/gemini.ts b/packages/cli/src/gemini.ts index 0d8b1ac7..56a44e30 100644 --- a/packages/cli/src/gemini.ts +++ b/packages/cli/src/gemini.ts @@ -9,24 +9,9 @@ import { render } from 'ink'; import { App } from './ui/App.js'; import { loadCliConfig } from './config/config.js'; import { readStdin } from './utils/readStdin.js'; -import { GeminiClient, ServerTool } from '@gemini-code/server'; - -import { PartListUnion } from '@google/genai'; +import { GeminiClient } from '@gemini-code/server'; async function main() { - let initialInput: string | undefined = undefined; - - // Check if input is being piped - if (!process.stdin.isTTY) { - try { - initialInput = await readStdin(); - } catch (error) { - console.error('Error reading from stdin:', error); - process.exit(1); - } - } - - // Load configuration const config = loadCliConfig(); // Render UI, passing necessary config values and initial input @@ -34,44 +19,35 @@ async function main() { render( React.createElement(App, { config, - initialInput, }), ); - } else if (initialInput) { - // If not a TTY and we have initial input, process it directly - const geminiClient = new GeminiClient( - config.getApiKey(), - config.getModel(), - ); - const toolRegistry = config.getToolRegistry(); - const availableTools: ServerTool[] = toolRegistry.getAllTools(); - const toolDeclarations = toolRegistry.getFunctionDeclarations(); - const chat = await geminiClient.startChat(toolDeclarations); + return; + } - const request: PartListUnion = [{ text: initialInput }]; - - try { - for await (const event of geminiClient.sendMessageStream( - chat, - request, - availableTools, - )) { - if (event.type === 'content') { - process.stdout.write(event.value); - } - // We might need to handle other event types later, but for now, just content. - } - process.stdout.write('\n'); // Add a newline at the end - process.exit(0); - } catch (error) { - console.error('Error processing piped input:', error); - process.exit(1); - } - } else { - // If not a TTY and no initial input, exit with an error + const input = await readStdin(); + if (!input) { console.error('No input provided via stdin.'); process.exit(1); } + + // If not a TTY and we have initial input, process it directly + const geminiClient = new GeminiClient(config); + const chat = await geminiClient.startChat(); + try { + for await (const event of geminiClient.sendMessageStream(chat, [ + { text: input }, + ])) { + if (event.type === 'content') { + process.stdout.write(event.value); + } + // We might need to handle other event types later, but for now, just content. + } + process.stdout.write('\n'); // Add a newline at the end + process.exit(0); + } catch (error) { + console.error('Error processing piped input:', error); + process.exit(1); + } } // --- Global Unhandled Rejection Handler --- diff --git a/packages/cli/src/ui/App.tsx b/packages/cli/src/ui/App.tsx index b99a57c3..43a1d1e6 100644 --- a/packages/cli/src/ui/App.tsx +++ b/packages/cli/src/ui/App.tsx @@ -4,7 +4,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -import React, { useState, useMemo, useEffect } from 'react'; // Added useEffect +import React, { useState, useMemo } from 'react'; import { Box, Text } from 'ink'; import { StreamingState, type HistoryItem } from './types.js'; import { useGeminiStream } from './hooks/useGeminiStream.js'; @@ -25,10 +25,9 @@ import { Colors } from './colors.js'; interface AppProps { config: Config; - initialInput?: string; // Added optional prop } -export const App = ({ config, initialInput }: AppProps) => { +export const App = ({ config }: AppProps) => { // Destructured prop const [history, setHistory] = useState([]); const [startupWarnings, setStartupWarnings] = useState([]); @@ -40,15 +39,6 @@ export const App = ({ config, initialInput }: AppProps) => { useStartupWarnings(setStartupWarnings); useInitializationErrorEffect(initError, history, setHistory); - // Effect to handle initial piped input - useEffect(() => { - if (initialInput && initialInput.trim() !== '') { - submitQuery(initialInput); - } - // Run only once on mount - // eslint-disable-next-line react-hooks/exhaustive-deps - }, []); - const userMessages = useMemo( () => history diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index b8d13269..c4d44749 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -7,13 +7,11 @@ import { exec as _exec } from 'child_process'; import { useState, useRef, useCallback, useEffect } from 'react'; import { useInput } from 'ink'; -// Import server-side client and types import { GeminiClient, GeminiEventType as ServerGeminiEventType, // Rename to avoid conflict getErrorMessage, isNodeError, - ToolResult, Config, ToolCallConfirmationDetails, ToolCallResponseInfo, @@ -23,12 +21,7 @@ import { ToolEditConfirmationDetails, ToolExecuteConfirmationDetails, } from '@gemini-code/server'; -import { - type Chat, - type PartListUnion, - type FunctionDeclaration, - type Part, -} from '@google/genai'; +import { type Chat, type PartListUnion, type Part } from '@google/genai'; import { StreamingState, HistoryItem, @@ -69,10 +62,7 @@ export const useGeminiStream = ( setInitError(null); if (!geminiClientRef.current) { try { - geminiClientRef.current = new GeminiClient( - config.getApiKey(), - config.getModel(), - ); + geminiClientRef.current = new GeminiClient(config); } catch (error: unknown) { setInitError( `Failed to initialize client: ${getErrorMessage(error) || 'Unknown error'}`, @@ -166,9 +156,7 @@ export const useGeminiStream = ( if (!chatSessionRef.current) { try { - // Use getFunctionDeclarations for startChat - const toolSchemas = toolRegistry.getFunctionDeclarations(); - chatSessionRef.current = await client.startChat(toolSchemas); + chatSessionRef.current = await client.startChat(); } catch (err: unknown) { setInitError(`Failed to start chat: ${getErrorMessage(err)}`); setStreamingState(StreamingState.Idle); @@ -196,15 +184,7 @@ export const useGeminiStream = ( abortControllerRef.current = new AbortController(); const signal = abortControllerRef.current.signal; - // Get ServerTool descriptions for the server call - const serverTools: ServerTool[] = toolRegistry.getAllTools(); - - const stream = client.sendMessageStream( - chat, - query, - serverTools, - signal, - ); + const stream = client.sendMessageStream(chat, query, signal); // Process the stream events from the server logic let currentGeminiText = ''; // To accumulate message content @@ -477,13 +457,3 @@ export const useGeminiStream = ( return { streamingState, submitQuery, initError, debugMessage }; }; - -// Define ServerTool interface here if not importing from server (circular dep issue?) -interface ServerTool { - name: string; - schema: FunctionDeclaration; - shouldConfirmExecute( - params: Record, - ): Promise; - execute(params: Record): Promise; -} diff --git a/packages/server/src/core/client.ts b/packages/server/src/core/client.ts index e65e58a7..e930be6c 100644 --- a/packages/server/src/core/client.ts +++ b/packages/server/src/core/client.ts @@ -12,15 +12,16 @@ import { SchemaUnion, PartListUnion, Content, - FunctionDeclaration, Tool, } from '@google/genai'; import { CoreSystemPrompt } from './prompts.js'; import process from 'node:process'; import { getFolderStructure } from '../utils/getFolderStructure.js'; -import { Turn, ServerTool, ServerGeminiStreamEvent } from './turn.js'; +import { Turn, ServerGeminiStreamEvent } from './turn.js'; +import { Config } from '../config/config.js'; export class GeminiClient { + private config: Config; private client: GoogleGenAI; private model: string; private generateContentConfig: GenerateContentConfig = { @@ -29,9 +30,10 @@ export class GeminiClient { }; private readonly MAX_TURNS = 100; - constructor(apiKey: string, model: string) { - this.client = new GoogleGenAI({ apiKey: apiKey }); - this.model = model; + constructor(config: Config) { + this.client = new GoogleGenAI({ apiKey: config.getApiKey() }); + this.config = config; + this.model = config.getModel(); } private async getEnvironment(): Promise { @@ -54,8 +56,11 @@ export class GeminiClient { return { text: context }; } - async startChat(toolDeclarations: FunctionDeclaration[]): Promise { + async startChat(): Promise { const envPart = await this.getEnvironment(); + const toolDeclarations = this.config + .getToolRegistry() + .getFunctionDeclarations(); const tools: Tool[] = [{ functionDeclarations: toolDeclarations }]; try { return this.client.chats.create({ @@ -86,10 +91,10 @@ export class GeminiClient { async *sendMessageStream( chat: Chat, request: PartListUnion, - availableTools: ServerTool[], signal?: AbortSignal, ): AsyncGenerator { let turns = 0; + const availableTools = this.config.getToolRegistry().getAllTools(); try { while (turns < this.MAX_TURNS) { turns++; diff --git a/packages/server/src/utils/BackgroundTerminalAnalyzer.ts b/packages/server/src/utils/BackgroundTerminalAnalyzer.ts index 625b06b6..c63fbb57 100644 --- a/packages/server/src/utils/BackgroundTerminalAnalyzer.ts +++ b/packages/server/src/utils/BackgroundTerminalAnalyzer.ts @@ -78,10 +78,7 @@ export class BackgroundTerminalAnalyzer { ) { try { // Initialize Gemini client using config - this.geminiClient = new GeminiClient( - config.getApiKey(), - config.getModel(), - ); + this.geminiClient = new GeminiClient(config); } catch (error) { console.error( 'Failed to initialize GeminiClient in BackgroundTerminalAnalyzer:',