Skip to content

Commit d473d39

Browse files
findleyrfindleyr
committed
WIP: adding a larger tool schema for benchmarking
1 parent 845c29f commit d473d39

File tree

3 files changed

+235
-0
lines changed

3 files changed

+235
-0
lines changed

‎examples/server/weather/main.go‎

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,96 @@
1+
// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
2+
// Use of this source code is governed by an MIT-style
3+
// license that can be found in the LICENSE file.
4+
5+
package main
6+
7+
import (
8+
"context"
9+
"log"
10+
"time"
11+
12+
"github.com/google/jsonschema-go/jsonschema"
13+
"github.com/modelcontextprotocol/go-sdk/mcp"
14+
)
15+
16+
type WeatherInput struct {
17+
Location Location `json:"location" jsonschema:"user location"`
18+
Days int `json:"days" jsonschema:"number of days to forecast"`
19+
}
20+
21+
type Location struct {
22+
Name string `json:"name"`
23+
Latitude *float64 `json:"latitude,omitempty"`
24+
Longitude *float64 `json:"longitude,omitempty"`
25+
}
26+
27+
type Forecast struct {
28+
Forecast string `json:"forecast" jsonschema:"description of the day's weather"`
29+
Type WeatherType `json:"type" jsonschema:"type of weather"`
30+
Rain float64 `json:"rain" jsonschema:"probability of rain, between 0 and 1"`
31+
High float64 `json:"high" jsonschema:"high temperature"`
32+
Low float64 `json:"low" jsonschema:"low temperature"`
33+
}
34+
35+
type WeatherType string
36+
37+
const (
38+
Sunny WeatherType = "sun"
39+
PartlyCloudy WeatherType = "partly_cloudy"
40+
Cloudy WeatherType = "clouds"
41+
Rainy WeatherType = "rain"
42+
Snowy WeatherType = "snow"
43+
)
44+
45+
type Probability float64
46+
47+
type WeatherOutput struct {
48+
Summary string `json:"summary" jsonschema:"a summary of the weather forecast"`
49+
Confidence Probability `json:"confidence" jsonschema:"confidence, between 0 and 1"`
50+
AsOf time.Time `json:"asOf" jsonschema:"the time the weather was computed"`
51+
DailyForecast []Forecast `json:"dailyForecast" jsonschema:"the daily forecast"`
52+
Source string `json:"source,omitempty" jsonschema:"the organization providing the weather forecast"`
53+
}
54+
55+
func WeatherTool(ctx context.Context, req *mcp.CallToolRequest, in WeatherInput) (*mcp.CallToolResult, WeatherOutput, error) {
56+
perfectWeather := WeatherOutput{
57+
Summary: "perfect",
58+
Confidence: 1.0,
59+
AsOf: time.Now(),
60+
}
61+
for range in.Days {
62+
perfectWeather.DailyForecast = append(perfectWeather.DailyForecast, Forecast{
63+
Forecast: "another perfect day",
64+
Type: Sunny,
65+
Rain: 0.0,
66+
High: 72.0,
67+
Low: 72.0,
68+
})
69+
}
70+
return nil, perfectWeather, nil
71+
}
72+
73+
func main() {
74+
customSchemas := map[any]*jsonschema.Schema{
75+
Probability(0): {Type: "number", Minimum: jsonschema.Ptr(0.0), Maximum: jsonschema.Ptr(1.0)},
76+
WeatherType(""): {Type: "string", Enum: []any{Sunny, PartlyCloudy, Cloudy, Rainy, Snowy}},
77+
}
78+
opts := &jsonschema.ForOptions{TypeSchemas: customSchemas}
79+
in, err := jsonschema.For[WeatherInput](opts)
80+
if err != nil {
81+
log.Fatal(err)
82+
}
83+
out, err := jsonschema.For[WeatherOutput](opts)
84+
if err != nil {
85+
log.Fatal(err)
86+
}
87+
88+
server := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, nil)
89+
mcp.AddTool(server, &mcp.Tool{
90+
Name: "weather",
91+
InputSchema: in,
92+
OutputSchema: out,
93+
}, WeatherTool)
94+
95+
server.Run(context.Background(), &mcp.StdioTransport{})
96+
}

‎mcp/streamable_bench_test.go‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
2+
// Use of this source code is governed by an MIT-style
3+
// license that can be found in the LICENSE file.
4+
5+
package mcp_test
6+
7+
import (
8+
"testing"
9+
)
10+
11+
func BenchmarkStreamableServing(b *testing.B) {
12+
// server := mcp.NewServer(testImpl, nil)
13+
// mcp.AddTool(server, &mcp.Tool{Name: "greet", Description: "say hi"}, sayHi)
14+
}

‎mcp/tool_example_test.go‎

