From 702ba92ab72afed1432c258b4d6329e02da094a0 Mon Sep 17 00:00:00 2001 From: Konstantin Aleev Date: Thu, 2 Oct 2025 06:37:25 +0200 Subject: [PATCH] Examples: AIA MCP client registry - place AIA examples into a separate directory --- examples/ai_assistant/chat_handler.ipynb | 226 ++++++++++++++++ examples/ai_assistant/mcp_client.ipynb | 325 +++++++++++++++++++++++ examples/aia_extensibility.ipynb | 140 ---------- 3 files changed, 551 insertions(+), 140 deletions(-) create mode 100644 examples/ai_assistant/chat_handler.ipynb create mode 100644 examples/ai_assistant/mcp_client.ipynb delete mode 100644 examples/aia_extensibility.ipynb diff --git a/examples/ai_assistant/chat_handler.ipynb b/examples/ai_assistant/chat_handler.ipynb new file mode 100644 index 0000000..a04a70c --- /dev/null +++ b/examples/ai_assistant/chat_handler.ipynb @@ -0,0 +1,226 @@ +{ + "cells": [ + { + "metadata": {}, + "cell_type": "markdown", + "source": [ + "# JetBrains AI Assistant Chat Handler\n", + "\n", + "This notebook presents how to extend the [JetBrains AI Assistant](https://plugins.jetbrains.com/plugin/22282-jetbrains-ai-assistant) plugin in the runtime by providing a custom message handler." + ] + }, + { + "metadata": { + "ExecuteTime": { + "end_time": "2025-10-02T04:22:19.265561Z", + "start_time": "2025-10-02T04:22:17.834847Z" + }, + "executionRelatedData": { + "compiledClasses": [ + "Line_1_jupyter", + "Line_2_jupyter" + ] + } + }, + "cell_type": "code", + "source": "%use intellij-platform", + "outputs": [ + { + "data": { + "text/plain": [ + "IntelliJ Platform integration is loaded" + ] + }, + "metadata": {}, + "output_type": "display_data", + "jetTransient": { + "display_id": null + } + } + ], + "execution_count": 1 + }, + { + "metadata": {}, + "cell_type": "markdown", + "source": "Load the JetBrains AI Assistant bundled plugin using its `com.intellij.ml.llm` ID." + }, + { + "metadata": { + "ExecuteTime": { + "end_time": "2025-10-02T04:22:22.365779Z", + "start_time": "2025-10-02T04:22:21.587057Z" + }, + "executionRelatedData": { + "compiledClasses": [ + "Line_4_jupyter", + "Line_5_jupyter", + "Line_6_jupyter", + "Line_7_jupyter", + "Line_8_jupyter" + ] + } + }, + "cell_type": "code", + "source": "loadPlugins(\"com.intellij.ml.llm\")", + "outputs": [], + "execution_count": 2 + }, + { + "metadata": { + "ExecuteTime": { + "end_time": "2025-10-02T03:55:51.653816Z", + "start_time": "2025-10-02T03:55:51.239719Z" + } + }, + "cell_type": "code", + "source": [ + "import com.intellij.ml.llm.agents.ChatAgent\n", + "import com.intellij.ml.llm.core.chat.messages.ChatMessage\n", + "import com.intellij.ml.llm.core.chat.messages.UserMessage\n", + "import com.intellij.ml.llm.core.chat.messages.impl.SimpleCompletableMessage\n", + "import com.intellij.ml.llm.core.chat.session.ChatKind\n", + "import com.intellij.ml.llm.core.chat.session.ChatMessageHandler\n", + "import com.intellij.ml.llm.core.chat.session.ChatSession\n", + "import com.intellij.ml.llm.smartChat.endpoints.SmartChatEndpoint\n", + "import com.intellij.openapi.application.ApplicationManager\n", + "import com.intellij.openapi.extensions.ExtensionPointName\n", + "import com.intellij.openapi.project.Project\n", + "\n", + "class MyChatMessageHandler : ChatMessageHandler {\n", + " override fun isApplicable(project: Project, kind: ChatKind, userMessage: UserMessage) = true\n", + "\n", + " override fun createAnswerMessage(\n", + " project: Project,\n", + " chat: ChatSession,\n", + " userMessage: UserMessage,\n", + " kind: ChatKind,\n", + " ) = SimpleCompletableMessage(chat)\n", + "\n", + " override suspend fun serveAnswerMessage(\n", + " project: Project,\n", + " chat: ChatSession,\n", + " answerMessage: ChatMessage,\n", + " smartChatEndpoints: List,\n", + " ) = ChatAgent.EP_NAME.extensionList.last().serveAnswerMessage(project, chat, answerMessage)\n", + "}\n", + "\n", + "registerExtension(ChatMessageHandler.EP, MyChatMessageHandler())" + ], + "outputs": [], + "execution_count": 6 + }, + { + "metadata": { + "ExecuteTime": { + "end_time": "2025-10-02T03:55:51.730786Z", + "start_time": "2025-10-02T03:55:51.654405Z" + } + }, + "cell_type": "code", + "source": [ + "import com.intellij.ml.llm.agents.ChatAgent\n", + "import com.intellij.ml.llm.core.chat.messages.ChatMessage\n", + "import com.intellij.ml.llm.core.chat.messages.impl.SimpleCompletableMessage\n", + "import com.intellij.ml.llm.core.chat.session.ChatSession\n", + "import com.intellij.ml.llm.privacy.trustedStringBuilders.privacyConst\n", + "import com.intellij.openapi.application.ApplicationManager\n", + "import com.intellij.openapi.project.Project\n", + "\n", + "class MyChatAgent : ChatAgent {\n", + "\n", + " override val id = \"groot\"\n", + " override val name = \"I am Groot\"\n", + "\n", + " override fun createAnswerMessage(\n", + " project: Project,\n", + " chat: ChatSession,\n", + " userMessage: UserMessage,\n", + " kind: ChatKind\n", + " ): ChatMessage = SimpleCompletableMessage(chat)\n", + "\n", + " override suspend fun serveAnswerMessage(\n", + " project: Project,\n", + " chat: ChatSession,\n", + " answerMessage: ChatMessage,\n", + " ) {\n", + " if (answerMessage !is SimpleCompletableMessage) return\n", + " answerMessage.appendText(\"*I am Groot*\".privacyConst)\n", + " }\n", + "}\n", + "\n", + "registerExtension(ChatAgent.EP_NAME, MyChatAgent())" + ], + "outputs": [ + { + "ename": "org.jetbrains.kotlinx.jupyter.exceptions.ReplCompilerException", + "evalue": "at Cell In[7], line 9, column 1: Class 'MyChatAgent' is not abstract and does not implement abstract member public abstract val icon: Icon? defined in com.intellij.ml.llm.agents.ChatAgent", + "output_type": "error", + "traceback": [ + "org.jetbrains.kotlinx.jupyter.exceptions.ReplCompilerException: at Cell In[7], line 9, column 1: Class 'MyChatAgent' is not abstract and does not implement abstract member public abstract val icon: Icon? defined in com.intellij.ml.llm.agents.ChatAgent", + "\tat org.jetbrains.kotlinx.jupyter.repl.impl.JupyterCompilerImpl.compileSync(JupyterCompilerImpl.kt:151)", + "\tat org.jetbrains.kotlinx.jupyter.repl.impl.InternalEvaluatorImpl.eval(InternalEvaluatorImpl.kt:126)", + "\tat org.jetbrains.kotlinx.jupyter.repl.impl.CellExecutorImpl.execute_W38Nk0s$lambda$0$1(CellExecutorImpl.kt:95)", + "\tat org.jetbrains.kotlinx.jupyter.repl.impl.ReplForJupyterImpl.withHost(ReplForJupyterImpl.kt:730)", + "\tat org.jetbrains.kotlinx.jupyter.repl.impl.CellExecutorImpl.execute-W38Nk0s(CellExecutorImpl.kt:93)", + "\tat org.jetbrains.kotlinx.jupyter.repl.execution.CellExecutor.execute-W38Nk0s$default(CellExecutor.kt:14)", + "\tat org.jetbrains.kotlinx.jupyter.repl.impl.ReplForJupyterImpl.evaluateUserCode-wNURfNM(ReplForJupyterImpl.kt:591)", + "\tat org.jetbrains.kotlinx.jupyter.repl.impl.ReplForJupyterImpl.evalExImpl(ReplForJupyterImpl.kt:472)", + "\tat org.jetbrains.kotlinx.jupyter.repl.impl.ReplForJupyterImpl.evalEx$lambda$0(ReplForJupyterImpl.kt:466)", + "\tat org.jetbrains.kotlinx.jupyter.repl.impl.ReplForJupyterImpl.withEvalContext(ReplForJupyterImpl.kt:448)", + "\tat org.jetbrains.kotlinx.jupyter.repl.impl.ReplForJupyterImpl.evalEx(ReplForJupyterImpl.kt:465)", + "\tat org.jetbrains.kotlinx.jupyter.messaging.IdeCompatibleMessageRequestProcessor.processExecuteRequest$lambda$0$0$0(IdeCompatibleMessageRequestProcessor.kt:161)", + "\tat org.jetbrains.kotlinx.jupyter.streams.NonBlockingSubstitutionEngine.withDataSubstitution(SubstitutionEngine.kt:124)", + "\tat org.jetbrains.kotlinx.jupyter.streams.StreamSubstitutionManager.withSubstitutedStreams(StreamSubstitutionManager.kt:118)", + "\tat org.jetbrains.kotlinx.jupyter.messaging.IdeCompatibleMessageRequestProcessor.withForkedIn(IdeCompatibleMessageRequestProcessor.kt:351)", + "\tat org.jetbrains.kotlinx.jupyter.messaging.IdeCompatibleMessageRequestProcessor.evalWithIO$lambda$0$0(IdeCompatibleMessageRequestProcessor.kt:364)", + "\tat org.jetbrains.kotlinx.jupyter.streams.NonBlockingSubstitutionEngine.withDataSubstitution(SubstitutionEngine.kt:124)", + "\tat org.jetbrains.kotlinx.jupyter.streams.StreamSubstitutionManager.withSubstitutedStreams(StreamSubstitutionManager.kt:118)", + "\tat org.jetbrains.kotlinx.jupyter.messaging.IdeCompatibleMessageRequestProcessor.withForkedErr(IdeCompatibleMessageRequestProcessor.kt:341)", + "\tat org.jetbrains.kotlinx.jupyter.messaging.IdeCompatibleMessageRequestProcessor.evalWithIO$lambda$0(IdeCompatibleMessageRequestProcessor.kt:363)", + "\tat org.jetbrains.kotlinx.jupyter.streams.NonBlockingSubstitutionEngine.withDataSubstitution(SubstitutionEngine.kt:124)", + "\tat org.jetbrains.kotlinx.jupyter.streams.StreamSubstitutionManager.withSubstitutedStreams(StreamSubstitutionManager.kt:118)", + "\tat org.jetbrains.kotlinx.jupyter.messaging.IdeCompatibleMessageRequestProcessor.withForkedOut(IdeCompatibleMessageRequestProcessor.kt:334)", + "\tat org.jetbrains.kotlinx.jupyter.messaging.IdeCompatibleMessageRequestProcessor.evalWithIO(IdeCompatibleMessageRequestProcessor.kt:362)", + "\tat org.jetbrains.kotlinx.jupyter.messaging.IdeCompatibleMessageRequestProcessor.processExecuteRequest$lambda$0$0(IdeCompatibleMessageRequestProcessor.kt:160)", + "\tat org.jetbrains.kotlinx.jupyter.execution.JupyterExecutorImpl$Task.execute(JupyterExecutorImpl.kt:41)", + "\tat org.jetbrains.kotlinx.jupyter.execution.JupyterExecutorImpl.executorThread$lambda$0(JupyterExecutorImpl.kt:81)", + "\tat kotlin.concurrent.ThreadsKt$thread$thread$1.run(Thread.kt:30)", + "" + ] + } + ], + "execution_count": 7 + }, + { + "metadata": {}, + "cell_type": "markdown", + "source": [ + "To use this newly registered chat agent, click *AI Chat* on the right toolbar to open the AI Assistant chat tool window.\n", + "\n", + "In the input field, type your question, for example `Who are you?`. AI Assistant will use the newly registered Chat Agent and provide the reply: `I am Groot`." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Kotlin", + "language": "kotlin", + "name": "kotlin" + }, + "language_info": { + "name": "kotlin", + "version": "1.9.23", + "mimetype": "text/x-kotlin", + "file_extension": ".kt", + "pygments_lexer": "kotlin", + "codemirror_mode": "text/x-kotlin", + "nbconvert_exporter": "" + }, + "ktnbPluginMetadata": { + "sessionRunMode": "IDE_PROCESS" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} diff --git a/examples/ai_assistant/mcp_client.ipynb b/examples/ai_assistant/mcp_client.ipynb new file mode 100644 index 0000000..c5a4c25 --- /dev/null +++ b/examples/ai_assistant/mcp_client.ipynb @@ -0,0 +1,325 @@ +{ + "cells": [ + { + "metadata": {}, + "cell_type": "markdown", + "source": [ + "# JetBrains AI Assistant MCP Client Registry\n", + "\n", + "This notebook demonstrates how to use a registry service to manage Model Context Protocol (MCP) server configurations at runtime for the [JetBrains AI Assistant](https://plugins.jetbrains.com/plugin/22282-jetbrains-ai-assistant) plugin. It shows how to add, update, and remove configurations that control how individual MCP servers are launched.\n", + "\n", + "## What this demo shows\n", + "\n", + "- How to get the registry service instance\n", + "- How to register a new MCP server configuration\n", + "- How to update an existing configuration\n", + "- How to remove a configuration" + ] + }, + { + "metadata": { + "executionRelatedData": { + "compiledClasses": [ + "Line_1_jupyter", + "Line_2_jupyter", + "Line_2_jupyter", + "Line_3_jupyter" + ] + }, + "ExecuteTime": { + "end_time": "2025-10-02T07:37:31.557228Z", + "start_time": "2025-10-02T07:37:31.152560Z" + } + }, + "cell_type": "code", + "source": "%use intellij-platform", + "outputs": [ + { + "data": { + "text/plain": [ + "IntelliJ Platform integration is loaded" + ] + }, + "metadata": {}, + "output_type": "display_data", + "jetTransient": { + "display_id": null + } + } + ], + "execution_count": 1 + }, + { + "metadata": {}, + "cell_type": "markdown", + "source": "Load the JetBrains AI Assistant bundled plugin using its `com.intellij.ml.llm` ID." + }, + { + "metadata": { + "executionRelatedData": { + "compiledClasses": [ + "Line_4_jupyter", + "Line_5_jupyter", + "Line_6_jupyter", + "Line_7_jupyter", + "Line_8_jupyter", + "Line_4_jupyter", + "Line_5_jupyter", + "Line_6_jupyter", + "Line_7_jupyter", + "Line_8_jupyter" + ] + }, + "ExecuteTime": { + "end_time": "2025-10-02T07:37:34.082746Z", + "start_time": "2025-10-02T07:37:33.892063Z" + } + }, + "cell_type": "code", + "source": "loadPlugins(\"com.intellij.ml.llm\")", + "outputs": [], + "execution_count": 2 + }, + { + "metadata": { + "executionRelatedData": { + "compiledClasses": [ + "Line_9_jupyter", + "Line_9_jupyter", + "Line_13_jupyter", + "Line_9_jupyter", + "Line_16_jupyter", + "Line_13_jupyter", + "Line_17_jupyter", + "Line_16_jupyter", + "Line_19_jupyter", + "Line_17_jupyter", + "Line_21_jupyter", + "Line_19_jupyter" + ] + }, + "ExecuteTime": { + "end_time": "2025-10-02T07:53:40.180627Z", + "start_time": "2025-10-02T07:53:40.057501Z" + } + }, + "cell_type": "code", + "source": [ + "import com.intellij.ml.llm.mcp.client.McpClientRegistry\n", + "import com.intellij.ml.llm.mcp.client.McpClientRegistryException\n", + "import com.intellij.ml.llm.mcp.client.McpServerCommandDescriptor\n", + "\n", + "val registry = McpClientRegistry.getInstance()\n", + "var id: Int = -1\n", + "\n", + "try {\n", + " // 1) Create and register a new MCP server configuration\n", + " // The configuration is persisted and created in a disabled state.\n", + " // A notification will be shown to the user,\n", + " // and the user will be able to enable the server via the notification action.\n", + " val initialDescriptor = McpServerCommandDescriptor(\n", + " name = \"demo-mcp-server\",\n", + " command = \"npx\",\n", + " arguments = listOf(\"@modelcontextprotocol/new-server\"),\n", + " environmentVariables = emptyMap(),\n", + " workingDirectory = \"\" // optional working directory, leave empty to inherit project dir or process default\n", + " )\n", + "\n", + " // Optionally pass `project` to register a project-level configuration (instead of application-level)\n", + " id = registry.addConfiguration(initialDescriptor, /*project = project */)\n", + "}\n", + "catch (e: McpClientRegistryException) {\n", + " // Handle registry-specific issues (e.g., duplicate names, invalid id)\n", + " // You may log the error, show a notification, or surface it to the user as appropriate.\n", + "}" + ], + "outputs": [], + "execution_count": 15 + }, + { + "metadata": { + "executionRelatedData": { + "compiledClasses": [ + "Line_15_jupyter", + "Line_20_jupyter", + "Line_15_jupyter" + ] + }, + "ExecuteTime": { + "end_time": "2025-10-02T07:51:22.487335Z", + "start_time": "2025-10-02T07:51:22.432267Z" + } + }, + "cell_type": "code", + "source": [ + "try {\n", + " // 2) Update the configuration if you need to change the startup command or its arguments\n", + " val updatedDescriptor = McpServerCommandDescriptor(\n", + " name = \"demo-mcp-server\",\n", + " command = \"npx\",\n", + " arguments = listOf(\"@modelcontextprotocol/new-server\", \"additional_argument\"),\n", + " environmentVariables = emptyMap(),\n", + " workingDirectory = \"\"\n", + " )\n", + " registry.updateConfiguration(id, updatedDescriptor)\n", + "}\n", + "catch (e: McpClientRegistryException) {\n", + "}" + ], + "outputs": [], + "execution_count": 14 + }, + { + "metadata": { + "ExecuteTime": { + "end_time": "2025-10-02T07:50:26.227968Z", + "start_time": "2025-10-02T07:50:26.198484Z" + }, + "executionRelatedData": { + "compiledClasses": [ + "Line_18_jupyter" + ] + } + }, + "cell_type": "code", + "source": [ + "try {\n", + " // 3) Remove the configuration when it's no longer needed\n", + " registry.removeConfiguration(id)\n", + "}\n", + "catch (e: McpClientRegistryException) {\n", + "}" + ], + "outputs": [], + "execution_count": 12 + }, + { + "metadata": {}, + "cell_type": "markdown", + "source": [ + "## API reference: MCP registry service and descriptors\n", + "\n", + "Public interfaces and data types for registering, updating, and removing MCP server configurations across application- and project-level scopes. This API is experimental.\n", + "\n", + "```kotlin\n", + "package com.intellij.ml.llm.mcp.client\n", + "\n", + "import com.intellij.openapi.application.ApplicationManager\n", + "import com.intellij.openapi.components.service\n", + "import com.intellij.openapi.project.Project\n", + "import org.jetbrains.annotations.ApiStatus\n", + "\n", + "/**\n", + " * An application service for managing Model Context Protocol (MCP) client configurations.\n", + " *\n", + " * Provides CRUD operations on persisted MCP server configurations for either:\n", + " * - application-level scope when [project] is null, or\n", + " * - project-level scope when a [Project] is provided.\n", + " *\n", + " * This API is experimental and may change in future versions.\n", + " */\n", + "@ApiStatus.Experimental\n", + "interface McpClientRegistry {\n", + " companion object {\n", + " @JvmStatic\n", + " fun getInstance(): McpClientRegistry = ApplicationManager.getApplication().service()\n", + " }\n", + "\n", + " /**\n", + " * Returns all persisted MCP server configurations for the given scope.\n", + " *\n", + " * @param project target project for project-level configurations; pass null for application-level configurations\n", + " * @return a map of configuration id to [McpServerDescriptor]\n", + " */\n", + " fun getConfigurations(project: Project? = null): Map\n", + "\n", + " /**\n", + " * Adds a new MCP server configuration to the specified scope.\n", + " *\n", + " * The configuration is persisted and created in the disabled state. If a configuration with the same name\n", + " * already exists in the scope, the operation fails.\n", + " *\n", + " * @param descriptor server descriptor to persist\n", + " * @param project target project for project-level persistence; pass null for application-level persistence\n", + " * @return the generated unique configuration id\n", + " * @throws McpClientRegistryException if validation fails or a configuration with the same name already exists\n", + " */\n", + " @Throws(McpClientRegistryException::class)\n", + " fun addConfiguration(descriptor: McpServerDescriptor, project: Project? = null): Int\n", + "\n", + " /**\n", + " * Updates an existing MCP server configuration in the specified scope.\n", + " *\n", + " * If the configuration type differs from the existing one, an exception is thrown and no changes are applied.\n", + " *\n", + " * @param id id of the configuration to update\n", + " * @param descriptor new values for the configuration\n", + " * @param project target project for project-level scope; pass null for application-level scope\n", + " * @throws McpClientRegistryException if the configuration with the given id does not exist, validation fails,\n", + " * or the descriptor type differs from the existing configuration type\n", + " */\n", + " @Throws(McpClientRegistryException::class)\n", + " fun updateConfiguration(id: Int, descriptor: McpServerDescriptor, project: Project? = null)\n", + "\n", + " /**\n", + " * Removes an existing MCP server configuration from the specified scope.\n", + " *\n", + " * @param id id of the configuration to remove\n", + " * @param project target project for project-level scope; pass null for application-level scope\n", + " * @throws McpClientRegistryException if the configuration with the given id does not exist\n", + " */\n", + " @Throws(McpClientRegistryException::class)\n", + " fun removeConfiguration(id: Int, project: Project? = null)\n", + "}\n", + "```\n", + "\n", + "`McpServerDescriptor`:\n", + "```kotlin\n", + "package com.intellij.ml.llm.mcp.client\n", + "\n", + "import org.jetbrains.annotations.ApiStatus\n", + "\n", + "@ApiStatus.Experimental\n", + "sealed class McpServerDescriptor(open val name: String)\n", + "\n", + "@ApiStatus.Experimental\n", + "data class McpServerCommandDescriptor(\n", + " override val name: String,\n", + " val command: String,\n", + " val arguments: List,\n", + " val environmentVariables: Map,\n", + " val workingDirectory: String,\n", + ") : McpServerDescriptor(name)\n", + "\n", + "@ApiStatus.Experimental\n", + "data class McpServerUrlDescriptor(\n", + " override val name: String,\n", + " val url: String,\n", + " val headers: Map?,\n", + ") : McpServerDescriptor(name)\n", + "```" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Kotlin", + "language": "kotlin", + "name": "kotlin" + }, + "language_info": { + "name": "kotlin", + "version": "1.9.23", + "mimetype": "text/x-kotlin", + "file_extension": ".kt", + "pygments_lexer": "kotlin", + "codemirror_mode": "text/x-kotlin", + "nbconvert_exporter": "" + }, + "ktnbPluginMetadata": { + "sessionRunMode": "IDE_PROCESS" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} diff --git a/examples/aia_extensibility.ipynb b/examples/aia_extensibility.ipynb deleted file mode 100644 index 745ca7a..0000000 --- a/examples/aia_extensibility.ipynb +++ /dev/null @@ -1,140 +0,0 @@ -{ - "cells": [ - { - "metadata": {}, - "cell_type": "markdown", - "source": [ - "# JetBrains AI Assistant Extensibility\n", - "\n", - "This notebook presents how to extend the [JetBrains AI Assistant](https://plugins.jetbrains.com/plugin/22282-jetbrains-ai-assistant) plugin in the runtime by providing a custom message handler." - ] - }, - { - "metadata": {}, - "cell_type": "code", - "source": "%use intellij-platform", - "outputs": [], - "execution_count": null - }, - { - "metadata": {}, - "cell_type": "markdown", - "source": "Load the JetBrains AI Assistant bundled plugin using its `com.intellij.ml.llm` ID." - }, - { - "metadata": {}, - "cell_type": "code", - "source": "loadPlugins(\"com.intellij.ml.llm\")", - "outputs": [], - "execution_count": null - }, - { - "metadata": {}, - "cell_type": "code", - "source": [ - "import com.intellij.ml.llm.agents.ChatAgent\n", - "import com.intellij.ml.llm.core.chat.messages.ChatMessage\n", - "import com.intellij.ml.llm.core.chat.messages.UserMessage\n", - "import com.intellij.ml.llm.core.chat.messages.impl.SimpleCompletableMessage\n", - "import com.intellij.ml.llm.core.chat.session.ChatKind\n", - "import com.intellij.ml.llm.core.chat.session.ChatMessageHandler\n", - "import com.intellij.ml.llm.core.chat.session.ChatSession\n", - "import com.intellij.ml.llm.smartChat.endpoints.SmartChatEndpoint\n", - "import com.intellij.openapi.application.ApplicationManager\n", - "import com.intellij.openapi.extensions.ExtensionPointName\n", - "import com.intellij.openapi.project.Project\n", - "\n", - "class MyChatMessageHandler : ChatMessageHandler {\n", - " override fun isApplicable(project: Project, kind: ChatKind, userMessage: UserMessage) = true\n", - "\n", - " override fun createAnswerMessage(\n", - " project: Project,\n", - " chat: ChatSession,\n", - " userMessage: UserMessage,\n", - " kind: ChatKind,\n", - " ) = SimpleCompletableMessage(chat)\n", - "\n", - " override suspend fun serveAnswerMessage(\n", - " project: Project,\n", - " chat: ChatSession,\n", - " answerMessage: ChatMessage,\n", - " smartChatEndpoints: List,\n", - " ) = ChatAgent.EP_NAME.extensionList.last().serveAnswerMessage(project, chat, answerMessage)\n", - "}\n", - "\n", - "registerExtension(ChatMessageHandler.EP, MyChatMessageHandler())" - ], - "outputs": [], - "execution_count": null - }, - { - "metadata": {}, - "cell_type": "code", - "source": [ - "import com.intellij.ml.llm.agents.ChatAgent\n", - "import com.intellij.ml.llm.core.chat.messages.ChatMessage\n", - "import com.intellij.ml.llm.core.chat.messages.impl.SimpleCompletableMessage\n", - "import com.intellij.ml.llm.core.chat.session.ChatSession\n", - "import com.intellij.ml.llm.privacy.trustedStringBuilders.privacyConst\n", - "import com.intellij.openapi.application.ApplicationManager\n", - "import com.intellij.openapi.project.Project\n", - "\n", - "class MyChatAgent : ChatAgent {\n", - "\n", - " override val id = \"groot\"\n", - " override val name = \"I am Groot\"\n", - "\n", - " override fun createAnswerMessage(\n", - " project: Project,\n", - " chat: ChatSession,\n", - " userMessage: UserMessage,\n", - " kind: ChatKind\n", - " ): ChatMessage = SimpleCompletableMessage(chat)\n", - "\n", - " override suspend fun serveAnswerMessage(\n", - " project: Project,\n", - " chat: ChatSession,\n", - " answerMessage: ChatMessage,\n", - " ) {\n", - " if (answerMessage !is SimpleCompletableMessage) return\n", - " answerMessage.appendText(\"*I am Groot*\".privacyConst)\n", - " }\n", - "}\n", - "\n", - "registerExtension(ChatAgent.EP_NAME, MyChatAgent())" - ], - "outputs": [], - "execution_count": null - }, - { - "metadata": {}, - "cell_type": "markdown", - "source": [ - "To use this newly registered chat agent, click *AI Chat* on the right toolbar to open the AI Assistant chat tool window.\n", - "\n", - "In the input field, type your question, for example `Who are you?`. AI Assistant will use the newly registered Chat Agent and provide the reply: `I am Groot`." - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Kotlin", - "language": "kotlin", - "name": "kotlin" - }, - "language_info": { - "name": "kotlin", - "version": "1.9.23", - "mimetype": "text/x-kotlin", - "file_extension": ".kt", - "pygments_lexer": "kotlin", - "codemirror_mode": "text/x-kotlin", - "nbconvert_exporter": "" - }, - "ktnbPluginMetadata": { - "sessionRunMode": "IDE_PROCESS" - } - }, - "nbformat": 4, - "nbformat_minor": 0 -}