diff --git a/src/main/kotlin/ee/carlrobert/codegpt/CodeGPTLookupListener.kt b/src/main/kotlin/ee/carlrobert/codegpt/CodeGPTLookupListener.kt index 292cb2f2..2fc44448 100644 --- 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().state.codeAssistantEnabled - || service().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 index 1503b2d3..c81900ba 100644 --- 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().state.codeAssistantEnabled - && service().countTokens(editor.document.text) <= 4098) { + && service().countTokens(editor.document.text) <= 4096) { ApplicationManager.getApplication().executeOnPooledThread { service().displayAutocompletePrediction( editor, diff --git a/src/main/kotlin/ee/carlrobert/codegpt/predictions/TriggerCustomPredictionAction.kt b/src/main/kotlin/ee/carlrobert/codegpt/predictions/TriggerCustomPredictionAction.kt index c143ba04..b7a9aecc 100644 --- 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().countTokens(editor.document.text) > 4098) { - OverlayUtil.showNotification("The file exceeds the token limit of 4,098.") + val encodingManager = service() + 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 }