From 1cd36002de0a3d86c37931dbfe14b47d5f217fd5 Mon Sep 17 00:00:00 2001
From: mondaychen <monday.chen@gmail.com>
Date: Tue, 17 Sep 2024 14:32:53 -0400
Subject: [PATCH] refactor: use agent mode in run task logic

---
 src/common/settings/ModelDropdown.tsx | 26 ++++++++------------------
 src/state/currentTask.ts              | 24 +++++++++++++++++++-----
 2 files changed, 27 insertions(+), 23 deletions(-)

diff --git a/src/common/settings/ModelDropdown.tsx b/src/common/settings/ModelDropdown.tsx
index 9e8e82a..5a0587d 100644
--- a/src/common/settings/ModelDropdown.tsx
+++ b/src/common/settings/ModelDropdown.tsx
@@ -3,10 +3,7 @@ import { useAppState } from "../../state/store";
 import {
   SupportedModels,
   DisplayName,
-  isOpenAIModel,
-  isAnthropicModel,
-  AgentMode,
-  hasVisionSupport,
+  isValidModelSettings,
 } from "../../helpers/aiSdkUtils";
 import { enumValues } from "../../helpers/utils";
 
@@ -22,19 +19,6 @@ const ModelDropdown = () => {
     anthropicKey: state.settings.anthropicKey,
   }));
 
-  const isModelSupported = (model: SupportedModels) => {
-    if (agentMode === AgentMode.VisionEnhanced && !hasVisionSupport(model)) {
-      return false;
-    }
-    if (isOpenAIModel(model)) {
-      return !!openAIKey;
-    }
-    if (isAnthropicModel(model)) {
-      return !!anthropicKey;
-    }
-    return false;
-  };
-
   return (
     <Select
       id="model-select"
@@ -44,7 +28,13 @@ const ModelDropdown = () => {
       }
     >
       {enumValues(SupportedModels).map((model) => (
-        <option key={model} value={model} disabled={!isModelSupported(model)}>
+        <option
+          key={model}
+          value={model}
+          disabled={
+            !isValidModelSettings(model, agentMode, openAIKey, anthropicKey)
+          }
+        >
           {DisplayName[model]}
         </option>
       ))}
diff --git a/src/state/currentTask.ts b/src/state/currentTask.ts
index b7aff95..1325431 100644
--- a/src/state/currentTask.ts
+++ b/src/state/currentTask.ts
@@ -27,7 +27,7 @@ import { MyStateCreator } from "./store";
 import buildAnnotatedScreenshots from "../helpers/buildAnnotatedScreenshots";
 import { voiceControl } from "../helpers/voiceControl";
 import { fetchKnowledge, type Knowledge } from "../helpers/knowledge";
-import { hasVisionSupport } from "../helpers/aiSdkUtils";
+import { isValidModelSettings, AgentMode } from "../helpers/aiSdkUtils";
 
 export type TaskHistoryEntry = {
   prompt: string;
@@ -77,6 +77,22 @@ export const createCurrentTaskSlice: MyStateCreator<CurrentTaskSlice> = (
     runTask: async (onError) => {
       const voiceMode = get().settings.voiceMode;
 
+      if (
+        !isValidModelSettings(
+          get().settings.selectedModel,
+          get().settings.agentMode,
+          get().settings.openAIKey,
+          get().settings.anthropicKey,
+        )
+      ) {
+        onError(
+          "The current model settings are not valid. Please verify your API keys, and note that some models are not compatible with certain agent modes.",
+        );
+        return;
+      }
+
+      const agentMode = get().settings.agentMode;
+
       const wasStopped = () => get().currentTask.status !== "running";
       const setActionStatus = (status: CurrentTaskSlice["actionStatus"]) => {
         set((state) => {
@@ -124,8 +140,6 @@ export const createCurrentTaskSlice: MyStateCreator<CurrentTaskSlice> = (
             );
           }
 
-          const isVisionModel = hasVisionSupport(get().settings.selectedModel);
-
           const performAction = async (
             query: QueryResult,
           ): Promise<boolean> => {
@@ -161,7 +175,7 @@ export const createCurrentTaskSlice: MyStateCreator<CurrentTaskSlice> = (
             ) {
               return false;
             }
-            if (isVisionModel) {
+            if (agentMode === AgentMode.VisionEnhanced) {
               await operateTool(tabId, query.action.operation);
             } else {
               await operateToolWithSimpliedDom(tabId, query.action.operation);
@@ -198,7 +212,7 @@ export const createCurrentTaskSlice: MyStateCreator<CurrentTaskSlice> = (
             .currentTask.history.map((entry) => entry.action)
             .filter(truthyFilter);
 
-          if (isVisionModel) {
+          if (agentMode === AgentMode.VisionEnhanced) {
             setActionStatus("fetching-knoweldge");
             const url = new URL(activeTab.url ?? "");
             const customKnowledgeBase = get().settings.customKnowledgeBase;