internal/mcp: support prompts
Add support for 'prompts/list' and 'prompts/get', and registering prompts on the server using AddPrompts. Add a 'MakePrompt' helper to construct prompts from a prompt handler using reflection. Change-Id: I479cb9c9b99313cca99640c1f5ea1939363759e6 Reviewed-on: https://go-review.googlesource.com/c/tools/+/669015 LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com> Reviewed-by: Jonathan Amsterdam <jba@google.com>
This commit is contained in:
Rob Findley
Robert Findley
parent
e082a340be
commit
af4d138732
2 changed files with 531 additions and 98 deletions
|
|
@ -13,13 +13,11 @@ package main
|
|||
|
||||
import (
|
||||
"bytes"
|
||||
"cmp"
|
||||
"encoding/json"
|
||||
"flag"
|
||||
"fmt"
|
||||
"go/format"
|
||||
"io"
|
||||
"iter"
|
||||
"log"
|
||||
"net/http"
|
||||
"os"
|
||||
|
|
@ -28,6 +26,7 @@ import (
|
|||
"strings"
|
||||
|
||||
"github.com/modelcontextprotocol/go-sdk/internal/jsonschema"
|
||||
"github.com/modelcontextprotocol/go-sdk/internal/util"
|
||||
)
|
||||
|
||||
var schemaFile = flag.String("schema_file", "", "if set, use this file as the persistent schema file")
|
||||
|
|
@ -54,31 +53,36 @@ var declarations = config{
|
|||
"CallToolRequest": {
|
||||
Fields: config{"Params": {Name: "CallToolParams"}},
|
||||
},
|
||||
"CallToolResult": {
|
||||
Name: "CallToolResult",
|
||||
},
|
||||
"CallToolResult": {Name: "CallToolResult"},
|
||||
"CancelledNotification": {
|
||||
Fields: config{"Params": {Name: "CancelledParams"}},
|
||||
},
|
||||
"ClientCapabilities": {Name: "ClientCapabilities"},
|
||||
"Implementation": {Name: "Implementation"},
|
||||
"GetPromptRequest": {
|
||||
Fields: config{"Params": {Name: "GetPromptParams"}},
|
||||
},
|
||||
"GetPromptResult": {Name: "GetPromptResult"},
|
||||
"Implementation": {Name: "Implementation"},
|
||||
"InitializeRequest": {
|
||||
Fields: config{"Params": {Name: "InitializeParams"}},
|
||||
},
|
||||
"InitializeResult": {
|
||||
Name: "InitializeResult",
|
||||
},
|
||||
"InitializeResult": {Name: "InitializeResult"},
|
||||
"InitializedNotification": {
|
||||
Fields: config{"Params": {Name: "InitializedParams"}},
|
||||
},
|
||||
"ListPromptsRequest": {
|
||||
Fields: config{"Params": {Name: "ListPromptsParams"}},
|
||||
},
|
||||
"ListPromptsResult": {Name: "ListPromptsResult"},
|
||||
"ListToolsRequest": {
|
||||
Fields: config{"Params": {Name: "ListToolsParams"}},
|
||||
},
|
||||
"ListToolsResult": {
|
||||
Name: "ListToolsResult",
|
||||
},
|
||||
"RequestId": {Substitute: "any"}, // null|number|string
|
||||
"Role": {Name: "Role"},
|
||||
"ListToolsResult": {Name: "ListToolsResult"},
|
||||
"Prompt": {Name: "Prompt"},
|
||||
"PromptMessage": {Name: "PromptMessage"},
|
||||
"PromptArgument": {Name: "PromptArgument"},
|
||||
"RequestId": {Substitute: "any"}, // null|number|string
|
||||
"Role": {Name: "Role"},
|
||||
"ServerCapabilities": {
|
||||
Name: "ServerCapabilities",
|
||||
Fields: config{
|
||||
|
|
@ -92,9 +96,7 @@ var declarations = config{
|
|||
Name: "Tool",
|
||||
Fields: config{"InputSchema": {Substitute: "*jsonschema.Schema"}},
|
||||
},
|
||||
"ToolAnnotations": {
|
||||
Name: "ToolAnnotations",
|
||||
},
|
||||
"ToolAnnotations": {Name: "ToolAnnotations"},
|
||||
}
|
||||
|
||||
func main() {
|
||||
|
|
@ -114,7 +116,7 @@ func main() {
|
|||
// writing types, we collect definitions and concatenate them later. This
|
||||
// also allows us to sort.
|
||||
named := make(map[string]*bytes.Buffer)
|
||||
for name, def := range sorted(schema.Definitions) {
|
||||
for name, def := range util.Sorted(schema.Definitions) {
|
||||
config := declarations[name]
|
||||
if config == nil {
|
||||
continue
|
||||
|
|
@ -142,7 +144,7 @@ import (
|
|||
`)
|
||||
|
||||
// Write out types.
|
||||
for _, b := range sorted(named) {
|
||||
for _, b := range util.Sorted(named) {
|
||||
fmt.Fprintln(buf)
|
||||
fmt.Fprint(buf, b.String())
|
||||
}
|
||||
|
|
@ -242,8 +244,8 @@ func writeType(w io.Writer, config *typeConfig, def *jsonschema.Schema, named ma
|
|||
// unmarshal them into a map[string]any, or delay unmarshalling with
|
||||
// json.RawMessage. For now, use json.RawMessage as it defers the choice.
|
||||
if def.Type == "object" && canHaveAdditionalProperties(def) {
|
||||
w.Write([]byte("json.RawMessage"))
|
||||
return nil
|
||||
w.Write([]byte("map[string]"))
|
||||
return writeType(w, nil, def.AdditionalProperties, named)
|
||||
}
|
||||
|
||||
if def.Type == "" {
|
||||
|
|
@ -269,7 +271,7 @@ func writeType(w io.Writer, config *typeConfig, def *jsonschema.Schema, named ma
|
|||
|
||||
case "object":
|
||||
fmt.Fprintf(w, "struct {\n")
|
||||
for name, fieldDef := range sorted(def.Properties) {
|
||||
for name, fieldDef := range util.Sorted(def.Properties) {
|
||||
if fieldDef.Description != "" {
|
||||
fmt.Fprintf(w, "%s\n", toComment(fieldDef.Description))
|
||||
}
|
||||
|
|
@ -385,28 +387,3 @@ func assert(cond bool, msg string) {
|
|||
panic(msg)
|
||||
}
|
||||
}
|
||||
|
||||
// Helpers below are copied from gopls' moremaps package.
|
||||
|
||||
// sorted returns an iterator over the entries of m in key order.
|
||||
func sorted[M ~map[K]V, K cmp.Ordered, V any](m M) iter.Seq2[K, V] {
|
||||
// TODO(adonovan): use maps.Sorted if proposal #68598 is accepted.
|
||||
return func(yield func(K, V) bool) {
|
||||
keys := keySlice(m)
|
||||
slices.Sort(keys)
|
||||
for _, k := range keys {
|
||||
if !yield(k, m[k]) {
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// keySlice returns the keys of the map M, like slices.Collect(maps.Keys(m)).
|
||||
func keySlice[M ~map[K]V, K comparable, V any](m M) []K {
|
||||
r := make([]K, 0, len(m))
|
||||
for k := range m {
|
||||
r = append(r, k)
|
||||
}
|
||||
return r
|
||||
}
|
||||
|
|
|
|||
|
|
@ -29,8 +29,8 @@ type Annotations struct {
|
|||
}
|
||||
|
||||
type CallToolParams struct {
|
||||
Arguments json.RawMessage `json:"arguments,omitempty"`
|
||||
Name string `json:"name"`
|
||||
Arguments map[string]json.RawMessage `json:"arguments,omitempty"`
|
||||
Name string `json:"name"`
|
||||
}
|
||||
|
||||
// The server's response to a tool call.
|
||||
|
|
@ -46,8 +46,8 @@ type CallToolParams struct {
|
|||
type CallToolResult struct {
|
||||
// This result property is reserved by the protocol to allow clients and servers
|
||||
// to attach additional metadata to their responses.
|
||||
Meta json.RawMessage `json:"_meta,omitempty"`
|
||||
Content []json.RawMessage `json:"content"`
|
||||
Meta map[string]json.RawMessage `json:"_meta,omitempty"`
|
||||
Content []json.RawMessage `json:"content"`
|
||||
// Whether the tool call ended in an error.
|
||||
//
|
||||
// If not set, this is assumed to be false (the call was successful).
|
||||
|
|
@ -70,14 +70,31 @@ type CancelledParams struct {
|
|||
// additional capabilities.
|
||||
type ClientCapabilities struct {
|
||||
// Experimental, non-standard capabilities that the client supports.
|
||||
Experimental json.RawMessage `json:"experimental,omitempty"`
|
||||
Experimental map[string]map[string]json.RawMessage `json:"experimental,omitempty"`
|
||||
// Present if the client supports listing roots.
|
||||
Roots *struct {
|
||||
// Whether the client supports notifications for changes to the roots list.
|
||||
ListChanged bool `json:"listChanged,omitempty"`
|
||||
} `json:"roots,omitempty"`
|
||||
// Present if the client supports sampling from an LLM.
|
||||
Sampling json.RawMessage `json:"sampling,omitempty"`
|
||||
Sampling map[string]json.RawMessage `json:"sampling,omitempty"`
|
||||
}
|
||||
|
||||
type GetPromptParams struct {
|
||||
// Arguments to use for templating the prompt.
|
||||
Arguments map[string]string `json:"arguments,omitempty"`
|
||||
// The name of the prompt or prompt template.
|
||||
Name string `json:"name"`
|
||||
}
|
||||
|
||||
// The server's response to a prompts/get request from the client.
|
||||
type GetPromptResult struct {
|
||||
// This result property is reserved by the protocol to allow clients and servers
|
||||
// to attach additional metadata to their responses.
|
||||
Meta map[string]json.RawMessage `json:"_meta,omitempty"`
|
||||
// An optional description for the prompt.
|
||||
Description string `json:"description,omitempty"`
|
||||
Messages []PromptMessage `json:"messages"`
|
||||
}
|
||||
|
||||
// Describes the name and version of an MCP implementation.
|
||||
|
|
@ -99,8 +116,8 @@ type InitializeParams struct {
|
|||
type InitializeResult struct {
|
||||
// This result property is reserved by the protocol to allow clients and servers
|
||||
// to attach additional metadata to their responses.
|
||||
Meta json.RawMessage `json:"_meta,omitempty"`
|
||||
Capabilities ServerCapabilities `json:"capabilities"`
|
||||
Meta map[string]json.RawMessage `json:"_meta,omitempty"`
|
||||
Capabilities ServerCapabilities `json:"capabilities"`
|
||||
// Instructions describing how to use the server and its features.
|
||||
//
|
||||
// This can be used by clients to improve the LLM's understanding of available
|
||||
|
|
@ -114,7 +131,24 @@ type InitializeResult struct {
|
|||
ServerInfo Implementation `json:"serverInfo"`
|
||||
}
|
||||
|
||||
type InitializedParams json.RawMessage
|
||||
type InitializedParams map[string]json.RawMessage
|
||||
|
||||
type ListPromptsParams struct {
|
||||
// An opaque token representing the current pagination position. If provided,
|
||||
// the server should return results starting after this cursor.
|
||||
Cursor string `json:"cursor,omitempty"`
|
||||
}
|
||||
|
||||
// The server's response to a prompts/list request from the client.
|
||||
type ListPromptsResult struct {
|
||||
// This result property is reserved by the protocol to allow clients and servers
|
||||
// to attach additional metadata to their responses.
|
||||
Meta map[string]json.RawMessage `json:"_meta,omitempty"`
|
||||
// An opaque token representing the pagination position after the last returned
|
||||
// result. If present, there may be more results available.
|
||||
NextCursor string `json:"nextCursor,omitempty"`
|
||||
Prompts []Prompt `json:"prompts"`
|
||||
}
|
||||
|
||||
type ListToolsParams struct {
|
||||
// An opaque token representing the current pagination position. If provided,
|
||||
|
|
@ -126,19 +160,48 @@ type ListToolsParams struct {
|
|||
type ListToolsResult struct {
|
||||
// This result property is reserved by the protocol to allow clients and servers
|
||||
// to attach additional metadata to their responses.
|
||||
Meta json.RawMessage `json:"_meta,omitempty"`
|
||||
Meta map[string]json.RawMessage `json:"_meta,omitempty"`
|
||||
// An opaque token representing the pagination position after the last returned
|
||||
// result. If present, there may be more results available.
|
||||
NextCursor string `json:"nextCursor,omitempty"`
|
||||
Tools []Tool `json:"tools"`
|
||||
}
|
||||
|
||||
// A prompt or prompt template that the server offers.
|
||||
type Prompt struct {
|
||||
// A list of arguments to use for templating the prompt.
|
||||
Arguments []PromptArgument `json:"arguments,omitempty"`
|
||||
// An optional description of what this prompt provides
|
||||
Description string `json:"description,omitempty"`
|
||||
// The name of the prompt or prompt template.
|
||||
Name string `json:"name"`
|
||||
}
|
||||
|
||||
// Describes an argument that a prompt can accept.
|
||||
type PromptArgument struct {
|
||||
// A human-readable description of the argument.
|
||||
Description string `json:"description,omitempty"`
|
||||
// The name of the argument.
|
||||
Name string `json:"name"`
|
||||
// Whether this argument must be provided.
|
||||
Required bool `json:"required,omitempty"`
|
||||
}
|
||||
|
||||
// Present if the server offers any prompt templates.
|
||||
type PromptCapabilities struct {
|
||||
// Whether this server supports notifications for changes to the prompt list.
|
||||
ListChanged bool `json:"listChanged,omitempty"`
|
||||
}
|
||||
|
||||
// Describes a message returned as part of a prompt.
|
||||
//
|
||||
// This is similar to `SamplingMessage`, but also supports the embedding of
|
||||
// resources from the MCP server.
|
||||
type PromptMessage struct {
|
||||
Content json.RawMessage `json:"content"`
|
||||
Role Role `json:"role"`
|
||||
}
|
||||
|
||||
// Present if the server offers any resources to read.
|
||||
type ResourceCapabilities struct {
|
||||
// Whether this server supports notifications for changes to the resource list.
|
||||
|
|
@ -155,11 +218,11 @@ type Role string
|
|||
// additional capabilities.
|
||||
type ServerCapabilities struct {
|
||||
// Present if the server supports argument autocompletion suggestions.
|
||||
Completions json.RawMessage `json:"completions,omitempty"`
|
||||
Completions map[string]json.RawMessage `json:"completions,omitempty"`
|
||||
// Experimental, non-standard capabilities that the server supports.
|
||||
Experimental json.RawMessage `json:"experimental,omitempty"`
|
||||
Experimental map[string]map[string]json.RawMessage `json:"experimental,omitempty"`
|
||||
// Present if the server supports sending log messages to the client.
|
||||
Logging json.RawMessage `json:"logging,omitempty"`
|
||||
Logging map[string]json.RawMessage `json:"logging,omitempty"`
|
||||
// Present if the server offers any prompt templates.
|
||||
Prompts *PromptCapabilities `json:"prompts,omitempty"`
|
||||
// Present if the server offers any resources to read.
|
||||
|
|
|
|||
Loading…
Reference in a new issue