Merge pull request #64 from beekeeper-studio/feat/context-window-management

Context Window Management

Author: azmy60

instructions/compact.txt

@@ -0,0 +1,4 @@
+Summarize the above conversation for a new session. The new session will
+continue the conversation and will not be able to access it. Include any useful
+context needed, such as what has been done, what is currently being done, and
+next steps if applicable.

src/assets/styles/_primevue.scss

@@ -136,13 +136,12 @@
   // Popover
   --p-popover-gutter: 0.5rem;
   --p-popover-arrow-offset: 1rem;
-  --p-popover-background: color-mix(in srgb,
-      var(--theme-base) 6%,
-      var(--query-editor-bg));
-  --p-popover-color: var(--text);
+  --p-popover-background: var(--p-menu-background);
+  --p-popover-color: var(--p-menu-item-color);
   --p-popover-border-color: transparent;
-  --p-popover-border-radius: 8px;
+  --p-popover-border-radius: var(--p-menu-border-radius);
   --p-popover-box-shadow: 0 0 0.5rem var(--theme-bg);
+  --p-popover-shadow: 0 0 0.5rem var(--query-editor-bg);
   // -----------------------------------
 
   // Progress

src/assets/styles/components/_message.scss

@@ -198,11 +198,6 @@
     font-weight: normal;
   }
 
