mirror of
https://github.com/carlrobertoh/ProxyAI.git
synced 2026-07-09 17:28:59 +00:00
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
This commit is contained in:
Roman Gromov
parent
9e8a437a0a
commit
f30c8ace7a
10 changed files with 135 additions and 6 deletions
|
|
@ -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<MessageWithContext, String, out AIAgent<MessageWithContext, String>>,
|
||||
val provider: ServiceType,
|
||||
val modelId: String,
|
||||
val modelSignature: ModelRuntimeSignature,
|
||||
val selectedServerIds: Set<String>,
|
||||
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<String, Job>()
|
||||
private val pendingMessages = ConcurrentHashMap<String, ArrayDeque<MessageWithContext>>()
|
||||
private val sessionAgents = ConcurrentHashMap<String, AIAgent<MessageWithContext, String>>()
|
||||
|
|
@ -256,14 +265,14 @@ class AgentService(private val project: Project) {
|
|||
provider: ServiceType,
|
||||
events: AgentEvents
|
||||
): SessionRuntime {
|
||||
val modelId = service<ModelSettings>().getAgentModel().id
|
||||
val modelSignature = currentModelSignature()
|
||||
val selectedServerIds = project.service<AgentMcpContextService>()
|
||||
.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<ModelSettings>().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}",
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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())
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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<CustomServicesSettings>()
|
||||
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<CustomServicesSettings>()
|
||||
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
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue