feat(js): implemented generate middleware

js/ai/src/generate-api.ts

@@ -0,0 +1,354 @@
+/**
+ * Copyright 2026 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 {
+  getContext,
+  run,
+  z,
+  type ActionContext,
+  type Operation,
+} from '@genkit-ai/core';
+import { type Registry } from '@genkit-ai/core/registry';
+import { cancelOperation } from './cancel-operation.js';
+import { checkOperation } from './check-operation.js';
+import { type DocumentData } from './document.js';
+import {
+  embed,
+  embedMany,
+  type EmbedderArgument,
+  type EmbedderParams,
+  type Embedding,
+  type EmbeddingBatch,
+} from './embedder.js';
+import {
+  generate,
+  generateStream,
+  type GenerateOptions,
+  type GenerateResponse,
+  type GenerateStreamOptions,
+  type GenerateStreamResponse,
+} from './generate.js';
+import { GenerationCommonConfigSchema, type Part } from './model-types.js';
+
+/**
+ * `GenerateAPI` encapsulates model generate APIs.
+ */
+export class GenerateAPI {
+  readonly registry: Registry;
+
+  constructor(registry: Registry) {
+    this.registry = registry;
+  }
+
+  /**
+   * Embeds the given `content` using the specified `embedder`.
+   */
+  embed<CustomOptions extends z.ZodTypeAny>(
+    params: EmbedderParams<CustomOptions>
+  ): Promise<Embedding[]> {
+    return embed(this.registry, params);
+  }
+
+  /**
+   * A veneer for interacting with embedder models in bulk.
+   */
+  embedMany<ConfigSchema extends z.ZodTypeAny = z.ZodTypeAny>(params: {
+    embedder: EmbedderArgument<ConfigSchema>;
+    content: string[] | DocumentData[];
+    metadata?: Record<string, unknown>;
+    options?: z.infer<ConfigSchema>;
+  }): Promise<EmbeddingBatch> {
+    return embedMany(this.registry, params);
+  }
+
+  /**
+   * Make a generate call to the default model with a simple text prompt.
+   *
+   * ```ts
+   * const ai = genkit({
+   *   plugins: [googleAI()],
+   *   model: googleAI.model('gemini-2.5-flash'), // default model
+   * })
+   *
+   * const { text } = await ai.generate('hi');
+   * ```
+   */
+  generate<O extends z.ZodTypeAny = z.ZodTypeAny>(
+    strPrompt: string
+  ): Promise<GenerateResponse<z.infer<O>>>;
+
+  /**
+   * Make a generate call to the default model with a multipart request.
+   *
+   * ```ts
+   * const ai = genkit({
+   *   plugins: [googleAI()],
+   *   model: googleAI.model('gemini-2.5-flash'), // default model
+   * })
+   *
+   * const { text } = await ai.generate([
+   *   { media: {url: 'http://....'} },
+   *   { text: 'describe this image' }
+   * ]);
+   * ```
+   */
+  generate<O extends z.ZodTypeAny = z.ZodTypeAny>(
+    parts: Part[]
+  ): Promise<GenerateResponse<z.infer<O>>>;
+
+  /**
+   * Generate calls a generative model based on the provided prompt and configuration. If
+   * `messages` is provided, the generation will include a conversation history in its
+   * request. If `tools` are provided, the generate method will automatically resolve
+   * tool calls returned from the model unless `returnToolRequests` is set to `true`.
+   *
+   * See {@link GenerateOptions} for detailed information about available options.
+   *
+   * ```ts
+   * const ai = genkit({
+   *   plugins: [googleAI()],
+   * })
+   *
+   * const { text } = await ai.generate({
+   *   system: 'talk like a pirate',
+   *   prompt: [
+   *     { media: { url: 'http://....' } },
+   *     { text: 'describe this image' }
+   *   ],
+   *   messages: conversationHistory,
+   *   tools: [ userInfoLookup ],
+   *   model: googleAI.model('gemini-2.5-flash'),
+   * });
+   * ```
+   */
+  generate<
+    O extends z.ZodTypeAny = z.ZodTypeAny,
+    CustomOptions extends z.ZodTypeAny = typeof GenerationCommonConfigSchema,
+  >(
+    opts:
+      | GenerateOptions<O, CustomOptions>
+      | PromiseLike<GenerateOptions<O, CustomOptions>>
+  ): Promise<GenerateResponse<z.infer<O>>>;
+
+  async generate<
+    O extends z.ZodTypeAny = z.ZodTypeAny,
+    CustomOptions extends z.ZodTypeAny = typeof GenerationCommonConfigSchema,
+  >(
+    options:
+      | string
+      | Part[]
+      | GenerateOptions<O, CustomOptions>
+      | PromiseLike<GenerateOptions<O, CustomOptions>>
+  ): Promise<GenerateResponse<z.infer<O>>> {
+    let resolvedOptions: GenerateOptions<O, CustomOptions>;
+    if (options instanceof Promise) {
+      resolvedOptions = await options;
+    } else if (typeof options === 'string' || Array.isArray(options)) {
+      resolvedOptions = {
+        prompt: options,
+      };
+    } else {
+      resolvedOptions = options as GenerateOptions<O, CustomOptions>;
+    }
+    return generate(this.registry, resolvedOptions);
+  }
+
+  /**
+   * Make a streaming generate call to the default model with a simple text prompt.
+   *
+   * ```ts
+   * const ai = genkit({
+   *   plugins: [googleAI()],
+   *   model: googleAI.model('gemini-2.5-flash'), // default model
+   * })
+   *
+   * const { response, stream } = ai.generateStream('hi');
+   * for await (const chunk of stream) {
+   *   console.log(chunk.text);
+   * }
+   * console.log((await response).text);
+   * ```
+   */
+  generateStream<O extends z.ZodTypeAny = z.ZodTypeAny>(
+    strPrompt: string
+  ): GenerateStreamResponse<z.infer<O>>;
+
+  /**
+   * Make a streaming generate call to the default model with a multipart request.
+   *
+   * ```ts
+   * const ai = genkit({
+   *   plugins: [googleAI()],
+   *   model: googleAI.model('gemini-2.5-flash'), // default model
+   * })
+   *
+   * const { response, stream } = ai.generateStream([
+   *   { media: {url: 'http://....'} },
+   *   { text: 'describe this image' }
+   * ]);
+   * for await (const chunk of stream) {
+   *   console.log(chunk.text);
+   * }
+   * console.log((await response).text);
+   * ```
+   */
+  generateStream<O extends z.ZodTypeAny = z.ZodTypeAny>(
+    parts: Part[]
+  ): GenerateStreamResponse<z.infer<O>>;
+
+  /**
+   * Streaming generate calls a generative model based on the provided prompt and configuration. If
+   * `messages` is provided, the generation will include a conversation history in its
+   * request. If `tools` are provided, the generate method will automatically resolve
+   * tool calls returned from the model unless `returnToolRequests` is set to `true`.
+   *
+   * See {@link GenerateOptions} for detailed information about available options.
+   *
+   * ```ts
+   * const ai = genkit({
+   *   plugins: [googleAI()],
+   * })
+   *
+   * const { response, stream } = ai.generateStream({
+   *   system: 'talk like a pirate',
+   *   prompt: [
+   *     { media: { url: 'http://....' } },
+   *     { text: 'describe this image' }
+   *   ],
+   *   messages: conversationHistory,
+   *   tools: [ userInfoLookup ],
+   *   model: googleAI.model('gemini-2.5-flash'),
+   * });
+   * for await (const chunk of stream) {
+   *   console.log(chunk.text);
+   * }
+   * console.log((await response).text);
+   * ```
+   */
+  generateStream<
+    O extends z.ZodTypeAny = z.ZodTypeAny,
+    CustomOptions extends z.ZodTypeAny = typeof GenerationCommonConfigSchema,
+  >(
+    parts:
+      | GenerateOptions<O, CustomOptions>
+      | PromiseLike<GenerateOptions<O, CustomOptions>>
+  ): GenerateStreamResponse<z.infer<O>>;
+
+  generateStream<
+    O extends z.ZodTypeAny = z.ZodTypeAny,
+    CustomOptions extends z.ZodTypeAny = typeof GenerationCommonConfigSchema,
+  >(
+    options:
+      | string
+      | Part[]
+      | GenerateStreamOptions<O, CustomOptions>
+      | PromiseLike<GenerateStreamOptions<O, CustomOptions>>
+  ): GenerateStreamResponse<z.infer<O>> {
+    if (typeof options === 'string' || Array.isArray(options)) {
+      options = { prompt: options };
+    }
+    return generateStream(this.registry, options);
+  }
+
+  /**
+   * Checks the status of of a given operation. Returns a new operation which will contain the updated status.
+   *
+   * ```ts
+   * let operation = await ai.generateOperation({
+   *   model: googleAI.model('veo-2.0-generate-001'),
+   *   prompt: 'A banana riding a bicycle.',
+   * });
+   *
+   * while (!operation.done) {
+   *   operation = await ai.checkOperation(operation!);
+   *   await new Promise((resolve) => setTimeout(resolve, 5000));
+   * }
+   * ```
+   *
+   * @param operation
+   * @returns
+   */
+  checkOperation<T>(operation: Operation<T>): Promise<Operation<T>> {
+    return checkOperation(this.registry, operation);
+  }
+
+  /**
+   * Cancels a given operation. Returns a new operation which will contain the updated status.
+   *
+   * @param operation
+   * @returns
+   */
+  cancelOperation<T>(operation: Operation<T>): Promise<Operation<T>> {
+    return cancelOperation(this.registry, operation);
+  }
+
+  /**
+   * A flow step that executes the provided function. Each run step is recorded separately in the trace.
+   *
+   * ```ts
+   * ai.defineFlow('hello', async() => {
+   *   await ai.run('step1', async () => {
+   *     // ... step 1
+   *   });
+   *   await ai.run('step2', async () => {
+   *     // ... step 2
+   *   });
+   *   return result;
+   * })
+   * ```
+   */
+  run<T>(name: string, func: () => Promise<T>): Promise<T>;
+
+  /**
+   * A flow step that executes the provided function. Each run step is recorded separately in the trace.
+   *
+   * ```ts
+   * ai.defineFlow('hello', async() => {
+   *   await ai.run('step1', async () => {
+   *     // ... step 1
+   *   });
+   *   await ai.run('step2', async () => {
+   *     // ... step 2
+   *   });
+   *   return result;
+   * })
+   */
+  run<T>(
+    name: string,
+    input: any,
+    func: (input?: any) => Promise<T>
+  ): Promise<T>;
+
+  run<T>(
+    name: string,
+    funcOrInput: () => Promise<T> | any,
+    maybeFunc?: (input?: any) => Promise<T>
+  ): Promise<T> {
+    if (maybeFunc) {
+      return run(name, funcOrInput, maybeFunc, this.registry);
+    }
+    return run(name, funcOrInput, this.registry);
+  }
+
+  /**
+   * Returns current action (or flow) invocation context. Can be used to access things like auth
+   * data set by HTTP server frameworks. If invoked outside of an action (e.g. flow or tool) will
+   * return `undefined`.
+   */
+  currentContext(): ActionContext | undefined {
+    return getContext();
+  }
+}