Function ExampleAddTool_complexSchema

mcp/tool_example_test.go:191–254

!-weathertool

()

// !-weathertool

func ExampleAddTool_complexSchema() {
	// This example demonstrates a tool with a more 'realistic' input and output
	// schema. We use a combination of techniques to tune our input and output
	// schemas.

	// !+customschemas

	// Distinguished Go types allow custom schemas to be reused during inference.
	customSchemas := map[reflect.Type]*jsonschema.Schema{
		reflect.TypeFor[Probability](): {Type: "number", Minimum: jsonschema.Ptr(0.0), Maximum: jsonschema.Ptr(1.0)},
		reflect.TypeFor[WeatherType](): {Type: "string", Enum: []any{Sunny, PartlyCloudy, Cloudy, Rainy, Snowy}},
	}
	opts := &jsonschema.ForOptions{TypeSchemas: customSchemas}
	in, err := jsonschema.For[WeatherInput](opts)
	if err != nil {
		log.Fatal(err)
	}

	// Furthermore, we can tweak the inferred schema, in this case limiting
	// forecasts to 0-10 days.
	daysSchema := in.Properties["days"]
	daysSchema.Minimum = jsonschema.Ptr(0.0)
	daysSchema.Maximum = jsonschema.Ptr(10.0)

	// Output schema inference can reuse our custom schemas from input inference.
	out, err := jsonschema.For[WeatherOutput](opts)
	if err != nil {
		log.Fatal(err)
	}

	// Now add our tool to a server. Since we've customized the schemas, we need
	// to override the default schema inference.
	server := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, nil)
	mcp.AddTool(server, &mcp.Tool{
		Name:         "weather",
		InputSchema:  in,
		OutputSchema: out,
	}, WeatherTool)

	// !-customschemas

	ctx := context.Background()
	session, err := connect(ctx, server) // create an in-memory connection
	if err != nil {
		log.Fatal(err)
	}
	defer session.Close()

	// Check that the client observes the correct schemas.
	for t, err := range session.Tools(ctx, nil) {
		if err != nil {
			log.Fatal(err)
		}
		// Formatting the entire schemas would be too much output.
		// Just check that our customizations were effective.
		fmt.Println("max days:", jsonPath(t.InputSchema, "properties", "days", "maximum"))
		fmt.Println("max confidence:", jsonPath(t.OutputSchema, "properties", "confidence", "maximum"))
		fmt.Println("weather types:", jsonPath(t.OutputSchema, "properties", "dailyForecast", "items", "properties", "type", "enum"))