88 "reflect"
99 "strconv"
1010
11+ jsonschemaGenerator "github.com/invopop/jsonschema"
1112 "github.com/santhosh-tekuri/jsonschema"
1213)
1314
@@ -482,13 +483,15 @@ type Tool struct {
482483 // A JSON Schema object defining the expected parameters for the tool.
483484 InputSchema ToolSchema `json:"inputSchema"`
484485 // A JSON Schema object defining the expected output for the tool.
485- OutputSchema ToolSchema `json:"outputSchema,omitempty"`
486- // Compiled JSON schema validator for output validation, cached for performance
487- compiledOutputSchema * jsonschema.Schema `json:"-"`
486+ OutputSchema json.RawMessage `json:"outputSchema,omitempty"`
488487 // Alternative to InputSchema - allows arbitrary JSON Schema to be provided
489488 RawInputSchema json.RawMessage `json:"-"` // Hide this from JSON marshaling
490489 // Optional properties describing tool behavior
491490 Annotations ToolAnnotation `json:"annotations"`
491+
492+ // Internal fields for output validation, not serialized
493+ outputValidator * jsonschema.Schema `json:"-"`
494+ outputType reflect.Type `json:"-"`
492495}
493496
494497// GetName returns the name of the tool.
@@ -499,45 +502,43 @@ func (t Tool) GetName() string {
499502// HasOutputSchema returns true if the tool has an output schema defined.
500503// This indicates that the tool can return structured content.
501504func (t Tool ) HasOutputSchema () bool {
502- return t .OutputSchema . Type != ""
505+ return t .OutputSchema != nil
503506}
504507
505508// validateStructuredOutput performs the actual validation using the compiled schema
506509func (t Tool ) validateStructuredOutput (result * CallToolResult ) error {
507- return t .compiledOutputSchema .ValidateInterface (result .StructuredContent )
510+ return t .outputValidator .ValidateInterface (result .StructuredContent )
508511}
509512
510513// ensureOutputSchemaValidator compiles and caches the JSON schema validator if not already done
511514func (t * Tool ) ensureOutputSchemaValidator () error {
512- if t .compiledOutputSchema != nil {
515+ if t .outputValidator != nil {
513516 return nil
514517 }
515518
516- schemaBytes , err := t .OutputSchema .MarshalJSON ()
517- if err != nil {
518- return err
519+ if t .OutputSchema == nil {
520+ return nil
519521 }
520522
521523 compiler := jsonschema .NewCompiler ()
522524
523- const validatorKey = "output-schema-validator"
524- if err := compiler .AddResource (validatorKey , bytes .NewReader (schemaBytes )); err != nil {
525+ if err := compiler .AddResource ("output-schema" , bytes .NewReader (t .OutputSchema )); err != nil {
525526 return err
526527 }
527528
528- compiledSchema , err := compiler .Compile (validatorKey )
529+ compiledSchema , err := compiler .Compile ("output-schema" )
529530 if err != nil {
530531 return err
531532 }
532533
533- t .compiledOutputSchema = compiledSchema
534+ t .outputValidator = compiledSchema
534535 return nil
535536}
536537
537538// ValidateStructuredOutput validates the structured content against the tool's output schema.
538539// Returns nil if the tool has no output schema or if validation passes.
539540// Returns an error if the tool has an output schema but the structured content is invalid.
540- func (t Tool ) ValidateStructuredOutput (result * CallToolResult ) error {
541+ func (t * Tool ) ValidateStructuredOutput (result * CallToolResult ) error {
541542 if ! t .HasOutputSchema () {
542543 return nil
543544 }
@@ -577,7 +578,7 @@ func (t Tool) MarshalJSON() ([]byte, error) {
577578 }
578579
579580 // Add output schema if defined
580- if t .HasOutputSchema () {
581+ if t .OutputSchema != nil {
581582 m ["outputSchema" ] = t .OutputSchema
582583 }
583584
@@ -645,19 +646,17 @@ func NewTool(name string, opts ...ToolOption) Tool {
645646 Properties : make (map [string ]any ),
646647 Required : nil , // Will be omitted from JSON if empty
647648 },
648- OutputSchema : ToolSchema {
649- Type : "" ,
650- Properties : make (map [string ]any ),
651- Required : nil , // Will be omitted from JSON if empty
652- },
653- compiledOutputSchema : nil ,
649+ OutputSchema : nil ,
650+ RawInputSchema : nil ,
654651 Annotations : ToolAnnotation {
655652 Title : "" ,
656653 ReadOnlyHint : ToBoolPtr (false ),
657654 DestructiveHint : ToBoolPtr (true ),
658655 IdempotentHint : ToBoolPtr (false ),
659656 OpenWorldHint : ToBoolPtr (true ),
660657 },
658+ outputValidator : nil ,
659+ outputType : nil ,
661660 }
662661
663662 for _ , opt := range opts {
@@ -1159,144 +1158,105 @@ func WithBooleanItems(opts ...PropertyOption) PropertyOption {
11591158
11601159// WithOutputSchema sets the output schema for the Tool.
11611160// This allows the tool to define the structure of its return data.
1162- func WithOutputSchema (schema ToolSchema ) ToolOption {
1161+ func WithOutputSchema (schema json. RawMessage ) ToolOption {
11631162 return func (t * Tool ) {
11641163 t .OutputSchema = schema
11651164 }
11661165}
11671166
1168- // WithOutputBoolean adds a boolean property to the tool's output schema.
1169- // It accepts property options to configure the boolean property's behavior and constraints.
1170- func WithOutputBoolean (name string , opts ... PropertyOption ) ToolOption {
1171- return func (t * Tool ) {
1172- // Initialize output schema if not set
1173- if t .OutputSchema .Type == "" {
1174- t .OutputSchema .Type = "object"
1175- }
1176-
1177- schema := map [string ]any {
1178- "type" : "boolean" ,
1179- }
1180-
1181- for _ , opt := range opts {
1182- opt (schema )
1183- }
1184-
1185- // Remove required from property schema and add to OutputSchema.required
1186- if required , ok := schema ["required" ].(bool ); ok && required {
1187- delete (schema , "required" )
1188- t .OutputSchema .Required = append (t .OutputSchema .Required , name )
1189- }
1190-
1191- t .OutputSchema .Properties [name ] = schema
1192- }
1193- }
1167+ //
1168+ // New Output Schema Functions
1169+ //
11941170
1195- // WithOutputNumber adds a number property to the tool's output schema .
1196- // It accepts property options to configure the number property's behavior and constraints .
1197- func WithOutputNumber ( name string , opts ... PropertyOption ) ToolOption {
1171+ // WithOutputType sets the output schema for the Tool using Go generics and struct tags .
1172+ // This replaces the builder pattern with a cleaner interface based on struct definitions .
1173+ func WithOutputType [ T any ]( ) ToolOption {
11981174 return func (t * Tool ) {
1199- // Initialize output schema if not set
1200- if t .OutputSchema .Type == "" {
1201- t .OutputSchema .Type = "object"
1202- }
1203-
1204- schema := map [string ]any {
1205- "type" : "number" ,
1206- }
1207-
1208- for _ , opt := range opts {
1209- opt (schema )
1175+ var zero T
1176+ validator , schemaBytes , err := compileOutputSchema (zero )
1177+ if err != nil {
1178+ // Skip setting output schema if compilation fails
1179+ // This allows the tool to work without validation
1180+ return
12101181 }
12111182
1212- // Remove required from property schema and add to OutputSchema.required
1213- if required , ok := schema ["required" ].(bool ); ok && required {
1214- delete (schema , "required" )
1215- t .OutputSchema .Required = append (t .OutputSchema .Required , name )
1216- }
1217-
1218- t .OutputSchema .Properties [name ] = schema
1183+ t .OutputSchema = schemaBytes
1184+ t .outputValidator = validator
1185+ t .outputType = reflect .TypeOf (zero )
12191186 }
12201187}
12211188
1222- // WithOutputString adds a string property to the tool's output schema.
1223- // It accepts property options to configure the string property's behavior and constraints.
1224- func WithOutputString (name string , opts ... PropertyOption ) ToolOption {
1225- return func (t * Tool ) {
1226- // Initialize output schema if not set
1227- if t .OutputSchema .Type == "" {
1228- t .OutputSchema .Type = "object"
1229- }
1230-
1231- schema := map [string ]any {
1232- "type" : "string" ,
1233- }
1234-
1235- for _ , opt := range opts {
1236- opt (schema )
1237- }
1238-
1239- // Remove required from property schema and add to OutputSchema.required
1240- if required , ok := schema ["required" ].(bool ); ok && required {
1241- delete (schema , "required" )
1242- t .OutputSchema .Required = append (t .OutputSchema .Required , name )
1243- }
1244-
1245- t .OutputSchema .Properties [name ] = schema
1189+ // compileOutputSchema generates JSON schema from a struct and compiles it for validation
1190+ func compileOutputSchema [T any ](sample T ) (* jsonschema.Schema , json.RawMessage , error ) {
1191+ // Generate JSON Schema from struct
1192+ reflector := jsonschemaGenerator.Reflector {
1193+ // Use Draft 7 which is widely supported
1194+ DoNotReference : true ,
12461195 }
1247- }
1248-
1249- // WithOutputObject adds an object property to the tool's output schema.
1250- // It accepts property options to configure the object property's behavior and constraints.
1251- func WithOutputObject (name string , opts ... PropertyOption ) ToolOption {
1252- return func (t * Tool ) {
1253- // Initialize output schema if not set
1254- if t .OutputSchema .Type == "" {
1255- t .OutputSchema .Type = "object"
1256- }
1196+ schema := reflector .Reflect (& sample )
12571197
1258- schema := map [string ]any {
1259- "type" : "object" ,
1260- "properties" : map [string ]any {},
1261- }
1198+ // Manually override the schema version to Draft-07.
1199+ // This is required because our validator, santhosh-tekuri/jsonschema,
1200+ // only supports up to Draft-07. This workaround ensures compatibility
1201+ // between the generator and the validator.
1202+ schema .Version = "http://json-schema.org/draft-07/schema#"
12621203
1263- for _ , opt := range opts {
1264- opt (schema )
1265- }
1204+ // Serialize to JSON
1205+ schemaBytes , err := json .Marshal (schema )
1206+ if err != nil {
1207+ return nil , nil , fmt .Errorf ("failed to marshal schema: %w" , err )
1208+ }
12661209
1267- // Remove required from property schema and add to OutputSchema.required
1268- if required , ok := schema [ "required" ].( bool ); ok && required {
1269- delete ( schema , "required" )
1270- t . OutputSchema . Required = append ( t . OutputSchema . Required , name )
1271- }
1210+ // Compile for validation using santhosh-tekuri/jsonschema
1211+ compiler := jsonschema . NewCompiler ()
1212+ if err := compiler . AddResource ( " schema" , bytes . NewReader ( schemaBytes )); err != nil {
1213+ return nil , nil , fmt . Errorf ( "failed to add schema resource: %w" , err )
1214+ }
12721215
1273- t .OutputSchema .Properties [name ] = schema
1216+ validator , err := compiler .Compile ("schema" )
1217+ if err != nil {
1218+ return nil , nil , fmt .Errorf ("failed to compile schema: %w" , err )
12741219 }
1220+
1221+ return validator , schemaBytes , nil
12751222}
12761223
1277- // WithOutputArray adds an array property to the tool's output schema.
1278- // It accepts property options to configure the array property's behavior and constraints.
1279- func WithOutputArray (name string , opts ... PropertyOption ) ToolOption {
1280- return func (t * Tool ) {
1281- // Initialize output schema if not set
1282- if t .OutputSchema .Type == "" {
1283- t .OutputSchema .Type = "object"
1284- }
1224+ // ValidateOutput validates the structured content against the tool's output schema.
1225+ // Skips validation when IsError is true or the tool has no output schema defined.
1226+ func (t * Tool ) ValidateOutput (result * CallToolResult ) error {
1227+ // Skip validation if IsError is true or no output schema defined
1228+ if result .IsError || ! t .HasOutputSchema () {
1229+ return nil
1230+ }
12851231
1286- schema := map [ string ] any {
1287- "type" : "array" ,
1288- }
1232+ if result . StructuredContent == nil {
1233+ return fmt . Errorf ( "tool %s requires structured output but got nil" , t . Name )
1234+ }
12891235
1290- for _ , opt := range opts {
1291- opt (schema )
1236+ // Ensure the validator is compiled
1237+ if err := t .ensureOutputSchemaValidator (); err != nil {
1238+ return fmt .Errorf ("failed to compile output schema for tool %s: %w" , t .Name , err )
1239+ }
1240+
1241+ // Convert structured content to JSON-compatible format for validation
1242+ var validationData any
1243+
1244+ // If it's already a map, slice, or primitive, use it as-is
1245+ // If it's a struct, convert it to JSON-compatible format
1246+ switch result .StructuredContent .(type ) {
1247+ case map [string ]any , []any , string , int , int64 , float64 , bool , nil :
1248+ validationData = result .StructuredContent
1249+ default :
1250+ // Convert struct to JSON-compatible format via JSON marshaling/unmarshaling
1251+ jsonBytes , err := json .Marshal (result .StructuredContent )
1252+ if err != nil {
1253+ return fmt .Errorf ("failed to marshal structured content for validation: %w" , err )
12921254 }
12931255
1294- // Remove required from property schema and add to OutputSchema.required
1295- if required , ok := schema ["required" ].(bool ); ok && required {
1296- delete (schema , "required" )
1297- t .OutputSchema .Required = append (t .OutputSchema .Required , name )
1256+ if err := json .Unmarshal (jsonBytes , & validationData ); err != nil {
1257+ return fmt .Errorf ("failed to unmarshal structured content for validation: %w" , err )
12981258 }
1299-
1300- t .OutputSchema .Properties [name ] = schema
13011259 }
1260+
1261+ return t .outputValidator .ValidateInterface (validationData )
13021262}
0 commit comments