From cfbaa2b0955676607b7fe019a85acf459f6b5309 Mon Sep 17 00:00:00 2001 From: Carl-Robert Linnupuu Date: Tue, 3 Feb 2026 10:30:12 +0000 Subject: [PATCH] fix: agent tool UI cosmetics --- .../java/ee/carlrobert/codegpt/ui/UIUtil.java | 9 + .../clients/CustomOpenAIChatCompletion.kt | 1 + .../agent/clients/CustomOpenAILLMClient.kt | 36 ++- .../agent/clients/ProxyAIChatCompletion.kt | 1 + .../codegpt/agent/clients/ProxyAILLMClient.kt | 35 ++- .../codegpt/agent/tools/GetLibraryDocsTool.kt | 34 +-- .../descriptor/ToolCallDescriptorFactory.kt | 232 +++++++++++++----- 7 files changed, 259 insertions(+), 89 deletions(-) diff --git a/src/main/java/ee/carlrobert/codegpt/ui/UIUtil.java b/src/main/java/ee/carlrobert/codegpt/ui/UIUtil.java index 657e5db9..9b1af157 100644 --- a/src/main/java/ee/carlrobert/codegpt/ui/UIUtil.java +++ b/src/main/java/ee/carlrobert/codegpt/ui/UIUtil.java @@ -85,6 +85,15 @@ public class UIUtil { return textArea; } + public static JBTextArea createReadOnlyTextArea(String text) { + var textArea = new JBTextArea(text); + textArea.setEditable(false); + textArea.setLineWrap(true); + textArea.setWrapStyleWord(true); + textArea.setFont(JBUI.Fonts.smallFont()); + return textArea; + } + public static JButton createIconButton(Icon icon) { var button = new JButton(icon); button.setBorder(BorderFactory.createEmptyBorder()); diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/CustomOpenAIChatCompletion.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/CustomOpenAIChatCompletion.kt index 509638e4..550e6867 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/CustomOpenAIChatCompletion.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/CustomOpenAIChatCompletion.kt @@ -98,6 +98,7 @@ public class CustomOpenAIStreamDelta( @Serializable public class CustomOpenAIToolCall( public val id: String? = "", + public val index: Int? = 0, public val function: CustomOpenAIFunction ) { /** The type of the tool. Currently, only `function` is supported. */ diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/CustomOpenAILLMClient.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/CustomOpenAILLMClient.kt index 8d0bebc2..6508b3f6 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/CustomOpenAILLMClient.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/CustomOpenAILLMClient.kt @@ -9,16 +9,20 @@ import ai.koog.prompt.executor.clients.openai.base.AbstractOpenAILLMClient import ai.koog.prompt.executor.clients.openai.base.OpenAIBaseSettings import ai.koog.prompt.executor.clients.openai.base.OpenAICompatibleToolDescriptorSchemaGenerator import ai.koog.prompt.executor.clients.openai.base.models.* +import ai.koog.prompt.executor.clients.openai.models.OpenAIChatCompletionStreamResponse import ai.koog.prompt.llm.LLMProvider import ai.koog.prompt.llm.LLModel import ai.koog.prompt.message.LLMChoice import ai.koog.prompt.message.Message import ai.koog.prompt.message.ResponseMetaInfo import ai.koog.prompt.params.LLMParams +import ai.koog.prompt.streaming.StreamFrame import ai.koog.prompt.streaming.StreamFrameFlowBuilder +import ai.koog.prompt.streaming.buildStreamFrameFlow import ee.carlrobert.codegpt.settings.service.custom.CustomServiceChatCompletionSettingsState import io.github.oshai.kotlinlogging.KotlinLogging import io.ktor.client.* +import kotlinx.coroutines.flow.Flow import kotlinx.datetime.Clock import java.net.URI import kotlin.io.encoding.ExperimentalEncodingApi @@ -141,17 +145,31 @@ public class CustomOpenAILLMClient( override fun decodeResponse(data: String): CustomOpenAIChatCompletionResponse = json.decodeFromString(data) - override suspend fun StreamFrameFlowBuilder.processStreamingChunk(chunk: CustomOpenAIChatCompletionStreamResponse) { - chunk.choices.firstOrNull()?.let { choice -> - choice.delta.content?.let { emitAppend(it) } - choice.delta.toolCalls?.forEachIndexed { index, openAIToolCall -> - val id = openAIToolCall.id - val name = openAIToolCall.function.name - val arguments = openAIToolCall.function.arguments - upsertToolCall(index, id, name, arguments) + override fun processStreamingResponse( + response: Flow + ): Flow = buildStreamFrameFlow { + var finishReason: String? = null + var metaInfo: ResponseMetaInfo? = null + + response.collect { chunk -> + chunk.choices.firstOrNull()?.let { choice -> + choice.delta.content?.let { emitAppend(it) } + + choice.delta.toolCalls?.forEach { openAIToolCall -> + val index = openAIToolCall.index ?: 0 + val id = openAIToolCall.id + val functionName = openAIToolCall.function.name + val functionArgs = openAIToolCall.function.arguments + upsertToolCall(index, id, functionName, functionArgs) + } + + choice.finishReason?.let { finishReason = it } } - choice.finishReason?.let { emitEnd(it, createMetaInfo(chunk.usage)) } + + chunk.usage?.let { metaInfo = createMetaInfo(it) } } + + emitEnd(finishReason, metaInfo) } override suspend fun moderate(prompt: Prompt, model: LLModel): ModerationResult { diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/ProxyAIChatCompletion.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/ProxyAIChatCompletion.kt index 1bcd99a3..5ee66220 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/ProxyAIChatCompletion.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/ProxyAIChatCompletion.kt @@ -66,6 +66,7 @@ public class ProxyAIStreamDelta( @Serializable public class ProxyAIToolCall( public val id: String? = "", + public val index: Int? = 0, public val function: ProxyAIFunction ) { /** The type of the tool. Currently, only `function` is supported. */ diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/ProxyAILLMClient.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/ProxyAILLMClient.kt index eab9199c..44bae851 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/ProxyAILLMClient.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/ProxyAILLMClient.kt @@ -17,9 +17,12 @@ import ai.koog.prompt.message.LLMChoice import ai.koog.prompt.message.Message import ai.koog.prompt.message.ResponseMetaInfo import ai.koog.prompt.params.LLMParams +import ai.koog.prompt.streaming.StreamFrame import ai.koog.prompt.streaming.StreamFrameFlowBuilder +import ai.koog.prompt.streaming.buildStreamFrameFlow import io.github.oshai.kotlinlogging.KotlinLogging import io.ktor.client.* +import kotlinx.coroutines.flow.Flow import kotlinx.datetime.Clock import kotlinx.serialization.json.JsonObject import kotlinx.serialization.json.JsonObjectBuilder @@ -128,17 +131,31 @@ public class ProxyAILLMClient( override fun decodeResponse(data: String): ProxyAIChatCompletionResponse = json.decodeFromString(data) - override suspend fun StreamFrameFlowBuilder.processStreamingChunk(chunk: ProxyAIChatCompletionStreamResponse) { - chunk.choices.firstOrNull()?.let { choice -> - choice.delta?.content?.let { emitAppend(it) } - choice.delta?.toolCalls?.forEachIndexed { index, toolCall -> - val id = toolCall.id - val name = toolCall.function.name - val arguments = toolCall.function.arguments - upsertToolCall(index, id, name, arguments) + override fun processStreamingResponse( + response: Flow + ): Flow = buildStreamFrameFlow { + var finishReason: String? = null + var metaInfo: ResponseMetaInfo? = null + + response.collect { chunk -> + chunk.choices.firstOrNull()?.let { choice -> + choice.delta?.content?.let { emitAppend(it) } + + choice.delta?.toolCalls?.forEach { openAIToolCall -> + val index = openAIToolCall.index ?: 0 + val id = openAIToolCall.id + val functionName = openAIToolCall.function.name + val functionArgs = openAIToolCall.function.arguments + upsertToolCall(index, id, functionName, functionArgs) + } + + choice.finishReason?.let { finishReason = it } } - choice.finishReason?.let { emitEnd(it, createProxyMetaInfo(chunk.usage)) } + + chunk.usage?.let { metaInfo = createProxyMetaInfo(it) } } + + emitEnd(finishReason, metaInfo) } override suspend fun moderate(prompt: Prompt, model: LLModel): ModerationResult { diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/GetLibraryDocsTool.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/GetLibraryDocsTool.kt index fc6cf941..361c58bc 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/GetLibraryDocsTool.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/GetLibraryDocsTool.kt @@ -180,26 +180,28 @@ class GetLibraryDocsTool( is Result.Error -> ("Failed to retrieve library documentation: ${result.error}").truncateToolResult() } - private fun parseLibraryId(libraryId: String): LibraryComponents { - val cleaned = libraryId.removePrefix("/") - val parts = cleaned.split("/") + companion object { + fun parseLibraryId(libraryId: String): LibraryComponents { + val cleaned = libraryId.removePrefix("/") + val parts = cleaned.split("/") - if (parts.size < 2) { - throw IllegalArgumentException( - "Invalid library ID format: $libraryId. Expected format: /username/library or /username/library/tag" + if (parts.size < 2) { + throw IllegalArgumentException( + "Invalid library ID format: $libraryId. Expected format: /username/library or /username/library/tag" + ) + } + + return LibraryComponents( + username = parts[0], + library = parts[1], + tag = if (parts.size > 2) parts[2] else null ) } - return LibraryComponents( - username = parts[0], - library = parts[1], - tag = if (parts.size > 2) parts[2] else null + data class LibraryComponents( + val username: String, + val library: String, + val tag: String? ) } - - private data class LibraryComponents( - val username: String, - val library: String, - val tag: String? - ) } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/ui/descriptor/ToolCallDescriptorFactory.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/ui/descriptor/ToolCallDescriptorFactory.kt index d37c0e79..32718b65 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/ui/descriptor/ToolCallDescriptorFactory.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/ui/descriptor/ToolCallDescriptorFactory.kt @@ -11,6 +11,7 @@ import ee.carlrobert.codegpt.toolwindow.agent.ui.approval.DiffViewAction import ee.carlrobert.codegpt.toolwindow.agent.ui.renderer.ChangeColors import ee.carlrobert.codegpt.toolwindow.agent.ui.renderer.DiffBadgeText import ee.carlrobert.codegpt.toolwindow.agent.ui.renderer.diffBadgeText +import ee.carlrobert.codegpt.ui.UIUtil import kotlinx.serialization.json.JsonArray import kotlinx.serialization.json.JsonObject import kotlinx.serialization.json.JsonPrimitive @@ -53,7 +54,7 @@ object ToolCallDescriptorFactory { ToolKind.LIBRARY_RESOLVE -> createLibraryResolveDescriptor(args, result, projectId) ToolKind.LIBRARY_DOCS -> createLibraryDocsDescriptor(args, result, projectId) ToolKind.ASK_QUESTION -> createAskDescriptor(args, result, projectId) - ToolKind.EXIT -> createExitDescriptor(toolName, args, result, projectId) + ToolKind.EXIT -> createExitDescriptor(args, result, projectId) ToolKind.OTHER -> createOtherDescriptor(toolName, args, result, projectId) } } @@ -88,7 +89,6 @@ object ToolCallDescriptorFactory { titlePrefix = "Clarify Requirements", titleMain = "", tooltip = "Ask the user clarifying questions", - supportsStreaming = false, args = args, result = result, projectId = projectId @@ -96,7 +96,6 @@ object ToolCallDescriptorFactory { } private fun createExitDescriptor( - toolName: String, args: Any, result: Any?, projectId: String? @@ -107,33 +106,39 @@ object ToolCallDescriptorFactory { titlePrefix = "Exit", titleMain = "", tooltip = "Agent task completed", - supportsStreaming = false, args = args, result = result, projectId = projectId ) } - private fun showTextDialog(content: String, title: String) { - val dialog = JDialog().apply { - this.title = title - isModal = true - } - val textArea = JTextArea(content).apply { - isEditable = false - lineWrap = true - wrapStyleWord = true - font = JBUI.Fonts.smallFont() - } - val scrollPane = JScrollPane(textArea).apply { + private fun createScrollPaneWithBorder(textArea: JTextArea): JScrollPane { + return JScrollPane(textArea).apply { preferredSize = JBUI.size(700, 400) border = JBUI.Borders.customLine( JBUI.CurrentTheme.CustomFrameDecorations.separatorForeground() ) } - val footer = JPanel(FlowLayout(FlowLayout.RIGHT)).apply { + } + + private fun createFooterButtonPanel(vararg buttons: JButton): JPanel { + return JPanel(FlowLayout(FlowLayout.RIGHT)).apply { isOpaque = false - add(JButton("Copy").apply { + for (button in buttons) { + add(button) + } + } + } + + private fun createDialogFooterPanel(dialog: JDialog): JPanel { + return createFooterButtonPanel( + JButton("Close").apply { addActionListener { dialog.dispose() } } + ) + } + + private fun createDialogFooterPanelWithCopy(dialog: JDialog, content: String): JPanel { + return createFooterButtonPanel( + JButton("Copy").apply { addActionListener { val selection = StringSelection(content) Toolkit.getDefaultToolkit().systemClipboard.setContents( @@ -141,18 +146,32 @@ object ToolCallDescriptorFactory { null ) } - }) - add(JButton("Close").apply { addActionListener { dialog.dispose() } }) - } + }, + JButton("Close").apply { addActionListener { dialog.dispose() } } + ) + } + + private fun showDialog(dialog: JDialog, scrollPane: JScrollPane, footerPanel: JPanel) { dialog.contentPane = BorderLayoutPanel().apply { add(scrollPane, BorderLayout.CENTER) - add(footer, BorderLayout.SOUTH) + add(footerPanel, BorderLayout.SOUTH) } dialog.pack() dialog.setLocationRelativeTo(null) dialog.isVisible = true } + private fun showTextDialog(content: String, title: String) { + val dialog = JDialog().apply { + this.title = title + isModal = true + } + val textArea = UIUtil.createReadOnlyTextArea(content) + val scrollPane = createScrollPaneWithBorder(textArea) + val footer = createDialogFooterPanelWithCopy(dialog, content) + showDialog(dialog, scrollPane, footer) + } + private fun createReadDescriptor( args: Any, result: Any?, @@ -183,8 +202,6 @@ object ToolCallDescriptorFactory { enabled = true ), secondaryBadges = listOfNotNull(lineBadge), - actions = emptyList(), - supportsStreaming = false, args = args, result = result, projectId = projectId @@ -206,7 +223,7 @@ object ToolCallDescriptorFactory { if (result is WriteTool.Result && writeArgs != null) { when (result) { is WriteTool.Result.Success -> { - badges.add(Badge("${result.bytesWritten} bytes", JBColor.GREEN)) + badges.add(Badge("[${writeArgs.content.lines().size} lines]", JBColor.GREEN)) actions.add( ToolAction("View Changes", AllIcons.Actions.Diff) { DiffViewAction.showDiff(writeArgs.filePath, project) @@ -234,7 +251,6 @@ object ToolCallDescriptorFactory { enabled = result != null ), actions = actions, - supportsStreaming = false, args = args, result = result, projectId = projectId @@ -322,7 +338,6 @@ object ToolCallDescriptorFactory { column = firstLocation?.column ), actions = actions, - supportsStreaming = false, args = args, result = result, projectId = projectId, @@ -387,6 +402,18 @@ object ToolCallDescriptorFactory { ) } + private fun buildSearchBadges(result: Any?): List { + return if (result is IntelliJSearchTool.Result) { + listOf(Badge( + "[${result.totalMatches} matches]", + JBColor.BLUE, + action = { showTextDialog(result.output, "Search Results") } + )) + } else { + emptyList() + } + } + private fun createSearchDescriptor( args: Any, result: Any?, @@ -396,29 +423,19 @@ object ToolCallDescriptorFactory { val pattern = searchArgs?.pattern ?: "" val scopeOrPath = searchArgs?.path?.substringAfterLast('/') ?: (searchArgs?.scope ?: "") val titleMain = buildSearchDisplay(truncatePattern(pattern), scopeOrPath) - val badges = mutableListOf() - val actions = mutableListOf() - - when (result) { - is IntelliJSearchTool.Result -> { - badges.add( - Badge( - "${result.totalMatches} matches", - JBColor.BLUE, - action = { showTextDialog(result.output, "Search Results") }) - ) - } - } return ToolCallDescriptor( kind = ToolKind.SEARCH, icon = AllIcons.Actions.Search, titlePrefix = "Search:", titleMain = titleMain, - tooltip = buildTooltipString("Search", pattern, scopeOrPath.ifBlank { null }), - secondaryBadges = badges, - actions = actions, - supportsStreaming = false, + tooltip = if (scopeOrPath.isBlank()) { + "Search: \"$pattern\"" + } else { + "Search: \"$pattern\" in $scopeOrPath" + }, + secondaryBadges = buildSearchBadges(result), + actions = emptyList(), args = args, result = result, projectId = projectId @@ -443,7 +460,7 @@ object ToolCallDescriptorFactory { titlePrefix = "Web:", titleMain = truncatedQuery, tooltip = "Web search: $query", - supportsStreaming = false, + secondaryBadges = buildWebBadges(args, result), args = args, result = result, projectId = projectId @@ -487,8 +504,6 @@ object ToolCallDescriptorFactory { titlePrefix = titlePrefix, titleMain = description, tooltip = "Task: $description", - secondaryBadges = emptyList(), - supportsStreaming = false, args = args, result = result, projectId = projectId, @@ -609,7 +624,13 @@ object ToolCallDescriptorFactory { val badges = mutableListOf() if (result is ResolveLibraryIdTool.Result.Success) { - badges.add(Badge("${result.libraries.size} found", JBColor.BLUE)) + badges.add( + Badge( + "[${result.libraries.size} found]", + JBColor.BLUE, + action = { showLibrariesDialog(result) } + ) + ) } return ToolCallDescriptor( @@ -619,7 +640,6 @@ object ToolCallDescriptorFactory { titleMain = libraryName, tooltip = "Resolve library: $libraryName", secondaryBadges = badges, - supportsStreaming = false, args = args, result = result, projectId = projectId @@ -642,7 +662,7 @@ object ToolCallDescriptorFactory { titlePrefix = "Docs:", titleMain = libraryId, tooltip = "Get library docs: $libraryId", - supportsStreaming = false, + secondaryBadges = buildDocsBadges(result), args = args, result = result, projectId = projectId @@ -661,7 +681,6 @@ object ToolCallDescriptorFactory { titlePrefix = "Tool:", titleMain = toolName, tooltip = "Tool: $toolName", - supportsStreaming = false, args = args, result = result, projectId = projectId @@ -704,11 +723,114 @@ object ToolCallDescriptorFactory { } } - private fun buildTooltipString(operation: String, pattern: String, scope: String?): String { - return if (scope.isNullOrBlank()) { - "$operation: \"$pattern\"" + private fun buildDocsBadges(result: Any?): List { + return if (result is GetLibraryDocsTool.Result.Success) { + listOf(Badge( + "[View Results]", + JBColor.BLUE, + action = { + showTextDialog( + result.documentation, + "Documentation: ${result.libraryId}" + ) + } + )) } else { - "$operation: \"$pattern\" in $scope" + emptyList() } } + + private fun buildWebBadges(args: Any, result: Any?): List { + if (result !is WebSearchTool.Result) return emptyList() + + val argsObj = args as? WebSearchTool.Args + val badges = mutableListOf(Badge( + "[${result.results.size} results]", + JBColor.BLUE, + action = { showWebResultsDialog(result) } + )) + + if (argsObj != null && !argsObj.allowedDomains.isNullOrEmpty()) { + badges.add(Badge("[${argsObj.allowedDomains.size} domains]", JBColor.GRAY)) + } + + return badges + } + + private fun showWebResultsDialog(result: WebSearchTool.Result) { + val dialog = JDialog().apply { + title = "Web Search Results" + isModal = true + } + + val content = buildString { + if (result.results.isEmpty()) { + appendLine("No search results found.") + } else { + result.results.forEachIndexed { index, searchResult -> + appendLine("${index + 1}. ${searchResult.title}") + appendLine(" URL: ${searchResult.url}") + appendLine(" ${searchResult.content}") + appendLine() + } + } + } + + val textArea = UIUtil.createReadOnlyTextArea(content) + val scrollPane = createScrollPaneWithBorder(textArea) + val footerPanel = createDialogFooterPanel(dialog) + showDialog(dialog, scrollPane, footerPanel) + } + + private fun showLibrariesDialog(result: ResolveLibraryIdTool.Result.Success) { + val content = buildString { + if (result.libraries.isEmpty()) { + appendLine("No libraries found for '${result.libraryName}'.") + appendLine() + appendLine("Please try with different search terms or check the library name spelling.") + } else { + appendLine("Available Libraries:") + appendLine() + result.libraries.forEachIndexed { index, library -> + appendLine("${index + 1}. ${library.name}") + appendLine(" Library ID: ${library.id}") + if (library.description.isNotBlank()) { + appendLine(" Description: ${library.description}") + } + appendLine(" Code Snippets: ${library.codeSnippets}") + appendLine(" Source Reputation: ${library.sourceReputation}") + appendLine(" Benchmark Score: ${library.benchmarkScore}") + if (!library.versions.isNullOrEmpty()) { + appendLine(" Available Versions: ${library.versions.joinToString(", ")}") + } + appendLine() + } + + val topLibrary = result.libraries.maxByOrNull { + (it.benchmarkScore * 0.4 + it.codeSnippets * 0.3 + when (it.sourceReputation.lowercase()) { + "high" -> 30 + "medium" -> 20 + "low" -> 10 + else -> 0 + } * 0.3).toInt() + } + if (topLibrary != null) { + appendLine("Recommended Selection:") + appendLine() + appendLine("Library ID: ${topLibrary.id}") + appendLine("Name: ${topLibrary.name}") + appendLine("Reasoning: Highest combined score of benchmark (${topLibrary.benchmarkScore}), code snippets (${topLibrary.codeSnippets}), and source reputation (${topLibrary.sourceReputation})") + } + } + } + + val textArea = UIUtil.createReadOnlyTextArea(content) + val scrollPane = createScrollPaneWithBorder(textArea) + val dialog = JDialog().apply { + title = "Library Search Results" + isModal = true + } + val footerPanel = createDialogFooterPanel(dialog) + showDialog(dialog, scrollPane, footerPanel) + } }