Add structured output support to Ollama plugin#4884
Add structured output support to Ollama plugin#4884HVbajoria wants to merge 9 commits intofirebase:mainfrom
Conversation
Summary of ChangesHello, 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 enhances the Ollama plugin for Genkit Go by introducing robust structured output capabilities. This allows models to return responses strictly conforming to a provided JSON schema, greatly improving the reliability and ease of parsing model outputs. The changes are thoroughly tested and documented, providing clear guidance and examples for developers to leverage this new feature effectively. 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
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces structured output support to the Ollama plugin, a valuable feature that leverages Ollama's native format parameter for JSON output. However, it also introduces a critical Denial of Service (DoS) vulnerability. The current implementation uses incorrect type assertions, attempting to assert interface variables containing value types as pointers, which will cause an immediate runtime panic in Go and crash the request handler when structured output is requested. While the feature is well-tested with unit, property-based, and integration tests, and includes clear documentation, the type assertion issue needs to be addressed to prevent the DoS vulnerability. Additionally, there is a suggestion for code simplification in the tests.
…Request; simplify string containment check
|
/gemini review |
|
HI @apascal07 , Tried resolving the issues as marked by Gemini. Happy to learn while innovating. |
There was a problem hiding this comment.
Code Review
This pull request introduces structured output support for the Ollama plugin, a valuable feature that includes core logic changes, extensive unit and integration tests, new examples, and documentation. However, a critical security concern has been identified: the JSON schemas are being double-encoded as a string when passed to the Ollama API due to incorrect type definitions in the request structs. This issue will likely cause Ollama to ignore schema enforcement, bypassing the intended security and integrity benefits of structured output. Additionally, my review suggests improving test quality by addressing duplicated logic and fixing a compilation issue in the integration tests, along with a minor refactoring in the core logic to enhance clarity and maintainability.
go/plugins/ollama/ollama.go
Outdated
| schemaJSON, err := json.Marshal(input.Output.Schema) | ||
| if err != nil { | ||
| // Return error before making HTTP request if schema serialization fails | ||
| return nil, fmt.Errorf("failed to serialize output schema: %v", err) | ||
| } | ||
|
|
||
| // Set the format field on the appropriate request type (chat or generate endpoint) | ||
| if isChatModel { | ||
| payload.(*ollamaChatRequest).Format = string(schemaJSON) | ||
| } else { | ||
| payload.(*ollamaModelRequest).Format = string(schemaJSON) | ||
| } |
There was a problem hiding this comment.
This section of code has a critical issue: the JSON schema is being double-encoded when passed to the Ollama API, which will likely cause Ollama to ignore schema enforcement, bypassing the intended security and integrity benefits of structured output. The Format field in ollamaChatRequest and ollamaModelRequest (lines 143 and 153) is currently a string. When input.Output.Schema is marshaled to a JSON string (line 292) and assigned to this field, it results in double-encoding (e.g., "format": "{\"type\":\"object\",...}") when the request payload is marshaled for the API call. Ollama expects a JSON Schema object, not an escaped string.
To resolve this, the Format field in both ollamaChatRequest and ollamaModelRequest structs should be changed to any or json.RawMessage. Then, assign the input.Output.Schema map directly to the Format field instead of marshaling it to a string first.
Additionally, the logic for setting the Format field is currently repeated for both chat and generate models. This can be refactored for better clarity and maintainability once the type issue is addressed.
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request adds structured output support to the Ollama plugin, which is a great enhancement. The implementation is well-structured, and I appreciate the comprehensive set of unit, property-based, and integration tests, as well as the new documentation and examples. I have a couple of suggestions: one critical fix for a missing import in the integration tests that causes a compilation error, and a minor refactoring suggestion to improve code clarity and reduce duplication in the payload building logic. Overall, this is a high-quality contribution.
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request adds structured output support to the Ollama plugin, which is a great enhancement. The implementation is well-structured, with the logic for payload creation refactored into a separate buildPayload function. The addition of comprehensive unit, property-based, and integration tests is excellent and ensures the new functionality is robust and backward compatible. The new documentation and examples are also clear and helpful. I've found a couple of minor issues: a formatting problem in ollama.go and a missing import in the new integration test file that will cause a compilation failure. Once these are addressed, this PR will be in great shape.
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
|
/gemini review |
There was a problem hiding this comment.
Code Review
This is a high-quality pull request that adds significant value by implementing structured output support for the Ollama plugin. The code is well-structured, particularly with the refactoring of payload creation into its own function. The addition of comprehensive documentation, examples, and an exceptionally thorough test suite—including unit, integration, and property-based tests—is commendable and greatly enhances the robustness and usability of the new feature. My feedback includes a couple of minor suggestions to further improve code robustness in the implementation and tests. Overall, excellent work.
…ance test case for format validation
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces structured output support for the Ollama plugin, which is a significant and valuable feature. The implementation is well-structured, with a clear separation of concerns in building the request payload. The addition of comprehensive unit, property-based, and integration tests ensures the feature is robust and backward compatible. The documentation and examples are thorough and will be very helpful for users. I have one suggestion to improve the reliability of an integration test case.
| t.Run("Ollama API error response", func(t *testing.T) { | ||
| // Use a valid model but with a very short timeout to potentially trigger errors | ||
| modelName := getChatModel() | ||
| checkModelAvailable(t, serverAddress, modelName) | ||
|
|
||
| shortTimeoutOllama := &Ollama{ | ||
| ServerAddress: serverAddress, | ||
| Timeout: 1, // 1 second timeout | ||
| } | ||
| shortTimeoutOllama.Init(ctx) | ||
|
|
||
| model := shortTimeoutOllama.DefineModel(g, ModelDefinition{ | ||
| Name: modelName, | ||
| Type: "chat", | ||
| }, nil) | ||
|
|
||
| // Make a request that might timeout | ||
| request := &ai.ModelRequest{ | ||
| Messages: []*ai.Message{ | ||
| { | ||
| Role: ai.RoleUser, | ||
| Content: []*ai.Part{ | ||
| ai.NewTextPart("Write a very long essay about the history of computing with at least 5000 words."), | ||
| }, | ||
| }, | ||
| }, | ||
| } | ||
|
|
||
| _, err := model.Generate(ctx, request, nil) | ||
| // This might or might not error depending on model speed | ||
| // If it errors, verify the error is properly formatted | ||
| if err != nil { | ||
| errMsg := err.Error() | ||
| t.Logf("Received error (expected with short timeout): %v", err) | ||
|
|
||
| // Verify error message is descriptive | ||
| if errMsg == "" { | ||
| t.Error("Error message is empty") | ||
| } | ||
| } else { | ||
| t.Log("Request completed within timeout (no error to test)") | ||
| } | ||
| }) |
There was a problem hiding this comment.
The test case Ollama API error response is non-deterministic because it relies on a request timeout, which may not occur consistently across different environments or under varying system loads. This can lead to flaky tests.
To make error handling tests more reliable, it's better to test against a deterministic failure. For instance, you could test how the client handles a connection error by attempting to connect to an invalid server address. This would provide consistent and reliable testing of the error propagation logic.
Here's a suggestion to replace the current test case with one that tests a connection error:
t.Run("Invalid server address connection error", func(t *testing.T) {
// Use a port that is unlikely to be in use.
invalidAddressOllama := &Ollama{
ServerAddress: "http://localhost:99999",
Timeout: 2,
}
invalidAddressOllama.Init(ctx)
model := invalidAddressOllama.DefineModel(g, ModelDefinition{
Name: "any-model",
Type: "chat",
}, nil)
request := &ai.ModelRequest{
Messages: []*ai.Message{
{
Role: ai.RoleUser,
Content: []*ai.Part{
ai.NewTextPart("Hello"),
},
},
},
}
_, err := model.Generate(ctx, request, nil)
if err == nil {
t.Fatal("Expected an error for an invalid server address, but got nil")
}
if !strings.Contains(err.Error(), "failed to send request") {
t.Errorf("Expected error message to contain 'failed to send request', but got: %v", err)
}
t.Logf("Successfully received expected connection error: %v", err)
})
|
Hi @apascal07 , Tried resolving all the issues as highlighted by Gemini. Not sure if last one should be implemented. Could you please assist? |
1 participant
Add Structured Output Support to Ollama Plugin (Go)
Fixes #4788
Overview
This PR adds structured output support to the Ollama plugin in Genkit Go, enabling developers to enforce JSON schemas on model responses. This leverages Ollama's native structured output capability via the
formatparameter as documented at https://docs.ollama.com/capabilities/structured-outputs.Changes
Core Implementation
Modified
go/plugins/ollama/ollama.go:generator.generate()methodformatparameter for both chat and generate endpointsKey Features:
ModelRequest.Output.Schemais provided, the schema is serialized to JSON and passed to Ollama'sformatparameterModelRequest.Output.Formatis set to"json"without a schema, enables generic JSON output modeTesting
Modified
go/plugins/ollama/ollama_test.go:gopterlibrary (100+ iterations per property)Added
go/plugins/ollama/ollama_integration_test.go:Documentation
Added
go/plugins/ollama/README.md:Added
go/plugins/ollama/INTEGRATION_TESTS.md:Examples
Added
go/samples/ollama-structured/main.go:Added
go/samples/ollama-json-mode/main.go:Dependencies
Modified
go/go.modandgo/go.sum:github.com/leanovate/gopter v0.2.11for property-based testingTesting
All tests pass successfully:
Backward Compatibility
This implementation maintains full backward compatibility:
Outputconfig work unchangedSchemawork unchangedUsage Example
Checklist
Related Documentation