@genkit-ai/googleai/lib/index.mjs

Genkit AI framework plugin for Google AI APIs, including Gemini APIs.

import {
  embedderActionMetadata,
  embedderRef,
  modelActionMetadata
} from "genkit";
import { logger } from "genkit/logging";
import { modelRef } from "genkit/model";
import { genkitPlugin } from "genkit/plugin";
import { getApiKeyFromEnvVar } from "./common.js";
import {
  SUPPORTED_MODELS as EMBEDDER_MODELS,
  GeminiEmbeddingConfigSchema,
  defineGoogleAIEmbedder,
  textEmbedding004,
  textEmbeddingGecko001
} from "./embedder.js";
import {
  GeminiConfigSchema,
  SUPPORTED_V15_MODELS,
  defineGoogleAIModel,
  gemini,
  gemini10Pro,
  gemini15Flash,
  gemini15Flash8b,
  gemini15Pro,
  gemini20Flash,
  gemini20FlashExp,
  gemini20FlashLite,
  gemini20ProExp0205,
  gemini25FlashPreview0417,
  gemini25ProExp0325,
  gemini25ProPreview0325
} from "./gemini.js";
import { listModels } from "./list-models.js";
async function initializer(ai, options) {
  let apiVersions = ["v1"];
  if (options?.apiVersion) {
    if (Array.isArray(options?.apiVersion)) {
      apiVersions = options?.apiVersion;
    } else {
      apiVersions = [options?.apiVersion];
    }
  }
  if (apiVersions.includes("v1beta")) {
    Object.keys(SUPPORTED_V15_MODELS).forEach(
      (name) => defineGoogleAIModel({
        ai,
        name,
        apiKey: options?.apiKey,
        apiVersion: "v1beta",
        baseUrl: options?.baseUrl,
        debugTraces: options?.experimental_debugTraces
      })
    );
  }
  if (apiVersions.includes("v1")) {
    Object.keys(SUPPORTED_V15_MODELS).forEach(
      (name) => defineGoogleAIModel({
        ai,
        name,
        apiKey: options?.apiKey,
        apiVersion: void 0,
        baseUrl: options?.baseUrl,
        debugTraces: options?.experimental_debugTraces
      })
    );
    Object.keys(EMBEDDER_MODELS).forEach(
      (name) => defineGoogleAIEmbedder(ai, name, { apiKey: options?.apiKey })
    );
  }
  if (options?.models) {
    for (const modelOrRef of options?.models) {
      const modelName = typeof modelOrRef === "string" ? modelOrRef : (
        // strip out the `googleai/` prefix
        modelOrRef.name.split("/")[1]
      );
      const modelRef2 = typeof modelOrRef === "string" ? gemini(modelOrRef) : modelOrRef;
      defineGoogleAIModel({
        ai,
        name: modelName,
        apiKey: options?.apiKey,
        baseUrl: options?.baseUrl,
        info: {
          ...modelRef2.info,
          label: `Google AI - ${modelName}`
        },
        debugTraces: options?.experimental_debugTraces
      });
    }
  }
}
async function resolver(ai, actionType, actionName, options) {
  switch (actionType) {
    case "model":
      resolveModel(ai, actionName, options);
      break;
    case "embedder":
      resolveEmbedder(ai, actionName, options);
      break;
    default:
  }
}
function resolveModel(ai, actionName, options) {
  const modelRef2 = gemini(actionName);
  defineGoogleAIModel({
    ai,
    name: modelRef2.name,
    apiKey: options?.apiKey,
    baseUrl: options?.baseUrl,
    info: {
      ...modelRef2.info,
      label: `Google AI - ${actionName}`
    },
    debugTraces: options?.experimental_debugTraces
  });
}
function resolveEmbedder(ai, actionName, options) {
  defineGoogleAIEmbedder(ai, `googleai/${actionName}`, {
    apiKey: options?.apiKey
  });
}
async function listActions(options) {
  const apiKey = options?.apiKey || getApiKeyFromEnvVar();
  if (!apiKey) {
    logger.error(
      "Pass in the API key or set the GEMINI_API_KEY or GOOGLE_API_KEY environment variable."
    );
    return [];
  }
  const models = await listModels(
    options?.baseUrl || "https://generativelanguage.googleapis.com",
    apiKey
  );
  return [
    // Models
    ...models.filter((m) => m.supportedGenerationMethods.includes("generateContent")).filter((m) => !m.description || !m.description.includes("deprecated")).map((m) => {
      const ref = gemini(
        m.name.startsWith("models/") ? m.name.substring("models/".length) : m.name
      );
      return modelActionMetadata({
        name: ref.name,
        info: ref.info,
        configSchema: GeminiConfigSchema
      });
    }),
    // Embedders
    ...models.filter((m) => m.supportedGenerationMethods.includes("embedContent")).filter((m) => !m.description || !m.description.includes("deprecated")).map((m) => {
      const name = "googleai/" + (m.name.startsWith("models/") ? m.name.substring("models/".length) : m.name);
      return embedderActionMetadata({
        name,
        configSchema: GeminiEmbeddingConfigSchema,
        info: {
          dimensions: 768,
          label: `Google Gen AI - ${name}`,
          supports: {
            input: ["text"]
          }
        }
      });
    })
  ];
}
function googleAIPlugin(options) {
  let listActionsCache;
  return genkitPlugin(
    "googleai",
    async (ai) => await initializer(ai, options),
    async (ai, actionType, actionName) => await resolver(ai, actionType, actionName, options),
    async () => {
      if (listActionsCache) return listActionsCache;
      listActionsCache = await listActions(options);
      return listActionsCache;
    }
  );
}
const googleAI = googleAIPlugin;
googleAI.model = (name, config) => {
  return modelRef({
    name: `googleai/${name}`,
    config,
    configSchema: GeminiConfigSchema
  });
};
googleAI.embedder = (name, config) => {
  return embedderRef({
    name: `googleai/${name}`,
    config,
    configSchema: GeminiEmbeddingConfigSchema
  });
};
var src_default = googleAI;
export {
  src_default as default,
  gemini,
  gemini10Pro,
  gemini15Flash,
  gemini15Flash8b,
  gemini15Pro,
  gemini20Flash,
  gemini20FlashExp,
  gemini20FlashLite,
  gemini20ProExp0205,
  gemini25FlashPreview0417,
  gemini25ProExp0325,
  gemini25ProPreview0325,
  googleAI,
  googleAIPlugin,
  textEmbedding004,
  textEmbeddingGecko001
};
//# sourceMappingURL=index.mjs.map