Skip to content

Comments

feat(go): added common middleware (e.g. tool approval, retry, fallback)#4719

Draft
apascal07 wants to merge 3 commits intoap/go-middlewarefrom
ap/go-middleware-impls
Draft

feat(go): added common middleware (e.g. tool approval, retry, fallback)#4719
apascal07 wants to merge 3 commits intoap/go-middlewarefrom
ap/go-middleware-impls

Conversation

@apascal07
Copy link
Collaborator

Checklist (if applicable):

@apascal07 apascal07 changed the base branch from main to ap/go-middleware February 17, 2026 20:43
@github-actions github-actions bot added docs Improvements or additions to documentation js go tooling python Python labels Feb 17, 2026
@gemini-code-assist
Copy link
Contributor

Summary of Changes

Hello @apascal07, 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 extensibility and control over AI generation flows by introducing a new, unified middleware system. This system allows developers to inject custom logic at various stages of the generation process, including before generation, around model calls, and during tool execution. The changes include new type definitions, schema updates, and the implementation of several practical middleware plugins, providing robust capabilities for handling errors, retries, fallbacks, and tool interactions.

Highlights

  • Unified Middleware System: Introduced a new, unified Middleware interface in Go, replacing the older ModelMiddleware and providing hooks for Generate, Model, and Tool stages of AI generation. This allows for more comprehensive control over the generation lifecycle.
  • Middleware Types and Integration: Added MiddlewareDesc and MiddlewareRef types in both TypeScript and Go, along with corresponding JSON schema updates, to describe and reference middleware configurations. The GenerateActionOptions now includes a use field to apply middleware to generation requests.
  • Built-in Middleware Plugins: Implemented several common middleware as Go plugins: Retry (for exponential backoff on model failures), Fallback (for trying alternative models), ToolApproval (for explicit tool execution approval), and ToolErrorHandler (to convert tool errors into model-readable responses).
  • Context Propagation and Cloning: Enhanced context propagation by introducing a genkitCtxKey to pass the Genkit instance to middleware, and added Clone methods for Part and Message types in Go to ensure independent copies during processing, preventing unintended mutations.
  • Reflection API for Middleware: Added a new reflection API endpoint /api/values to list registered values, specifically enabling discovery and inspection of defined middleware.

🧠 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
  • genkit-tools/common/src/types/index.ts
    • Exported new middleware types.
  • genkit-tools/common/src/types/middleware.ts
    • Added TypeScript definitions for MiddlewareDesc and MiddlewareRef schemas.
  • genkit-tools/common/src/types/model.ts
    • Imported MiddlewareRefSchema.
    • Added use field to GenerateActionOptionsSchema.
  • genkit-tools/genkit-schema.json
    • Updated JSON schema to include MiddlewareDesc and MiddlewareRef definitions.
    • Added use property to GenerateActionOptions schema.
  • genkit-tools/scripts/schema-exporter.ts
    • Added middleware.ts to the list of exported type modules.
  • go/ai/document.go
    • Imported maps and slices packages.
    • Added Clone methods for Part and Message structs to create shallow copies.
  • go/ai/document_test.go
    • Imported bytes package.
    • Added TestPartClone and TestMessageClone to verify cloning functionality.
  • go/ai/gen.go
    • Added Use field of type []*MiddlewareRef to GenerateActionOptions.
    • Added MiddlewareDesc and MiddlewareRef struct definitions.
  • go/ai/generate.go
    • Deprecated ModelMiddleware in favor of the new Middleware interface.
    • Integrated middleware handling into GenerateWithRequest for Generate, Model, and Tool hooks.
    • Updated handleToolRequests and handleResumedToolRequest to accept and utilize middleware handlers.
    • Removed the generic clone function, replaced by specific Clone methods on Part and Message.
  • go/ai/middleware.go
    • Added new file defining the Middleware interface and related types (GenerateParams, ModelParams, ToolParams).
    • Provided BaseMiddleware for default pass-through implementations.
    • Added helper functions NewMiddleware, DefineMiddleware, and LookupMiddleware for managing middleware descriptors.
    • Defined MiddlewarePlugin interface for plugins providing middleware.
  • go/ai/middleware_test.go
    • Added new file with comprehensive tests for the middleware system, including testMiddleware, stableStateMiddleware, and orderMiddleware.
  • go/ai/option.go
    • Deprecated WithMiddleware option.
    • Added WithUse option to apply the new Middleware interface to generation requests.
    • Updated commonGenOptions to include the Use field.
  • go/ai/prompt.go
    • Updated Prompt.Execute to process and apply middleware specified via the Use option.
  • go/ai/tools.go
    • Added NewToolInterruptError function to create tool interrupt errors.
    • Added IsToolResumed function to check if a tool execution is resumed.
  • go/core/schemas.config
    • Added documentation for GenerateActionOptions.use.
    • Added documentation for new MiddlewareDesc and MiddlewareRef types and their fields.
  • go/genkit/genkit.go
    • Added genkitCtxKey and FromContext function to retrieve the Genkit instance from context.
    • Updated Init to register middleware provided by plugins.
    • Modified Generate and related functions to pass the Genkit instance into the context for middleware access.
  • go/genkit/reflection.go
    • Added a new HTTP handler handleListValues for the /api/values endpoint to list registered values, including middleware.
  • go/plugins/middleware/fallback.go
    • Added new file implementing the Fallback middleware, which retries model calls with alternative models on specific errors.
  • go/plugins/middleware/fallback_test.go
    • Added new file with tests for the Fallback middleware.
  • go/plugins/middleware/plugin.go
    • Added new file defining the Middleware plugin, which registers built-in middleware like Retry, Fallback, ToolApproval, and ToolErrorHandler.
  • go/plugins/middleware/retry.go
    • Added new file implementing the Retry middleware, providing exponential backoff for failed model calls.
  • go/plugins/middleware/retry_test.go
    • Added new file with tests for the Retry middleware.
  • go/plugins/middleware/subagent_design.md
    • Added a design document outlining a conceptual subagent middleware architecture.
  • go/plugins/middleware/tool_approval.go
    • Added new file implementing the ToolApproval middleware, which requires explicit approval for tool execution.
  • go/plugins/middleware/tool_approval_test.go
    • Added new file with tests for the ToolApproval middleware.
  • go/plugins/middleware/tool_error_handler.go
    • Added new file implementing the ToolErrorHandler middleware, which converts tool errors into model responses.
  • py/packages/genkit/src/genkit/core/typing.py
    • Added MiddlewareDesc and MiddlewareRef Pydantic models.
    • Added use field to GenerateActionOptions model.
