From cacf0cc0ef97f781ec742ff883c70ee7b0a04cee Mon Sep 17 00:00:00 2001 From: Jaana Dogan Date: Mon, 21 Apr 2025 17:15:20 -0700 Subject: Simplify GeminiClient (#101) Doing some more clean-up: * Remove confusing continue/break * Handle empty result * Rename the file just client.js --- packages/server/src/core/client.ts | 162 +++++++++++++++++++ packages/server/src/core/gemini-client.ts | 173 --------------------- packages/server/src/index.ts | 2 +- .../server/src/utils/BackgroundTerminalAnalyzer.ts | 2 +- 4 files changed, 164 insertions(+), 175 deletions(-) create mode 100644 packages/server/src/core/client.ts delete mode 100644 packages/server/src/core/gemini-client.ts (limited to 'packages/server/src') diff --git a/packages/server/src/core/client.ts b/packages/server/src/core/client.ts new file mode 100644 index 00000000..e65e58a7 --- /dev/null +++ b/packages/server/src/core/client.ts @@ -0,0 +1,162 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { + GenerateContentConfig, + GoogleGenAI, + Part, + Chat, + 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'; + +export class GeminiClient { + private client: GoogleGenAI; + private model: string; + private generateContentConfig: GenerateContentConfig = { + temperature: 0, + topP: 1, + }; + private readonly MAX_TURNS = 100; + + constructor(apiKey: string, model: string) { + this.client = new GoogleGenAI({ apiKey: apiKey }); + this.model = model; + } + + private async getEnvironment(): Promise { + const cwd = process.cwd(); + const today = new Date().toLocaleDateString(undefined, { + weekday: 'long', + year: 'numeric', + month: 'long', + day: 'numeric', + }); + const platform = process.platform; + const folderStructure = await getFolderStructure(cwd); + const context = ` + Okay, just setting up the context for our chat. + Today is ${today}. + My operating system is: ${platform} + I'm currently working in the directory: ${cwd} + ${folderStructure} + `.trim(); + return { text: context }; + } + + async startChat(toolDeclarations: FunctionDeclaration[]): Promise { + const envPart = await this.getEnvironment(); + const tools: Tool[] = [{ functionDeclarations: toolDeclarations }]; + try { + return this.client.chats.create({ + model: this.model, + config: { + systemInstruction: CoreSystemPrompt, + ...this.generateContentConfig, + tools: tools, + }, + history: [ + { + role: 'user', + parts: [envPart], + }, + { + role: 'model', + parts: [{ text: 'Got it. Thanks for the context!' }], + }, + ], + }); + } catch (error) { + console.error('Error initializing Gemini chat session:', error); + const message = error instanceof Error ? error.message : 'Unknown error.'; + throw new Error(`Failed to initialize chat: ${message}`); + } + } + + async *sendMessageStream( + chat: Chat, + request: PartListUnion, + availableTools: ServerTool[], + signal?: AbortSignal, + ): AsyncGenerator { + let turns = 0; + try { + while (turns < this.MAX_TURNS) { + turns++; + const turn = new Turn(chat, availableTools); + const resultStream = turn.run(request, signal); + for await (const event of resultStream) { + yield event; + } + + const confirmations = turn.getConfirmationDetails(); + if (confirmations.length > 0) { + break; + } + + // What do we do when we have both function responses and confirmations? + const fnResponses = turn.getFunctionResponses(); + if (fnResponses.length == 0) { + break; // user's turn to respond + } + request = fnResponses; + } + if (turns >= this.MAX_TURNS) { + console.warn( + 'sendMessageStream: Reached maximum tool call turns limit.', + ); + } + } catch (error: unknown) { + if (error instanceof Error && error.name === 'AbortError') { + console.log('Gemini stream request aborted by user.'); + throw error; + } + console.error(`Error during Gemini stream or tool interaction:`, error); + throw error; + } + } + + async generateJson( + contents: Content[], + schema: SchemaUnion, + ): Promise> { + try { + const result = await this.client.models.generateContent({ + model: this.model, + config: { + ...this.generateContentConfig, + systemInstruction: CoreSystemPrompt, + responseSchema: schema, + responseMimeType: 'application/json', + }, + contents, + }); + if (!result || !result.text) { + throw new Error('API returned an empty response.'); + } + try { + return JSON.parse(result.text); + } catch (parseError) { + console.error('Failed to parse JSON response:', result.text); + throw new Error( + `Failed to parse API response as JSON: ${parseError instanceof Error ? parseError.message : String(parseError)}`, + ); + } + } catch (error) { + console.error('Error generating JSON content:', error); + const message = + error instanceof Error ? error.message : 'Unknown API error.'; + throw new Error(`Failed to generate JSON content: ${message}`); + } + } +} diff --git a/packages/server/src/core/gemini-client.ts b/packages/server/src/core/gemini-client.ts deleted file mode 100644 index 5829afa8..00000000 --- a/packages/server/src/core/gemini-client.ts +++ /dev/null @@ -1,173 +0,0 @@ -/** - * @license - * Copyright 2025 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import { - GenerateContentConfig, - GoogleGenAI, - Part, - Chat, - 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'; - -export class GeminiClient { - private ai: GoogleGenAI; - private model: string; - private generateContentConfig: GenerateContentConfig = { - temperature: 0, - topP: 1, - }; - private readonly MAX_TURNS = 100; - - constructor(apiKey: string, model: string) { - this.ai = new GoogleGenAI({ apiKey: apiKey }); - this.model = model; - } - - private async getEnvironment(): Promise { - const cwd = process.cwd(); - const today = new Date().toLocaleDateString(undefined, { - weekday: 'long', - year: 'numeric', - month: 'long', - day: 'numeric', - }); - const platform = process.platform; - const folderStructure = await getFolderStructure(cwd); - const context = ` - Okay, just setting up the context for our chat. - Today is ${today}. - My operating system is: ${platform} - I'm currently working in the directory: ${cwd} - ${folderStructure} - `.trim(); - return { text: context }; - } - - async startChat(toolDeclarations: FunctionDeclaration[]): Promise { - const envPart = await this.getEnvironment(); - // const tools: Tool[] = toolDeclarations.map((declaration) => ({ - // functionDeclarations: [declaration], - // })); - // merge all functions into a single tool, as seems to be required for gemini 2.5 series - // can test by asking "what tools do you have?", which lists single function unless merged - const tools: Tool[] = [{ functionDeclarations: toolDeclarations }]; - try { - const chat = this.ai.chats.create({ - model: this.model, - config: { - systemInstruction: CoreSystemPrompt, - ...this.generateContentConfig, - tools: tools, - }, - history: [ - { - role: 'user', - parts: [envPart], - }, - { - role: 'model', - parts: [{ text: 'Got it. Thanks for the context!' }], - }, - ], - }); - return chat; - } catch (error) { - console.error('Error initializing Gemini chat session:', error); - const message = error instanceof Error ? error.message : 'Unknown error.'; - throw new Error(`Failed to initialize chat: ${message}`); - } - } - - async *sendMessageStream( - chat: Chat, - request: PartListUnion, - availableTools: ServerTool[], - signal?: AbortSignal, - ): AsyncGenerator { - let turns = 0; - try { - while (turns < this.MAX_TURNS) { - turns++; - const turn = new Turn(chat, availableTools); - const resultStream = turn.run(request, signal); - for await (const event of resultStream) { - yield event; - } - - const confirmations = turn.getConfirmationDetails(); - if (confirmations.length > 0) { - break; - } - - // What do we do when we have both function responses and confirmations? - - const fnResponses = turn.getFunctionResponses(); - if (fnResponses.length > 0) { - request = fnResponses; - continue; - } else { - break; - } - } - if (turns >= this.MAX_TURNS) { - console.warn( - 'sendMessageStream: Reached maximum tool call turns limit.', - ); - } - } catch (error: unknown) { - if (error instanceof Error && error.name === 'AbortError') { - console.log('Gemini stream request aborted by user.'); - throw error; - } - console.error(`Error during Gemini stream or tool interaction:`, error); - throw error; - } - } - - async generateJson( - contents: Content[], - schema: SchemaUnion, - ): Promise> { - try { - const result = await this.ai.models.generateContent({ - model: this.model, - config: { - ...this.generateContentConfig, - systemInstruction: CoreSystemPrompt, - responseSchema: schema, - responseMimeType: 'application/json', - }, - contents, - }); - const responseText = result.text; - if (!responseText) { - throw new Error('API returned an empty response.'); - } - try { - const parsedJson = JSON.parse(responseText); - return parsedJson; - } catch (parseError) { - console.error('Failed to parse JSON response:', responseText); - throw new Error( - `Failed to parse API response as JSON: ${parseError instanceof Error ? parseError.message : String(parseError)}`, - ); - } - } catch (error) { - console.error('Error generating JSON content:', error); - const message = - error instanceof Error ? error.message : 'Unknown API error.'; - throw new Error(`Failed to generate JSON content: ${message}`); - } - } -} diff --git a/packages/server/src/index.ts b/packages/server/src/index.ts index 2cbfbe6b..998ed584 100644 --- a/packages/server/src/index.ts +++ b/packages/server/src/index.ts @@ -8,7 +8,7 @@ export * from './config/config.js'; // Export Core Logic -export * from './core/gemini-client.js'; +export * from './core/client.js'; export * from './core/prompts.js'; export * from './core/turn.js'; // Potentially export types from turn.ts if needed externally diff --git a/packages/server/src/utils/BackgroundTerminalAnalyzer.ts b/packages/server/src/utils/BackgroundTerminalAnalyzer.ts index 54ad9d50..625b06b6 100644 --- a/packages/server/src/utils/BackgroundTerminalAnalyzer.ts +++ b/packages/server/src/utils/BackgroundTerminalAnalyzer.ts @@ -6,7 +6,7 @@ import { Content, SchemaUnion, Type } from '@google/genai'; import { getErrorMessage, isNodeError } from '../utils/errors.js'; -import { GeminiClient } from '../core/gemini-client.js'; +import { GeminiClient } from '../core/client.js'; import { Config } from '../config/config.js'; import { promises as fs } from 'fs'; import { exec as _exec } from 'child_process'; -- cgit v1.2.3