Skip to content

Commit 7423b7e

Browse files
committed
feat(mcp): enhance tool integration and UI in LLM config #371
Refactor enabledTools to use Tool objects instead of strings for better type safety. Move system prompt generation to McpLlmConfig class. Replace result label with text area for better output display. These changes improve tool handling and user experience.
1 parent 0c92bed commit 7423b7e

File tree

2 files changed

+53
-59
lines changed

2 files changed

+53
-59
lines changed

core/src/main/kotlin/cc/unitmesh/devti/mcp/editor/McpLlmConfigDialog.kt

Lines changed: 41 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import com.intellij.ui.dsl.builder.TopGap
99
import com.intellij.ui.dsl.builder.panel
1010
import com.intellij.util.ui.JBUI
1111
import com.intellij.util.ui.UIUtil
12+
import com.sun.java.accessibility.util.AWTEventMonitor.addActionListener
1213
import io.modelcontextprotocol.kotlin.sdk.Tool
1314
import kotlinx.serialization.encodeToString
1415
import kotlinx.serialization.json.Json
@@ -17,14 +18,42 @@ import javax.swing.JSlider
1718

1819
data class McpLlmConfig(
1920
var temperature: Double = 0.0,
20-
var enabledTools: MutableList<String> = mutableListOf(),
21+
var enabledTools: MutableList<Tool> = mutableListOf(),
2122
var systemPrompt: String = ""
22-
)
23+
) {
24+
fun createSystemPrompt(): String {
25+
val systemPrompt = """
26+
In this environment you have access to a set of tools you can use to answer the user's question.
27+
You can invoke functions by writing a "<antml:function_calls>" block like the following as part of your reply to the user:
28+
<antml:function_calls>
29+
<antml:invoke name="${'$'}FUNCTION_NAME">
30+
<antml:parameter name="${'$'}PARAMETER_NAME">${'$'}PARAMETER_VALUE</antml:parameter>
31+
...
32+
</antml:invoke>
33+
<antml:invoke name="${'$'}FUNCTION_NAME2">
34+
...
35+
</antml:invoke>
36+
</antml:function_calls>
37+
38+
String and scalar parameters should be specified as is, while lists and objects should use JSON format.
39+
40+
Here are the functions available in JSONSchema format:
41+
<functions>
42+
${enabledTools.joinToString("\n") { tool -> "<function>" + Json.encodeToString(tool) } + "</function>"} }
43+
</functions>
44+
45+
Answer the user's request using the relevant tool(s), if they are available. Check that all the required parameters for each tool call are provided or can reasonably be inferred from context. IF there are no relevant tools or there are missing values for required parameters, ask the user to supply these values; otherwise proceed with the tool calls. If the user provides a specific value for a parameter (for example provided in quotes), make sure to use that value EXACTLY. DO NOT make up values for or ask about optional parameters. Carefully analyze descriptive terms in the request as they may indicate required parameter values that should be included even if not explicitly quoted.
46+
47+
If you intend to call multiple tools and there are no dependencies between the calls, make all of the independent calls in the same <antml:function_calls></antml:function_calls> block.
48+
""".trimIndent()
49+
return systemPrompt
50+
}
51+
}
2352

