Tools can be built with predefined Input Schema #30
Conversation
Existing LLM projects will have already developed JSON Schemas to define their tool functions. Some LLM projects use libraries such as invopop/jsonschema, to build JSON Schemas automatically. This change honors the existing Options-based mechanism for building schemas, while opening a door to allow the use of JSON Schemas defined elsewhere. In so doing, users of this package are better enabled to validate the parity between Tool JSON schemas and the Go structs that receive the argument.
WalkthroughThis pull request updates the tool schema handling in the mcp package. The Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant Tool
Client->>Tool: Create tool instance (using NewToolWithRawSchema or standard constructor)
Client->>Tool: Call MarshalJSON()
alt Both InputSchema and RawInputSchema provided
Tool-->>Client: Return errToolSchemaConflict error
else RawInputSchema is provided
Tool-->>Client: Return JSON serialized with RawInputSchema
else
Tool-->>Client: Return JSON serialized with InputSchema
end
Poem
✨ Finishing Touches
🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (4)
mcp/tools_test.go (2)
19-20: Consider adding validation of initial successful marshal.You're correctly checking that the initial marshaling succeeds, but consider asserting on the structure of the marshaled JSON to ensure it contains the expected schema before adding the conflicting raw schema.
- _, err := json.Marshal(tool) - assert.Nil(t, err) + initialJSON, err := json.Marshal(tool) + assert.Nil(t, err) + + // Verify initial schema was properly included + var initialResult map[string]interface{} + err = json.Unmarshal(initialJSON, &initialResult) + assert.NoError(t, err) + assert.Contains(t, initialResult, "inputSchema")
44-46: Consider testing NewToolWithRawSchema's error handling.While you're testing the successful case, consider also testing how
NewToolWithRawSchemahandles invalid JSON in the raw schema parameter. This would ensure robust error handling for malformed schemas.+func TestToolWithInvalidRawSchema(t *testing.T) { + // Create an invalid raw schema + rawSchema := json.RawMessage(`{invalid json`) + + // Create a tool with invalid raw schema + tool := NewToolWithRawSchema("invalid-schema-tool", "Invalid schema", rawSchema) + + // Marshal to JSON should fail + _, err := json.Marshal(tool) + assert.Error(t, err) +}mcp/tools.go (2)
82-104: Robust schema conflict detection in MarshalJSON.The implementation correctly handles schema selection and conflict detection. When both schemas are present, it returns a descriptive error that includes the tool name.
Consider adding a check for the case where neither schema is provided, which could provide a more helpful error or default behavior:
} else if t.InputSchema.Type != "" { // Use the structured InputSchema m["inputSchema"] = t.InputSchema +} else { + // Neither schema is provided, using a default empty object schema + m["inputSchema"] = ToolInputSchema{ + Type: "object", + Properties: make(map[string]interface{}), + } }
144-159: Well-documented constructor for raw schema tools.The constructor is well-documented with a clear NOTE about incompatibility with ToolOption. The implementation is straightforward and correctly sets up a tool with a raw schema.
Consider adding validation to ensure the provided raw schema is valid JSON:
func NewToolWithRawSchema(name, description string, schema json.RawMessage) Tool { + // Validate that the schema is valid JSON + var js interface{} + if err := json.Unmarshal(schema, &js); err != nil { + // Since we can't return an error from this constructor, log a warning + // or consider panicking in development builds + // For now, we'll continue but the MarshalJSON will fail later + } + tool := Tool{ Name: name, Description: description, RawInputSchema: schema, } return tool }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
mcp/tools.go(3 hunks)mcp/tools_test.go(1 hunks)
🔇 Additional comments (4)
mcp/tools_test.go (2)
10-31: Thorough test case for schema conflict detection.This test effectively verifies that an error is returned when both schema types are used simultaneously. It creates a realistic scenario where a developer might add both schema types accidentally.
33-75: Comprehensive test for raw schema functionality.This test thoroughly validates the raw schema functionality by checking both the marshaling process and the resulting structure. The assertions properly verify that all components of the schema are correctly preserved.
mcp/tools.go (2)
9-9: Well-defined error with clear message.This error message clearly communicates the problem to developers when both schema types are used simultaneously.
75-78: Clear implementation of dual schema support.The implementation correctly hides both schema fields from direct JSON marshaling, allowing the custom MarshalJSON method to handle the schema logic.
NotesI gave consideration to alternative routes that involve changing the Landed on this proposal not out of affinity towards the internal representation but out of the simplicity of the changes and the maintaining a fairly straightforward API. |
|
I like this idea. Also I don't have a detailed roadmap at the moment. The main goal at the moment is to have feature parity with the official SDKs like TS and Python. Any and all PRs are welcome in that regard. |
Existing LLM projects will have already developed JSON Schemas to define their tool functions.
Some LLM projects use libraries such as invopop/jsonschema, to build JSON Schemas automatically.
This change honors the existing Options-based mechanism for building schemas, while opening a door to allow the use of JSON Schemas defined elsewhere.
In so doing, users of this package are better enabled to validate the parity between Tool JSON schemas and the Go structs that receive the argument.
Let me know if you'd prefer a different approach, etc. Your implementation is the best in class so far as I can see in the Golang world. I'd love to see what your road map looks like, in hopes of contributing in places where the goals align with what my primary project is doing.
Summary by CodeRabbit
New Features
Tests