Activity
  • The pull request title adheres to the Conventional Commits specification.
  • The author has manually tested the changes.
  • Unit tests have been implemented for the new features.
Using Gemini Code Assist

The 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 /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature Command Description
Code Review /gemini review Performs a code review for the current pull request in its current state.
Pull Request Summary /gemini summary Provides a summary of the current pull request in its current state.
Comment @gemini-code-assist Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help /gemini help Displays a list of available commands.

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 .gemini/ folder in the base of the repository. Detailed instructions can be found here.

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

  1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution.

Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces a comprehensive middleware system for Go, adding valuable features like retry, fallback, and tool approval. The implementation is well-structured, with each middleware in its own file accompanied by thorough tests. Key improvements include replacing the generic JSON-based clone function with more efficient and type-safe Clone methods, and a new context-based mechanism for middleware to access the Genkit registry. The overall changes significantly enhance the framework's flexibility and robustness. I have a couple of minor suggestions for improving code clarity and hygiene.


// GenerateWithRequest is the central generation implementation for ai.Generate(), prompt.Execute(), and the GenerateAction direct call.
func GenerateWithRequest(ctx context.Context, r api.Registry, opts *GenerateActionOptions, mw []ModelMiddleware, cb ModelStreamCallback) (*ModelResponse, error) {
func GenerateWithRequest(ctx context.Context, r api.Registry, opts *GenerateActionOptions, mmws []ModelMiddleware, cb ModelStreamCallback) (*ModelResponse, error) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The parameter name mmws is a bit cryptic. While it was likely changed from mw to avoid shadowing, a more descriptive name like modelMiddlewares would improve readability for future maintainers. This change would need to be applied to its usages within this function as well.

Comment on lines 843 to 844
// clone creates a deep copy of the provided object using JSON marshaling and unmarshaling.
func clone[T any](obj *T) *T {
if obj == nil {
return nil
}

bytes, err := json.Marshal(obj)
if err != nil {
panic(fmt.Sprintf("clone: failed to marshal object: %v", err))
}

var newObj T
if err := json.Unmarshal(bytes, &newObj); err != nil {
panic(fmt.Sprintf("clone: failed to unmarshal object: %v", err))
}

return &newObj
}

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

This comment describes the clone function which has been removed. This dangling comment should be removed to avoid confusion.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

docs Improvements or additions to documentation go js python Python tooling

Projects

Status: No status

Development

Successfully merging this pull request may close these issues.

1 participant