From b76cd67b2606c32f11345b9a5a657ca9ed45cb90 Mon Sep 17 00:00:00 2001 From: Alex Pascal Date: Tue, 3 Feb 2026 12:53:39 -0800 Subject: [PATCH 1/4] Create go-middleware-design.md --- docs/go-middleware-design.md | 349 +++++++++++++++++++++++++++++++++++ 1 file changed, 349 insertions(+) create mode 100644 docs/go-middleware-design.md diff --git a/docs/go-middleware-design.md b/docs/go-middleware-design.md new file mode 100644 index 0000000000..028daf2d87 --- /dev/null +++ b/docs/go-middleware-design.md @@ -0,0 +1,349 @@ +# Middleware Design for Genkit Go + +## Problem + +The current `ModelMiddleware` only wraps the raw model call. We need middleware that can: + +1. Wrap the entire generation (including tool loop) +2. Wrap individual tool executions +3. Share state across these hooks within a single `ai.Generate()` invocation + +## Design + +### Core Interface + +```go +// Middleware provides hooks for different stages of generation. +type Middleware interface { + // New returns a fresh instance for each ai.Generate() call, enabling per-invocation state. + New() Middleware + // Generate wraps each iteration of the tool loop. + Generate(ctx context.Context, state *GenerateState, next GenerateNext) (*ModelResponse, error) + // Model wraps each model API call. + Model(ctx context.Context, state *ModelState, next ModelNext) (*ModelResponse, error) + // Tool wraps each tool execution. + Tool(ctx context.Context, state *ToolState, next ToolNext) (*ToolResponse, error) +} + +// State structs - can be extended without breaking the interface +type GenerateState struct { + Options *GenerateActionOptions // original options passed to ai.Generate() + Request *ModelRequest // current request for this iteration (has accumulated messages) + Iteration int // current loop iteration (0-indexed) +} + +type ModelState struct { + Request *ModelRequest + Callback ModelStreamCallback +} + +type ToolState struct { + Request *ToolRequest + Tool Tool // provides Name(), Definition(), etc. +} + +// Next function types for each hook +type GenerateNext func(ctx context.Context, state *GenerateState) (*ModelResponse, error) +type ModelNext func(ctx context.Context, state *ModelState) (*ModelResponse, error) +type ToolNext func(ctx context.Context, state *ToolState) (*ToolResponse, error) +``` + +Note: `GenerateActionOptions`, `ModelRequest`, `ModelResponse`, `ToolRequest`, `ToolResponse`, and `ToolDefinition` are existing types in `ai/gen.go`. + +### Base Implementation + +Embed this to get default pass-through behavior for hooks you don't need. You must still implement `New()` yourself. + +```go +type BaseMiddleware struct{} + +func (b *BaseMiddleware) Generate(ctx context.Context, state *GenerateState, next GenerateNext) (*ModelResponse, error) { + return next(ctx, state) +} + +func (b *BaseMiddleware) Model(ctx context.Context, state *ModelState, next ModelNext) (*ModelResponse, error) { + return next(ctx, state) +} + +func (b *BaseMiddleware) Tool(ctx context.Context, state *ToolState, next ToolNext) (*ToolResponse, error) { + return next(ctx, state) +} +``` + +### Usage + +```go +resp, err := ai.Generate(ctx, r, + ai.WithModel(myModel), + ai.WithPrompt("Hello"), + ai.WithUse(RetryMiddleware{MaxRetries: 3}), +) +``` + +### Registration + +```go +// Registers middleware, exposing schema to Dev UI +ai.DefineMiddlewareFor[RetryMiddleware](r, "retry", "Retries failed tool calls") +``` + +Registration stores metadata and a factory that knows how to create middleware from raw JSON config: + +```go +type MiddlewareDesc struct { + Name string `json:"name"` + Description string `json:"description"` + ConfigSchema map[string]any `json:"configSchema"` + createFromConfig func(configJSON []byte) (Middleware, error) // internal, not serialized +} + +// Register registers the middleware with the registry. +func (d *MiddlewareDesc) Register(r api.Registry) { + r.RegisterValue("/middleware/"+d.Name, d) +} + +// NewMiddlewareFor creates a middleware descriptor without registering it. +func NewMiddlewareFor[Config Middleware](name, description string) *MiddlewareDesc { + return &MiddlewareDesc{ + Name: name, + Description: description, + ConfigSchema: core.InferSchemaMap(*new(Config)), + createFromConfig: func(configJSON []byte) (Middleware, error) { + var cfg Config + if err := json.Unmarshal(configJSON, &cfg); err != nil { + return nil, err + } + return cfg.New(), nil + }, + } +} + +// DefineMiddlewareFor creates and registers a middleware descriptor. +func DefineMiddlewareFor[Config Middleware](r api.Registry, name, description string) *MiddlewareDesc { + def := NewMiddlewareFor[Config](name, description) + def.Register(r) + return def +} +``` + +The generic `Config` type parameter is captured in the closure, allowing the definition to unmarshal raw JSON into the correct typed config at runtime. + +### MiddlewarePlugin Interface + +Plugins that provide middleware implement `MiddlewarePlugin`: + +```go +type MiddlewarePlugin interface { + ListMiddleware(ctx context.Context) []*MiddlewareDesc +} +``` + +During `genkit.Init()`, the framework calls `ListMiddleware` on plugins that implement this interface and registers each returned descriptor. + +**Plugin example:** +```go +func (p *MyPlugin) ListMiddleware(ctx context.Context) []*ai.MiddlewareDesc { + return []*ai.MiddlewareDesc{ + ai.NewMiddlewareFor[TracingMiddleware]("tracing", "Distributed tracing"), + ai.NewMiddlewareFor[RetryMiddleware]("retry", "Retries failed operations"), + } +} +``` + +### API/Schema Integration + +To support middleware via the API and Dev UI, we add schemas for middleware descriptors and references. + +**Zod schema (genkit-tools/common/src/types/middleware.ts):** +```ts +import { z } from 'zod'; +import { JSONSchema7Schema } from './action'; + +/** Descriptor for a registered middleware, returned by reflection API. */ +export const MiddlewareDescSchema = z.object({ + /** Unique name of the middleware. */ + name: z.string(), + /** Human-readable description of what the middleware does. */ + description: z.string().optional(), + /** JSON Schema for the middleware's configuration. */ + configSchema: JSONSchema7Schema.optional(), +}); +export type MiddlewareDesc = z.infer; + +/** Reference to a registered middleware with optional configuration. */ +export const MiddlewareRefSchema = z.object({ + /** Name of the registered middleware. */ + name: z.string(), + /** Configuration for the middleware (schema defined by the middleware). */ + config: z.any().optional(), +}); +export type MiddlewareRef = z.infer; +``` + +**GenerateActionOptions addition (genkit-tools/common/src/types/model.ts):** +```ts +export const GenerateActionOptionsSchema = z.object({ + // ... existing fields ... + + /** Middleware to apply to this generation. */ + use: z.array(MiddlewareRefSchema).optional(), +}); +``` + +**Go types:** +```go +type GenerateActionOptions struct { + // ... existing fields ... + Use []*MiddlewareRef `json:"use,omitempty"` +} + +type MiddlewareRef struct { + Name string `json:"name"` + Config any `json:"config,omitempty"` +} +``` + +### Reflection API + +New endpoint to list registered middleware: + +``` +GET /api/middlewares → map[string]MiddlewareDesc +``` + +Dev UI uses this to list middleware and render config forms from `configSchema`. + +## Example: Tool Retry Middleware + +```go +type RetryMiddleware struct { + ai.BaseMiddleware + MaxRetries int `json:"maxRetries"` + Backoff time.Duration `json:"backoff"` + totalRetries int // per-invocation state +} + +// New returns a fresh instance with the same config +func (c RetryMiddleware) New() ai.Middleware { + return &RetryMiddleware{ + MaxRetries: c.MaxRetries, + Backoff: c.Backoff, + } +} + +// Override Generate to log iteration info +func (m *RetryMiddleware) Generate(ctx context.Context, state *ai.GenerateState, next ai.GenerateNext) (*ai.ModelResponse, error) { + if state.Iteration == 0 { + log.Printf("Starting generation with model %s", state.Options.Model) + } + resp, err := next(ctx, state) + if m.totalRetries > 0 { + log.Printf("Iteration %d complete, total retries so far: %d", state.Iteration, m.totalRetries) + } + return resp, err +} + +// Override Tool to add retry logic +func (m *RetryMiddleware) Tool(ctx context.Context, state *ai.ToolState, next ai.ToolNext) (*ai.ToolResponse, error) { + var lastErr error + for attempt := 0; attempt <= m.MaxRetries; attempt++ { + if attempt > 0 { + m.totalRetries++ + time.Sleep(m.Backoff * time.Duration(attempt)) + } + resp, err := next(ctx, state) + if err == nil { + return resp, nil + } + lastErr = err + } + return nil, lastErr +} +``` + +**Usage:** + +```go +resp, err := ai.Generate(ctx, r, + ai.WithModel(myModel), + ai.WithTools(weatherTool), + ai.WithPrompt("What's the weather?"), + ai.WithUse(RetryMiddleware{MaxRetries: 3, Backoff: time.Second}), +) +``` + +## Multiple Middleware + +```go +ai.WithUse( + LoggingMiddleware{}, + RetryMiddleware{MaxRetries: 3}, +) +``` + +First middleware is outermost: `Logging.Tool(Retry.Tool(actual))`. + +## Plugin Integration + +Plugins implement `MiddlewarePlugin` to register middleware automatically, and provide factory methods to inject dependencies: + +```go +type TracingMiddleware struct { + ai.BaseMiddleware + SampleRate float64 `json:"sampleRate"` + exporter trace.Exporter // injected by plugin, not serializable +} + +func (c TracingMiddleware) New() ai.Middleware { + return &TracingMiddleware{SampleRate: c.SampleRate, exporter: c.exporter} +} + +func (m *TracingMiddleware) Model(ctx context.Context, state *ai.ModelState, next ai.ModelNext) (*ai.ModelResponse, error) { + ctx, span := m.exporter.Start(ctx, "model-call") + defer span.End() + return next(ctx, state) +} + +// Plugin implements MiddlewarePlugin to register middleware +func (p *MyPlugin) ListMiddleware(ctx context.Context) []*ai.MiddlewareDesc { + return []*ai.MiddlewareDesc{ + ai.NewMiddlewareFor[TracingMiddleware]("tracing", "Distributed tracing"), + } +} + +// Plugin provides a factory that injects its exporter +func (p *MyPlugin) Tracing(cfg TracingMiddleware) TracingMiddleware { + cfg.exporter = p.exporter + return cfg +} +``` + +**Usage:** + +```go +// Use plugin's factory to get middleware with injected dependencies +ai.WithUse(myPlugin.Tracing(TracingMiddleware{SampleRate: 0.1})) +``` + +## Deprecation of Existing Middleware + +The existing `core.Middleware` and `ai.ModelMiddleware` types will be deprecated: + +**This release:** +- Deprecate `core.Middleware`, `ai.ModelMiddleware`, and `WithMiddleware()` option +- Refactor internal middleware to use new `Middleware` type: + - `addAutomaticTelemetry()` → uses Model hook + - `simulateSystemPrompt()` → uses Model hook + - `validateSupport()` → uses Model hook + - `augmentWithContext()` → uses Model hook + - `DownloadRequestMedia()` → uses Model hook +- New user-facing middleware uses `Middleware` exclusively + +**Next major version:** +- Remove deprecated `core.Middleware` and `ai.ModelMiddleware` types + +## Implementation Notes + +- `New()` called once at start of each `ai.Generate()` +- Tool hooks run in parallel—implementations must be thread-safe if mutating shared state +- Hooks that just pass through can embed `BaseMiddleware` or return `next()` directly From 42b8bd6bd766e9f4832b6e6d5d51ad4081906922 Mon Sep 17 00:00:00 2001 From: Alex Pascal Date: Thu, 5 Feb 2026 11:06:09 -0800 Subject: [PATCH 2/4] Update go-middleware-design.md --- docs/go-middleware-design.md | 83 +++++++++++++++++------------------- 1 file changed, 39 insertions(+), 44 deletions(-) diff --git a/docs/go-middleware-design.md b/docs/go-middleware-design.md index 028daf2d87..d4444c4e3b 100644 --- a/docs/go-middleware-design.md +++ b/docs/go-middleware-design.md @@ -15,6 +15,8 @@ The current `ModelMiddleware` only wraps the raw model call. We need middleware ```go // Middleware provides hooks for different stages of generation. type Middleware interface { + // Name returns the middleware's unique identifier. + Name() string // New returns a fresh instance for each ai.Generate() call, enabling per-invocation state. New() Middleware // Generate wraps each iteration of the tool loop. @@ -52,7 +54,7 @@ Note: `GenerateActionOptions`, `ModelRequest`, `ModelResponse`, `ToolRequest`, ` ### Base Implementation -Embed this to get default pass-through behavior for hooks you don't need. You must still implement `New()` yourself. +Embed this to get default pass-through behavior for hooks you don't need. You must still implement `Name()` and `New()` yourself. ```go type BaseMiddleware struct{} @@ -83,51 +85,45 @@ resp, err := ai.Generate(ctx, r, ### Registration ```go -// Registers middleware, exposing schema to Dev UI -ai.DefineMiddlewareFor[RetryMiddleware](r, "retry", "Retries failed tool calls") +ai.DefineMiddleware(r, "Retries failed tool calls", &RetryMiddleware{}) ``` -Registration stores metadata and a factory that knows how to create middleware from raw JSON config: - ```go type MiddlewareDesc struct { - Name string `json:"name"` - Description string `json:"description"` - ConfigSchema map[string]any `json:"configSchema"` - createFromConfig func(configJSON []byte) (Middleware, error) // internal, not serialized + Name string `json:"name"` + Description string `json:"description,omitempty"` + ConfigSchema map[string]any `json:"configSchema,omitempty"` + configFromJSON func([]byte) (Middleware, error) // not serialized } -// Register registers the middleware with the registry. func (d *MiddlewareDesc) Register(r api.Registry) { r.RegisterValue("/middleware/"+d.Name, d) } -// NewMiddlewareFor creates a middleware descriptor without registering it. -func NewMiddlewareFor[Config Middleware](name, description string) *MiddlewareDesc { +func NewMiddleware[T Middleware](description string, prototype T) *MiddlewareDesc { return &MiddlewareDesc{ - Name: name, + Name: prototype.Name(), Description: description, - ConfigSchema: core.InferSchemaMap(*new(Config)), - createFromConfig: func(configJSON []byte) (Middleware, error) { - var cfg Config - if err := json.Unmarshal(configJSON, &cfg); err != nil { - return nil, err + ConfigSchema: core.InferSchemaMap(*new(T)), + configFromJSON: func(configJSON []byte) (Middleware, error) { + inst := prototype.New() + if len(configJSON) > 0 { + if err := json.Unmarshal(configJSON, inst); err != nil { + return nil, err + } } - return cfg.New(), nil + return inst, nil }, } } -// DefineMiddlewareFor creates and registers a middleware descriptor. -func DefineMiddlewareFor[Config Middleware](r api.Registry, name, description string) *MiddlewareDesc { - def := NewMiddlewareFor[Config](name, description) - def.Register(r) - return def +func DefineMiddleware[T Middleware](r api.Registry, description string, prototype T) *MiddlewareDesc { + d := NewMiddleware(description, prototype) + d.Register(r) + return d } ``` -The generic `Config` type parameter is captured in the closure, allowing the definition to unmarshal raw JSON into the correct typed config at runtime. - ### MiddlewarePlugin Interface Plugins that provide middleware implement `MiddlewarePlugin`: @@ -144,8 +140,7 @@ During `genkit.Init()`, the framework calls `ListMiddleware` on plugins that imp ```go func (p *MyPlugin) ListMiddleware(ctx context.Context) []*ai.MiddlewareDesc { return []*ai.MiddlewareDesc{ - ai.NewMiddlewareFor[TracingMiddleware]("tracing", "Distributed tracing"), - ai.NewMiddlewareFor[RetryMiddleware]("retry", "Retries failed operations"), + ai.NewMiddleware("Distributed tracing", &TracingMiddleware{exporter: p.exporter}), } } ``` @@ -223,11 +218,12 @@ type RetryMiddleware struct { totalRetries int // per-invocation state } -// New returns a fresh instance with the same config -func (c RetryMiddleware) New() ai.Middleware { +func (m *RetryMiddleware) Name() string { return "retry" } + +func (m *RetryMiddleware) New() ai.Middleware { return &RetryMiddleware{ - MaxRetries: c.MaxRetries, - Backoff: c.Backoff, + MaxRetries: m.MaxRetries, + Backoff: m.Backoff, } } @@ -294,8 +290,10 @@ type TracingMiddleware struct { exporter trace.Exporter // injected by plugin, not serializable } -func (c TracingMiddleware) New() ai.Middleware { - return &TracingMiddleware{SampleRate: c.SampleRate, exporter: c.exporter} +func (m *TracingMiddleware) Name() string { return "tracing" } + +func (m *TracingMiddleware) New() ai.Middleware { + return &TracingMiddleware{exporter: m.exporter} } func (m *TracingMiddleware) Model(ctx context.Context, state *ai.ModelState, next ai.ModelNext) (*ai.ModelResponse, error) { @@ -304,25 +302,22 @@ func (m *TracingMiddleware) Model(ctx context.Context, state *ai.ModelState, nex return next(ctx, state) } -// Plugin implements MiddlewarePlugin to register middleware +// Plugin registers middleware with injected dependencies via prototype func (p *MyPlugin) ListMiddleware(ctx context.Context) []*ai.MiddlewareDesc { return []*ai.MiddlewareDesc{ - ai.NewMiddlewareFor[TracingMiddleware]("tracing", "Distributed tracing"), + ai.NewMiddleware("Distributed tracing", &TracingMiddleware{exporter: p.exporter}), } } - -// Plugin provides a factory that injects its exporter -func (p *MyPlugin) Tracing(cfg TracingMiddleware) TracingMiddleware { - cfg.exporter = p.exporter - return cfg -} ``` **Usage:** ```go -// Use plugin's factory to get middleware with injected dependencies -ai.WithUse(myPlugin.Tracing(TracingMiddleware{SampleRate: 0.1})) +// Inline: exporter comes from prototype via New() +ai.WithUse(TracingMiddleware{SampleRate: 0.1}) + +// Dev UI: sends {"name": "tracing", "config": {"sampleRate": 0.1}} +// → configFromJSON calls prototype.New() (preserves exporter), unmarshals config on top ``` ## Deprecation of Existing Middleware From ce8dc7eb19fe6b4e575a43b269e5b92b8bedeb65 Mon Sep 17 00:00:00 2001 From: Alex Pascal Date: Thu, 5 Feb 2026 11:42:29 -0800 Subject: [PATCH 3/4] Update go-middleware-design.md --- docs/go-middleware-design.md | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/docs/go-middleware-design.md b/docs/go-middleware-design.md index d4444c4e3b..df3e8312b4 100644 --- a/docs/go-middleware-design.md +++ b/docs/go-middleware-design.md @@ -130,7 +130,7 @@ Plugins that provide middleware implement `MiddlewarePlugin`: ```go type MiddlewarePlugin interface { - ListMiddleware(ctx context.Context) []*MiddlewareDesc + ListMiddleware(ctx context.Context) ([]*MiddlewareDesc, error) } ``` @@ -138,10 +138,10 @@ During `genkit.Init()`, the framework calls `ListMiddleware` on plugins that imp **Plugin example:** ```go -func (p *MyPlugin) ListMiddleware(ctx context.Context) []*ai.MiddlewareDesc { +func (p *MyPlugin) ListMiddleware(ctx context.Context) ([]*ai.MiddlewareDesc, error) { return []*ai.MiddlewareDesc{ ai.NewMiddleware("Distributed tracing", &TracingMiddleware{exporter: p.exporter}), - } + }, nil } ``` @@ -303,10 +303,10 @@ func (m *TracingMiddleware) Model(ctx context.Context, state *ai.ModelState, nex } // Plugin registers middleware with injected dependencies via prototype -func (p *MyPlugin) ListMiddleware(ctx context.Context) []*ai.MiddlewareDesc { +func (p *MyPlugin) ListMiddleware(ctx context.Context) ([]*ai.MiddlewareDesc, error) { return []*ai.MiddlewareDesc{ ai.NewMiddleware("Distributed tracing", &TracingMiddleware{exporter: p.exporter}), - } + }, nil } ``` From c0f8c3632967f02d0f07d2a4f8e21d7869ed2797 Mon Sep 17 00:00:00 2001 From: Alex Pascal Date: Thu, 5 Feb 2026 13:23:13 -0800 Subject: [PATCH 4/4] Update go-middleware-design.md --- docs/go-middleware-design.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/go-middleware-design.md b/docs/go-middleware-design.md index df3e8312b4..baa9807ae0 100644 --- a/docs/go-middleware-design.md +++ b/docs/go-middleware-design.md @@ -200,10 +200,10 @@ type MiddlewareRef struct { ### Reflection API -New endpoint to list registered middleware: +General values endpoint (matches JS), filtered by type: ``` -GET /api/middlewares → map[string]MiddlewareDesc +GET /api/values?type=middleware → map[string]MiddlewareDesc ``` Dev UI uses this to list middleware and render config forms from `configSchema`.