2453
class McpLlmConfigDialog(
2554
private val project: Project,
2655
private val config: McpLlmConfig,
27-
private val allTools: Map<String, List<io.modelcontextprotocol.kotlin.sdk.Tool>>
56+
private val allTools: Map<String, List<Tool>>
2857
) : DialogWrapper(project) {
2958
private lateinit var temperatureSlider: JSlider
3059
private val toolCheckboxes = mutableMapOf<String, JBCheckBox>()
@@ -33,12 +62,7 @@ class McpLlmConfigDialog(
3362
init {
3463
title = "Chatbot Configuration"
3564
allTools.forEach { (serverName, tools) ->
36-
tools.forEach { tool ->
37-
val toolId = "${serverName}:${tool.name}"
38-
if (!config.enabledTools.contains(toolId)) {
39-
config.enabledTools.add(toolId)
40-
}
41-
}
65+
config.enabledTools.addAll(tools)
4266
}
4367

4468
init()
@@ -48,8 +72,7 @@ class McpLlmConfigDialog(
4872
* Based on https://github.com/jujumilk3/leaked-system-prompts/blob/main/anthropic-claude-api-tool-use_20250119.md
4973
*/
5074
override fun createCenterPanel(): JComponent {
51-
val systemPrompt = createSystemPrompt()
52-
75+
val systemPrompt = config.createSystemPrompt()
5376
promptField = MarkdownLanguageField(project, systemPrompt, "", "instructions.md").also {
5477
it.preferredSize = JBUI.size(480, 540)
5578
}
@@ -80,19 +103,16 @@ class McpLlmConfigDialog(
80103
group("Enabled Tools") {
81104
allTools.forEach { (serverName, tools) ->
82105
tools.forEach { tool ->
83-
val toolId = "${serverName}:${tool.name}"
84106
row {
85107
label("${tool.name} (${serverName})")
86108
checkBox("").apply {
87-
component.isSelected = config.enabledTools.contains(toolId)
88-
toolCheckboxes[toolId] = component
89-
component.addActionListener {
90-
if (component.isSelected) {
91-
if (!config.enabledTools.contains(toolId)) {
92-
config.enabledTools.add(toolId)
93-
}
109+
toolCheckboxes[tool.name] = this.component
110+
this.component.isSelected = config.enabledTools.any { it.name == tool.name }
111+
addActionListener {
112+
if (this.component.isSelected) {
113+
config.enabledTools.add(tool)
94114
} else {
95-
config.enabledTools.remove(toolId)
115+
config.enabledTools.remove(tool)
96116
}
97117
}
98118
}
@@ -110,36 +130,8 @@ class McpLlmConfigDialog(
110130
}.withPreferredSize(500, 600)
111131
}
112132

113-
private fun createSystemPrompt(): String {
114-
val systemPrompt = """
115-
In this environment you have access to a set of tools you can use to answer the user's question.
116-
You can invoke functions by writing a "<antml:function_calls>" block like the following as part of your reply to the user:
117-
<antml:function_calls>
118-
<antml:invoke name="${'$'}FUNCTION_NAME">
119-
<antml:parameter name="${'$'}PARAMETER_NAME">${'$'}PARAMETER_VALUE</antml:parameter>
120-
...
121-
</antml:invoke>
122-
<antml:invoke name="${'$'}FUNCTION_NAME2">
123-
...
124-
</antml:invoke>
125-
</antml:function_calls>
126-
127-
String and scalar parameters should be specified as is, while lists and objects should use JSON format.
128-
129-
Here are the functions available in JSONSchema format:
130-
<functions>
131-
${allTools.values.flatten().joinToString("\n") { tool -> "<function>" + Json.encodeToString(tool) } + "</function>"} }
132-
</functions>
133-
134-
Answer the user's request using the relevant tool(s), if they are available. Check that all the required parameters for each tool call are provided or can reasonably be inferred from context. IF there are no relevant tools or there are missing values for required parameters, ask the user to supply these values; otherwise proceed with the tool calls. If the user provides a specific value for a parameter (for example provided in quotes), make sure to use that value EXACTLY. DO NOT make up values for or ask about optional parameters. Carefully analyze descriptive terms in the request as they may indicate required parameter values that should be included even if not explicitly quoted.
135-
136-
If you intend to call multiple tools and there are no dependencies between the calls, make all of the independent calls in the same <antml:function_calls></antml:function_calls> block.
137-
""".trimIndent()
138-
return systemPrompt
139-
}
140-
141133
fun getConfig(): McpLlmConfig {
142-
config.systemPrompt = promptField.text
134+
config.systemPrompt = config.createSystemPrompt()
143135
return config
144136
}
145137
}

core/src/main/kotlin/cc/unitmesh/devti/mcp/editor/McpPreviewEditor.kt

Lines changed: 12 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,7 @@ open class McpPreviewEditor(
6161
private lateinit var chatInput: JBTextField
6262
private lateinit var testButton: ActionButton
6363
private lateinit var configButton: JButton
64-
private lateinit var resultLabel: JLabel
64+
private lateinit var resultTextArea: JTextArea
6565
private lateinit var resultPanel: JPanel
6666
private val config = McpLlmConfig()
6767
private val borderColor = JBColor(0xE5E7EB, 0x3C3F41) // Equivalent to Tailwind gray-200
@@ -154,16 +154,18 @@ open class McpPreviewEditor(
154154

155155
resultPanel = JPanel(BorderLayout()).apply {
156156
background = UIUtil.getPanelBackground()
157-
border = CompoundBorder(
158-
BorderFactory.createMatteBorder(1, 0, 1, 0, borderColor),
159-
JBUI.Borders.empty(16)
160-
)
161157
isVisible = false
162158
}
163159

164-
resultLabel = JLabel()
160+
resultTextArea = JTextArea().apply {
161+
isEditable = false
162+
wrapStyleWord = true
163+
lineWrap = true
164+
background = UIUtil.getPanelBackground()
165+
border = JBUI.Borders.empty(4)
166+
}
165167

166-
val resultScrollPane = JBScrollPane(resultLabel).apply {
168+
val resultScrollPane = JBScrollPane(resultTextArea).apply {
167169
border = BorderFactory.createEmptyBorder()
168170
background = UIUtil.getPanelBackground()
169171
}
@@ -300,9 +302,9 @@ open class McpPreviewEditor(
300302
val llmProvider = CustomLLMProvider(project, llmConfig)
301303
val message = chatInput.text.trim()
302304
val result = StringBuilder()
303-
val stream: Flow<String> = llmProvider.stream(message, systemPrompt = config.systemPrompt)
305+
val stream: Flow<String> = llmProvider.stream(message, systemPrompt = config.createSystemPrompt())
304306

305-
resultLabel.text = "Loading response..."
307+
resultTextArea.text = "Loading response..."
306308
resultPanel.isVisible = true
307309
mainPanel.revalidate()
308310
mainPanel.repaint()
@@ -311,7 +313,7 @@ open class McpPreviewEditor(
311313
stream.cancellable().collect { chunk ->
312314
result.append(chunk)
313315
SwingUtilities.invokeLater {
314-
resultLabel.text = result.toString()
316+
resultTextArea.text = result.toString()
315317
mainPanel.revalidate()
316318
mainPanel.repaint()
317319
}

0 commit comments

Comments
 (0)