1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66

package mcpserver

import (
	"context"

	"github.com/mark3labs/mcp-go/mcp"
	"github.com/mark3labs/mcp-go/server"
	"github.com/samber/lo"

	"{{{ .Package }}}/app"
	"{{{ .Package }}}/app/util"
)

type PromptHandler func(ctx context.Context, as *app.State, req mcp.GetPromptRequest, args util.ValueMap, logger util.Logger) (string, error)

type Prompt struct {
	Name        string `json:"name"`
	Description string `json:"description,omitzero"`
	Icon        string `json:"icon,omitzero"`
	Content     string `json:"content,omitzero"`
}

func NewStaticPrompt(name string, description string, content string) *Prompt {
	return &Prompt{Name: name, Description: description, Content: content}
}

func (p *Prompt) ToMCP() (mcp.Prompt, error) {
	return mcp.NewPrompt(p.Name, mcp.WithPromptDescription(p.Description)), nil
}

func (p *Prompt) IconSafe() string {
	return util.Choose(p.Icon == "", "gift", p.Icon)
}

func (p *Prompt) Handler(as *app.State, logger util.Logger) server.PromptHandlerFunc {
	return func(ctx context.Context, req mcp.GetPromptRequest) (*mcp.GetPromptResult, error) {
		args := util.ValueMap{}
		for k, v := range req.Params.Arguments {
			args[k] = v
		}
		return &mcp.GetPromptResult{
			Description: p.Description,
			Messages:    []mcp.PromptMessage{{Role: mcp.RoleUser, Content: mcp.TextContent{Type: "text", Text: p.Content}}},
		}, nil
	}
}

type Prompts []*Prompt

func (r Prompts) Get(n string) *Prompt {
	return lo.FindOrElse(r, nil, func(x *Prompt) bool {
		return x.Name == n
	})
}

func (s *Server) AddPrompts(as *app.State, logger util.Logger, prompts ...*Prompt) error {
	for _, r := range prompts {
		s.Prompts = append(s.Prompts, r)
		m, err := r.ToMCP()
		if err != nil {
			return err
		}
		s.MCP.AddPrompt(m, r.Handler(as, logger))
	}
	return nil
}