@genkit-ai/ai/lib/embedder.mjs

Genkit AI framework generative AI APIs.

import {
  action,
  defineAction,
  z
} from "@genkit-ai/core";
import { toJsonSchema } from "@genkit-ai/core/schema";
import { Document, DocumentDataSchema } from "./document.js";
const EmbeddingSchema = z.object({
  embedding: z.array(z.number()),
  metadata: z.record(z.string(), z.unknown()).optional()
});
const EmbedRequestSchema = z.object({
  input: z.array(DocumentDataSchema),
  options: z.any().optional()
});
const EmbedResponseSchema = z.object({
  embeddings: z.array(EmbeddingSchema)
  // TODO: stats, etc.
});
function withMetadata(embedder2, configSchema) {
  const withMeta = embedder2;
  withMeta.__configSchema = configSchema;
  return withMeta;
}
function embedder(options, runner) {
  const embedder2 = action(
    embedderActionOptions(options),
    (i, opts) => runner(
      {
        input: i.input.map((dd) => new Document(dd)),
        options: i.options
      },
      opts
    )
  );
  const ewm = withMetadata(
    embedder2,
    options.configSchema
  );
  return ewm;
}
function embedderActionOptions(options) {
  return {
    actionType: "embedder",
    name: options.name,
    inputSchema: EmbedRequestSchema,
    outputSchema: EmbedResponseSchema,
    metadata: {
      type: "embedder",
      info: options.info,
      embedder: {
        customOptions: options.configSchema ? toJsonSchema({ schema: options.configSchema }) : void 0
      }
    }
  };
}
function defineEmbedder(registry, options, runner) {
  const embedder2 = defineAction(
    registry,
    embedderActionOptions(options),
    (i) => runner(
      i.input.map((dd) => new Document(dd)),
      i.options
    )
  );
  const ewm = withMetadata(
    embedder2,
    options.configSchema
  );
  return ewm;
}
async function embed(registry, params) {
  const embedder2 = await resolveEmbedder(registry, params);
  if (!embedder2.embedderAction) {
    let embedderId;
    if (typeof params.embedder === "string") {
      embedderId = params.embedder;
    } else if (params.embedder?.__action?.name) {
      embedderId = params.embedder.__action.name;
    } else {
      embedderId = params.embedder.name;
    }
    throw new Error(`Unable to resolve embedder ${embedderId}`);
  }
  const response = await embedder2.embedderAction({
    input: typeof params.content === "string" ? [Document.fromText(params.content, params.metadata)] : [params.content],
    options: {
      version: embedder2.version,
      ...embedder2.config,
      ...params.options
    }
  });
  return response.embeddings;
}
async function resolveEmbedder(registry, params) {
  if (typeof params.embedder === "string") {
    return {
      embedderAction: await registry.lookupAction(
        `/embedder/${params.embedder}`
      )
    };
  } else if (Object.hasOwnProperty.call(params.embedder, "__action")) {
    return {
      embedderAction: params.embedder
    };
  } else if (Object.hasOwnProperty.call(params.embedder, "name")) {
    const ref = params.embedder;
    return {
      embedderAction: await registry.lookupAction(
        `/embedder/${params.embedder.name}`
      ),
      config: {
        ...ref.config
      },
      version: ref.version
    };
  }
  throw new Error(`failed to resolve embedder ${params.embedder}`);
}
async function embedMany(registry, params) {
  let embedder2;
  if (typeof params.embedder === "string") {
    embedder2 = await registry.lookupAction(`/embedder/${params.embedder}`);
  } else if (Object.hasOwnProperty.call(params.embedder, "info")) {
    embedder2 = await registry.lookupAction(
      `/embedder/${params.embedder.name}`
    );
  } else {
    embedder2 = params.embedder;
  }
  if (!embedder2) {
    throw new Error("Unable to utilize the provided embedder");
  }
  const response = await embedder2({
    input: params.content.map(
      (i) => typeof i === "string" ? Document.fromText(i, params.metadata) : i
    ),
    options: params.options
  });
  return response.embeddings;
}
const EmbedderInfoSchema = z.object({
  /** Friendly label for this model (e.g. "Google AI - Gemini Pro") */
  label: z.string().optional(),
  /** Supported model capabilities. */
  supports: z.object({
    /** Model can input this type of data. */
    input: z.array(z.enum(["text", "image", "video"])).optional(),
    /** Model can support multiple languages */
    multilingual: z.boolean().optional()
  }).optional(),
  /** Embedding dimension */
  dimensions: z.number().optional()
});
function embedderRef(options) {
  let name = options.name;
  if (options.namespace && !name.startsWith(options.namespace + "/")) {
    name = `${options.namespace}/${name}`;
  }
  return { ...options, name };
}
function embedderActionMetadata({
  name,
  info,
  configSchema
}) {
  return {
    actionType: "embedder",
    name,
    inputJsonSchema: toJsonSchema({ schema: EmbedRequestSchema }),
    outputJsonSchema: toJsonSchema({ schema: EmbedResponseSchema }),
    metadata: {
      embedder: {
        ...info,
        customOptions: configSchema ? toJsonSchema({ schema: configSchema }) : void 0
      }
    }
  };
}
export {
  EmbedderInfoSchema,
  EmbeddingSchema,
  defineEmbedder,
  embed,
  embedMany,
  embedder,
  embedderActionMetadata,
  embedderRef
};
//# sourceMappingURL=embedder.mjs.map