internal/mcp: simplify prompt API
Remove NewServerPrompt. Schema inference is irrelevant for prompts. All arguments are strings. So using a struct is not as valuable as with tools. The only advantage is catching misspellings at compile time. It's not worth it. Change-Id: I45aaa5d55c8940cc9e32dc277bd543f79d932519 Reviewed-on: https://go-review.googlesource.com/c/tools/+/682797 Reviewed-by: Robert Findley <rfindley@google.com> LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
This commit is contained in:
Jonathan Amsterdam
Robert Findley
parent
e228d78d19
commit
a6ec3f5f4e
5 changed files with 37 additions and 177 deletions
|
|
@ -32,11 +32,11 @@ func SayHi(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParam
|
|||
|
||||
// TODO(jba): it should be OK for args to be a pointer, but this fails in
|
||||
// jsonschema. Needs investigation.
|
||||
func PromptHi(ctx context.Context, ss *mcp.ServerSession, args HiArgs, _ *mcp.GetPromptParams) (*mcp.GetPromptResult, error) {
|
||||
func PromptHi(ctx context.Context, ss *mcp.ServerSession, params *mcp.GetPromptParams) (*mcp.GetPromptResult, error) {
|
||||
return &mcp.GetPromptResult{
|
||||
Description: "Code review prompt",
|
||||
Messages: []*mcp.PromptMessage{
|
||||
{Role: "user", Content: mcp.NewTextContent("Say hi to " + args.Name)},
|
||||
{Role: "user", Content: mcp.NewTextContent("Say hi to " + params.Arguments["name"])},
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
|
@ -48,7 +48,10 @@ func main() {
|
|||
server.AddTools(mcp.NewServerTool("greet", "say hi", SayHi, mcp.Input(
|
||||
mcp.Property("name", mcp.Description("the name to say hi to")),
|
||||
)))
|
||||
server.AddPrompts(mcp.NewServerPrompt("greet", "", PromptHi))
|
||||
server.AddPrompts(&mcp.ServerPrompt{
|
||||
Prompt: &mcp.Prompt{Name: "greet"},
|
||||
Handler: PromptHi,
|
||||
})
|
||||
server.AddResources(&mcp.ServerResource{
|
||||
Resource: &mcp.Resource{
|
||||
Name: "info",
|
||||
|
|
|
|||
|
|
@ -88,9 +88,9 @@ func TestList(t *testing.T) {
|
|||
})
|
||||
|
||||
t.Run("prompts", func(t *testing.T) {
|
||||
promptA := mcp.NewServerPrompt("apple", "apple prompt", testPromptHandler[struct{}])
|
||||
promptB := mcp.NewServerPrompt("banana", "banana prompt", testPromptHandler[struct{}])
|
||||
promptC := mcp.NewServerPrompt("cherry", "cherry prompt", testPromptHandler[struct{}])
|
||||
promptA := newServerPrompt("apple", "apple prompt")
|
||||
promptB := newServerPrompt("banana", "banana prompt")
|
||||
promptC := newServerPrompt("cherry", "cherry prompt")
|
||||
wantPrompts := []*mcp.Prompt{promptA.Prompt, promptB.Prompt, promptC.Prompt}
|
||||
prompts := []*mcp.ServerPrompt{promptA, promptB, promptC}
|
||||
server.AddPrompts(prompts...)
|
||||
|
|
@ -122,3 +122,15 @@ func testIterator[T any](ctx context.Context, t *testing.T, seq iter.Seq2[*T, er
|
|||
t.Fatalf("mismatch (-want +got):\n%s", diff)
|
||||
}
|
||||
}
|
||||
|
||||
// testPromptHandler is used for type inference newServerPrompt.
|
||||
func testPromptHandler(context.Context, *mcp.ServerSession, *mcp.GetPromptParams) (*mcp.GetPromptResult, error) {
|
||||
panic("not implemented")
|
||||
}
|
||||
|
||||
func newServerPrompt(name, desc string) *mcp.ServerPrompt {
|
||||
return &mcp.ServerPrompt{
|
||||
Prompt: &mcp.Prompt{Name: name, Description: desc},
|
||||
Handler: testPromptHandler,
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -441,18 +441,27 @@ var (
|
|||
}
|
||||
|
||||
prompts = map[string]*ServerPrompt{
|
||||
"code_review": NewServerPrompt("code_review", "do a code review",
|
||||
func(_ context.Context, _ *ServerSession, params struct{ Code string }, _ *GetPromptParams) (*GetPromptResult, error) {
|
||||
"code_review": {
|
||||
Prompt: &Prompt{
|
||||
Name: "code_review",
|
||||
Description: "do a code review",
|
||||
Arguments: []*PromptArgument{{Name: "Code", Required: true}},
|
||||
},
|
||||
Handler: func(_ context.Context, _ *ServerSession, params *GetPromptParams) (*GetPromptResult, error) {
|
||||
return &GetPromptResult{
|
||||
Description: "Code review prompt",
|
||||
Messages: []*PromptMessage{
|
||||
{Role: "user", Content: NewTextContent("Please review the following code: " + params.Code)},
|
||||
{Role: "user", Content: NewTextContent("Please review the following code: " + params.Arguments["Code"])},
|
||||
},
|
||||
}, nil
|
||||
}),
|
||||
"fail": NewServerPrompt("fail", "", func(_ context.Context, _ *ServerSession, args struct{}, _ *GetPromptParams) (*GetPromptResult, error) {
|
||||
return nil, errTestFailure
|
||||
}),
|
||||
},
|
||||
},
|
||||
"fail": {
|
||||
Prompt: &Prompt{Name: "fail"},
|
||||
Handler: func(_ context.Context, _ *ServerSession, _ *GetPromptParams) (*GetPromptResult, error) {
|
||||
return nil, errTestFailure
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
resource1 = &Resource{
|
||||
|
|
|
|||
111
mcp/prompt.go
111
mcp/prompt.go
|
|
@ -6,13 +6,6 @@ package mcp
|
|||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"reflect"
|
||||
"slices"
|
||||
|
||||
"github.com/modelcontextprotocol/go-sdk/internal/util"
|
||||
"github.com/modelcontextprotocol/go-sdk/jsonschema"
|
||||
)
|
||||
|
||||
// A PromptHandler handles a call to prompts/get.
|
||||
|
|
@ -23,107 +16,3 @@ type ServerPrompt struct {
|
|||
Prompt *Prompt
|
||||
Handler PromptHandler
|
||||
}
|
||||
|
||||
// NewServerPrompt is a helper that uses reflection to create a prompt for the
|
||||
// given handler.
|
||||
//
|
||||
// The arguments for the prompt are extracted from the request type for the
|
||||
// handler. The handler request type must be a struct consisting only of fields
|
||||
// of type string or *string. The argument names for the resulting prompt
|
||||
// definition correspond to the JSON names of the request fields, and any
|
||||
// fields that are not marked "omitempty" are considered required.
|
||||
//
|
||||
// The handler is passed [GetPromptParams] so it can have access to prompt
|
||||
// parameters other than name and arguments. At present, there are no such
|
||||
// parameters.
|
||||
func NewServerPrompt[In any](name, description string, handler func(context.Context, *ServerSession, In, *GetPromptParams) (*GetPromptResult, error), opts ...PromptOption) *ServerPrompt {
|
||||
schema, err := jsonschema.For[In]()
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
if schema.Type != "object" || !reflect.DeepEqual(schema.AdditionalProperties, &jsonschema.Schema{Not: &jsonschema.Schema{}}) {
|
||||
panic(fmt.Sprintf("invalid input type %q: handler input type must be a struct", schema.Type))
|
||||
}
|
||||
resolved, err := schema.Resolve(nil)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
prompt := &ServerPrompt{
|
||||
Prompt: &Prompt{
|
||||
Name: name,
|
||||
Description: description,
|
||||
},
|
||||
}
|
||||
required := make(map[string]bool)
|
||||
for _, p := range schema.Required {
|
||||
required[p] = true
|
||||
}
|
||||
for name, prop := range util.Sorted(schema.Properties) {
|
||||
if prop.Type != "string" {
|
||||
panic(fmt.Sprintf("invalid property type %q: handler type must consist only of string fields", prop.Type))
|
||||
}
|
||||
prompt.Prompt.Arguments = append(prompt.Prompt.Arguments, &PromptArgument{
|
||||
Name: name,
|
||||
Description: prop.Description,
|
||||
Required: required[name],
|
||||
})
|
||||
}
|
||||
prompt.Handler = func(ctx context.Context, ss *ServerSession, params *GetPromptParams) (*GetPromptResult, error) {
|
||||
// For simplicity, just marshal and unmarshal the arguments.
|
||||
// This could be avoided in the future.
|
||||
rawArgs, err := json.Marshal(params.Arguments)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var args In
|
||||
if err := unmarshalSchema(rawArgs, resolved, &args); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return handler(ctx, ss, args, params)
|
||||
}
|
||||
for _, opt := range opts {
|
||||
opt.set(prompt)
|
||||
}
|
||||
return prompt
|
||||
}
|
||||
|
||||
// A PromptOption configures the behavior of a Prompt.
|
||||
type PromptOption interface {
|
||||
set(*ServerPrompt)
|
||||
}
|
||||
|
||||
type promptSetter func(*ServerPrompt)
|
||||
|
||||
func (s promptSetter) set(p *ServerPrompt) { s(p) }
|
||||
|
||||
// Argument configures the 'schema' of a prompt argument.
|
||||
// If the argument does not exist, it is added.
|
||||
//
|
||||
// Since prompt arguments are not a full JSON schema, Argument only accepts
|
||||
// Required and Description, and panics when encountering any other option.
|
||||
func Argument(name string, opts ...SchemaOption) PromptOption {
|
||||
return promptSetter(func(p *ServerPrompt) {
|
||||
i := slices.IndexFunc(p.Prompt.Arguments, func(arg *PromptArgument) bool {
|
||||
return arg.Name == name
|
||||
})
|
||||
var arg *PromptArgument
|
||||
if i < 0 {
|
||||
i = len(p.Prompt.Arguments)
|
||||
arg = &PromptArgument{Name: name}
|
||||
p.Prompt.Arguments = append(p.Prompt.Arguments, arg)
|
||||
} else {
|
||||
arg = p.Prompt.Arguments[i]
|
||||
}
|
||||
for _, opt := range opts {
|
||||
switch v := opt.(type) {
|
||||
case required:
|
||||
arg.Required = bool(v)
|
||||
case description:
|
||||
arg.Description = string(v)
|
||||
default:
|
||||
panic(fmt.Sprintf("unsupported prompt argument schema option %T", opt))
|
||||
}
|
||||
}
|
||||
p.Prompt.Arguments[i] = arg
|
||||
})
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,53 +0,0 @@
|
|||
// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
|
||||
// Use of this source code is governed by an MIT-style
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
package mcp_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/google/go-cmp/cmp"
|
||||
"github.com/modelcontextprotocol/go-sdk/mcp"
|
||||
)
|
||||
|
||||
// testPromptHandler is used for type inference in TestNewServerPrompt.
|
||||
func testPromptHandler[T any](context.Context, *mcp.ServerSession, T, *mcp.GetPromptParams) (*mcp.GetPromptResult, error) {
|
||||
panic("not implemented")
|
||||
}
|
||||
|
||||
func TestNewServerPrompt(t *testing.T) {
|
||||
tests := []struct {
|
||||
prompt *mcp.ServerPrompt
|
||||
want []*mcp.PromptArgument
|
||||
}{
|
||||
{
|
||||
mcp.NewServerPrompt("empty", "", testPromptHandler[struct{}]),
|
||||
nil,
|
||||
},
|
||||
{
|
||||
mcp.NewServerPrompt("add_arg", "", testPromptHandler[struct{}], mcp.Argument("x")),
|
||||
[]*mcp.PromptArgument{{Name: "x"}},
|
||||
},
|
||||
{
|
||||
mcp.NewServerPrompt("combo", "", testPromptHandler[struct {
|
||||
Name string `json:"name"`
|
||||
Country string `json:"country,omitempty"`
|
||||
State string
|
||||
}],
|
||||
mcp.Argument("name", mcp.Description("the person's name")),
|
||||
mcp.Argument("State", mcp.Required(false))),
|
||||
[]*mcp.PromptArgument{
|
||||
{Name: "State"},
|
||||
{Name: "country"},
|
||||
{Name: "name", Required: true, Description: "the person's name"},
|
||||
},
|
||||
},
|
||||
}
|
||||
for _, test := range tests {
|
||||
if diff := cmp.Diff(test.want, test.prompt.Prompt.Arguments); diff != "" {
|
||||
t.Errorf("NewServerPrompt(%v) mismatch (-want +got):\n%s", test.prompt.Prompt.Name, diff)
|
||||
}
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue