TypeScript Embedding Adapters

Copyexport interface Embedder {
  embed(
    text: string,
    memoryAction?: "add" | "update" | "search",
  ): Promise<number[]>;
  embedBatch(
    texts: string[],
    memoryAction?: "add" | "update" | "search",
  ): Promise<number[][]>;
}

Copyexport * from "./embeddings/base";
export * from "./embeddings/aws_bedrock";
export * from "./embeddings/huggingface";
export * from "./embeddings/openai";
export * from "./embeddings/ollama";
export * from "./embeddings/lmstudio";
export * from "./embeddings/together";
export * from "./embeddings/google";
export * from "./embeddings/azure";
export * from "./embeddings/langchain";
export * from "./embeddings/vertexai";
export * from "./embeddings/fastembed";

Copyexport class OpenAIEmbedder implements Embedder {
  private openai: OpenAI;
  private model: string;
  private embeddingDims: number | undefined;
 
  constructor(config: EmbeddingConfig) {
    this.openai = new OpenAI({
      apiKey: config.apiKey,
      baseURL: config.baseURL || config.url,
    });
    this.model = config.model || "text-embedding-3-small";
    this.embeddingDims = config.embeddingDims;
  }
 
  async embed(text: string): Promise<number[]> {
    const response = await this.openai.embeddings.create({
      model: this.model,
      input: text,
      encoding_format: "float",
      ...(this.embeddingDims !== undefined && {
        dimensions: this.embeddingDims,
      }),
    });
    return response.data[0].embedding;
  }
 
  async embedBatch(texts: string[]): Promise<number[][]> {
    const MAX_BATCH = 100;
    const allEmbeddings: number[][] = [];
    for (let i = 0; i < texts.length; i += MAX_BATCH) {
      const chunk = texts.slice(i, i + MAX_BATCH);
      const response = await this.openai.embeddings.create({
        model: this.model,
        input: chunk,
        encoding_format: "float",
        ...(this.embeddingDims !== undefined && {
          dimensions: this.embeddingDims,
        }),
      });
      allEmbeddings.push(
        ...response.data
          .sort((a, b) => a.index - b.index)
          .map((item) => item.embedding),
      );
    }
    if (allEmbeddings.length !== texts.length) {
      throw new Error(
        `OpenAI embedBatch() returned ${allEmbeddings.length} embeddings for ${texts.length} texts using model '${this.model}'`,
      );
    }
    return allEmbeddings;
  }

Copyexport class AzureOpenAIEmbedder implements Embedder {
  private client: AzureOpenAI;
  private model: string;
  private embeddingDims: number | undefined;
 
  constructor(config: EmbeddingConfig) {
    if (!config.apiKey || !config.modelProperties?.endpoint) {
      throw new Error("Azure OpenAI requires both API key and endpoint");
    }
 
    const { endpoint, ...rest } = config.modelProperties;
 
    this.client = new AzureOpenAI({
      apiKey: config.apiKey,
      endpoint: endpoint as string,
      ...rest,
    });
    this.model = config.model || "text-embedding-3-small";
    this.embeddingDims = config.embeddingDims;
  }

Copyexport class TogetherEmbedder extends OpenAIEmbedder {
  constructor(config: EmbeddingConfig) {
    const openAICompatibleConfig = { ...config };
    delete openAICompatibleConfig.embeddingDims;
 
    const apiKey = config.apiKey || process.env.TOGETHER_API_KEY;
    if (!apiKey) {
      throw new Error("Together API key is required");
    }
 
    super({
      ...openAICompatibleConfig,
      apiKey,
      // Honor TOGETHER_API_BASE like the Together LLM does, so a user behind a
      // gateway who sets it gets it for embeddings too instead of silently
      // hitting api.together.ai.
      baseURL:
        config.baseURL ||
        config.url ||
        process.env.TOGETHER_API_BASE ||
        DEFAULT_BASE_URL,
      model: config.model || DEFAULT_MODEL,
    });
  }

Copyexport class HuggingFaceEmbedder implements Embedder {
  private openai: OpenAI;
  private model: string;
 
  constructor(config: EmbeddingConfig) {
    const baseURL =
      config.huggingfaceBaseUrl ||
      config.baseURL ||
      config.url ||
      process.env.HUGGINGFACE_BASE_URL;
 
    if (!baseURL) {
      throw new Error(
        "HuggingFace embedder requires an inference endpoint. Set " +
          "`huggingfaceBaseUrl` (or `baseURL`) in the embedder config, or the " +
          "HUGGINGFACE_BASE_URL environment variable (e.g. a TEI server at " +
          "http://localhost:8080/v1).",
      );
    }
 
    this.openai = new OpenAI({
      apiKey: config.apiKey || process.env.HUGGINGFACE_API_KEY || "hf",
      baseURL,
    });
    // TEI ignores the model field; default mirrors the Python provider.
    this.model = config.model || "tei";
  }

Copyexport class LMStudioEmbedder implements Embedder {
  private openai: OpenAI;
  private model: string;
 
  constructor(config: EmbeddingConfig) {
    const baseURL = config.baseURL ?? config.url ?? DEFAULT_BASE_URL;
    const apiKey = config.apiKey || DEFAULT_LMSTUDIO_API_KEY;
    this.openai = new OpenAI({ apiKey, baseURL: String(baseURL) });
    this.model = config.model || DEFAULT_MODEL;
  }

Copy * Supports the Amazon Titan and Cohere embedding model families. The
 * `@aws-sdk/client-bedrock-runtime` dependency is lazily imported so the
 * package stays optional: importing this module never forces the SDK to be
 * installed until a Bedrock embedder actually embeds something.
 */
export class AWSBedrockEmbedder implements Embedder {
  private readonly model: string;
  private readonly region: string;
  private readonly embeddingDims?: number;
  private readonly credentials?: BedrockCredentials;
  private clientPromise?: Promise<{
    sdk: BedrockRuntimeModule;
    client: { send: (command: any) => Promise<{ body?: Uint8Array }> };
  }>;
 
  constructor(config: EmbeddingConfig) {
    this.model = config.model || DEFAULT_MODEL;
    this.region = config.awsRegion || process.env.AWS_REGION || DEFAULT_REGION;
    this.embeddingDims = config.embeddingDims;
 
    const hasKeyPair = Boolean(
      config.awsAccessKeyId && config.awsSecretAccessKey,
    );
    const hasAnyCredential = Boolean(
      config.awsAccessKeyId ||
      config.awsSecretAccessKey ||
      config.awsSessionToken,
    );
 
    // Partially configured credentials would silently fall back to the default
    // chain, embedding under an identity the caller never chose.
    if (hasAnyCredential && !hasKeyPair) {
      throw new Error(
        "AWS Bedrock requires both awsAccessKeyId and awsSecretAccessKey when any explicit credential is configured. " +
          "Omit all credential fields to use the AWS default credential chain.",
      );
    }
 
    // Leaving `credentials` unset lets the AWS SDK resolve them from its
    // default chain: environment, shared config, SSO, or the instance role.
    if (hasKeyPair) {
      this.credentials = {
        accessKeyId: config.awsAccessKeyId!,
        secretAccessKey: config.awsSecretAccessKey!,
        ...(config.awsSessionToken && { sessionToken: config.awsSessionToken }),
      };
    }
  }

Copyexport class GoogleEmbedder implements Embedder {
  private google!: GoogleGenAI;
  private model: string;
  private embeddingDims: number | undefined;
  private readonly apiKey: string | undefined;
 
  constructor(config: EmbeddingConfig) {
    this.apiKey = config.apiKey || process.env.GOOGLE_API_KEY;
    this.model = config.model || "gemini-embedding-001";
    this.embeddingDims = config.embeddingDims;
  }
 
  private async ensureClient(): Promise<void> {
    if (this.google) return;
    const sdk = await loadPeer(
      "@google/genai",
      "Google embedder",
      () => import("@google/genai"),
    );
    this.google = new sdk.GoogleGenAI({ apiKey: this.apiKey });
  }

Copyasync embed(text: string): Promise<number[]> {
  await this.ensureClient();
  const response = await this.google.models.embedContent({
    model: this.model,
    contents: text,
    ...(this.embeddingDims !== undefined && {
      config: { outputDimensionality: this.embeddingDims },
    }),
  });
  return response.embeddings![0].values!;
}
 
async embedBatch(texts: string[]): Promise<number[][]> {
  await this.ensureClient();
  const response = await this.google.models.embedContent({
    model: this.model,
    contents: texts,
    ...(this.embeddingDims !== undefined && {
      config: { outputDimensionality: this.embeddingDims },
    }),
  });
  return response.embeddings!.map((item) => item.values!);
}

Copyexport class VertexAIEmbedder implements Embedder {
  private client: PredictionServiceClient | undefined;
  private helpers: AIPlatform["helpers"] | undefined;
  private initPromise: Promise<void> | undefined;
  private clientOptions: ClientOptions;
  private model: string;
  private embeddingDims: number;
  private location: string;
  private projectId: string;
  private embeddingTypes: {
    add: string;
    update: string;
    search: string;
  };
 
  constructor(config: VertexAIConfig) {
    this.model = config.model || "gemini-embedding-001";
    this.embeddingDims = config.embeddingDims || 256;
    this.location =
      config.location || process.env.GCP_LOCATION || "us-central1";
 
    // Left empty when unset: initClient() resolves it from Application Default
    // Credentials or the service account key file, the way the Python SDK does.
    this.projectId =
      config.googleProjectId ||
      process.env.GCP_PROJECT_ID ||
      process.env.GOOGLE_CLOUD_PROJECT ||
      process.env.GCLOUD_PROJECT ||
      "";
 
    this.embeddingTypes = {
      add: config.memoryAddEmbeddingType || "RETRIEVAL_DOCUMENT",
      update: config.memoryUpdateEmbeddingType || "RETRIEVAL_DOCUMENT",
      search: config.memorySearchEmbeddingType || "RETRIEVAL_QUERY",
    };
 
    const endpoint = `${this.location}-aiplatform.googleapis.com`;
    this.clientOptions = { apiEndpoint: endpoint };
 
    if (config.vertexCredentialsJson) {
      this.clientOptions.keyFilename = config.vertexCredentialsJson;
    } else if (config.googleServiceAccountJson) {
      try {
        this.clientOptions.credentials =
          typeof config.googleServiceAccountJson === "string"
            ? JSON.parse(config.googleServiceAccountJson)
            : config.googleServiceAccountJson;
      } catch (err) {
        throw new Error(
          "Failed to parse googleServiceAccountJson: " + (err as Error).message,
        );
      }
    }
  }

Copyfunction maxInstancesPerRequest(model: string): number {
  return model.startsWith("gemini-embedding") ? 1 : 250;
}

Copyprivate async ensureClient(): Promise<void> {
  if (this.ollama) return;
  const sdk = await loadPeer(
    "ollama",
    "Ollama embedder",
    () => import("ollama"),
  );
  this.ollama = new sdk.Ollama({ host: this.host });
}

Copyprivate async ensureModelExists(): Promise<boolean> {
  if (this.initialized) {
    return true;
  }
  await this.ensureClient();
  const local_models = await this.ollama.list();
  const target = OllamaEmbedder.normalizeModelName(this.model);
  if (
    !local_models.models.find(
      (m: any) => OllamaEmbedder.normalizeModelName(m.name) === target,
    )
  ) {
    logger.info(`Pulling model ${this.model}...`);
    await this.ollama.pull({ model: this.model });
  }
  this.initialized = true;
  return true;
}

Copyconst SUPPORTED_MODELS = [
  "fast-all-MiniLM-L6-v2",
  "fast-bge-base-en",
  "fast-bge-base-en-v1.5",
  "fast-bge-small-en",
  "fast-bge-small-en-v1.5",
  "fast-bge-small-zh-v1.5",
  "fast-multilingual-e5-large",
] as const;
type FastEmbedModel = (typeof SUPPORTED_MODELS)[number];
const DEFAULT_MODEL: FastEmbedModel = "fast-bge-small-en-v1.5";
 
export class FastEmbedEmbedder implements Embedder {
  private readonly modelName: FastEmbedModel;
  private embeddingModel?: Promise<FlagEmbedding>;
 
  constructor(config: EmbeddingConfig) {
    if (typeof config.model === "string" && config.model.length > 0) {
      if (!SUPPORTED_MODELS.includes(config.model as FastEmbedModel)) {
        throw new Error(
          `Unsupported FastEmbed model "${config.model}". ` +
            `Supported models: ${SUPPORTED_MODELS.join(", ")}.`,
        );
      }
      this.modelName = config.model as FastEmbedModel;
    } else {
      this.modelName = DEFAULT_MODEL;
    }
  }

Copyprivate getEmbeddingModel(): Promise<FlagEmbedding> {
  if (!this.embeddingModel) {
    this.embeddingModel = this.initEmbeddingModel().catch((error) => {
      this.embeddingModel = undefined;
      throw error;
    });
  }
 
  return this.embeddingModel;
}
 
/**
 * Lazily import the optional `fastembed` peer and initialize the model, so
 * consumers that never touch FastEmbed don't need it installed.
 */
private async initEmbeddingModel(): Promise<FlagEmbedding> {
  const sdk = await loadPeer(
    "fastembed",
    "FastEmbed embedder",
    () => import("fastembed"),
  );
 
  return sdk.FlagEmbedding.init({ model: this.modelName });
}

Copyconstructor(config: EmbeddingConfig) {
  // Check if config.model is provided and is an object (the instance)
  if (!config.model || typeof config.model !== "object") {
    throw new Error(
      "Langchain embedder provider requires an initialized Langchain Embeddings instance passed via the 'model' field in the embedder config.",
    );
  }
  // Basic check for embedding methods
  if (
    typeof (config.model as any).embedQuery !== "function" ||
    typeof (config.model as any).embedDocuments !== "function"
  ) {
    throw new Error(
      "Provided Langchain 'instance' in the 'model' field does not appear to be a valid Langchain Embeddings instance (missing embedQuery or embedDocuments method).",
    );
  }
  this.embedderInstance = config.model as Embeddings;

Copy  // Store batch size if the instance has it (optional)
  this.batchSize = (this.embedderInstance as any).batchSize;
}
 
async embed(text: string): Promise<number[]> {
  try {
    // Use embedQuery for single text embedding
    return await this.embedderInstance.embedQuery(text);
  } catch (error) {
    console.error("Error embedding text with Langchain Embedder:", error);
    throw error;
  }
}
 
async embedBatch(texts: string[]): Promise<number[][]> {
  try {
    // Use embedDocuments for batch embedding
    // Langchain's embedDocuments handles batching internally if needed/supported
    return await this.embedderInstance.embedDocuments(texts);
  } catch (error) {
    console.error("Error embedding batch with Langchain Embedder:", error);
    throw error;
  }
}