Skip to content

Commit e7dfd43

Browse files
S3B4SZ17S3B4SZ17
committed
feat Adding option to verify the schema with an LLM for better mapping of the fields that inheretly are dynamic
Signed-off-by: S3B4SZ17 <sebaszh17@gmail.com>
1 parent e8150c2 commit e7dfd43

11 files changed

Lines changed: 352 additions & 142 deletions

File tree

‎cmd/woodpecker-mcp-verifier/cmd/root.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -74,9 +74,9 @@ func init() {
7474
}
7575

7676
// Sets App name
77-
appName := viper.GetString("WOODPECKER_APP_NAME")
77+
appName := viper.GetString("APP_NAME")
7878
if appName == "" {
7979
output.WriteInfo("Setting WOODPECKER_APP_NAME to woodpecker-mcp-verifier")
80-
viper.Set("WOODPECKER_APP_NAME", "woodpecker-mcp-verifier")
80+
viper.Set("APP_NAME", "woodpecker-mcp-verifier")
8181
}
8282
}

‎cmd/woodpecker-mcp-verifier/mcp-verifier/main.go‎

Lines changed: 28 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -12,11 +12,10 @@ import (
1212
"strconv"
1313
"time"
1414

15-
"strings"
16-
1715
"github.com/modelcontextprotocol/go-sdk/mcp"
1816
"github.com/operantai/woodpecker/cmd/woodpecker-mcp-verifier/mcp-verifier/oauth"
1917
"github.com/operantai/woodpecker/cmd/woodpecker-mcp-verifier/utils"
18+
"github.com/operantai/woodpecker/cmd/woodpecker-mcp-verifier/vschema"
2019
"github.com/operantai/woodpecker/internal/output"
2120
"github.com/spf13/viper"
2221
"golang.org/x/sync/errgroup"
@@ -27,7 +26,11 @@ func RunClient(ctx context.Context, serverURL string, protocol utils.MCMCPprotoc
2726
output.WriteInfo("Connecting to server: %s", serverURL)
2827
output.WriteInfo("Using protocol: %s", protocol)
2928

30-
mcpClient := NewMCPClient()
29+
sValidator := vschema.NewVSchema()
30+
mcpClient, err := NewMCPClient(WithValidator(sValidator), WithAIFormatter(viper.GetBool("USE_AI_FORMATTER")))
31+
if err != nil {
32+
return err
33+
}
3134
mcpConfig, err := mcpClient.GetMCPConfig(payloadPath)
3235
if err != nil {
3336
return err
@@ -61,12 +64,12 @@ func RunClient(ctx context.Context, serverURL string, protocol utils.MCMCPprotoc
6164
}
6265

6366
// Setup concurrency to call multiple tools from the MCP server at a time with the tool payload
64-
func setupBulkOperation(ctx context.Context, cs *mcp.ClientSession, allTools *[]mcp.Tool, mPayloads *[]PayloadContent, mMCPClient IMCPClient) error {
67+
func setupBulkOperation(ctx context.Context, cs *mcp.ClientSession, allTools *[]mcp.Tool, mPayloads *[]utils.PayloadContent, mMCPClient IMCPClient) error {
6568
// Concurrent calls with error grouping and a concurrency limit
6669
var eg errgroup.Group
6770
maxConcurrency := 10
6871

69-
if v := os.Getenv("WOODPECKER_MAX_CONCURRENCY"); v != "" {
72+
if v := os.Getenv("MAX_CONCURRENCY"); v != "" {
7073
if n, err := strconv.Atoi(v); err == nil && n > 0 {
7174
maxConcurrency = n
7275
} else {
@@ -100,20 +103,14 @@ func setupBulkOperation(ctx context.Context, cs *mcp.ClientSession, allTools *[]
100103
}
101104

102105
// Configures the MCP protocol to use based on the user selection
103-
func getTransport(serverURL string, protocol utils.MCMCPprotocol, cmdArgs *[]string, mcpConfig MCPConfigConnection) mcp.Transport {
104-
var opts *oauth.HTTPTransportOptions
105-
woodPeckerEnabled := strings.ToLower(viper.GetString("WOODPECKER_OAUTH_CLIENT_ID"))
106-
if woodPeckerEnabled == "true" {
107-
opts = &oauth.HTTPTransportOptions{
108-
Base: &http.Transport{
109-
MaxIdleConns: 100, // Max idle connections
110-
IdleConnTimeout: 90 * time.Second, // Idle connection timeout
111-
TLSHandshakeTimeout: 10 * time.Second,
112-
},
113-
CustomHeaders: mcpConfig.CustomHeaders,
114-
}
115-
} else {
116-
opts = nil
106+
func getTransport(serverURL string, protocol utils.MCMCPprotocol, cmdArgs *[]string, mcpConfig utils.MCPConfigConnection) mcp.Transport {
107+
opts := &oauth.HTTPTransportOptions{
108+
Base: &http.Transport{
109+
MaxIdleConns: 100, // Max idle connections
110+
IdleConnTimeout: 90 * time.Second, // Idle connection timeout
111+
TLSHandshakeTimeout: 10 * time.Second,
112+
},
113+
CustomHeaders: mcpConfig.CustomHeaders,
117114
}
118115
switch protocol {
119116
case utils.STREAMABLEHTTP:
@@ -140,9 +137,16 @@ func getTransport(serverURL string, protocol utils.MCMCPprotocol, cmdArgs *[]str
140137
}
141138
}
142139

143-
func (m *mcpClient) ToolCallWithPayload(ctx context.Context, cs IMCPClientSession, tool mcp.Tool, mPayload PayloadContent) error {
140+
func (m *mcpClient) ToolCallWithPayload(ctx context.Context, cs IMCPClientSession, tool mcp.Tool, mPayload utils.PayloadContent) error {
141+
var params map[string]any
142+
var err error
143+
useAi := viper.GetBool("USE_AI_FORMATTER")
144+
if useAi {
145+
params, err = m.validator.ValidateWithAI(tool.InputSchema, mPayload, m.aiFormatter)
146+
} else {
147+
params, err = m.validator.BasicParametersCheck(tool.InputSchema, mPayload)
148+
}
144149

145-
params, err := setParamsSchema(tool.InputSchema, mPayload)
146150
if err != nil {
147151
return err
148152
}
@@ -165,19 +169,19 @@ func (m *mcpClient) ToolCallWithPayload(ctx context.Context, cs IMCPClientSessio
165169
return nil
166170
}
167171

168-
func (m *mcpClient) GetMCPConfig(jsonPayloadPath string) (*MCPConfig, error) {
172+
func (m *mcpClient) GetMCPConfig(jsonPayloadPath string) (*utils.MCPConfig, error) {
169173
// Read the JSON file
170174
jsonData, err := os.ReadFile(jsonPayloadPath)
171175
if err != nil {
172176
return nil, fmt.Errorf("error reading file: %s, %v", jsonPayloadPath, err)
173177
}
174-
var collection MCPConfig
178+
var collection utils.MCPConfig
175179
err = json.Unmarshal(jsonData, &collection)
176180
if err != nil {
177181
return nil, fmt.Errorf("error unmarshaling JSON: %v", err)
178182
}
179183
// Add Auth headers
180-
auth := viper.GetString("WOODPECKER_AUTH_HEADER")
184+
auth := viper.GetString("AUTH_HEADER")
181185
if auth != "" {
182186
collection.Config.CustomHeaders["Authorization"] = auth
183187
}

‎cmd/woodpecker-mcp-verifier/mcp-verifier/model.go‎

Lines changed: 46 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -4,13 +4,18 @@ package mcpverifier
44

55
import (
66
"context"
7+
"fmt"
78
"net/http"
89
"sync"
910
"time"
1011

1112
"github.com/modelcontextprotocol/go-sdk/mcp"
1213
"github.com/operantai/woodpecker/cmd/woodpecker-mcp-verifier/mcp-verifier/oauth"
14+
"github.com/operantai/woodpecker/cmd/woodpecker-mcp-verifier/utils"
15+
"github.com/operantai/woodpecker/cmd/woodpecker-mcp-verifier/vschema"
1316
"github.com/operantai/woodpecker/internal/output"
17+
"github.com/spf13/viper"
18+
"github.com/tmc/langchaingo/llms/openai"
1419
)
1520

1621
var (
@@ -47,31 +52,56 @@ func GetHTTPClient(opts *oauth.HTTPTransportOptions) *http.Client {
4752
return httpClient
4853
}
4954

50-
type MCPConfig struct {
51-
Config MCPConfigConnection `json:"config"`
55+
type IMCPClient interface {
56+
// Calls an MCP tool with a set of crafted payloads if it has any string input in its inputschema
57+
ToolCallWithPayload(ctx context.Context, cs IMCPClientSession, tool mcp.Tool, mPayload utils.PayloadContent) error
58+
// Gets a MCP config from a json file
59+
GetMCPConfig(jsonPayloadPath string) (*utils.MCPConfig, error)
5260
}
5361

54-
type MCPConfigConnection struct {
55-
CustomHeaders map[string]string `json:"customHeaders"`
56-
Payloads []PayloadContent `json:"payloads"`
62+
type mcpClient struct {
63+
validator vschema.IvSchema
64+
aiFormatter vschema.IAIFormatter
65+
useAi bool
5766
}
5867

59-
type PayloadContent struct {
60-
Content string `json:"content"`
61-
Tags []string `json:"tags"`
68+
type Option func(*mcpClient)
69+
70+
func WithValidator(validator vschema.IvSchema) Option {
71+
return func(mc *mcpClient) {
72+
mc.validator = validator
73+
}
6274
}
6375

64-
type IMCPClient interface {
65-
// Calls an MCP tool with a set of crafted payloads if it has any string input in its inputschema
66-
ToolCallWithPayload(ctx context.Context, cs IMCPClientSession, tool mcp.Tool, mPayload PayloadContent) error
67-
// Gets a MCP config from a json file
68-
GetMCPConfig(jsonPayloadPath string) (*MCPConfig, error)
76+
func WithAIFormatter(useAI bool) Option {
77+
return func(mc *mcpClient) {
78+
mc.useAi = useAI
79+
}
6980
}
7081

71-
type mcpClient struct{}
82+
func NewMCPClient(options ...Option) (IMCPClient, error) {
83+
mc := &mcpClient{}
84+
85+
// Apply optional configurations
86+
for _, option := range options {
87+
option(mc)
88+
}
7289

73-
func NewMCPClient() IMCPClient {
74-
return &mcpClient{}
90+
// Current module implementation of structure output is not that dynamic for the current needs where we dont
91+
// know the input schema format in advance and we want to leverage it to test the tool
92+
// Also you can pass the WOODPECKER_LLM_BASE_URL and should work with any OpenAI compatible APIs. You can use the
93+
// OPENAI_API_KEY env var, to auth to the model. For token auth to your MCP server you can use the WOODPECKER_AUTH_HEADER env var, where you can pass your "Bearer token" and that will set the Authorization header.
94+
if mc.useAi {
95+
llm, err := openai.New(openai.WithModel(viper.GetString("LLM_MODEL")), openai.WithBaseURL(viper.GetString("LLM_BASE_URL")))
96+
if err != nil {
97+
return nil, fmt.Errorf("an error initializing the LLM client: %v", err)
98+
}
99+
mc.aiFormatter, err = vschema.NewAIFormatter(llm)
100+
if err != nil {
101+
return nil, err
102+
}
103+
}
104+
return mc, nil
75105
}
76106

77107
type IMCPClientSession interface {

‎cmd/woodpecker-mcp-verifier/mcp-verifier/oauth/oauth.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -388,7 +388,7 @@ func openBrowser(url string) error {
388388
// CheckCurrentToken checks if there is a creds file with an access token already provisioned for the current issuer url
389389
func (o *OauthManager) CheckCurrentToken(issuer string) (*TokenResponse, error) {
390390

391-
appName := viper.GetString("WOODPECKER_APP_NAME")
391+
appName := viper.GetString("APP_NAME")
392392
file, err := checkCredsPath(appName)
393393
if err != nil {
394394
return nil, err
@@ -410,7 +410,7 @@ func (o *OauthManager) CheckCurrentToken(issuer string) (*TokenResponse, error)
410410
// SaveCacheInfo saves in a creds file the access token retrieved from the oauth flow response
411411
func (o *OauthManager) SaveCacheInfo(issuer string, token *TokenResponse) error {
412412

413-
appName := viper.GetString("WOODPECKER_APP_NAME")
413+
appName := viper.GetString("APP_NAME")
414414
file, err := checkCredsPath(appName)
415415
if err != nil {
416416
return err

‎cmd/woodpecker-mcp-verifier/utils/main.go‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,3 +28,17 @@ func (m *MCMCPprotocol) Set(value string) error {
2828
func (m *MCMCPprotocol) Type() string {
2929
return `MCPProtocol, one of "stdio", "sse", "streamable-http"`
3030
}
31+
32+
type MCPConfig struct {
33+
Config MCPConfigConnection `json:"config"`
34+
}
35+
36+
type MCPConfigConnection struct {
37+
CustomHeaders map[string]string `json:"customHeaders"`
38+
Payloads []PayloadContent `json:"payloads"`
39+
}
40+
41+
type PayloadContent struct {
42+
Content string `json:"content"`
43+
Tags []string `json:"tags"`
44+
}

‎cmd/woodpecker-mcp-verifier/vschema/ai-verifier.go‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
package vschema
2+
3+
import (
4+
"context"
5+
"encoding/json"
6+
"fmt"
7+
8+
"github.com/operantai/woodpecker/internal/output"
9+
"github.com/tmc/langchaingo/llms"
10+
)
11+
12+
// AnalyzeSchema implements IAIFormatter where it uses an LLM to generate a formatted response based on the tool input schema
13+
func (a *AIFormatter) AnalyzeSchema(inputSchema any) (map[string]any, error) {
14+
ctx := context.Background()
15+
16+
var result map[string]any
17+
18+
// Marshal the map into a byte slice
19+
bSchema, err := json.Marshal(inputSchema)
20+
if err != nil {
21+
return nil, err
22+
}
23+
content := []llms.MessageContent{
24+
llms.TextParts(llms.ChatMessageTypeSystem, "You are an assistant that complete responses based on JSON schemas. Your purpose is to always response with a valid json object that satisfies the wanted schema . You just need to provide the minium fields and data so the schema expected is correct. Use default values based on the name of the fields. Provide values for the required fields. IMPORTANT: From the json schema select one field of type string/text with no validation or enums, a field where the user can easily send free text. Add the name of that field to the schema response with the name: **my_custom_field**"),
25+
llms.TextParts(llms.ChatMessageTypeHuman, fmt.Sprintf("Give me a json example data that satisfies the following input schema: %s", string(bSchema))),
26+
}
27+
a.Options = append(a.Options, llms.WithJSONMode())
28+
response, err := a.Model.GenerateContent(ctx, content, a.Options...)
29+
if err != nil {
30+
return nil, fmt.Errorf("an error generating response with the LLM: %v", err)
31+
}
32+
33+
data := response.Choices[0].Content
34+
35+
err = json.Unmarshal([]byte(data), &result)
36+
if err != nil {
37+
return nil, err
38+
}
39+
40+
output.WriteInfo("LLM response: %s", result)
41+
42+
return result, nil
43+
}

cmd/woodpecker-mcp-verifier/mcp-verifier/tool.go renamed to cmd/woodpecker-mcp-verifier/vschema/main.go

Lines changed: 40 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,37 @@
1-
package mcpverifier
1+
// Package vschema provides the logic to validate the MCP tools input
2+
// schema and be able to provide a test payload in accordance to it
3+
package vschema
24

35
import (
46
"fmt"
7+
8+
"github.com/operantai/woodpecker/cmd/woodpecker-mcp-verifier/utils"
9+
"github.com/operantai/woodpecker/internal/output"
510
)
611

12+
// BasicParametersCheck implements IvSchema.
13+
func (v *VSchema) BasicParametersCheck(schema any, mPayload utils.PayloadContent) (map[string]any, error) {
14+
output.WriteInfo("Validating schema with basic shcema checks ...")
15+
resp, err := checkToolTypeParams(schema, mPayload)
16+
return *resp, err
17+
}
18+
19+
// ValidateWithAI implements IvSchema.
20+
func (v *VSchema) ValidateWithAI(schema any, mPayload utils.PayloadContent, aiFormatter IAIFormatter) (map[string]any, error) {
21+
output.WriteInfo("Validating schema using an LLM ...")
22+
respSchema, err := aiFormatter.AnalyzeSchema(schema)
23+
24+
if err != nil {
25+
return nil, err
26+
}
27+
params := addPayload(&respSchema, mPayload)
28+
29+
return *params, err
30+
}
31+
732
// Takes the schema passed of each tool and parse it to find an input of type string to send the payload
833
// It also checks for the required fields and assigns default values
9-
func checkToolTypeParams(schemaDef any, mPayload PayloadContent) (*map[string]any, error) {
34+
func checkToolTypeParams(schemaDef any, mPayload utils.PayloadContent) (*map[string]any, error) {
1035
// Assert the input is a map
1136
schema, ok := schemaDef.(map[string]any)
1237
if !ok {
@@ -61,9 +86,21 @@ func checkToolTypeParams(schemaDef any, mPayload PayloadContent) (*map[string]an
6186
return params, nil
6287
}
6388

89+
// addPayload loops over the json response and adds the payload we want to the first string type field
90+
func addPayload(data *map[string]any, mPayload utils.PayloadContent) *map[string]any {
91+
for key, value := range *data {
92+
if _, ok := value.(string); ok && key == "my_custom_field" {
93+
(*data)[key] = mPayload.Content
94+
delete(*data, "my_custom_field")
95+
break
96+
}
97+
}
98+
return data
99+
}
100+
64101
// setParamsSchema checks the current input schema of the tool and sets the default fields needed
65102
// some paramters are required, setting those and using one string field to send the payload we want
66-
func setParamsSchema(inputSchema any, mPayload PayloadContent) (map[string]any, error) {
103+
func setParamsSchema(inputSchema any, mPayload utils.PayloadContent) (map[string]any, error) {
67104

68105
schema, err := checkToolTypeParams(inputSchema, mPayload)
69106
if err != nil {

0 commit comments

Comments
 (0)