Lines changed: 125 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,125 @@
1+
// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
2+
// Use of this source code is governed by an MIT-style
3+
// license that can be found in the LICENSE file.
4+
5+
package mcp_test
6+
7+
import (
8+
"context"
9+
"encoding/json"
10+
"fmt"
11+
"log"
12+
"time"
13+
14+
"github.com/google/jsonschema-go/jsonschema"
15+
"github.com/modelcontextprotocol/go-sdk/mcp"
16+
)
17+
18+
type WeatherInput struct {
19+
Location Location `json:"location" jsonschema:"user location"`
20+
Days int `json:"days" jsonschema:"number of days to forecast"`
21+
}
22+
23+
type Location struct {
24+
Name string `json:"name"`
25+
Latitude *float64 `json:"latitude,omitempty"`
26+
Longitude *float64 `json:"longitude,omitempty"`
27+
}
28+
29+
type Forecast struct {
30+
Forecast string `json:"forecast" jsonschema:"description of the day's weather"`
31+
Type WeatherType `json:"type" jsonschema:"type of weather"`
32+
Rain float64 `json:"rain" jsonschema:"probability of rain, between 0 and 1"`
33+
High float64 `json:"high" jsonschema:"high temperature"`
34+
Low float64 `json:"low" jsonschema:"low temperature"`
35+
}
36+
37+
type WeatherType string
38+
39+
const (
40+
Sunny WeatherType = "sun"
41+
PartlyCloudy WeatherType = "partly_cloudy"
42+
Cloudy WeatherType = "clouds"
43+
Rainy WeatherType = "rain"
44+
Snowy WeatherType = "snow"
45+
)
46+
47+
type Probability float64
48+
49+
type WeatherOutput struct {
50+
Summary string `json:"summary" jsonschema:"a summary of the weather forecast"`
51+
Confidence Probability `json:"confidence" jsonschema:"confidence, between 0 and 1"`
52+
AsOf time.Time `json:"asOf" jsonschema:"the time the weather was computed"`
53+
DailyForecast []Forecast `json:"dailyForecast" jsonschema:"the daily forecast"`
54+
Source string `json:"source,omitempty" jsonschema:"the organization providing the weather forecast"`
55+
}
56+
57+
func WeatherTool(ctx context.Context, req *mcp.CallToolRequest, in WeatherInput) (*mcp.CallToolResult, WeatherOutput, error) {
58+
perfectWeather := WeatherOutput{
59+
Summary: "perfect",
60+
Confidence: 1.0,
61+
AsOf: time.Now(),
62+
}
63+
for range in.Days {
64+
perfectWeather.DailyForecast = append(perfectWeather.DailyForecast, Forecast{
65+
Forecast: "another perfect day",
66+
Type: Sunny,
67+
Rain: 0.0,
68+
High: 72.0,
69+
Low: 72.0,
70+
})
71+
}
72+
return nil, perfectWeather, nil
73+
}
74+
75+
func ExampleAddTool_complexSchema() {
76+
customSchemas := map[any]*jsonschema.Schema{
77+
Probability(0): {Type: "number", Minimum: jsonschema.Ptr(0.0), Maximum: jsonschema.Ptr(1.0)},
78+
WeatherType(""): {Type: "string", Enum: []any{Sunny, PartlyCloudy, Cloudy, Rainy, Snowy}},
79+
}
80+
opts := &jsonschema.ForOptions{TypeSchemas: customSchemas}
81+
in, err := jsonschema.For[WeatherInput](opts)
82+
if err != nil {
83+
log.Fatal(err)
84+
}
85+
out, err := jsonschema.For[WeatherOutput](opts)
86+
if err != nil {
87+
log.Fatal(err)
88+
}
89+
90+
server := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, nil)
91+
mcp.AddTool(server, &mcp.Tool{
92+
Name: "weather",
93+
InputSchema: in,
94+
OutputSchema: out,
95+
}, WeatherTool)
96+
97+
ctx := context.Background()
98+
session, err := connect(ctx, server) // create an in-memory connection
99+
if err != nil {
100+
log.Fatal(err)
101+
}
102+
defer session.Close()
103+
104+
for t, err := range session.Tools(ctx, nil) {
105+
if err != nil {
106+
log.Fatal(err)
107+
}
108+
schemaJSON, err := json.MarshalIndent(t.OutputSchema, "", "\t")
109+
if err != nil {
110+
log.Fatal(err)
111+
}
112+
fmt.Println(t.Name, string(schemaJSON))
113+
}
114+
// Output:
115+
// nothing
116+
}
117+
118+
func connect(ctx context.Context, server *mcp.Server) (*mcp.ClientSession, error) {
119+
t1, t2 := mcp.NewInMemoryTransports()
120+
if _, err := server.Connect(ctx, t1, nil); err != nil {
121+
return nil, err
122+
}
123+
client := mcp.NewClient(&mcp.Implementation{Name: "client", Version: "v0.0.1"}, nil)
124+
return client.Connect(ctx, t2, nil)
125+
}

0 commit comments

Comments
 (0)