Fix the bug to check gemini model in production env

chrome-extension/src/background/agent/agents/base.ts

@@ -51,8 +51,9 @@ export abstract class BaseAgent<T extends z.ZodType, M = unknown> {
     this.chatLLM = options.chatLLM;
     this.prompt = options.prompt;
     this.context = options.context;
+    // TODO: fix this, the name is not correct in production environment
     this.chatModelLibrary = this.chatLLM.constructor.name;
-    this.modelName = this.setModelNames();
+    this.modelName = this.getModelName();
     this.withStructuredOutput = this.setWithStructuredOutput();
     // extra options
     this.id = extraOptions?.id || 'agent';
@@ -62,7 +63,10 @@ export abstract class BaseAgent<T extends z.ZodType, M = unknown> {
   }
 
   // Set the model name
-  private setModelNames(): string {
+  private getModelName(): string {
+    if ('modelName' in this.chatLLM) {
+      return this.chatLLM.modelName as string;
+    }
     if ('model_name' in this.chatLLM) {
       return this.chatLLM.model_name as string;
     }

chrome-extension/src/background/agent/agents/navigator.ts

@@ -49,6 +49,7 @@ export interface NavigatorResult {
 
 export class NavigatorAgent extends BaseAgent<z.ZodType, NavigatorResult> {
   private actionRegistry: NavigatorActionRegistry;
+  private jsonSchema: Record<string, unknown>;
 
   constructor(
     actionRegistry: NavigatorActionRegistry,
@@ -58,19 +59,14 @@ export class NavigatorAgent extends BaseAgent<z.ZodType, NavigatorResult> {
     super(actionRegistry.setupModelOutputSchema(), options, { ...extraOptions, id: 'navigator' });
 
     this.actionRegistry = actionRegistry;
+
+    this.jsonSchema = this.modelName.startsWith('gemini') ? geminiNavigatorOutputSchema : jsonNavigatorOutputSchema;
   }
 
   async invoke(inputMessages: BaseMessage[]): Promise<this['ModelOutput']> {
     // Use structured output
     if (this.withStructuredOutput) {
-      // For Google Generative AI, we need to use the modelOutputSchema directly
-      // but make sure it doesn't have any 'default' properties that cause issues
-
-      const schema =
-        this.chatModelLibrary === 'ChatGoogleGenerativeAI' ? geminiNavigatorOutputSchema : jsonNavigatorOutputSchema;
-
-      // TODO: don't know why zod can not generate the same schema. Use the json schema exported from browser-use as a workaround for now, need to fix it
-      const structuredLlm = this.chatLLM.withStructuredOutput(schema, {
+      const structuredLlm = this.chatLLM.withStructuredOutput(this.jsonSchema, {
         includeRaw: true,
       });
 
@@ -271,7 +267,7 @@ export class NavigatorAgent extends BaseAgent<z.ZodType, NavigatorResult> {
         await new Promise(resolve => setTimeout(resolve, 1000));
       } catch (error) {
         const errorMessage = error instanceof Error ? error.message : String(error);
-        logger.error('doAction error', errorMessage);
+        logger.error('doAction error', actionName, actionArgs, errorMessage);
         // unexpected error, emit event
         this.context.emitEvent(Actors.NAVIGATOR, ExecutionState.ACT_FAIL, errorMessage);
         errCount++;