From f30c8ace7a8126fe23f1aa44afe02426557c7ffc Mon Sep 17 00:00:00 2001 From: Roman Gromov <35900229+sapphirepro@users.noreply.github.com> Date: Sun, 15 Mar 2026 01:20:06 +0100 Subject: [PATCH] Added custom fields for context size and output size to Custom OpenAI Added fields in settings to manually set max overall context and max output context for calculation in agent mode. Default values were too low and didn't count specific model specs. Now it's solved --- .../carlrobert/codegpt/agent/AgentService.kt | 29 ++++++++-- .../codegpt/settings/models/ModelProvider.kt | 4 +- .../service/custom/CustomServiceSettings.kt | 4 ++ .../service/custom/form/CustomServiceForm.kt | 39 ++++++++++++++ .../form/model/CustomServiceSettingsData.kt | 4 ++ .../custom/form/model/DataToStateMapper.kt | 2 + .../custom/form/model/StateToDataMapper.kt | 2 + .../resources/messages/codegpt.properties | 2 + .../resources/messages/codegpt_zh.properties | 2 + .../settings/models/ModelSettingsTest.kt | 53 ++++++++++++++++++- 10 files changed, 135 insertions(+), 6 deletions(-) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/AgentService.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/AgentService.kt index e185f2a5..fa898ad7 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/AgentService.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/AgentService.kt @@ -4,6 +4,7 @@ import ai.koog.agents.core.agent.AIAgent import ai.koog.agents.core.agent.AIAgentService import ai.koog.agents.snapshot.feature.AgentCheckpointData import ai.koog.agents.snapshot.providers.file.JVMFilePersistenceStorageProvider +import ai.koog.prompt.llm.LLMCapability import com.intellij.openapi.components.Service import com.intellij.openapi.components.service import com.intellij.openapi.diagnostic.thisLogger @@ -46,11 +47,19 @@ class AgentService(private val project: Project) { private data class SessionRuntime( val service: AIAgentService>, val provider: ServiceType, - val modelId: String, + val modelSignature: ModelRuntimeSignature, val selectedServerIds: Set, val events: AgentEvents ) + private data class ModelRuntimeSignature( + val selectionId: String, + val modelId: String, + val contextLength: Long?, + val maxOutputTokens: Long?, + val usesResponsesApi: Boolean, + ) + private val sessionJobs = ConcurrentHashMap() private val pendingMessages = ConcurrentHashMap>() private val sessionAgents = ConcurrentHashMap>() @@ -256,14 +265,14 @@ class AgentService(private val project: Project) { provider: ServiceType, events: AgentEvents ): SessionRuntime { - val modelId = service().getAgentModel().id + val modelSignature = currentModelSignature() val selectedServerIds = project.service() .get(sessionId) ?.selectedServerIds ?: emptySet() val existing = sessionRuntimes[sessionId] if (existing != null && existing.provider == provider && - existing.modelId == modelId && + existing.modelSignature == modelSignature && existing.selectedServerIds == selectedServerIds && existing.events === events ) { @@ -287,7 +296,7 @@ class AgentService(private val project: Project) { pendingMessages = pendingMessages ), provider = provider, - modelId = modelId, + modelSignature = modelSignature, selectedServerIds = selectedServerIds, events = events ) @@ -295,6 +304,18 @@ class AgentService(private val project: Project) { return created } + private fun currentModelSignature(): ModelRuntimeSignature { + val selection = service().getModelSelectionForFeature(FeatureType.AGENT) + val model = selection.llmModel + return ModelRuntimeSignature( + selectionId = selection.selectionId, + modelId = model.id, + contextLength = model.contextLength, + maxOutputTokens = model.maxOutputTokens, + usesResponsesApi = model.supports(LLMCapability.OpenAIEndpoint.Responses) + ) + } + private fun logCheckpointLoadFailure(sessionId: String, agentId: String, ex: Throwable) { logger.error( "Agent checkpoints: failed to load for sessionId=$sessionId agentId=$agentId error=${ex.message}", diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelProvider.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelProvider.kt index 3ee72c99..5208cb33 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelProvider.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelProvider.kt @@ -747,7 +747,9 @@ private class CustomOpenAIModelProvider : ModelProvider { LLMCapability.Document, LLMCapability.Completion, LLMCapability.MultipleChoices, - ) + additionalOpenAICapability + ) + additionalOpenAICapability, + contextLength = svc.contextWindowSize.coerceAtLeast(1L), + maxOutputTokens = svc.maxOutputTokens.coerceAtLeast(1L) ), displayName = displayName, serviceId = serviceId, diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/CustomServiceSettings.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/CustomServiceSettings.kt index 2f371172..80a1c3b7 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/CustomServiceSettings.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/CustomServiceSettings.kt @@ -17,6 +17,8 @@ import java.util.UUID import ee.carlrobert.codegpt.util.MapConverter private const val DEFAULT_SERVICE_SETTINGS_NAME = "Default" +const val DEFAULT_CUSTOM_OPENAI_CONTEXT_WINDOW_SIZE = 200_000L +const val DEFAULT_CUSTOM_OPENAI_MAX_OUTPUT_TOKENS = 32_768L @Service @State( @@ -153,6 +155,8 @@ class CustomServiceSettingsState : BaseState() { var id by string(UUID.randomUUID().toString()) var name by string(DEFAULT_SERVICE_SETTINGS_NAME) var template by enum(CustomServiceTemplate.OPENAI) + var contextWindowSize by property(DEFAULT_CUSTOM_OPENAI_CONTEXT_WINDOW_SIZE) + var maxOutputTokens by property(DEFAULT_CUSTOM_OPENAI_MAX_OUTPUT_TOKENS) var chatCompletionSettings by property(CustomServiceChatCompletionSettingsState()) var codeCompletionSettings by property(CustomServiceCodeCompletionSettingsState()) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/form/CustomServiceForm.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/form/CustomServiceForm.kt index 9e7c6f43..d6e5348d 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/form/CustomServiceForm.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/form/CustomServiceForm.kt @@ -22,6 +22,8 @@ import com.intellij.util.ui.components.BorderLayoutPanel import ee.carlrobert.codegpt.CodeGPTBundle import ee.carlrobert.codegpt.credentials.CredentialsStore import ee.carlrobert.codegpt.credentials.CredentialsStore.CredentialKey +import ee.carlrobert.codegpt.settings.service.custom.DEFAULT_CUSTOM_OPENAI_CONTEXT_WINDOW_SIZE +import ee.carlrobert.codegpt.settings.service.custom.DEFAULT_CUSTOM_OPENAI_MAX_OUTPUT_TOKENS import ee.carlrobert.codegpt.settings.service.custom.CustomServiceSettingsState import ee.carlrobert.codegpt.settings.service.custom.CustomServicesSettings import ee.carlrobert.codegpt.settings.service.custom.form.model.CustomServiceSettingsData @@ -40,6 +42,7 @@ import kotlinx.coroutines.flow.launchIn import kotlinx.coroutines.flow.onEach import kotlinx.coroutines.flow.update import kotlinx.coroutines.launch +import java.text.ParseException import java.awt.BorderLayout import java.awt.Dimension import java.awt.FlowLayout @@ -125,6 +128,10 @@ class CustomServiceForm( private val nameField = JBTextField().apply { columns = 30 } + private val contextWindowSizeSpinner = + createTokenSpinner(DEFAULT_CUSTOM_OPENAI_CONTEXT_WINDOW_SIZE) + private val maxOutputTokensSpinner = + createTokenSpinner(DEFAULT_CUSTOM_OPENAI_MAX_OUTPUT_TOKENS) private val templateHelpText = JBLabel(General.ContextHelp) private val templateComboBox = ComboBox(EnumComboBoxModel(CustomServiceTemplate::class.java)) private val chatCompletionsForm: CustomServiceChatCompletionForm @@ -217,6 +224,8 @@ class CustomServiceForm( apiKeyField.text = selectedItem.apiKey ?: "" nameField.text = selectedItem.name + contextWindowSizeSpinner.value = selectedItem.contextWindowSize + maxOutputTokensSpinner.value = selectedItem.maxOutputTokens templateComboBox.selectedItem = selectedItem.template updateTemplateHelpTextTooltip(selectedItem.template) } finally { @@ -233,6 +242,14 @@ class CustomServiceForm( name = nameField.text, template = templateComboBox.item, apiKey = getApiKey(), + contextWindowSize = readSpinnerValue( + contextWindowSizeSpinner, + editedItem.contextWindowSize + ), + maxOutputTokens = readSpinnerValue( + maxOutputTokensSpinner, + editedItem.maxOutputTokens + ), chatCompletionSettings = editedItem.chatCompletionSettings.copy( url = chatCompletionsForm.url, body = chatCompletionsForm.body, @@ -444,6 +461,14 @@ class CustomServiceForm( .addComponentToRightColumn( UIUtil.createComment("settingsConfigurable.service.custom.openai.apiKey.comment") ) + .addLabeledComponent( + CodeGPTBundle.get("settingsConfigurable.service.custom.openai.contextWindowSize.label"), + contextWindowSizeSpinner + ) + .addLabeledComponent( + CodeGPTBundle.get("settingsConfigurable.service.custom.openai.maxOutputTokens.label"), + maxOutputTokensSpinner + ) .addVerticalGap(4) .addComponent(tabbedPane) .addComponentFillVertically(JPanel(), 0) @@ -626,4 +651,18 @@ class CustomServiceForm( throw RuntimeException(e) } } + + private fun createTokenSpinner(initialValue: Long): JSpinner = + JSpinner(SpinnerNumberModel(initialValue, 1L, Long.MAX_VALUE, 1L)).apply { + (editor as? JSpinner.DefaultEditor)?.textField?.columns = 30 + } + + private fun readSpinnerValue(spinner: JSpinner, fallback: Long): Long { + return try { + spinner.commitEdit() + (spinner.value as? Number)?.toLong()?.coerceAtLeast(1L) ?: fallback + } catch (_: ParseException) { + fallback + } + } } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/form/model/CustomServiceSettingsData.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/form/model/CustomServiceSettingsData.kt index 76761d28..6fbd80c2 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/form/model/CustomServiceSettingsData.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/form/model/CustomServiceSettingsData.kt @@ -1,5 +1,7 @@ package ee.carlrobert.codegpt.settings.service.custom.form.model +import ee.carlrobert.codegpt.settings.service.custom.DEFAULT_CUSTOM_OPENAI_CONTEXT_WINDOW_SIZE +import ee.carlrobert.codegpt.settings.service.custom.DEFAULT_CUSTOM_OPENAI_MAX_OUTPUT_TOKENS import ee.carlrobert.codegpt.settings.service.custom.template.CustomServiceTemplate import java.util.UUID @@ -8,6 +10,8 @@ data class CustomServiceSettingsData( val name: String?, val template: CustomServiceTemplate, val apiKey: String?, + val contextWindowSize: Long = DEFAULT_CUSTOM_OPENAI_CONTEXT_WINDOW_SIZE, + val maxOutputTokens: Long = DEFAULT_CUSTOM_OPENAI_MAX_OUTPUT_TOKENS, val chatCompletionSettings: CustomServiceChatCompletionSettingsData, val codeCompletionSettings: CustomServiceCodeCompletionSettingsData ) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/form/model/DataToStateMapper.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/form/model/DataToStateMapper.kt index 21ad4902..320aa827 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/form/model/DataToStateMapper.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/form/model/DataToStateMapper.kt @@ -11,6 +11,8 @@ fun CustomServiceSettingsData.mapToState(): CustomServiceSettingsState = serviceState.id = if (id.isBlank()) UUID.randomUUID().toString() else id serviceState.name = name serviceState.template = template + serviceState.contextWindowSize = contextWindowSize + serviceState.maxOutputTokens = maxOutputTokens serviceState.chatCompletionSettings = chatCompletionSettings.mapToState() serviceState.codeCompletionSettings = codeCompletionSettings.mapToState() } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/form/model/StateToDataMapper.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/form/model/StateToDataMapper.kt index 76ebf851..6614f713 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/form/model/StateToDataMapper.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/form/model/StateToDataMapper.kt @@ -18,6 +18,8 @@ fun CustomServiceSettingsState.mapToData(): CustomServiceSettingsData = apiKey = if (!id.isNullOrEmpty()) CredentialsStore.getCredential(CredentialsStore.CredentialKey.CustomServiceApiKeyById(id!!)) else null, + contextWindowSize = contextWindowSize, + maxOutputTokens = maxOutputTokens, chatCompletionSettings = chatCompletionSettings.mapToData(), codeCompletionSettings = codeCompletionSettings.mapToData() ) diff --git a/src/main/resources/messages/codegpt.properties b/src/main/resources/messages/codegpt.properties index 480846be..e6d9c887 100644 --- a/src/main/resources/messages/codegpt.properties +++ b/src/main/resources/messages/codegpt.properties @@ -183,6 +183,8 @@ settingsConfigurable.service.custom.openai.importSettings=Import settings... settingsConfigurable.service.custom.openai.exportSettings=Export settings settingsConfigurable.service.custom.openai.marketing.text=Handling sensitive data? Connect to self-hosted models and manage configuration centrally with a private, configurable ProxyAI extension. settingsConfigurable.service.custom.openai.marketing.learnMore=Learn more about our custom extension +settingsConfigurable.service.custom.openai.contextWindowSize.label=Context window size: +settingsConfigurable.service.custom.openai.maxOutputTokens.label=Max output tokens: settingsConfigurable.prompts.import=Import settings... settingsConfigurable.prompts.export=Export settings settingsConfigurable.prompts.exportDialog.saveTo=Save to: diff --git a/src/main/resources/messages/codegpt_zh.properties b/src/main/resources/messages/codegpt_zh.properties index e59997da..639110ac 100644 --- a/src/main/resources/messages/codegpt_zh.properties +++ b/src/main/resources/messages/codegpt_zh.properties @@ -178,6 +178,8 @@ settingsConfigurable.service.custom.openai.connectionSuccess=\u8FDE\u63A5\u6210\ settingsConfigurable.service.custom.openai.connectionFailed=\u8FDE\u63A5\u5931\u8D25\u3002 settingsConfigurable.service.custom.openai.importSettings=\u5BFC\u5165\u8BBE\u7F6E... settingsConfigurable.service.custom.openai.exportSettings=\u5BFC\u51FA\u8BBE\u7F6E +settingsConfigurable.service.custom.openai.contextWindowSize.label=\u4E0A\u4E0B\u6587\u7A97\u53E3\u5927\u5C0F: +settingsConfigurable.service.custom.openai.maxOutputTokens.label=\u6700\u5927\u8F93\u51FA\u6807\u8BB0\u6570: settingsConfigurable.prompts.import=\u5BFC\u5165\u8BBE\u7F6E... settingsConfigurable.prompts.export=\u5BFC\u51FA\u8BBE\u7F6E settingsConfigurable.prompts.exportDialog.saveTo=\u4FDD\u5B58\u5230: diff --git a/src/test/kotlin/ee/carlrobert/codegpt/settings/models/ModelSettingsTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/settings/models/ModelSettingsTest.kt index 10e365c7..8c8982f8 100644 --- a/src/test/kotlin/ee/carlrobert/codegpt/settings/models/ModelSettingsTest.kt +++ b/src/test/kotlin/ee/carlrobert/codegpt/settings/models/ModelSettingsTest.kt @@ -1,5 +1,6 @@ package ee.carlrobert.codegpt.settings.models +import ai.koog.prompt.llm.LLMCapability import com.intellij.openapi.application.ApplicationManager import com.intellij.openapi.components.service import com.intellij.util.messages.MessageBusConnection @@ -228,6 +229,47 @@ class ModelSettingsTest : IntegrationTest() { .containsExactly("Second (chat-second)", "First (chat-first)") } + fun `test custom openai agent model uses configured token limits`() { + val customServicesSettings = service() + val customService = createCustomService( + name = "Custom Agent", + model = "custom-agent-model", + contextWindowSize = 512_000L, + maxOutputTokens = 12_000L + ) + + customServicesSettings.state.services.clear() + customServicesSettings.state.services.add(customService) + + val selection = modelSettings.getAvailableModels(FeatureType.AGENT) + .single { it.id == customService.id } + + assertThat(selection.llmModel.contextLength).isEqualTo(512_000L) + assertThat(selection.llmModel.maxOutputTokens).isEqualTo(12_000L) + assertThat(selection.llmModel.supports(LLMCapability.OpenAIEndpoint.Responses)).isFalse() + } + + fun `test custom openai responses agent model uses configured token limits`() { + val customServicesSettings = service() + val customService = createCustomService( + name = "Custom Responses", + model = "custom-responses-model", + contextWindowSize = 1_048_576L, + maxOutputTokens = 65_536L, + path = "https://example.com/v1/responses" + ) + + customServicesSettings.state.services.clear() + customServicesSettings.state.services.add(customService) + + val selection = modelSettings.getAvailableModels(FeatureType.AGENT) + .single { it.id == customService.id } + + assertThat(selection.llmModel.contextLength).isEqualTo(1_048_576L) + assertThat(selection.llmModel.maxOutputTokens).isEqualTo(65_536L) + assertThat(selection.llmModel.supports(LLMCapability.OpenAIEndpoint.Responses)).isTrue() + } + fun `test migrateMissingProviderInformation updates missing providers`() { val state = ModelSettingsState() val detailsState = ModelDetailsState() @@ -260,9 +302,18 @@ class ModelSettingsTest : IntegrationTest() { .map { it.displayName } } - private fun createCustomService(name: String, model: String): CustomServiceSettingsState { + private fun createCustomService( + name: String, + model: String, + contextWindowSize: Long = 200_000L, + maxOutputTokens: Long = 32_768L, + path: String = "https://example.com/v1/chat/completions" + ): CustomServiceSettingsState { return CustomServiceSettingsState().apply { this.name = name + this.contextWindowSize = contextWindowSize + this.maxOutputTokens = maxOutputTokens + chatCompletionSettings.url = path chatCompletionSettings.body.clear() chatCompletionSettings.body["model"] = model }