feat: unify schema validation using jsonschema library#595
feat: unify schema validation using jsonschema library#595GuyGoldenberg wants to merge 1 commit intogoogle:mainfrom
Conversation
Summary of ChangesHello @GuyGoldenberg, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly refactors the schema validation logic within the ADK by standardizing on the Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
There was a problem hiding this comment.
Code Review
The pull request successfully unifies schema validation across the ADK by leveraging the jsonschema library. This is a significant improvement that ensures consistent validation behavior, particularly for enums which were previously handled inconsistently. I have identified a critical bug in the schema converter regarding the interaction between Nullable and AnyOf fields, as well as some opportunities to improve performance by caching resolved schemas in hot paths.
internal/schemautil/convert.go
Outdated
| if gs.Nullable != nil && *gs.Nullable { | ||
| js.AnyOf = []*jsonschema.Schema{ | ||
| {Type: strings.ToLower(string(gs.Type))}, | ||
| {Type: "null"}, | ||
| } | ||
| // Clear the main type when using AnyOf | ||
| js.Type = "" | ||
| } | ||
|
|
||
| // Convert AnyOf (recursive) | ||
| if len(gs.AnyOf) > 0 { | ||
| js.AnyOf = make([]*jsonschema.Schema, len(gs.AnyOf)) | ||
| for i, s := range gs.AnyOf { | ||
| js.AnyOf[i] = GenaiToJSONSchema(s) | ||
| } | ||
| } |
There was a problem hiding this comment.
There is a bug in the handling of Nullable and AnyOf. The current implementation of AnyOf (lines 111-116) completely overwrites the AnyOf slice created by the Nullable logic (lines 101-108). If a schema is both nullable and contains AnyOf, the nullability will be lost. The logic should be merged to ensure null is added to the AnyOf list if Nullable is true.
// Convert AnyOf (recursive)
if len(gs.AnyOf) > 0 {
js.AnyOf = make([]*jsonschema.Schema, len(gs.AnyOf))
for i, s := range gs.AnyOf {
js.AnyOf[i] = GenaiToJSONSchema(s)
}
}
// Handle Nullable via AnyOf with null type
if gs.Nullable != nil && *gs.Nullable {
if len(js.AnyOf) > 0 {
js.AnyOf = append(js.AnyOf, &jsonschema.Schema{Type: "null"})
} else {
js.AnyOf = []*jsonschema.Schema{
{Type: js.Type},
{Type: "null"},
}
}
// Clear the main type when using AnyOf for nullability
js.Type = ""
}
internal/schemautil/convert.go
Outdated
| js := &jsonschema.Schema{ | ||
| Type: strings.ToLower(string(gs.Type)), | ||
| Description: gs.Description, | ||
| Title: gs.Title, | ||
| Pattern: gs.Pattern, | ||
| Required: gs.Required, | ||
| } |
There was a problem hiding this comment.
The converter is missing the Format field. JSON Schema supports formats like date-time, email, etc., which are often defined in genai.Schema and are useful for validation.
| js := &jsonschema.Schema{ | |
| Type: strings.ToLower(string(gs.Type)), | |
| Description: gs.Description, | |
| Title: gs.Title, | |
| Pattern: gs.Pattern, | |
| Required: gs.Required, | |
| } | |
| js := &jsonschema.Schema{ | |
| Type: strings.ToLower(string(gs.Type)), | |
| Format: gs.Format, | |
| Description: gs.Description, | |
| Title: gs.Title, | |
| Pattern: gs.Pattern, | |
| Required: gs.Required, | |
| } |
internal/schemautil/convert.go
Outdated
| // Convert numeric constraints | ||
| if gs.Minimum != nil { | ||
| js.Minimum = gs.Minimum | ||
| } | ||
| if gs.Maximum != nil { | ||
| js.Maximum = gs.Maximum | ||
| } |
internal/schemautil/convert.go
Outdated
| // Convert array constraints | ||
| if gs.MinItems != nil { | ||
| minItems := int(*gs.MinItems) | ||
| js.MinItems = &minItems | ||
| } | ||
| if gs.MaxItems != nil { | ||
| maxItems := int(*gs.MaxItems) | ||
| js.MaxItems = &maxItems | ||
| } |
There was a problem hiding this comment.
The converter is missing the UniqueItems array constraint.
// Convert array constraints
if gs.MinItems != nil {
minItems := int(*gs.MinItems)
js.MinItems = &minItems
}
if gs.MaxItems != nil {
maxItems := int(*gs.MaxItems)
js.MaxItems = &maxItems
}
if gs.UniqueItems != nil {
js.UniqueItems = *gs.UniqueItems
}| var err error | ||
| if agentInputSchema != nil { | ||
| if err = utils.ValidateMapOnSchema(margs, agentInputSchema, true); err != nil { | ||
| resolved, resolveErr := schemautil.GenaiToResolvedJSONSchema(agentInputSchema) |
There was a problem hiding this comment.
Calling GenaiToResolvedJSONSchema inside the Run method is inefficient. This function performs a recursive conversion of the schema and then resolves (compiles) it, which is an expensive operation. Since the agent's input schema is typically static, it should be resolved once (e.g., during tool initialization or lazily cached) rather than on every execution.
| if err := json.Unmarshal([]byte(outputText), &parsedOutput); err != nil { | ||
| return nil, fmt.Errorf("failed to parse output JSON for sub-agent %s: %w", t.agent.Name(), err) | ||
| } | ||
| resolved, resolveErr := schemautil.GenaiToResolvedJSONSchema(agentOutputSchema) |
| return nil, fmt.Errorf("unexpected args type for set_model_response: %T", args) | ||
| } | ||
| if err := utils.ValidateMapOnSchema(m, t.schema, false); err != nil { | ||
| resolved, err := schemautil.GenaiToResolvedJSONSchema(t.schema) |
There was a problem hiding this comment.
This change unifies schema validation across all tool types by using the jsonschema library consistently, which properly validates enum values and other schema constraints. Previously, FunctionTool used jsonschema.Resolved.Validate() while AgentTool and output schema validation used a custom ValidateMapOnSchema function that lacked enum validation support. Changes: - Add internal/schemautil package with GenaiToJSONSchema converter - Update AgentTool to use jsonschema validation - Update setModelResponseTool to use jsonschema validation - Remove unused custom validation functions from internal/utils This aligns with the TODO at llmagent.go:238 to migrate toward jsonschema.
7ddaee1 to
734823a
Compare
Summary
Unifies schema validation by using the
jsonschemalibrary consistently across all tool types, adding enum validation support forAgentTooland output schemas.Problem
FunctionToolusedjsonschema.Resolved.Validate()which validates enums, butAgentTooland output schema used a customValidateMapOnSchema()that did not.Solution
Add a
genai.Schema→jsonschema.Schemaconverter so all validation uses jsonschema.Addresses TODO at
agent/llmagent/llmagent.go:238.Changes
internal/schemautilwith converter functionsAgentToolandsetModelResponseToolto use jsonschema validationinternal/utils/schema_utils.go