feat(core): Add BaseLlmClient.generateContent. (#13591)

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',
+      );
+    });
+  });
 });

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,35 +88,17 @@ export class BaseLlmClient {
     options: GenerateJsonOptions,
   ): Promise<Record<string, unknown>> {
     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',
-    };
-
-    try {
-      const apiCall = () =>
-        this.contentGenerator.generateContent(
-          {
-            model,
-            config: requestConfig,
-            contents,
-          },
-          promptId,
-        );
 
     const shouldRetryOnContent = (response: GenerateContentResponse) => {
       const text = getResponseText(response)?.trim();
@@ -99,51 +106,36 @@ export class BaseLlmClient {
         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;
+        return false; // It's valid, don't retry
       } catch (_e) {
-          return true;
+        return true; // It's not valid, retry
       }
     };
 
-      const result = await retryWithBackoff(apiCall, {
+    const result = await this._generateWithRetry(
+      {
+        model,
+        contents,
+        config: {
+          ...generateContentConfig,
+          ...(systemInstruction && { systemInstruction }),
+          responseJsonSchema: schema,
+          responseMimeType: 'application/json',
+          abortSignal,
+        },
+      },
+      promptId,
+      maxAttempts,
       shouldRetryOnContent,
-        maxAttempts: maxAttempts ?? DEFAULT_MAX_ATTEMPTS,
+      'generateJson',
-      });
+    );
 
     // 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)}`,
-      );
-    }
   }
 
   async generateEmbedding(texts: string[]): Promise<number[][]> {
@@ -193,4 +185,87 @@ export class BaseLlmClient {
     }
     return text;
   }
+
+  async generateContent(
+    options: GenerateContentOptions,
+  ): Promise<GenerateContentResponse> {
+    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<GenerateContentResponse> {
+    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)}`);
+    }
+  }
 }