feat(js/plugins/compat-oai): allow passing api key at runtime instead of config time (#3946)

Author: pavelgj

js/ai/src/model-types.ts

@@ -201,6 +201,12 @@ export const GenerationCommonConfigSchema = z
         'Set of character sequences (up to 5) that will stop output generation.'
       )
       .optional(),
+    apiKey: z
+      .string()
+      .describe(
+        'API Key to use for the model call, overrides API key provided in plugin config.'
+      )
+      .optional(),
   })
   .passthrough();
 

js/genkit/src/embedder.ts

@@ -17,6 +17,7 @@
 export {
   EmbedderInfoSchema,
   embedderRef,
+  type EmbedRequest,
   type EmbedderAction,
   type EmbedderArgument,
   type EmbedderInfo,

js/plugins/compat-oai/src/audio.ts

@@ -22,13 +22,15 @@ import type {
 import { GenerationCommonConfigSchema, Message, modelRef, z } from 'genkit';
 import type { ModelAction, ModelInfo } from 'genkit/model';
 import { model } from 'genkit/plugin';
-import type OpenAI from 'openai';
+import OpenAI from 'openai';
 import { Response } from 'openai/core.mjs';
 import type {
   SpeechCreateParams,
   Transcription,
   TranscriptionCreateParams,
 } from 'openai/resources/audio/index.mjs';
+import { PluginOptions } from './index.js';
+import { maybeCreateRequestScopedOpenAIClient } from './utils.js';
 
 export type SpeechRequestBuilder = (
   req: GenerateRequest,
@@ -185,10 +187,16 @@ export function defineCompatOpenAISpeechModel<
   client: OpenAI;
   modelRef?: ModelReference<CustomOptions>;
   requestBuilder?: SpeechRequestBuilder;
+  pluginOptions: PluginOptions;
 }): ModelAction {
-  const { name, client, modelRef, requestBuilder } = params;
+  const {
+    name,
+    client: defaultClient,
+    pluginOptions,
+    modelRef,
+    requestBuilder,
+  } = params;
   const modelName = name.substring(name.indexOf('/') + 1);
-
   return model(
     {
       name,
@@ -197,6 +205,11 @@ export function defineCompatOpenAISpeechModel<
     },
     async (request, { abortSignal }) => {
       const ttsRequest = toTTSRequest(modelName!, request, requestBuilder);
+      const client = maybeCreateRequestScopedOpenAIClient(
+        pluginOptions,
+        request,
+        defaultClient
+      );
       const result = await client.audio.speech.create(ttsRequest, {
         signal: abortSignal,
       });
@@ -338,11 +351,17 @@ export function defineCompatOpenAITranscriptionModel<
 >(params: {
   name: string;
   client: OpenAI;
+  pluginOptions?: PluginOptions;
   modelRef?: ModelReference<CustomOptions>;
   requestBuilder?: TranscriptionRequestBuilder;
 }): ModelAction {
-  const { name, client, modelRef, requestBuilder } = params;
-
+  const {
+    name,
+    pluginOptions,
+    client: defaultClient,
+    modelRef,
+    requestBuilder,
+  } = params;
   return model(
     {
       name,
@@ -353,6 +372,11 @@ export function defineCompatOpenAITranscriptionModel<
       const modelName = name.substring(name.indexOf('/') + 1);
 
       const params = toSttRequest(modelName!, request, requestBuilder);
+      const client = maybeCreateRequestScopedOpenAIClient(
+        pluginOptions,
+        request,
+        defaultClient
+      );
       // Explicitly setting stream to false ensures we use the non-streaming overload
       const result = await client.audio.transcriptions.create(
         {

js/plugins/compat-oai/src/deepseek/index.ts

@@ -36,26 +36,25 @@ import {
 
 export type DeepSeekPluginOptions = Omit<PluginOptions, 'name' | 'baseURL'>;
 
-const resolver = async (
-  client: OpenAI,
-  actionType: ActionType,
-  actionName: string
-) => {
-  if (actionType === 'model') {
-    const modelRef = deepSeekModelRef({
-      name: actionName,
-    });
-    return defineCompatOpenAIModel({
-      name: modelRef.name,
-      client,
-      modelRef,
-      requestBuilder: deepSeekRequestBuilder,
-    });
-  } else {
-    logger.warn('Only model actions are supported by the DeepSeek plugin');
-    return undefined;
-  }
-};
+function createResolver(pluginOptions: PluginOptions) {
+  return async (client: OpenAI, actionType: ActionType, actionName: string) => {
+    if (actionType === 'model') {
+      const modelRef = deepSeekModelRef({
+        name: actionName,
+      });
+      return defineCompatOpenAIModel({
+        name: modelRef.name,
+        client,
+        pluginOptions,
+        modelRef,
+        requestBuilder: deepSeekRequestBuilder,
+      });
+    } else {
+      logger.warn('Only model actions are supported by the DeepSeek plugin');
+      return undefined;
+    }
+  };
+}
 
 const listActions = async (client: OpenAI): Promise<ActionMetadata[]> => {
   return await client.models.list().then((response) =>
@@ -87,6 +86,7 @@ export function deepSeekPlugin(
         'Please pass in the API key or set the DEEPSEEK_API_KEY environment variable.',
     });
   }
+  const pluginOptions = { name: 'deepseek', ...options };
   return openAICompatible({
     name: 'deepseek',
     baseURL: 'https://api.deepseek.com',
@@ -97,12 +97,13 @@ export function deepSeekPlugin(
         defineCompatOpenAIModel({
           name: modelRef.name,
           client,
+          pluginOptions,
           modelRef,
           requestBuilder: deepSeekRequestBuilder,
         })
       );
     },
-    resolver,
+    resolver: createResolver(pluginOptions),
     listActions,
   });
 }

js/plugins/compat-oai/src/embedder.ts

@@ -20,6 +20,8 @@
 import type { EmbedderAction, EmbedderReference } from 'genkit';
 import { embedder } from 'genkit/plugin';
 import OpenAI from 'openai';
+import { PluginOptions } from './index.js';
+import { maybeCreateRequestScopedOpenAIClient } from './utils.js';
 
 /**
  * Method to define a new Genkit Embedder that is compatibale with the Open AI
@@ -37,11 +39,11 @@ import OpenAI from 'openai';
 export function defineCompatOpenAIEmbedder(params: {
   name: string;
   client: OpenAI;
+  pluginOptions?: PluginOptions;
   embedderRef?: EmbedderReference;
 }): EmbedderAction {
-  const { name, client, embedderRef } = params;
+  const { name, client: defaultClient, pluginOptions, embedderRef } = params;
   const modelName = name.substring(name.indexOf('/') + 1);
-
   return embedder(
     {
       name,
@@ -50,6 +52,11 @@ export function defineCompatOpenAIEmbedder(params: {
     },
     async (req) => {
       const { encodingFormat: encoding_format, ...restOfConfig } = req.options;
+      const client = maybeCreateRequestScopedOpenAIClient(
+        pluginOptions,
+        req,
+        defaultClient
+      );
       const embeddings = await client.embeddings.create({
         model: modelName!,
         input: req.input.map((d) => d.text),

js/plugins/compat-oai/src/image.ts

@@ -27,6 +27,8 @@ import type {
   ImageGenerateParams,
   ImagesResponse,
 } from 'openai/resources/images.mjs';
+import { PluginOptions } from './index.js';
+import { maybeCreateRequestScopedOpenAIClient } from './utils.js';
 
 export type ImageRequestBuilder = (
   req: GenerateRequest,
@@ -122,10 +124,17 @@ export function defineCompatOpenAIImageModel<
 >(params: {
   name: string;
   client: OpenAI;
+  pluginOptions?: PluginOptions;
   modelRef?: ModelReference<CustomOptions>;
   requestBuilder?: ImageRequestBuilder;
 }): ModelAction<CustomOptions> {
-  const { name, client, modelRef, requestBuilder } = params;
+  const {
+    name,
+    client: defaultClient,
+    pluginOptions,
+    modelRef,
+    requestBuilder,
+  } = params;
   const modelName = name.substring(name.indexOf('/') + 1);
 
   return model(
@@ -135,6 +144,11 @@ export function defineCompatOpenAIImageModel<
       configSchema: modelRef?.configSchema,
     },
     async (request, { abortSignal }) => {
+      const client = maybeCreateRequestScopedOpenAIClient(
+        pluginOptions,
+        request,
+        defaultClient
+      );
       const result = await client.images.generate(
         toImageGenerateParams(modelName!, request, requestBuilder),
         { signal: abortSignal }

js/plugins/compat-oai/src/index.ts

@@ -17,7 +17,7 @@
 import { ActionMetadata } from 'genkit';
 import { ResolvableAction, genkitPluginV2 } from 'genkit/plugin';
 import { ActionType } from 'genkit/registry';
-import { OpenAI, type ClientOptions } from 'openai';
+import OpenAI, { type ClientOptions } from 'openai';
 import { compatOaiModelRef, defineCompatOpenAIModel } from './model.js';
 
 export {
@@ -45,7 +45,8 @@ export {
   type ModelRequestBuilder,
 } from './model.js';
 
-export interface PluginOptions extends Partial<ClientOptions> {
+export interface PluginOptions extends Partial<Omit<ClientOptions, 'apiKey'>> {
+  apiKey?: ClientOptions['apiKey'] | false;
   name: string;
   initializer?: (client: OpenAI) => Promise<ResolvableAction[]>;
   resolver?: (
@@ -110,24 +111,33 @@ export interface PluginOptions extends Partial<ClientOptions> {
  */
 export const openAICompatible = (options: PluginOptions) => {
   let listActionsCache;
+  var client: OpenAI;
+  function createClient() {
+    if (client) return client;
+    const { apiKey, ...restofOptions } = options;
+    client = new OpenAI({
+      ...restofOptions,
+      apiKey: apiKey === false ? 'placeholder' : apiKey,
+    });
+    return client;
+  }
   return genkitPluginV2({
     name: options.name,
     async init() {
       if (!options.initializer) {
         return [];
       }
-      const client = new OpenAI(options);
-      return await options.initializer(client);
+      return await options.initializer(createClient());
     },
     async resolve(actionType: ActionType, actionName: string) {
-      const client = new OpenAI(options);
       if (options.resolver) {
-        return await options.resolver(client, actionType, actionName);
+        return await options.resolver(createClient(), actionType, actionName);
       } else {
         if (actionType === 'model') {
           return defineCompatOpenAIModel({
             name: actionName,
-            client,
+            client: createClient(),
+            pluginOptions: options,
             modelRef: compatOaiModelRef({
               name: actionName,
             }),
@@ -136,14 +146,15 @@ export const openAICompatible = (options: PluginOptions) => {
         return undefined;
       }
     },
-    list: options.listActions
-      ? async () => {
-          if (listActionsCache) return listActionsCache;
-          const client = new OpenAI(options);
-          listActionsCache = await options.listActions!(client);
-          return listActionsCache;
-        }
-      : undefined,
+    list:
+      // Don't attempt to list models if apiKey set to false
+      options.listActions && options.apiKey !== false
+        ? async () => {
+            if (listActionsCache) return listActionsCache;
+            listActionsCache = await options.listActions!(createClient());
+            return listActionsCache;
+          }
+        : undefined,
   });
 };