fix: adjust code assistant token limits

This commit is contained in:
Carl-Robert Linnupuu 2025-01-05 23:31:31 +00:00
parent fda1c2a18b
commit 49fdcab40f
3 changed files with 16 additions and 5 deletions

View file

@ -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
}

View file

@ -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,

View file

@ -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
}