Skip to content

Commit 954a0d8

Browse files
loretoloreto
authored
[ai] Use mcp jsonschema package (#541)
## Summary ## How was it tested? ## Community Contribution License All community contributions in this pull request are licensed to the project maintainers under the terms of the [Apache 2 License](https://www.apache.org/licenses/LICENSE-2.0). By creating this pull request I represent that I have the right to license the contributions to the project maintainers under the Apache 2 License as stated in the [Community Contribution License](https://github.com/jetify-com/opensource/blob/main/CONTRIBUTING.md#community-contribution-license).
1 parent e28f7d2 commit 954a0d8

18 files changed

+1263
-141
lines changed

‎aisdk/ai/api/jsonschema_helpers.go

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
package api
2+
3+
import (
4+
"github.com/modelcontextprotocol/go-sdk/jsonschema"
5+
)
6+
7+
// FalseSchema returns a new Schema that fails to validate any value.
8+
// This represents "additionalProperties": false in JSON Schema.
9+
// In JSON Schema 2020-12, false can also be represented as {"not": {}}
10+
// and that's what the `jsonschema` package does.
11+
func FalseSchema() *jsonschema.Schema {
12+
return &jsonschema.Schema{Not: &jsonschema.Schema{}}
13+
}

‎aisdk/ai/api/llm_call_options.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ import (
55
"fmt"
66
"net/http"
77

8-
jsonschema "github.com/sashabaranov/go-openai/jsonschema"
8+
"github.com/modelcontextprotocol/go-sdk/jsonschema"
99
"github.com/tidwall/gjson"
1010
)
1111

@@ -131,7 +131,7 @@ type ResponseFormat struct {
131131
Type string `json:"type"`
132132

133133
// Schema optionally provides a JSON schema to guide the model's output
134-
Schema *jsonschema.Definition `json:"schema,omitzero"`
134+
Schema *jsonschema.Schema `json:"schema,omitzero"`
135135

136136
// Name optionally provides a name for the output to guide the model
137137
Name string `json:"name,omitzero"`

‎aisdk/ai/api/llm_call_options_test.go

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ func TestCallOptions_JSON(t *testing.T) {
5454
"age": {"type": "integer"}
5555
},
5656
"required": ["name"],
57-
"additionalProperties": false
57+
"additionalProperties": {"not": {}}
5858
},
5959
"name": "user_info",
6060
"description": "User information structure"
@@ -72,7 +72,7 @@ func TestCallOptions_JSON(t *testing.T) {
7272
"units": {"type": "string", "enum": ["metric", "imperial"]}
7373
},
7474
"required": ["location"],
75-
"additionalProperties": false
75+
"additionalProperties": {"not": {}}
7676
}
7777
},
7878
{
@@ -144,7 +144,7 @@ func BenchmarkCallOptionsUnmarshal(b *testing.B) {
144144
"age": {"type": "integer"}
145145
},
146146
"required": ["name"],
147-
"additionalProperties": false
147+
"additionalProperties": {"not": {}}
148148
},
149149
"name": "user_info",
150150
"description": "User information structure"
@@ -162,7 +162,7 @@ func BenchmarkCallOptionsUnmarshal(b *testing.B) {
162162
"units": {"type": "string", "enum": ["metric", "imperial"]}
163163
},
164164
"required": ["location"],
165-
"additionalProperties": false
165+
"additionalProperties": {"not": {}}
166166
}
167167
},
168168
{

‎aisdk/ai/api/llm_tool.go

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,9 @@
11
package api
22

3-
// TODO: This schema package is pretty small.
4-
// It might be best to just in line it into our AI SDK.
53
import (
64
"encoding/json"
75

8-
"github.com/sashabaranov/go-openai/jsonschema"
6+
"github.com/modelcontextprotocol/go-sdk/jsonschema"
97
)
108

119
// ToolChoice specifies how tools should be selected by the model.
@@ -44,7 +42,7 @@ type FunctionTool struct {
4442
// InputSchema defines the expected inputs. The language model uses this to understand
4543
// the tool's input requirements and provide matching suggestions.
4644
// InputSchema should be defined using a JSON schema.
47-
InputSchema *jsonschema.Definition `json:"input_schema,omitzero"`
45+
InputSchema *jsonschema.Schema `json:"input_schema,omitzero"`
4846
}
4947

5048
var _ ToolDefinition = &FunctionTool{}

‎aisdk/ai/go.mod

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,8 @@ go 1.24.0
55
require (
66
github.com/anthropics/anthropic-sdk-go v1.6.2
77
github.com/k0kubun/pp/v3 v3.5.0
8+
github.com/modelcontextprotocol/go-sdk v0.2.0
89
github.com/openai/openai-go v1.12.0
9-
github.com/sashabaranov/go-openai v1.40.5
1010
github.com/stretchr/testify v1.10.0
1111
github.com/tidwall/gjson v1.18.0
1212
go.jetify.com/pkg v0.0.0-20250801151358-5c7132d29664

‎aisdk/ai/go.sum

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,8 @@ github.com/anthropics/anthropic-sdk-go v1.6.2/go.mod h1:3qSNQ5NrAmjC8A2ykuruSQtt
33
github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
44
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
55
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
6+
github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
7+
github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
68
github.com/k0kubun/pp/v3 v3.5.0 h1:iYNlYA5HJAJvkD4ibuf9c8y6SHM0QFhaBuCqm1zHp0w=
79
github.com/k0kubun/pp/v3 v3.5.0/go.mod h1:5lzno5ZZeEeTV/Ky6vs3g6d1U3WarDrH8k240vMtGro=
810
github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
@@ -13,14 +15,14 @@ github.com/mattn/go-colorable v0.1.14 h1:9A9LHSqF/7dyVVX6g0U9cwm9pG3kP9gSzcuIPHP
1315
github.com/mattn/go-colorable v0.1.14/go.mod h1:6LmQG8QLFO4G5z1gPvYEzlUgJ2wF+stgPZH1UqBm1s8=
1416
github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
1517
github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
18+
github.com/modelcontextprotocol/go-sdk v0.2.0 h1:PESNYOmyM1c369tRkzXLY5hHrazj8x9CY1Xu0fLCryM=
19+
github.com/modelcontextprotocol/go-sdk v0.2.0/go.mod h1:0sL9zUKKs2FTTkeCCVnKqbLJTw5TScefPAzojjU459E=
1620
github.com/openai/openai-go v1.12.0 h1:NBQCnXzqOTv5wsgNC36PrFEiskGfO5wccfCWDo9S1U0=
1721
github.com/openai/openai-go v1.12.0/go.mod h1:g461MYGXEXBVdV5SaR/5tNzNbSfwTBBefwc+LlDCK0Y=
1822
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
1923
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
2024
github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ=
2125
github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc=
22-
github.com/sashabaranov/go-openai v1.40.5 h1:SwIlNdWflzR1Rxd1gv3pUg6pwPc6cQ2uMoHs8ai+/NY=
23-
github.com/sashabaranov/go-openai v1.40.5/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg=
2426
github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA=
2527
github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
2628
github.com/tidwall/gjson v1.14.2/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk=

‎aisdk/ai/options_test.go

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ import (
55
"net/http"
66
"testing"
77

8-
"github.com/sashabaranov/go-openai/jsonschema"
8+
"github.com/modelcontextprotocol/go-sdk/jsonschema"
99
"github.com/stretchr/testify/assert"
1010
"go.jetify.com/ai/api"
1111
"go.jetify.com/pkg/pointer"
@@ -77,15 +77,15 @@ func TestCallOptionBuilders(t *testing.T) {
7777
name: "WithResponseFormat",
7878
option: WithResponseFormat(&api.ResponseFormat{
7979
Type: "json",
80-
Schema: &jsonschema.Definition{},
80+
Schema: &jsonschema.Schema{},
8181
Name: "test",
8282
Description: "test desc",
8383
}),
8484
expected: GenerateOptions{
8585
CallOptions: api.CallOptions{
8686
ResponseFormat: &api.ResponseFormat{
8787
Type: "json",
88-
Schema: &jsonschema.Definition{},
88+
Schema: &jsonschema.Schema{},
8989
Name: "test",
9090
Description: "test desc",
9191
},

‎aisdk/ai/provider/anthropic/codec/encode_tools.go

Lines changed: 0 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@ import (
55
"fmt"
66

77
"github.com/anthropics/anthropic-sdk-go"
8-
"github.com/sashabaranov/go-openai/jsonschema"
98
"go.jetify.com/ai/api"
109
)
1110

@@ -41,45 +40,6 @@ func convertArgs[T any](args any) (*T, error) {
4140
return &result, nil
4241
}
4342

44-
// encodeInputSchema converts the JSON schema definition to Anthropic's schema format
45-
func encodeInputSchema(schema *jsonschema.Definition) (anthropic.BetaToolInputSchemaParam, error) {
46-
// Verify the schema type is "object"
47-
if schema.Type != "" && schema.Type != "object" {
48-
return anthropic.BetaToolInputSchemaParam{}, fmt.Errorf("unsupported schema type: %s, only 'object' is supported", schema.Type)
49-
}
50-
51-
// Create the input schema with the type field
52-
inputSchema := anthropic.BetaToolInputSchemaParam{
53-
Type: "object",
54-
}
55-
56-
// Add properties only if they exist
57-
if len(schema.Properties) > 0 {
58-
// Convert the properties to a map[string]any by marshaling and unmarshaling
59-
var properties map[string]any
60-
// Marshal to JSON
61-
propsJSON, err := json.Marshal(schema.Properties)
62-
if err != nil {
63-
return anthropic.BetaToolInputSchemaParam{}, fmt.Errorf("failed to marshal properties: %w", err)
64-
}
65-
66-
// Unmarshal back to map[string]any
67-
if err := json.Unmarshal(propsJSON, &properties); err != nil {
68-
return anthropic.BetaToolInputSchemaParam{}, fmt.Errorf("failed to unmarshal properties: %w", err)
69-
}
70-
71-
// Set the properties field
72-
inputSchema.Properties = properties
73-
}
74-
75-
// Add the required field if it's present in the original schema
76-
if len(schema.Required) > 0 {
77-
inputSchema.Required = schema.Required
78-
}
79-
80-
return inputSchema, nil
81-
}
82-
8343
// EncodeFunctionTool converts an API FunctionTool to Anthropic's tool format
8444
func EncodeFunctionTool(tool api.FunctionTool) (anthropic.BetaToolUnionParam, error) {
8545
inputSchema, err := encodeInputSchema(tool.InputSchema)

‎aisdk/ai/provider/anthropic/codec/encode_tools_test.go

Lines changed: 36 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ import (
55
"testing"
66

77
"github.com/anthropics/anthropic-sdk-go"
8-
"github.com/sashabaranov/go-openai/jsonschema"
8+
"github.com/modelcontextprotocol/go-sdk/jsonschema"
99
"github.com/stretchr/testify/assert"
1010
"github.com/stretchr/testify/require"
1111
"go.jetify.com/ai/api"
@@ -23,9 +23,9 @@ func TestEncodeFunctionTool(t *testing.T) {
2323
input: api.FunctionTool{
2424
Name: "test_function",
2525
Description: "A test function",
26-
InputSchema: &jsonschema.Definition{
26+
InputSchema: &jsonschema.Schema{
2727
Type: "object",
28-
Properties: map[string]jsonschema.Definition{
28+
Properties: map[string]*jsonschema.Schema{
2929
"param1": {
3030
Type: "string",
3131
Description: "First parameter",
@@ -50,11 +50,39 @@ func TestEncodeFunctionTool(t *testing.T) {
5050
}
5151
}`,
5252
},
53+
{
54+
name: "function tool with additionalProperties false",
55+
input: api.FunctionTool{
56+
Name: "test_function",
57+
Description: "A test function",
58+
InputSchema: &jsonschema.Schema{
59+
Type: "object",
60+
Properties: map[string]*jsonschema.Schema{
61+
"param1": {Type: "string"},
62+
},
63+
AdditionalProperties: api.FalseSchema(),
64+
},
65+
},
66+
want: `{
67+
"type": "tool",
68+
"name": "test_function",
69+
"description": "A test function",
70+
"input_schema": {
71+
"type": "object",
72+
"properties": {
73+
"param1": {
74+
"type": "string"
75+
}
76+
},
77+
"additionalProperties": false
78+
}
79+
}`,
80+
},
5381
{
5482
name: "function tool with minimal fields",
5583
input: api.FunctionTool{
5684
Name: "minimal_function",
57-
InputSchema: &jsonschema.Definition{
85+
InputSchema: &jsonschema.Schema{
5886
Type: "object",
5987
},
6088
},
@@ -538,9 +566,9 @@ func TestEncodeTools(t *testing.T) {
538566
functionTool := api.FunctionTool{
539567
Name: "test_function",
540568
Description: "A test function",
541-
InputSchema: &jsonschema.Definition{
569+
InputSchema: &jsonschema.Schema{
542570
Type: "object",
543-
Properties: map[string]jsonschema.Definition{
571+
Properties: map[string]*jsonschema.Schema{
544572
"param1": {
545573
Type: "string",
546574
Description: "First parameter",
@@ -673,7 +701,7 @@ func TestEncodeTools(t *testing.T) {
673701
api.FunctionTool{
674702
Name: "test_tool",
675703
Description: "A test tool",
676-
InputSchema: &jsonschema.Definition{Type: "object"},
704+
InputSchema: &jsonschema.Schema{Type: "object"},
677705
},
678706
},
679707
choice: &api.ToolChoice{Type: "tool", ToolName: "test_tool"},
@@ -682,7 +710,7 @@ func TestEncodeTools(t *testing.T) {
682710
mustEncodeFunctionTool(api.FunctionTool{
683711
Name: "test_tool",
684712
Description: "A test tool",
685-
InputSchema: &jsonschema.Definition{Type: "object"},
713+
InputSchema: &jsonschema.Schema{Type: "object"},
686714
}),
687715
},
688716
ToolChoice: specificChoice("test_tool"),

0 commit comments

Comments
 (0)