Merge pull request #23 from beekeeper-studio/rc-11

Framework migration, Add Gemini & OpenAI, token limits

Author: azmy60

instructions.txt

@@ -1,3 +1,5 @@
+The current date is {current_date}.
+
 You are an interactive SQL assistant embedded in Beekeeper Studio. A SQL GUI and database manager used by software engineers, DBAs, and other technically skilled people.
 Your job is to help users write, debug, and understand SQL queries based on their natural language questions, schema context, and query history.
 
@@ -26,6 +28,11 @@ Your job is to help users write, debug, and understand SQL queries based on thei
 * Don't comment code unless asked.
 * Running `SELECT` queries can be expensive, add LIMITS of 10, 100, or 1000 as needed unless explicitly asked for the full query.
 
+## Synthetic Data
+
+Sometimes you may see tools resulting in "User rejected tool call". These are actually messages added by the system to inform you that the user has cancelled what you were doing.
+
+## Database Context
 
 The information of the database you are working on is as follows:
 

manifest.json

@@ -1,7 +1,7 @@
 {
   "id": "bks-ai-shell",
   "name": "AI Shell",
-  "version": "1.0.29",
+  "version": "1.1.0",
   "minAppVersion": "5",
   "description": "Ask AI to analyze your database and generate SQL queries.",
   "author": {

package.json

@@ -13,11 +13,14 @@
     "plugin:build": "vite build && cp manifest.json dist/"
   },
   "dependencies": {
-    "@anthropic-ai/sdk": "^0.39.0",
+    "@ai-sdk/anthropic": "^1.2.12",
+    "@ai-sdk/google": "^1.2.22",
+    "@ai-sdk/openai": "^1.3.22",
+    "@ai-sdk/vue": "^1.2.12",
     "@beekeeperstudio/plugin": "^1.3.0-beta",
-    "@langchain/anthropic": "^0.3.21",
-    "@langchain/core": "^0.3.48",
+    "@langchain/core": "^0.3.61",
     "@material-symbols/font-400": "^0.31.2",
+    "ai": "^4.3.16",
     "highlight.js": "^11.11.1",
     "langchain": "^0.3.19",
     "lodash": "^4.17.21",
@@ -27,7 +30,7 @@
     "pluralize": "^8.0.0",
     "typeface-roboto": "^0.0.75",
     "vue": "^3.4.21",
-    "zod": "^3.22.4"
+    "zod": "^3.25.32"
   },
   "devDependencies": {
     "@types/lodash": "^4.17.16",

src/App.vue

@@ -1,28 +1,38 @@
 <template>
   <div class="shell-app">
-    <!-- API Key Form -->
-    <ApiKeyForm
-      v-if="page === 'api-key-form'"
-      :initial-provider-id="providerId"
-      :initial-api-key="apiKey"
-      @submit="handleApiKeySubmit"
-      :disabled="disabledApiKeyForm"
-      :error="error"
-    />
+    <div v-if="!appReady && showLoading" class="not-ready">
+      <h1>AI Shell</h1>
+      <div class="progress-bar"></div>
+    </div>
+    <template v-else>
+      <!-- API Key Form -->
+      <ApiKeyForm
+        v-if="!apiKeyExists || page === 'api-key-form'"
+        @submit="page = 'chat-interface'"
+        @cancel="page = 'chat-interface'"
+        :cancelable="apiKeyExists"
+      />
 
-    <!-- Chat Interface -->
-    <ChatInterface
-      v-else-if="page === 'chat-interface'"
-      @navigate-to-api-form="page = 'api-key-form'"
-    />
+      <ChatInterface
+        v-else
+        :initialMessages="messages"
+        :openaiApiKey="openaiApiKey"
+        :anthropicApiKey="anthropicApiKey"
+        :googleApiKey="googleApiKey"
+        @manage-models="page = 'api-key-form'"
+      />
+    </template>
   </div>
 </template>
 
-<script>
+<script lang="ts">
 import ApiKeyForm from "./components/ApiKeyForm.vue";
 import ChatInterface from "./components/ChatInterface.vue";
-import { useProviderStore } from "./store";
-import { mapState, mapActions } from "pinia";
+import { useChatStore } from "@/stores/chat";
+import { useConfigurationStore } from "@/stores/configuration";
+import { useInternalDataStore } from "@/stores/internalData";
+import { useTabState } from "@/stores/tabState";
+import { mapState, mapActions, mapGetters } from "pinia";
 
 export default {
   components: {
@@ -34,52 +44,58 @@ export default {
     return {
       page: "", // Will be set in mounted based on API key availability
       disabledApiKeyForm: false,
-      error: "",
+      error: "" as unknown,
+      appReady: false,
+      showLoading: false,
     };
   },
 
   async mounted() {
-    await this.initializeChat();
-    // Check if API key exists and auto-navigate to appropriate page
-    if (this.apiKey && this.providerId) {
-      try {
-        await this.initializeProvider();
-        this.page = "chat-interface";
-      } catch (e) {
-        // If initialization fails, go to API key form
+    // Show loading bar after 500ms if not ready
+    const loadingTimer = setTimeout(() => {
+      this.showLoading = true;
+    }, 1000);
+
+    try {
+      await this.initialize();
+
+      const configuration = useConfigurationStore();
+      const apiKey =
+        configuration[`providers.${this.lastUsedProviderId}.apiKey`] ?? "";
+
+      // Check if API key exists and auto-navigate to appropriate page
+      if (apiKey && this.lastUsedProviderId) {
+        try {
+          this.page = "chat-interface";
+        } catch (e) {
+          // If initialization fails, go to API key form
+          this.page = "api-key-form";
+          this.error = e;
+        }
+      } else {
         this.page = "api-key-form";
-        this.error = e;
       }
-    } else {
-      this.page = "api-key-form";
+    } finally {
+      clearTimeout(loadingTimer);
+      this.appReady = true;
     }
   },
 
   computed: {
-    ...mapState(useProviderStore, ["providerId", "apiKey"]),
+    ...mapState(useInternalDataStore, ["lastUsedProviderId"]),
+    ...mapState(useTabState, ["messages"]),
+    ...mapState(useConfigurationStore, {
+      openaiApiKey: "providers.openai.apiKey",
+      anthropicApiKey: "providers.anthropic.apiKey",
+      googleApiKey: "providers.google.apiKey",
+    }),
+    ...mapGetters(useConfigurationStore, ["apiKeyExists"]),
   },
 
   methods: {
-    ...mapActions(useProviderStore, [
-      "setApiKey",
-      "setProviderId",
-      "initializeChat",
-      "initializeProvider",
-    ]),
-    async handleApiKeySubmit(data) {
-      this.error = "";
-      this.disabledApiKeyForm = true;
-      await this.$nextTick()
-      this.setApiKey(data.key);
-      this.setProviderId(data.provider);
-      try {
-        await this.initializeProvider();
-        this.page = "chat-interface";
-      } catch (e) {
-        this.error = e;
-      }
-      this.disabledApiKeyForm = false;
-    },
+    ...mapActions(useConfigurationStore, ["configure"]),
+    ...mapActions(useInternalDataStore, ["setInternal"]),
+    ...mapActions(useChatStore, ["initialize"]),
   },
 };
 </script>

src/assets/styles/components/_all.scss

@@ -2,3 +2,4 @@
 @use "spinner";
 @use "dropdown";
 @use "run-query-message";
+@use "progress-bar";

src/assets/styles/components/_progress-bar.scss

@@ -0,0 +1,23 @@
+.progress-bar {
+  height: 4px;
+  width: 100%;
+  background: var(--border-color);
+  position: relative;
+  overflow: hidden;
+}
+
+.progress-bar::before {
+  content: '';
+  position: absolute;
+  top: 0;
+  left: -100%;
+  height: 100%;
+  width: 100%;
+  background: linear-gradient(90deg, transparent, var(--theme-base), transparent);
+  animation: progress 1.5s infinite;
+}
+
+@keyframes progress {
+  0% { left: -100%; }
+  100% { left: 100%; }
+}

src/assets/styles/main.scss

@@ -19,3 +19,11 @@
   overflow: auto;
 }
 
+.not-ready {
+  width: 100%;
+  max-width: 20rem;
+  display: flex;
+  flex-direction: column;
+  align-items: center;
+  justify-content: center;
+}

src/assets/styles/pages/_api-key-form.scss

@@ -17,7 +17,16 @@
   display: flex;
   flex-direction: column;
   gap: 1rem;
+  margin-inline: auto;
   margin-bottom: 1.5rem;
+  max-width: 20rem;
+
+  .input-group {
+    display: flex;
+    gap: 0.5rem;
+    justify-content: space-between;
+    align-items: center;
+  }
 
   input {
     padding: 0.65rem 1rem;
@@ -32,10 +41,13 @@
     }
   }
 
-  button[type="submit"] {
-    padding-block: 1.15rem;
-    font-size: 0.85rem;
+  .actions {
+    display: flex;
+    justify-content: center;
+    gap: 0.5rem;
+  }
 
+  button {
     &:disabled {
       opacity: 0.6;
     }
@@ -45,10 +57,7 @@
     font-size: 0.85rem;
     color: var(--text-dark);
     opacity: 0.7;
-
-    p {
-      margin-bottom: 0.5rem;
-    }
+    margin-bottom: 0.5rem;
   }
 
   .error-message {