Bug fix: Update BaseLlm constructor to use a parameter object.

PiperOrigin-RevId: 828077838

Author: kalenkevich

core/src/common.ts

@@ -30,6 +30,7 @@ export {Gemini} from './models/google_llm.js';
 export type {GeminiParams} from './models/google_llm.js';
 export type {LlmRequest} from './models/llm_request.js';
 export type {LlmResponse} from './models/llm_response.js';
+export {LLMRegistry} from './models/registry.js';
 export {BasePlugin} from './plugins/base_plugin.js';
 export {LoggingPlugin} from './plugins/logging_plugin.js';
 export {PluginManager} from './plugins/plugin_manager.js';

core/src/models/base_llm.ts

@@ -12,13 +12,17 @@ import {LlmResponse} from './llm_response.js';
  * The BaseLLM class.
  */
 export abstract class BaseLlm {
+  readonly model: string;
+
   /**
    * Creates an instance of BaseLLM.
-   *
-   * @param model The name of the LLM, e.g. gemini-1.5-flash or
+   * @param params The parameters for creating a BaseLlm instance.
+   * @param params.model The name of the LLM, e.g. gemini-1.5-flash or
    *     gemini-1.5-flash-001.
    */
-  constructor(readonly model: string) {}
+  constructor({model}: {model: string}) {
+    this.model = model;
+  }
 
   /**
    * List of supported models in regex for LlmRegistry.

core/src/models/google_llm.ts

@@ -66,14 +66,18 @@ export class Gemini extends BaseLlm {
    * @param params The parameters for creating a Gemini instance.
    */
   constructor({
-    model = 'gemini-2.5-flash',
+    model,
     apiKey,
     vertexai,
     project,
     location,
-    headers
-  }: GeminiParams = {}) {
-    super(model);
+    headers,
+  }: GeminiParams) {
+    if (!model) {
+      model = 'gemini-2.5-flash';
+    }
+
+    super({model});
 
     this.project = project;
     this.location = location;

core/src/models/registry.ts

@@ -13,7 +13,7 @@ import {Gemini} from './google_llm.js';
  * type[BaseLlm] equivalent in TypeScript, represents a class that can be new-ed
  * to create a BaseLlm instance.
  */
-export type BaseLlmType = (new (model: string) => BaseLlm)&{
+export type BaseLlmType = (new (params: {model: string}) => BaseLlm)&{
   readonly supportedModels: Array<string|RegExp>;
 };
 
@@ -68,7 +68,7 @@ export class LLMRegistry {
    * @returns The LLM instance.
    */
   static newLlm(model: string): BaseLlm {
-    return new (LLMRegistry.resolve(model))(model);
+    return new (LLMRegistry.resolve(model))({model});
   }
 
   private static _register(modelNameRegex: string|RegExp, llmCls: BaseLlmType) {
@@ -85,9 +85,10 @@ export class LLMRegistry {
    * Registers a new LLM class.
    * @param llmCls The class that implements the model.
    */
-  static register<T extends BaseLlm>(llmCls: (new(model: string) => T)&{
-    readonly supportedModels: Array<string|RegExp>;
-  }) {
+  static register<T extends BaseLlm>(
+      llmCls: (new(params: {model: string}) => T)&{
+        readonly supportedModels: Array<string|RegExp>;
+      }) {
     for (const regex of llmCls.supportedModels) {
       LLMRegistry._register(regex, llmCls);
     }

core/test/agents/llm_agent_test.ts

@@ -33,7 +33,7 @@ class MockLlm extends BaseLlm {
   error: Error|null;
 
   constructor(response: LlmResponse|null, error: Error|null = null) {
-    super('mock-llm');
+    super({model: 'mock-llm'});
     this.response = response;
     this.error = error;
   }

core/test/models/registry_test.ts

@@ -0,0 +1,70 @@
+import {BaseLlm, BaseLlmConnection, LlmAgent, LLMRegistry, LlmRequest, LlmResponse} from '@google/adk';
+import {Blob, Content, createModelContent, GenerateContentResponse} from '@google/genai';
+
+class TestLlmConnection implements BaseLlmConnection {
+  async sendHistory(history: Content[]): Promise<void> {
+    return Promise.resolve();
+  }
+
+  async sendContent(content: Content): Promise<void> {}
+
+  async sendRealtime(blob: Blob): Promise<void> {}
+
+  async * receive(): AsyncGenerator<LlmResponse, void, void> {}
+
+  async close(): Promise<void> {}
+}
+
+class TestLlmModel extends BaseLlm {
+  constructor({model}: {model: string}) {
+    super({model});
+  }
+
+  static override readonly supportedModels = ['test-llm-model'];
+
+  async *
+      generateContentAsync(llmRequest: LlmRequest, stream?: boolean):
+          AsyncGenerator<LlmResponse, void> {
+    const generateContentResponse = new GenerateContentResponse();
+
+    generateContentResponse.candidates =
+        [{content: createModelContent('test-llm-model-response')}];
+    const candidate = generateContentResponse.candidates[0];
+
+    yield {
+      content: candidate.content,
+      groundingMetadata: candidate.groundingMetadata,
+      usageMetadata: generateContentResponse.usageMetadata,
+      finishReason: candidate.finishReason,
+    };
+  }
+
+  async connect(llmRequest: LlmRequest): Promise<BaseLlmConnection> {
+    return new TestLlmConnection();
+  }
+}
+
+describe('LLMRegistry', () => {
+  beforeAll(() => {
+    LLMRegistry.register(TestLlmModel);
+  });
+
+  it('resolves model to LLM class', () => {
+    expect(LLMRegistry.newLlm('test-llm-model')).toBeInstanceOf(TestLlmModel);
+  });
+
+  it('resolves the provided as a string model correctly in LlmAgent', () => {
+    const agent = new LlmAgent({name: 'test_agent', model: 'test-llm-model'});
+
+    expect(agent.canonicalModel).toBeInstanceOf(TestLlmModel);
+  });
+
+  it('resolves the provided as class model correctly in LlmAgent', () => {
+    const agent = new LlmAgent({
+      name: 'test_agent',
+      model: new TestLlmModel({model: 'test-llm-model'})
+    });
+
+    expect(agent.canonicalModel).toBeInstanceOf(TestLlmModel);
+  })
+});