From 8c36b106828bb54243f3771cf89fa144942c1280 Mon Sep 17 00:00:00 2001 From: joshualitt Date: Fri, 21 Nov 2025 17:27:57 -0800 Subject: [PATCH] feat(core): Add `BaseLlmClient.generateContent`. (#13591) --- packages/core/src/core/baseLlmClient.test.ts | 141 +++++++++++- packages/core/src/core/baseLlmClient.ts | 215 +++++++++++++------ 2 files changed, 280 insertions(+), 76 deletions(-) diff --git a/packages/core/src/core/baseLlmClient.test.ts b/packages/core/src/core/baseLlmClient.test.ts index d924fabacb..d1596451c1 100644 --- a/packages/core/src/core/baseLlmClient.test.ts +++ b/packages/core/src/core/baseLlmClient.test.ts @@ -12,6 +12,7 @@ import { beforeEach, afterEach, type Mocked, + type Mock, } from 'vitest'; import type { GenerateContentResponse } from '@google/genai'; @@ -299,6 +300,45 @@ describe('BaseLlmClient', () => { expect(result).toEqual({ color: 'orange' }); expect(logMalformedJsonResponse).not.toHaveBeenCalled(); }); + + it('should use the resolved model name when logging malformed JSON telemetry', async () => { + const aliasModel = 'fast-alias'; + const resolvedModel = 'gemini-1.5-flash'; + + // Override the mock for this specific test to simulate resolution + ( + mockConfig.modelConfigService.getResolvedConfig as unknown as Mock + ).mockReturnValue({ + model: resolvedModel, + generateContentConfig: { + temperature: 0, + topP: 1, + }, + }); + + const malformedResponse = '```json\n{"color": "red"}\n```'; + mockGenerateContent.mockResolvedValue( + createMockResponse(malformedResponse), + ); + + const options = { + ...defaultOptions, + modelConfigKey: { model: aliasModel }, + }; + + const result = await client.generateJson(options); + + expect(result).toEqual({ color: 'red' }); + + expect(logMalformedJsonResponse).toHaveBeenCalled(); + const calls = vi.mocked(logMalformedJsonResponse).mock.calls; + const lastCall = calls[calls.length - 1]; + const event = lastCall[1] as MalformedJsonResponseEvent; + + // This is the key assertion: it should be the resolved model, not the alias + expect(event.model).toBe(resolvedModel); + expect(event.model).not.toBe(aliasModel); + }); }); describe('generateJson - Error Handling', () => { @@ -306,14 +346,14 @@ describe('BaseLlmClient', () => { mockGenerateContent.mockResolvedValue(createMockResponse('')); await expect(client.generateJson(defaultOptions)).rejects.toThrow( - 'Failed to generate JSON content: Retry attempts exhausted for invalid content', + 'Failed to generate content: Retry attempts exhausted for invalid content', ); // Verify error reporting details expect(reportError).toHaveBeenCalledTimes(1); expect(reportError).toHaveBeenCalledWith( expect.any(Error), - 'API returned invalid content (empty or unparsable JSON) after all retries.', + 'API returned invalid content after all retries.', defaultOptions.contents, 'generateJson-invalid-content', ); @@ -324,13 +364,13 @@ describe('BaseLlmClient', () => { mockGenerateContent.mockResolvedValue(createMockResponse(invalidJson)); await expect(client.generateJson(defaultOptions)).rejects.toThrow( - 'Failed to generate JSON content: Retry attempts exhausted for invalid content', + 'Failed to generate content: Retry attempts exhausted for invalid content', ); expect(reportError).toHaveBeenCalledTimes(1); expect(reportError).toHaveBeenCalledWith( expect.any(Error), - 'API returned invalid content (empty or unparsable JSON) after all retries.', + 'API returned invalid content after all retries.', defaultOptions.contents, 'generateJson-invalid-content', ); @@ -342,14 +382,14 @@ describe('BaseLlmClient', () => { mockGenerateContent.mockRejectedValue(apiError); await expect(client.generateJson(defaultOptions)).rejects.toThrow( - 'Failed to generate JSON content: Service Unavailable (503)', + 'Failed to generate content: Service Unavailable (503)', ); // Verify generic error reporting expect(reportError).toHaveBeenCalledTimes(1); expect(reportError).toHaveBeenCalledWith( apiError, - 'Error generating JSON content via API.', + 'Error generating content via API.', defaultOptions.contents, 'generateJson-api', ); @@ -464,4 +504,93 @@ describe('BaseLlmClient', () => { ); }); }); + + describe('generateContent', () => { + it('should call generateContent with correct parameters and utilize retry mechanism', async () => { + const mockResponse = createMockResponse('This is the content.'); + mockGenerateContent.mockResolvedValue(mockResponse); + + const options = { + modelConfigKey: { model: 'test-model' }, + contents: [{ role: 'user', parts: [{ text: 'Give me content.' }] }], + abortSignal: abortController.signal, + promptId: 'content-prompt-id', + }; + + const result = await client.generateContent(options); + + expect(result).toBe(mockResponse); + + // Ensure the retry mechanism was engaged + expect(retryWithBackoff).toHaveBeenCalledTimes(1); + expect(retryWithBackoff).toHaveBeenCalledWith( + expect.any(Function), + expect.objectContaining({ + shouldRetryOnContent: expect.any(Function), + }), + ); + + // Validate the parameters passed to the underlying generator + expect(mockGenerateContent).toHaveBeenCalledTimes(1); + expect(mockGenerateContent).toHaveBeenCalledWith( + { + model: 'test-model', + contents: options.contents, + config: { + abortSignal: options.abortSignal, + temperature: 0, + topP: 1, + }, + }, + 'content-prompt-id', + ); + }); + + it('should validate content using shouldRetryOnContent function', async () => { + const mockResponse = createMockResponse('Some valid content.'); + mockGenerateContent.mockResolvedValue(mockResponse); + + const options = { + modelConfigKey: { model: 'test-model' }, + contents: [{ role: 'user', parts: [{ text: 'Give me content.' }] }], + abortSignal: abortController.signal, + promptId: 'content-prompt-id', + }; + + await client.generateContent(options); + + const retryCall = vi.mocked(retryWithBackoff).mock.calls[0]; + const shouldRetryOnContent = retryCall[1]?.shouldRetryOnContent; + + // Valid content should not trigger retry + expect(shouldRetryOnContent!(mockResponse)).toBe(false); + + // Empty response should trigger retry + expect(shouldRetryOnContent!(createMockResponse(''))).toBe(true); + expect(shouldRetryOnContent!(createMockResponse(' '))).toBe(true); + }); + + it('should throw and report error for empty response after retry exhaustion', async () => { + mockGenerateContent.mockResolvedValue(createMockResponse('')); + const options = { + modelConfigKey: { model: 'test-model' }, + contents: [{ role: 'user', parts: [{ text: 'Give me content.' }] }], + abortSignal: abortController.signal, + promptId: 'content-prompt-id', + }; + + await expect(client.generateContent(options)).rejects.toThrow( + 'Failed to generate content: Retry attempts exhausted for invalid content', + ); + + // Verify error reporting details + expect(reportError).toHaveBeenCalledTimes(1); + expect(reportError).toHaveBeenCalledWith( + expect.any(Error), + 'API returned invalid content after all retries.', + options.contents, + 'generateContent-invalid-content', + ); + }); + }); }); diff --git a/packages/core/src/core/baseLlmClient.ts b/packages/core/src/core/baseLlmClient.ts index adaa3e3c39..166579b166 100644 --- a/packages/core/src/core/baseLlmClient.ts +++ b/packages/core/src/core/baseLlmClient.ts @@ -6,10 +6,10 @@ import type { Content, - GenerateContentConfig, Part, EmbedContentParameters, GenerateContentResponse, + GenerateContentParameters, } from '@google/genai'; import type { Config } from '../config/config.js'; import type { ContentGenerator } from './contentGenerator.js'; @@ -50,6 +50,31 @@ export interface GenerateJsonOptions { maxAttempts?: number; } +/** + * Options for the generateContent utility function. + */ +export interface GenerateContentOptions { + /** The desired model config. */ + modelConfigKey: ModelConfigKey; + /** The input prompt or history. */ + contents: Content[]; + /** + * Task-specific system instructions. + * If omitted, no system instruction is sent. + */ + systemInstruction?: string | Part | Part[] | Content; + /** Signal for cancellation. */ + abortSignal: AbortSignal; + /** + * A unique ID for the prompt, used for logging/telemetry correlation. + */ + promptId: string; + /** + * The maximum number of attempts for the request. + */ + maxAttempts?: number; +} + /** * A client dedicated to stateless, utility-focused LLM calls. */ @@ -63,87 +88,54 @@ export class BaseLlmClient { options: GenerateJsonOptions, ): Promise> { const { + schema, modelConfigKey, contents, - schema, - abortSignal, systemInstruction, + abortSignal, promptId, maxAttempts, } = options; const { model, generateContentConfig } = this.config.modelConfigService.getResolvedConfig(modelConfigKey); - const requestConfig: GenerateContentConfig = { - abortSignal, - ...generateContentConfig, - ...(systemInstruction && { systemInstruction }), - responseJsonSchema: schema, - responseMimeType: 'application/json', + + const shouldRetryOnContent = (response: GenerateContentResponse) => { + const text = getResponseText(response)?.trim(); + if (!text) { + return true; // Retry on empty response + } + try { + // We don't use the result, just check if it's valid JSON + JSON.parse(this.cleanJsonResponse(text, model)); + return false; // It's valid, don't retry + } catch (_e) { + return true; // It's not valid, retry + } }; - try { - const apiCall = () => - this.contentGenerator.generateContent( - { - model, - config: requestConfig, - contents, - }, - promptId, - ); + const result = await this._generateWithRetry( + { + model, + contents, + config: { + ...generateContentConfig, + ...(systemInstruction && { systemInstruction }), + responseJsonSchema: schema, + responseMimeType: 'application/json', + abortSignal, + }, + }, + promptId, + maxAttempts, + shouldRetryOnContent, + 'generateJson', + ); - const shouldRetryOnContent = (response: GenerateContentResponse) => { - const text = getResponseText(response)?.trim(); - if (!text) { - return true; // Retry on empty response - } - try { - JSON.parse(this.cleanJsonResponse(text, model)); - return false; - } catch (_e) { - return true; - } - }; - - const result = await retryWithBackoff(apiCall, { - shouldRetryOnContent, - maxAttempts: maxAttempts ?? DEFAULT_MAX_ATTEMPTS, - }); - - // If we are here, the content is valid (not empty and parsable). - return JSON.parse( - this.cleanJsonResponse(getResponseText(result)!.trim(), model), - ); - } catch (error) { - if (abortSignal.aborted) { - throw error; - } - - // Check if the error is from exhausting retries, and report accordingly. - if ( - error instanceof Error && - error.message.includes('Retry attempts exhausted') - ) { - await reportError( - error, - 'API returned invalid content (empty or unparsable JSON) after all retries.', - contents, - 'generateJson-invalid-content', - ); - } else { - await reportError( - error, - 'Error generating JSON content via API.', - contents, - 'generateJson-api', - ); - } - - throw new Error( - `Failed to generate JSON content: ${getErrorMessage(error)}`, - ); - } + // If we are here, the content is valid (not empty and parsable). + return JSON.parse( + this.cleanJsonResponse(getResponseText(result)!.trim(), model), + ); } async generateEmbedding(texts: string[]): Promise { @@ -193,4 +185,87 @@ export class BaseLlmClient { } return text; } + + async generateContent( + options: GenerateContentOptions, + ): Promise { + const { + modelConfigKey, + contents, + systemInstruction, + abortSignal, + promptId, + maxAttempts, + } = options; + + const { model, generateContentConfig } = + this.config.modelConfigService.getResolvedConfig(modelConfigKey); + + const shouldRetryOnContent = (response: GenerateContentResponse) => { + const text = getResponseText(response)?.trim(); + return !text; // Retry on empty response + }; + + return this._generateWithRetry( + { + model, + contents, + config: { + ...generateContentConfig, + ...(systemInstruction && { systemInstruction }), + abortSignal, + }, + }, + promptId, + maxAttempts, + shouldRetryOnContent, + 'generateContent', + ); + } + + private async _generateWithRetry( + requestParams: GenerateContentParameters, + promptId: string, + maxAttempts: number | undefined, + shouldRetryOnContent: (response: GenerateContentResponse) => boolean, + errorContext: 'generateJson' | 'generateContent', + ): Promise { + const abortSignal = requestParams.config?.abortSignal; + + try { + const apiCall = () => + this.contentGenerator.generateContent(requestParams, promptId); + + return await retryWithBackoff(apiCall, { + shouldRetryOnContent, + maxAttempts: maxAttempts ?? DEFAULT_MAX_ATTEMPTS, + }); + } catch (error) { + if (abortSignal?.aborted) { + throw error; + } + + // Check if the error is from exhausting retries, and report accordingly. + if ( + error instanceof Error && + error.message.includes('Retry attempts exhausted') + ) { + await reportError( + error, + `API returned invalid content after all retries.`, + requestParams.contents as Content[], + `${errorContext}-invalid-content`, + ); + } else { + await reportError( + error, + `Error generating content via API.`, + requestParams.contents as Content[], + `${errorContext}-api`, + ); + } + + throw new Error(`Failed to generate content: ${getErrorMessage(error)}`); + } + } }