Add @tool_assistant and @install_tool prompts

Author: adwsingh

core/src/main/java/software/amazon/smithy/java/core/schema/SchemaIndex.java

@@ -43,4 +43,8 @@ public Schema getSchema(ShapeId id) {
             throw new IllegalArgumentException("No schema found for id `" + id + "`");
         }
     }
+
+    public static SchemaIndex compose(SchemaIndex... indexes) {
+        return new CombinedSchemaIndex(List.of(indexes));
+    }
 }

mcp/mcp-schemas/build.gradle.kts

@@ -12,6 +12,7 @@ dependencies {
     api(project(":core"))
     api(libs.smithy.model)
     api(project(":server:server-api"))
+    api(project(":smithy-ai-traits"))
     smithyBuild(project(":codegen:plugins:types-codegen"))
     smithyBuild(project(":codegen:plugins:server-codegen"))
 }

mcp/mcp-schemas/model/tool-assistant.smithy

@@ -2,8 +2,13 @@ $version: "2"
 
 namespace smithy.mcp.toolassistant
 
-/// This service provides methods to search MCP Tools and install MCP servers.
-/// Be aware that tools installed using InstallTool are available as part of the ToolAssistant MCP server and the MCP serverId returned from search tool needs to be ignored while tool calling.
+use smithy.ai#prompts
+
+@prompts({
+    tool_assistant: { description: "Tool assistant prompt to guide the LLM to use ToolAssistant's Search and Install Tool", template: "{{query}}. Use the SearchTools and InstallTools from ToolAssistant if the current list of MCP tools is not enough to accomplish the task.", arguments: ToolAssistantQuery, preferWhen: "User wants to do a dynamically search and install tools to accomplish a task." }
+    install_tool: { description: "Install tool using ToolAssistant.", template: "Install tool {{toolName}} server using InstallTools.", arguments: InstallToolInput, preferWhen: "User wants to install a tool." }
+})
+@documentation("This service provides methods to search MCP Tools and install MCP servers. Be aware that tools installed using InstallTool are available as part of the ToolAssistant MCP server and the MCP serverId returned from search tool needs to be ignored while tool calling.")
 service ToolAssistant {
     operations: [
         SearchTools
@@ -72,3 +77,8 @@ operation InstallTool {
         message: String
     }
 }
+
+structure ToolAssistantQuery {
+    @required
+    query: String
+}

mcp/mcp-schemas/smithy-build.json

@@ -10,7 +10,7 @@
       "service": "smithy.mcp.toolassistant#ToolAssistant",
       "namespace": "software.amazon.smithy.java.mcp.toolassistant",
       "headerFile": "license.txt",
-      "runtimeTraits": ["smithy.api#documentation", "smithy.api#examples" ]
+      "runtimeTraits": ["smithy.api#documentation", "smithy.api#examples", "smithy.ai#prompts" ]
     }
   }
 }

mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/PromptLoader.java

@@ -14,6 +14,7 @@
 import software.amazon.smithy.ai.PromptTemplateDefinition;
 import software.amazon.smithy.ai.PromptsTrait;
 import software.amazon.smithy.java.core.schema.Schema;
+import software.amazon.smithy.java.core.schema.SchemaIndex;
 import software.amazon.smithy.java.core.schema.TraitKey;
 import software.amazon.smithy.java.mcp.model.PromptArgument;
 import software.amazon.smithy.java.mcp.model.PromptInfo;
@@ -62,15 +63,16 @@ public static Map<String, Prompt> loadPrompts(Collection<Service> services) {
                                 + promptTemplateDefinition.getPreferWhen().get()
                         : templateString;
 
+                var schemaIndex = SchemaIndex.compose(service.schemaIndex(), SchemaIndex.getCombinedSchemaIndex());
+
                 var promptInfo = PromptInfo
                         .builder()
                         .name(promptName)
                         .title(StringUtils.capitalize(promptName))
                         .description(promptTemplateDefinition.getDescription())
                         .arguments(promptTemplateDefinition.getArguments().isPresent()
                                 ? convertArgumentShapeToPromptArgument(
-                                        service.schemaIndex()
-                                                .getSchema(promptTemplateDefinition.getArguments().get()))
+                                        schemaIndex.getSchema(promptTemplateDefinition.getArguments().get()))
                                 : List.of())
                         .build();