Moved middleware up to actions

js/ai/src/model.ts

@@ -18,6 +18,7 @@ import {
   Action,
   defineAction,
   getStreamingCallback,
+  Middleware,
   StreamingCallback,
 } from '@genkit-ai/core';
 import { toJsonSchema } from '@genkit-ai/core/schema';
@@ -234,38 +235,10 @@ export type ModelAction<
   __configSchema: CustomOptionsSchema;
 };
 
-export interface ModelMiddleware {
-  (
-    req: GenerateRequest,
-    next: (req?: GenerateRequest) => Promise<GenerateResponseData>
-  ): Promise<GenerateResponseData>;
-}
-
-/**
- *
- */
-export function modelWithMiddleware(
-  model: ModelAction,
-  middleware: ModelMiddleware[]
-): ModelAction {
-  const wrapped = (async (req: GenerateRequest) => {
-    const dispatch = async (index: number, req: GenerateRequest) => {
-      if (index === middleware.length) {
-        // end of the chain, call the original model action
-        return await model(req);
-      }
-
-      const currentMiddleware = middleware[index];
-      return currentMiddleware(req, async (modifiedReq) =>
-        dispatch(index + 1, modifiedReq || req)
-      );
-    };
-
-    return await dispatch(0, req);
-  }) as ModelAction;
-  wrapped.__action = model.__action;
-  return wrapped;
-}
+export type ModelMiddleware = Middleware<
+  z.infer<typeof GenerateRequestSchema>,
+  z.infer<typeof GenerateResponseSchema>
+>;
 
 /**
  * Defines a new model and adds it to the registry.
@@ -291,6 +264,11 @@ export function defineModel<
   ) => Promise<GenerateResponseData>
 ): ModelAction<CustomOptionsSchema> {
   const label = options.label || `${options.name} GenAI model`;
+  const middleware = [
+    ...(options.use || []),
+    validateSupport(options),
+    conformOutput(),
+  ];
   const act = defineAction(
     {
       actionType: 'model',
@@ -308,6 +286,7 @@ export function defineModel<
           supports: options.supports,
         },
       },
+      use: middleware,
     },
     (input) => {
       const startTimeMs = performance.now();
@@ -331,16 +310,7 @@ export function defineModel<
   Object.assign(act, {
     __configSchema: options.configSchema || z.unknown(),
   });
-  const middleware = [
-    ...(options.use || []),
-    validateSupport(options),
-    conformOutput(),
-  ];
-  const ma = modelWithMiddleware(
-    act as ModelAction,
-    middleware
-  ) as ModelAction<CustomOptionsSchema>;
-  return ma;
+  return act as ModelAction<CustomOptionsSchema>;
 }
 
 export interface ModelReference<CustomOptions extends z.ZodTypeAny> {

js/core/src/action.ts

@@ -72,8 +72,40 @@ type ActionParams<
   outputSchema?: O;
   outputJsonSchema?: JSONSchema7;
   metadata?: M;
+  use?: Middleware<z.infer<I>, z.infer<O>>[];
 };
 
+export interface Middleware<I = any, O = any> {
+  (req: I, next: (req?: I) => Promise<O>): Promise<O>;
+}
+
+export function actionWithMiddleware<
+  I extends z.ZodTypeAny,
+  O extends z.ZodTypeAny,
+  M extends Record<string, any> = Record<string, any>,
+>(
+  action: Action<I, O, M>,
+  middleware: Middleware<z.infer<I>, z.infer<O>>[]
+): Action<I, O, M> {
+  const wrapped = (async (req: z.infer<I>) => {
+    const dispatch = async (index: number, req: z.infer<I>) => {
+      if (index === middleware.length) {
+        // end of the chain, call the original model action
+        return await action(req);
+      }
+
+      const currentMiddleware = middleware[index];
+      return currentMiddleware(req, async (modifiedReq) =>
+        dispatch(index + 1, modifiedReq || req)
+      );
+    };
+
+    return await dispatch(0, req);
+  }) as Action<I, O, M>;
+  wrapped.__action = action.__action;
+  return wrapped;
+}
+
 /**
  * Creates an action with the provided config.
  */
@@ -140,6 +172,10 @@ export function action<
     outputJsonSchema: config.outputJsonSchema,
     metadata: config.metadata,
   } as ActionMetadata<I, O, M>;
+
+  if (config.use) {
+    return actionWithMiddleware(actionFn, config.use);
+  }
   return actionFn;
 }
 

js/core/tests/action_test.ts

@@ -0,0 +1,47 @@
+/**
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import assert from 'node:assert';
+import { beforeEach, describe, it } from 'node:test';
+import { z } from 'zod';
+import { action } from '../src/action.js';
+import { __hardResetRegistryForTesting } from '../src/registry.js';
+
+describe('action', () => {
+  beforeEach(__hardResetRegistryForTesting);
+
+  it('applies middleware', async () => {
+    const act = action(
+      {
+        name: 'foo',
+        inputSchema: z.string(),
+        outputSchema: z.number(),
+        use: [
+          async (input, next) => (await next(input + 'middle1')) + 1,
+          async (input, next) => (await next(input + 'middle2')) + 2,
+        ],
+      },
+      async (input) => {
+        return input.length;
+      }
+    );
+
+    assert.strictEqual(
+      await act('foo'),
+      20 // "foomiddle1middle2".length + 1 + 2
+    );
+  });
+});