feat(go/ai): added DefineMiddleware (Middleware V2)

[apascal07]

Checklist (if applicable):

• PR title is following https://www.conventionalcommits.org/en/v1.0.0/
• Tested (manually, unit tested, etc.)
• Docs updated (updated docs or a docs bug required)

[gemini-code-assist]

Summary of Changes

Hello @apascal07, I'm Gemini Code Assist¹! 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 of the AI generation pipeline by introducing a comprehensive middleware system. It moves beyond simple model-level interception to offer granular control over the entire generation lifecycle, including tool usage and the iterative generation loop. This change provides a powerful mechanism for developers to inject custom logic, such as logging, metrics, or content filtering, at various critical points within the AI workflow, making the system more adaptable and observable.

Highlights

• New Middleware System: Introduced a new, more flexible Middleware interface in Go, replacing the older ModelMiddleware. This new system provides distinct hooks for different stages of the AI generation process: Generate (for the overall tool loop), Model (for individual model API calls), and Tool (for tool executions).
• TypeScript Type Definitions: Added new TypeScript type definitions (MiddlewareDescSchema, MiddlewareRefSchema) to genkit-tools/common/src/types/middleware.ts to support the new middleware concept, allowing middleware to be described and referenced with optional configuration.
• Integration into Generation Options: Integrated the new middleware into GenerateActionOptions in both TypeScript and Go, allowing users to specify middleware to apply to a generation using a new Use field. The older WithMiddleware option is now deprecated.
• Middleware Resolution and Chaining: Implemented logic in go/ai/generate.go to resolve middleware references, unmarshal their configurations, and chain their respective Model, Generate, and Tool hooks around the core generation logic. A new helper runToolWithMiddleware was added for tool hook application.
• Plugin and Reflection API Support: Updated the Genkit Init function to register middleware provided by plugins and added a new reflection API endpoint (GET /api/values?type=middleware) to list registered middleware, enhancing discoverability and extensibility.

🧠 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
  • Exports the newly defined middleware types.
• genkit-tools/common/src/types/middleware.ts
  • Added new file defining TypeScript types for MiddlewareDescSchema and MiddlewareRefSchema using zod for schema validation. These describe middleware descriptors (name, description, config schema) and references (name, config).
• genkit-tools/common/src/types/model.ts
  • Imports MiddlewareRefSchema.
  • Adds a new optional use field (an array of MiddlewareRefSchema) to GenerateActionOptionsSchema, allowing middleware to be specified for generation actions.
• go/ai/gen.go
  • Adds a Use field (slice of *MiddlewareRef) to the GenerateActionOptions struct, mirroring the TypeScript changes and enabling middleware configuration.
• go/ai/generate.go
  • Marks ModelMiddleware as deprecated, advising use of the new Middleware interface with WithUse.
  • Introduces logic to resolve MiddlewareRefs from GenerateActionOptions.Use into executable Middleware handlers.
  • Implements chaining for Model hooks from the new middleware, applying them before legacy ModelMiddleware.
  • Modifies handleToolRequests to accept and pass middlewareHandlers for tool hook processing.
  • Adds chaining for Generate hooks, wrapping the main generation function.
  • Introduces runToolWithMiddleware function to execute tools while applying Tool hooks from the middleware chain.
  • Updates handleResumeOption to correctly propagate the Use field.
• go/ai/middleware.go
  • Added new file defining the Middleware interface with Name(), New(), Generate(), Model(), and Tool() hooks.
  • Defines GenerateState, ModelState, ToolState structs to pass context to middleware hooks.
  • Introduces GenerateNext, ModelNext, ToolNext type aliases for chaining functions.
  • Provides BaseMiddleware for easy implementation of middleware with default pass-through behavior.
  • Defines MiddlewareDesc for registering middleware and helper functions NewMiddleware, DefineMiddleware, LookupMiddleware.
  • Defines MiddlewareRef for serializable middleware references.
  • Introduces MiddlewarePlugin interface for plugins to provide middleware.
• go/ai/middleware_test.go
  • Added new file containing comprehensive unit tests for the new middleware system, covering its definition, lookup, configuration, and hook invocation order for Model and Tool hooks.
• go/ai/option.go
  • Adds a Use field (slice of Middleware) to commonGenOptions.
  • Updates applyCommonGen to handle the new Use option.
  • Deprecates WithMiddleware and introduces WithUse as the new option for applying the enhanced Middleware.
• go/ai/prompt.go
  • Adds logic within Prompt.Execute to register dynamic middleware and build MiddlewareRefs from execOpts.Use, similar to the Generate function.
• go/genkit/genkit.go
  • Modifies the Init function to check if a plugin implements ai.MiddlewarePlugin and, if so, registers its provided middleware descriptors.
• go/genkit/reflection.go
  • Adds a new HTTP handler for GET /api/values to list registered values, specifically supporting filtering by type query parameter (e.g., ?type=middleware).

Activity

• The pull request was created by apascal07.
• The author has completed the checklist items for conventional commit title and testing.

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. ↩

[gemini-code-assist]

Code Review

This pull request introduces a new, more powerful Middleware system for Go, which is a great enhancement. It adds hooks for Generate, Model, and Tool execution stages, deprecating the older ModelMiddleware. The implementation is solid and includes good test coverage. I've identified one bug related to data loss in tool responses and an opportunity to deduplicate some code. Overall, this is a well-executed feature.