feat(js): implemented generate middleware

[pavelgj]

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 @pavelgj, 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 introduces a powerful new middleware system for Genkit's generate functionality. This system enables developers to inject custom logic at various stages of the generation process, including before and after the main generate action, model calls, and individual tool executions. It provides a flexible mechanism for cross-cutting concerns like logging, caching, and dynamic tool injection, enhancing the extensibility and observability of AI applications built with Genkit.

Highlights

• Introduction of Generate Middleware: A new system is implemented allowing interception and modification of the generate action, model calls, and tool calls within Genkit AI.
• Middleware Definition and Usage: Developers can now define middleware using the generateMiddleware helper with optional configuration schemas and apply them via the use option in generate calls.
• Middleware Hooks: Middleware can provide specific hooks for different stages: generate (wrapping the entire process), model (intercepting model execution), and tool (intercepting individual tool calls).
• Dynamic Tool Injection: Middleware can dynamically inject tools into the generation request, expanding the capabilities of AI models on the fly.
• Plugin Integration and Backward Compatibility: The new middleware system is fully integrated with the Genkit plugin architecture for global registration and maintains backward compatibility with older functional middleware patterns.
• Reflection API Support: The Genkit reflection server now exposes registered middleware, allowing for introspection and discovery of available 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

• js/ai/src/generate.ts
  • Added imports for new middleware-related types and functions.
  • Updated GenerateOptions.use to accept new middleware types.
  • Introduced normalizeMiddleware to standardize middleware representations.
  • Modified the generate function to create a child registry, normalize and resolve middleware, and pass resolved middleware to generateHelper.
  • Added maybeRegisterDynamicMiddlewareTools to register tools provided by middleware.
  • Refactored maybeRegisterDynamicTools and maybeRegisterDynamicResources to operate on the child registry.
  • Updated toGenerateActionOptions to include use in the generated parameters.
• js/ai/src/generate/action.ts
  • Updated imports for Registry, maybeRegisterDynamicMiddlewareTools, normalizeMiddleware, GenerateMiddlewareDef, and resolveMiddleware.
  • Removed old ModelMiddlewareArgument and ModelMiddlewareWithOptions types.
  • Modified defineGenerateAction to create a child registry, normalize and resolve middleware, register dynamic middleware tools, and pass resolved middleware to generateActionImpl.
  • Renamed generate to generateActionImpl and introduced a new generateActionTurn function.
  • Implemented middleware dispatch logic within generateActionImpl for the generate hook.
  • Updated generateHelper and generateActionImpl to use GenerateMiddlewareDef[] for middleware.
  • Added logic to append tools supplied by middleware to the request.
  • Modified resolveResumeOption to accept middleware.
  • Updated dispatchModel logic to use currentMiddleware.model for model interception.
  • Modified resolveToolRequests to accept tools and middleware arguments.
• js/ai/src/generate/middleware.ts
  • New file: Defines MiddlewareDescSchema, MiddlewareRefSchema, GenerateMiddleware interface, and GenerateMiddlewareDef interface.
  • Introduces generateMiddleware function for creating configurable middleware.
  • Introduces resolveMiddleware function to instantiate middleware definitions from references.
• js/ai/src/generate/resolve-tool-requests.ts
  • Added import for ActionRunOptions and GenerateMiddlewareDef.
  • Modified resolveToolRequest to accept middleware and implemented dispatch logic for the tool hook.
  • Modified resolveToolRequests to accept tools and middleware arguments, and removed registry parameter.
  • Modified resolveResumedToolRequest, resolveResumeOption, and resolveRestartedTools to accept middleware.
• js/ai/src/index.ts
  • Exported new middleware-related types and functions from generate/middleware.ts.
  • Exported GenkitPluginV2.
• js/ai/src/model-types.ts
  • Imported MiddlewareRefSchema.
  • Updated GenerateActionOptionsSchema to include an optional use array of MiddlewareRefSchema.
• js/ai/src/plugin.ts
  • New file: Defines GenkitPluginV2 interface with generateMiddleware method.
• js/ai/tests/generate/action_test.ts
  • Added a new test suite generateAction middleware injection.
  • Added a test case to verify that tools injected via middleware are passed to the model.
• js/ai/tests/generate/middleware_test.ts
  • New file: Comprehensive tests for generateMiddleware functionality.
  • Tests execution order of generate, model, and tool hooks.
  • Tests configuration and old-style functional middleware, and pre-registered middleware via plugins.
  • Tests error handling for incorrect middleware usage, stream interception and modification, and multiple middleware execution order.
  • Tests tool injection and execution via middleware.
• js/core/src/plugin.ts
  • Updated BaseGenkitPluginV2 interface to include generateMiddleware method.
  • Added ResolvableAction type.
• js/core/src/reflection.ts
  • Modified /api/values endpoint to support middleware type.
  • Added logic to call toJson() method on middleware objects if available for serialization.
• js/core/tests/reflection_test.ts
  • New file: Tests for the ReflectionServer API.
  • Tests handling of missing and unsupported type parameters.
  • Tests retrieval of defaultModel values and middleware values, including toJson() serialization.
• js/docs/generate-middleware.md
  • New file: Documentation explaining how to define and use generate middleware, including plugin integration.
• js/genkit/src/common.ts
  • Exported new middleware-related types and functions: MiddlewareDescSchema, MiddlewareRefSchema, generateMiddleware, GenerateMiddleware, GenerateMiddlewareDef, MiddlewareDesc, and MiddlewareRef.
• js/genkit/src/genkit.ts
  • Modified Genkit class to process generateMiddleware from plugins and register them in the registry.
• js/genkit/src/plugin.ts
  • Updated GenkitPluginV2 interface to include generateMiddleware method.
  • Updated GenkitPluginV2Instance to implement generateMiddleware.
• js/genkit/tests/generate_test.ts
  • Added a test case to verify middleware plugin integration with genkit.generate.

Activity

• The author, pavelgj, has implemented a new feature: generate middleware.
• The pull request includes comprehensive unit tests for the new middleware system.
• New documentation for generate middleware has been added.
• The PR title follows conventional commits specification.
• The changes have been manually and unit tested.

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 comprehensive middleware system for the generate function, which is a great feature for adding cross-cutting concerns like logging, caching, and telemetry. The implementation is well-structured, with clear separation of concerns and good test coverage. I've left a few comments on minor improvements and a potential design smell in the plugin interface. Overall, this is a solid contribution.

[pavelgj]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces a powerful middleware system for the generate function, allowing interception of different stages of the generation process (generate, model, tool). This is a significant feature that enables cross-cutting concerns like logging, caching, and telemetry. The implementation is well-structured, with a new generateMiddleware helper, clear definitions for middleware hooks, and support for both new object-based middleware and backward compatibility with old function-based middleware.

The PR also includes a significant and beneficial refactoring, extracting core AI functionalities into a new GenerateAPI base class, which the main Genkit class now extends. This improves code organization and separation of concerns.

The changes are accompanied by comprehensive tests covering various middleware scenarios and documentation to guide users on the new feature.

I have one minor suggestion to improve an error message in the reflection API to make it more accurate. Overall, this is an excellent contribution.

[apascal07]

☹️

What use cases were you thinking of?

[pavelgj]

Some use-cases:

• Automatic Context compression middleware (requires access to an llm)
• LLM classifier based model routing (simple tasks use small model, complex models use large model)
• agentic RAG use-case require embedders

[apascal07]

This is just for backwards compat, right?