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 }