refactor(py): widen wrap_model middleware signature to ModelResponse | Operation - #6022
refactor(py): widen wrap_model middleware signature to ModelResponse | Operation#6022huangjeff5 wants to merge 2 commits into
Conversation
![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 updates the wrap_model method in _middleware.py to support background operations by widening its return type and its next_fn parameter's return type to ModelResponse | Operation. The reviewer notes that wrap_generate and its next_fn should also be widened to ModelResponse | Operation to prevent static type-checker errors, as background operations returned by wrap_model will propagate up through wrap_generate.
| next_fn: Callable[[ModelHookParams, GenerateMiddlewareContext], Awaitable[ModelResponse | Operation]], | ||
| ) -> ModelResponse | Operation: |
There was a problem hiding this comment.
While widening wrap_model to return ModelResponse | Operation is necessary to support background operations, wrap_generate (defined at line 258) and its next_fn are still typed to return only ModelResponse. Since wrap_generate wraps the tool loop iteration which executes the model call, any background operation returned by wrap_model will propagate up through wrap_generate's next_fn. This will cause static type-checker errors because next_fn is expected to return ModelResponse but will actually return ModelResponse | Operation at runtime. To ensure type safety across the entire pipeline, wrap_generate in both BaseMiddleware and MiddlewareDef should also be widened to return ModelResponse | Operation.
huangjeff5
changed the title
Coerce persisted generate() dumps, raise INVALID_ARGUMENT for bad Veo config, advertise the remaining GenerateVideosConfig fields, and default background models to longRunning. wrap_model next() stays a bare Operation; genkit-ai#6022 should drop its assert. Co-authored-by: jeffdh5 <jeffdh5@users.noreply.github.com>
1 participant
Product Context & Motivation
Genkit supports both standard synchronous/streaming models and long-running background models (such as video generation with Veo or Deep Research). Model middleware (such as logging, telemetry, or safety filters) needs to operate seamlessly across both foreground model responses (
ModelResponse) and long-running operations (Operation[ModelResponse, Any]).This PR widens the
wrap_modelmiddleware type annotations to allow background operation handles to pass through middleware pipelines cleanly without static type-checker errors.Key Changes
wrap_modeltype annotations ingenkit._core._middlewareto acceptModelResponse | Operation[ModelResponse, Any].Verification & Safety
ModelResponseremain fully valid.pytest py/packages/genkit/tests/genkit/ai/(537 passed).