fix: adjust code assistant token limits

carlrobertoh · Jan 5, 2025 · 49fdcab · 49fdcab
1 parent fda1c2a
commit 49fdcab
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 5 deletions.
diff --git a/src/main/kotlin/ee/carlrobert/codegpt/CodeGPTLookupListener.kt b/src/main/kotlin/ee/carlrobert/codegpt/CodeGPTLookupListener.kt
@@ -8,6 +8,8 @@ import com.intellij.codeInsight.lookup.impl.LookupImpl
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.application.runReadAction
 import com.intellij.openapi.components.service
+import ee.carlrobert.codegpt.credentials.CredentialsStore.CredentialKey.CODEGPT_API_KEY
+import ee.carlrobert.codegpt.credentials.CredentialsStore.isCredentialSet
 import ee.carlrobert.codegpt.predictions.PredictionService
 import ee.carlrobert.codegpt.settings.GeneralSettings
 import ee.carlrobert.codegpt.settings.service.ServiceType
@@ -32,10 +34,11 @@ class CodeGPTLookupListener : LookupManagerListener {
 
                 override fun itemSelected(event: LookupEvent) {
                     val editor = newLookup.editor
-
+                    val encodingManager = EncodingManager.getInstance()
                     if (GeneralSettings.getSelectedService() != ServiceType.CODEGPT
                         || !service<CodeGPTServiceSettings>().state.codeAssistantEnabled
-                        || service<EncodingManager>().countTokens(editor.document.text) > 4098
+                        || encodingManager.countTokens(editor.document.text) > 4096
+                        || !isCredentialSet(CODEGPT_API_KEY) && encodingManager.countTokens(editor.document.text) > 2048
                     ) {
                         return
                     }

diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionInsertAction.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionInsertAction.kt
@@ -52,7 +52,7 @@ class CodeCompletionInsertAction :
 
                 if (GeneralSettings.getSelectedService() == ServiceType.CODEGPT
                     && service<CodeGPTServiceSettings>().state.codeAssistantEnabled
-                    && service<EncodingManager>().countTokens(editor.document.text) <= 4098) {
+                    && service<EncodingManager>().countTokens(editor.document.text) <= 4096) {
                     ApplicationManager.getApplication().executeOnPooledThread {
                         service<PredictionService>().displayAutocompletePrediction(
                             editor,

diff --git a/src/main/kotlin/ee/carlrobert/codegpt/predictions/TriggerCustomPredictionAction.kt b/src/main/kotlin/ee/carlrobert/codegpt/predictions/TriggerCustomPredictionAction.kt
@@ -13,6 +13,9 @@ import com.intellij.openapi.editor.actionSystem.EditorAction
 import com.intellij.openapi.editor.actionSystem.EditorWriteActionHandler
 import ee.carlrobert.codegpt.CodeGPTKeys
 import ee.carlrobert.codegpt.EncodingManager
+import ee.carlrobert.codegpt.credentials.CredentialsStore
+import ee.carlrobert.codegpt.credentials.CredentialsStore.CredentialKey.CODEGPT_API_KEY
+import ee.carlrobert.codegpt.credentials.CredentialsStore.isCredentialSet
 import ee.carlrobert.codegpt.settings.GeneralSettings
 import ee.carlrobert.codegpt.settings.service.ServiceType
 import ee.carlrobert.codegpt.settings.service.codegpt.CodeGPTServiceSettings
@@ -47,8 +50,13 @@ class TriggerCustomPredictionAction : EditorAction(Handler()), HintManagerImpl.A
                 return
             }
 
-            if (service<EncodingManager>().countTokens(editor.document.text) > 4098) {
-                OverlayUtil.showNotification("The file exceeds the token limit of 4,098.")
+            val encodingManager = service<EncodingManager>()
+            if (!isCredentialSet(CODEGPT_API_KEY) && encodingManager.countTokens(editor.document.text) > 2048) {
+                OverlayUtil.showNotification("The file exceeds the token limit of 2,048. Please upgrade your plan to access higher limits.")
+                return
+            }
+            if (encodingManager.countTokens(editor.document.text) > 4096) {
+                OverlayUtil.showNotification("The file exceeds the token limit of 4,096.")
                 return
             }