-  .message-content.literally-empty {
-    font-style: italic;
-    --theme-text-message-system: var(--text-muted);
-  }
-
   [data-tool-state="input-streaming"] {
     &>.markdown code .hljs-comment::after {
       content: "";

src/assets/styles/pages/configuration/_main.scss

@@ -48,7 +48,7 @@
     width: 100%;
     box-shadow: none;
     text-align: left;
-    height: 2rem;
+    height: 1.75rem;
     justify-content: flex-start;
     background-color: transparent;
 

src/components/ChatInterface.vue

@@ -1,8 +1,15 @@
 <template>
-  <div class="chat-container" :class="{ 'empty-chat': messages.length === 0 }" :data-status="status">
+  <div
+    class="chat-container"
+    :class="{ 'empty-chat': messages.length === 0 }"
+    :data-status="status"
+  >
     <div class="scroll-container" ref="scrollContainerRef">
       <div class="header">
-        <button class="btn btn-flat-2 settings-btn" @click="$emit('open-configuration')">
+        <button
+          class="btn btn-flat-2 settings-btn"
+          @click="$emit('open-configuration')"
+        >
           <span class="material-symbols-outlined">settings</span>
           <span class="title-popup">Settings</span>
         </button>
@@ -12,30 +19,51 @@
         <p>
           The AI Shell can see your table schemas, and (with your permission)
           run {{ sqlOrCode }} to answer your questions.
-          <ExternalLink href="https://docs.beekeeperstudio.io/user_guide/sql-ai-shell/">Learn more</ExternalLink>.
+          <ExternalLink
+            href="https://docs.beekeeperstudio.io/user_guide/sql-ai-shell/"
+            >Learn more</ExternalLink
+          >.
         </p>
       </div>
       <div class="chat-messages">
         <template v-for="(message, index) in messages" :key="message.id">
           <message
             v-if="
-              !(message.role === 'assistant' && message.parts.find((p) => p.type === 'data-userEditedToolCall'))
-                && !(message.role === 'user' && message.parts.find((p) => p.type === 'data-editedQuery'))
-              "
+              !(
+                message.role === 'assistant' &&
+                message.parts.find((p) => p.type === 'data-userEditedToolCall')
+              ) &&
+              !(
+                message.role === 'user' &&
+                message.parts.find((p) => p.type === 'data-editedQuery')
+              )
+            "
             :message="message"
-            :status="index === messages.length - 1 ? (status === 'ready' || status === 'error' ? 'ready' : 'processing') : 'ready'"
-            @accept-permission="acceptPermission" @reject-permission="handleRejectPermission" />
+            :status="
+              index === messages.length - 1
+                ? status === 'ready' || status === 'error'
+                  ? 'ready'
+                  : 'processing'
+                : 'ready'
+            "
+            @accept-permission="acceptPermission"
+            @reject-permission="handleRejectPermission"
+          />
         </template>
         <div class="message error" v-if="isUnexpectedError">
           <div class="message-content">
             Something went wrong.
             <div v-if="isOllamaToolError" class="error-hint">
-              💡 <strong>Hint:</strong> This might be because your Ollama model doesn't support tools. Try using a
-              different model, or switch to a different provider.
+              💡 <strong>Hint:</strong> This might be because your Ollama model
+              doesn't support tools. Try using a different model.
             </div>
             <pre v-if="!isErrorTruncated || showFullError" v-text="error" />
             <pre v-else v-text="truncatedError" />
-            <button v-if="isErrorTruncated" @click="showFullError = !showFullError" class="btn show-more-btn">
+            <button
+              v-if="isErrorTruncated"
+              @click="showFullError = !showFullError"
+              class="btn show-more-btn"
+            >
               {{ showFullError ? "Show less" : "Show more" }}
             </button>
             <button class="btn" @click="() => reload()">
@@ -47,18 +75,46 @@
         <div class="message error" v-if="noModelError">
           <div class="message-content">No model selected</div>
         </div>
-        <div class="spinner-container" :style="{ visibility: showSpinner ? 'visible' : 'hidden' }">
+        <div
+          class="spinner-container"
+          :style="{ visibility: showSpinner ? 'visible' : 'hidden' }"
+        >
           <span class="spinner" />
+          <span class="label" v-if="compacting">Compacting</span>
         </div>
       </div>
-      <button v-if="!isAtBottom" @click="scrollToBottom({ smooth: true })" class="btn scroll-down-btn"
-        title="Scroll to bottom">
+      <button
+        v-if="!isAtBottom"
+        @click="scrollToBottom({ smooth: true })"
+        class="btn scroll-down-btn"
+        title="Scroll to bottom"
+      >
         <span class="material-symbols-outlined">keyboard_arrow_down</span>
       </button>
     </div>
     <div class="chat-input-container-container">
-      <PromptInput ref="promptInput" storage-key="inputHistory" :processing="processing" :selected-model="model"
-        @select-model="selectModel" @manage-models="$emit('manage-models')" @submit="submit" @stop="stop" />
+      <PromptInput
+        ref="promptInput"
+        storage-key="inputHistory"
+        :processing="processing"
+        :selected-model="model"
+        @select-model="selectModel"
+        @manage-models="$emit('manage-models')"
+        @submit="submit"
+        @stop="stop"
+      />
+      <div
+        v-if="enableAutoCompact && contextLeftUntilAutoCompact <= 15"
+        class="auto-compact-notice"
+      >
+        <template v-if="contextOverflow">
+          Auto-compacting on next message
+        </template>
+        <template v-else>
+          Context left until auto-compact:
+          {{ contextLeftUntilAutoCompact.toFixed(1) }}%
+        </template>
+      </div>
     </div>
   </div>
 </template>
@@ -79,6 +135,7 @@ import PromptInput from "@/components/common/PromptInput.vue";
 import { getConnectionInfo } from "@beekeeperstudio/plugin";
 import ExternalLink from "@/components/common/ExternalLink.vue";
 import { log } from "@beekeeperstudio/plugin";
+import { useConfigurationStore } from "@/stores/configuration";
 
 export default {
   name: "ChatInterface",
@@ -101,20 +158,7 @@ export default {
   },
 
   setup(props) {
-    const ai = useAI({ initialMessages: props.initialMessages });
-
-    return {
-      send: ai.send,
-      abort: ai.abort,
-      messages: ai.messages,
-      error: ai.error,
-      status: ai.status,
-      acceptPermission: ai.acceptPermission,
-      rejectPermission: ai.rejectPermission,
-      rejectAllPendingApprovals: ai.rejectAllPendingApprovals,
-      hasPendingApprovals: ai.hasPendingApprovals,
-      retry: ai.retry,
-    };
+    return useAI({ initialMessages: props.initialMessages });
   },
 
   data() {
@@ -127,13 +171,22 @@ export default {
   },
 
   computed: {
-    ...mapGetters(useChatStore, ["systemPrompt"]),
+    ...mapGetters(useConfigurationStore, ["enableAutoCompact"]),
+    ...mapGetters(useChatStore, [
+      "systemPrompt",
+      "contextOverflow",
+      "contextLeftUntilAutoCompact",
+    ]),
     ...mapWritableState(useChatStore, ["model"]),
     processing() {
       if (this.hasPendingApprovals) return false;
       return this.status !== "ready" && this.status !== "error";
     },
     showSpinner() {
+      if (this.compacting) {
+        return true;
+      }
+
       return (
         !this.hasPendingApprovals &&
         (this.status === "submitted" || this.status === "streaming")
@@ -148,12 +201,12 @@ export default {
         return true;
       }
 
-      if (this.error.message.includes('User rejected tool call')) {
+      if (this.error.message.includes("User rejected tool call")) {
         return false;
       }
 
       // User aborted request before AI got a chance to respond
-      if (this.error.message.includes('aborted without reason')) {
+      if (this.error.message.includes("aborted without reason")) {
         return false;
       }
 
@@ -168,8 +221,8 @@ export default {
     isOllamaToolError() {
       if (!this.error || !this.model) return false;
       const errorStr = this.error.toString().toLowerCase();
-      const isOllama = this.model.provider === 'ollama';
-      const hasToolError = errorStr.includes('bad request');
+      const isOllama = this.model.provider === "ollama";
+      const hasToolError = errorStr.includes("bad request");
       return isOllama && hasToolError;
     },
 
@@ -184,7 +237,7 @@ export default {
             }
           },
         },
-      ]
+      ];
     },
   },
 
@@ -207,9 +260,11 @@ export default {
 
   async mounted() {
     getConnectionInfo().then((connection) => {
-      if (connection.databaseType === "mongodb"
-        || connection.connectionType === "surrealdb"
-        || connection.connectionType === "redis") {
+      if (
+        connection.databaseType === "mongodb" ||
+        connection.connectionType === "surrealdb" ||
+        connection.connectionType === "redis"
+      ) {
         this.sqlOrCode = "Code";
       }
     });
@@ -219,8 +274,8 @@ export default {
       // Calculate if we're near bottom (within 50px of bottom)
       const isNearBottom =
         scrollContainer.scrollHeight -
-        scrollContainer.scrollTop -
-        scrollContainer.clientHeight <
+          scrollContainer.scrollTop -
+          scrollContainer.clientHeight <
         50;
 
       this.isAtBottom = isNearBottom;
@@ -233,7 +288,7 @@ export default {
   methods: {
     ...mapActions(useInternalDataStore, ["setInternal"]),
 
-    submit(input: string) {
+    async submit(input: string) {
       if (!this.model) {
         // FIXME we should catch this and show it on screen
         this.noModelError = true;
@@ -246,7 +301,11 @@ export default {
         this.rejectAllPendingApprovals();
       }
 
-      this.send(input);
+      if (this.contextOverflow) {
+        await this.compact(input);
+      } else {
+        await this.send(input);
+      }
     },
 
     async reload() {
@@ -268,7 +327,7 @@ export default {
       if (options?.smooth) {
         this.$refs.scrollContainerRef.scrollTo({
           top: this.$refs.scrollContainerRef.scrollHeight,
-          behavior: 'smooth'
+          behavior: "smooth",
         });
       } else {
         this.$refs.scrollContainerRef.scrollTop =
@@ -299,3 +358,15 @@ export default {
   },
 };
 </script>
+
+<style scoped>
+.spinner-container .label {
+  padding-left: 1ch;
+}
+.auto-compact-notice {
+  width: 100%;
+  margin-top: 0.5rem;
+  text-align: right;
+  color: var(--text);
+}
+</style>

src/components/OnboardingScreen.vue

@@ -62,7 +62,7 @@ export default {
   methods: {
     handleChangeProvider(provider: AvailableProviders) {
       this.providerChanged = true;
-      if (provider === "ollama") {
+      if (provider === "ollama" || provider === "mock") {
         this.changed = true;
       }
     },

src/components/common/BaseInput.vue

@@ -92,7 +92,9 @@ export default {
 
   emits: ["update:modelValue", "input", "change", "click"],
 
-  expose: ["focus"],
+  // FIXME: Strip this out for now cause vue-tsc isn't happy
+  // See https://github.com/vuejs/language-tools/issues/5069
+  // expose: ["focus"],
 
   methods: {
     emitInput(event: Event) {

src/components/common/PromptInput.vue

@@ -47,7 +47,9 @@ export default defineComponent({
 
   emits: ["submit", "stop", "manage-models", "select-model"],
 
-  expose: ["focus"],
+  // FIXME: Strip this out for now cause vue-tsc isn't happy
+  // See https://github.com/vuejs/language-tools/issues/5069
+  // expose: ["focus"],
 
   props: {
     processing: Boolean,